BrushNet 图像修复从零上手：5 分钟跑通文生图修复

BrushNet 图像修复从零上手：5 分钟跑通文生图修复

【免费下载链接】BrushNet[ECCV 2024] The official implementation of paper "BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion"项目地址: https://gitcode.com/GitHub_Trending/br/BrushNet

BrushNet 是 ECCV 2024 提出的即插即用图像修复（image inpainting）模型，它以一个附加的双分支网络接到现成的 Stable Diffusion 权重上，让你不动基础模型就能做文本引导的图像修复：给一张图加一个掩码、再给一句提示词，掩码区域就会被补出与周围环境风格一致的内容。读完本文，你能在自己的机器上把 BrushNet 的推理例子跑起来，换用自己的图片和掩码做修复，并知道显存不够、VAE 出 NaN 这类常见问题的处理办法。

🧩 先搞懂 BrushNet 在修什么

传统 inpainting 往往只补"像素"，不管补出来的东西符不符合文字描述、和画面风格搭不搭。BrushNet 的思路是把"被遮住的图像特征"和"带噪 latent"分开喂给一个小网络，再通过零卷积把它的输出逐层加到冻结的 UNet 上——相当于在原模型旁边插了一个只负责修复的旁路，学得快、不伤原模型的生成能力。

整个仓库其实是一个带 BrushNet 扩展的 diffusers 分支，核心代码在src/diffusers/models/brushnet.py（网络结构）和src/diffusers/pipelines/brushnet/（SD 1.5 与 SDXL 两条 pipeline），示例和训练脚本在examples/brushnet/。

🛠️ 把环境拉起来

仓库官方按 Python 3.9 + PyTorch 1.12.1 测试，环境搭建分四步：建环境、装 PyTorch、装本仓库的 diffusers、装 brushnet 示例依赖。

第一步，克隆仓库并建一个干净的 conda 环境：

git clone https://gitcode.com/GitHub_Trending/br/BrushNet cd BrushNet conda create -n brushnet python=3.9 -y conda activate brushnet python -m pip install --upgrade pip

第二步，装 PyTorch（按自己的 CUDA 版本选 index-url，cu116 / cu118 都行）：

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 以可编辑模式装进去——注意必须装仓库内这一版，因为BrushNetModel和两条 BrushNet pipeline 是加在这里面的，官方 pip 上的 diffusers 没有这些类：

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

examples/brushnet/requirements.txt里锁定了Pillow==9.5.0和accelerate==0.20.3等版本，如果你环境里已有别的版本，按报错提示降级或升级对应包即可。

📦 模型和测试图放哪

推理需要两样东西：一个基础 SD 模型（比如 SD 1.5 的写实风 checkpoint）和一个 BrushNet 权重。仓库 README 给出的目录约定如下，下载后照这个结构放即可：

data/ └── ckpt/ ├── realisticVisionV60B1_v51VAE/ # 基础 SD1.5 模型（含 model_index.json、vae/ 等） ├── segmentation_mask_brushnet_ckpt/ # 分割掩码版 BrushNet └── random_mask_brushnet_ckpt/ # 随机掩码版 BrushNet

两个 BrushNet 权重的区别：segmentation_mask_brushnet_ckpt用带分割先验的数据训练，掩码形状贴合物体轮廓，适合"抠掉某个物体再补"；random_mask_brushnet_ckpt用随机形状掩码训练，形状适应性更广。没有分割掩码就用后者。

测试图不用自己准备，examples/brushnet/src/下已有一套test_image.jpg和test_mask.jpg（白色区域是要修复的区域），仓库里还有多组示例效果图可作参考：

🏃🏼 第一个例子跑起来

最小可运行的推理逻辑就是examples/brushnet/test_brushnet.py，拆成三段看。

第一段，加载模型。BrushNet 权重单独加载，再作为brushnet参数注入 pipeline：

from diffusers import StableDiffusionBrushNetPipeline, BrushNetModel, UniPCMultistepScheduler import torch brushnet = BrushNetModel.from_pretrained("data/ckpt/segmentation_mask_brushnet_ckpt", torch_dtype=torch.float16) pipe = StableDiffusionBrushNetPipeline.from_pretrained( "data/ckpt/realisticVisionV60B1_v51VAE", brushnet=brushnet, torch_dtype=torch.float16, low_cpu_mem_usage=False)

第二段，准备输入。掩码要做成二值图（白=修复区），原图则把掩码区域置黑，这样原图信息不会"漏"进修复区：

import cv2 import numpy as np from PIL import Image init_image = cv2.imread("examples/brushnet/src/test_image.jpg")[:, :, ::-1] mask_image = 1. * (cv2.imread("examples/brushnet/src/test_mask.jpg").sum(-1) > 255)[:, :, np.newaxis] init_image = Image.fromarray((init_image * (1 - mask_image)).astype(np.uint8)).convert("RGB") mask_image = Image.fromarray((mask_image.astype(np.uint8).repeat(3, -1) * 255)).convert("RGB")

第三段，调 pipeline 出图。换上 UniPC 调度器可以明显减少出图时间：

pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) pipe.enable_model_cpu_offload() # 显存紧张时保留这行 generator = torch.Generator("cuda").manual_seed(1234) image = pipe( "A cake on the table.", init_image, mask_image, num_inference_steps=50, generator=generator, brushnet_conditioning_scale=1.0, ).images[0] image.save("output.png")

在仓库根目录执行python examples/brushnet/test_brushnet.py即可得到output.png。

🎛️ 关键参数速查

推理时真正值得动的参数就几个，都通过pipe(...)传入：

🔍 三个高频用法

跑 SDXL 版本。流程几乎一样，只是换 pipeline 类和调度器，参考examples/brushnet/test_brushnet_sdxl.py：

from diffusers import StableDiffusionXLBrushNetPipeline, DPMSolverMultistepScheduler pipe = StableDiffusionXLBrushNetPipeline.from_pretrained( base_model_path, brushnet=brushnet, torch_dtype=torch.float16, use_safetensors=True) pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

注意 SDXL 在 fp16 下做 VAE 编码容易出 NaN，示例里用madebyollin/sdxl-vae-fp16-fix这个修复版 VAE 替换掉了默认的。

开一个可点选掩码的网页。examples/brushnet/app_brushnet.py是一个 Gradio 应用，内置 Segment Anything：在图上点几下就能生成掩码，还能切换 blended（模糊掩码融合）模式，适合不想自己写掩码图的用户：

python examples/brushnet/app_brushnet.py

批量跑评测指标。有 BrushBench 数据集的话，examples/brushnet/evaluate_brushnet.py会输出 CLIPScore、LPIPS、PSNR、SSIM、HPSv2 等指标，用来横向对比不同掩码权重或自己训练的版本。

⚠️ 踩坑实录

现象：加载 pipeline 时 CUDA OOM。原因：SD 1.5 本体 + BrushNet 旁路网络同时常驻显存，8GB 卡很紧张。解法：保留pipe.enable_model_cpu_offload()，它会把不活跃的模块挪回 CPU；再不行就加pipe.enable_attention_slicing()。low_cpu_mem_usage这里要显式设False，否则权重装载路径和 BrushNet 的注入方式会打架。

现象：修复结果没变化，像是把整张原图透出来了。原因：掩码极性弄反——BrushNet 约定白色是"要修复的区域"，而很多修图软件（PS 选区等）导出的是黑色=修复区。解法：读入掩码后做一次255 - mask反转，或者直接检查mask.sum(-1) > 255这一行的判断是否符合你的掩码极性。

现象：SDXL + fp16 推理出全黑图或 NaN。原因：SDXL 的 VAE 在 fp16 下数值不稳定。解法：按test_brushnet_sdxl.py的做法换sdxl-vae-fp16-fix权重，或把 VAE 单独放 fp32。

现象：pip install -e .后from diffusers import BrushNetModel报 ImportError。原因：环境里同时存在官方 diffusers，import 被旧包抢走。解法：确认pip show diffusers指向的是本地仓库路径，必要时先pip uninstall diffusers再装。

🚀 进阶入口

• 用自己的数据微调：训练脚本是examples/brushnet/train_brushnet.py（SD1.5）和train_brushnet_sdxl.py（SDXL），数据按 BrushData 的 tar 格式组织后传--train_data_dir；随机掩码训练加--random_mask即可。官方训练命令示例写在 README.md 的 Running Scripts 一节，直接照着改路径。
• 指令引导编辑玩法：仓库里的 InstructionGuidedEditing/editing_pipeline.ipynb 展示了 BrushNet 做 CVPR2024 Workshop 夺冠方案时的组合用法。
• 想改网络结构或逐层注入逻辑，读src/diffusers/models/brushnet.py里的BrushNetModel即可，输出是 down/mid/up 三组特征，通过零卷积进 UNet。

🎬 接下来可以试试

1. 把brushnet_conditioning_scale从 1.0 降到 0.8 再出一张图，对比修复区与边缘的融合差异。
2. 用同一张test_image.jpg分别跑 segmentation 和 random 两个掩码权重，看两种训练的泛化差别。
3. 打开app_brushnet.py的网页，点选一个物体当掩码，体验不用手画 mask 的完整流程。

【免费下载链接】BrushNet[ECCV 2024] The official implementation of paper "BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion"项目地址: https://gitcode.com/GitHub_Trending/br/BrushNet