app.py

import gc
import threading

import gradio as gr
import spaces
import torch
from PIL import Image, ImageOps
from transformers import AutoModelForMultimodalLM, AutoProcessor


MODEL_ID = "Qwen/Qwen3-VL-4B-Instruct"
MAX_IMAGE_EDGE = 2048

model = None
processor = None
model_lock = threading.Lock()


STYLE_GUIDANCE = {
    "Auto-detect": "Infer the most faithful medium and visual style from the image.",
    "Photorealistic": "Write for a photorealistic image model; emphasize optics, lighting, materials, and natural detail.",
    "Cinematic": "Write for a cinematic frame; emphasize shot design, lens language, lighting, atmosphere, and color grade.",
    "Anime / illustration": "Write for an illustration model; emphasize line work, rendering, palette, stylization, and composition.",
    "Product photography": "Write for product photography; emphasize object geometry, materials, surface finish, studio light, and backdrop.",
    "Architecture / interior": "Write for architectural visualization; emphasize space, structure, materials, lighting, perspective, and decor.",
}

DETAIL_TOKENS = {
    "Balanced": 700,
    "Very detailed": 1100,
    "Maximum detail": 1500,
}


def load_model():
    """Load the model once, only when the first generation request arrives."""
    global model, processor
    if model is not None:
        return model, processor

    with model_lock:
        if model is not None:
            return model, processor
        if not torch.cuda.is_available():
            raise gr.Error("This Space needs a CUDA GPU. Select ZeroGPU or a GPU in the Space settings.")

        processor = AutoProcessor.from_pretrained(MODEL_ID)
        model = AutoModelForMultimodalLM.from_pretrained(
            MODEL_ID,
            dtype=torch.bfloat16,
            device_map="auto",
            low_cpu_mem_usage=True,
        ).eval()
        return model, processor


def prepare_image(image: Image.Image) -> Image.Image:
    image = ImageOps.exif_transpose(image).convert("RGB")
    if max(image.size) > MAX_IMAGE_EDGE:
        image.thumbnail((MAX_IMAGE_EDGE, MAX_IMAGE_EDGE), Image.Resampling.LANCZOS)
    return image


def build_instruction(style: str, detail: str, focus: str) -> str:
    extra_focus = (focus or "").strip()
    return f"""
Act as an expert prompt engineer performing image-to-prompt reconstruction.

Study the supplied image carefully and return a faithful, highly descriptive prompt that could recreate it in a modern text-to-image model. Describe only visually supported details; do not invent names, brands, hidden facts, or a backstory. If a detail is uncertain, use neutral visual language. Do not identify a real person.

Requested style treatment: {STYLE_GUIDANCE[style]}
Requested detail level: {detail}.
User's optional focus: {extra_focus if extra_focus else "No extra focus; cover the whole image evenly."}

Analyze and incorporate, where visible:
- primary subject(s), count, appearance, pose, expression, gaze, and interaction
- clothing, accessories, objects, textures, materials, and fine details
- environment, foreground, midground, background, and spatial relationships
- composition, crop, framing, viewpoint, perspective, symmetry, and depth
- lighting direction, quality, intensity, shadows, highlights, time-of-day cues, and atmosphere
- colors, contrast, palette, medium, rendering technique, and overall aesthetic
- camera/lens cues such as shot type, focal-length feel, aperture/depth of field, focus, and motion
- any legible text, reproduced exactly in quotation marks; omit text if it is not clearly readable

Output exactly these two sections and nothing else:

DETAILED PROMPT
A cohesive, generator-ready prompt in natural language. Prefer precise visual terms over vague praise. Do not mention this analysis, the source image, or uncertainty.

NEGATIVE PROMPT
A concise comma-separated list of defects and unwanted changes that are specifically useful for preserving this image's composition and quality.
""".strip()


def split_response(text: str):
    cleaned = text.strip()
    marker = "NEGATIVE PROMPT"
    if marker in cleaned:
        positive, negative = cleaned.split(marker, 1)
        positive = positive.replace("DETAILED PROMPT", "", 1).strip(" \n:#")
        negative = negative.strip(" \n:#")
    else:
        positive = cleaned.replace("DETAILED PROMPT", "", 1).strip(" \n:#")
        negative = ""
    return positive, negative


@spaces.GPU(duration=90)
def generate_prompt(image, style, detail, focus, progress=gr.Progress()):
    if image is None:
        raise gr.Error("Upload an image first.")

    progress(0.05, desc="Loading vision model…")
    vision_model, vision_processor = load_model()
    image = prepare_image(image)
    instruction = build_instruction(style, detail, focus)
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": instruction},
            ],
        }
    ]

    progress(0.2, desc="Reading image details…")
    inputs = vision_processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(vision_model.device)

    try:
        with model_lock, torch.inference_mode():
            generated = vision_model.generate(
                **inputs,
                max_new_tokens=DETAIL_TOKENS[detail],
                do_sample=False,
                repetition_penalty=1.05,
            )
        generated = generated[:, inputs.input_ids.shape[1] :]
        response = vision_processor.batch_decode(
            generated,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )[0]
    except torch.cuda.OutOfMemoryError as exc:
        gc.collect()
        torch.cuda.empty_cache()
        raise gr.Error("The GPU ran out of memory. Try a smaller image or a shorter detail level.") from exc

    progress(1, desc="Prompt ready")
    return split_response(response)


CSS = """
:root { --paper: #f7f4ed; --ink: #191816; --muted: #6f6a61; --accent: #6750e8; }
.gradio-container { max-width: 1180px !important; margin: 0 auto !important; background: var(--paper); }
.hero { padding: 2.4rem 0 1.2rem; }
.hero .kicker { color: var(--accent); font-weight: 800; letter-spacing: .14em; text-transform: uppercase; font-size: .75rem; }
.hero h1 { color: var(--ink); font-size: clamp(2.5rem, 6vw, 5.2rem); line-height: .94; letter-spacing: -.055em; margin: .45rem 0 .85rem; }
.hero p { color: var(--muted); max-width: 720px; font-size: 1.05rem; }
.panel { border: 1px solid #dcd6ca !important; border-radius: 20px !important; background: rgba(255,255,255,.58) !important; }
.run-btn { background: var(--accent) !important; border: 0 !important; color: white !important; font-weight: 800 !important; }
.output textarea { font-family: ui-monospace, SFMono-Regular, Consolas, monospace !important; line-height: 1.55 !important; }
.note { color: var(--muted); font-size: .82rem; padding: .8rem 0 1.5rem; text-align: center; }
"""


with gr.Blocks() as demo:
    gr.HTML("""
        <section class="hero">
          <div class="kicker">Image → Prompt</div>
          <h1>See every detail.<br>Write the whole scene.</h1>
          <p>Upload any image and turn it into a precise, generator-ready prompt covering subject, composition, lighting, camera, materials, palette, and style.</p>
        </section>
    """)

    with gr.Row(equal_height=False):
        with gr.Column(scale=5, elem_classes="panel"):
            source = gr.Image(
                type="pil",
                image_mode="RGB",
                label="Source image",
                sources=["upload", "clipboard"],
                height=470,
            )
            with gr.Row():
                style = gr.Dropdown(list(STYLE_GUIDANCE), value="Auto-detect", label="Prompt style")
                detail = gr.Radio(list(DETAIL_TOKENS), value="Very detailed", label="Detail level")
            focus = gr.Textbox(
                label="Optional focus",
                placeholder="Example: emphasize the lighting, outfit, exact composition, or product materials",
            )
            run = gr.Button("Generate detailed prompt", variant="primary", size="lg", elem_classes="run-btn")

        with gr.Column(scale=6):
            positive = gr.Textbox(
                label="Detailed prompt",
                lines=18,
                elem_classes="output",
            )
            negative = gr.Textbox(
                label="Negative prompt",
                lines=5,
                elem_classes="output",
            )

    gr.HTML("<div class='note'>Powered by Qwen3-VL-4B-Instruct · Generated descriptions may need human review for tiny or ambiguous details.</div>")

    run.click(
        fn=generate_prompt,
        inputs=[source, style, detail, focus],
        outputs=[positive, negative],
        api_name="generate_prompt",
    )
    source.upload(
        fn=generate_prompt,
        inputs=[source, style, detail, focus],
        outputs=[positive, negative],
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(css=CSS, show_error=True)