基于寒武纪自研推理引擎magicmind在StableDiffusion当中的应用示例(MLU 370X4)

• 概述

MagicMind是面向寒武纪MLU的推理加速引擎。MagicMind能将AI框架(Tensorflow,PyTorch,ONNX 等)训练好的算法模型转换成MagicMind统一计算图表示,并提供端到端的模型优化、代码生成以及推理业务部署能力。本贴探讨如何使用将StableDiffusion网络的PyTorch模型转换为MagicMind模型，进而部署在寒武纪MLU板卡上。

本例使用的StableDiffusion实现代码来自github开源项目https://github.com/Stability-AI/stablediffusion.git，StableDiffuison模型包含三个主要模块,UNet,AutoEncoderKL和CLIP der，目前仓库支持选择将UNet或者AutoEncoderKL使用MagicMind引擎来推理(暂不支持CLIPEmbedder)，其余模块则使用寒武纪PyTorch来推理。通过use_mm_net,use_mm_ae,use_mm_clip控制是(true)否(false)使用。

MLU代码实现:https://gitee.com/xiaoqi25478/stablediffusion_mlu,代码仓库内部README.md提供了详细的使用步骤，大家可参考该README.md来部署应用，接下来简单介绍一下使用流程。

• 环境准备

示例基于MagicMind推理引擎部署，为方便快速部署和体验该示例，用户可联系寒武纪技术人员获取MagicMind Docker最新镜像包。后续操作步骤默认在Docker内部运行。

• 下载仓库

git clone https://gitee.com/xiaoqi25478/stablediffusion_mlu.git
cd stablediffusion_mlu
source env.sh

• 准备数据集，模型

cd ${PROJ_ROOT_PATH}/export_model
use_mm_unet=true
use_mm_ae=true
use_mm_clip=false
bash run.sh ${use_mm_unet} ${use_mm_ae} ${use_mm_clip}

• 编译MagicMind模型

precision=force_float16
dynamic_shape=false
batch_size=4
use_mm_unet=true
use_mm_ae=true
use_mm_clip=false
magicmind_model=${MODEL_PATH}/stablediffusion_pytorch_model_${precision}_${dynamic_shape}
cd ${PROJ_ROOT_PATH}/gen_model
bash run.sh ${magicmind_model} ${precision} ${batch_size} ${dynamic_shape} ${use_mm_unet} ${use_mm_ae} ${use_mm_clip}

• 执行推理

#代码将读取export_model/stablediffusion/prompt.txt当中的若干条prompt进行推理，生成图片路径存放于data/output/下面
cd ${PROJ_ROOT_PATH}/infer_python
use_mm_unet=true
use_mm_ae=true
use_mm_clip=false
magicmind_model=${MODEL_PATH}/stablediffusion_pytorch_model_${precision}_${dynamic_shape}
bash run.sh  ${magicmind_model} ${batch_size} ${use_mm_unet} ${use_mm_ae} ${use_mm_clip} 4 128

• 图片展示

prompt "protein"

prompt "a capybara"

prompt "a panda"