Spaces:
Running on Zero
Running on Zero
| """VideoJev-4B — temporal grounding + multiple-choice video QA in one forward pass. | |
| Demo for https://huggingface.co/ZhuYuhan/VideoJev-4B (Apache-2.0), the video System One | |
| model from "VideoJev: System One Model for Video Understanding and Interaction" | |
| (arXiv:2610.18615). VideoJev scores time intervals and answer options directly against | |
| the video — no autoregressive decoding — and answers a whole group of questions in a | |
| single forward pass. | |
| Inference follows the authors' reference path exactly (the `modeling_videojev.py` / | |
| `processing_videojev.py` shipped inside the checkpoint, byte-identical to | |
| MCG-NJU/VideoJev `videojev_hf/`): `prepare_video` → `encode_video` → `predict_features`, | |
| which is precisely what `model.predict_video` does internally — split here only so the | |
| demo can also show the per-timestamp probability timeline. The model runs with | |
| `attn_implementation="sdpa"`: transformers 5.x falls back to a pure-torch chunked | |
| `gated_delta_rule` for the hybrid linear-attention layers when flash-linear-attention | |
| kernels are absent, which keeps the Space dependency-light. | |
| """ | |
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") | |
| os.environ.setdefault("GRADIO_SSR_MODE", "false") | |
| # Match the authors' env (scripts/env.sh) so native recurrence numerics stay put. | |
| os.environ.setdefault("FLA_INTRACARD_CP", "0") | |
| import spaces # noqa: E402 -- MUST come before torch / transformers / any CUDA-touching import | |
| import html # noqa: E402 | |
| import string # noqa: E402 | |
| import sys # noqa: E402 | |
| import time # noqa: E402 | |
| import gradio as gr # noqa: E402 | |
| import torch # noqa: E402 | |
| from transformers import AutoModel, AutoProcessor # noqa: E402 | |
| MODEL_ID = "ZhuYuhan/VideoJev-4B" | |
| LETTERS = string.ascii_uppercase | |
| print("[videojev] loading processor + model (bf16, sdpa)…", flush=True) | |
| processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True) | |
| model = AutoModel.from_pretrained( | |
| MODEL_ID, | |
| trust_remote_code=True, | |
| dtype=torch.bfloat16, | |
| attn_implementation="sdpa", | |
| ).to("cuda") | |
| model.prepare_for_inference() | |
| # Pull the checkpoint's own preprocessing module — the exact code revision the weights | |
| # ship with — from the sys.modules entry trust_remote_code already created for it. | |
| _processing = sys.modules[model.__class__.__module__] | |
| prepare_video = _processing.prepare_video | |
| print(f"[videojev] {MODEL_ID} ready — grounding threshold {model.config.decoder['threshold']}", flush=True) | |
| # --------------------------------------------------------------------------- | |
| # Input helpers | |
| # --------------------------------------------------------------------------- | |
| def _normalize_rows(rows) -> list: | |
| """Accept whatever gr.Dataframe hands over (pandas / numpy / nested list).""" | |
| if rows is None: | |
| return [] | |
| if hasattr(rows, "values"): # pandas DataFrame | |
| rows = rows.astype(object).where(rows.notna(), "").values.tolist() | |
| elif hasattr(rows, "tolist"): # numpy array | |
| rows = rows.tolist() | |
| return [[("" if c is None else c) for c in list(r) + [""] * (3 - len(list(r)))] for r in rows] | |
| def _parse_options(text) -> list: | |
| """Split the options cell on ';' or newlines (commas are kept inside options).""" | |
| if text is None or not str(text).strip(): | |
| return [] | |
| return [o.strip() for o in str(text).replace(";", "\n").split("\n") if o.strip()] | |
| def _label_options(options: list) -> list: | |
| """Keep author-format options ('A. …') untouched; add letter labels otherwise.""" | |
| out = [] | |
| for i, option in enumerate(options): | |
| if i >= len(LETTERS): | |
| raise gr.Error(f"At most {len(LETTERS)} options are supported.") | |
| letter = LETTERS[i] | |
| out.append(option if option.startswith(f"{letter}.") else f"{letter}. {option}") | |
| return out | |
| def _build_questions(rows) -> list: | |
| """Turn the questions table into the authors' question-dict format.""" | |
| questions = [] | |
| for i, row in enumerate(_normalize_rows(rows)): | |
| task, query, options_cell = row[0], row[1], row[2] | |
| task = str(task).strip().lower() | |
| query = str(query or "").strip() | |
| if not query and not task: | |
| continue | |
| if task not in ("grounding", "choice"): | |
| raise gr.Error( | |
| f"Row {i + 1}: task must be 'grounding' or 'choice' (got '{row[0]}')." | |
| ) | |
| if not query: | |
| raise gr.Error(f"Row {i + 1}: the query text is required.") | |
| if "<|" in query or "<|" in str(options_cell or ""): | |
| raise gr.Error("Special control tokens (<|…|>) are not allowed in questions or options.") | |
| options = _label_options(_parse_options(options_cell)) | |
| if task == "choice" and len(options) < 2: | |
| raise gr.Error(f"Row {i + 1}: choice questions need at least two options, separated by ';' or new lines.") | |
| question = {"id": f"q{i + 1}", "task": task, "query": query} | |
| if task == "choice": | |
| question["options"] = options | |
| questions.append(question) | |
| if not questions: | |
| raise gr.Error("Add at least one question (a query text is required).") | |
| return questions | |
| # --------------------------------------------------------------------------- | |
| # Rendering | |
| # --------------------------------------------------------------------------- | |
| _BAR = ( | |
| '<div style="display:flex;align-items:center;gap:8px;margin:2px 0">' | |
| '<span style="flex:0 0 46px;text-align:right;font-size:12px;opacity:.65">{t:.1f}s</span>' | |
| '<div style="flex:1;height:22px;background:var(--border-color-primary);border-radius:3px;position:relative;overflow:hidden">' | |
| '<div style="width:{fill:.1f}%;position:absolute;inset:0;' | |
| 'background:var(--color-accent);opacity:.25"></div>' | |
| '<div style="position:absolute;left:{ls:.2f}%;width:{ws:.2f}%;top:3px;bottom:3px;' | |
| 'background:var(--color-accent);border-radius:2px"></div>' | |
| "</div>" | |
| '<span style="flex:0 0 44px;font-size:12px;opacity:.75">{p:.2f}</span></div>' | |
| ) | |
| def _grounding_html(item: dict, record: dict) -> str: | |
| start, end = item["interval"] | |
| duration = max(record["sampled_duration"], 1e-6) | |
| left = 100.0 * start / duration | |
| width = 100.0 * max(end - start, 0.0) / duration | |
| bars = "".join( | |
| _BAR.format(t=t, fill=100.0 * p, ls=left, ws=width, p=p) | |
| for t, p in zip(record["score_times"], item["probabilities"]) | |
| ) | |
| return ( | |
| '<div style="padding:10px 0;border-bottom:1px solid var(--border-color-primary)">' | |
| '<div style="font-size:12.5px;opacity:.7;margin-bottom:4px"><code>grounding</code> · ' | |
| f"{html.escape(item['_query'])}</div>" | |
| '<div style="font-size:20px;font-weight:600;margin:2px 0">' | |
| f"⏱ {start:.1f}s – {end:.1f}s" | |
| '<span style="font-weight:400;font-size:12.5px;opacity:.65;margin-left:8px">' | |
| f"highlighted over the {duration:.1f}s video</span></div>" | |
| f'<div style="margin-top:6px">{bars}</div></div>' | |
| ) | |
| def _choice_html(item: dict) -> str: | |
| winner = item["answer_index"] | |
| rows = "".join( | |
| '<div style="display:flex;align-items:center;gap:8px;margin:2px 0">' | |
| f'<span style="flex:1;font-size:12.5px;opacity:{1.0 if i == winner else 0.5};' | |
| f'font-weight:{600 if i == winner else 400}">{html.escape(option)}</span>' | |
| '<div style="flex:0 0 170px;background:var(--border-color-primary);height:12px;border-radius:3px">' | |
| f'<div style="width:{100.0 * p:.1f}%;height:12px;border-radius:3px;background:var(--color-accent);' | |
| f'opacity:{1.0 if i == winner else 0.4}"></div></div>' | |
| f'<span style="flex:0 0 50px;font-size:12px;text-align:right;' | |
| f'opacity:{1.0 if i == winner else 0.5}">{p * 100:.1f}%</span></div>' | |
| for i, (option, p) in enumerate(zip(item["_options"], item["probabilities"])) | |
| ) | |
| return ( | |
| '<div style="padding:10px 0;border-bottom:1px solid var(--border-color-primary)">' | |
| '<div style="font-size:12.5px;opacity:.7;margin-bottom:4px"><code>choice</code> · ' | |
| f"{html.escape(item['_query'])}</div>" | |
| '<div style="font-size:20px;font-weight:600;margin:2px 0">' | |
| f"✔ {html.escape(item['answer'])}" | |
| '<span style="font-weight:400;font-size:12.5px;opacity:.65;margin-left:8px">' | |
| f"softmax P = {item['probabilities'][winner]:.3f}</span></div>" | |
| f'<div style="margin-top:6px">{rows}</div></div>' | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # GPU handler | |
| # --------------------------------------------------------------------------- | |
| def _estimate_duration(video, rows, *args, **kwargs): | |
| """Duration sizing from measured runs (2.7s GPU for a 10s video at 2.2K tokens, | |
| 7.9s for a 27s video at 21K tokens — both including first-call weight streaming). | |
| Context caps tokens regardless of video length, so a flat term plus a small | |
| per-second allowance covers decode + encode + prefill with healthy margin.""" | |
| try: | |
| import av | |
| with av.open(str(video)) as container: | |
| stream = container.streams.video[0] | |
| if stream.duration is not None: | |
| seconds = float(stream.duration * stream.time_base) | |
| else: | |
| seconds = container.duration / 1e6 | |
| except Exception: | |
| seconds = 30.0 | |
| return min(120, 25 + int(seconds * 0.5)) | |
| def analyze( | |
| video, | |
| rows, | |
| fps=2.0, | |
| max_frames=512, | |
| context=32768, | |
| threshold=0.25, | |
| progress=gr.Progress(track_tqdm=True), | |
| ): | |
| """Answer a group of temporal-grounding and multiple-choice questions about one video in a single forward pass. | |
| Args: | |
| video: input video file (any container PyAV can decode). | |
| rows: question table — one row per question with columns | |
| [task ('grounding' | 'choice'), query, options]. Options apply to 'choice' | |
| questions only and are separated by ';' or new lines; letters are added | |
| automatically unless already present ("A. …"). | |
| fps: frame sampling rate (authors use 2 for QA, 4 for grounding). | |
| max_frames: cap on sampled frames across the whole video. | |
| context: token budget for video + queries; video resolution adapts to fit. | |
| threshold: sigmoid threshold used to cut the grounding interval. | |
| progress: Gradio progress bar. | |
| Returns: | |
| (results_html, stats_markdown, raw_json): per-question rendering, run stats, | |
| and the raw model output (interval / answer / probabilities per question). | |
| """ | |
| if not video: | |
| raise gr.Error("Upload a video or pick one of the examples first.") | |
| video = str(video) | |
| questions = _build_questions(rows) | |
| fps, max_frames, context, threshold = float(fps), int(max_frames), int(context), float(threshold) | |
| t0 = time.perf_counter() | |
| try: | |
| with torch.inference_mode(): | |
| # The authors' exact predict_video pipeline, split open so the per-timestamp | |
| # probability timeline can be shown in the UI. | |
| pixels, record = prepare_video(processor, video, fps=fps, max_frames=max_frames, context=context) | |
| features = model.encode_video(pixels, torch.tensor([record["grid_thw"]])) | |
| results = model.predict_features( | |
| processor, features, record, questions, context=context, threshold=threshold | |
| ) | |
| except gr.Error: | |
| raise | |
| except Exception as e: # surface a readable message instead of a stack trace | |
| raise gr.Error(f"VideoJev inference failed: {type(e).__name__}: {e}") | |
| elapsed = time.perf_counter() - t0 | |
| parts = [] | |
| for item, question in zip(results, questions): | |
| item = dict(item) | |
| item["_query"] = question["query"] | |
| if question["task"] == "choice": | |
| item["_options"] = question["options"] | |
| parts.append(_choice_html(item)) | |
| else: | |
| parts.append(_grounding_html(item, record)) | |
| n_ground = sum(1 for q in questions if q["task"] == "grounding") | |
| n_choice = len(questions) - n_ground | |
| stats = ( | |
| f"**{len(questions)} question{'s' if len(questions) > 1 else ''}** " | |
| f"({n_ground} grounding · {n_choice} choice) answered in " | |
| f"**{elapsed:.1f}s** on GPU · {len(record['score_times'])} video timestamps · " | |
| f"{record['tokens']:,} video tokens" | |
| ) | |
| payload = [ | |
| {k: v for k, v in dict(item, query=q["query"]).items()} | |
| for item, q in zip(results, questions) | |
| ] | |
| return "".join(parts), stats, payload | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| DEFAULT_ROWS = [ | |
| ["grounding", "The woman in green eats pasta.", ""], | |
| ["choice", "What is the woman in green eating?", "A. Pasta.; B. Rice.; C. Soup."], | |
| ] | |
| with gr.Blocks() as demo: | |
| gr.Markdown( | |
| """ | |
| # 🎬 VideoJev-4B — video decisions without generation | |
| **Temporal grounding + multiple-choice video QA in one forward pass.** VideoJev is a video *System One* | |
| model: instead of decoding text, it scores candidate time intervals and answer options directly against | |
| the video — orders of magnitude faster than generative answering on the same tasks. All questions below | |
| are answered together, sharing one video encoding. | |
| [Paper](https://arxiv.org/abs/2610.18615) · [Code](https://github.com/MCG-NJU/VideoJev) · | |
| [Model](https://huggingface.co/ZhuYuhan/VideoJev-4B) · [Homepage](https://mcg-nju.github.io/VideoJev) | |
| """ | |
| ) | |
| with gr.Column(elem_id="col-container"): | |
| with gr.Row(): | |
| video_in = gr.Video(label="Video", sources=["upload"]) | |
| with gr.Column(scale=7): | |
| questions = gr.Dataframe( | |
| headers=["task", "query", "options (choice only: ';' or newline separated)"], | |
| datatype=["str", "str", "str"], | |
| column_count=(3, "fixed"), | |
| row_count=(1, "dynamic"), | |
| value=DEFAULT_ROWS, | |
| wrap=True, | |
| label="Questions — all answered in ONE shared forward pass", | |
| interactive=True, | |
| ) | |
| run = gr.Button("Answer questions", variant="primary") | |
| with gr.Accordion("Advanced settings", open=False): | |
| with gr.Row(): | |
| fps = gr.Radio([2.0, 4.0], value=2.0, label="Sampling fps (2 = QA, 4 = grounding per authors)") | |
| max_frames = gr.Slider(16, 2048, value=512, step=16, info="Caps sampling; covers 4+ min at 2 fps", | |
| label="Max frames over the whole video") | |
| context = gr.Slider(8192, 262144, value=32768, step=2048, | |
| info="Video resolution adapts to fit this budget", | |
| label="Context tokens") | |
| threshold = gr.Slider(0.05, 0.95, value=0.25, step=0.05, | |
| info="Sigmoid cut for grounding intervals", | |
| label="Grounding threshold") | |
| stats = gr.Markdown("") | |
| results = gr.HTML(label="Results") | |
| raw_json = gr.JSON(label="Raw model output") | |
| with gr.Accordion("How to read the output", open=False): | |
| gr.Markdown( | |
| """ | |
| - **grounding** → an event span `[start, end]` in seconds. The bar chart shows the model's sigmoid | |
| probability at each video timestamp (a translucent bar = below threshold, solid = inside the | |
| returned span). | |
| - **choice** → the selected option, with softmax probabilities over all options. | |
| - The raw JSON is exactly what the checkpoint returns from `predict_features`: `interval`, | |
| `answer`, `answer_index` and `probabilities`. | |
| """ | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| # The checkpoint's README quick-start case: grounding + QA at 2 fps. | |
| ["eating_spaghetti.mp4", DEFAULT_ROWS, 2.0], | |
| # Charades case from the paper's project page; grounding runs at the | |
| # authors' recommended 4 fps (fp numbers match their published 4B output). | |
| ["case-01.mp4", [ | |
| ["grounding", "A man is sitting on a chair.", ""], | |
| ["grounding", "The man gets up from the chair.", ""], | |
| ["grounding", "The man runs toward the door.", ""], | |
| ], 4.0], | |
| ["case-01.mp4", [ | |
| ["choice", "What does the man do first?", | |
| "A. He sits on a chair; B. He runs to the door; C. He eats pasta"], | |
| ], 2.0], | |
| ], | |
| inputs=[video_in, questions, fps], | |
| outputs=[results, stats, raw_json], | |
| fn=analyze, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| run.click( | |
| analyze, | |
| inputs=[video_in, questions, fps, max_frames, context, threshold], | |
| outputs=[results, stats, raw_json], | |
| api_name="analyze", | |
| ) | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) |