BrushNet 快速上手：5 分钟跑通 ECCV 2024 即插即用图像修复模型-平芜编程栈

BrushNet 快速上手：5 分钟跑通 ECCV 2024 即插即用图像修复模型

【免费下载链接】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 是论文《BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion》（ECCV 2024）的官方实现：一个即插即用的文本引导图像修复模型。这篇教程带你从环境搭建到跑通第一次图像修复只要 3 步、约 5 分钟，全程只给跑通所必需的命令和参数，并附高频报错避坑手册。

一、价值主张：它替你解决什么

还在为图像修复头疼吗？——擦掉照片里的路人，边缘留下一圈鬼影；想换掉背景一角，补出来的纹理和周围对不上。传统做法要么重训整个扩散模型，要么忍受边缘不自然、内容不连贯的结果。

BrushNet 的思路是把修复任务拆成两条支路：一条负责"读懂"原图里要保留的内容，另一条负责往掩码区域填新内容，两支路并联进现有的 Stable Diffusion UNet。因为只是"插上"一个独立模块，你已有的基础模型（SD 1.5 / SDXL）一行都不用动，这就是它"即插即用"的含义。

读完本文你能带走：

✅ 一套 3 步跑通流程：装环境 → 放模型 → 出第一张修复图
✅ 5 个关键参数（条件强度、推理步数、种子、混合融合等）的调节依据
✅ 4 个高频报错的"现象 → 原因 → 解决"排查路径
✅ 从 Gradio 演示界面、Docker 部署到自定义训练的落地路线

弄清了"它是什么"，下面先把前置条件核一遍，避免装到一半才发现硬件不达标。

二、上手前准备：一张表核对前置条件

开始之前，对照下表确认你的环境。最容易被忽略的是磁盘空间——基础模型加 BrushNet 权重加起来约 10GB。

项目	最低要求	推荐配置
Python	3.8+	3.9 ~ 3.10
PyTorch	1.12+（含 CUDA 支持）	2.0+（CUDA 11.8）
GPU 显存	8GB（需开启 CPU 卸载）	16GB+
磁盘空间	10GB（SD 1.5 模型组）	30GB+（含 SDXL 与训练数据）
操作系统	Linux / Windows / macOS	Ubuntu 20.04+

条件齐了，就可以进入真正动手的 3 步。

三、跑通主流程：3 步出第一张修复图

步骤 1：环境准备（2 条命令装好依赖）

BrushNet 是基于diffusers修改的版本，必须先装仓库里的diffusers，再装示例依赖，顺序反了容易出依赖冲突。

# 克隆仓库 git clone https://gitcode.com/GitHub_Trending/br/BrushNet.git cd BrushNet # 安装本仓库的 diffusers（核心） pip install -e . # 安装 BrushNet 示例依赖（含 accelerate、gradio、SAM 等） cd examples/brushnet && pip install -r requirements.txt

步骤 2：放模型（目录结构一次放对）

代码里默认从data/ckpt/读模型，把模型按下面的结构放好，后面所有脚本都不用改路径：

data/ └── ckpt/ ├── realisticVisionV60B1_v51VAE/ # 基础模型（SD 1.5 格式，需含 model_index.json） ├── segmentation_mask_brushnet_ckpt/ # BrushNet 权重（分割掩码版，SD 1.5） ├── segmentation_mask_brushnet_ckpt_sdxl_v0/ # BrushNet 权重（SDXL 版） └── sam_vit_h_4b8939.pth # Gradio 应用点选掩码用的 SAM 权重

说明：基础模型也可以直接用 Hugging Face 上的runwayml/stable-diffusion-v1-5等官方模型，把脚本里的base_model_path改成对应 id 即可。BrushNet 权重则必须指向上面第 2 类目录。

模型就位后，直接跑仓库自带的测试脚本，它用的就是仓库里现成的示例图和掩码。

步骤 3：第一次运行（1 条命令出图）

# 回到仓库根目录 python examples/brushnet/test_brushnet.py

跑完当前目录会生成output.png，这就是第一次修复结果。脚本做的事可以用这张流程图概括：

脚本里 5 个关键参数，第一次运行保持默认即可，调优时再回来改：

参数	默认值	作用
`num_inference_steps`	50	采样步数，越大质量越好、越慢
`brushnet_conditioning_scale`	1.0	BrushNet 条件强度，越小越依赖原图
`generator`种子	1234	固定种子保证结果可复现
`blended`	False	是否用模糊掩码融合结果，去接缝
`torch_dtype`	float16	半精度推理，显存减半

看到第一张图了，接下来聊聊怎么判断效果、怎么再快一点。

四、看效果与调优

以仓库示例 4（野外开车场景）为例，输入、掩码、输出三件套是这样的：

看效果的要点：掩码白色区域按 prompt 重绘，非掩码区域应基本保持原样；如果修复区和保留区交界处有明显"贴片感"，把blended打开（脚本里已有高斯模糊融合的现成实现），边缘会自然很多。

跑通之后，常见优化手段按这张表选：

方法	收益	适用场景
`UniPCMultistepScheduler`（测试脚本已启用）	⚡ 采样加速 2-3 倍	所有推理场景
`pipe.enable_model_cpu_offload()`（测试脚本已启用）	💾 显存占用大幅降低	8GB 显存卡
`pipe.enable_xformers_memory_efficient_attention()`	显存再降约 20%	已装 xformers 且 Torch < 2.0
`pipe.enable_attention_slicing()`	注意力显存分段加载	高分辨率输入
`blended=True`模糊融合	去除接缝伪影	掩码边界与背景反差大时

速度之外，真正影响观感的其实还是参数选择，下面几个坑基本覆盖了 90% 的翻车现场。

五、避坑手册：4 个高频问题排雷

坑 1：CUDA 内存不足

现象：torch.cuda.OutOfMemoryError: CUDA out of memory
原因：UNet + BrushNet + VAE 同时在显存里，小显存卡放不下
解决：测试脚本已默认enable_model_cpu_offload()，确认这行没被注释；再叠加enable_attention_slicing()；仍不够就降到 SD 1.5（比 SDXL 省一半显存）

坑 2：加载基础模型报错或权重乱码

现象：from_pretrained抛OSError，或加载成功但出图全噪
原因：base_model_path指向的不是 diffusers 格式目录——单个.ckpt/.safetensors文件不行，必须是含model_index.json、vae/、unet/子目录的完整文件夹
解决：用仓库 scripts/ 下的转换脚本（如convert_ldm_original_checkpoint_to_diffusers.py）把单文件权重转成 diffusers 格式再放入data/ckpt/

坑 3：依赖冲突，装完 import 就报错

现象：pip install -r requirements.txt后import diffusers或import accelerate失败
原因：requirements.txt固定了accelerate==0.20.3、Pillow==9.5.0、gradio==3.50.0，与你环境里已有版本打架
解决：按"纯净 conda 环境 →pip install -e .→ 再装 requirements"的顺序重来；不要先装最新 accelerate 再降级，直接装 requirements 里的固定版本

坑 4：SDXL 版出图全是 NaN 黑斑

现象：跑test_brushnet_sdxl.py时 VAE 解码后图像出现大片黑色/NaN
原因：SDXL 的官方 VAE 在 fp16 下编码不稳定
解决：test_brushnet_sdxl.py 已经给出做法——换用madebyollin/sdxl-vae-fp16-fix作为pipe.vae

排完这些坑，你就可以把"脚本级演示"升级成日常能用的工具了。

六、从演示到落地

演示脚本跑通只是开始，这几条路径能帮你把 BrushNet 变成日常工具：

① Gradio 交互界面（最快出活）

python examples/brushnet/app_brushnet.py

浏览器里上传照片、点选要修复的区域（背后是 SAM 自动分割）、填 prompt 即可出图，注意提前把 SAM 权重放到data/ckpt/sam_vit_h_4b8939.pth。

② Docker 部署

仓库 docker/ 目录已提供多套 diffusers 官方 Dockerfile（CUDA / CPU / xformers / ONNX 等），可以基于docker/diffusers-pytorch-cuda/Dockerfile叠加自己的模型层：

FROM pytorch/pytorch:2.0.1-cuda11.8-cudnn8-runtime WORKDIR /app COPY . . RUN pip install -e . && pip install -r examples/brushnet/requirements.txt CMD ["python", "examples/brushnet/app_brushnet.py"]

③ 量化评测

用 evaluate_brushnet.py 对整批结果打分（Image Reward、HPS V2.1、PSNR、LPIPS、CLIP Similarity 等），输出 CSV 汇总表，适合调参数前后做 A/B 对比：

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

④ 用自己的数据微调

训练脚本支持分割掩码和随机掩码两种模式（数据为*.tar的 WebDataset 格式）：

accelerate launch examples/brushnet/train_brushnet.py \ --pretrained_model_name_or_path runwayml/stable-diffusion-v1-5 \ --output_dir runs/logs/custom_brushnet \ --train_data_dir data/custom_dataset \ --resolution 512 \ --learning_rate 5e-6 \ --train_batch_size 2

加上--random_mask即切换到随机掩码训练（通用修复），不加则默认分割掩码训练（物体移除更稳）。想要更高分辨率？仓库同样提供了 SDXL 版推理 test_brushnet_sdxl.py 与训练 train_brushnet_sdxl.py。

工具链到这里就齐了，最后用 3 句话收个尾。