双分支扩散结构实现即插即用图像修复模型
BrushNet是一种新型图像修复模型,采用双分支扩散结构,可集成到预训练扩散模型中。通过分离遮罩图像特征和噪声潜在表示,并对预训练模型实施像素级控制,提高了图像修复效果。该模型适用于通用场景和特定应用,为图像编辑和生成提供了新的解决方案。
本仓库包含ECCV2024论文"BrushNet: 具有分解双分支扩散的即插即用图像修复模型"的实现
关键词:图像修复、扩散模型、图像生成
<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>Xuan Ju<sup>12</sup>, Xian Liu<sup>12</sup>, Xintao Wang<sup>1*</sup>, Yuxuan Bian<sup>2</sup>, Ying Shan<sup>1</sup>, Qiang Xu<sup>2*</sup><br> <sup>1</sup>腾讯PCG ARC实验室 <sup>2</sup>香港中文大学 <sup>*</sup>通讯作者
📖 目录
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文件夹包含:
segmentation_mask_brushnet_ckpt
和random_mask_brushnet_ckpt
)scripts/convert_original_stable_diffusion_to_diffusers.py
来处理从Civitai下载的其他模型。segmentation_mask_brushnet_ckpt_sdxl_v1
和random_mask_brushnet_ckpt_sdxl_v0
)。yuanhang将很快发布一个更好的版本。敬请期待!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修改的,感谢所有贡献者!
AI小说写作助手,一站式润色、改写、扩写
蛙蛙写作—国内先进的AI写作平台,涵盖小说、学术、社交媒体等多场景。提供续写、改写、润色等功能,助力创作者高效优化写作流程。界面简洁,功能全面,适合各类写作者提升内容品质和工作效率。
字节跳动发布的AI编程神器IDE
Trae是一种自适应的集成开发环境(IDE),通过自动化和多元协作改变开发流程。利用Trae,团队能够更快速、精确地编写和部署代码,从而提高编程效率和项目交付速度。Trae具备上下文感知和代码自动完成功能,是提升开发效率的理想工具。