pipeline代码适配
onnx pipeline的主要作用是将onnx模型进行一系列编排，并在onnx Runtime上按照编排顺序执行。因此，需要将转换得到的mindir模型按照相同的逻辑进行编排，并在MindSpore Lite上执行。只需要将原始onnx的pipeline中涉及到onnx模型初始化及推理的接口替换为MindSpore Lite的接口即可。
MindSpore Lite提供了Python、C++以及JAVA三种应用开发接口，此处以Python接口为例，介绍如何使用MindSpore Lite Python API构建并推理Stable Diffusion模型，更多信息请参考MindSpore Lite应用开发。
以官方onnx pipeline代码为例，其提供的onnx pipeline代码路径在“${diffusers}/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py”，其中${diffusers}表示diffusers包的安装路径，可以通过pip进行查看。
# shell pip show diffusers
修改代码依赖
新建并进入/home_host/work/pipeline目录。
mkdir -p /home_host/work/pipeline cd /home_host/work/pipeline
将onnx pipeline依赖的图生图源码“pipeline_onnx_stable_diffusion_img2img.py”复制到该目录下，名称改为“pipeline_onnx_stable_diffusion_img2img_mslite.py”，以便与源文件名称区分。但是这样也会导致无法正确找到源码中相对路径下的依赖，需要将对于diffusers包内的相对路径修改为绝对路径的形式。
将推理代码“modelarts-ascend/examples/AIGC/stable_diffusion/onnx_pipeline.py”也复制一份到该目录，名称改为“mslite_pipeline.py”，迁移后的推理代码中的pipeline需要修改为从复制的onnx pipeline文件导入：
# onnx_pipeline.py from pipeline_onnx_stable_diffusion_img2img_mslite import OnnxStableDiffusionImg2ImgPipeline
模型初始化
使用MindSpore Lite进行推理时一般需要先设置目标设备的上下文信息，然后构建推理模型，获取输入数据，模型预测并得到最终的结果。一个基础的推理框架写法如下所示：
# base_mslite_demo.py
import mindspore_lite as mslite
# 设置目标设备上下文为Ascend，指定device_id为0
context = mslite.Context()
context.target = ["ascend"]
context.ascend.device_id = 0
# 构建模型
model = mslite.Model()
model.build_from_file("./model.mindir", mslite.ModelType.MINDIR, context)
# 输入数据到Device侧，针对于多输入场景可以通过list来指定输入
in_data = [np.array(data1), np.array(data2)]
inputs = model.get_inputs()
for i, _inputs in enumerate(inputs):
_input.set_data_from_numpy(in_data[i])
# 前向推理，并将结果从device侧传到host侧
outputs = model.predict(inputs)
outputs = [output.get_data_to_numpy() for output in outputs]
# 后处理...
为了同时兼容onnx模型和mindir模型都能够在适配后的pipeline中运行，需要对于Model进行封装，MsliteModel各参数模型说明已给出，根据模型初始化参数设置当前模型使用onnx模型（运行在CPU上）或mindir模型（运行在昇腾设备上），也能够方便进行精度的校验。
# mslite_model_proxy.py
import os
import mindspore_lite as mslite
class MsliteModel:
def __init__(self, model_path, model_name='ms model', device_type='ascend', use_ascend=True,
onnx_runtime_model=None, get_shape=False, resize_shape=False) -> None:
"""
mindir模型代理类
Args:
model_path: mindir文件路径
model_name: 模型名称
device_type: 设备类型
use_ascend: 是否使用Ascend
onnx_runtime_model: onnx模型对象
get_shape: 是否获取模型shape信息、输入数据shape信息
resize_shape: resize shape开关，分档模型需开启
"""
print('model_path:{}'.format(model_path))
self.model_name = model_name
self.context = MsliteModel.init_context(device_type)
self.model = mslite.Model()
self.model.build_from_file(model_path, mslite.ModelType.MINDIR, self.context)
self.ms_inputs = self.model.get_inputs()
self.use_ascend = use_ascend
self.onnx_runtime_model = onnx_runtime_model
self.get_shape = get_shape
self.resize_shape = resize_shape
@staticmethod
def init_context(device_type='ascend'):
context = mslite.Context()
context.target = [device_type]
context.ascend.device_id = int(os.getenv('DEVICE_ID') or 0)
context.cpu.thread_num = 1 if device_type == 'ascend' else 32
context.cpu.thread_affinity_mode = 2
return context
def __call__(self, **kwargs):
if not self.use_ascend:
return self.onnx_runtime_model(**kwargs)
inputs = list(kwargs.values())
if len(inputs) <= 0:
raise Exception('get tensor input info failed')
ms_input = self.model.get_inputs()
if self.get_shape:
print(f'{self.model_name} shape info:')
for index, val in enumerate(self.model.get_inputs()):
print(f'{self.model_name}: param{index} shape -> {val.shape}')
shapes = [list(input.shape) for input in inputs]
print(f"inputs: input_shape -> {shapes}")
if self.resize_shape:
self.model.resize(ms_input, [list(input.shape) for input in inputs])
for index, val in enumerate(ms_input):
val.set_data_from_numpy(inputs[index])
outputs = self.model.predict(ms_input)
outputs = [output.get_data_to_numpy() for output in outputs]
return outputs
适配MindSpore Lite Runtime到onnx pipeline，首先需要初始化MindSpore LiteModel对象，通过在OnnxStableDiffusionImg2ImgPipeline中增加mindir模型初始化函数，然后在pipeline类的__init__方法调用该函数，在pipeline初始化的时候直接初始化模型。可以参照如下样例，可以通过修改use_ascend去修改该模型是否使用mindir运行，也可以编写代码通过环境变量指定。
# pipeline_onnx_stable_diffusion_img2img_mslite.py
class OnnxStableDiffusionImg2ImgPipeline(DiffusionPipeline):
...
def mslite_modules_init(self):
self.text_encoder_ms = MsliteModel(
model_path=os.environ['TEXT_ENCODER_PATH'],
model_name='text_encoder', use_ascend=True,
onnx_runtime_model=self.text_encoder)
self.vae_encoder_ms = MsliteModel(
model_path=os.environ['VAE_ENCODER_PATH'],
model_name='vae_encoder', use_ascend=True,
onnx_runtime_model=self.vae_encoder)
self.unet_ms = MsliteModel(model_path=os.environ['UNET_PATH'],
model_name='unet', use_ascend=True,
onnx_runtime_model=self.unet)
self.vae_decoder_ms = MsliteModel(model_path=os.environ['VAE_DECODER_PATH'],
model_name='vae_decoder', use_ascend=True,
onnx_runtime_model=self.vae_decoder)
self.safety_checker_ms = MsliteModel(model_path=os.environ['SAFETY_CHECKER_PATH'],
model_name='safety_checker', use_ascend=True,
onnx_runtime_model=self.safety_checker)
...
模型推理适配
完成模型初始化后，需要将onnx模型推理的代码等价替换为对应的mindir模型推理接口。以vae_encoder模型为例，在pipeline代码中查找vae_encoder推理调用的地方，然后修改为对应的MindSpore Lite版本的推理接口模型。
- 使用MindSpore Lite Runtime接口替换onnx Runtime接口
# pipeline_onnx_stable_diffusion_img2img_mslite.py … # onnx模型 # init_latents = self.vae_encoder(sample=image)[0] # ----------------修改点----------------- # mslite模型 init_latents = self.vae_encoder_ms(sample=image)[0] ...
- 替换内嵌模型
# pipeline_onnx_stable_diffusion_img2img_mslite.py … # onnx模型 # image = np.concatenate([self.vae_decoder(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])]) # ----------------修改点----------------- # mslite模型 image = np.concatenate([self.vae_decoder_ms(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])]) ...
修改后的文件参考Gitee代码库中的如下两个文件：
- pipeline_onnx_stable_diffusion_img2img_mslite.py
- mslite_model_proxy.py
运行pipeline代码
pipeline代码如下：
# mslite_pipeline.py
import os
import requests
import torch
import numpy as np
from PIL import Image
from io import BytesIO
from pipeline_onnx_stable_diffusion_img2img_mslite import OnnxStableDiffusionImg2ImgPipeline
def setup_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
torch.backends.cudnn.deterministic = True
setup_seed(0)
# 指定mindir和onnx模型路径
mindir_dir = "/home_host/work/static_shape_convert/mindir_models"
onnx_model_path = "/home_host/work/runwayml/onnx_models"
os.environ['DEVICE_ID'] = "0"
os.environ['TEXT_ENCODER_PATH'] = f"{mindir_dir}/text_encoder.mindir"
os.environ['VAE_ENCODER_PATH'] = f"{mindir_dir}/vae_encoder.mindir"
os.environ['UNET_PATH'] = f"{mindir_dir}/unet_graph.mindir"
os.environ['VAE_DECODER_PATH'] = f"{mindir_dir}/vae_decoder.mindir"
os.environ['SAFETY_CHECKER_PATH'] = f"{mindir_dir}/safety_checker.mindir"
pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(onnx_model_path,
torch_dtype=torch.float32).to("cpu")
url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
response = requests.get(url, verify=False)
init_image = Image.open(BytesIO(response.content)).convert("RGB")
init_image = init_image.resize((512, 512))
prompt = "A fantasy landscape, trending on artstation"
images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images
images[0].save("fantasy_landscape_npu.png")
在运行pipeline时，默认的加速卡为0号卡，当机器有多人使用时，可能存在资源占用而无法正常运行的情况，可以通过环境变量指定加速卡ID，如指定5号卡进行执行。
# mslite_pipeline.py … os.environ['DEVICE_ID'] = "5" …
最后执行python脚本进行推理：
#shell python mslite_pipeline.py
