stable-fast:一个高性能的Stable Diffusion推理优化框架

stable-fast

stable-fast:加速Stable Diffusion推理的利器

在人工智能图像生成领域,Stable Diffusion模型以其强大的生成能力而备受关注。然而,Stable Diffusion模型的推理速度一直是一个痛点,尤其是在处理高分辨率图像或视频时。为了解决这个问题,GitHub上的开源项目stable-fast应运而生。本文将深入介绍stable-fast的特性、优势以及如何使用它来优化Stable Diffusion模型的推理性能。

stable-fast简介

stable-fast是一个专为HuggingFace Diffusers设计的超轻量级推理优化框架。它能够在NVIDIA GPU上实现最先进的推理性能,甚至可以优化最新的StableVideoDiffusionPipeline。与TensorRT或AITemplate等需要数十分钟来编译模型的工具不同,stable-fast只需几秒钟即可完成模型编译。此外,stable-fast还支持动态形状、LoRA和ControlNet等功能。

stable-fast性能对比图

stable-fast的核心特性

stable-fast采用了多项先进技术来实现高性能推理:

CUDNN卷积融合: 实现了一系列全功能、全兼容的CUDNN卷积融合算子,可处理各种Conv + Bias + Add + Act计算模式的组合。
低精度和融合GEMM: 实现了一系列使用fp16精度计算的融合GEMM算子,比PyTorch默认的方法(读写fp16,计算fp32)更快。
融合线性GEGLU: 能够将GEGLU(x, W, V, b, c) = GELU(xW + b) ⊗ (xV + c)融合到一个CUDA内核中。
NHWC和融合GroupNorm: 使用OpenAI的Triton实现了高度优化的融合NHWC GroupNorm + Silu算子,消除了内存格式置换算子的需求。
完全跟踪模型: 改进了torch.jit.trace接口,使其更适合跟踪复杂模型。StableDiffusionPipeline/StableVideoDiffusionPipeline的几乎每个部分都可以被跟踪并转换为TorchScript。
CUDA图: 可以将UNet、VAE和TextEncoder捕获为CUDA图格式,在批处理大小较小时可以减少CPU开销。这种实现还支持动态形状。
融合多头注意力: 使用xformers并使其与TorchScript兼容。

stable-fast与其他加速库的区别

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

如何安装stable-fast

stable-fast目前仅在Linux和Windows的WSL2上进行了测试。首先需要安装支持CUDA的PyTorch(建议使用1.12到2.1版本)。

安装预构建的wheel:

从Releases页面下载与您系统对应的wheel文件。
使用以下命令安装:

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>'

从源代码安装:

pip3 install wheel 'torch>=2.1.0' 'xformers>=0.0.22' 'triton>=2.1.0' 'diffusers>=0.19.3'
pip3 install ninja
pip3 install -v -U git+https://github.com/chengzeyi/stable-fast.git@main#egg=stable-fast

如何使用stable-fast

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()
config.enable_xformers = True
config.enable_triton = True
config.enable_cuda_graph = True

model = compile(model, config)

kwargs = 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(**kwargs).images[0]

# 测试推理时间
begin = time.time()
output_image = model(**kwargs).images[0]
print(f'Inference time: {time.time() - begin:.3f}s')

stable-fast的性能表现

stable-fast在各种硬件上都展现出了优秀的性能。以下是在不同GPU上的性能对比:

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

框架	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 (with SDPA)	24.5 it/s	4.3 it/s
TensorRT (AUTO1111 WebUI)	40.8 it/s
TensorRT Official Demo	52.6 it/s
stable-fast (with xformers & Triton)	51.6 it/s	9.1 it/s	36.7 it/s

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 (with xformers & Triton)	104.6 it/s	21.6 it/s	72.6 it/s