当MMEngine遇到Diffusers,带你轻松玩转扩散模型 !

2023-09-13 16:26:39 浏览数 (1)

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

代码语言:javascript复制
pip install openmim
pip install git https://github.com/okotaku/diffengine.git

使用预定义配置进行训练

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

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

代码语言:javascript复制
mim train diffengine stable_diffusion_v15_dreambooth_lora_dog.py

监控进度和获取结果

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

代码语言:javascript复制
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 进行推理。

代码语言:javascript复制
import torch
from diffusers import DiffusionPipeline
checkpoint = '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,谢谢您的阅读。

0 人点赞