Diffusers 更换 scheduler 后保持采样步数可比的方法

Diffusers 更换 scheduler 后，相同的 num_inference_steps 只能保证相同的请求步数，不能自动保证相同去噪网络调用次数、耗时或噪声时间网格。要让对比有意义，应固定模型、提示词、分辨率、引导强度和随机种子，记录实际网络调用次数，再明确采用“同请求步数”还是“同计算预算”的比较口径。

如果我要比较两种采样器，首先会放弃“两个进度条都是 20 就公平”的直觉：进度条是界面表达，模型真正做了多少次计算还要单独记录。下面给出我的比较组织方式，代码是未执行的实验示例，不提供虚构的速度或质量排名。

Diffusers 官方地址：https://huggingface.co/docs/diffusers/

我会先明确要比较的是步数还是计算预算

同请求步数回答“把相同步数参数交给两个采样器，会得到什么”；同去噪网络调用预算回答“用近似相同的主要推理计算次数，谁更适合当前任务”。两种问题都合理，结论却不能混写。

高阶求解方法、额外校正或初始化过程可能改变时间步条目与网络评估之间的关系。scheduler.order、某些算法的 solver_order 也不是可以通用互换的“每步网络次数”。我更愿意在报告里分列请求步数、时间步条目数、实际 UNet 调用次数和耗时，而不是挑其中一个代替全部指标。

替换时保留模型对应的配置

官方支持通过 from_config 从已有配置创建兼容的 scheduler。这里“兼容”比类名更重要：模型的预测目标、训练噪声配置和采样器能力需要一致。普通扩散、flow matching 与少步蒸馏模型不能不加区分地套用同一组替换类。

我会在任何替换之前保存原始 pipe.scheduler.config，每个候选都从这份配置建立新对象。这样不会把上一轮采样器的运行状态带进下一轮。但 from_config 不是自动证明所有设置等价的工具；目标类的专属默认值仍可能不同，应把实际生效的配置一并记录。

第一轮只替换 scheduler，不同时开启 Karras sigmas、变更 timestep spacing 或改动预测类型。如果之后想比较另一种噪声网格，就把它作为单独变量。不同算法未必支持相同的自定义 timesteps 或 sigmas，不应为了“完全相同”而传入不受支持的参数。

每轮重新建立随机数状态

生成器会随着随机数消费向前推进。把同一个 Generator 对象连续传给两次生成，并不等于两次从同一个随机状态开始。示例在每轮内部重新创建 CPU Generator 并设置相同 seed；这是控制随机输入的起点，不是承诺不同采样器产生相同图片。

CPU Generator 便于控制跨设备随机采样差异，但不同硬件、依赖版本及算子仍可能使结果发生变化。使用相同的模型修订、Diffusers 与 PyTorch 版本、设备、精度和 attention 实现，才能减少额外干扰。含随机附加噪声的算法还可能在采样过程中消费不同数量的随机数。

记录实际去噪网络调用次数

下面针对标准 StableDiffusionPipeline 的 UNet 文生图路径，用 PyTorch forward hook 记录实际调用。准备好可用的 Diffusers、PyTorch、模型相关依赖，并将环境变量 SD_MODEL_DIR 指向本地完整模型目录；CPU 执行可能很慢，示例不用于速度排名。候选 DDIM 与 Heun 必须先确认适用于所选模型。

import os
import json
import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler, HeunDiscreteScheduler

# 指向本地、已获授权的标准 Stable Diffusion 文生图模型目录；
# 本例不是 SDXL、Flux 或蒸馏模型的通用替换脚本。
model_dir = os.environ["SD_MODEL_DIR"]
pipe = StableDiffusionPipeline.from_pretrained(
    model_dir, local_files_only=True
).to("cpu")
base_config = dict(pipe.scheduler.config)
prompt = "a ceramic teapot on a linen cloth, soft daylight"

def observe(scheduler_class, requested_steps, seed):
    # 每轮从同一份原配置创建新对象，避免继承上一轮采样器状态。
    pipe.scheduler = scheduler_class.from_config(base_config)
    generator = torch.Generator(device="cpu").manual_seed(seed)
    counter = {"unet_calls": 0}

    def count_forward(module, inputs, output):
        # 统计实际 UNet forward 调用，不把进度条或 scheduler 阶数当作次数。
        counter["unet_calls"] += 1

    handle = pipe.unet.register_forward_hook(count_forward)
    try:
        with torch.inference_mode():
            pipe(
                prompt=prompt,
                negative_prompt="",
                height=512,
                width=512,
                guidance_scale=7.0,
                num_images_per_prompt=1,
                num_inference_steps=requested_steps,
                generator=generator,
                output_type="latent",
            )
        return {
            "scheduler": scheduler_class.__name__,
            "requested_steps": requested_steps,
            "timestep_entries": len(pipe.scheduler.timesteps),
            "unet_calls": counter["unet_calls"],
            "seed": seed,
            "scheduler_config": dict(pipe.scheduler.config),
        }
    finally:
        # 异常也卸下 hook，防止下一轮重复计数。
        handle.remove()

# 先比较相同请求步数；实际数据由读者在自己的固定环境运行取得。
for scheduler_class in (DDIMScheduler, HeunDiscreteScheduler):
    record = observe(scheduler_class, requested_steps=20, seed=731)
    print(json.dumps(record, ensure_ascii=False, default=str))

代码输出的 unet_calls 是 forward 调用次数，不是严格 FLOPs。classifier-free guidance 可能把条件与无条件样本合成一次更大的 batch，因此需要固定 guidance、batch 和分辨率。对使用 transformer、多个去噪器或额外网络的 pipeline，应在真正的计算模块上统计，不能照搬这个 UNet hook。

输出类型设为 latent，目的是隔离这段观察所关心的去噪过程；它不提供用于质量评估的最终图片。如果要做图像质量对照，再使用一致的解码与后处理设置生成样本，并分别注明统计耗时是否包含文本编码、VAE 解码和保存。

把结论限制在同一实验口径

同预算实验可以为候选选择不同的请求步数，再按实际调用次数配对；若不能恰好相等，就明确记录预算差距，不用隐含折算掩盖它。比较延迟时先做一致的预热，多次运行并记录中位数；CUDA 计时需要处理异步执行，不能拿未同步的普通计时直接下结论。

我的取舍是先用一组提示词和多组固定种子观察稳定性，再根据任务需要决定看同预算质量，还是看达到可接受质量所需的耗时。单张图更合口味，可以支持局部选择，不能证明一种 scheduler 在所有模型上更好。

更换 scheduler 时的常见问题

为什么相同步数下 timesteps 长度不同？

不同方法可能使用不同的内部调度或校正安排。记录实际长度并查看对应类的实现，不要把列表长度差异直接判为传参失败。

把 scheduler.order 乘以步数能代替统计吗？

不能作为所有 pipeline 的通用公式。求解器阶数、进度更新方式和实际网络调用不是同一个概念，额外分支也可能影响次数。

保持相同 seed 为什么图像仍然变化？

采样轨迹、噪声尺度与算法本身已经改变，相同种子只控制随机输入的来源。更换算法后期待图像完全相同，反而与实验目的冲突。

调度器替换：https://huggingface.co/docs/diffusers/using-diffusers/schedulers。随机性控制：https://huggingface.co/docs/diffusers/using-diffusers/reusing_seeds。