BrushNet/README.md

BrushNet

This repository contains a fork of the implementation of the paper "BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion" that was used for generating PinPoint counterfactuals. Please, refer to the Original Project for details on the environment installation and a complete description of BrushNet's features.

Getting Started

Environment Requirement 🌍

BrushNet has been implemented and tested on Pytorch 1.12.1 with python 3.9.

Clone the repo:

git clone https://github.com/TencentARC/BrushNet.git

We recommend you first use conda to create virtual environment, and install pytorch following official instructions. For example:

conda create -n diffusers python=3.9 -y
conda activate diffusers
python -m pip install --upgrade pip
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

Then, you can install diffusers (implemented in this repo) with:

pip install -e .

After that, you can install required packages thourgh:

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

Data Download ⬇️

Dataset

After downloading the datasets (CC3M: https://ai.google.com/research/ConceptualCaptions/download and FACET https://ai.meta.com/datasets/facet-downloads/), edit the paths in ./examples/brushnet/inpaint_cc3m.py, and ./examples/brushnet/inpaint_facet.py.

Checkpoints

Checkpoints of BrushNet can be downloaded from here. The ckpt folder needs to contain pretrained checkpoints and pretrained Stable Diffusion checkpoint (realisticVisionV60B1_v51VAE from Civitai). The data structure should be like:

|-- data
    |-- BrushData
    |-- BrushDench
    |-- EditBench
    |-- ckpt
        |-- realisticVisionV60B1_v51VAE
            |-- model_index.json
            |-- vae
            |-- ...
        |-- segmentation_mask_brushnet_ckpt

🏃🏼 Running Scripts

Inference 📜

You can in-paint FACET with the script:

python examples/brushnet/inpaint_facet.py

You can experiment with different prompts for in-painting different values of protected attributes. For in-painting PP*, WB*, and C&A* setups, change the corresponding prompts to f"A photo of a man/woman who is a {category}.", otherwise leave it as f"A photo of a man/woman."

For in-paintinf CC3M use the script:

python examples/brushnet/inpaint_cc3m.py

Evaluation 📏

You can evaluate the image realism using the following script, first installing the required dependencies:

python examples/brushnet/evaluate_dir.py

Make sure to set the paths to the image directories before running the script.

Note that the evaluation script requires the following additional dependencies:

• open_clip
• hpsv2
• ImageReward
• CLIPScore