Flash Diffusion API参考Embedder、UNet、VAE与Adapter核心组件完全清单【免费下载链接】flash-diffusionFlash Diffusion — accelerating conditional diffusion models (AAAI 2025 Oral)项目地址: https://gitcode.com/gh_mirrors/fl/flash-diffusionFlash Diffusion 是 AAAI 2025 Oral 论文Flash Diffusion: Accelerating Any Conditional Diffusion Model的官方实现它通过一步蒸馏让 SD1.5、SDXL、Pixart 等条件扩散模型只需4 步NFE就能出图。本文以 API 参考的角度为你完整梳理src/flash/models/下的四大核心组件——Embedder、UNet、VAE 与 Adapter 的类名、职责与配置文件位置帮你快速读懂这份扩散模型加速源码 。组件总览一张表看懂 Flash Diffusion 架构Flash Diffusion 的训练流程可概括为Embedder 生成条件 → UNet/Transformer 预测噪声 → VAE 做像素与潜变量互转 → Adapter 注入结构条件。官方方法示意图如下组件核心类源码位置作用EmbedderClipEmbedder、T5TextEmbedder、TimestepsEmbedder、TorchNNEmbeddersrc/flash/models/embedders/将文本、时间步、图像编码为 UNet 可用的条件UNetDiffusersUNet2DCondWrapper、DiffusersUNet2DWrappersrc/flash/models/unets/unet.pySD1.5 / SDXL 的噪声预测骨干TransformerDiffusersTransformer2DWrappersrc/flash/models/transformers/tranformers.pyDiT 骨干Pixart-α的去噪器VAEAutoencoderKLDiffuserssrc/flash/models/vae/autoencoderKL.py像素 ↔ 潜变量编解码支持分块解码AdapterDiffusersT2IAdapterWrappersrc/flash/models/adapters/t2i_adapter.pyCanny 等结构条件注入蒸馏训练器FlashDiffusionsrc/flash/models/flash/flash_diffusion_model.py组织教师-学生蒸馏、DMD 损失与 GAN 损失Embedder 组件条件编码器完全清单Embedder 家族继承自 BaseConditioner所有 Embedder 的输出维度会自动映射为三种条件类型DIM2CONDITIONING2 维 →vector向量条件如 pooled 文本嵌入3 维 →crossattn交叉注意力条件如文本序列4 维 →concat与潜变量拼接的图像条件具体成员及位置类文件说明ClipEmbedder/ClipEmbedderWithProjectionclip_embedder_model.pyCLIP 文本编码器SD 系标配可配置返回 pooled 向量T5TextEmbeddert5_embedder_model.pyT5 文本编码器用于 SDXL 的 text_encoder_2 等TimestepsEmbeddertimesteps_embedding.py时间步编码TorchNNEmbeddertorch_nn/embedders.py用任意 PyTorch 模块如 Conv2d处理图像输入多个 Embedder 通过 ConditionerWrapper 组合成统一的条件字典并支持unconditional_conditioning_rate按概率丢弃条件用于 classifier-free guidance 训练。UNet 与 Transformer两种去噪骨干Flash Diffusion 兼容 UNet 与 DiT 两类骨干均封装了freeze()方法教师模型训练时冻结参数DiffusersUNet2DCondWrapperunet.pySD1.5 / SDXL 使用。forward接收sample、timestep、conditioning三类输入并支持down_intrablock_additional_residuals注入 Adapter 残差。DiffusersUNet2DWrapper不带交叉注意力的 UNet2D适合轻量场景。DiffusersTransformer2DWrapperDiT 版 Transformer用于 Pixart-α 蒸馏。VAE 组件AutoencoderKLDiffusers 详解AutoencoderKLDiffusers 是对 diffusersAutoencoderKL的封装关键 APIencode(x)图像编码为潜变量带缩放系数分批复用显存友好decode(z)潜变量解码回像素自动分块tiling大分辨率推理不爆显存配置类AutoencoderKLDiffusersConfig只需传入 Hugging Face 仓库名如 SDXL VAE并支持tiling_size、tiling_overlap参数。初始化时它会用一张 32×32 的探针图自动探测下采样因子与潜通道数新手无需手动配置 ✅。Adapter 组件DiffusersT2IAdapterWrapperDiffusersT2IAdapterWrapper 封装了 diffusers 的T2IAdapter用于 Canny 等结构条件。它在训练/采样时输出down_intrablock_additional_residuals按adapter_conditioning_scale缩放后注入 UNet 的 down 块中。官方蒸馏出的 Canny Adapter 效果如下训练入口examples 目录与 YAML 配置四大训练脚本与配置一一对应改好数据路径即可启动蒸馏脚本配置蒸馏目标train_flash_sd.pyflash_sd.yamlSD1.5train_flash_sdxl.pyflash_sdxl.yamlSDXLtrain_flash_pixart.pyflash_pixart.yamlPixart-αDiTtrain_flash_canny_adapter.pyflash_canny_adapter.yamlCanny Adapter以 flash_sd.yaml 为例关键超参包括K教师采样步数调度、DISTILL_LOSS_TYPEL1/L2/LPIPS、USE_DMD_LOSS分布匹配蒸馏、ADVERSARIAL_LOSS_SCALEGAN 损失权重按阶段递增与TIMESTEP_DISTRIBUTION: mixture可调时间步混合分布这是加速的核心技巧。数据格式遵循 webdataset每个样本含jpg图像与带caption的json。实际效果4 NFE 出图速度一览蒸馏后的模型用 4 步即可达到接近 50 步的质量。Flash SD由 SD1.5 教师蒸馏的结果Flash SDXL 与 Flash PixartDiT 骨干同样支持 4 步推理且 LoRA 兼容可继续做风格微调。新手上手三步走装环境Python ≥ 3.10依次pip install -r requirements.txt与pip install -e .跑推理用 Hugging Face pipeline 加载 Flash 权重LoRA 形式把num_inference_steps设为 4、guidance_scale设为 0 即可改数据再蒸馏修改 YAML 中的SHARDS_PATH_OR_URLS指向自己的 webdataset运行对应训练脚本仅需数小时 GPU 时间。测试用例可参考 tests/test_embedders/ 与 tests/test_unet/数据管道测试在 tests/test_dataset/帮助你验证环境配置是否正确。掌握以上组件清单后你就可以按需替换 UNet/Transformer、自由组合 Embedder 条件快速复现或扩展 Flash Diffusion 的蒸馏实验 ⚡。【免费下载链接】flash-diffusionFlash Diffusion — accelerating conditional diffusion models (AAAI 2025 Oral)项目地址: https://gitcode.com/gh_mirrors/fl/flash-diffusion创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考