diffusers 优化加速文生图效率-CSDN博客

本文链接：https://blog.csdn.net/zhilaizhiwang/article/details/140645848

参考自：Accelerate inference of text-to-image diffusion models

1.bfloat16

使用 torch.bfloat16 或者torch.float16，降低数据精度能加快推理速度，并且对结果的影响也很小。
如果GPU的内存不足，也可以使用torch.bfloat16 或者torch.float16，能降低内存占用。

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
).to("cuda")

2.SDPA

scaled_dot_product_attention，在diffusers中已经默认使用。

scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)
将scores除以d_k的平方根（np.sqrt(d_k)），这就是所谓的缩放，已经是transformers的默认操作了。

3.torch.compile

需要PyTorch 2 以上版本。第一次编译会很慢，编译好了推理会提速很多。

from diffusers import StableDiffusionXLPipeline
import torch

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True

pipe.unet.to(memory_format=torch.channels_last)
pipe.vae.to(memory_format=torch.channels_last)

# Compile the UNet and VAE.
pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=True)
pipe.vae.decode = torch.compile(pipe.vae.decode, mode="max-autotune", fullgraph=True)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

# First call to `pipe` is slow, subsequent ones are faster.
image = pipe(prompt, num_inference_steps=30).images[0]

4.Combine the attention block’s projection matrices

在计算transformers的时候，Q，K，V如果是通过三个独立的投影创建的，不如合并到一起，一步完成投影，然后再拆分，这样更有效率。

pipe.fuse_qkv_projections()

5.Dynamic quantization

使用 torchao 对UNet和VAE模型应用动态 int8 量化。
量化会在模型中增加额外的转换开销，如果矩阵很小，转换成本会抵消量化的效益。

torch._inductor.config.force_fuse_int_mm_with_mul = True
torch._inductor.config.use_mixed_mm = True

#一些线性层并不能从int8量化中受益，可以先过滤掉
def dynamic_quant_filter_fn(mod, *args):
    return (
        isinstance(mod, torch.nn.Linear)
        and mod.in_features > 16
        and (mod.in_features, mod.out_features)
        not in [
            (1280, 640),
            (1920, 1280),
            (1920, 640),
            (2048, 1280),
            (2048, 2560),
            (2560, 1280),
            (256, 128),
            (2816, 1280),
            (320, 640),
            (512, 1536),
            (512, 256),
            (512, 512),
            (640, 1280),
            (640, 1920),
            (640, 320),
            (640, 5120),
            (640, 640),
            (960, 320),
            (960, 640),
        ]
    )


def conv_filter_fn(mod, *args):
    return (
        isinstance(mod, torch.nn.Conv2d) and mod.kernel_size == (1, 1) and 128 in [mod.in_channels, mod.out_channels]
    )
#把 pointwise convolution layers 转换为线性层，用来量化。  
from torchao import swap_conv2d_1x1_to_linear

swap_conv2d_1x1_to_linear(pipe.unet, conv_filter_fn)
swap_conv2d_1x1_to_linear(pipe.vae, conv_filter_fn)

from torchao import apply_dynamic_quant

apply_dynamic_quant(pipe.unet, dynamic_quant_filter_fn)
apply_dynamic_quant(pipe.vae, dynamic_quant_filter_fn)