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/BrushNetBrushNet 是 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 python3.9 -y conda activate brushnet python -m pip install --upgrade pip第二步装 PyTorch按自己的 CUDA 版本选 index-urlcu116 / cu118 都行pip install torch1.12.1cu116 torchvision0.13.1cu116 torchaudio0.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里锁定了Pillow9.5.0和accelerate0.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参数注入 pipelinefrom diffusers import StableDiffusionBrushNetPipeline, BrushNetModel, UniPCMultistepScheduler import torch brushnet BrushNetModel.from_pretrained(data/ckpt/segmentation_mask_brushnet_ckpt, torch_dtypetorch.float16) pipe StableDiffusionBrushNetPipeline.from_pretrained( data/ckpt/realisticVisionV60B1_v51VAE, brushnetbrushnet, torch_dtypetorch.float16, low_cpu_mem_usageFalse)第二段准备输入。掩码要做成二值图白修复区原图则把掩码区域置黑这样原图信息不会漏进修复区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_steps50, generatorgenerator, brushnet_conditioning_scale1.0, ).images[0] image.save(output.png)在仓库根目录执行python examples/brushnet/test_brushnet.py即可得到output.png。️ 关键参数速查推理时真正值得动的参数就几个都通过pipe(...)传入参数推荐值作用备注brushnet_conditioning_scale1.0修复区跟随 BrushNet 条件的强度1.0 时修复更依赖原模型融合更自然Gradio 界面里对应 control_strengthnum_inference_steps20~50去噪步数配合 UniPC 调度器20 步左右已有可用效果guidance_scale7.5默认文本提示词的引导强度调高更像 prompt调低更贴原图风格generator种子任意固定随机种子同种子可复现结果换种子可重roll 三个高频用法跑 SDXL 版本。流程几乎一样只是换 pipeline 类和调度器参考examples/brushnet/test_brushnet_sdxl.pyfrom diffusers import StableDiffusionXLBrushNetPipeline, DPMSolverMultistepScheduler pipe StableDiffusionXLBrushNetPipeline.from_pretrained( base_model_path, brushnetbrushnet, torch_dtypetorch.float16, use_safetensorsTrue) 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。原因环境里同时存在官方 diffusersimport 被旧包抢走。解法确认pip show diffusers指向的是本地仓库路径必要时先pip uninstall diffusers再装。 进阶入口用自己的数据微调训练脚本是examples/brushnet/train_brushnet.pySD1.5和train_brushnet_sdxl.pySDXL数据按 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。 接下来可以试试把brushnet_conditioning_scale从 1.0 降到 0.8 再出一张图对比修复区与边缘的融合差异。用同一张test_image.jpg分别跑 segmentation 和 random 两个掩码权重看两种训练的泛化差别。打开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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考