高性能推理优化技巧

摘要

Stable Diffusion WebUI 在图像生成过程中需要大量的计算资源，尤其是在消费级硬件上运行时，性能优化显得尤为重要。本文将深入探讨 WebUI 中的各种性能优化技术，包括 Token Merging (ToMe)、交叉注意力优化、批处理优化、模型精度调整等。我们将分析这些优化技术的实现原理和应用场景，并提供实际使用建议，帮助用户在保证生成质量的前提下最大化推理速度。

关键词： Stable Diffusion, 性能优化, Token Merging, 注意力机制, 推理加速

1. 引言

Stable Diffusion 模型虽然在图像生成方面表现出色，但其计算复杂度也相当高。在实际应用中，特别是在消费级 GPU 上，推理速度往往成为用户体验的关键瓶颈。Stable Diffusion WebUI 作为一个广泛使用的前端界面，集成了多种性能优化技术，旨在提升推理速度而不显著降低图像质量。

本文将从以下几个方面详细介绍高性能推理优化技巧：

1. Token Merging (ToMe) 优化技术
2. 交叉注意力机制优化
3. 批处理和条件缓存优化
4. 模型精度和量化优化
5. 其他性能优化策略

2. Token Merging (ToMe) 优化技术

2.1 ToMe 原理简介

Token Merging (ToMe) 是一种通过合并视觉 Transformer 中相似的 token 来减少计算量的优化技术。在扩散模型中，这种方法特别适用于减少 UNet 网络在去噪过程中的计算负担。

ToMe 的核心思想是识别并合并语义上相似的 token，从而在保持视觉保真度的同时显著减少计算量。这种方法在图像生成任务中表现良好，因为它保留了重要的全局结构信息，同时减少了冗余计算。

2.2 WebUI 中的 ToMe 实现

在 WebUI 中，ToMe 通过 [apply_token_merging](file:///e:/project/stable-diffusion-webui/modules/sd_models.py#L998-L1031) 函数实现：

def apply_token_merging(sd_model, token_merging_ratio):
    """
    Applies speed and memory optimizations from tomesd.
    """

    current_token_merging_ratio = getattr(sd_model, 'applied_token_merged_ratio', 0)

    if current_token_merging_ratio == token_merging_ratio:
        return

    if current_token_merging_ratio > 0:
        tomesd.remove_patch(sd_model)

    if token_merging_ratio > 0:
        tomesd.apply_patch(
            sd_model,
            ratio=token_merging_ratio,
            use_rand=False,  # can cause issues with some samplers
            merge_attn=True,
            merge_crossattn=False,
            merge_mlp=False
        )

    sd_model.applied_token_merged_ratio = token_merging_ratio

2.3 ToMe 参数配置

WebUI 提供了多个与 ToMe 相关的配置选项：

1. Token merging ratio：控制 ToMe 的合并比例，范围为 0.0-0.9，数值越高合并越多，速度越快但质量可能下降
2. Token merging ratio for img2img：专门针对图像到图像任务的合并比例
3. Token merging ratio for high-res pass：针对高分辨率修复过程的合并比例

"token_merging_ratio": OptionInfo(0.0, "Token merging ratio", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}, infotext='Token merging ratio').link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9256").info("0=disable, higher=faster"),
"token_merging_ratio_img2img": OptionInfo(0.0, "Token merging ratio for img2img", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}).info("only applies if non-zero and overrides above"),
"token_merging_ratio_hr": OptionInfo(0.0, "Token merging ratio for high-res pass", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}, infotext='Token merging ratio hr').info("only applies if non-zero and overrides above"),

2.4 ToMe 使用建议

1. 保守设置：对于一般用途，建议将比率设置为 0.2-0.4，可以在保持较好质量的同时获得明显的速度提升
2. 激进设置：对于草图或快速预览，可以尝试 0.5-0.7 的比率
3. 禁用场景：对于最终输出或对质量要求极高的场景，建议将比率设为 0.0（禁用）

3. 交叉注意力机制优化

3.1 注意力机制优化概述

交叉注意力机制是 Transformer 架构的核心组件，但也往往是计算瓶颈所在。WebUI 提供了多种优化交叉注意力计算的方法：

1. xformers：使用 Facebook Research 的 xformers 库优化注意力计算
2. Sub-quadratic attention：使用次二次复杂度的注意力算法
3. Memory efficient attention：内存高效的注意力实现

3.2 xformers 优化

xformers 是一个高效的 Transformer 库，它提供了多种优化的注意力实现。在 WebUI 中启用 xformers 非常简单，只需在启动时添加 --xformers 参数。

# 在 import_hook.py 中
if "--xformers" not in "".join(sys.argv):
    sys.modules["xformers"] = None

3.3 次二次注意力优化

对于不支持 xformers 的环境，WebUI 还提供了次二次注意力优化：

# 在 sd_hijack_optimizations.py 中
class SdOptimization:
    def apply(self):
        pass

    def undo(self):
        ldm.modules.attention.CrossAttention.forward = hypernetwork.attention_CrossAttention_forward
        ldm.modules.diffusionmodules.model.AttnBlock.forward = diffusionmodules_model_AttnBlock_forward

        sgm.modules.attention.CrossAttention.forward = hypernetwork.attention_CrossAttention_forward

4. 批处理和条件缓存优化

4.1 批处理优化

批处理是一种常见的优化技术，通过并行处理多个样本提高硬件利用率：

"batch_cond_uncond": OptionInfo(True, "Batch cond/uncond").info("do both conditional and unconditional denoising in one batch; uses a bit more VRAM during sampling, but improves speed; previously this was controlled by --always-batch-cond-uncond commandline argument"),

这个选项控制是否在单个批次中同时处理条件和无条件去噪，虽然会略微增加显存使用，但能显著提升速度。

4.2 条件缓存优化

为了避免重复计算相同的提示词条件，WebUI 实现了条件缓存机制：

"persistent_cond_cache": OptionInfo(True, "Persistent cond cache").info("do not recalculate conds from prompts if prompts have not changed since previous calculation"),

当启用此选项时，如果提示词未发生变化，系统将重用之前的条件计算结果，而不是重新计算。

4.3 提示词填充优化

为了优化不同长度提示词的处理，WebUI 提供了填充选项：

"pad_cond_uncond": OptionInfo(False, "Pad prompt/negative prompt", infotext='Pad conds').info("improves performance when prompt and negative prompt have different lengths; changes seeds"),
"pad_cond_uncond_v0": OptionInfo(False, "Pad prompt/negative prompt (v0)", infotext='Pad conds v0').info("alternative implementation for the above; used prior to 1.6.0 for DDIM sampler; overrides the above if set; WARNING: truncates negative prompt if it's too long; changes seeds"),

这些选项通过填充较短的提示词序列来匹配较长的序列，从而提高批处理效率。

5. 模型精度和量化优化

5.1 半精度和混合精度

现代 GPU 支持半精度 (FP16) 和混合精度计算，可以显著提升计算速度并减少显存占用。WebUI 默认启用这些优化：

devices.dtype_inference = torch.float32 if cmd_opts.precision == 'full' else devices.dtype

if cmd_opts.precision == "half":
    msg = "--no-half and --no-half-vae conflict with --precision half"
    assert devices.dtype == torch.float16, msg
    assert devices.dtype_vae == torch.float16, msg
    assert devices.dtype_inference == torch.float16, msg
    devices.force_fp16 = True
    devices.force_model_fp16()

5.2 FP8 量化

最新的 WebUI 版本支持 FP8 量化，进一步减少模型大小和计算需求：

"fp8_storage": OptionInfo("Disable", "FP8 weight", gr.Radio, {"choices": ["Disable", "Enable for SDXL", "Enable"]}).info("Use FP8 to store Linear/Conv layers' weight. Require pytorch>=2.1.0."),
"cache_fp16_weight": OptionInfo(False, "Cache FP16 weight for LoRA").info("Cache fp16 weight when enabling FP8, will increase the quality of LoRA. Use more system ram."),

FP8 量化可以显著减少显存使用，但需要 PyTorch 2.1.0 或更高版本支持。

6. 其他性能优化策略

6.1 负面引导最小 Sigma 优化

通过跳过小 sigma 值的负面引导计算来提升性能：

"s_min_uncond": OptionInfo(0.0, "Negative Guidance minimum sigma", gr.Slider, {"minimum": 0.0, "maximum": 15.0, "step": 0.01}, infotext='NGMS').link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9177").info("skip negative prompt for some steps when the image is almost ready; 0=disable, higher=faster"),
"s_min_uncond_all": OptionInfo(False, "Negative Guidance minimum sigma all steps", infotext='NGMS all steps').info("By default, NGMS above skips every other step; this makes it skip all steps"),

6.2 随机数生成器优化

不同的随机数生成器源会影响性能和跨平台一致性：

"randn_source": OptionInfo("GPU", "Random number generator source.", gr.Radio, {"choices": ["GPU", "CPU", "NV"]}, infotext="RNG").info("changes seeds drastically; use CPU to produce the same picture across different videocard vendors; use NV to produce same picture as on NVidia videocards"),

6.3 VAE 优化

VAE（变分自编码器）是图像编解码的关键组件，优化 VAE 可以提升整体性能：

# 在 sd_vae_taesd.py 中
class TAESDDecoder(nn.Module):
    # Tiny AutoEncoder 实现，用于快速预览

WebUI 支持多种 VAE 优化模式，包括使用更小更快的 TAESD 模型进行实时预览。

7. 性能优化配置推荐

7.1 入门级配置（8-12GB 显存）

1. 启用 xformers 或 sub-quadratic attention
2. 设置 Token merging ratio 为 0.3-0.4
3. 启用 batch_cond_uncond
4. 启用 persistent_cond_cache
5. 使用半精度而非全精度

7.2 中端配置（12-16GB 显存）

1. 启用 xformers
2. 设置 Token merging ratio 为 0.4-0.5
3. 启用 batch_cond_uncond
4. 启用 persistent_cond_cache
5. 启用 pad_cond_uncond
6. 考虑使用 FP8 量化

7.3 高端配置（16GB+ 显存）

1. 启用 xformers
2. 设置 Token merging ratio 为 0.5-0.6（视质量要求而定）
3. 启用所有批处理和缓存优化
4. 根据需要选择精度模式

8. 性能监控和调试

8.1 性能指标

在优化过程中，应该关注以下关键指标：

1. 每秒迭代次数 (it/s)：反映模型处理速度
2. 显存使用量：确保不超过硬件限制
3. 生成时间：端到端的图像生成时间
4. 图像质量：主观评估优化对质量的影响

8.2 调试技巧

1. 逐步调整：一次只调整一个参数，观察其影响
2. 对比测试：在相同条件下比较不同配置的效果
3. 日志分析：查看 WebUI 输出的日志信息，了解各阶段耗时

9. 总结

Stable Diffusion WebUI 提供了丰富的性能优化选项，用户可以根据自己的硬件配置和质量要求进行调整。通过合理配置 Token Merging、注意力优化、批处理策略和精度设置，可以在保证图像质量的前提下显著提升推理速度。

关键要点包括：

1. Token Merging 是最有效的通用优化技术之一，适度使用可以获得很好的性价比
2. 交叉注意力优化（如 xformers）对高端显卡有显著效果
3. 合理使用批处理和缓存机制可以避免重复计算
4. 模型量化技术（如 FP8）是未来的趋势，值得尝试
5. 不同的优化策略可以组合使用，但需要注意它们之间的相互影响

在实际应用中，建议从保守设置开始，逐步调整参数直到找到质量和速度的最佳平衡点。通过这些优化技术，即使是消费级硬件也能获得流畅的图像生成体验。

参考资料

1. Stable Diffusion WebUI GitHub 仓库: https://github.com/AUTOMATIC1111/stable-diffusion-webui
2. Token Merging 论文: https://arxiv.org/abs/2303.17604
3. xformers 库: https://github.com/facebookresearch/xformers
4. Memory-Efficient Attention: https://github.com/AminRezaei0x443/memory-efficient-attention
5. PyTorch 混合精度训练文档: https://pytorch.org/docs/stable/amp.html