更新时间:2024-04-30 GMT+08:00
分享

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”,其中${diffuers}表示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包内的相对路径修改为绝对路径的形式。

图1 代码依赖修改前与修改后

将推理代码“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
图2 执行推理脚本
图3 MindSpore Lite pipeline输出的结果图片
分享:

    相关文档

    相关产品