• ⚙️ Installation
  • 📦 Dataset
  • 📚 Training and Inference
  • ⚽️ Score and Evaluation
• ✨ Basic Usage
  • 📂 Generate Candidates
  • 🏷️ Construct Preference Data
  • 🏋️ Preference Alignment Training
  • 💯 Evaluation
• 🚀 Try Our Image Inpainting Model

⚙️ Installation

1) 📦 Dataset

1. Download BrushData
   • Hugging Face dataset: https://huggingface.co/datasets/random123123/BrushData
   • Each item in BrushData contains fields like:
     • 000279227.aesthetic_score (unused)
     • 000279227.caption (string)
     • 000279227.height (int)
     • 000279227.image (binary; can be read as PNG — we highly recommend saving images locally as .png)
     • 000279227.original_key (unused)
     • 000279227.segmentation (dict)
     • 000279227.url (unused)
     • 000279227.width (int)
2. Extract and build JSON
   • Untar all archives from BrushData, gather the useful fields, and create a JSON list like:

     [
       {
         "gt_image_path": "/your/path/to/000279227.png",
         "caption": "000279227.caption",
         "height": 000279227.height,
         "width": 000279227.width,
         "segmentation": 000279227.segmentation,
         "image_id": "00027_000279227"
       }
       // ...
     ]

2) 📚 Training and Inference

3) ⚽️ Score and Evaluation

• Place the extra/bpe_simple_vocab_16e6.txt.gz into ~/miniconda3/envs/test_eval/lib/python3.12/site-packages/hpsv2/src/open_clip/
• We modified score functions in ~/miniconda3/envs/test_eval/lib/python3.12/site-packages/hpsv2/__init__.py and ~/miniconda3/envs/test_eval/lib/python3.12/site-packages/hpsv2/img_score.py for multi-gpu evaluation, we also provide our modified version under extra/

✨ Basic Usage

1) 📂 Generate Candidates

1. Configure paths and seed
   • Set the JSON path from Installation and the random seed in either:
     • Sample part in /path/to/configs_flux.yaml
     • Sample part in /path/to/configs_brushnet.yaml
2. Run the sampler
   • Use one of:
     • scripts/sample_flux.py
     • scripts/sample_brushnet.py
   • For sampling, we use config files to manage everthing. Specify generator_seed in configs/configs_<model>.yaml to control diversity.
     bash

     1# BrushNet 2accelerate launch \ 3 --config_file /root/.cache/huggingface/accelerate/accelerate_default.yaml \ 4 scripts/sample_brushnet.py 5# FLUX.1 Fill 6accelerate launch \ 7 --config_file /root/.cache/huggingface/accelerate/accelerate_default.yaml \ 8 scripts/sample_flux.py

2) 🏷️ Construct Preference Data

1. Create annotation list
   • Follow Merge the annotations for sampled images in scripts/merge_score_jsons.ipynb to produce a list-style JSON for scoring.
2. Start vLLM server Following UnifiedReward to create the env for CodeGoat24/UnifiedReward-qwen-7b
   bash

   1conda activate vllm 2bash scripts/vllm_server.sh
3. Run scoring
   • General metrics:
   bash

   1conda activate eval 2accelerate launch 3 --config_file /root/.cache/huggingface/accelerate/accelerate_default.yaml \ 4 scripts/score.py \ 5 --metric <ensemble> \ 6 --annotation_path /path/to/annotation_list \ 7 --output_dir /path/to/output \
   • UnifiedReward only:
   bash

   1python /root/test_env/scripts/score_unifiedreward.py 2 --annotation_path /path/to/annotation_list 3 --output_dir /path/to/output 4 --seed 0 5 --port <port_you_specify_when_launching_vllm>
4. Merge scores into one json file
   • Continue in scripts/merge_score_jsons.ipynb to attach scores to annotations for training.

3) 🏋️ Preference Alignment Training

• For training, we use config files to manage everthing. You could find it in configs/configs_<model>.yaml.
• FLUX.1 Fill
  bash

  1conda activate train 2accelerate launch \ 3 --config_file /root/test_env/configs/accelerate_flux.yaml \ 4 scripts/train_flux_dpo.py
• BrushNet
  bash

  1conda activate train 2accelerate launch \ 3 --config_file /root/.cache/huggingface/accelerate/accelerate_default.yaml \ 4 /root/test_env/scripts/train_brushnet_dpo.py

4) 💯 Evaluation

1. Generate images for testing
   bash

   1conda activate train 2accelerate launch \ 3 --num_processes 8 \ 4 --mixed_precision bf16 \ 5 script gen.py \ 6 --ckpt_path /path/to/ckpt \ 7 --base_model_path [black-forest-labs/FLUX.1-Fill-dev|runwayml/stable-diffusion-v1-5] \ 8 --image_save_path /path/to/save/images \ 9 --mapping_file /path/to/BrushBench/mapping_file_list.json \ 10 --base_dir /path/to/BrushBench \ 11 --use_blended \ 12 --model flux \ 13 --benchmark brushbench \ 14 --num_steps 50 # To reproduce our results, 20 for FLUX.1 Fill and 10 for BrushNet
2. Evaluate images Remember to launch the vLLM server first
   bash

   1conda activate eval 2accelerate launch 3 --num_processes 8 \ 4 --mixed_precision bf16 \ 5 scripts/eval.py \ 6 --image_save_path /path/to/save/images \ 7 --benchmark brushbench \ 8 --mapping_file /path/to/BrushBench/mapping_file_list.json \ 9 --base_dir /path/to/BrushBench \ 10 --use_blend
3. Evaluate images with GPT4
   bash

   1python scripts/eval_gpt4.py \ 2 --save_path /path/to/save/results \ 3 --mapping_file /path/to/BrushBench/mapping_file_list.json \ 4 --mask_key inpainting_mask \ 5 --base_dir /path/to/bench \ 6 --image_dir /path/to/images/to/eval

🚀 Try Our Image Inpainting Model

1. BrushNet
   python

   1from diffusers import StableDiffusionBrushNetPipeline, BrushNetModel, UniPCMultistepScheduler 2import torch 3import cv2 4import numpy as np 5from PIL import Image 6 7brushnet = BrushNetModel.from_pretrained(ckpt_path, torch_dtype=torch.float16) 8pipe = StableDiffusionBrushNetPipeline.from_pretrained( 9 "runwayml/stable-diffusion-v1-5", brushnet=brushnet, torch_dtype=torch.float16 10).to("cuda") 11pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) 12 13init_image = cv2.imread(...)[:,:,::-1] 14mask_image = 1.*(cv2.imread(...).sum(-1)>255)[:,:,np.newaxis] 15init_image = init_image * (1-mask_image) 16init_image = Image.fromarray(init_image.astype(np.uint8)).convert("RGB") 17mask_image = Image.fromarray(mask_image.astype(np.uint8).repeat(3,-1)*255).convert("RGB") 18 19image = pipe( 20 caption, 21 init_image, 22 mask_image, 23 num_inference_steps=50, 24 generator=generator, 25 brushnet_conditioning_scale=brushnet_conditioning_scale 26).images[0] 27image.save(f"output.png")
2. FLUX.1 Fill
   python

   1import torch 2from diffusers import FluxFillPipeline, FluxTransformer2DModel 3from PIL import Image 4 5image = Image.open(...).convert("RGB") 6mask = Image.open(...).convert("RGB") 7ckpt_path = ... 8 9transformer = FluxTransformer2DModel.from_pretrained( 10 ckpt_path, torch_dtype=torch.bfloat16 11).to(device) 12pipe = FluxFillPipeline.from_pretrained( 13 "black-forest-labs/FLUX.1-Fill-dev", 14 torch_dtype=torch.bfloat16, 15 transformer=transformer 16).to("cuda") 17 18image = pipe( 19 prompt="...", 20 image=image, 21 mask_image=mask, 22 height=512, 23 width=512, 24 guidance_scale=30, 25 num_inference_steps=20, 26 generator=torch.Generator("cpu").manual_seed(0) 27).images[0] 28image.save(f"output.png")