"""Jev-Omni on ZeroGPU: text, image, audio or video in; a probability per option out.

The model is placed on `cuda` at module level on purpose. ZeroGPU emulates CUDA outside `@spaces.GPU`
and swaps in a real GPU inside it, and its docs say startup placement is the efficient path.
Weights are the single bf16 checkpoint at the root of the model repo, so one ~24 GB download
covers text, vision and audio.
"""
import os, subprocess, tempfile, time
import gradio as gr
import numpy as np
import spaces
import torch
torch.backends.cuda.enable_cudnn_sdp(False)   # cuDNN attention hit a version mismatch on one CUDA image; the other SDPA backends are fine
from huggingface_hub import hf_hub_download, snapshot_download
from PIL import Image
from transformers import AutoConfig, AutoProcessor

REPO = "akhilaaa3/Jev-Omni"
DEVICE = "cuda"


class Head256(torch.nn.Module):
    def __init__(self, hidden):
        super().__init__()
        self.register_buffer("mu", torch.zeros(1, hidden))
        self.register_buffer("sd", torch.ones(1, hidden))
        self.linear = torch.nn.Linear(hidden, 256, dtype=torch.float32)

    def forward(self, features, counts):
        z = self.linear((features.float() - self.mu) / self.sd)
        return z.masked_fill(torch.arange(256, device=z.device)[None] >= counts[:, None], -1e30)


def _prompt(state, question, options):
    choices = "\n".join(f"{i + 1}. {v}" for i, v in enumerate(options))
    return (f"{state}\n\n---\n\nQUESTION: {question}\n\nOPTIONS:\n{choices}\n\n"
            f"Reply with only the number of the correct option (1-{len(options)}).\n"
            "Output a single number and nothing else.")


def _video_frames(path, count=16):
    import cv2
    cap = cv2.VideoCapture(str(path)); total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    wanted = sorted({int(round((total - 1) * (k + .5) / count)) for k in range(count)})
    frames, cur = [], 0
    for idx in wanted:
        while cur < idx and cap.grab(): cur += 1
        ok, fr = cap.read(); cur += 1
        if not ok: break
        frames.append(Image.fromarray(cv2.cvtColor(fr, cv2.COLOR_BGR2RGB)))
    cap.release()
    if not frames: raise gr.Error("Could not decode that video.")
    return frames


# ---- load once, at startup, onto (emulated) cuda
t0 = time.time()
local = snapshot_download(REPO, allow_patterns=[
    "config.json", "generation_config.json", "model*.safetensors*", "processor_config.json",
    "tokenizer.json", "tokenizer_config.json", "chat_template.jinja", "head.pt", "decision_config.json"])
import transformers, json
cfg = AutoConfig.from_pretrained(local)
MODEL = getattr(transformers, cfg.architectures[0]).from_pretrained(
    local, dtype=torch.bfloat16, device_map=DEVICE).eval()
PROC = AutoProcessor.from_pretrained(local)
dc = json.load(open(os.path.join(local, "decision_config.json")))
HEAD = Head256(dc["hidden_size"]).to(DEVICE).eval()
HEAD.load_state_dict(torch.load(os.path.join(local, "head.pt"), map_location=DEVICE, weights_only=True))
CAP = {}
for path in ("model.language_model", "language_model.model", "model.text_model", "model"):
    node = MODEL
    for part in path.split("."):
        node = getattr(node, part, None)
        if node is None: break
    if node is not None and hasattr(node, "layers"):
        node.register_forward_hook(lambda _m, _a, out: CAP.__setitem__(
            "h", (out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0])[:, -1].float()))
        break
else:
    raise RuntimeError("text backbone not found")
print(f"loaded in {time.time() - t0:.0f}s", flush=True)


def _duration(state, question, options_text, media, modality):
    return 90 if modality == "video" else 45


@spaces.GPU(duration=_duration)
@torch.inference_mode()
def decide(state, question, options_text, media, modality):
    options = [o.strip() for o in options_text.splitlines() if o.strip()]
    if not 2 <= len(options) <= 256: raise gr.Error("Give between 2 and 256 options, one per line.")
    if modality != "text" and not media: raise gr.Error(f"Upload a file for {modality} input.")
    content, tmp = [], None
    if modality == "image": content.append({"type": "image", "image": Image.open(media).convert("RGB")})
    elif modality == "video": content.extend({"type": "image", "image": f} for f in _video_frames(media))
    elif modality == "audio":
        tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False); tmp.close()
        subprocess.run(["ffmpeg", "-v", "error", "-y", "-i", str(media), "-t", "30", "-ac", "1", "-ar", "16000", tmp.name], check=True)
        content.append({"type": "audio", "audio": tmp.name})
    content.append({"type": "text", "text": _prompt(state, question, options)})
    try:
        inputs = PROC.apply_chat_template([{"role": "user", "content": content}], add_generation_prompt=True,
                                          tokenize=True, return_dict=True, return_tensors="pt", enable_thinking=False)
    finally:
        if tmp: os.unlink(tmp.name)
    inputs = {k: (v.to(DEVICE, dtype=torch.bfloat16) if torch.is_floating_point(v) else v.to(DEVICE)) for k, v in inputs.items()}
    CAP.clear(); t = time.time()
    with torch.autocast("cuda", dtype=torch.bfloat16):
        MODEL(**inputs, use_cache=False, logits_to_keep=1)
        probs = HEAD(CAP["h"], torch.tensor([len(options)], device=DEVICE))[0, :len(options)].softmax(-1)
    ms = (time.time() - t) * 1000
    p = probs.float().cpu().tolist()
    best = int(np.argmax(p))
    return {o: v for o, v in zip(options, p)}, f"**{options[best]}** · {p[best]*100:.1f}% · {ms:.0f} ms"


EXAMPLES = [
    ["The meeting starts at 10 AM. It is now 9 AM.", "Has the meeting started?", "Yes\nNo", None, "text"],
    ["Customer: I was charged twice.\nAgent: I've refunded $29 to your card.\nCustomer: Got it. All sorted, thanks!",
     "Was the issue actually resolved?", "Yes\nNo", None, "text"],
]

with gr.Blocks(title="Jev-Omni") as demo:
    gr.Markdown("# Jev-Omni\nA decision classifier for text, images, audio and video. "
                "Give it a situation and a question with options; it returns a probability for each option — "
                "no generated text, one forward pass.")
    with gr.Row():
        with gr.Column():
            modality = gr.Radio(["text", "image", "audio", "video"], value="text", label="Input type")
            media = gr.File(label="Image / audio / video file (not needed for text)", file_types=["image", "audio", "video"])
            state = gr.Textbox(label="Situation", lines=6, placeholder="What the model should know. For media, this can be short.")
            question = gr.Textbox(label="Question")
            options = gr.Textbox(label="Options, one per line", lines=4, value="Yes\nNo")
            go = gr.Button("Decide", variant="primary")
        with gr.Column():
            verdict = gr.Markdown()
            probs = gr.Label(label="Probabilities", num_top_classes=10)
    go.click(decide, [state, question, options, media, modality], [probs, verdict])
    gr.Examples(EXAMPLES, [state, question, options, media, modality])
    gr.Markdown("Best supported at ≤20 options; audio is capped at 30 s; video is sampled to 16 frames. "
                "Weights: [akhilaaa3/Jev-Omni](https://huggingface.co/akhilaaa3/Jev-Omni), Apache-2.0.")

demo.queue().launch()