DiffSynth-Studio 训练框架导读：一文读懂 Diffusion 模型基本原理与 Flow Matching 工程实现

DiffSynth-Studio 训练框架导读一文读懂 Diffusion 模型基本原理与 Flow Matching 工程实现【免费下载链接】DiffSynth-StudioEnjoy the magic of Diffusion models!项目地址: https://gitcode.com/GitHub_Trending/dif/DiffSynth-Studio导读本文是 DiffSynth-Studio 训练系列文档的起点面向想理解训练框架为什么这样构建的开发者。我们将抛弃复杂的随机微分方程用一组简洁的数学定义数据定义、模型定义、迭代公式重构 Diffusion 模型的理论骨架并逐一对照 diffsynth/diffusion/flow_match.py 与 diffsynth/diffusion/loss.py 中的源码说明噪声含量如何定义、迭代去噪如何计算、模型如何训练、现代架构如何三段式组装。读完本文你将掌握阅读 DiffSynth-Studio 任意模型训练脚本所需的理论基础并清楚理解scheduler.add_noise、scheduler.training_target、FlowMatchSFTLoss等核心 API 背后的数学含义。引言从随机噪声到清晰图像的多步迭代Diffusion 模型通过多步迭代式地去噪denoise生成清晰的图像或视频内容。以数据样本 $x_0$ 的生成过程为例在完整的一轮 denoise 过程中我们从随机高斯噪声 $x_T$ 开始通过迭代依次得到 $x_{T-1}$、$x_{T-2}$、$x_{T-3}$、$\cdots$每一步逐渐减少噪声含量最终得到不含噪声的数据样本 $x_0$。这个过程直观易懂但若要深入理解细节需要回答五个问题本文后续小节将逐一展开每一步的噪声含量是如何定义的迭代去噪的计算是如何进行的如何训练这样的 Diffusion 模型现代 Diffusion 模型的架构是什么样的DiffSynth-Studio 如何封装和实现模型训练每一步的噪声含量是如何定义的在 Diffusion 模型的理论体系中噪声的含量由一系列参数 $\sigma_T$、$\sigma_{T-1}$、$\sigma_{T-2}$、$\cdots$、$\sigma_0$ 决定其约束关系如下$\sigma_T1$对应的 $x_T$ 为纯粹的高斯噪声$\sigma_T\sigma_{T-1}\sigma_{T-2}\cdots\sigma_0$迭代过程中噪声含量逐渐减小$\sigma_00$对应的 $x_0$ 为不含任何噪声的数据样本。至于中间 $\sigma_{T-1}$、$\sigma_{T-2}$、$\cdots$、$\sigma_1$ 的具体数值则不是固定的只需满足严格递减的条件即可。只要 $\sigma_t$ 序列确定在中间的任意一步我们都可以直接合成含噪声的数据样本$$x_t(1-\sigma_t)x_0\sigma_t x_T$$这条公式就是 Flow Matching 理论中的数据定义噪声数据 $x_t$ 是干净数据 $x_0$ 与纯高斯噪声 $x_T$ 按 $1-\sigma_t$ 与 $\sigma_t$ 比例的线性插值。它直接对应 DiffSynth-Studio 中FlowMatchScheduler.add_noise的实现见 flow_match.pydef add_noise(self, original_samples, noise, timestep): ... sigma self.sigmas[timestep_id] sample (1 - sigma) * original_samples sigma * noise return sample代码中的original_samples即 $x_0$noise即 $x_T$sigma即 $\sigma_t$与理论公式逐项对应。在 DiffSynth-Studio 中$\sigma_t$ 序列由FlowMatchScheduler针对不同预训练模型生成。FlowMatchScheduler.__init__通过template参数注册了 FLUX.1、Wan、Qwen-Image、FLUX.2、Z-Image、LTX-2、ERNIE-Image、ACE-Step、Ideogram4、Krea-2、MiniMax-H3、LingBot-Video、SenseNova-U1 等十余套时间步方案见 flow_match.py。以默认的set_timesteps_flux为例flow_match.py它先在 $[\sigma_{min},\sigma_{max}]$ 内做torch.linspace等距采样再施加 shift 变换 $\sigma\frac{shift\cdot\sigma}{1(shift-1)\sigma}$ 使采样点向低噪声区域倾斜最后乘以num_train_timesteps1000得到离散的时间步。不同的模板拥有各自的 $\sigma$ 生成策略如 Qwen-Image 使用指数 shift 与动态 shift 长度、Wan 默认 shift5这也印证了文档中中间 $\sigma$ 数值不固定满足递减即可的论断——$\sigma$ 序列本身就是模型定制的重要自由度。迭代去噪的计算是如何进行的去噪模型的输入与输出要理解迭代去噪的计算首先要搞清去噪模型的输入和输出。我们把模型抽象为符号 $\hat \epsilon$其输入通常包含三部分时间步 $t$模型需要理解当前处于去噪过程的哪个阶段含噪声的数据样本 $x_t$模型需要理解要对什么数据进行去噪引导条件 $c$模型需要理解要通过去噪生成什么样的数据样本。其中引导条件 $c$ 由用户输入可以是描述图像内容的文本也可以是勾勒图像结构的线稿图。而模型的输出 $\hat \epsilon(x_t,c,t)$近似地等于 $x_T-x_0$——即整个扩散过程去噪过程的反向过程的方向。这就是 Flow Matching 理论中的模型定义。一步迭代的数学推导在时间步 $t$模型计算出近似的 $x_T-x_0$ 后我们计算下一步的 $x_{t-1}$$$ \begin{aligned} x_{t-1}x_t (\sigma_{t-1} - \sigma_t) \cdot \hat \epsilon(x_t,c,t)\ \approx x_t (\sigma_{t-1} - \sigma_t) \cdot (x_T-x_0)\ (1-\sigma_t)x_0\sigma_t x_T (\sigma_{t-1} - \sigma_t) \cdot (x_T-x_0)\ (1-\sigma_{t-1})x_0\sigma_{t-1}x_T \end{aligned} $$推导结果与时间步 $t-1$ 时的噪声含量定义完美契合。这条公式对应FlowMatchScheduler.step的实现见 flow_match.pydef step(self, model_output, timestep, sample, to_finalFalse, **kwargs): ... sigma self.sigmas[timestep_id] if to_final or timestep_id 1 len(self.timesteps): sigma_ 0 else: sigma_ self.sigmas[timestep_id 1] prev_sample sample model_output * (sigma_ - sigma) return prev_sample其中sample即 $x_t$model_output即 $\hat \epsilon(x_t,c,t)$sigma_为 $\sigma_{t-1}$。注意这里sigma_ - sigma是负数$\sigma$ 单调递减等价于公式中的 $\sigma_{t-1}-\sigma_t$一次前向即完成一步去噪。两个定义 一个迭代公式完备的理论框架这部分可能有点难懂首次阅读时建议跳过不影响后文阅读。完成推导后我们思考一个问题为什么模型的输出要近似地等于 $x_T-x_0$可以设定成其他值吗实际上Diffusion 模型依赖两个定义形成完备的理论。从以上公式中可以提炼出数据定义$x_t(1-\sigma_t)x_0\sigma_t x_T$模型定义$\hat \epsilon(x_t,c,t)x_T-x_0$导出迭代公式$x_{t-1}x_t (\sigma_{t-1} - \sigma_t) \cdot \hat \epsilon(x_t,c,t)$这三个数学公式是完备的——把数据定义和模型定义代入迭代公式可以得到与数据定义吻合的 $x_{t-1}$。在 DiffSynth-Studio 中模型定义对应FlowMatchScheduler.training_target见 flow_match.py它正是训练时模型输出需要拟合的目标def training_target(self, sample, noise, timestep): target noise - sample # 即 x_T - x_0 return target以上是基于 Flow Matching 理论构建的两个定义但 Diffusion 模型也可以用其他两个定义实现。例如早期基于 DDPMDenoising Diffusion Probabilistic Models的模型其两个定义及导出的迭代公式为数据定义$x_t\sqrt{\alpha_t}x_0\sqrt{1-\alpha_t}x_T$模型定义$\hat \epsilon(x_t,c,t)x_T$导出迭代公式$x_{t-1}\sqrt{\alpha_{t-1}}\left(\frac{x_t-\sqrt{1-\alpha_t}\hat \epsilon(x_t,c,t)}{\sqrt{\sigma_t}}\right)\sqrt{1-\alpha_{t-1}}\hat \epsilon(x_t,c,t)$更一般地可以用矩阵描述迭代公式的导出过程。对于任意数据定义和模型定义有数据定义$x_tC_T(x_0,x_T)^T$模型定义$\hat \epsilon(x_t,c,t)C_T^{[\epsilon]}(x_0,x_T)^T$导出迭代公式$x_{t-1}C_{t-1}(C_t,C_t^{[\epsilon]})^{-T}(x_t,\hat \epsilon(x_t,c,t))^T$其中 $C_t$、$C_t^{[\epsilon]}$ 是 $1\times 2$ 的系数矩阵。构造两个定义时需保证矩阵 $(C_t,C_t^{[\epsilon]})^T$ 是可逆的。尽管 Flow Matching 与 DDPM 已被大量预训练模型广泛验证但这并不代表它们是最优方案DiffSynth-Studio 鼓励开发者设计新的 Diffusion 模型理论以取得更好的训练效果——FlowMatchScheduler中高度模块化的set_timesteps_fn注册机制正是为了方便开发者挂载自定义的数据定义与迭代方案。如何训练这样的 Diffusion 模型为什么训练要随机采样单个时间步搞清楚迭代去噪过程后我们考虑如何训练这样的模型。训练过程不同于生成过程如果在训练中保留多步迭代梯度需要经过多步回传带来的时间和空间复杂度是灾难性的。为了提高计算效率训练中我们随机选择某一时间步 $t$ 进行训练。这样每一步训练只需一次模型前向与一次反向代价是单步损失只监督模型的局部行为——这与文档前述理论一致只要模型在任意时间步都能准确预测 $x_T-x_0$多步迭代即可正确完成去噪。训练过程的伪代码从数据集获取数据样本 $x_0$ 和引导条件 $c$随机采样时间步 $t\in(0,T]$随机采样高斯噪声 $x_T\in \mathcal N(0,I)$计算 $x_t(1-\sigma_t)x_0\sigma_t x_T$前向计算 $\hat \epsilon(x_t,c,t)$计算损失函数 $\mathcal L||\hat \epsilon(x_t,c,t)-(x_T-x_0)||_2^2$梯度回传并更新模型参数源码对照FlowMatchSFTLoss 就是这份伪代码的逐行实现DiffSynth-Studio 中这份伪代码由FlowMatchSFTLoss完整实现见 diffsynth/diffusion/loss.py。对照如下def FlowMatchSFTLoss(pipe: BasePipeline, **inputs): ... max_timestep_boundary int(inputs.get(max_timestep_boundary, 1) * len(pipe.scheduler.timesteps)) min_timestep_boundary int(inputs.get(min_timestep_boundary, 0) * len(pipe.scheduler.timesteps)) # 随机采样时间步 t timestep_id torch.randint(min_timestep_boundary, max_timestep_boundary, (1,)) timestep pipe.scheduler.timesteps[timestep_id].to(dtypepipe.torch_dtype, devicepipe.device) # 随机采样高斯噪声 x_T并按数据定义合成 x_t noise torch.randn_like(inputs[input_latents]) * inputs.get(noise_scale, 1.0) inputs[latents] pipe.scheduler.add_noise(inputs[input_latents], noise, timestep) # 计算训练目标 x_T - x_0 training_target pipe.scheduler.training_target(inputs[input_latents], noise, timestep) # 前向噪声预测 models {name: getattr(pipe, name) for name in pipe.in_iteration_models} noise_pred pipe.model_fn(**models, **inputs, timesteptimestep) # MSE 损失并乘以时间步权重 loss torch.nn.functional.mse_loss(noise_pred.float(), training_target.float()) loss loss * pipe.scheduler.training_weight(timestep) return loss这段实现有几个值得注意的工程细节时间步范围可控max_timestep_boundary与min_timestep_boundary将采样范围限制在 $[0,1]$ 的指定比例内用于实现时间步裁剪等训练技巧如跳过接近纯噪声或接近干净数据的时间步。噪声缩放noise_scale参数允许对噪声幅度做缩放适应不同模型的训练需求。时间步权重pipe.scheduler.training_weight(timestep)乘以基础 MSE 损失。该权重由set_training_weight见 flow_match.py基于高斯形式的经验公式生成set_timesteps(1000, trainingTrue)时启用见 flow_match.py用于平衡不同噪声水平时间步的学习强度。现代 Diffusion 模型的架构是什么样的从理论到实践还需要填充更多细节。现代 Diffusion 模型架构已经发展成熟主流架构沿用了 Latent Diffusion 提出的三段式架构包括数据编解码器、引导条件编码器、去噪模型三部分。在 DiffSynth-Studio 中每个模型的 pipeline如StableDiffusionXLPipeline、FluxPipeline、WanVideoPipeline都以这三个子模块为核心组装而成对应文件可参考 diffsynth/models 目录。数据编解码器Data Encoder-Decoder在前文中我们一直将 $x_0$ 称为数据样本而不是图像或视频这是因为现代 Diffusion 模型通常不会直接在图像或视频上进行处理而是用编码器Encoder-解码器Decoder架构的模型——通常是 VAEVariational Auto-Encoders模型——将图像或视频编码为 Embedding 张量得到 $x_0$。数据经过编码器编码后再经解码器解码重建的内容与原始内容近似一致仅有少量误差。为什么要在编码后的 Embedding 张量上处理而不是在图像或视频上直接处理主要原因有两点编码的同时对数据进行了压缩编码后处理的计算量更小编码后的数据分布与高斯分布更相似更容易用去噪模型对数据进行建模。在生成过程中编码器部分不参与计算迭代完成后用解码器部分解码 $x_0$ 即可得到清晰的图像或视频。在训练过程中解码器部分不参与计算仅编码器用于计算 $x_0$。DiffSynth-Studio 中各模型的 VAE 实现即为该角色的落地例如 stable_diffusion_vae.py、flux_vae.py、wan_video_vae.py 等。引导条件编码器Guidance Condition Encoder用户输入的引导条件 $c$ 可能复杂多样需要由专门的编码器模型将其处理成 Embedding 张量。按照引导条件的类型引导条件编码器可以分为以下几类文本类型例如 CLIP、Qwen-VL图像类型例如 ControlNet、IP-Adapter视频类型例如 VAE。前文中的模型 $\hat \epsilon$ 指代此处所有引导条件编码器和去噪模型这一整体。文档把引导条件编码器单独拆分列出是因为这类模型在 Diffusion 训练中通常是冻结的且其输出值与时间步 $t$ 无关因此引导条件编码器的计算可以离线进行例如通过 sft:data_process 任务预先缓存特征训练时直接读取缓存避免重复前向。这一设计在训练框架中有直接体现BasePipeline.freeze_except见 diffsynth/diffusion/base_pipeline.py先将整个 pipeline 置为eval()并冻结所有参数requires_grad_(False)再仅对trainable_models指定的模型调用module.train()与module.requires_grad_(True)解冻。SDXL 全量训练脚本中--trainable_models unet见 stable-diffusion-xl-base-1.0.sh即表示只训练 UNet 去噪模型两个文本编码器与 VAE 均被冻结。去噪模型Denoising Model去噪模型是 Diffusion 模型真正的本体其结构多种多样例如 UNet、DiT模型开发者可以在此结构上自由发挥。DiffSynth-Studio 中既有经典的 UNet 实现如 stable_diffusion_xl_unet.py也有大量 DiT 架构实现如 flux_dit.py、wan_video_dit.py、qwen_image_dit.py 等印证了文档模型结构多种多样的论述。DiffSynth-Studio 如何封装和实现模型训练将上述理论落到工程上DiffSynth-Studio 的封装分为四个层次调度器Scheduler、训练模块DiffusionTrainingModule、损失函数Loss与启动器Runner。1. 调度器理论定义的工程载体FlowMatchSchedulerdiffsynth/diffusion/flow_match.py封装了理论中的全部三个核心公式理论定义源码方法位置数据定义 $x_t(1-\sigma_t)x_0\sigma_t x_T$add_noiseflow_match.py模型定义 $\hat \epsilonx_T-x_0$training_targetflow_match.py迭代公式 $x_{t-1}x_t(\sigma_{t-1}-\sigma_t)\hat\epsilon$stepflow_match.py$\sigma_t$ 序列生成set_timesteps_*系列flow_match.py训练开始前DiffusionTrainingModule.switch_pipe_to_training_mode会调用pipe.scheduler.set_timesteps(1000, trainingTrue)见 diffsynth/diffusion/training_module.py将训练时间步设为 1000 并启用训练权重。2. 训练模块与损失函数DiffusionTrainingModulediffsynth/diffusion/training_module.py负责加载模型、解析模型配置--model_paths/--model_id_with_origin_paths、注入 LoRA、切分计算图split_pipeline_units见 training_module.py以及切换训练模式。各模型训练入口如 examples/stable_diffusion_xl/model_training/train.py通过task_to_loss字典把任务名映射到损失函数sft任务使用本文讲解的FlowMatchSFTLossdirect_distill任务使用DirectDistillLoss见 diffsynth/diffusion/loss.py后者在训练时完整执行多步迭代并拟合教师轨迹是多步迭代的另一种训练范式。3. 训练循环与命令行参数训练启动由 diffsynth/diffusion/runner.py 中的launch_training_task驱动构建 AdamW 优化器默认学习率 1e-4、权重衰减 0.01、DataLoader、调用accelerator.backward(loss)回传梯度、optimizer.step()更新参数——与伪代码的梯度回传并更新模型参数一致。全部命令行参数由 diffsynth/diffusion/parsers.py 统一注册主要包括数据配置add_dataset_base_config/add_image_size_config--dataset_base_path、--dataset_metadata_path、--dataset_repeat、--height、--width、--max_pixels等模型配置add_model_config--model_pathsJSON 格式路径、--model_id_with_origin_paths如stabilityai/stable-diffusion-xl-base-1.0:text_encoder/model.safetensors、--fp8_models、--offload_models、--quant_options等训练配置add_training_config--learning_rate默认 1e-4、--num_epochs默认 1、--trainable_models如dit、vae、text_encoder、--weight_decay默认 0.01、--task默认sft等LoRA 配置add_lora_config--lora_base_model、--lora_target_modules默认q,k,v,o,ffn.0,ffn.2、--lora_rank默认 32等梯度配置add_gradient_config--use_gradient_checkpointing、--gradient_accumulation_steps默认 1等输出与日志配置add_output_config/add_logger_config--output_path、--save_steps、--remove_prefix_in_ckpt默认pipe.dit.以及 TensorBoard / SwanLab / WandB / CSV 日志开关。一个完整的全量 SFT 训练命令示例如下完整脚本见 stable-diffusion-xl-base-1.0.shaccelerate launch examples/stable_diffusion_xl/model_training/train.py \ --dataset_base_path data/diffsynth_example_dataset/stable_diffusion_xl/stable-diffusion-xl-base-1.0 \ --dataset_metadata_path data/diffsynth_example_dataset/stable_diffusion_xl/stable-diffusion-xl-base-1.0/metadata.csv \ --height 1024 \ --width 1024 \ --dataset_repeat 10 \ --model_id_with_origin_paths stabilityai/stable-diffusion-xl-base-1.0:text_encoder/model.safetensors,stabilityai/stable-diffusion-xl-base-1.0:text_encoder_2/model.safetensors,stabilityai/stable-diffusion-xl-base-1.0:unet/diffusion_pytorch_model.safetensors,stabilityai/stable-diffusion-xl-base-1.0:vae/diffusion_pytorch_model.safetensors \ --learning_rate 1e-5 \ --num_epochs 2 \ --trainable_models unet \ --remove_prefix_in_ckpt pipe.unet. \ --output_path ./models/train/stable-diffusion-xl-base-1.0_full \ --use_gradient_checkpointing结合本文理论可以看到--trainable_models unet指定只训练去噪模型UNet文本编码器引导条件编码器与 VAE数据编解码器保持冻结--model_id_with_origin_paths中的四个条目分别对应三段式架构中的三类组件--use_gradient_checkpointing以少量额外计算换取显存节省缓解梯度回传时间空间复杂度高的训练压力。继续深入标准监督训练本文完成了 Diffusion 模型理论框架与 DiffSynth-Studio 训练封装原理的对接。下一文档《标准监督训练》将在此基础上具体讲解数据集的构建、训练脚本的编写、LoRA 与全量训练的实操流程帮助你真正跑通自己的第一个训练任务。【免费下载链接】DiffSynth-StudioEnjoy the magic of Diffusion models!项目地址: https://gitcode.com/GitHub_Trending/dif/DiffSynth-Studio创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考