DragDiffusion：利用扩散模型实现基于点的交互式图像编辑

史宇俊    薛楚辉    刘俊豪    潘嘉纯   
颜汉书    张文清    谭文杰    白松

免责声明

本项目为研究项目，而非商业产品。用户可以自由使用该工具生成图像，但应遵守当地法律法规，并以负责任的方式使用。开发者对用户的任何不当使用不承担任何责任。

新闻与更新

• [1月29日] 更新以支持 diffusers==0.24.0！
• [10月23日] DragBench 的代码和数据已发布！详情请查看“drag_bench_evaluation”下的 README。
• [10月16日] 在拖拽生成图像时集成 FreeU。
• [10月3日] 加速真实图像编辑时的 LoRA 训练。（现在在 A100 上仅需约 20 秒！）
• [9月3日] v0.1.0 发布。
  • 支持 拖拽扩散生成的图像。
  • 引入新的引导机制，显著提升拖拽结果的质量。（受 MasaCtrl 启发）
  • 支持任意长宽比的图像拖拽。
  • 添加对 DPM++Solver 的支持（生成图像）。
• [7月18日] v0.0.1 发布。
  • 将 LoRA 训练集成到用户界面中。无需使用训练脚本，所有操作均可在界面中便捷完成！
  • 优化了用户界面布局。
  • 支持使用更好的 VAE 来处理眼睛和面部细节（参见 此链接）。
• [7月8日] v0.0.0 发布。
  • 实现了 DragDiffusion 的基本功能。

安装

建议在配备 NVIDIA GPU 和 Linux 系统的环境中运行我们的代码。我们尚未在其他配置上进行测试。目前，运行该方法大约需要 14 GB 的 GPU 内存。我们将继续优化内存效率。

要安装所需的库，只需运行以下命令：

conda env create -f environment.yaml
conda activate dragdiff

运行 DragDiffusion

首先，在命令行中运行以下命令启动 Gradio 用户界面：

python3 drag_ui.py

您可以通过我们上面的 GIF 动画 按步骤了解如何使用该界面。

基本上，操作流程如下：

场景 1：拖拽输入的真实图像

1) 训练 LoRA

• 将输入图像拖放到最左侧的框中。
• 在“prompt”字段中输入描述该图像的提示词。
• 点击“Train LoRA”按钮，根据输入图像训练一个 LoRA。

2) 执行“拖拽”编辑

• 在最左侧的框中绘制遮罩，以指定可编辑区域。
• 在中间的框中点击控制点和目标点。您也可以通过点击“Undo point”来重置所有点。
• 点击“Run”按钮运行我们的算法。编辑后的结果将显示在最右侧的框中。

场景 2：拖拽扩散生成的图像

1) 生成图像

• 填写生成参数（例如，正面/负面提示词、生成配置和 FreeU 参数）。
• 点击“Generate Image”。

2) 对生成的图像进行“拖拽”

• 在最左侧的框中绘制遮罩，以指定可编辑区域。
• 在中间的框中点击控制点和目标点。
• 点击“Run”按钮运行我们的算法。编辑后的结果将显示在最右侧的框中。

许可证

DragDiffusion 算法相关的代码采用 Apache 2.0 许可证。

BibTeX

如果您觉得我们的仓库有用，请考虑给个 Star 或引用我们的论文 :)

@article{shi2023dragdiffusion,
  title={DragDiffusion: Harnessing Diffusion Models for Interactive Point-based Image Editing},
  author={Shi, Yujun and Xue, Chuhui and Pan, Jiachun and Zhang, Wenqing and Tan, Vincent YF and Bai, Song},
  journal={arXiv preprint arXiv:2306.14435},
  year={2023}
}

联系方式

如对本项目有任何疑问，请联系 Yujun (shi.yujun@u.nus.edu)

致谢

本工作受到令人惊叹的 DragGAN 的启发。LoRA 训练代码基于 diffusers 的一个 示例 修改而来。图像样本收集自 unsplash、pexels 和 pixabay。最后，向所有优秀的开源扩散模型和库致以衷心的感谢。

相关链接

• Drag Your GAN：生成图像流形上的交互式点操控
• MasaCtrl：无需调优的一致性图像合成与编辑的互斥自注意力控制
• 从图像扩散中涌现的对应关系
• DragonDiffusion：在扩散模型上实现拖拽式操控
• FreeDrag：交互式点操控并不需要点跟踪技术

常见问题及解决方案

1. 对于因网络限制而无法从 Hugging Face 加载模型的用户，请按以下步骤操作：1) 参考此 链接 下载模型至 “local_pretrained_models” 目录；2) 运行 “drag_ui.py”，并在“算法参数 -> 基础模型配置 -> 扩散模型路径” 中选择您预训练模型所在的目录。

DragDiffusion 快速上手指南

DragDiffusion 是一个基于扩散模型的交互式图像编辑工具，允许用户通过拖拽点的方式对真实图像或生成图像进行精准操控。

环境准备

• 操作系统: Linux (推荐)
• 硬件要求: NVIDIA GPU，显存至少 14 GB
• 软件依赖:
  • Conda (用于环境管理)
  • Python (通过 conda 自动安装)
  • CUDA (需与 PyTorch 版本匹配)

注意: 本项目目前仅在 Linux + NVIDIA GPU 环境下经过测试。

安装步骤

1. 创建并激活虚拟环境 在项目根目录下运行以下命令，将自动安装所有必要的依赖库：

   conda env create -f environment.yaml
   conda activate dragdiff
   

2. 模型下载（可选，针对网络受限用户） 如果无法直接从 Hugging Face 加载模型，可手动下载预训练模型至 local_pretrained_models 目录，并在启动后于界面中指定路径（详见“基本使用”中的参数配置）。

基本使用

1. 启动界面

在终端中运行以下命令启动 Gradio 用户界面：

python3 drag_ui.py

启动成功后，终端会显示本地访问地址（通常为 http://127.0.0.1:7860），请在浏览器中打开该地址。

2. 操作流程

根据图片来源不同，分为两种使用模式：

模式一：编辑真实图片 (Dragging Input Real Images)

1. 训练 LoRA:
   • 将待编辑的真实图片拖入左侧第一个输入框。
   • 在 "prompt" 字段中输入描述该图片的提示词。
   • 点击 "Train LoRA" 按钮，系统将基于该图片训练一个轻量级 LoRA 模型（A100 上约需 20 秒）。
2. 执行拖拽编辑:
   • 在左侧图片框中绘制蒙版（Mask），圈定允许编辑的区域。
   • 在中间的预览框中，依次点击设置 控制点 (Handle Point) 和 目标点 (Target Point)。如需重置可点击 "Undo point"。
   • 点击 "Run" 按钮，编辑后的结果将显示在右侧输出框。

模式二：编辑扩散模型生成的图片 (Dragging Diffusion-Generated Images)

1. 生成图片:
   • 填写生成参数（如正向/负向提示词、Generation Config 及 FreeU 参数等）。
   • 点击 "Generate Image" 生成底图。
2. 执行拖拽编辑:
   • 在左侧生成的图片上绘制蒙版，指定编辑区域。
   • 在中间框设置 控制点 和 目标点。
   • 点击 "Run" 按钮获取编辑结果。

3. 关键参数说明（界面内折叠项）

• Diffusion Model Path: 默认使用 runwayml/stable-diffusion-v1-5。若使用本地下载的模型，请在此处修改路径。
• VAE Choice:
  • default: 使用原始 VAE。
  • stabilityai/sd-vae-ft-mse: 推荐用于包含人脸或眼睛的图像，可提升细节质量。
• n_pix_step: 运动监督的最大步数。如果控制点未移动到预期位置，可适当增加此数值。