app.py

import os

# HF Spaces often set HTTP_PROXY; without NO_PROXY, Gradio's localhost health check fails.
os.environ.setdefault("NO_PROXY", "localhost,127.0.0.1,::1")
os.environ.setdefault("no_proxy", "localhost,127.0.0.1,::1")

import gradio as gr
from PIL import Image

from pipeline import run_pipeline
from pipeline.background_replace import warmup as warmup_rembg
from pipeline.planner import warmup
from pipeline.text_regions import warmup as warmup_ocr

_warmed_up = False


def _ensure_warmup():
    global _warmed_up
    if _warmed_up:
        return
    print("Loading embedding planner, RapidOCR, and rembg...")
    try:
        warmup()
        warmup_ocr()
        warmup_rembg()
        print("Models loaded.")
    except Exception as exc:
        print(f"Warmup failed (will retry on next request): {exc}")
        return
    _warmed_up = True

EXAMPLE_PROMPTS = [
    ["rotate randomly"],
    ["rotate left"],
    ["rotate right"],
    ["flip upside down"],
    ["rotate 180 degrees"],
    ["gamma brighten"],
    ["make it brighter"],
    ["more contrast"],
    ["vintage warm look"],
    ["replace background"],
    ["replace background with beach sunset"],
    ["replace background with snowy mountains"],
    ["replace background with anything"],
    ["replace background with neon city at night and rotate"],
    ["cover a piece but not text"],
    ["just cover"],
    ["cover and add cutout"],
    ["cover and several cutouts avoid text"],
    ["add cutout avoid text"],
    ["add cutout"],
    ["transparent cutout"],
    ["perspective transform"],
    ["cartoon style"],
    ["anime style"],
    ["starry night style"],
    ["pencil sketch"],
    ["stylize"],
    ["mosaic painting"],
    ["impressionist look"],
    ["flip horizontally"],
    ["blur everything softly"],
]

SUPPORTED_EXAMPLES = "\n".join(f"- `{row[0]}`" for row in EXAMPLE_PROMPTS)


def _format_spatial(applied_spatial: list[dict]) -> str:
    if not applied_spatial:
        return ""
    lines = []
    for item in applied_spatial:
        op = item.get("op", "")
        if op == "replace_background":
            q = item.get("query") or "(random)"
            src = item.get("source", "?")
            lines.append(f"- replace_background: \"{q}\" via {src}")
            continue
        if op in ("cover_and_cutout", "cover_and_cutout_avoid_text"):
            covers = item.get("cover_rects", [])
            cutouts = item.get("cutouts", [])
            fallback = " (fallback position)" if item.get("used_fallback") else ""
            lines.append(
                f"- {op}: {len(covers)} cover(s){fallback}, {len(cutouts)} cutout(s)"
            )
            for idx, rect in enumerate(covers, 1):
                lines.append(f"  - cover {idx} at {rect}")
            for idx, detail in enumerate(cutouts, 1):
                shape = detail.get("shape", "?")
                style = detail.get("style", "?")
                lines.append(f"  - cutout {idx} ({shape}, {style}) at {detail.get('rect')}")
            continue
        if op in ("add_cutout", "cutout") and item.get("cutouts"):
            cutouts = item.get("cutouts", [])
            lines.append(f"- {op}: {len(cutouts)} cutout(s)")
            for idx, detail in enumerate(cutouts, 1):
                shape = detail.get("shape", "?")
                style = detail.get("style", "?")
                lines.append(f"  - cutout {idx} ({shape}, {style}) at {detail.get('rect')}")
            continue
        rect = item.get("rect", ())
        style = item.get("cutout_style")
        shape = item.get("shape")
        style_note = f", {style}" if style else ""
        shape_note = f", {shape}" if shape else ""
        fallback = " (fallback position)" if item.get("used_fallback") else ""
        lines.append(f"- {op}{style_note}{shape_note} at {rect}{fallback}")
    return "\n".join(lines)


def process(image, instruction, strength):
    if image is None:
        return None, None, "Please upload an image."

    _ensure_warmup()
    result = run_pipeline(image, instruction, strength=strength)

    if not result["supported"]:
        return (
            image,
            image,
            f"**Not supported**\n\n{result['reason']}\n\n**Try these prompts:**\n{SUPPORTED_EXAMPLES}",
        )

    tags = ", ".join(result["applied_tags"]) or "(none)"
    spatial_info = _format_spatial(result.get("applied_spatial", []))
    scores = result.get("planned_scores", [])
    score_info = ", ".join(f"{name} ({score:.2f})" for name, score in scores) if scores else "(none)"
    info = (
        f"**Instruction:** {result['instruction']}\n\n"
        f"**Matched keywords:** {score_info}\n\n"
        f"**Applied transforms:** {tags}\n\n"
    )
    if spatial_info:
        info += f"**Spatial ops:**\n{spatial_info}\n\n"
    if result.get("text_regions_found", 0) > 0:
        info += f"**Text regions detected:** {result['text_regions_found']}\n"

    return image, result["image"], info


with gr.Blocks(title="Text-Driven Image Augmentation") as demo:
    gr.Markdown(
        "# Text-Driven Image Augmentation\n"
        "Upload an image and describe what you want. Transforms are chosen by **semantic similarity** "
        "to augmentation keywords (1–5 matches above a similarity threshold). "
        "Background replacement uses CC-licensed photos from Openverse "
        "(web only — describe any scene, e.g. `replace background with sunset over Paris`)."
    )

    with gr.Row():
        input_image = gr.Image(type="pil", label="Input image")
        instruction = gr.Textbox(
            label="Instruction",
            placeholder='e.g. "vintage warm look", "replace background with snowy mountains", "blur softly"',
            lines=2,
        )

    strength = gr.Slider(
        minimum=0.5,
        maximum=1.5,
        value=1.0,
        step=0.1,
        label="Augmentation strength",
    )

    gr.Examples(
        examples=EXAMPLE_PROMPTS,
        inputs=[instruction],
        label="Example prompts (click to use)",
    )

    run_btn = gr.Button("Augment", variant="primary")

    with gr.Row():
        original_out = gr.Image(label="Original", interactive=False)
        augmented_out = gr.Image(label="Augmented", interactive=False)

    info_out = gr.Markdown()

    run_btn.click(
        fn=process,
        inputs=[input_image, instruction, strength],
        outputs=[original_out, augmented_out, info_out],
    )

if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        share=False,
    )