🚀稳定快速

注意

对 stable-fast 的积极开发已暂停。我目前正在基于新的 torch._dynamo 项目，目标是如 stable-cascade、SD3 和 Sora 等新的模型。 它将会更快、更灵活，并且支持更多的硬件后端，而不是 CUDA。

欢迎联系。

Discord 频道

stable-fast 在__所有__类型的扩散器模型上实现了 SOTA 推理性能，即使是最新的 StableVideoDiffusionPipeline。 而且与需要几十分钟编译模型的 TensorRT 或 AITemplate 不同，stable-fast 只需几秒钟即可编译模型。 stable-fast 还支持开箱即用的 动态形状、LoRA 和 ControlNet。

注意: 在基准测试期间，TensorRT 使用 static batch size 和 CUDA Graph enabled 测试，而 stable-fast 运行时使用动态形状。

• 🚀稳定快速
  • 简介
    • 这是什么?
    • 与其他加速库的不同之处
  • 安装
    • 安装预构建包
    • 从源码安装
  • 使用
    • 优化 StableDiffusionPipeline
    • 优化 LCM Pipeline
    • 优化 StableVideoDiffusionPipeline
    • 动态切换 LoRA
    • 模型量化
    • 一些加速 PyTorch 的常用方法
  • 性能比较
    • RTX 4080 (512x512，batch size 1，fp16, in WSL2)
    • H100
    • A100
  • 兼容性
  • 故障排除

简介

这是什么?

stable-fast 是一个超轻量的推理优化框架，针对 HuggingFace Diffusers 在 NVIDIA GPUs 上进行优化。 stable-fast 通过利用一些关键技术和功能提供超快速的推理优化：

• CUDNN 卷积融合: stable-fast 实现了一系列全功能和完全兼容的 CUDNN 卷积融合算子，可用于所有类型的 Conv + Bias + Add + Act 计算模式组合。
• 低精度和融合 GEMM: stable-fast 实现了一系列融合 GEMM 操作符，用 fp16 精度计算，比 PyTorch 的默认设置要快（读取和写入使用 fp16，计算使用 fp32）。
• 融合线性 GEGLU: stable-fast 可以将 GEGLU(x, W, V, b, c) = GELU(xW + b) ⊗ (xV + c) 融合成一个 CUDA 内核。
• NHWC 和融合 GroupNorm: stable-fast 使用 OpenAI 的 Triton 实现了高度优化的融合 NHWC GroupNorm + Silu 操作符，消除了内存格式变换操作的需求。
• 完全跟踪模型: stable-fast 改进了 torch.jit.trace 接口，使其更适用于跟踪复杂模型。几乎所有 StableDiffusionPipeline/StableVideoDiffusionPipeline 的部分都可以被跟踪并转换为 TorchScript。它比 torch.compile 更稳定，并且比 torch.compile 拥有显著更低的 CPU 开销，并支持 ControlNet 和 LoRA。
• CUDA 图: stable-fast 可以将 UNet、VAE 和 TextEncoder 捕获为 CUDA 图格式，从而减少 batch size 较小时的 CPU 开销。该实现还支持动态形状。
• 融合多头注意力: stable-fast 使用 xformers 并使其与 TorchScript 兼容。

我的下一个目标是保持 stable-fast 作为 diffusers 的最快推理优化框架之一，同时 为 transformers 提供速度提升和 VRAM 减少。 实际上，我已经使用 stable-fast 优化了大语言模型 (LLMs) 并实现了显著的速度提升。 但我仍需要一些工作来使其更加稳定和易用，并提供一个稳定的用户接口。

与其他加速库的不同之处

• 快速: stable-fast 专门针对 HuggingFace Diffusers 进行了优化。它在许多库中实现了高性能表现。并且它提供非常快的编译速度，只需几秒钟。它在编译时间上显著比 torch.compile、TensorRT 和 AITemplate 更快。
• 最小化: stable-fast 作为 PyTorch 的插件框架工作。它利用现有的 PyTorch 功能和基础设施，兼容其他加速技术，以及流行的微调技术和部署解决方案。
• 最大兼容性: stable-fast 与各种类型的 HuggingFace Diffusers 和 PyTorch 版本兼容。它还兼容 ControlNet 和 LoRA。甚至开箱即用地支持最新的 StableVideoDiffusionPipeline！

安装

注意: stable-fast 目前仅在 Linux 和 Windows WSL2 上测试。 首先需要安装具有 CUDA 支持的 PyTorch（建议版本从 1.12 到 2.1）。

我只在 torch>=2.1.0、xformers>=0.0.22 和 triton>=2.1.0 上测试了 stable-fast，运行在 CUDA 12.1 和 Python 3.10 上。 其他版本可能也可以构建和运行成功，但不保证。

安装预构建包

从 Releases Page 下载与系统对应的轮子文件，并使用 pip3 install <wheel file> 安装。

目前有 Linux 和 Windows 版的轮子文件。

# 将 cu121 替换为你的 CUDA 版本，并将 <wheel file> 替换为轮子文件的路径。
# 确保轮子文件与你的 PyTorch 版本兼容。
pip3 install --index-url https://download.pytorch.org/whl/cu121 \
    'torch>=2.1.0' 'xformers>=0.0.22' 'triton>=2.1.0' 'diffusers>=0.19.3' \
    '<wheel file>'

从源码安装

# 确保已安装 CUDNN/CUBLAS。
# https://developer.nvidia.com/cudnn
# https://developer.nvidia.com/cublas

# 首先安装具有 CUDA 的 PyTorch 及其他包。
# Windows 用户：可能无法使用 Triton，你可以跳过它。
# 注意：'wheel' 是必须项，否则在构建时会遇到 `No module named 'torch'` 错误。
pip3 install wheel 'torch>=2.1.0' 'xformers>=0.0.22' 'triton>=2.1.0' 'diffusers>=0.19.3'

# （可选）让构建更快。
pip3 install ninja

# 如果在不同的 GPU 类型上运行和构建，请设置 TORCH_CUDA_ARCH_LIST。
# 你也可以从 PyPI 安装最新的稳定版本。
# pip3 install -v -U stable-fast
pip3 install -v -U git+https://github.com/chengzeyi/stable-fast.git@main#egg=stable-fast
# (这可能需要几十分钟)

注意: 在 sfast.compilers 之外的任何使用不保证向后兼容。

注意: 为了获得最佳性能，xformers 和 OpenAI 的 triton>=2.1.0 需要安装和启用。 你可能需要从源码构建 xformers 以使其与 PyTorch 兼容。

使用

优化 StableDiffusionPipeline

stable-fast 可以直接优化 StableDiffusionPipeline 和 StableDiffusionPipelineXL。

import time
import torch
from diffusers import (StableDiffusionPipeline,
                       EulerAncestralDiscreteScheduler)
from sfast.compilers.diffusion_pipeline_compiler import (compile,
                                                         CompilationConfig)

def load_model():
    model = StableDiffusionPipeline.from_pretrained(
        'runwayml/stable-diffusion-v1-5',
        torch_dtype=torch.float16)

    model.scheduler = EulerAncestralDiscreteScheduler.from_config(
        model.scheduler.config)
    model.safety_checker = None
    model.to(torch.device('cuda'))
    return model

model = load_model()

config = CompilationConfig.Default()
# xformers 和 Triton 建议启用以获得最佳性能。
try:
    import xformers
    config.enable_xformers = True
except ImportError:
    print('xformers 未安装，跳过')
try:
    import triton
    config.enable_triton = True
except ImportError:
    print('Triton 未安装，跳过')
# CUDA 图建议用于小 batch size 和小分辨率以减少 CPU 开销。
# 但它可能会增加 GPU 内存使用量。
# 对于 StableVideoDiffusionPipeline，这不是必须的。
config.enable_cuda_graph = True

model = compile(model, config)

kwarg_inputs = dict(
    prompt=
    '(杰作:1.2), 最优质量, 杰作, 最佳详细面部, 一个漂亮的女孩',
    height=512,
    width=512,
    num_inference_steps=30,
    num_images_per_prompt=1,
)

# 注意: 预热它。
# 最初的调用将触发编译，可能非常慢。
# 之后，它应该非常快。
for _ in range(3):
    output_image = model(**kwarg_inputs).images[0]

# 看看它！
# 注意: 由于 CUDA 的异步特性，进度条可能不正确工作。
begin = time.time()
output_image = model(**kwarg_inputs).images[0]
print(f'推理时间: {time.time() - begin:.3f}s')

# 在终端中查看它！
from sfast.utils.term_image import print_image

print_image(output_image, max_width=80)

更多详情参见 examples/optimize_stable_diffusion_pipeline.py。

你可以查看这个 Colab，了解它在 T4 GPU 上的工作方式：

优化 LCM Pipeline

stable-fast 可以优化最新的 latent consistency model pipeline 并实现显著的速度提升。

更多详情参见 examples/optimize_lcm_pipeline.py，了解如何使用 LCM LoRA 优化普通的 SD 模型。 更多详情参见 examples/optimize_lcm_pipeline.py，了解如何优化独立的 LCM 模型。

优化 StableVideoDiffusionPipeline

stable-fast 可以优化最新的 StableVideoDiffusionPipeline 并实现 2x 速度提升

更多详情参见 examples/optimize_stable_video_diffusion_pipeline.py

动态切换 LoRA

动态切换 LoRA 是支持的，但需要一些额外的工作。 这是可能的，因为编译的图和 CUDA 图 与原始的 UNet 模型共享相同的底层数据（指针）。所以你需要做的 就是在原始 UNet 模型的参数上进行就地更新。

以下代码假设你已经加载了一个 LoRA 并编译了模型， 并且你想切换到另一个 LoRA。

如果你没有启用 CUDA 图并保持 `preserve

将“另一个”LoRA切换到UNet中

def switch_lora(unet, lora): # 存储原始UNet参数 state_dict = unet.state_dict() # 将另一个LoRA加载到unet中 unet.load_attn_procs(lora) # 就地复制当前UNet参数到原始unet参数中 update_state_dict(state_dict, unet.state_dict()) # 加载回原始UNet参数 # 由于我们仍然希望保留原始UNet参数的引用，因此使用assign=True unet.load_state_dict(state_dict, assign=True)

switch_lora(compiled_model.unet, lora_b_path)

模型量化

stable-fast扩展了PyTorch的quantize_dynamic功能，并在CUDA后台提供了动态量化的线性运算符。 启用此功能后，你可以为diffusers获得轻微的VRAM减少，为transformers获得显著的VRAM减少，并可能获得潜在的加速（并不总是如此）。

对于SD XL，在图像大小为1024x1024时，预计会减少2GB的VRAM。

def quantize_unet(m):
    from diffusers.utils import USE_PEFT_BACKEND
    assert USE_PEFT_BACKEND
    m = torch.quantization.quantize_dynamic(m, {torch.nn.Linear},
                                            dtype=torch.qint8,
                                            inplace=True)
    return m

model.unet = quantize_unet(model.unet)
if hasattr(model, 'controlnet'):
    model.controlnet = quantize_unet(model.controlnet)

有关更多详细信息，请参阅examples/optimize_stable_diffusion_pipeline.py。

加速PyTorch的一些常用方法

# 强烈建议使用TCMalloc减少CPU开销
# https://github.com/google/tcmalloc
LD_PRELOAD=/path/to/libtcmalloc.so python3 ...

import packaging.version
import torch

if packaging.version.parse(torch.__version__) >= packaging.version.parse('1.12.0'):
    torch.backends.cuda.matmul.allow_tf32 = True

性能比较

性能在不同的硬件/软件/平台/驱动程序配置之间会有很大差异。 准确的基准测试非常困难。而准备基准测试的环境也是一项艰巨的任务。 我之前在一些平台上进行了测试，但结果可能仍然不准确。 请注意，在基准测试时，由于CUDA的异步性质，tqdm显示的进度条可能不准确。 为了解决这个问题，我使用CUDA Event来准确测量每秒迭代的速度。

stable-fast在较新的GPU和较新的CUDA版本上预期效果更好。 在较旧的GPU上，性能提升可能有限。 在基准测试期间，由于CUDA的异步性质，进度条可能无法正常工作。

RTX 4080 (512x512, batch size 1, fp16, 在WSL2中)

这是我的个人游戏PC😄。它的CPU比云服务器提供商的那些更强大。

H100

感谢__@Consceleratus__和__@harishp__的帮助，我已经测试了H100上的速度。

A100

感谢__@SuperSecureHuman__和__@jon-chuang__的帮助，现已在A100上进行基准测试。

兼容性

故障排除

有关更多详细信息，请参阅doc/troubleshooting.md。

你还可以加入Discord频道寻求帮助。