app.py

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / any CUDA-touching import
import torch
import gradio as gr
from threading import Thread
from transformers import (
    Qwen3VLForConditionalGeneration,
    AutoProcessor,
    TextIteratorStreamer,
)
from qwen_vl_utils import process_vision_info

MODEL_ID = "prithivMLmods/OpenCaption-4B-VL-SFT-v1.0"

SYSTEM_PROMPT = """You are a detailed image captioning assistant.

Structure every caption as follows:

1. Open with one sentence naming the shot type (e.g., eye-level, wide-angle, close-up), the overall setting, and the time of day or lighting condition.

2. Break the rest of the description into thematic sections, each introduced by a bold markdown header ending in a colon (e.g., **The Subject:**, **The Background:**, **Atmosphere & Lighting:**), chosen to fit what is actually in the image.

3. Within each section, use bullet points to list specific, concrete details, including positions, colors, textures, materials, actions, spatial relationships, and any legible text or fine-grained visual elements. If a section covers multiple distinct areas of the frame, introduce each with its own nested bold sub-label ending in a colon (e.g., **Foreground Right:**, **Background:**) before its bullet points.

4. Close with a short, unheaded paragraph beginning with "In summary," that ties the scene together and conveys its overall mood or narrative.

Use precise, sensory, and fluent language throughout. Describe only what is visible in the image, avoid speculation or unsupported inferences, and do not use emojis."""

DEFAULT_PROMPT = "Provide a detailed caption for this image with fine-grained visual details."

print("Loading processor...")
processor = AutoProcessor.from_pretrained(MODEL_ID)

print("Loading model...")
model = Qwen3VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
).to("cuda").eval()
print("Model loaded.")


@spaces.GPU(duration=60)
def generate_caption(
    image,
    user_prompt=DEFAULT_PROMPT,
    max_new_tokens=1024,
    temperature=0.7,
    top_p=0.8,
    top_k=20,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a dense, structured image caption from an uploaded image.

    Args:
        image: Input image to caption.
        user_prompt: Instruction prompt for the captioning task.
        max_new_tokens: Maximum number of tokens to generate.
        temperature: Sampling temperature for generation.
        top_p: Nucleus sampling probability threshold.
        top_k: Top-k sampling limit.
    """
    if image is None:
        yield "Please upload an image first."
        return

    messages = [
        {
            "role": "system",
            "content": [{"type": "text", "text": SYSTEM_PROMPT}],
        },
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": user_prompt or DEFAULT_PROMPT},
            ],
        },
    ]

    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    ).to("cuda")

    streamer = TextIteratorStreamer(
        processor.tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    generation_kwargs = {
        **inputs,
        "streamer": streamer,
        "max_new_tokens": int(max_new_tokens),
        "do_sample": True,
        "temperature": float(temperature),
        "top_p": float(top_p),
        "top_k": int(top_k),
    }

    generation_error = {"error": None}

    def _run_generation():
        try:
            model.generate(**generation_kwargs)
        except Exception as e:
            generation_error["error"] = e
            try:
                streamer.end()
            except Exception:
                pass

    thread = Thread(target=_run_generation, daemon=True)
    thread.start()

    buffer = ""
    for new_text in streamer:
        buffer += new_text
        yield buffer

    thread.join(timeout=2.0)

    if generation_error["error"] is not None:
        if buffer.strip():
            yield buffer + f"\n\n[ERROR] {generation_error['error']}"
        else:
            yield f"[ERROR] Inference failed: {generation_error['error']}"
        return

    if not buffer.strip():
        yield "[ERROR] No output was generated."
        return


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

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 📝 OpenCaption-4B-VL-SFT\n"
            "Dense, fine-grained image captioning powered by "
            "[OpenCaption-4B-VL-SFT-v1.0](https://huggingface.co/prithivMLmods/OpenCaption-4B-VL-SFT-v1.0), "
            "a fine-tune of Qwen3-VL-4B-Instruct."
        )

        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="Input Image",
                    type="pil",
                    height=360,
                )
                prompt_input = gr.Textbox(
                    label="Prompt",
                    value=DEFAULT_PROMPT,
                    placeholder="Describe what you want the model to do with the image...",
                )
                run_btn = gr.Button("Generate Caption", variant="primary")

                with gr.Accordion("Advanced Settings", open=False):
                    max_tokens = gr.Slider(
                        label="Max New Tokens",
                        minimum=128,
                        maximum=4096,
                        value=1024,
                        step=64,
                    )
                    temp_slider = gr.Slider(
                        label="Temperature",
                        minimum=0.0,
                        maximum=2.0,
                        value=0.7,
                        step=0.05,
                    )
                    top_p_slider = gr.Slider(
                        label="Top-p",
                        minimum=0.1,
                        maximum=1.0,
                        value=0.8,
                        step=0.05,
                    )
                    top_k_slider = gr.Slider(
                        label="Top-k",
                        minimum=1,
                        maximum=100,
                        value=20,
                        step=1,
                    )

            with gr.Column(scale=1):
                output = gr.Markdown(
                    label="Caption",
                    value="Upload an image and click **Generate Caption** to get a detailed, structured description.",
                )

        gr.Examples(
            examples=[
                ["examples/cafe_interior.jpg"],
                ["examples/dog.jpg"],
                ["examples/cherry_blossom.jpg"],
                ["examples/hot_air_balloon.jpg"],
                ["examples/sushi_platter.jpg"],
            ],
            inputs=[image_input],
            fn=generate_caption,
            outputs=output,
            cache_examples=True,
            cache_mode="lazy",
        )

        run_btn.click(
            fn=generate_caption,
            inputs=[
                image_input,
                prompt_input,
                max_tokens,
                temp_slider,
                top_p_slider,
                top_k_slider,
            ],
            outputs=output,
        )

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