BrushNet

本仓库包含ECCV2024论文"BrushNet: 具有分解双分支扩散的即插即用图像修复模型"的实现

关键词：图像修复、扩散模型、图像生成

Xuan Ju¹², Xian Liu¹², Xintao Wang^1*, Yuxuan Bian², Ying Shan¹, Qiang Xu^2*<br> ¹腾讯PCG ARC实验室 ²香港中文大学 ^*通讯作者

<p align="center"> <a href="https://tencentarc.github.io/BrushNet/">🌐项目主页</a> | <a href="https://arxiv.org/abs/2403.06976">📜Arxiv论文</a> | <a href="https://forms.gle/9TgMZ8tm49UYsZ9s5">🗄️数据</a> | <a href="https://drive.google.com/file/d/1IkEBWcd2Fui2WHcckap4QFPcCI0gkHBh/view">📹视频</a> | <a href="https://huggingface.co/spaces/TencentARC/BrushNet">🤗Hugging Face演示</a> | </p>

📖 目录

• BrushNet
  • 待办事项
  • 🛠️ 方法概述
  • 🚀 开始使用
    • 环境要求 🌍
    • 数据下载 ⬇️
  • 🏃🏼 运行脚本
    • 训练 🤯
    • 推理 📜
    • 评估 📏
  • 🤝🏼 引用我们
  • 💖 致谢

待办事项

• 发布训练和推理代码
• 发布检查点 (sdv1.5)
• 发布检查点 (sdxl)。遗憾的是，我只有V100用于训练这个检查点，只能以1的批量大小进行训练，速度很慢。当前的检查点只训练了很少的步数，因此表现不佳。但幸运的是，yuanhang自愿帮助训练一个更好的版本。敬请期待！感谢yuanhang的努力！
• 发布评估代码
• 发布Gradio演示
• 发布ComfyUI演示。感谢nullquant (ConfyUI-BrushNet)和kijai (ComfyUI-BrushNet-Wrapper)的帮助！
• 发布训练数据。感谢random123123的帮助！
• 我们使用BrushNet参加了CVPR2024 GenAI媒体生成挑战工作坊并获得了最高奖！解决方案在InstructionGuidedEditing中提供
• 发布新版本的检查点 (sdxl)。

🛠️ 方法概述

BrushNet是一个基于扩散的文本引导图像修复模型，可以即插即用到任何预训练的扩散模型中。我们的架构设计包含两个关键见解：(1)将掩码图像特征和噪声潜在表示分开可以减少模型的学习负担，(2)利用整个预训练模型的密集逐像素控制可以增强其在图像修复任务中的适用性。更多分析可以在主论文中找到。

🚀 开始使用

环境要求 🌍

BrushNet已在Python 3.9和PyTorch 1.12.1上实现和测试。

克隆仓库：

git clone https://github.com/TencentARC/BrushNet.git

我们建议您先使用conda创建虚拟环境，并按照官方说明安装pytorch。例如：

conda create -n diffusers python=3.9 -y
conda activate diffusers
python -m pip install --upgrade pip
pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116

然后，您可以通过以下方式安装diffusers（在本仓库中实现）：

pip install -e .

之后，您可以通过以下方式安装所需的包：

cd examples/brushnet/
pip install -r requirements.txt

数据下载 ⬇️

数据集

您可以在这里下载BrushData和BrushBench（以及我们重新处理的EditBench），这些数据用于训练和测试BrushNet。通过下载数据，您同意许可条款和条件。数据结构应如下所示：

|-- data
    |-- BrushData
        |-- 00200.tar
        |-- 00201.tar
        |-- ...
    |-- BrushDench
        |-- images
        |-- mapping_file.json
    |-- EditBench
        |-- images
        |-- mapping_file.json

注意：由于空间限制，我们在Google Drive上只提供了部分BrushData。random123123已帮助在Hugging Face上上传了完整的数据集，可以在这里找到。感谢他的帮助！

检查点

BrushNet的检查点可以从这里下载。ckpt文件夹包含：

• 用于Stable Diffusion v1.5的BrushNet预训练检查点（segmentation_mask_brushnet_ckpt和random_mask_brushnet_ckpt）
• 预训练的Stable Diffusion v1.5检查点（例如，来自Civitai的realisticVisionV60B1_v51VAE）。您可以使用scripts/convert_original_stable_diffusion_to_diffusers.py来处理从Civitai下载的其他模型。
• 用于Stable Diffusion XL的BrushNet预训练检查点（segmentation_mask_brushnet_ckpt_sdxl_v1和random_mask_brushnet_ckpt_sdxl_v0）。yuanhang将很快发布一个更好的版本。敬请期待！
• 预训练的Stable Diffusion XL检查点（例如，来自Civitai的juggernautXL_juggernautX）。您可以使用StableDiffusionXLPipeline.from_single_file("safetensors的路径").save_pretrained("保存路径",safe_serialization=False)来处理从Civitai下载的其他模型。

数据结构应如下所示：

|-- data
    |-- BrushData
    |-- BrushDench
    |-- EditBench
    |-- ckpt
        |-- realisticVisionV60B1_v51VAE
            |-- model_index.json
            |-- vae
            |-- ...
        |-- segmentation_mask_brushnet_ckpt
        |-- segmentation_mask_brushnet_ckpt_sdxl_v0
        |-- random_mask_brushnet_ckpt
        |-- random_mask_brushnet_ckpt_sdxl_v0
        |-- ...

segmentation_mask_brushnet_ckpt和segmentation_mask_brushnet_ckpt_sdxl_v0中的检查点提供了在BrushData上训练的检查点，具有分割先验（掩码与对象形状相同）。random_mask_brushnet_ckpt和random_mask_brushnet_ckpt_sdxl提供了一个更通用的检查点，适用于随机掩码形状。

🏃🏼 运行脚本

训练 🤯

您可以使用以下脚本训练带有分割掩码的模型：

# sd v1.5
accelerate launch examples/brushnet/train_brushnet.py \
--pretrained_model_name_or_path runwayml/stable-diffusion-v1-5 \
--output_dir runs/logs/brushnet_segmentationmask \
--train_data_dir data/BrushData \
--resolution 512 \
--learning_rate 1e-5 \
--train_batch_size 2 \
--tracker_project_name brushnet \
--report_to tensorboard \
--resume_from_checkpoint latest \
--validation_steps 300
--checkpointing_steps 10000 

# sdxl
accelerate launch examples/brushnet/train_brushnet_sdxl.py \
--pretrained_model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \
--output_dir runs/logs/brushnetsdxl_segmentationmask \
--train_data_dir data/BrushData \
--resolution 1024 \
--learning_rate 1e-5 \
--train_batch_size 1 \
--gradient_accumulation_steps 4 \
--tracker_project_name brushnet \
--report_to tensorboard \
--resume_from_checkpoint latest \
--validation_steps 300 \
--checkpointing_steps 10000 

要使用自定义数据集，您可以将自己的数据处理成BrushData的格式，并修改--train_data_dir。

您可以使用以下脚本进行随机掩码训练（通过添加--random_mask）：

# sd v1.5
accelerate launch examples/brushnet/train_brushnet.py \
--pretrained_model_name_or_path runwayml/stable-diffusion-v1-5 \
--output_dir runs/logs/brushnet_randommask \
--train_data_dir data/BrushData \
--resolution 512 \
--learning_rate 1e-5 \
--train_batch_size 2 \
--tracker_project_name brushnet \
--report_to tensorboard \
--resume_from_checkpoint latest \
--validation_steps 300 \
--random_mask

# sdxl
accelerate launch examples/brushnet/train_brushnet_sdxl.py \
--pretrained_model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \
--output_dir runs/logs/brushnetsdxl_randommask \
--train_data_dir data/BrushData \
--resolution 1024 \
--learning_rate 1e-5 \
--train_batch_size 1 \
--gradient_accumulation_steps 4 \
--tracker_project_name brushnet \
--report_to tensorboard \
--resume_from_checkpoint latest \
--validation_steps 300 \
--checkpointing_steps 10000 \
--random_mask

推理 📜

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

# sd v1.5
python examples/brushnet/test_brushnet.py
# sdxl
python examples/brushnet/test_brushnet_sdxl.py

由于BrushNet是在Laion上训练的，它只能保证在一般场景下的性能。如果您有高质量的工业应用需求，我们建议您在自己的数据（例如，产品展示、虚拟试穿）上进行训练。如果您愿意贡献您训练的模型，我们也将不胜感激！

您还可以通过gradio演示进行推理：

# sd v1.5
python examples/brushnet/app_brushnet.py

评估 📏

您可以使用以下脚本进行评估：

python examples/brushnet/evaluate_brushnet.py \
--brushnet_ckpt_path data/ckpt/segmentation_mask_brushnet_ckpt \
--image_save_path runs/evaluation_result/BrushBench/brushnet_segmask/inside \
--mapping_file data/BrushBench/mapping_file.json \
--base_dir data/BrushBench \
--mask_key inpainting_mask

--mask_key表示使用哪种类型的掩码，inpainting_mask用于内部修复，outpainting_mask用于外部修复。评估结果（图像和指标）将保存在--image_save_path中。

请注意，您需要在src/diffusers/pipelines/brushnet/pipeline_brushnet.py#1261中忽略nsfw检测器以获得正确的评估结果。此外，我们发现不同的机器可能会生成不同的图像，因此在这里提供了我们机器上的结果。

🤝🏼 引用我们

@misc{ju2024brushnet,
  title={BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion}, 
  author={Xuan Ju and Xian Liu and Xintao Wang and Yuxuan Bian and Ying Shan and Qiang Xu},
  year={2024},
  eprint={2403.06976},
  archivePrefix={arXiv},
  primaryClass={cs.CV}
}

💖 致谢

<span id="acknowledgement"></span>

我们的代码是基于diffusers修改的，感谢所有贡献者！