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)