README.md

---
license: apache-2.0
tags:
  - image-editing
  - diffusion
  - flux
  - multi-instance
---

# 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

```bash
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

```bash
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:

```bash
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)

```bash
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)

```bash
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`).

```bash
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)

```bash
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/
```