BitResEdit

Training-free image editing with the Infinity visual autoregressive model.

BitResEdit edits a source image to match a target prompt without any fine-tuning. For each next-scale generation step it applies per-bit source-negative AR-CFG guidance and composites a per-scale gated source-anchored residual through the edit mask, so background regions are preserved while the edited region follows the target prompt.

This repository contains:

• BitResEdit — the method (ResidualEdit/bitresedit.py).
• AREdit — a re-implementation of the AREdit baseline (Wang et al., 2025) used for comparison (AREdit/aredit.py).
• The shared Infinity inference backbone (infinity/, tools/run_infinity.py).
• The PIE-Bench evaluation harness (evaluation/).

---

1. Results (PIE-Bench v1, Infinity-2B)

Under the default GT-mask localization:

Method	PSNR ↑	LPIPS ×10³ ↓	MSE ×10⁴ ↓	SSIM ×10² ↑	CLIP-Whole ↑	CLIP-Edited ↑
BitResEdit (ours)	26.67	43.21	31.61	91.71	27.13	24.54
AREdit (baseline)	24.19	87.00	—	83.70	25.42	22.77

CLIP-Whole = CLIP similarity to the target caption over the whole image; CLIP-Edited = CLIP similarity over the edited region only. Background-preservation metrics (PSNR/LPIPS/MSE/SSIM) are computed over the unedited region.

---

2. Environment

Python 3.12 with a CUDA-enabled PyTorch (≥ 2.5.1, required for FlexAttention). A single modern GPU is sufficient for the Infinity-2B model at 512 px.

conda create -n bitresedit python=3.12 -y
conda activate bitresedit
pip install -r requirements.txt

flash-attn requires a CUDA toolchain matching your PyTorch build; if a prebuilt wheel is unavailable, install it following the flash-attention instructions.

Measured GPU memory: the Infinity-2B evaluation reserves ≈ 37 GB per GPU at the default 1M-token scale schedule, so it fits comfortably on a 48 GB (or 40 GB A100) card. A single-image edit needs less; reduce runtime.pn/runtime.size if you are memory-constrained.

---

3. Model weights

Download the Infinity-2B weights and the flan-t5-xl text encoder into weights/ (git-ignored):

File / dir	Source	Target path
infinity_2b_reg.pth	FoundationVision/Infinity	weights/infinity_2b_reg.pth
infinity_vae_d32reg.pth	FoundationVision/Infinity	weights/infinity_vae_d32_reg.pth
flan-t5-xl/	google/flan-t5-xl	weights/flan-t5-xl/

⚠️ VAE filename: the Hugging Face file is named infinity_vae_d32reg.pth, but the config expects infinity_vae_d32_reg.pth. Rename (or symlink) it, or override model.vae_path on the command line.

Example:

mkdir -p weights
huggingface-cli download FoundationVision/infinity infinity_2b_reg.pth --local-dir weights
huggingface-cli download FoundationVision/infinity infinity_vae_d32reg.pth --local-dir weights
mv weights/infinity_vae_d32reg.pth weights/infinity_vae_d32_reg.pth
huggingface-cli download google/flan-t5-xl --local-dir weights/flan-t5-xl

The model paths are configured in evaluation/conf/model/default.yaml (model_base_dir: weights).

---

4. Data (PIE-Bench v1)

BitResEdit is evaluated on the PIE-Bench v1 benchmark (700 images, 10 edit categories). Obtain it from the PnP-Inversion / Direct-Inversion release and place it at data/PIE-Bench_v1/ (git-ignored):

data/PIE-Bench_v1/
├── annotation_images/        # <category>/<id>.jpg  (700 images, 10 categories)
└── mapping_file.json         # per-image prompts + RLE edit mask

The ground-truth edit mask is the run-length-encoded mask field inside mapping_file.json (decoded automatically) — no separate mask download is needed. The data path is set in evaluation/conf/config.yaml (data.path: data/PIE-Bench_v1).

---

5. Reproduce the PIE-Bench numbers

Run from the repository root with the environment active. Each command runs inference, computes metrics, and writes an HTML report.

# BitResEdit (apex preset; GT-mask localization)
python evaluation/run_full_eval.py \
    --method bitresedit --edit-config bitresedit --gpus 7 \
    --output-dir outputs/bitresedit_eval

# AREdit baseline (GT-mask localization, as reported in the paper)
python evaluation/run_full_eval.py \
    --method aredit --gpus 7 --edit-override use_gt_mask=true \
    --output-dir outputs/aredit_eval

⚠️ --edit-config takes a bare preset name (bitresedit), not a path. Passing a path resolves to a doubled .yaml.yaml and raises FileNotFoundError.

--gpus selects the physical GPU(s); the harness sets CUDA_VISIBLE_DEVICES per worker, so you do not need to set it yourself.

Useful flags:

• --gpus 6,7 — shard the run across multiple GPUs.
• --num-samples 20 — quick smoke run on a subset.
• --example-key 000000000010 — edit a single image.
• --skip-comparisons — skip the per-image figure grid (faster).
• --skip-inference — recompute metrics only from existing outputs.

Reading the results

Each --output-dir contains:

• annotation_images/ — edited images (mirrors the input tree).
• metrics.csv — per-image metrics.
• summary.json — per-category ("0".."9") and "overall" means.
• evaluation_report.html — a browsable report with best/worst examples.
• comparisons/ — per-scale source-vs-edited figures.

Metric-name mapping (summary.json → table above): clip_similarity_target_image → CLIP-Whole, clip_similarity_target_image_edit_part → CLIP-Edited, ssim_unedit_part → SSIM, psnr_unedit_part → PSNR, lpips_unedit_part → LPIPS, structure_distance → Structure.

A run is skipped if summary.json + evaluation_report.html already exist in the output dir — delete them or use a fresh dir to re-run.

---

6. Single-image demo

# AREdit (mask-free)
python scripts/demo.py --method aredit \
    --image path/to/source.jpg \
    --source-prompt "a cat sitting on a sofa" \
    --target-prompt "a [dog] sitting on a sofa" \
    --out edited.png

# BitResEdit (mask-only — provide a binary edit mask, white = edit region)
python scripts/demo.py --method bitresedit --edit-config bitresedit \
    --image path/to/source.jpg --mask path/to/mask.png \
    --source-prompt "a cat sitting on a sofa" \
    --target-prompt "a [dog] sitting on a sofa" \
    --out edited.png

---

7. Tests

python -m pytest unit_test -q

---

8. Repository layout

AREdit/          AREdit method + shared InfinityEditor base
ResidualEdit/    BitResEdit method + residual editor base + AR-CFG helpers
infinity/        Infinity VAR transformer + BSQ-VAE (inference subset)
tools/           run_infinity.py — backbone loader bridge
evaluation/      run_full_eval.py harness, metrics, configs
scripts/         demo.py single-image demo
unit_test/       unit tests

---

9. License & acknowledgements

This project is released under the terms in LICENSE and builds on:

• Infinity — the visual autoregressive backbone.
• PIE-Bench — the evaluation benchmark.
• AREdit (Wang et al., 2025) — the training-free VAR editing baseline.

If you use BitResEdit, please cite the accompanying paper.