app.py

import os

# Allocator config helps with transient spikes from the pixel-space / large
# attention-mask ops in the region-attention editor.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  (must precede torch / CUDA-touching imports)
import time  # noqa: E402
import random  # noqa: E402

import torch  # noqa: E402
import gradio as gr  # noqa: E402
from PIL import Image  # noqa: E402

from replan.pipelines.replan import RePlanPipeline  # noqa: E402

# -----------------------------------------------------------------------------
# Model configuration
# -----------------------------------------------------------------------------
# RePlan couples a reasoning VLM planner (Qwen2.5-VL-7B, fine-tuned by the
# authors) with a diffusion editor. We use FLUX.2 [klein] 4B as the editor:
# it is Apache-2.0 (ungated), compact (~16 GB bf16), distilled (4 steps), and
# is one of the backbones the RePlan authors explicitly support.
VLM_CKPT = "TainU/RePlan-Qwen2.5-VL-7B"
DIFFUSION_MODEL = "black-forest-labs/FLUX.2-klein-4B"
PIPELINE_TYPE = "klein"

# klein-specific defaults from RePlan's run_replan.py
DEFAULT_STEPS = 4
DEFAULT_EXPAND_VALUE = 0.15
DEFAULT_ATTENTION_SWITCH_STEP = 0.05
MAX_SEED = 2**31 - 1

# Load both stages once at module scope, eagerly moved to CUDA. ZeroGPU packs
# the weights to disk at startup and streams them into VRAM on the first
# @spaces.GPU call. enable_flex_attn is left False (see below), so the editor
# uses a dense attention mask via SDPA and never needs torch.compile.
pipeline = RePlanPipeline(
    vlm_ckpt_path=VLM_CKPT,
    diffusion_model_name=DIFFUSION_MODEL,
    pipeline_type=PIPELINE_TYPE,
    output_dir=None,
    device="cuda",
    torch_dtype=torch.bfloat16,
    vlm_prompt_template_path="replan.txt",
    init_vlm=True,
    image_dir=None,
    stage_offload="none",
)


# Measured GPU time for a single-region 4-step edit is ~5-6s. Multi-region
# instructions produce longer VLM reasoning (up to 2048 tokens) + more diffusion
# work, so we keep a comfortable-but-lean margin.
@spaces.GPU(duration=60)
def edit_image(
    image: Image.Image,
    instruction: str,
    seed: int = 0,
    randomize_seed: bool = True,
    steps: int = DEFAULT_STEPS,
    guidance_scale: float = 4.0,
    expand_value: float = DEFAULT_EXPAND_VALUE,
    progress=gr.Progress(track_tqdm=True),
):
    """Edit an image following a complex natural-language instruction.

    RePlan first runs a reasoning VLM planner that decomposes the instruction
    into region-specific bounding boxes + local hints and a global prompt, then
    executes the edit with a diffusion backbone using training-free
    region-constrained attention injection.

    Args:
        image: The input image to edit.
        instruction: The natural-language editing instruction.
        seed: RNG seed for reproducibility.
        randomize_seed: If True, pick a fresh random seed each run.
        steps: Number of diffusion denoising steps.
        guidance_scale: Classifier-free guidance scale for the editor.
        expand_value: Ratio by which planned bounding boxes are expanded.

    Returns:
        A tuple of (edited image, VLM plan/reasoning text, the seed used).
    """
    if image is None:
        raise gr.Error("Please provide an input image.")
    if not instruction or not instruction.strip():
        raise gr.Error("Please provide an editing instruction.")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)

    if isinstance(image, str):
        image = Image.open(image)
    image = image.convert("RGB")

    t0 = time.perf_counter()
    results = pipeline.edit_single(
        image=image,
        instruction=instruction.strip(),
        expand_value=float(expand_value),
        expand_mode="ratio",
        attention_switch_step=DEFAULT_ATTENTION_SWITCH_STEP,
        bboxes_attend_to_each_other=True,
        symmetric_masking=False,
        enable_flex_attn=False,          # dense-mask SDPA path (no torch.compile)
        flex_attn_use_bitmask=True,
        skip_save=True,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance_scale),
        generator=torch.manual_seed(seed),
    )
    elapsed = time.perf_counter() - t0

    edited = results["edited_image"]
    response = results.get("vlm_response", "") or ""
    global_prompt = results.get("global_prompt", "") or ""
    bbox_data = results.get("bbox_data") or []

    # Build a compact, readable plan summary from the raw VLM response.
    plan_lines = [f"Inference time: {elapsed:.1f}s  |  seed: {seed}", ""]
    if global_prompt.strip():
        plan_lines.append(f"Global edit: {global_prompt.strip()}")
    if bbox_data:
        plan_lines.append("")
        plan_lines.append(f"Planned regions ({len(bbox_data)}):")
        for i, b in enumerate(bbox_data, 1):
            hint = b.get("hint", "")
            box = b.get("bbox_2d", "")
            plan_lines.append(f"  {i}. {hint}  @ bbox {box}")
    plan_lines.append("")
    plan_lines.append("--- Full VLM reasoning ---")
    plan_lines.append(response.strip())
    plan_text = "\n".join(plan_lines)

    return edited, plan_text, seed


# -----------------------------------------------------------------------------
# UI
# -----------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES = [
    ["examples/cup.png", "Replace the glass that has been used and left on the desk with a small potted plant"],
    ["examples/crowd.png", "Find the woman with light blue backpack and change the color of her shoes to red"],
    ["examples/festival.png", "The sun has now set, and someone has decorated the large bush on the right with colorful fairy lights for the evening's festivities."],
    ["examples/sunglasses.png", "Remove the sunglasses from the person sitting on the left."],
]

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # RePlan — Reasoning-guided Region Planning for Image Editing

            RePlan tackles **complex instruction-based image editing** by first having a
            reasoning VLM planner decompose your instruction into region-specific guidance
            (bounding boxes + local hints), then executing the edit with training-free
            **region-constrained attention injection** on a diffusion editor.

            Planner: [`TainU/RePlan-Qwen2.5-VL-7B`](https://huggingface.co/TainU/RePlan-Qwen2.5-VL-7B) ·
            Editor: [`FLUX.2-klein-4B`](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B) ·
            [Paper](https://huggingface.co/papers/2512.16864) ·
            [Code](https://github.com/dvlab-research/RePlan)
            """
        )

        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Input image", type="pil", height=380)
                instruction = gr.Textbox(
                    label="Editing instruction",
                    placeholder="e.g. Find the woman with the blue backpack and change her shoes to red",
                    lines=2,
                )
                run_btn = gr.Button("Edit image", variant="primary")
                with gr.Accordion("Advanced settings", open=False):
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                    seed = gr.Number(label="Seed", value=0, precision=0)
                    steps = gr.Slider(
                        label="Denoising steps", minimum=2, maximum=12, step=1,
                        value=DEFAULT_STEPS,
                    )
                    guidance_scale = gr.Slider(
                        label="Guidance scale", minimum=1.0, maximum=6.0, step=0.5,
                        value=4.0,
                    )
                    expand_value = gr.Slider(
                        label="Region expansion (ratio)", minimum=0.0, maximum=0.5,
                        step=0.05, value=DEFAULT_EXPAND_VALUE,
                    )
            with gr.Column():
                output_image = gr.Image(label="Edited image", height=380)
                plan_out = gr.Textbox(
                    label="Planner reasoning & regions", lines=12,
                )

        gr.Examples(
            examples=EXAMPLES,
            inputs=[input_image, instruction],
            outputs=[output_image, plan_out, seed],
            fn=edit_image,
            cache_examples=True,
            cache_mode="lazy",
        )

    inputs = [input_image, instruction, seed, randomize_seed, steps, guidance_scale, expand_value]
    outputs = [output_image, plan_out, seed]
    run_btn.click(edit_image, inputs=inputs, outputs=outputs, api_name="edit")
    instruction.submit(edit_image, inputs=inputs, outputs=outputs, api_name=False)

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)