Jev-Omni demo: text, image, audio and video decisions on ZeroGPU from the unified bf16 checkpoint

README.md

@@ -1,13 +1,21 @@
 ---
-title: Jev Omni
-emoji: 📈
-colorFrom: indigo
-colorTo: red
+title: Jev-Omni
+emoji: 🎯
+colorFrom: blue
+colorTo: gray
 sdk: gradio
-sdk_version: 6.28.0
-python_version: '3.12'
+sdk_version: 5.49.1
 app_file: app.py
 pinned: false
+license: apache-2.0
+short_description: Multimodal decision classifier — text, image, audio, video
+models:
+- akhilaaa3/Jev-Omni
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Jev-Omni
+
+Give it a situation and a question with options. It returns a probability for each option from one
+forward pass — no generated text. Built on Gemma 4 12B IT with a 30,000-question fine-tune.
+
+Model card and results: [akhilaaa3/Jev-Omni](https://huggingface.co/akhilaaa3/Jev-Omni).

app.py

@@ -0,0 +1,142 @@
+"""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 unified bf16 checkpoint under `unified/` in 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=["unified/*", "head.pt", "decision_config.json"])
+import transformers, json
+cfg = AutoConfig.from_pretrained(local, subfolder="unified")
+MODEL = getattr(transformers, cfg.architectures[0]).from_pretrained(
+    local, subfolder="unified", dtype=torch.bfloat16, device_map=DEVICE).eval()
+PROC = AutoProcessor.from_pretrained(local, subfolder="unified")
+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()

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -0,0 +1,11 @@
+torch>=2.10
+torchvision
+transformers==5.17.0
+accelerate
+huggingface_hub
+safetensors
+numpy
+Pillow
+opencv-python-headless
+soundfile
+librosa