当MMEngine遇到Diffusers，带你轻松玩转扩散模型 ！

OpenMMLab 社区贡献者@takuoko 近期开源了一个用于训练最先进的扩散模型的工具箱 DiffEngine，DiffEngine 使用了 Diffusers 和 MMEngine。

本文是@takuoko 关于 DiffEngine 的具体介绍，原文链接：https://medium.com/@to78314910/introduction-to-diffengine-cad272e900c4

OpenMMLab 经授权后发布，以下是原文译文。

近期 takuoko 发布了 DiffEngine，这是一个开源工具箱，旨在简化最先进的扩散模型的训练。扩散模型是一类生成模型，可以从随机噪声中合成逼真的图像、文本、音频和视频。

来源：https://stability.ai/stablediffusion

在 Diffusers 和 MMEngine 的帮助下，DiffEngine 使您能够自定义和优化您的扩散模型。无论您是该领域的新手还是专家，DiffEngine 都可以通过简单统一的接口帮助您创建和改进扩散模型。

DiffEngine GitHub：

https://github.com/okotaku/diffengine

DiffEngine 文档：

https://diffengine.readthedocs.io/en/latest/

关于 MMEngine

MMEngine 是一个基于 PyTorch 实现的，用于训练深度学习模型的基础库。它功能强大，是所有 OpenMMLab 项目的训练引擎，涵盖了广泛的研究领域。

借助 MMEngine，您可以用最少的代码轻松训练最先进的模型，用高级功能定制和优化您的模型，并用各种工具评估和可视化您的结果。

如果您想了解有关 MMEngine 的更多信息，请访问：

MMEngine GitHub：

https://github.com/open-mmlab/mmengine

OpenMMLab ：

https://openmmlab.com/

关于 Diffusers

Diffusers 可帮助您使用扩散模型生成逼真多样的图像、音频和其他类型的数据样本。

Diffusers 提供以下功能：

• 仅使用几行代码就能进行推理的扩散流程

• 各种预训练模型

如果您想了解更多关于 Diffusers 的信息，请访问：

Diffusers GitHub：

https://github.com/huggingface/diffusers

Huggingface：

https://huggingface.co/huggingface

DiffEngine 特点

• 训练最先进的扩散模型

DiffEngine 支持最先进的扩散模型，这些模型已取得了令人惊讶的效果。您可以使用 Stable Diffusion、Stable Diffusion XL、DreamBooth、LoRA 等模型生成高质量的图像。您还可以使用 ControlNet 来训练条件扩散模型，该模型可根据文本提示或其他图像生成图像。

• 统一的配置系统和模块设计

DiffEngine 借助 MMEngine，可为您的项目提供统一的配置系统和模块化设计。您可以使用预定义配置或创建自己的配置，轻松调整扩散模型的超参数、损失函数、数据集和其他设置。您还可以重复使用和组合不同的模块，创建复杂而灵活的模型。

• 使用 diffusers.pipeline 进行推理

DiffEngine 通过 diffusers.pipeline 简化了将训练过的扩散模型部署到推理任务中的过程。您只需加载您的模型，然后用它根据给定条件生成样本即可。

DiffEngine 用法

在安装 DiffEngine 之前，请确保已根据官方指南（

https://pytorch.org/get-started/locally/

）成功安装了 PyTorch。

安装 DiffEngine

pip install openmimpip install git+https://github.com/okotaku/diffengine.git

使用预定义配置进行训练

在 DiffEngine 仓库的 configs 目录下可以找到各种预定义配置。例如，如果您希望使用 Stable Diffusion 算法训练 DreamBooth模型，访问文件configs/stable_diffusion_dreambooth/stable_diffusion_v15_dreambooth_lora_dog.py。

要使用所选配置进行训练，请打开终端并运行以下命令：

mim train diffengine stable_diffusion_v15_dreambooth_lora_dog.py

监控进度和获取结果

启动训练后，您可以跟踪其进度。例如，如果您使用 stable_diffusion_v15_dreambooth_lora_dog 配置，那么训练的输出将位于 work_dirs/stable_diffusion_v15_dreambooth_lora_dog 目录。

work_dirs/stable_diffusion_v15_dreambooth_lora_dog├── 20230802_033741| ├── 20230802_033741.log # log file| └── vis_data| ├── 20230802_033741.json # log json file| ├── config.py # config file for each experiment| └── vis_image # visualized image from each step├── step999| └── pytorch_lora_weights.bin # weight for inferencing with diffusers.pipeline├── iter_1000.pth # checkpoint from each step├── last_checkpoint # last checkpoint, it can be used for resuming└── stable_diffusion_v15_dreambooth_lora_dog.py # latest config file

下面提供了一个输出示例：

利用 diffusers.pipeline 进行推理

完成训练后，您只需指定模型路径，即可使用 diffusers.pipeline 进行推理。

import torchfrom diffusers import DiffusionPipelinecheckpoint = 'work_dirs/stable_diffusion_v15_dreambooth_lora_dog/step999'prompt = 'A photo of sks dog in a bucket'pipe = DiffusionPipeline.from_pretrained( 'runwayml/stable-diffusion-v1–5', torch_dtype=torch.float16)pipe.to('cuda')pipe.load_lora_weights(checkpoint)image = pipe( prompt, num_inference_steps=50,).images[0]image.save('demo.png')

示例 Notebook

如果您想更深入地了解 DiffEngine，您可以在 Colaboratory 上运行示例 Notebook。该 Notebook 演示了使用 SDV1.5 和 SDV2.1 DreamBooth 配置进行训练的过程。 

Colab 链接：https://colab.research.google.com/gist/okotaku/cfca8baf57938f193c212dcb43c02fac/diffengine-example.ipynb

DiffEngine 展望

Diffusers 提供各种最先进的扩散模型，而 MMEngine 则提供出色的统一配置系统和模块化设计。通过利用这两个优势，DiffEngine 将为 Diffusers 创建一个统一的训练库。

DiffEngine 还打算支持其他的训练方法，例如：

• Distill SD

• Instruct Pix2Pix

• ControlNet Small

• T2I Adapter

• IP Adapter

此外，DiffEngine 还计划开发 LoRAs。DiffEngine 正在寻找对这些开发项感兴趣的核心开发人员。

takuoko 的链接：https://twitter.com/takuoko1，谢谢您的阅读。

@tokuoko 是 OpenMMLab SIG 社区的活跃贡献者之一，为 OpenMMLab 共建提交了诸多 PR。如果你也想成为 MMSIG 的一员，与大佬一起交流，欢迎扫描下面的二维码加入我们！

2023-09-07

点击下方“阅读原文”直达DiffEngine