从零拆解 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 官方教程《Understanding pipelines, models and schedulers》为主线讲解如何拆解DiffusionPipeline将模型Model与调度器Scheduler解耦从零手写一套可运行的扩散采样系统。内容覆盖基础 DDPM 管线的完整去噪循环以及 Stable Diffusion 这种多组件条件扩散模型的逐段实现读完你可以不依赖DiffusionPipeline仅用UNet2DModel、UNet2DConditionModel、AutoencoderKL、CLIP 文本编码器与各类调度器自主组装出任意扩散系统。核心思路管线只是模型 调度器的打包 Diffusers 被设计成一个面向扩散系统的灵活工具箱其核心只有两类构件模型models与调度器schedulers。模型负责学习去噪能力例如 UNet2DModel、UNet2DConditionModel 与 AutoencoderKL调度器负责定义扩散过程的数学规则例如 DDPMScheduler、PNDMScheduler、UniPCMultistepScheduler 等它们位于 src/diffusers/schedulers 目录下。DiffusionPipeline只是把二者打包以提供开箱即用的便利但它完全可以被解绑——这也正是官方教程倡导的把管线拆开用模型和调度器分别组装从而创建属于你自己的扩散系统。整个教程的核心结论只有一句话无论多复杂的扩散系统本质都是一个去噪循环denoising loop。拆解基础管线DDPMPipeline 内部发生了什么用DDPMPipeline生成一张图片只需四行代码 from diffusers import DDPMPipeline ddpm DDPMPipeline.from_pretrained(google/ddpm-cat-256).to(cuda) image ddpm(num_inference_steps25).images[0] image从源码看DDPMPipeline 的构造函数只注册了两个模块class DDPMPipeline(DiffusionPipeline): model_cpu_offload_seq unet def __init__(self, unet: UNet2DModel, scheduler: DDPMScheduler): super().__init__() self.register_modules(unetunet, schedulerscheduler)也就是说这个管线内部只有UNet2DModel模型和DDPMScheduler调度器。它的工作方式是生成一张与期望输出尺寸相同的随机噪声图让模型反复过这张图多次在每一个 timestep模型预测出noise residual噪声残差调度器利用该残差预测出噪声更少的上一张图如此循环直到跑完设定的推理步数。对照 pipeline_ddpm.py 的__call__实现可以看到这 4 行代码背后其实就是一个标准去噪循环# 1. 采样高斯噪声作为起点 image randn_tensor(image_shape, generatorgenerator, devicedevice, dtypeself.unet.dtype) # 2. 设置推理步数 self.scheduler.set_timesteps(num_inference_steps) for t in self.progress_bar(self.scheduler.timesteps): # 3. 模型预测噪声 model_output self.unet(image, t).sample # 4. 调度器计算上一时刻图像 x_t - x_t-1 image self.scheduler.step(model_output, t, image, generatorgenerator).prev_sample # 5. 将张量归一化到 [0,1] 并转为 PIL 图像 image (image / 2 0.5).clamp(0, 1)这个模式就是所有扩散系统的通用骨架。下面我们把它手动拆开写一遍。手写去噪循环DDPM 六步走第 1 步加载模型与调度器与直接加载管线不同这里把两个组件分开加载 from diffusers import DDPMScheduler, UNet2DModel scheduler DDPMScheduler.from_pretrained(google/ddpm-cat-256) model UNet2DModel.from_pretrained(google/ddpm-cat-256).to(cuda)UNet2DModel是无条件图像去噪模型输出与输入同尺寸DDPMScheduler负责噪声调度。第 2 步设置去噪的 timestep 数量 scheduler.set_timesteps(50)set_timesteps会在调度器上生成一条等间距的 timestep 序列其长度等于推理步数。可以在 scheduling_ddpm.py 中看到其实现逻辑它会根据timestep_spacing配置linspace/leading/trailing从训练时的总步数默认 1000中等间隔抽取推理步并反转成严格递减的序列同时校验num_inference_steps不得超过训练步数num_train_timesteps。第 3 步查看生成的 timesteps 张量 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])每个元素代表模型对图像去噪的一个时间刻度从接近完全噪声的 980 一直递减到 0。后续去噪循环就是迭代这个张量。第 4 步生成与目标输出同形状的随机噪声 import torch sample_size model.config.sample_size noise torch.randn((1, 3, sample_size, sample_size), devicecuda)model.config.sample_size来自 UNet 的配置文件定义了模型期望的输入分辨率本案例为 256通道数 3 对应 RGB。第 5 步编写去噪循环 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循环每一步做两件事模型前向model(input, t)调用UNet2DModel.forward预测当前噪声图像中的噪声残差调度器反向一步scheduler.step(noisy_residual, t, input)返回prev_sample即上一时刻噪声更少的图像作为下一轮模型输入。DDPMScheduler.step的内部实现见 scheduling_ddpm.py揭示了其数学本质先由预测的噪声反推预测原始样本pred_original_sample对应 DDPM 论文公式 15依据prediction_type为epsilon/sample/v_prediction采用不同反推公式再按公式 (7) 组合出上一时刻样本的均值最后按variance_type决定是否叠加随机方差项。整个过程即是对扩散 SDE 的反向求解。第 6 步把去噪结果转成图片 from PIL import Image import numpy as np image (input / 2 0.5).clamp(0, 1) image image.cpu().permute(0, 2, 3, 1).numpy()[0] image Image.fromarray((image * 255).round().astype(uint8)) image先将张量从[-1, 1]归一化到[0, 1]再转换通道顺序NCHW → NHWC并映射到uint8的 0~255 范围最后包装成 PIL 图像。至此整个去噪循环完成。这个模式可以复用到任意扩散系统中——差异只在于组件的种类和数量。拆解 Stable Diffusion 管线多组件条件扩散模型Stable Diffusion 是 text-to-image 的latent diffusion潜在扩散模型。它不是在真实像素空间而是在图像的低维潜在表示上工作因而内存效率更高。其组件结构为编码器VAE Encoder把图像压缩成更小的潜在表示解码器VAE Decoder把压缩表示还原成图像Tokenizer 文本编码器为 text-to-image 生成文本嵌入text embeddingsUNet 模型与调度器在潜在空间中执行去噪。相比只有 UNet 的 DDPM 管线Stable Diffusion 包含三个独立的预训练模型VAE、文本编码器与 UNet。加载方法相同——用from_pretrained从预训练检查点的各子文件夹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, subfoldervae) tokenizer CLIPTokenizer.from_pretrained(CompVis/stable-diffusion-v1-4, subfoldertokenizer) text_encoder CLIPTextModel.from_pretrained(CompVis/stable-diffusion-v1-4, subfoldertext_encoder) unet UNet2DConditionModel.from_pretrained(CompVis/stable-diffusion-v1-4, subfolderunet)调度器可以随意替换把默认的PNDMScheduler换成UniPCMultistepScheduler只需一行代码体现了调度器即插即用的设计 from diffusers import UniPCMultistepScheduler scheduler UniPCMultistepScheduler.from_pretrained(CompVis/stable-diffusion-v1-4, subfolderscheduler) 调度器没有可训练权重可以留在 CPU 上而模型有训练权重为了加快推理应显式移到 GPU torch_device cuda vae.to(torch_device) text_encoder.to(torch_device) unet.to(torch_device)生成文本嵌入Text Embeddings文本嵌入用于条件化conditionUNet把扩散过程牵引向与输入提示词相似的方向。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)对文本做 tokenize 并编码为嵌入 text_input tokenizer( ... prompt, paddingmax_length, max_lengthtokenizer.model_max_length, truncationTrue, return_tensorspt ... ) 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必须与条件嵌入一致 max_length text_input.input_ids.shape[-1] uncond_input tokenizer([] * batch_size, paddingmax_length, max_lengthmax_length, return_tensorspt) uncond_embeddings text_encoder(uncond_input.input_ids.to(torch_device))[0]将条件与无条件嵌入拼接成一个 batch从而在一次前向中同时计算两者避免两次前向 text_embeddings torch.cat([uncond_embeddings, text_embeddings])生成随机噪声初始潜在变量扩散过程的起点是随机噪声也就是图像潜在表示的初始形态它将逐步被去噪。此时潜在图尺寸小于最终图片——这是有意为之因为 VAE 有 3 个下采样层高度和宽度被除以 8 可以运行下面的表达式验证下采样倍数2 ** (len(vae.config.block_out_channels) - 1) 8 latents torch.randn( ... (batch_size, unet.config.in_channels, height // 8, width // 8), ... generatorgenerator, ... devicetorch_device, ... )图像去噪含无分类器引导先用初始噪声分布的标准差sigma即init_noise_sigma缩放输入。这是UniPCMultistepScheduler这类改进型调度器所必需的见 scheduling_unipc_multistep.py 中init_noise_sigma的定义对 DDPM 其值恒为 1.0见 scheduling_ddpm.py latents latents * scheduler.init_noise_sigma去噪循环需要完成三件事① 设置调度器的 timesteps② 迭代 timesteps③ 在每个 timestep 调用 UNet 预测噪声残差交给调度器计算上一时刻的噪声样本。加上无分类器引导CFG后完整代码如下 from tqdm.auto import tqdm scheduler.set_timesteps(num_inference_steps) for t in tqdm(scheduler.timesteps): ... # 做 classifier-free guidance 时把 latents 复制一份避免两次前向 ... latent_model_input torch.cat([latents] * 2) ... latent_model_input scheduler.scale_model_input(latent_model_input, timestept) ... # 预测噪声残差 ... with torch.no_grad(): ... noise_pred unet(latent_model_input, t, encoder_hidden_statestext_embeddings).sample ... # 执行 guidance ... 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这段循环中有三个值得留意的点torch.cat([latents] * 2)把潜在变量复制成两份分别对应无条件分支与条件分支与前面拼接后的text_embeddings同样两份在 batch 维度对齐scheduler.scale_model_input用于兼容那些需要按当前 timestep 缩放模型输入的调度器对DDPMSchedulerscheduling_ddpm.py与UniPCMultistepSchedulerscheduling_unipc_multistep.py而言缩放因子为 1即原样返回接口的存在是为了保证调度器之间的互换性CFG 公式noise_pred_uncond guidance_scale * (noise_pred_text - noise_pred_uncond)是经典的无分类器引导形式guidance_scale越大图像与提示词的贴合度越高。用 VAE 解码图像最后一步把潜在表示从 VAE 的缩放系数还原再交给vae.decode解码成图像# 缩放 latents 并用 vae 解码 latents 1 / 0.18215 * latents with torch.no_grad(): image vae.decode(latents).sample 这里的0.18215是 Stable Diffusion VAE 的默认scaling_factor缩放因子可以在 pipeline_stable_diffusion.py 中看到官方实现同样使用1 / self.vae.config.scaling_factor * latents的写法。该默认值定义在 convert_from_ckpt.py 中。最后把张量转成PIL.Image查看生成结果 image (image / 2 0.5).clamp(0, 1) image image.detach().cpu().permute(0, 2, 3, 1).numpy() images (image * 255).round().astype(uint8) pil_images [Image.fromarray(image) for image in images] pil_images[0]总结你只需要一个去噪循环从最简单的 DDPM 到复杂的 Stable Diffusion可以发现编写自己的扩散系统全部核心就是一个去噪循环。这个循环需要设置调度器的 timesteps迭代 timesteps交替执行——调用 UNet 模型预测噪声残差把残差交给调度器计算上一时刻的噪声样本。这正是 Diffusers 的设计初衷让开发者能够基于模型与调度器直观、轻松地搭建属于自己的扩散系统。如果想要继续深入可以尝试阅读 src/diffusers/pipelines/ddpm/pipeline_ddpm.py 与 src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py对照官方实现检查自己手写的循环是否正确浏览 src/diffusers/schedulers 目录下的全部调度器DDPM、PNDM、UniPC、DDIM、DPM-Solver 等尝试把同一个去噪循环中的调度器逐个替换体会换调度器如换插件的设计哲学在 src/diffusers/models/unets 中研究不同 UNet 变体UNet2DModel与UNet2DConditionModel的输入差异理解条件扩散模型如何通过encoder_hidden_states接收文本条件参考官方文档 docs/source/ko/using-diffusers 目录下的其他教程进一步了解管线构建与社区贡献的完整流程。【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考