"""Gradio Workflow app that runs MiniCPM-V 4.7 35B-A3B in-process on ZeroGPU.

The canvas's MiniCPM-V nodes are `kind: "fn"` operators backed by the
functions below: the model is loaded once at module scope (the `spaces`
package packs the weights and transfers them when a GPU is allocated) and
each node call books its own GPU slot via `@spaces.GPU`. The graph lives in
workflow.json — the app's built-in toggle switches between the app view (run
the pipelines) and the Workflow canvas (re-wire nodes, edit settings, save
the graph).

Requires the Space hardware to be set to ZeroGPU (Settings -> Hardware).
"""

import os
import urllib.parse

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

# `spaces` has to be imported before anything initializes CUDA. It only exists
# on Hugging Face Spaces, so shim it for local runs.
try:
    import spaces
except ImportError:

    class _SpacesShim:
        @staticmethod
        def GPU(duration=60, **kwargs):
            return lambda fn: fn

    spaces = _SpacesShim()

import torch
from PIL import Image

import gradio as gr

MODEL_ID = "openbmb/MiniCPM-V-4.7-35B-A3B"
# "16x" merges visual tokens for efficiency; "4x" keeps 4x as many for finer
# detail. Must be passed to BOTH apply_chat_template and generate (model card).
DOWNSAMPLE_MODE = "16x"
MAX_NEW_TOKENS = 512

# "Built with ML Intern" badge, injected into the page by launch(css=, js=).
# gr.Workflow must be the top-level app (it cannot nest inside gr.Blocks), so
# there is no gr.HTML slot — the badge is a fixed-position element appended to
# <body> on page load, visible in both the app view and the canvas view.
_BADGE_CSS = """
#ml-intern-badge {
    position: fixed;
    left: 12px;
    bottom: 12px;
    z-index: 1000;
    line-height: 0;
}
#ml-intern-badge img { height: 28px; }
"""

_BADGE_JS = """
function addMlInternBadge() {
    if (document.getElementById('ml-intern-badge')) return;
    const badge = document.createElement('a');
    badge.id = 'ml-intern-badge';
    badge.href = 'https://hf.co/chat?mode=ml-intern';
    badge.target = '_blank';
    badge.rel = 'noopener';
    const img = document.createElement('img');
    img.src = 'https://huggingface.co/datasets/huggingface/badges/resolve/main/built-with-ml-intern-sm.svg';
    img.alt = 'Built with ML Intern';
    badge.appendChild(img);
    document.body.appendChild(badge);
}
if (document.body) {
    addMlInternBadge();
} else {
    document.addEventListener('DOMContentLoaded', addMlInternBadge);
}
"""

# Load once at module scope and move to cuda here, not inside the GPU
# functions: the spaces package intercepts the move, packs the weights, and
# transfers them when a GPU is actually allocated — module scope costs no GPU
# time. The guard matters because spaces forces torch.cuda.is_available() to
# True under ZeroGPU's CUDA emulation, while plain CPU hardware reports False —
# loading ~71GB there would OOM the small container, so CPU runs open the
# canvas and ask for the hardware flip instead.
print(f"[load] {MODEL_ID} …", flush=True)
if torch.cuda.is_available():
    from transformers import AutoModelForImageTextToText, AutoProcessor

    processor = AutoProcessor.from_pretrained(MODEL_ID)
    model = AutoModelForImageTextToText.from_pretrained(
        MODEL_ID, torch_dtype="auto"
    ).to("cuda")
else:
    processor = None
    model = None
print("[load] model ready", flush=True)


def to_pil(value):
    """Normalize whatever the canvas image port sends to a PIL image (or None).

    Canvas file values carry only a url (FileValue has no path): same-server
    /gradio_api/file= urls resolve to local paths, http(s) urls are fetched.
    """
    if value is None:
        return None
    if isinstance(value, Image.Image):
        return value
    if isinstance(value, str):
        src = value
    elif isinstance(value, dict):
        src = value.get("path") or value.get("url")
    else:
        return None
    if not isinstance(src, str) or not src:
        return None
    if src.startswith("/gradio_api/file="):
        src = urllib.parse.unquote(src.removeprefix("/gradio_api/file="))
    if src.startswith(("http://", "https://")):
        import httpx

        resp = httpx.get(src, timeout=30, follow_redirects=True)
        resp.raise_for_status()
        return Image.open(__import__("io").BytesIO(resp.content))
    return Image.open(src) if os.path.exists(src) else None


def _chat(messages, max_new_tokens=MAX_NEW_TOKENS, temperature=0.0):
    """One MiniCPM-V chat round, following the model card's recipe."""
    if model is None:
        raise gr.Error(
            "MiniCPM-V 4.7 needs a GPU — switch this Space's hardware to "
            "ZeroGPU (Settings → Hardware) and restart it."
        )
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
        downsample_mode=DOWNSAMPLE_MODE,
        max_slice_nums=36,
        enable_thinking=False,
    ).to(model.device)
    gen_kwargs = dict(downsample_mode=DOWNSAMPLE_MODE, max_new_tokens=max_new_tokens)
    if temperature and float(temperature) > 0:
        gen_kwargs.update(do_sample=True, temperature=float(temperature))
    generated = model.generate(**inputs, **gen_kwargs)
    trimmed = [
        out_ids[len(in_ids):]
        for in_ids, out_ids in zip(inputs.input_ids, generated)
    ]
    return processor.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]


def compose_extraction_prompt(caption: str) -> str:
    """CPU-only node between the two GPU calls: turn a caption into a JSON
    extraction instruction for the second MiniCPM-V call."""
    caption = (caption or "").strip()
    if not caption:
        raise gr.Error("Connect a caption to this node.")
    return (
        "Based on the following image description, extract the key objects, "
        "actions, and scene attributes as a JSON object with the keys "
        '"objects", "actions", and "setting". Return only the JSON.\n\n'
        f"Description: {caption}"
    )


# duration is per node, not per pipeline: each call books and releases its own
# GPU slot. The first call on a fresh GPU slice also transfers ~71GB of packed
# weights, so the image nodes book 120s requested; warm calls finish far
# sooner and ZeroGPU bills actual usage.
#
# size="xlarge" is load-bearing, not an optimization: the default ZeroGPU
# slice is ~48GB and the 70.4GB model OOMs while spaces transfers the packed
# weights onto it. xlarge requests the whole GPU (141GB H200 / 96GB RTX PRO
# 6000). It also skips the hardware duration multiplier, so 120s books exactly
# 120s — but if an xlarge call is rejected for quota, ZeroGPU reports the
# figure doubled.
@spaces.GPU(duration=120, size="xlarge")
def ask(image, prompt, temperature=0.7):
    """Answer a free-form question about an image."""
    img = to_pil(image)
    if img is None:
        raise gr.Error("Connect an image to this node.")
    if not (prompt or "").strip():
        raise gr.Error("Connect a prompt to this node.")
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": img},
                {"type": "text", "text": prompt.strip()},
            ],
        }
    ]
    return _chat(messages, temperature=temperature)


@spaces.GPU(duration=120, size="xlarge")
def caption(image, prompt):
    """Caption an image with the given instruction."""
    img = to_pil(image)
    if img is None:
        raise gr.Error("Connect an image to this node.")
    if not (prompt or "").strip():
        raise gr.Error("Connect a prompt to this node.")
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": img},
                {"type": "text", "text": prompt.strip()},
            ],
        }
    ]
    return _chat(messages)


@spaces.GPU(duration=60, size="xlarge")
def extract_json(prompt):
    """Text-only MiniCPM-V call: the LLM backbone answers without an image."""
    if not (prompt or "").strip():
        raise gr.Error("Connect a prompt to this node.")
    messages = [
        {"role": "user", "content": [{"type": "text", "text": prompt.strip()}]}
    ]
    return _chat(messages)


demo = gr.Workflow(
    graph="workflow.json",
    bind={
        "ask": ask,
        "caption": caption,
        "compose_extraction_prompt": compose_extraction_prompt,
        "extract_json": extract_json,
    },
)

if __name__ == "__main__":
    demo.launch(css=_BADGE_CSS, js=_BADGE_JS)