MICE: Multi-Instance Controlled Editing

Code for spatially-controlled multi-instance image editing with FLUX.2-klein (4B). MICE constrains cross-instance attention leakage during inference so that each edit stays within its intended spatial region.

Requirements

pip install torch diffusers transformers pillow tqdm loguru numpy

The model weights are downloaded automatically from HuggingFace on first run:

black-forest-labs/FLUX.2-klein-4B

Quick start — single image

python infer.py \
  --image input.jpg \
  --edits "Replace the cat with a dog" "Replace the lamp with a vase" \
  --masks mask_cat.png mask_lamp.png \
  --output result.png

Provide one --edits string and one --masks path per instance you want to edit. Masks are binary images (white = instance region, black = background).

If you do not have masks, use normalized bounding boxes instead:

python infer.py \
  --image input.jpg \
  --edits "Replace the cat with a dog" \
  --bboxes "0.1 0.2 0.5 0.8" \
  --output result.png

All options

Flag	Default	Description
--image	(required)	Path to the input image
--edits	(required)	Edit instruction per instance
--masks	—	Binary mask image per instance (mutually exclusive with --bboxes)
--bboxes	—	Normalized x1 y1 x2 y2 bbox per instance (mutually exclusive with --masks)
--output	result.png	Output path
--model	black-forest-labs/FLUX.2-klein-4B	HuggingFace model ID or local path
--num_inference_steps	4	Number of denoising steps
--kernel_size	11	Spatial kernel size for attention falloff
--temperature	3.0	Steepness of spatial falloff
--double_layers	0,5	Double-stream block range for attention binding
--single_layers	0,20	Single-stream block range for attention binding
--strict	flag	Stricter cross-instance isolation

Benchmark inference

Scripts for running inference over the full MICE-Bench and LoMOE-Bench datasets are also included.

MICE-Bench (ground-truth masks)

python infer_flux2_mice.py \
  --exp_name my_experiment \
  --dataset_root ../mice_bench \
  --num_inference_steps 4 \
  --kernel_size 11 \
  --temperature 3.0 \
  --hard_image_attribute_binding_list_double 0,5 \
  --hard_image_attribute_binding_list_single 0,20 \
  --use_masks

Results are saved to results_micebench/mice_flux2_klein_<exp_name>/.

MICE-Bench (SAM3 masks)

python infer_flux2_mice_sam3.py \
  --exp_name my_experiment \
  --dataset_root /path/to/mice_bench_sam3 \
  --num_inference_steps 4 \
  --kernel_size 11 \
  --temperature 3.0 \
  --use_masks

Results are saved to results_micebench_sam3/mice_flux2_sam3_<exp_name>/.

LoMOE-Bench (ground-truth masks)

LoMOE-Bench is expected at benchmark/data/LoMOE-Bench/ relative to the repository root (path is hardcoded in lomoe_dataset.py).

python infer_flux2_lomoe.py \
  --exp_name my_experiment \
  --num_inference_steps 4 \
  --kernel_size 11 \
  --temperature 3.0 \
  --hard_image_attribute_binding_list_double 0,5 \
  --hard_image_attribute_binding_list_single 0,20 \
  --use_masks

Results are saved to results/lomoe_flux2_klein_<exp_name>/.

LoMOE-Bench (SAM3 masks)

python infer_flux2_lomoe_sam3.py \
  --exp_name my_experiment \
  --dataset_root /path/to/lomoe_bench_sam3 \
  --num_inference_steps 4 \
  --kernel_size 11 \
  --temperature 3.0 \
  --use_masks

Results are saved to results_sam3/lomoe_flux2_sam3_<exp_name>/.

Dataset structure

MICE-Bench (--dataset_root):

mice_bench/
  LoMOE.json
  images/
  masks/

SAM3 datasets (--dataset_root):

<dataset_root>/
  LoMOE_mobius.json   # mask paths point to SAM3-generated masks
  images/
  masks/

LoMOE-Bench (hardcoded path benchmark/data/LoMOE-Bench/):

benchmark/data/LoMOE-Bench/
  LoMOE.json
  LoMOE_multi_turn_faithful.json
  LoMOE_multi_turn_enhanced.json
  utils/
    mask_orig_prompts.txt
    text_prompts.txt
  images/
  masks/