🚀稳定快速

使用LoRA的极速FLUX-dev

使用LoRA的极速Wan 2.1 T2V

使用LoRA的极速Wan 2.1 I2V

🎉重要公告🎉

经过一年的延迟，我很高兴地宣布我计划构建一个新项目Comfy-WaveSpeed，以提供使用ComfyUI运行的所有模型的最快推理速度。 它刚刚开始，我希望它会成为一个很棒的项目👏。请继续关注它，并给我反馈👍！

注意

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

欢迎联系。

Discord频道

stable-fast在__所有__类型的扩散模型上都实现了最先进的推理性能，即使是使用最新的StableVideoDiffusionPipeline也是如此。 与需要花费数十分钟来编译模型的TensorRT或AITemplate不同，stable-fast只需几秒钟即可完成模型编译。 stable-fast还开箱即用地支持动态形状、LoRA和ControlNet。

模型	torch	torch.compile	AIT	oneflow	TensorRT	stable-fast
SD 1.5 (毫秒)	1897	1510	1158	1003	991	995
SVD-XT (秒)	83	70				47

注意: 在基准测试期间，TensorRT是在静态批量大小和启用CUDA图的情况下进行测试的，而stable-fast则是在动态形状下运行的。

• 🚀稳定快速
  • 简介
    • 这是什么？
    • 与其他加速库的区别
  • 安装
    • 安装预构建的轮子
    • 从源代码安装
  • 使用方法
    • 优化StableDiffusionPipeline
    • 优化LCM Pipeline
    • 优化StableVideoDiffusionPipeline
    • 动态切换LoRA
    • 模型量化
    • 一些加快PyTorch速度的常用方法
  • 性能对比
    • RTX 4080 (512x512, 批量大小1, fp16, 在WSL2中)
    • H100
    • A100
  • 兼容性
  • 故障排除

简介

这是什么？

stable-fast 是一个针对 HuggingFace Diffusers 在 NVIDIA GPU 上的超轻量级推理优化框架。 stable-fast 通过利用一些关键技术与特性，提供了极快的推理优化：

• CUDNN 卷积融合: stable-fast 实现了一系列功能完备且完全兼容的 CUDNN 卷积融合算子，适用于各种 Conv + Bias + Add + Act 计算模式的组合。
• 低精度与融合 GEMM: stable-fast 实现了一组以 fp16 精度进行计算的融合 GEMM 算子，其速度比 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 更加稳定，且 CPU 开销显著低于 torch.compile，同时支持 ControlNet 和 LoRA。
• CUDA 图: stable-fast 可以将 UNet、VAE 和 TextEncoder 捕获为 CUDA 图格式，从而在小批量情况下降低 CPU 开销。该实现还支持动态形状。
• 融合多头注意力: stable-fast 直接使用 xformers，并使其与 TorchScript 兼容。

我的下一个目标是让 stable-fast 继续保持作为 diffusers 最快的推理优化框架之一，同时为 transformers 提供加速和显存占用减少的功能。 事实上，我已经使用 stable-fast 来优化 LLM，并取得了显著的加速效果。 不过，我还需要进一步完善，使其更加稳定、易于使用，并提供一个稳定的用户界面。

与其他加速库的区别

• 快速: 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 版本）。

我仅在 CUDA 12.1 和 Python 3.10 环境下，使用 torch>=2.1.0、xformers>=0.0.22 和 triton>=2.1.0 对 stable-fast 进行了测试。 其他版本或许也能成功构建和运行，但无法保证。

安装预编译的 wheel 包

从 Releases 页面 下载适合您系统的 wheel 文件，并使用 pip3 install <wheel file> 进行安装。

目前已有适用于 Linux 和 Windows 的 wheel 包。

# 将 cu121 替换为您使用的 CUDA 版本，<wheel file> 替换为 wheel 文件的路径。
# 请确保 wheel 文件与您的 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 图以减少 CPU 开销。
# 但这样做可能会增加 GPU 显存的使用量。
# 对于 StableVideoDiffusionPipeline，则无需启用。
config.enable_cuda_graph = True

model = compile(model, config)

kwarg_inputs = dict(
    prompt=
    '(masterpiece:1,2), best quality, masterpiece, best detailed face, a beautiful girl',
    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}秒')

# 让我们在终端中查看它！
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 流水线

stable-fast 能够优化最新的 latent consistency model 流水线，并实现显著的速度提升。

有关如何使用 LCM LoRA 优化普通 SD 模型的详细信息，请参阅 examples/optimize_lcm_lora.py。有关如何优化独立的 LCM 模型的详细信息，请参阅 examples/optimize_lcm_pipeline.py。

优化 StableVideoDiffusionPipeline

stable-fast 能够优化最新的 StableVideoDiffusionPipeline，并实现 2 倍 的速度提升。

更多详情请参阅 examples/optimize_stable_video_diffusion_pipeline.py。

动态切换 LoRA

支持动态切换 LoRA，但需要进行一些额外的操作。 之所以可行，是因为编译后的图和 CUDA Graph 与原始 UNet 模型共享相同的底层数据（指针）。因此，您只需就地更新原始 UNet 模型的参数即可。

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

如果您未启用 CUDA 图且保持 preserve_parameters = True，操作会简单得多。甚至可能不需要执行以下代码。

# load_state_dict with assign=True 需要 torch >= 2.1.0

def update_state_dict(dst, src):
    for key, value in src.items():
        # 进行就地复制。
        # 由于追踪的前向函数共享相同的底层数据（指针），此修改将反映在追踪的前向函数中。
        dst[key].copy_(value)

# 将“另一个”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 参数重新加载回模型。
    # 我们使用 assign=True，因为我们仍希望保留原始 UNet 参数的引用
    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 时，VRAM 可减少 2GB。

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，批次大小 1，fp16，在 WSL2 中）

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

框架	SD 1.5	SD XL (1024x1024)	SD 1.5 ControlNet
Vanilla PyTorch (2.1.0)	29.5 it/s	4.6 it/s	19.7 it/s
torch.compile (2.1.0，max-autotune)	40.0 it/s	6.1 it/s	21.8 it/s
AITemplate	44.2 it/s
OneFlow	53.6 it/s
AUTO1111 WebUI	17.2 it/s	3.6 it/s
AUTO1111 WebUI（带 SDPA）	24.5 it/s	4.3 it/s
TensorRT（AUTO1111 WebUI）	40.8 it/s
TensorRT 官方演示	52.6 it/s
stable-fast（带 xformers & Triton）	51.6 it/s	9.1 it/s	36.7 it/s

H100

感谢 @Consceleratus 和 @harishp 的帮助，我在 H100 上进行了速度测试。

框架	SD 1.5	SD XL (1024x1024)	SD 1.5 ControlNet
Vanilla PyTorch (2.1.0)	54.5 it/s	14.9 it/s	35.8 it/s
torch.compile (2.1.0，max-autotune)	66.0 it/s	18.5 it/s
stable-fast（带 xformers & Triton）	104.6 it/s	21.6 it/s	72.6 it/s

A100

感谢 @SuperSecureHuman 和 @jon-chuang 的帮助，现在可以在 A100 上进行基准测试了。

框架	SD 1.5	SD XL (1024×1024)	SD 1.5 ControlNet
原生 PyTorch (2.1.0)	35.6 it/s	8.7 it/s	25.1 it/s
torch.compile (2.1.0, max-autotune)	41.9 it/s	10.0 it/s
stable-fast (使用 xformers 和 Triton)	61.8 it/s	11.9 it/s	41.1 it/s

兼容性

模型	支持
Hugging Face Diffusers (1.5/2.1/XL)	是
使用 ControlNet	是
使用 LoRA	是
隐式一致性模型	是
SDXL Turbo	是
Stable Video Diffusion	是

功能	支持
动态形状	是
文本到图像	是
图像到图像	是
图像修复	是

UI 框架	支持	链接
AUTOMATIC1111	开发中
SD Next	是	SD Next
ComfyUI	是	ComfyUI_stable_fast

操作系统	支持
Linux	是
Windows	是
Windows WSL	是

故障排除

更多详情请参阅 doc/troubleshooting.md。

您也可以加入 Discord 频道 寻求帮助。

Stable Fast 快速上手指南

stable-fast 是一个专为 NVIDIA GPU 设计的超轻量级推理优化框架，旨在加速 HuggingFace Diffusers 库。相比 TensorRT 或 AITemplate 需要数十分钟的编译时间，stable-fast 仅需数秒即可完成模型编译，并原生支持动态形状（Dynamic Shape）、LoRA 和 ControlNet。

环境准备

在开始之前，请确保满足以下系统和依赖要求：

• 操作系统：Linux 或 Windows (WSL2)。
• GPU：NVIDIA GPU。
• Python 版本：推荐 Python 3.10。
• CUDA 版本：推荐 CUDA 12.1（支持范围 11.x - 12.x）。
• 核心依赖：
  • torch >= 2.1.0 (必须带 CUDA 支持)
  • xformers >= 0.0.22 (强烈建议安装以获得最佳性能)
  • triton >= 2.1.0 (强烈建议安装)
  • diffusers >= 0.19.3
  • ninja (可选，用于加速构建过程)

注意：虽然其他版本可能也能运行，但作者仅在上述特定版本组合下进行了充分测试。Windows 用户若无法安装 Triton 可跳过，但性能可能受影响。

安装步骤

方法一：安装预编译包（推荐）

从 Releases 页面 下载对应你系统环境和 PyTorch/CUDA 版本的 .whl 文件，然后使用 pip 安装。

# 请将 <wheel file> 替换为实际下载的文件路径
# 确保 --index-url 对应的 CUDA 版本与你安装的 PyTorch 版本一致 (例如 cu121)
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>'

方法二：从源码安装

如果你需要最新功能或没有匹配的预编译包，可以从源码安装。

# 1. 安装基础依赖 (确保已安装 CUDNN 和 CUBLAS)
# Windows 用户如果无法安装 triton 可移除该项
pip3 install wheel 'torch>=2.1.0' 'xformers>=0.0.22' 'triton>=2.1.0' 'diffusers>=0.19.3'

# 2. (可选) 安装 ninja 以加速编译
pip3 install ninja

# 3. (可选) 如果你的运行环境和编译环境 GPU 架构不同，需设置环境变量
# export TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0" 

# 4. 从 GitHub 安装
pip3 install -v -U git+https://github.com/chengzeyi/stable-fast.git@main#egg=stable-fast

基本使用

stable-fast 可以直接优化 StableDiffusionPipeline 和 StableDiffusionXLPipeline。以下是优化 SD 1.5 模型的最简示例：

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

def load_and_optimize_model():
    # 1. 加载原始模型
    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'))

    # 2. 配置编译选项
    config = CompilationConfig.Default()
    # 启用 xformers 和 Triton 以获得最佳性能
    config.enable_xformers = True
    config.enable_triton = True
    
    # 3. 编译模型 (仅需数秒)
    # compile 函数会原地修改 model 对象，使其获得加速
    compiled_model = compile(model, config)
    
    return compiled_model

# 初始化优化后的模型
pipe = load_and_optimize_model()

# 像往常一样使用
prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]
image.save("astronaut_rides_horse.png")

关键特性说明

• 动态形状支持：编译后的模型支持不同的分辨率和 Batch Size，无需重新编译。
• LoRA 支持：优化后的管道可以直接加载和切换 LoRA 权重，无需特殊处理。
• 显存优化：除了速度提升，该框架还能有效降低显存占用。