Diffree

论文"Diffree: 基于扩散模型的文本引导无形状约束物体修复"的官方PyTorch实现

[🌐 项目主页] [🗞️ 数据集] [🎥 视频] [📜 论文] [🤗 Hugging Face 演示]

摘要

点击查看完整摘要

本文解决了仅通过文本指导在图像中添加物体这一重要问题。这一任务具有挑战性，因为新添加的物体必须与图像的视觉上下文（如光照、纹理和空间位置）无缝集成。虽然现有的文本引导图像修复方法可以添加物体，但它们要么无法保持背景一致性，要么需要繁琐的人工干预来指定边界框或用户涂鸦的掩码。为了解决这一挑战，我们提出了Diffree，一个仅需文本控制就能实现文本引导物体添加的文本到图像（T2I）模型。为此，我们通过先进的图像修复技术移除物体，精心制作了OABench合成数据集。OABench包含74K个真实世界的元组，每个元组包含原始图像、移除物体后的修复图像、物体掩码和物体描述。Diffree在OABench上使用Stable Diffusion模型进行训练，并添加了一个额外的掩码预测模块，独特地预测新物体的位置，仅通过文本指导实现物体添加。大量实验表明，Diffree在添加新物体方面表现出色，具有高成功率，同时保持背景一致性、空间适当性以及物体相关性和质量。

我们欢迎任何建议和讨论，请随时通过liruizhao@stu.xmu.edu.cn与我们联系。

新闻

[2024/07] 发布推理代码和检查点
[2024/07] 发布🤗 Hugging Face 演示
[2024/08] 发布ConfyUI演示。感谢smthemex（ComfyUI_Diffree）的帮助！
[2024/08] 在Hugging Face上发布训练数据集OABench
[2024/08] 发布训练代码
[2024/08] 更新🤗 演示，现在支持通过文本列表进行迭代生成

安装

克隆此仓库并进入Diffree文件夹

git clone https://github.com/OpenGVLab/Diffree.git

cd Diffree

安装包

conda create -n diffree python=3.8.5

conda activate diffree

pip install -r requirements.txt

推理

从Huggingface下载Diffree模型。

pip install huggingface_hub

huggingface-cli download LiruiZhao/Diffree --local-dir ./checkpoints

你可以使用以下脚本进行推理：

python app.py

具体来说，--resolution定义了调整大小后的输入图像和输出图像的最大尺寸。对于我们的Hugging Face演示，我们将--resolution设置为512，以提高用户体验和生成更高分辨率的结果。而在Diffree的训练过程中，--resolution设置为256。因此，降低--resolution可能会改善结果（例如，可以尝试将320作为一个潜在值）。

数据下载

你可以在这里下载OABench，用于训练Diffree。

从Huggingface下载OABench数据集。

huggingface-cli download --repo-type dataset LiruiZhao/OABench --local-dir ./dataset --local-dir-use-symlinks False

cd dataset

ls *.tar.gz | xargs -n1 tar xvf

数据结构应该如下所示：

|-- dataset
    |-- original_images
        |-- 58134.jpg
        |-- 235791.jpg
        |-- ...
    |-- inpainted_images
        |-- 58134
          |-- 634757.jpg
          |-- 634761.jpg
          |-- ...
        |-- 235791
        |-- ...
    |-- mask_images
        |-- 58134
          |-- 634757.png
          |-- 634761.png
          |-- ...
        |-- 235791
        |-- ...
    |-- annotations.json

在inpainted_images和mask_images目录中，顶层文件夹对应原始图像，每个文件夹的内容是这些图像的修复图像和掩码。

训练

Diffree通过从初始StableDiffusion检查点微调来训练。

下载Stable Diffusion检查点并将其移动到checkpoints目录。对于我们的训练模型，我们使用v1.5检查点作为起点。你也可以使用以下命令：

curl -L https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt -o checkpoints/v1-5-pruned-emaonly.ckpt

接下来，你可以开始训练。

python main.py --name diffree --base config/train.yaml --train --gpus 0,1,2,3

所有配置都存储在YAML文件中。如果你需要使用自定义配置设置，可以修改--base以指向你的自定义配置文件。

引用

如果你觉得这项工作有用，请考虑引用：

@article{zhao2024diffree,
  title={Diffree: Text-Guided Shape Free Object Inpainting with Diffusion Model},
  author={Zhao, Lirui and Yang, Tianshuo and Shao, Wenqi and Zhang, Yuxin and Qiao, Yu and Luo, Ping and Zhang, Kaipeng and Ji, Rongrong},
  journal={arXiv preprint arXiv:2407.16982},
  year={2024}
}