用 Diffusers 手写自己的扩散管道：理解模型与调度器的去噪协作机制

用 Diffusers 手写自己的扩散管道：理解模型与调度器的去噪协作机制

【免费下载链接】diffusers🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers

🧨 Diffusers 被设计成一个面向扩散系统的"积木工具箱"：模型（models）负责预测噪声，调度器（schedulers）负责把噪声一步步"还原"成图像。虽然DiffusionPipeline把这些组件打包成一行代码即可调用，但你完全可以拆开管道，分别使用模型与调度器，亲手组装属于自己的扩散系统。本篇指南将带你从零拆解一个基础管道，再逐步进阶到 Stable Diffusion 这种复杂的文本生成图像管道，最终掌握"去噪循环（denoising loop）"这一贯穿所有扩散系统的核心模式。

阅读本文后，你将能够：理解UNet2DModel+DDPMScheduler的最小去噪流程、手工复现DDPMPipeline的全部推理逻辑、并具备独立拆解 Stable Diffusion 管道（VAE + 文本编码器 + UNet + 调度器）并自由更换调度器的实战能力。

先拆解一个基础管道

管道（pipeline）是运行模型推理的最快方式，生成一张图像只需要不到四行代码：

>>> from diffusers import DDPMPipeline >>> ddpm = DDPMPipeline.from_pretrained("google/ddpm-cat-256", use_safetensors=True).to("cuda") # 或 "mps"、"xpu"、"cpu" >>> image = ddpm(num_inference_steps=25).images[0] >>> image

如此简单，但管道内部到底做了什么？我们把它拆开来看。

在上述例子中，管道里包含一个UNet2DModel模型和一个DDPMScheduler调度器。管道的去噪逻辑是：取一个与期望输出尺寸相同的随机噪声，将其反复送入模型。在每个时间步（timestep），模型预测出噪声残差（noise residual），调度器则利用该残差预测出一个噪声更少的图像。管道不断重复这一过程，直到到达指定的推理步数。

在 DDPMPipeline 的官方实现 中，这段逻辑清晰可见：先采样高斯噪声作为起点，调用self.scheduler.set_timesteps(num_inference_steps)设定步数，然后遍历self.scheduler.timesteps，每一步先由self.unet(image, t).sample预测噪声输出，再由self.scheduler.step(model_output, t, image).prev_sample计算上一时间步的图像，如此循环往复。

为了用模型和调度器分别复现这条管道，我们手写自己的去噪过程。

第 1 步：加载模型与调度器

>>> from diffusers import DDPMScheduler, UNet2DModel >>> scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256") >>> model = UNet2DModel.from_pretrained("google/ddpm-cat-256", use_safetensors=True).to("cuda") # 或 "mps"、"xpu"、"cpu"

from_pretrained是DiffusionPipeline、各模型与调度器通用的加载入口。注意调度器与模型不同：它没有可训练权重，只是一个计算"上一时间步样本"的纯算法模块，因此无需移动到 GPU（不过移动也不会出错）。

第 2 步：设定去噪时间步数

>>> scheduler.set_timesteps(50)

第 3 步：查看生成的时间步张量

set_timesteps会创建一个元素均匀分布的张量（本例为 50 个），每个元素对应一个模型执行去噪的时间步。之后的去噪循环会遍历这个张量：

>>> scheduler.timesteps tensor([980, 960, 940, 920, 900, 880, 860, 840, 820, 800, 780, 760, 740, 720, 700, 680, 660, 640, 620, 600, 580, 560, 540, 520, 500, 480, 460, 440, 420, 400, 380, 360, 340, 320, 300, 280, 260, 240, 220, 200, 180, 160, 140, 120, 100, 80, 60, 40, 20, 0])

从 DDPMScheduler.set_timesteps 的实现 可以看到时间步的三种生成策略：默认的linspace在[0, num_train_timesteps - 1]区间内均匀取点后倒序排列；leading与trailing对应 2305.08891 论文表 2 的标注，通过step_ratio换算整数时间步。此外该方法还做了参数校验：num_inference_steps不能超过训练时间步数num_train_timesteps（google/ddpm-cat-256为 1000 步），并支持传入自定义的timesteps列表实现任意时间步间距。

第 4 步：创建随机噪声

创建与期望输出形状相同的随机噪声：

>>> import torch >>> sample_size = model.config.sample_size >>> noise = torch.randn((1, 3, sample_size, sample_size), device="cuda") # 或 "mps"、"xpu"、"cpu"

形状中的model.config.sample_size来自 UNet2DModel 的配置（google/ddpm-cat-256为 256），3是 RGB 通道数。

第 5 步：编写去噪循环

在每个时间步，模型执行一次UNet2DModel.forward前向传播并返回噪声残差；调度器的DDPMScheduler.step方法接收噪声残差、时间步与当前输入，预测出上一时间步的图像。该输出成为下一轮循环的输入，直到遍历完整个timesteps数组：

>>> input = noise >>> for t in scheduler.timesteps: ... with torch.no_grad(): ... noisy_residual = model(input, t).sample ... previous_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample ... input = previous_noisy_sample

这就是完整的去噪过程。你可以用同样的模式去编写任何扩散系统。

从源码层面看，DDPMScheduler.step内部按以下步骤工作：先根据prediction_type（epsilon/sample/v_prediction）从预测噪声反推预测的原始样本pred_original_sample（对应 DDPM 论文2006.11239 的公式 15）；再按需做 clip 或动态阈值裁剪（_threshold_sample）；随后用公式 (7) 的系数组合pred_original_sample与当前样本x_t得到上一时间步均值pred_prev_sample；最后当t > 0时通过_get_variance计算方差并叠加高斯噪声，完成x_t -> x_{t-1}的完整一步。

第 6 步：把去噪结果转换为图像

>>> from PIL import Image >>> import numpy as np >>> image = (input / 2 + 0.5).clamp(0, 1).squeeze() >>> image = (image.permute(1, 2, 0) * 255).round().to(torch.uint8).cpu().numpy() >>> image = Image.fromarray(image) >>> image

(input / 2 + 0.5).clamp(0, 1)将模型输出的[-1, 1]像素范围映射回[0, 1]，squeeze()去掉批次维度，最后转置通道布局并量化到uint8，交给PIL.Image显示。

进阶：拆解 Stable Diffusion 管道

Stable Diffusion 是一个文本生成图像的*潜在扩散（latent diffusion）*模型。之所以叫"潜在扩散"，是因为它作用于图像的低维表示（潜变量空间）而非真实像素空间，因此内存效率更高。编码器（encoder）把图像压缩成更小的表示，解码器（decoder）再把压缩表示还原成图像；对于文本生成图像模型，还需要一个分词器（tokenizer）和编码器（encoder）来生成文本嵌入（text embeddings）。结合上一节你已经知道的知识，还需要一个 UNet 模型和一个调度器。

可见它比只含一个 UNet 的 DDPM 管道复杂得多——Stable Diffusion 包含三个独立的预训练模型（VAE、UNet、文本编码器），外加分词器与调度器。

[!TIP] 💡 关于 VAE、UNet 与文本编码器各自工作原理的更多细节，可以参考 Hugging Face 官方博客How does Stable Diffusion work?。

现在用from_pretrained方法加载所有这些组件。它们都存放在预训练检查点CompVis/stable-diffusion-v1-4中，每个组件存储在一个独立的子文件夹（subfolder）中：

>>> from PIL import Image >>> import torch >>> from transformers import CLIPTextModel, CLIPTokenizer >>> from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler >>> vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", use_safetensors=True) >>> tokenizer = CLIPTokenizer.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="tokenizer") >>> text_encoder = CLIPTextModel.from_pretrained( ... "CompVis/stable-diffusion-v1-4", subfolder="text_encoder", use_safetensors=True ... ) >>> unet = UNet2DConditionModel.from_pretrained( ... "CompVis/stable-diffusion-v1-4", subfolder="unet", use_safetensors=True ... )

注意到这里每个组件都通过subfolder参数指定了检查点仓库内的子目录（vae、tokenizer、text_encoder、unet、scheduler），这与管道统一加载时使用的model_index.json索引（见 DiffusionPipeline.register_modules 与config_name = "model_index.json"）机制一致。

把默认的PNDMScheduler换成UniPCMultistepScheduler，看看更换调度器有多简单：

>>> from diffusers import UniPCMultistepScheduler >>> scheduler = UniPCMultistepScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")

为了加速推理，把模型移到 GPU——注意调度器没有可训练权重，不需要移动：

>>> torch_device = "cuda" # 或 "mps"、"xpu"、"cpu" >>> vae.to(torch_device) >>> text_encoder.to(torch_device) >>> unet.to(torch_device)

[!NOTE]from_pretrained加载模型后，组件之间的"打包"工作（例如register_modules把模块注册进配置、to()批量迁移设备）正是DiffusionPipeline基类提供的便利，这也是为什么手写管道时需要自己逐个.to(torch_device)。

创建文本嵌入

下一步是对文本做分词并生成嵌入。文本用于对 UNet 模型施加条件（conditioning），引导扩散过程生成与提示词相符的内容。

[!TIP] 💡guidance_scale参数决定生成图像时提示词所占的权重。

可以自由选择任何想要的提示词：

>>> prompt = ["a photograph of an astronaut riding a horse"] >>> height = 512 # Stable Diffusion 默认高度 >>> width = 512 # Stable Diffusion 默认宽度 >>> num_inference_steps = 25 # 去噪步数 >>> guidance_scale = 7.5 # 无分类器引导（classifier-free guidance）的缩放系数 >>> generator = torch.manual_seed(0) # 随机数种子，用于生成初始潜变量噪声 >>> batch_size = len(prompt)

对文本分词并生成嵌入：

>>> text_input = tokenizer( ... prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt" ... ) >>> with torch.no_grad(): ... text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]

还需要生成无条件文本嵌入（unconditional text embeddings），即填充（padding）token 的嵌入。其形状（batch_size与seq_length）必须与条件text_embeddings一致：

>>> max_length = text_input.input_ids.shape[-1] >>> uncond_input = tokenizer([""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt") >>> uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]

把条件与无条件嵌入拼接成一个批次，避免做两次前向传播：

>>> text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

创建随机噪声

接下来生成一些初始随机噪声，作为扩散过程的起点。这就是图像的潜在表示，它会在后续步骤中被逐步去噪。此刻latents的尺寸比最终图像小，这没问题——因为模型稍后会把它变换成最终的 512x512 图像。

[!TIP] 💡 高度和宽度除以 8，是因为vae模型有 3 个下采样层。可以通过下面的代码验证：

2 ** (len(vae.config.block_out_channels) - 1) == 8

>>> latents = torch.randn( ... (batch_size, unet.config.in_channels, height // 8, width // 8), ... generator=generator, ... device=torch_device, ... )

去噪图像

先用初始噪声分布的sigma（噪声缩放值）缩放输入。这是UniPCMultistepScheduler这类改进调度器所必需的：

>>> latents = latents * scheduler.init_noise_sigma

UniPCMultistepScheduler的初始化 将init_noise_sigma设为 1.0（对于DDPMScheduler同样如此），该值表示初始噪声分布的标准差，用于保证初始潜变量的噪声量级与调度器预期一致。

最后一步是创建去噪循环，把latents中的纯噪声逐步变换成提示词所描述的图像。记住，去噪循环需要做三件事：

1. 设定调度器去噪时使用的时间步。
2. 遍历时间步。
3. 在每个时间步，调用 UNet 模型预测噪声残差，并把它交给调度器计算上一噪声样本。

>>> from tqdm.auto import tqdm >>> scheduler.set_timesteps(num_inference_steps) >>> for t in tqdm(scheduler.timesteps): ... # 若进行无分类器引导，则扩展 latents，避免做两次前向传播。 ... latent_model_input = torch.cat([latents] * 2) ... latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t) ... # 预测噪声残差 ... with torch.no_grad(): ... noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample ... # 执行引导 ... noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) ... noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) ... # 计算上一噪声样本 x_t -> x_t-1 ... latents = scheduler.step(noise_pred, t, latents).prev_sample

这段手写循环与 StableDiffusionPipeline 的官方去噪循环 逐行对应：torch.cat([latents] * 2)把潜变量复制两份分别用于条件与无条件预测（即官方代码中的do_classifier_free_guidance分支）；scale_model_input供需要按时间步缩放模型输入的调度器使用（DDPMScheduler中为恒等操作，而 UniPC 等调度器会做实际缩放）；引导公式noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)与官方第 1054-1055 行完全一致；最终同样调用scheduler.step(noise_pred, t, latents)推进一步。

解码图像

最后一步是使用vae把潜在表示解码成图像，并通过sample取得解码输出：

# 缩放并用 vae 解码图像潜变量 latents = 1 / 0.18215 * latents with torch.no_grad(): image = vae.decode(latents).sample

这里的1 / 0.18215是 VAE 的缩放因子——在官方实现中写作latents / self.vae.config.scaling_factor（见 pipeline_stable_diffusion.py 第 1085 行），0.18215正是CompVis/stable-diffusion-v1-4检查点中vae.config.scaling_factor的取值，用于在解码前把潜变量恢复到训练时的量级。

最后把图像转换成PIL.Image查看生成结果：

>>> image = (image / 2 + 0.5).clamp(0, 1).squeeze() >>> image = (image.permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy() >>> image = Image.fromarray(image) >>> image

下一步

从基础管道到复杂管道，你会发现编写自己的扩散系统真正需要的只是一个去噪循环：设定调度器的时间步、遍历它们、交替调用 UNet 模型预测噪声残差并将其交给调度器计算上一噪声样本。

这正是 🧨 Diffusers 的设计初衷——让使用者能够借助模型与调度器，直观而轻松地搭建自己的扩散系统。官方在 DiffusionPipeline 基类 中强调，它负责存储所有组件（模型、调度器、处理器），并提供加载、下载、保存以及设备迁移、进度条开关等通用方法；而当你需要更细粒度控制时，完全可以绕过基类，像本文这样用裸循环组装组件。

接下来的进阶方向：

• 阅读 构建并向 Diffusers 贡献管道，了解如何把你自己组装的管道固化为库内的一等公民组件；
• 浏览现有管道一览，尝试逐个拆解它们，验证你是否能用模型和调度器独立复现出同样的流程。

【免费下载链接】diffusers🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers