stitch-flux-modular

Stitch for FLUX — training-free bounding-box position control (Modular Diffusers custom block)

🤗 Hub: remyxai/stitch-flux-modular · 📄 Paper: arXiv:2509.26644 · 💻 Reference: ExplainableML/Stitch (unlicensed — see the clean-room note) · 📦 Monorepo: flux-recipes

Put this object at this location — no fine-tuning, no detector, no extra weights — as a Modular Diffusers custom block. Implements Stitch (arXiv:2509.26644): you give a global prompt plus a few (bounding box, sub-prompt) regions, and each object is generated confined to its box, cut out, composited onto the background, and blended. In the paper this lifts FLUX's GenEval-Position score from ~0.22 to ~0.70 with zero training.

Off-the-shelf FLUX routinely fails simple spatial requests — ask for "a red cube to the left of a blue sphere" and the objects swap, merge, or drift. Stitch fixes exactly that, and it is hard placement, unlike our regional-prompting block, which does soft per-region prompt routing (a base prompt plus per-region prompts in one pass, no boxes and no compositing). Use Stitch when you know where things should go.

Usage

import torch
from diffusers import ModularPipeline

pipe = ModularPipeline.from_pretrained("remyxai/stitch-flux-modular", trust_remote_code=True)
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda")

img = pipe(
    prompt="a red cube to the left of a blue sphere, on a plain studio backdrop",
    regions=[
        {"box": [0.05, 0.3, 0.45, 0.75], "prompt": "a red cube"},      # left
        {"box": [0.55, 0.3, 0.95, 0.75], "prompt": "a blue sphere"},   # right
    ],
    height=1024, width=1024,
).images[0]
img.save("stitch.png")

box is normalized [x0, y0, x1, y1] in [0, 1], origin top-left, y down. Pass regions=None (or []) and the block is a bit-exact no-op against stock FluxPipeline — no processor is installed and nothing is mutated.

How it works

Three phases, all training-free on the frozen FLUX weights:

1. Region Binding (region_bind_steps = S, default 4). For the first S steps the transformer runs once per region, plus once for the background prompt — K+1 passes per step, all started from the same noise so they stay registered token-for-token. In each object pass an additive joint-attention mask confines generation to the box: inside-box image tokens can't attend outside the box, outside-box tokens can't attend to the sub-prompt, and the sub-prompt can't attend outside the box. The mask is added to the attention logits before softmax; Q/K/V projections are untouched. The background pass runs unmasked.
2. Cutout + composite (at τ = S). For each object pass the block reads the text→image attention of one fixed head (cutout_head, default block 14 / head 20 for FLUX.1-dev), averages it over the non-pad text tokens, sorts descending, and keeps tokens until their cumulative mass reaches cutout_eta = 0.65 — then 2-D max-pools that token mask with cutout_kernel = 5 to close holes. Each object's foreground latent tokens are written into the background latent at the box.
3. Refine (steps S…T). The composite is denoised by a single unmasked pass on the full global prompt to the end, which blends the objects into one coherent scene. Decode.

Two implementation notes that matter on FLUX specifically: the joint attention matrix is (n_txt + n_img)-square with text first, so the box mask is built on the image block only and every box→token index is offset by n_txt; and packed FLUX latents are (B, seq, 64), so box mapping, foreground selection and compositing all index tokens on dim 1.

The attention processors are swapped for the duration of the call and restored in finally; with regions empty none are installed at all.

Compute: Phase A costs K+1 transformer passes per step for S steps (K = number of regions). At K ≤ 3 this is comfortable on a single A100; the passes are independent, so a VRAM-tight setup can batch them instead of running them sequentially. A 2-region 50-step run is ≈ 10 steps × 3 passes + 40 steps × 1 pass ≈ 1.4× a stock 50-step generation.

Key parameters

Tuning (defaults tuned on FLUX): region_bind_steps=4 + cutout_eta=0.65 were chosen because the paper's 10 / 0.95 pasted the whole box-fill and produced rectangular slabs on FLUX. Sweet spot is 2–3 objects (4+ may show minor per-region duplication). If objects come out boxy or duplicated, lower cutout_eta or region_bind_steps further (fewer confined steps let the refine de-boxify); if an object drifts out of its box, raise region_bind_steps. cutout_head is the one genuinely per-build knob — head_probe.ipynb sweeps cutout_head/region_bind_steps/cutout_eta and scores each so you can re-pick if your diffusers build indexes heads differently than the paper's 14/20.

Fallback: if the Cutout is unstable on your build, note that Region Binding alone already gives strong position accuracy (the paper's ablation reports ~81% of the full method's spatial benefit) at the cost of blend quality. Set cutout_eta=1.0 to keep the full-mass prefix and you get approximately that behaviour with the composite still in place.

Dependencies

diffusers (main / ≥ 0.40), transformers, accelerate, sentencepiece, protobuf. No extra weights, no detector, no LLM — prompt → box decomposition is out of scope by design: the caller supplies regions, like a layout-to-image API.

Attribution & AI assistance

Stitch by Jessica Bader, Mateusz Pach, Maria A. Bravo, Serge Belongie and Zeynep Akata. Clean-room implementation: the reference repo ExplainableML/Stitch carries no license, so none of its code was read or copied — this block was written from the paper's description alone (the same discipline as our panorama and regional-prompting ports; the method itself is not copyrightable).

Authored with AI assistance (Claude) and validated by the Remyx AI team; method credit to the Stitch authors. Uses FLUX.1-dev under its non-commercial license — this derivative inherits that restriction.

Citation

@misc{bader2025stitch,
  title={Stitch: Training-Free Position Control in Multimodal Diffusion Transformers},
  author={Bader, Jessica and Pach, Mateusz and Bravo, Maria A. and Belongie, Serge and Akata, Zeynep},
  year={2025}, eprint={2509.26644}, archivePrefix={arXiv}, primaryClass={cs.CV},
  url={https://arxiv.org/abs/2509.26644}
}