开发者指南

相关源文件

• README.md
• ldm/models/diffusion/ddpm.py
• main.py

目的与范围

本指南为希望扩展、修改或深入理解 Stable Diffusion 实现细节的开发者提供了技术信息。它涵盖了项目结构、核心组件、训练流程以及模型自定义指南。有关基本用法说明，请参阅使用指南，有关模型架构本身的详细信息，请参阅模型架构。

项目结构

Stable Diffusion 的组织结构清晰，将系统的不同组件分开。理解这种组织方式对于希望扩展或修改模型的开发者至关重要。

来源： README.md1-52 ldm/models/diffusion/ddpm.py1-33

关键目录

目录	目的
scripts/	包含用于图像生成的入口脚本（txt2img.py, img2img.py）
ldm/	潜扩散模型的核心实现
ldm/models/	模型实现，包括扩散模型
ldm/modules/	构建块，如 UNet、注意力等
ldm/data/	数据集处理代码
configs/	模型变体的配置文件
models/	用于存储模型检查点的目录

该存储库遵循模块化设计，可以独立扩展系统的不同部分。

核心组件与架构

Stable Diffusion 基于分层类结构构建，实现了扩散模型框架。理解这些类及其关系对于扩展模型至关重要。

来源： ldm/models/diffusion/ddpm.py44-424 ldm/models/diffusion/ddpm.py1395-1421

关键类

1. DDPM（ldm/models/diffusion/ddpm.py）：实现去噪扩散概率模型的基类。它扩展了 PyTorch Lightning 的LightningModule，并提供了

   • 通过register_schedule()进行的噪声调度管理
   • 前向和反向扩散过程
   • 训练的损失计算

2. LatentDiffusion（ldm/models/diffusion/ddpm.py）：实现 Stable Diffusion 的主类，扩展 DDPM 以在潜在空间中工作

   • 像素空间与潜在空间之间的编码/解码
   • 条件机制（文本、图像等）
   • 采样接口

3. DiffusionWrapper（ldm/models/diffusion/ddpm.py）：封装 UNet 模型并根据条件键（concat、交叉注意力等）处理条件输入。

4. 采样器：各种采样算法作为独立类实现（DDIM、PLMS、DPM-Solver），每个类都有专门的采样策略。

来源： ldm/models/diffusion/ddpm.py424-1393

开发工作流

设置开发环境

设置开发环境

1. 克隆仓库

2. 创建并激活 conda 环境

3. 以开发模式安装包

来源： README.md22-38

修改模型

扩展或修改模型时，请关注这些常见的扩展点

1. 新的条件方法：扩展 LatentDiffusion 并修改 get_learned_conditioning() 方法。

2. 自定义采样器：创建一个实现采样接口的新采样器类，遵循现有采样器（如 DDIMSampler）的模式。

3. UNet 架构更改：修改 ldm/modules/diffusionmodules 目录中的 UNet 架构。

4. 自定义训练目标：在 LatentDiffusion 或 DDPM 中重写 p_losses() 方法。

训练流程

Stable Diffusion 在其训练流程中使用 PyTorch Lightning，这使得定制训练过程相对容易。

来源： main.py418-459 main.py513-569

配置系统

Stable Diffusion 使用基于 OmegaConf 的分层配置系统，允许灵活的模型配置。

1. 基础配置：configs/ 目录中的 YAML 文件定义了基础模型配置。
2. 命令行覆盖：可以通过命令行参数覆盖参数。
3. 配置合并：可以按顺序合并多个配置文件。

配置示例结构

来源： main.py513-517 main.py425-458

训练模型

训练或微调模型

1. 准备您的配置文件，指定模型架构、数据集和训练参数。
2. 运行训练脚本

训练脚本支持各种选项，包括

• --base：基础配置文件路径
• -t / --train：启用训练模式
• --gpus：指定要使用的 GPU
• --scale_lr：根据批次大小和 GPU 缩放学习率

来源： main.py717-724 main.py36-122

检查点和日志记录

训练流程包括

1. 检查点：根据监控的指标自动保存模型检查点。
2. 日志记录：支持 TensorBoard 和 WandB 日志记录。
3. 图像可视化：训练期间定期生成样本图像。

检查点保存在 logs/{run_name}/checkpoints/ 目录中。

来源： main.py569-613 main.py289-392

扩展模型

添加新采样器

实现新的采样方法

1. 创建一个实现所需接口的新采样器类（sample() 方法）。
2. 通过在适当的方法中调用它，将其与 LatentDiffusion 集成。

自定义采样器的示例骨架

实现自定义条件

添加新的条件方法

1. 修改 LatentDiffusion 中的 get_learned_conditioning() 方法。
2. 更新 DiffusionWrapper 中的 forward() 方法以处理新的条件类型。
3. 如有必要，更新 UNet 模型以处理新的条件信息。

来源： ldm/models/diffusion/ddpm.py551-562 ldm/models/diffusion/ddpm.py1402-1420

安全特性

Stable Diffusion 实现了两个关键的安全功能，开发者应了解这些功能

1. 安全检查器：一个用于过滤可能包含 NSFW 内容的输出的分类器。
2. 隐形水印：一个在生成的图像中嵌入隐形水印的系统。

这些功能已集成到采样脚本中（txt2img.py 和 img2img.py）。

来源： README.md86-91

使用安全功能

开发新扩展或修改时

1. 保持安全检查器集成，以实现负责任的 AI 部署。
2. 在推理过程中使用安全检查器，方法是将生成的图像传递给它。
3. 考虑为生成的内容添加水印以保持可追溯性。

安全检查器可以针对特定的研究目的禁用，但建议在面向公众的应用程序中保持启用状态。

调试和故障排除

调试模型时

1. 使用 --debug 标志启用调试模式，该标志会激活事后调试。
2. 使用不同详细级别的日志记录来跟踪模型的行为。
3. 检查模型检查点与您的代码版本之间的兼容性。

来源： main.py85-93 main.py726-731

贡献指南

当为 Stable Diffusion 做出贡献时

1. 遵循现有的代码风格和架构模式。
2. 为新功能添加测试。
3. 彻底记录您的更改。
4. 在可能的情况下，确保与现有模型检查点的向后兼容性。

遵循这些指南，您可以有效地扩展和修改 Stable Diffusion 模型以满足您的特定需求，同时保持与核心代码库的兼容性。