BrushNet：双分支扩散图像修复模型

BrushNet是一款基于扩散模型的图像修复工具，它采用了独特的双分支结构来处理被遮挡的图像特征和噪声潜变量，极大地减轻了模型的学习负担。此外，BrushNet能够在任何预训练的扩散模型上即插即用，提供了密集的像素级控制，使其非常适合图像修复任务。

1. 功能与方法概览

BrushNet主要通过以下方法实现高效的图像修复：

双分支结构：通过分离遮挡图像特征和噪声潜变量来降低学习难度。
密集像素控制：利用预训练模型的密集像素级控制优化图像修复效果。

home_open_BrushNet_1

2. 开始使用

环境需求：

Python 3.9
Pytorch 1.12.1

安装步骤：

克隆代码仓库：

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

创建虚拟环境并安装依赖：

    
    conda create -n diffusers python=3.9 -y
conda activate diffusers
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
pip install -e .
cd examples/brushnet/
pip install -r requirements.txt

3. 数据和训练

数据下载：

下载BrushData和BrushBench数据集，用于训练和测试BrushNet。

模型训练：

使用以下脚本开始训练模型：

    
    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

4. 推理和评估

推理：

使用以下命令进行图像修复推理：

    
    python examples/brushnet/test_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

BrushNet的安装和设置过程非常简单明了，即使是对深度学习不太熟悉的用户也能轻松上手。在使用过程中，BrushNet显示出了其强大的图像修复能力，特别是在处理复杂图像遮挡问题时，它的双分支结构设计显著提高了修复质量。