使用 Diffusers 训练文生图模型train_text_to_image.py 全流程实战指南【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers本文以 docs/source/en/training/text2image.md 为骨架结合仓库中的训练脚本 examples/text_to_image/train_text_to_image.py 源码系统讲解基于 Stable Diffusion 的文生图text-to-image全量微调流程从环境安装、数据集准备、脚本参数解析到训练循环的底层原理与推理部署。读完本文你将能够在单卡或分布式环境下使用公开数据集或自建数据集完成一次完整的文生图模型微调并掌握 checkpoint 恢复、Min-SNR 加权、EMA 等关键进阶技巧。概述文生图微调能做什么文生图模型如 Stable Diffusion以文本提示text prompt为条件生成对应图像。其训练目标是让模型学会「给定文本 → 生成符合语义的图像」。本文介绍的train_text_to_image.py脚本对模型进行全量微调fine-tune 整个 UNet 权重与仅训练 LoRA 适配器的方式相比可以更充分地适应目标数据分布但也更容易过拟合。⚠️ 注意该脚本是实验性的。全量微调很容易过拟合并可能引发灾难性遗忘catastrophic forgetting等问题。建议针对你的数据集尝试不同的超参数组合以获得最佳效果。硬件要求与显存优化训练模型对硬件有一定要求但通过两个关键开关可以显著降低门槛gradient_checkpointing梯度检查点以更慢的反向传播换取显存把激活值按需重算而非全部驻留显存mixed_precision混合精度将计算精度降为 fp16/bf16减少显存占用并加速。在这两项开启的情况下在单张 24GB 显存的 GPU 上即可完成训练。如果需要更大的 batch size 或更快的训练速度建议使用 30GB 以上显存的 GPU。此外还可以通过启用 xFormers 内存高效注意力memory-efficient attention进一步压缩显存足迹详见 xFormers 优化指南。环境准备与安装从源码安装 Diffusers由于训练脚本迭代频繁官方强烈建议从源码安装并保持更新。在新建的虚拟环境中执行git clone https://github.com/huggingface/diffusers cd diffusers pip install .说明以上命令中的仓库地址为公开的 diffusers 上游仓库若在本地环境中已有本仓库副本可直接在仓库根目录执行pip install .。本文所引用的全部脚本路径均以当前仓库为准。然后进入示例目录并安装该训练脚本所需的依赖cd examples/text_to_image pip install -r requirements.txtrequirements.txt 中的核心依赖包括accelerate0.16.0负责多 GPU/TPU 分布式训练与混合精度配置transformers4.25.1提供 CLIP 文本编码器与分词器datasets2.19.1加载与预处理数据集torchvision图像变换resize、crop、flip 等tensorboard默认的训练日志记录后端peft0.17.0作为 LoRA 训练的底层后端本脚本虽为全量微调但共享环境依赖ftfy、Jinja2文本修复与模型卡渲染辅助库。配置 Accelerate 环境 Accelerate 是帮助你进行多 GPU/TPU 训练或混合精度训练的库它会根据你的硬件和环境自动配置训练方案。初始化方式有三种方式一交互式配置accelerate config按提示选择设备类型、显存、混合精度、分布式策略等。方式二使用默认配置accelerate config default跳过所有交互选项直接生成默认配置。方式三在 Notebook 等非交互环境中配置from accelerate.utils import write_basic_config write_basic_config() 若要训练自有数据集请先阅读 创建训练数据集指南了解如何构造能被训练脚本直接消费的数据集格式如imagefoldermetadata.jsonl或直接上传到 Hub。脚本参数解析parse_args()训练脚本提供了大量参数用于定制训练过程。全部参数及其说明集中在parse_args()函数中源码见 train_text_to_image.py 的parse_args定义约 L201-L525。每个参数都带默认值如训练 batch size、学习率等你也可以在启动命令中通过--参数名 值覆盖。例如要使用 fp16 混合精度加速训练accelerate launch train_text_to_image.py \ --mixed_precisionfp16基础且重要的参数参数说明默认值--pretrained_model_name_or_pathHub 上的模型名或本地预训练模型路径必填无--dataset_nameHub 上的数据集名或本地数据集路径无--image_column数据集中存放图像的列名image--caption_column数据集中存放文本描述的列名text--output_dir训练模型与 checkpoint 的保存目录sd-model-finetuned--push_to_hub是否将训练好的模型推送到 Hub关闭--checkpointing_steps每隔多少步保存一次 checkpoint训练中断时可结合--resume_from_checkpoint断点续训500训练过程核心参数以下参数在源码中都有明确的默认值与语义见parse_args--resolution默认512输入图像统一缩放到的分辨率。注意若使用 768×768 版本的 stable-diffusion-2需要改为768--center_crop/--random_flip先缩放到目标分辨率再做中心裁剪或随机水平翻转作为数据增强--train_batch_size默认16每设备per device的 batch size--num_train_epochs默认100与--max_train_steps二者指定其一--max_train_steps优先--gradient_accumulation_steps默认1累积多少步再执行一次参数更新等效放大 batch size--gradient_checkpointing开启梯度检查点以节省显存--learning_rate默认1e-4--scale_lr按 GPU 数、累积步数与 batch size 缩放学习率--lr_scheduler可选[linear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmup]默认constant--lr_warmup_steps默认500--snr_gammaMin-SNR 损失加权的 γ 值推荐5.0详见下文--use_ema/--offload_ema/--foreach_emaEMA 权重跟踪及显存优化选项--use_8bit_adam使用 bitsandbytes 的 8-bit Adam 优化器--allow_tf32在 Ampere 架构 GPU 上允许 TF32 加速矩阵运算--adam_beta1/--adam_beta2/--adam_weight_decay/--adam_epsilonAdam 优化器超参--max_grad_norm默认1.0梯度裁剪范数--prediction_typeepsilon或v_prediction缺省时沿用 scheduler 配置--report_to日志后端默认tensorboard可选wandb、comet_ml或all--checkpoints_total_limit最多保留的 checkpoint 数量超出自动删除最旧的--resume_from_checkpoint断点续训传 checkpoint 路径或latest自动选择最新--enable_xformers_memory_efficient_attention启用 xFormers 内存高效注意力--noise_offset默认0推荐0.1为低噪声时间步增加偏移噪声改善暗部细节--input_perturbation默认0推荐0.1输入扰动提升小数据量下的收敛质量--validation_prompts与--validation_epochs默认5训练过程中周期性生成验证图像并记录日志--dream_training与--dream_detail_preservation默认1.0启用 DREAM 训练策略--image_interpolation_mode默认lanczosresize 时的插值方式--seed随机种子保证训练可复现--max_train_samples调试时截断训练样本数--dataloader_num_workers默认0数据加载子进程数--hub_model_id/--hub_token推送 Hub 时的仓库名与令牌--tracker_project_name默认text2image-fine-tune日志跟踪项目名。参数校验逻辑位于 parse_args 末尾--dataset_name与--train_data_dir二者必须提供其一否则直接抛出ValueError--non_ema_revision缺省时复用--revision。Min-SNR 加权加速收敛的损失重平衡Min-SNRMinimum Signal-to-Noise Ratio加权策略通过重新平衡损失帮助模型更快收敛。其思想是扩散模型在不同时间步的信噪比差异巨大直接对所有时间步等权求 MSE 会主导低信噪比高噪声区域的损失Min-SNR 通过min(SNR, γ)截断高 SNR 步的权重抑制其对梯度的支配。该脚本既支持预测epsilon噪声也支持v_prediction而 Min-SNR 对两种预测类型都兼容该加权策略仅由 PyTorch 支持对应 PyTorch 版训练脚本。启用方式推荐值 5.0accelerate launch train_text_to_image.py \ --snr_gamma5.0从源码看snr_gamma传入后在训练循环中生效train_text_to_image.py 约 L1022-L1039当--snr_gamma为空时使用普通F.mse_loss(..., reductionmean)否则调用 src/diffusers/training_utils.py 中导出的compute_snr()定义于 L81-L114基于 scheduler 的alphas_cumprod计算每个时间步的 SNR (α/σ)²再取min(snr, γ)作为逐样本损失权重当预测类型为epsilon时权重再除以snr为v_prediction时除以snr 1随后对逐元素 MSE 加权取平均。值得注意的实践结论对较小数据集Min-SNR 的效果可能不如大数据集明显社区在 Weights and Biases 上有不同snr_gamma取值如 1.0 与 5.0的损失面对比实验可据此观察收敛差异该策略在epsilon与v_prediction两种预测目标下均有配套的数学形式详见论文 Section 3.4 与 Section 4.2 的讨论。训练脚本内部结构main() 全流程拆解数据集预处理代码与训练循环集中在main()函数中train_text_to_image.py 约 L528-L1165。若需改造训练脚本这里是主要修改点。整体流程可拆解为以下阶段1. 加载 Scheduler 与 Tokenizer训练脚本首先加载噪声调度器noise scheduler与分词器tokenizer见 L592-L595。你可以在此处替换为其他 schedulernoise_scheduler DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolderscheduler) tokenizer CLIPTokenizer.from_pretrained( args.pretrained_model_name_or_path, subfoldertokenizer, revisionargs.revision )2. 加载 VAE、文本编码器与 UNet脚本在deepspeed_zero_init_disabled_context_manager()的上下文管理器中加载CLIPTextModel与AutoencoderKLL616-L622再加载UNet2DConditionModelL624-L626。随后vae与text_encoder冻结requires_grad_(False)只有unet.train()参与训练L629-L631若开启--use_ema为 UNet 参数创建EMAModelL634-L643训练结束时最终保存的权重使用 EMA 版本若开启--enable_xformers_memory_efficient_attention校验 xformers 版本0.0.16 在部分 GPU 上无法用于训练建议 ≥0.0.17并启用L645-L656。文档中给出的 UNet 加载写法用于 checkpoint 恢复场景load_model UNet2DConditionModel.from_pretrained(input_dir, subfolderunet) model.register_to_config(**load_model.config) model.load_state_dict(load_model.state_dict())这与脚本中load_model_hook的实现一致L684-L693从 checkpoint 目录加载unet子文件夹先同步 config 再加载权重。因为checkpoint 只保存 UNet从 checkpoint 恢复推理时只需单独加载 UNet详见下文推理小节。3. 数据预处理tokenize_captions 与 train_transforms接下来对数据集的文本列与图像列进行预处理tokenize_captions对每个 caption 做分词max_lengthtokenizer.model_max_lengthpaddingmax_length、truncationTrue当某样本含多个 caption列表形式时训练时随机取一个验证时取第一个train_transforms组合Resize(resolution)、中心/随机裁剪、随机水平翻转、ToTensor()与Normalize([0.5], [0.5])插值方式由--image_interpolation_mode控制preprocess_trainL817-L821将二者打包为数据集变换def preprocess_train(examples): images [image.convert(RGB) for image in examples[image_column]] examples[pixel_values] [train_transforms(image) for image in images] examples[input_ids] tokenize_captions(examples) return examplescollate_fnL829-L833将样本堆叠为连续的pixel_values与input_ids张量供 DataLoader 消费。4. 训练循环latent 编码 → 加噪 → 条件嵌入 → 更新参数训练循环是脚本的核心约 L963-L1116单步迭代逻辑如下编码到 latent 空间vae.encode(...).latent_dist.sample()得到潜在表示并乘上vae.config.scaling_factor采样噪声并加噪torch.randn_like(latents)采样噪声若开启--noise_offset则叠加偏移噪声再为每个样本随机采样时间步timesteps通过noise_scheduler.add_noise(latents, noise, timesteps)完成前向扩散加噪若开启--input_perturbation则先对噪声做扰动L967-L990计算文本嵌入text_encoder(batch[input_ids])得到条件向量encoder_hidden_states确定损失目标根据prediction_type选择target noiseepsilon或target noise_scheduler.get_velocity(latents, noise, timesteps)v_prediction若开启--dream_training则调用 src/diffusers/training_utils.py 中的compute_dream_and_update_latents执行 DREAM 策略L996-L1017前向与损失unet(noisy_latents, timesteps, encoder_hidden_states)预测噪声残差按上文是否启用 Min-SNR 计算 MSE 损失反向传播与更新accelerator.backward(loss)→ 梯度裁剪accelerator.clip_grad_norm_(..., args.max_grad_norm)→optimizer.step()→lr_scheduler.step()→ 清零梯度L1045-L1051EMA 更新与日志每个同步梯度步accelerator.sync_gradients后更新 EMA 权重、推进进度条、记录train_loss每--checkpointing_steps步保存checkpoint-{global_step}L1054-L1090。如果希望深入理解「管线、模型与调度器」在去噪过程中的基本模式可以参考官方教程 理解 Pipelines、Models 与 Schedulers。启动训练以 Naruto 数据集为例完成参数调整或确认默认配置后即可启动训练。下面以 lambdalabs/naruto-blip-captionsNaruto 角色 BLIP 自动生成描述数据集为例训练一个能生成火影角色的模型。首先设置环境变量MODEL_NAME与dataset_name分别指向预训练模型与数据集Hub 名称或本地路径export MODEL_NAMEstable-diffusion-v1-5/stable-diffusion-v1-5 export dataset_namelambdalabs/naruto-blip-captions accelerate launch --mixed_precisionfp16 train_text_to_image.py \ --pretrained_model_name_or_path$MODEL_NAME \ --dataset_name$dataset_name \ --use_ema \ --resolution512 --center_crop --random_flip \ --train_batch_size1 \ --gradient_accumulation_steps4 \ --gradient_checkpointing \ --max_train_steps15000 \ --learning_rate1e-05 \ --max_grad_norm1 \ --enable_xformers_memory_efficient_attention \ --lr_schedulerconstant --lr_warmup_steps0 \ --output_dirsd-naruto-model \ --push_to_hub 训练本地数据集时将TRAIN_DIR与OUTPUT_DIR环境变量分别指向数据集目录与模型保存目录并把命令中的--dataset_name替换为--train_data_dir$TRAIN_DIR。本地目录需符合imagefolder结构图片文件 metadata.jsonl每行一条{file_name: ..., text: ...}详见 创建训练数据集指南。多 GPU 训练时在accelerate launch命令中追加--multi_gpu参数即可。参数组合速查场景建议组合单卡 24GB 显存--gradient_checkpointing--mixed_precisionfp16--train_batch_size1--gradient_accumulation_steps4更大 batch / 更快训练使用 30GB 显存 GPU适当增大--train_batch_size进一步省显存追加--enable_xformers_memory_efficient_attention需安装 xformers加速收敛追加--snr_gamma5.0提升生成稳定性追加--use_ema需额外一份全精度参数内存显存紧张可用--offload_ema将 EMA 权重放到 CPU 固定内存关于 EMA 的补充来自 examples/text_to_image/README.mdEMA 通过对模型参数维护指数移动平均来平滑更新噪声、提升性能--foreach_ema使用更快的 foreach 实现--offload_ema将 EMA 权重驻留于 CPU 固定内存每个参数更新步非阻塞地搬回 GPU 更新后再搬回 CPU在主机-设备带宽充足时可做到几乎零额外开销。训练完成后加载模型进行推理训练结束后模型会被保存到--output_dir上述例子为sd-naruto-model。加载微调后的模型进行推理from diffusers import StableDiffusionPipeline import torch pipeline StableDiffusionPipeline.from_pretrained(path/to/saved_model, dtypetorch.float16, use_safetensorsTrue).to(cuda) # 也可用 mps、xpu、cpu image pipeline(promptyoda).images[0] image.save(yoda-naruto.png)训练脚本在结束时会通过StableDiffusionPipeline.from_pretrained(...)组装完整管线并save_pretrained(output_dir)L1125-L1133因此保存目录可直接作为完整 pipeline 加载。从 checkpoint 恢复推理由于训练过程的 checkpoint 只保存 UNet 权重从checkpoint-N恢复推理时需单独加载 UNet 再注入管线import torch from diffusers import StableDiffusionPipeline, UNet2DConditionModel model_path path_to_saved_model unet UNet2DConditionModel.from_pretrained(model_path /checkpoint-N/unet, dtypetorch.float16) pipe StableDiffusionPipeline.from_pretrained(initial model, unetunet, dtypetorch.float16) pipe.to(cuda) image pipe(promptyoda).images[0] image.save(yoda-naruto.png)中断续训若训练因意外中断可在启动命令中追加--resume_from_checkpointlatest自动选择output_dir中最新的checkpoint-*或显式指定--resume_from_checkpointcheckpoint-3000。源码中恢复逻辑见 L928-L953脚本会扫描output_dir下以checkpoint开头的目录并按步数排序通过accelerator.load_state恢复优化器、调度器与模型状态并同步重置global_step与起始 epoch。进阶方向与延伸阅读完成基础微调后可进一步探索LoRA 微调若训练 LoRA 权重对应train_text_to_image_lora.py脚本本文不做展开推理时加载 LoRA 权重的方法可参考 使用 PEFT 进行推理加载 LoRA 权重。LoRA 只需训练新增的低秩分解矩阵权重极小、不易灾难性遗忘可在 T4/V100 等消费级 GPU 上运行且可使用比全量微调高一个量级的学习率如1e-4而非1e-5推理控制关于 guidance scale、prompt weighting 等如何控制生成结果参见 文生图任务指南DREAM 训练通过--dream_training启用以多一次无梯度的 UNet 前向为代价换取更高的模型保真度--dream_detail_preservation默认 1.0控制细节保留因子SDXL 微调仓库还提供了面向 Stable Diffusion XL 的train_text_to_image_sdxl.py与train_text_to_image_lora_sdxl.py脚本详见 examples/text_to_image/README_sdxl.md。总而言之train_text_to_image.py是一个结构清晰、参数完备的文生图全量微调参考实现理解其parse_args的参数面、main()的「加载 → 预处理 → 训练循环 → 保存」链路以及 Min-SNR、EMA、checkpoint 恢复等机制你就能把它改造成适合自己数据集与硬件条件的定制化训练方案。【免费下载链接】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），仅供参考