如何使用“LoRa”的方式加载ONNX模型：StableDiffusion相关模型 的C++推理

如何使用“LoRa”的方式加载Onnx模型：StableDiffusion相关模型的C++推理

本文主要干了以下几个事：

1.基于 onnxruntime，将 StableDiffusionInpaintPipeline、StableDiffusionControlNetImg2ImgPipeline(stablediffusion + controlnet + LoRa) C++工程化；

2.输出一个 C++版本的 ddim-schduler 库；

3.提供一种“LoRa”的 onnx 模型加载方式；

4.所有相关代码、模型开源

项目地址: https://github.com/TalkUHulk/ai.deploy.box

模型地址: https://huggingface.co/TalkUHulk/AiDB

StableDiffusionInpaint

模型导出

StableDiffusionInpaint 的 onnx 导出非常简单，optimum 已经做好了集成，支持命令行直接导出，具体参考可参考optimum-cli：

这样得到了四个 onnx 模型（unet、 vae encoder、decoder 和 text encoder）。

tokenizer&scheduler

与检测、分类等传统 cv 方法不同，我们如果想在 c++中串起整个 pipeline，还缺少 c++版本的 tokenizer 和 scheduler。有很多优秀的开源 C++版本的 tokenizer，这里我选用了tokenizers_cpp，地址：https://github.com/mlc-ai/tokenizers-cpp。tokenizers-cpp 接口简单，并且可直接使用 🤗hugging face 中开源的的 tokenizer.json 配置文件。

而对于 scheduler，目前没找到很好用的 c++版本，所以作者实现了一个 C++版本的 ddim_scheduler，并做了开源ddim_scheduler_cpp，rep地址：https://github.com/TalkUHulk/ddim_scheduler_cpp。ddim_scheduler_cpp 底层基于 Eigen 实现，与 diffusers 接口保持一致，可直接替换。

C++推理

目前，我们将所有必须的 C++物料都集齐了。借助作者之前开源的一个开源工具AiDB（rep地址），只需要简单配置，直接可以使用 C++加载并推理 onnx 模型。

至此，我们已经成功搭建起 stablediffusioninpaint 的 C++ pipeline，但更常用、更有趣的是 controlnet 和 lora 与 stablediffusion 的结合。下面我们尝试搭建 StableDiffusionControlNetImg2ImgPipeline 的 C++推理代码，并支持 LoRa 加载。

StableDiffusionControlNetImg2ImgPipeline

模型导出

目前 optimum 还未提供 stablediffusion + controlnet +LoRa 的 onnx 模型导出选项，所以这里我们先将模型导出。 这里我们有两种导出方案，分别导出 controlNet 和 Unet，以及将二者合并为一个模型。

先看一下 controlNet 的整体架构，controlNet 和 Unet 的耦合比较深，如果我们分开导出，两个模型的输出和输入数量都会非常多，比入 Unet 部分有 down_block_res_sample_0 ～ down_block_res_sample_11、mid_block_res_sample 等 16 个输入，这样在写 inference 代码的时候就会比较繁琐。所以我们选择将两个模型合并为一个。但这样也有有另一个问题，比如我首先使用 controlNet-canny 导出了 onnx 模型，同时又想使用 controlNet-hed，那 unet 部分是不是要重复导出？这里有几个方法解决，我们后面再说明。

此处使用Yoji Shinkawa Style LoRA（🤗 https://civitai.com/models/12324/yoji-shinkawa-style-lora)

导出代码：

这里有几个点需要注意。

OP 问题

pytorch2.0 以上，需要做以下设置才可以成功导出

具体可以参考 diffusers->models->attention_processor.py 中的相关代码。Pytorch2.0 以上 scaled dot-product attention 计算会默认使用torch.nn.functional.scaled_dot_product_attention，而 onnx 导出时不支持该 OP。

因此需要做替换，diffusers 很贴心的把相关代码实现好，我们直接使用即可。

模型大小>2GB

ONNX 模型本质就是一个 Protobuf 序列化后的二进制文件，而 Protobuf 的文件大小限制为 2GB。因此对于 Unet 相关模型来说，存储大小已经超过了限制。onnx 给出的方案是单独存储 weights、bias 这些权重。 这里做下详细说明。 先来看下onnx.proto(文件地址：https://github.com/onnx/onnx/blob/main/onnx/onnx.proto)中的定义：

我们可以通过 data_location 来判断某个参数的位置，然后读取 external_data 参数加载权重，接下来我们在代码中手动加载：

摘出其中一个 tensor 的 external_data 详细说明：

location 记录了权重存储的文件名，offset 是该权重在文件中的偏移量，length 是权重的长度。有了以上信息，onnx 内部就可以直接 load 权重，解决 2GB 限制问题。

仔细的同学会观察到，导出的 uent 目录下有，除了.onnx 模型，还有非常非常多的 weight/bias 等文件。这其实就是每一个权重数据。如此碎片化，我们使用或者版本管理起来非常不方便。我们使用以下代码，将所有的权重合并到一个文件中：

这样所有的权重就会保存到一个 model.onnx_data 文件里。

C++推理

与上文类似，借助 AiDB，使用 C++串起整个 pipeline

LoRA 方式加载

回到上文提到的问题，以上例子使用 controlNet-canny 导出 onnx 模型，如果我们又想使用 controlNet-hed，或者使用更多的 LoRa 呢？是否一定必须重新导出整个模型， 是否可以用“LoRa”的方式加载模型呢。答案是肯定的，查看 onnruntime 的接口，官方提供了如下接口:

利用此接口，我们可以实现“LoRa 方式”的模型加载。这里以“LoRa”举例，controlNet 同理。

先做一点简单的知识储备，ONNX 模型本质就是一个 Protobuf 序列化后的二进制文件，所以理论上我们可以做任意合理的修改。根据 onnx.proto 的定义，首先来看一下 onnx 模型的结构。 onnx 主要包含以下几个类：ModelProto，NodeProto，GraphProto，TensorProto 等。ModelProto 作为 top-level 类，用于绑定 ML 模型并将其计算图与元数据相关联。NodeProto 用来描述了 graph 中的 node。TensorProto 则用来组织 tensor 的具体信息。所以 onnx 的结构大概可以用下图表示：

这样我们就有了一个大概的思路，读取 LoRa 模型，解析 LoRa 模型中 tensor，因为网络结构都是相同的，我们直接通过 onnxruntime 的 AddExternalInitializers 接口，来替换原始网络中的 LoRa 部分。

onnx 模型读取

使用 protobuf 读取 onnx 模型，而不是使用 ort：

OP名称

原始模型与 onnx 导出的模型的名字是不一致的，我们需要找到映射关系，才能正确加载。 首先加载 🤗safetensors 格式的模型

此时 state_dict 中的 key 并不是模型 onnx 导出前的 key，这里需要做一个转换。直接参考 diffusers 的代码：

执行以上代码，可以得到 torch.onnx.export 前模型的 key:value。接下来就是和 onnx 模型中的 name 找到对应关系。

其实 onnx 模型中已经储存了对应的对应关系，我们使用以下代码先观察下 onnx 模型中村了什么信息(这里只输出了 lora 相关的):

部分输出：

可以看到每个node的对应关系，格式如下torch-op-name = OP(param, onnx-tensor-name)。按照以上规则，可以找到两种模型opname的映射，将这种关系保存下来：

LoRa保存

最后就是如何组织新的LoRa模型了。这里为了方便，我们构造一个“假的”onnx模型，仅仅存储LoRa的权重，name以上一节映射后为准。

LoRa校验

以上3步已经得到了新的模型，但为了确认我们的方式是否正确，我们拿一个已经导出的Unet模型和对应的LoRa权重做一下校验

确认没问题，我们的准备工作也算完成。下面完成C++代码部分。

LoRa加载

读取新的LoRa模型，将权重的name和raw_data读取出来，然后创建对应的tensor，最后调用session_options.AddExternalInitializers一起初始化即可。需要注意的是，onnxruntime的CreateTensor操作是浅拷贝，所以在写法上注意局部变量的生存周期。

作者在C站找了几个相同结构的LoRa，分别为blindbox、mix4和moxin，测试一下效果

以上代码和模型都已开源，更多详情，敬请登陆github，欢迎Star。