新闻 📣
- 2024年1月14日:添加了LLM聊天应用(TinyLlama 1.1B和Mistral 7B),初始支持GPU!更多信息 这里。
- 2023年12月14日:添加了对Stable Diffusion XL Turbo 1.0的支持!(感谢@AeroX2)
- 2023年10月3日:添加了对Stable Diffusion XL 1.0 Base的支持!
目录 👇
- 简介
- Stable Diffusion 1.5
- Stable Diffusion XL 1.0 Base
- Stable Diffusion XL Turbo 1.0
- TinyLlama 1.1B和Mistral 7B
- OnnxStream的特性
- 性能
- 注意力切片和量化
- OnnxStream的工作原理
- 构建指南
- 如何转换SD 1.5模型
- 相关项目
- 鸣谢
OnnxStream
挑战是在树莓派Zero 2上运行Stable Diffusion 1.5,该设备是一台拥有512MB RAM的微型电脑,且不增加交换空间并且不将中间结果卸载到磁盘上。Stable Diffusion 1.5的推荐最小RAM/VRAM通常是8GB。
通常,主要的机器学习框架和库都专注于最小化推理延迟和/或最大化吞吐量,而这些通常以RAM使用为代价。因此,我决定编写一个超级小且可调试的推理库,专注于最小化内存消耗:OnnxStream。
OnnxStream是基于将推理引擎与负责提供模型权重的组件解耦的理念,这是一种从WeightsProvider
派生的类。一个WeightsProvider
专门化可以实现任何类型的加载、缓存和预取模型参数。例如,一个自定义的WeightsProvider
可以直接从HTTP服务器下载数据,而无需将任何东西加载或写入到磁盘(这也是“Stream”在“OnnxStream”中的含义)。有三种默认的WeightsProvider
:DiskNoCache
,DiskPrefetch
和Ram
。
OnnxStream可以比OnnxRuntime消耗少55倍的内存,只增加50%到200%的延迟(在CPU上,具有良好的SSD,参考SD 1.5的UNET - 见下文的性能部分)。
Stable Diffusion 1.5
这些图像是由这个仓库中的Stable Diffusion示例实现使用OnnxStream生成的,使用不同精度的VAE解码器。VAE解码器是Stable Diffusion 1.5中唯一无法以单精度或半精度适应树莓派Zero 2的RAM的模型。这是由于模型中存在残差连接以及非常大的张量和卷积。唯一的解决方案是静态量化(8位)。第三张图像是由我的RPI Zero 2在约3小时 1.5小时内生成的(编译时使用MAX_SPEED选项)。第一张图像是在我的PC上使用相同的RPI Zero 2生成的潜在空间进行比较的:
W16A16精度的VAE解码器:
W8A32精度的VAE解码器:
W8A8精度的VAE解码器,由我的RPI Zero 2在约3小时 1.5小时内生成(编译时使用MAX_SPEED选项):
Stable Diffusion XL 1.0(基础版)
现在,OnnxStream的Stable Diffusion示例实现支持SDXL 1.0(不含Refiner)。ONNX文件是从Hugging Face的Diffusers库(版本0.19.3)的SDXL 1.0实现中导出的。
SDXL 1.0比SD 1.5的计算开销显著增加。最显著的区别是在生成1024x1024图像而不是512x512图像。举个例子,使用HF的Diffusers生成10步图像需要在我的12核心、32GB RAM的PC上花费26分钟。SDXL的最低推荐VRAM通常是12GB。
OnnxStream可以在少于300MB的RAM中运行SDXL 1.0,因此能够在不增加交换空间且在推理期间不向磁盘写入任何数据的情况下,在RPI Zero 2上舒适运行。生成10步图像在我的RPI Zero 2上大约需要11个小时。
SDXL专用优化
对SD 1.5所做的一组优化也同样适用于SDXL 1.0,但存在以下差异。
对于UNET模型,为了在RPI Zero 2中以少于300MB的RAM运行,使用了UINT8动态量化,但仅限于特定的大型中间张量子集。
对于VAE解码器,SDXL 1.0的情况比SD 1.5更加复杂。SDXL 1.0的VAE解码器是SD 1.5的4倍大,当使用OnnxStream以FP32精度运行时,消耗4.4GB的RAM。
在SD 1.5的情况下,VAE解码器被静态量化(UINT8精度),这足以将内存消耗减少到260MB。而SDXL 1.0的VAE解码器在FP16运算下会溢出,并且其激活的数值范围太大,无法通过UINT8量化获得高质量的图像。
因此,我们面临一个消耗4.4GB RAM的模型,该模型不能在FP16精度下运行,也不能在UINT8精度下量化。(注意:至少有一个解决方案来解决FP16问题,但我没有进一步研究,因为即使在FP16精度下运行VAE解码器,总内存消耗将减少一半,所以模型最终消耗2.2GB而不是4.4GB,这对于RPI Zero 2来说仍然太多)
解决方案的灵感来自实施Hugging Face的Diffusers库的VAE解码器,即使用平铺解码。最终结果完全无法分辨出与完整解码器解码的图像有任何不同,这样可以将RAM内存消耗从4.4GB减少到298MB!
想法很简单。扩散过程的结果是形状为(1,4,128,128)的张量。想法是将这个张量分割成5x5(因此是25个)交叠的形状为(1,4,32,32)的张量,并分别解码这些张量。每个张量与其左侧和上方的平铺重叠25%。解码结果是一个形状为(1,3,256,256)的张量,然后将其适当地融合到最终图像中。
例如,这是在代码中手动关闭融合的情况下由平铺解码器生成的图像。你可以明显看到图像中的平铺:
而这是相同的图像,开启融合后的效果。这是最终结果:
这是一张由我的RPI Zero 2在约11小时内生成的图像:(10步,Euler Ancestral)
Stable Diffusion XL Turbo 1.0
SDXL Turbo的支持由善良的@AeroX2贡献。
SDXL和SDXL Turbo的主要区别在于Turbo版本生成512x512的图像而不是1024x1024,但所需的步数要少得多。即使只用一步也能生成出好质量的图像!
无需额外的优化即可在RPI Zero 2上运行SDXL Turbo。SDXL和SDXL Turbo共享相同的文本编码器和VAE解码器:需要使用平铺解码以将内存消耗控制在300MB以下。
这张图像是由我的树莓派Zero 2在29分钟(1步)生成的:
这张图像是3步生成的示例,耗时50分钟我的RPI Zero 2生成。质量和一步生成的图像相同:
在模型卡中有关于OnnxStream上运行SDXL 1.0和SDXL Turbo的步数比较(由@AeroX2提供)。
OnnxStream的特性
- 推理引擎与
WeightsProvider
解耦 WeightsProvider
可以是DiskNoCache
、DiskPrefetch
、Ram
或自定义- 注意力切片
- 动态量化(8 位无符号,不对称,百分位)
- 静态量化(W8A8 无符号,不对称,百分位)
- 简单校准量化模型
- 支持 FP16(无论是否使用 FP16 算术)
- 实现了 39 个 ONNX 操作符(最常见的)
- 操作按顺序执行,但
全部大多数操作符是多线程的 - 单个实现文件 + 头文件
- XNNPACK 调用包装在
XnnPack
类中(未来替换用) - 使用 cuBLAS 初始 GPU 支持(仅限 FP16 和 FP32,且仅用于 LLM 应用程序)
OnnxStream 依赖于 XNNPACK 提供一些(加速的)基础操作:矩阵乘法,卷积,逐元素加/减/乘/除,Sigmoid 和 Softmax。
性能
Stable Diffusion 1.5 由三个模型组成:文本编码器(672 个操作和 1.23 亿参数)、UNET 模型(2050 个操作和 8.54 亿参数)以及 VAE 解码器(276 个操作和 4900 万参数)。假设批量大小为 1,一个完整图像生成共 10 步(使用 Euler Ancestral 调度器),需要文本编码器运行 2 次,UNET 模型运行 20 次(即 2*10),以及 VAE 解码器运行 1 次。
下表展示了 Stable Diffusion 1.5 的三类模型的不同推理时间和内存消耗(即 Windows 上的 峰值工作集大小
或 Linux 上的 最大驻留集大小
)。
模型 / 库 | 第一次运行 | 第二次运行 | 第三次运行 |
---|---|---|---|
FP16 UNET / OnnxStream | 0.133 GB - 18.2 秒 | 0.133 GB - 18.7 秒 | 0.133 GB - 19.8 秒 |
FP16 UNET / OnnxRuntime | 5.085 GB - 12.8 秒 | 7.353 GB - 7.28 秒 | 7.353 GB - 7.96 秒 |
FP32 文本编码器 / OnnxStream | 0.147 GB - 1.26 秒 | 0.147 GB - 1.19 秒 | 0.147 GB - 1.19 秒 |
FP32 文本编码器 / OnnxRuntime | 0.641 GB - 1.02 秒 | 0.641 GB - 0.06 秒 | 0.641 GB - 0.07 秒 |
FP32 VAE 解码器 / OnnxStream | 1.004 GB - 20.9 秒 | 1.004 GB - 20.6 秒 | 1.004 GB - 21.2 秒 |
FP32 VAE 解码器 / OnnxRuntime | 1.330 GB - 11.2 秒 | 2.026 GB - 10.1 秒 | 2.026 GB - 11.1 秒 |
对于 UNET 模型(以 FP16 精度运行,并在 OnnxStream 中启用 FP16 算术),OnnxStream 的内存消耗比 OnnxRuntime 少 55 倍,但延迟增加 50% 到 200%。
注意:
- OnnxRuntime 的首次运行是热身推理,因为其
InferenceSession
在第一次运行前创建,并在后续运行中重复使用。OnnxStream 没有热身,因其设计为纯即时运行(然而后续运行可从操作系统缓存的权重文件中受益)。 - 目前 OnnxStream 不支持批量大小不等于 1 的输入,OnnxRuntime 可以通过批量大小为 2 来显著加速整个扩散过程,特别是在运行 UNET 模型时。
- 在我的测试中,更改 OnnxRuntime 的
SessionOptions
(如EnableCpuMemArena
和ExecutionMode
)对结果没有显著影响。 - OnnxRuntime 的性能非常接近 NCNN(我评估的另一个框架),无论是在内存消耗还是推理时间方面。若有需要,我会在未来包含 NCNN 基准测试。
- 测试在我的开发机器上进行:Windows Server 2019, 16GB RAM, 8750H CPU (AVX2), 970 EVO Plus SSD, VMWare 上的 8 个虚拟核。
注意力切片和量化
在运行 UNET 模型时使用“注意力切片”和在运行 VAE 解码器时使用 W8A8 量化是减少内存消耗到可以在 RPI Zero 2 上执行的关键。
虽然互联网上关于量化神经网络的信息很多,但关于“注意力切片”的资料却很少。其理念很简单:目标是在计算 UNET 模型中的多头注意机制的缩放点积注意时,避免物化整个 Q @ K^T
矩阵。UNET 模型中有 8 个注意头,Q
的形状为 (8,4096,40),而 K^T
的形状为 (8,40,4096):因此第一个矩阵乘法的结果形状为 (8,4096,4096),这是一个 512MB 的张量(以 FP32 精度计算):
解决方案是将 Q
垂直拆分,然后对每个 Q
的切片按常规进行注意操作。Q_sliced
形状为 (1,x,40),其中 x 为 4096(在本例中)除以 onnxstream::Model::m_attention_fused_ops_parts
(默认为 2,但可自定义)。这一简单技巧使 UNET 模型在 FP32 精度下的整体内存消耗从 1.1GB 降低到 300MB。可能的替代方案,当然更有效率,是使用 FlashAttention,然而 FlashAttention 需要为每个支持的架构(AVX, NEON 等)编写自定义内核,在我们的案例中则需绕过 XnnPack。
OnnxStream 如何工作
以下代码可运行定义在 path_to_model_folder/model.txt
中的模型(所有模型操作都定义在 model.txt
文件中;OnnxStream 期望在同一文件夹中找到所有权重文件,以 .bin 文件系列的形式存在)
#include "onnxstream.h"
using namespace onnxstream;
int main()
{
Model model;
//
// 可选参数,可以在 Model 对象上设置:
//
// model.set_weights_provider( ... ); // 指定不同的权重提供者(默认是 DiskPrefetchWeightsProvider)
// model.read_range_data( ... ); // 读取范围数据文件(包含量化模型的激活裁剪范围)
// model.write_range_data( ... ); // 写入范围数据文件(校准后有用)
// model.m_range_data_calibrate = true; // 校准模型
// model.m_use_fp16_arithmetic = true; // 推理期间使用 FP16 算术(权重为 FP16 精度时有用)
// model.m_use_uint8_arithmetic = true; // 推理期间使用 UINT8 算术
// model.m_use_uint8_qdq = true; // 使用 UINT8 动态量化(可减少某些模型的内存消耗)
// model.m_fuse_ops_in_attention = true; // 启用注意力切片
// model.m_attention_fused_ops_parts = ... ; // 参见上面的“注意力切片”部分
//
model.read_file("path_to_model_folder/model.txt");
tensor_vector<float> data;
... // 填充 tensor_vector 的张量数据。"tensor_vector" 只是一个带有自定义分配器的 std::vector 的别名。
Tensor t;
t.m_name = "input";
t.m_shape = { 1, 4, 64, 64 };
t.set_vector(std::move(data));
model.push_tensor(std::move(t));
model.run();
auto& result = model.m_data[0].get_vector<float>();
... // 处理结果:"result" 是对推理的第一个结果的引用(也是一个 tensor_vector<float>)。
return 0;
}
model.txt
文件以 ASCII 格式包含了所有模型操作,从原始 ONNX 文件导出。每行代表一个操作:例如这一行表示量化模型中的卷积:
Conv_4:Conv*input:input_2E_1(1,4,64,64);post_5F_quant_5F_conv_2E_weight_nchw.bin(uint8[0.0035054587850383684,134]:4,4,1,1);post_5F_quant_5F_conv_2E_bias.bin(float32:4)*output:input(1,4,64,64)*dilations:1,1;group:1;kernel_shape:1,1;pads:0,0,0,0;strides:1,1
为了从 ONNX 文件导出 model.txt
文件及其权重(以 .bin
文件系列的形式)用于 OnnxStream,提供了一个 Jupyter Notebook(单个单元格)(onnx2txt.ipynb
)。
在将 Pytorch 的 nn.Module
(在我们的案例中)导出为 ONNX 以供 OnnxStream 使用时,必须考虑以下事项:
- 调用
torch.onnx.export
时,dynamic_axes
应为空,因为 OnnxStream 不支持具有动态形状的输入。 - 强烈推荐在导出的 ONNX 文件转换为
model.txt
文件之前,运行优秀的 ONNX Simplifier。
如何在 Linux/Mac/Windows/Termux 上构建 Stable Diffusion 示例
- Windows 特有:启动如下提示命令:
Visual Studio Tools
>x64 Native Tools Command Prompt
。 - Mac 特有:确保安装 cmake:
brew install cmake
。
首先你需要构建 XNNPACK。
由于 XnnPack 的函数原型随时可能改变,我包含了一个 git checkout
,确保 OnnxStream 能与编写时的兼容版本的 XnnPack 正确编译:
git clone https://github.com/google/XNNPACK.git
cd XNNPACK
git checkout 1c8ee1b68f3a3e0847ec3c53c186c5909fa3fbd3
mkdir build
cd build
cmake -DXNNPACK_BUILD_TESTS=OFF -DXNNPACK_BUILD_BENCHMARKS=OFF ..
cmake --build . --config Release
然后你可以构建 Stable Diffusion 示例。
<DIRECTORY_WHERE_XNNPACK_WAS_CLONED>
例如 /home/vito/Desktop/XNNPACK
或 C:\Projects\SD\XNNPACK
(在 Windows 上):
git clone https://github.com/vitoplantamura/OnnxStream.git
cd OnnxStream
cd src
mkdir build
cd build
cmake -DMAX_SPEED=ON -DOS_LLM=OFF -DOS_CUDA=OFF -DXNNPACK_DIR=<DIRECTORY_WHERE_XNNPACK_WAS_CLONED> ..
cmake --build . --config Release
重要:MAX_SPEED 选项允许性能在 Windows 上增加约 10%,但在 Raspberry Pi 上超过 50%。此选项在构建时消耗更多内存,并且生成的可执行文件可能无法运行(如我的 Termux 测试中)。因此若有问题,首要尝试是将 MAX_SPEED 设置为 OFF。
现在你可以运行 Stable Diffusion 示例。
应用程序的最新版本在首次运行时自动下载所选模型的权重。点击此处了解如何手动下载权重。
在 Stable Diffusion 1.5 的情况下,权重可以在此处下载(约 2GB)。
git lfs install
git clone --depth=1 https://huggingface.co/vitoplantamura/stable-diffusion-1.5-onnxstream
对于 Stable Diffusion XL 1.0 Base,可以在这里下载权重(大约8GB):
git lfs install
git clone --depth=1 https://huggingface.co/vitoplantamura/stable-diffusion-xl-base-1.0-onnxstream
对于 Stable Diffusion XL Turbo 1.0,可以在这里下载权重(大约8GB):
git lfs install
git clone --depth=1 https://huggingface.co/AeroX2/stable-diffusion-xl-turbo-1.0-onnxstream
以下是Stable Diffusion示例的命令行选项:
--xl 运行 Stable Diffusion XL 1.0 而不是 Stable Diffusion 1.5。
--turbo 运行 Stable Diffusion Turbo 1.0 而不是 Stable Diffusion 1.5。
--models-path 设置包含 Stable Diffusion 模型的文件夹。
--ops-printf 推理期间,将当前操作写入stdout。
--output 设置输出的PNG文件。
--decode-latents 跳过扩散,并解码指定的潜在文件。
--prompt 设置正向提示。
--neg-prompt 设置负向提示。
--steps 设置扩散步骤数。
--seed 设置种子。
--save-latents 在扩散后,将潜在变量保存在指定文件中。
--decoder-calibrate (仅SD 1.5)校准量化版本的VAE解码器。
--not-tiled (仅SDXL 1.0)不要使用瓦片化VAE解码器。
--ram 使用RAM WeightsProvider(实验性)。
--download 强制(重新)下载当前模型。
--curl-parallel 设置使用CURL的并行下载次数,默认是4。
--rpi 配置模型在树莓派上运行。
--rpi-lowmem 配置模型在树莓派 Zero 2 上运行。
你可能感兴趣的选项:--xl
, --turbo
, --prompt
, --steps
, --rpi
。
如何使用OnnxStream转换和运行自定义Stable Diffusion 1.5模型(由@GaelicThunder撰写)
点击展开
本指南旨在帮助您将自定义的Stable Diffusion模型转换用于OnnxStream。无论您是从 .safetensors
还是 .onnx
开始,该指南都能为您提供帮助。
先决条件
- Python 3.x
- ONNX
- ONNX Simplifier
- Linux环境(在Ubuntu上测试,Windows WSL也适用)
- 交换空间(视您的方法而定)
为什么要特定步骤?
了解Einsum和其他操作
AUTO1111的Stable Diffusion实现使用了例如Einsum之类的操作,这些操作尚未被OnnxStream支持。因此,建议使用Hugging Face的实现,该实现更具兼容性。
可选:将 .safetensors 转换为 ONNX
如果您从 .safetensors
文件开始,可以使用此GitHub仓库中的工具将其转换为 .onnx
。
但是,推荐按照下面“选项A”部分中的方法进行操作。
导出您的模型
选项A:从Hugging Face导出(推荐)
from diffusers import StableDiffusionPipeline
import torch
pipe = StableDiffusionPipeline.from_single_file("https://huggingface.co/YourUsername/YourModel/blob/main/Model.safetensors")
dummy_input = (torch.randn(1, 4, 64, 64), torch.randn(1), torch.randn(1, 77, 768))
input_names = ["sample", "timestep", "encoder_hidden_states"]
output_names = ["out_sample"]
torch.onnx.export(pipe.unet, dummy_input, "/path/to/save/unet_temp.onnx", verbose=False, input_names=input_names, output_names=output_names, opset_version=14, do_constant_folding=True, export_params=True)
选项B:手动修复输入形状
python -m onnxruntime.tools.make_dynamic_shape_fixed --input_name sample --input_shape 1,4,64,64 model.onnx model_fixed1.onnx
python -m onnxruntime.tools.make_dynamic_shape_fixed --input_name timestep --input_shape 1 model_fixed1.onnx model_fixed2.onnx
python -m onnxruntime.tools.make_dynamic_shape_fixed --input_name encoder_hidden_states --input_shape 1,77,768 model_fixed2.onnx model_fixed3.onnx
Vito提示:这可以通过按照上面推荐的“选项A”来实现。修复输入形状可能会很有用,如果您的起点已经是一个ONNX文件。
运行ONNX Simplifier
python -m onnx_simplifier model_fixed3.onnx model_simplified.onnx
使用大模型可能会遇到简化器的问题。在这种情况下,这个工具有时可以提供帮助:https://github.com/luchangli03/onnxsim_large_model
注意:
- 如果您从Hugging Face导出您的模型,您需要大约100GB的交换空间。
- 如果您手动修复了输入形状,16GB的RAM应该就够了。
- 这个过程可能需要一些时间,请耐心等待。
最后的步骤和运行模型
一旦您得到了 onnx2txt
过程中的最终模型,将其移动到OnnxStream的标准SD 1.5模型的 unet_fp16
文件夹中。
运行模型的命令可能是这样的:
./sd --models-path ./Converted/ --prompt "space landscape" --steps 28 --rpi
关于“Shape”操作符的说明
如果在Onnx Simplifier或onnx2txt.ipynb
的输出中看到“Shape”操作符,这表明Onnx Simplifier可能没有正常工作。这个问题通常不是由Onnx Simplifier本身引起的,而是由Onnx的Shape推断引起的。
替代方案
在这种情况下,您可以选择通过修改torch.onnx.export
的参数重新导出模型。在您的计算机上找到这个文件:
并确保:
- 将
opset_version
设置为14 - 删除
dynamic_axes
在这些更改之后,您可以重新运行Onnx Simplifier和onnx2txt
。
Vito提示:尽管这个解决方案有效,但生成带有Einsum操作的ONNX文件。当OnnxStream支持Einsum操作符时,这个解决方案将变得推荐。
结论
这份指南旨在为那些希望使用OnnxStream运行自定义Stable Diffusion 1.5模型的人提供全面的资源。欢迎更多的贡献!
相关项目
- OnnxStreamGui by @ThomAce:OnnxStream的网页和桌面用户界面。
- Auto epaper art by @rvdveen: 显示新闻的自包含图像生成图片框。
致谢