BrushNet 图像修复实践指南：用最小示例跑通文本引导修复-平芜编程栈

BrushNet 图像修复实践指南：用最小示例跑通文本引导修复

【免费下载链接】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 管线（SD v1.5 / SDXL）上。本文带你完成环境搭建、检查点准备，并运行一个最小修复示例，跑通后你可以直接把它扩展到自己的数据场景。适合刚接触扩散模型修复方向的工程师。

它解决了什么问题

给定一张源图像、一张掩码（白色区域表示待修复区域）和一句文本描述，BrushNet 会在保持非掩码区域与原图一致的前提下，重建掩码区域的内容。仓库内置的示例展示了两张车的修复前后对比：

与微调整个 U-Net 不同，BrushNet 不改动基座模型的权重：它冻结预训练的 UNet，另加一个小型的 BrushNet 分支，将带噪潜变量、被掩码图像的 VAE 特征与下采样掩码拼接后输入，分支特征经过零卷积（zero convolution）逐层加回到冻结 UNet 中。相当于把"修复先验的学习"和"基座生成能力"解耦，因此任何现成的 Stable Diffusion 模型都能直接复用。推理结束后，生成结果还会与原始图像按模糊掩码做混合（blending），保证边缘平滑。

准备环境

项目	要求
Python	3.9（支持 3.8+）
PyTorch	1.12.1+cu116（官方测试版本）
GPU	CUDA 可用，显存建议 10GB 以上
操作系统	Linux

仓库本身是 diffusers 的修改版（新增了BrushNetModel与StableDiffusionBrushNetPipeline），因此必须用可编辑模式安装本仓库，而不是 pip 安装官方 PyPI 版本的 diffusers：

# 克隆仓库 git clone https://gitcode.com/GitHub_Trending/br/BrushNet cd BrushNet # 创建环境并安装官方推荐的 PyTorch 版本 conda create -n brushnet python=3.9 -y conda activate brushnet 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 修改版（必须用 -e） pip install -e . # 安装 BrushNet 示例依赖 pip install -r examples/brushnet/requirements.txt

校验是否安装成功：

# 校验 GPU 是否可用 python -c "import torch; print(torch.__version__, torch.cuda.is_available())" # 校验修改版 diffusers 中的 BrushNet 类可导入 python -c "from diffusers import StableDiffusionBrushNetPipeline, BrushNetModel; print('ok')"

推理前还需要下载检查点（仓库 README 中给出了官方下载入口），放在如下目录结构中：

data/ └── ckpt/ ├── realisticVisionV60B1_v51VAE/ # 基座 Stable Diffusion v1.5（含 model_index.json、vae/、unet/ 等） └── segmentation_mask_brushnet_ckpt/ # BrushNet 分支模型（分割掩码训练）

如果你有来自 Civitai 等其他渠道的 SD v1.5 模型，可以用仓库内的 scripts/convert_original_stable_diffusion_to_diffusers.py 转换后再使用。

跑通最小示例

仓库自带的 examples/brushnet/test_brushnet.py 可直接运行，示例输入为仓库内的examples/brushnet/src/test_image.jpg与对应的黑白掩码test_mask.jpg。下面是等价的精简代码：

import torch, cv2, numpy as np from PIL import Image from diffusers import StableDiffusionBrushNetPipeline, BrushNetModel, UniPCMultistepScheduler 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) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) pipe.enable_model_cpu_offload() init_image = cv2.imread("examples/brushnet/src/test_image.jpg")[:, :, ::-1] mask = 1.0 * (cv2.imread("examples/brushnet/src/test_mask.jpg").sum(-1) > 255)[:, :, np.newaxis] init_image = Image.fromarray((init_image * (1 - mask)).astype(np.uint8)).convert("RGB") mask_image = Image.fromarray((mask.astype(np.uint8) * 255).repeat(3, -1)).convert("RGB") image = pipe("A cake on the table.", init_image, mask_image, num_inference_steps=50, generator=torch.Generator("cuda").manual_seed(1234), brushnet_conditioning_scale=1.0).images[0] image.save("output.png")

预期输出：当前目录下生成output.png。原图中蛋糕被掩码遮盖的部分会按提示词补全，掩码之外的内容与原图保持一致；512 分辨率下一次推理约需十几秒到一分钟（取决于显卡）。如果结果边缘有轻微拼缝感，可参考官方脚本中的blended选项，用模糊掩码把生成区域与原图做加权合成。

拆解关键步骤

把 BrushNet 分支挂载进管线

BrushNetModel单独加载，作为brushnet参数传入管线；基座 UNet 在整个过程中保持冻结，这正是"即插即用"的含义。

brushnet = BrushNetModel.from_pretrained(brushnet_path, torch_dtype=torch.float16) pipe = StableDiffusionBrushNetPipeline.from_pretrained( base_model_path, brushnet=brushnet, torch_dtype=torch.float16)

预处理图像与掩码

掩码的白色（像素值大于 255）对应待修复区域。输入管线前，需要先把源图像的掩码区域置黑，并把单通道掩码扩成三通道：

mask = 1.0 * (cv2.imread(mask_path).sum(-1) > 255) # 白色区域置 1 init_image = init_image * (1 - mask) # 待修复区域置黑 mask_image = Image.fromarray((mask.astype(np.uint8) * 255).repeat(3, -1))

配置推理参数

两个最影响效果的参数：num_inference_steps控制去噪步数，30~50 之间按质量需求取舍；brushnet_conditioning_scale控制 BrushNet 条件的注入强度，1.0 为默认值，调低会让生成结果更保守、更贴近原图。

image = pipe(prompt, init_image, mask_image, num_inference_steps=50, # 推理步数 brushnet_conditioning_scale=1.0).images[0] # 条件强度

典型问题与排查

现象	可能原因	处理方法
导入`BrushNetModel`报`ModuleNotFoundError`	装的是 PyPI 上的官方 diffusers，不含 BrushNet	在仓库根目录执行`pip install -e .`
推理时 CUDA 显存不足	fp16 下基座 UNet 加 BrushNet 常驻约需 10GB	启用`pipe.enable_model_cpu_offload()`，或追加`pipe.enable_attention_slicing()`，并降低输入分辨率
单次推理明显偏慢	仍在默认 LMS 调度器上跑完整步数	换用`UniPCMultistepScheduler.from_config(pipe.scheduler.config)`，可在不牺牲太多质量的前提下提速
修复内容与提示词无关、或整张图被重绘	掩码极性与检查点类型不匹配：掩码白色应为待修复区域；分割掩码与随机掩码两种检查点适用范围不同	核对掩码极性，并按掩码形状选择`segmentation_mask_brushnet_ckpt`或`random_mask_brushnet_ckpt`

进阶方向

在自己的数据上训练：使用 examples/brushnet/train_brushnet.py，--train_data_dir指向 BrushData 格式（tar 文件加 mapping）的数据目录，追加--random_mask即可按随机形状掩码训练，学习率建议从1e-5起步。
SDXL 支持：推理与训练分别对应 examples/brushnet/test_brushnet_sdxl.py 和 examples/brushnet/train_brushnet_sdxl.py，--resolution需设为 1024。
交互演示与批量评估：examples/brushnet/app_brushnet.py 提供 Gradio 界面，可基于 Segment Anything 点选生成掩码；examples/brushnet/evaluate_brushnet.py 支持--mask_key inpainting_mask / outpainting_mask做基准评估。