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

BrushNet通过独创的双分支扩散技术,提供高效精准的图像修复解决方案,适用于各种图像处理场景。
直达下载
回到上一页
clickgpt_line.png_noView
介绍

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

1. 功能与方法概览

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

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

github_ai_tool_brushnet_1

2. 开始使用

环境需求

  • Python 3.9
  • Pytorch 1.12.1

安装步骤

  1. 克隆代码仓库:

        
    git clone https://github.com/TencentARC/BrushNet.git
    
  2. 创建虚拟环境并安装依赖:

        
    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显示出了其强大的图像修复能力,特别是在处理复杂图像遮挡问题时,它的双分支结构设计显著提高了修复质量。

编程学习
编程学习 免费领取编程学习资料 进编程学习交流群
订阅号
视频号
公众号 关注公众号,回复关键字python领取大厂最新面试题
×
编程学习
免费领取编程学习资料 进编程学习交流群