"""TransVLM — detect any shot transition in a video. A faithful port of the reference whole-video pipeline (``infer_video.py::_process_video`` in https://github.com/heygen-com/TransVLM) onto ZeroGPU: fps -> 25 -> smart_resize (ffmpeg) -> whole-video NeuFlow v2 optical flow -> sliding windows (10 s / 9 s stride, as time ranges) -> 6-channel (RGB (+) flow) Qwen3-VL generation -> JSON parse -> merge Nothing about the model path is simplified: the same prompt, the same ``smart_resize`` pixel budget, the same whole-video ``flow_to_image`` normalisation, the same half-frame window compensation and index rebasing. The only deviation from the CLI is that the input is trimmed to a bounded number of seconds so a single request fits inside a ZeroGPU slot. """ from __future__ import annotations import json import logging import os import subprocess import sys import tempfile import time from fractions import Fraction from pathlib import Path # `import spaces` MUST precede any torch / CUDA-touching import. import spaces # noqa: F401 (isort: skip) import gradio as gr import torch sys.path.insert(0, str(Path(__file__).resolve().parent)) from transvlm._paths import DEFAULT_PROMPT_FILE from transvlm.data.flow_computer import OnlineFlowComputer from transvlm.data.flow_config import FlowConfig from transvlm.data.resize_video_helper import ( probe_video_hw, resize_video_ffmpeg, smart_resize_for_qwen3vl, ) from transvlm.data.sliding_window import get_sliding_window_segments from transvlm.data.video_backend import load_process_vision_info, selected_video_backend from transvlm.inference.clip_engine import DEFAULT_FPS, DEFAULT_MAX_PIXELS from transvlm.inference.hf_clip_engine import ( ATTN_IMPLEMENTATION, DTYPE, HFClipInferenceEngine, ) from transvlm.inference.merger import merge_shots from transvlm.inference.parser import parse_model_output from transvlm.models.processor_patch import apply_timestamp_format logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s", force=True, ) logger = logging.getLogger("transvlm-space") MODEL_ID = "HeyGenAI/TransVLM-Qwen3-VL-4B-Instruct" PAPER_URL = "https://huggingface.co/papers/2604.27975" CODE_URL = "https://github.com/heygen-com/TransVLM" # Paper protocol. Do not change: the prompt, the flow signal and every predicted # timestamp are tied to these. TARGET_FPS = Fraction(25, 1) NFRAMES_FOR_RESIZE = 250 IMAGE_PATCH_SIZE = 16 MAX_PIXELS_OVERRIDE = DEFAULT_MAX_PIXELS # 524288 FLOW_CODEC = "libx264" FLOW_MINI_BATCH = 32 DEFAULT_MAX_SECONDS = 12.0 DEFAULT_WINDOW_SIZE = 10.0 DEFAULT_STRIDE = 9.0 DEFAULT_MERGE_EPS = 0.02 DEFAULT_MAX_NEW_TOKENS = 1024 # A shot "cut" is instantaneous; anything longer is a gradual transition # (dissolve / fade / wipe). 2 frames at 25 fps. CUT_THRESHOLD_SEC = 0.08 # --------------------------------------------------------------------------- # ffprobe helpers — ported verbatim from infer_video.py # --------------------------------------------------------------------------- def _ffprobe(path: Path, entries: str) -> list[str]: out = subprocess.run( [ "ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", entries, "-of", "default=nw=1:nk=1", str(path), ], capture_output=True, text=True, check=True, ) return out.stdout.split() def _frame_rates(path: Path) -> tuple[Fraction, Fraction]: avg, real = _ffprobe(path, "stream=avg_frame_rate,r_frame_rate")[:2] return Fraction(avg), Fraction(real) def _duration_sec(path: Path) -> float: return float(_ffprobe(path, "format=duration")[0]) def _needs_resample(path: Path, target: Fraction) -> tuple[bool, str]: """True when the video is not exactly ``target`` CFR (rational-exact).""" avg, real = _frame_rates(path) if avg != real: return True, f"VFR (avg {avg} != r {real})" if avg != target: return True, f"{avg} fps" return False, f"{target} CFR" def _trim_and_resample(src: Path, work: Path, max_seconds: float, target: Fraction) -> tuple[Path, str | None]: """Single ffmpeg pass: cap the duration and land the stream on ``target`` CFR. The reference CLI runs the whole video; a Space request has to fit in one ZeroGPU slot, so the input is capped first. Trimming before the fps filter is equivalent to trimming after it — ``fps=`` is a per-timestamp decision — so the frames the model sees for ``[0, max_seconds]`` are the ones the CLI would have produced. """ src_duration = _duration_sec(src) needs_fps, why = _needs_resample(src, target) needs_trim = src_duration > max_seconds + 1e-3 if not needs_fps and not needs_trim: logger.info("[prepare] %s already %s and <= %.1f s, skipping re-encode", src.name, why, max_seconds) return src, None dst = work / f"{src.stem}_prep.mp4" cmd = ["ffmpeg", "-y", "-v", "error"] if needs_trim: cmd += ["-t", f"{max_seconds:.3f}"] cmd += [ "-i", str(src), "-filter:v", f"fps={target}", "-c:v", "libx264", "-crf", "17", "-preset", "fast", "-an", str(dst), ] logger.info( "[prepare] %s: %s%s -> %s fps", src.name, why, f", trimming {src_duration:.1f}s -> {max_seconds:.1f}s" if needs_trim else "", target, ) result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: raise RuntimeError(f"ffmpeg prepare failed on {src}\n{result.stderr}") note = why if needs_fps else None return dst, note def _maybe_resize(src: Path, work: Path) -> tuple[Path, tuple[int, int], tuple[int, int]]: src_h, src_w = probe_video_hw(src) target_h, target_w = smart_resize_for_qwen3vl( src_h, src_w, nframes=NFRAMES_FOR_RESIZE, image_patch_size=IMAGE_PATCH_SIZE, max_pixels_override=MAX_PIXELS_OVERRIDE, ) if (target_h, target_w) == (src_h, src_w): logger.info("[resize] %s already %dx%d, skipping", src.name, src_h, src_w) return src, (src_h, src_w), (target_h, target_w) dst = work / f"{src.stem}_resized.mp4" logger.info("[resize] %s %dx%d -> %dx%d", src.name, src_h, src_w, target_h, target_w) resize_video_ffmpeg(src, dst, target_h=target_h, target_w=target_w, overwrite=True) return dst, (src_h, src_w), (target_h, target_w) def _faststart(src: Path) -> Path: """Remux (no re-encode) so a browser can start playing before the file ends.""" dst = src.with_name(f"{src.stem}_web.mp4") result = subprocess.run( ["ffmpeg", "-y", "-v", "error", "-i", str(src), "-c", "copy", "-movflags", "+faststart", str(dst)], capture_output=True, text=True, ) if result.returncode != 0 or not dst.exists(): logger.warning("faststart remux failed, serving the raw flow mp4: %s", result.stderr[:400]) return src return dst # --------------------------------------------------------------------------- # ZeroGPU adapters # --------------------------------------------------------------------------- class ZeroGPUEngine(HFClipInferenceEngine): """``HFClipInferenceEngine`` that loads from a Hub id with an eager ``.to("cuda")``. The upstream ``load()`` uses ``device_map={'': 'cuda:0'}``, which initialises CUDA in the main process and so bypasses ZeroGPU's lazy-allocation hijack, and it insists on a local checkpoint directory. Everything else — dtype, attention implementation, the timestamp-format patch, the 6-channel assertion — is kept. """ def load(self) -> None: # noqa: D102 — see class docstring if self._model is not None: return from transformers import AutoModelForImageTextToText, AutoProcessor apply_timestamp_format(self.timestamp_format, include_vllm=False) started = time.perf_counter() self._processor = AutoProcessor.from_pretrained(MODEL_ID) model = AutoModelForImageTextToText.from_pretrained( MODEL_ID, dtype=getattr(torch, DTYPE), attn_implementation=ATTN_IMPLEMENTATION, ) self._model = model.to("cuda").eval() self.model_load_sec = time.perf_counter() - started in_channels = getattr(self._model.config.vision_config, "in_channels", None) if in_channels != 6: raise RuntimeError( f"checkpoint reports vision_config.in_channels={in_channels}, expected 6; " "this is not the 6-channel TransVLM v1 checkpoint" ) logger.info("TransVLM ready in %.1f s (in_channels=%d)", self.model_load_sec, in_channels) class ZeroGPUFlowComputer(OnlineFlowComputer): """``OnlineFlowComputer`` whose weights are prepared on CPU, then moved to CUDA. ``OnlineFlowComputer.__init__`` moves NeuFlow to the device and *then* runs ``_fuse_bn``, which does real arithmetic (``torch.mm`` / ``torch.diag``) on the weights. Under ZeroGPU that arithmetic would run at module scope against tensors that only present as CUDA, so the fold happens on CPU first and the fused, half-precision model is moved afterwards — which is exactly the eager module-scope ``.to("cuda")`` the hijack expects. """ def __init__(self, cfg: FlowConfig, **kwargs) -> None: super().__init__(cfg, device="cpu", **kwargs) self._model = self._model.to("cuda") self._device_str = "cuda" self._device = torch.device("cuda") self._viz_on_gpu = True self._bhwd = None # --------------------------------------------------------------------------- # Module-scope load # --------------------------------------------------------------------------- logger.info("torch %s | video decoder settled below", torch.__version__) try: meminfo = dict( line.split(":", 1) for line in Path("/proc/meminfo").read_text().splitlines() if ":" in line ) logger.info("host MemTotal=%s MemAvailable=%s", meminfo.get("MemTotal", "?").strip(), meminfo.get("MemAvailable", "?").strip()) except Exception: # noqa: BLE001 pass # Settle the decoder choice (and log it) before anything imports the vendored # vision-process module, which freezes the backend at import time. load_process_vision_info() logger.info("video decoder: %s", selected_video_backend()) ENGINE = ZeroGPUEngine( ckpt_dir=MODEL_ID, prompt_file=DEFAULT_PROMPT_FILE, fps=DEFAULT_FPS, max_new_tokens=DEFAULT_MAX_NEW_TOKENS, device="cuda", ) ENGINE.load() FLOW = ZeroGPUFlowComputer( FlowConfig(mode="online", flow_codec=FLOW_CODEC), mini_batch_size=FLOW_MINI_BATCH, max_pixels_override=MAX_PIXELS_OVERRIDE, viz_on_gpu=True, ) # --------------------------------------------------------------------------- # Presentation helpers # --------------------------------------------------------------------------- def _kind(seg_start: float, seg_end: float) -> str: return "cut" if (seg_end - seg_start) <= CUT_THRESHOLD_SEC else "gradual" def _timeline_html(segments: list[dict], duration: float) -> str: if duration <= 0: duration = 1.0 marks = [] for i, seg in enumerate(segments, 1): left = max(0.0, min(1.0, seg["start_time"] / duration)) * 100 width = max(0.0, min(1.0, (seg["end_time"] - seg["start_time"]) / duration)) * 100 width = max(width, 0.45) # a hard cut is zero-width; keep it visible colour = "#ef4444" if seg["kind"] == "cut" else "#f59e0b" marks.append( f'
' ) ticks = [] n_ticks = 6 for t in range(n_ticks + 1): pos = t / n_ticks * 100 ticks.append( f'{duration * t / n_ticks:.1f}s' ) return ( '
' '
' + "".join(marks) + "
" '
' + "".join(ticks) + "
" '
' 'hard cut' 'gradual transition
' "
" ) # --------------------------------------------------------------------------- # Inference # --------------------------------------------------------------------------- def _estimate_duration(video=None, max_seconds: float = DEFAULT_MAX_SECONDS, *args, **kwargs) -> int: """Size the ZeroGPU slot to the request. Measured on this Space (zero-a10g, the three bundled examples): 11.0 s of video / 2 windows -> 14.1 s wall 13.2 s of video / 2 windows -> 16.0 s wall 14.0 s of video / 2 windows -> 18.0 s wall (worst observed) i.e. roughly ``3 + 1.05 x seconds``. The formula below keeps ~2.5x headroom for cold weight-streaming and for uploads that are slower to decode than the H.264 examples (4K HEVC, variable frame rate), and no more -- an inflated duration burns every visitor's quota and lowers queue priority. """ try: secs = float(max_seconds) except (TypeError, ValueError): secs = DEFAULT_MAX_SECONDS secs = max(1.0, min(60.0, secs)) return int(min(90, max(25, 20 + 1.8 * secs))) @spaces.GPU(duration=_estimate_duration) def detect( video: str | None, max_seconds: float = DEFAULT_MAX_SECONDS, window_size: float = DEFAULT_WINDOW_SIZE, stride: float = DEFAULT_STRIDE, merge_eps: float = DEFAULT_MERGE_EPS, max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS, progress=gr.Progress(), ) -> tuple[str, str, str, list[list], dict]: """Detect every shot transition in a video and return their timestamps. Runs the TransVLM pipeline: resample to 25 fps, resize to the model's pixel budget, compute whole-video NeuFlow v2 optical flow, then run the 6-channel Qwen3-VL model over 10-second sliding windows and merge the results. Args: video: path to the video file to analyse. max_seconds: analyse only the first N seconds of the video. window_size: sliding-window length in seconds (paper protocol: 10). stride: sliding-window stride in seconds (paper protocol: 9). merge_eps: tolerance in seconds when merging overlapping predictions. max_new_tokens: generation cap per window. Returns: The optical-flow visualisation video, a markdown summary, an HTML timeline, a table of detected transitions, and the full run record. """ if not video: raise gr.Error("Please upload a video first.") src = Path(video) if not src.is_file(): raise gr.Error(f"Video not found: {video}") max_seconds = float(max_seconds) window_size = float(window_size) stride = float(stride) merge_eps = float(merge_eps) max_new_tokens = int(max_new_tokens) ENGINE.generate_kwargs["max_new_tokens"] = max_new_tokens work = Path(tempfile.mkdtemp(prefix="transvlm_")) wall_started = time.perf_counter() timings: dict[str, float] = {} # 1. fps -> 25 (and trim to the analysis budget) progress(0.02, desc="Normalising to 25 fps…") started = time.perf_counter() model_src, resampled_from = _trim_and_resample(src, work, max_seconds, TARGET_FPS) timings["prepare_sec"] = round(time.perf_counter() - started, 3) # 2. smart_resize to the model's pixel budget progress(0.08, desc="Resizing to the model's pixel budget…") started = time.perf_counter() model_rgb, src_hw, target_hw = _maybe_resize(model_src, work) timings["resize_sec"] = round(time.perf_counter() - started, 3) # 3. whole-video optical flow, in ONE NeuFlow call (flow_to_image normalises # over its entire input, so this cannot be chunked without changing the signal) progress(0.16, desc="Computing NeuFlow v2 optical flow…") model_flow = work / f"{model_rgb.stem}_flow.mp4" started = time.perf_counter() _flow_out, flow_parts = FLOW.compute_flow_only(model_rgb, model_flow) timings["flow_sec"] = round(time.perf_counter() - started, 3) timings.update(flow_parts.as_dict()) # 4. window plan — time ranges, nothing is cut on disk duration = _duration_sec(model_rgb) windows = get_sliding_window_segments( start=0.0, end=duration, window_size=window_size, stride=stride, strict_tail=False, ) logger.info("[windows] %.2f s -> %d window(s)", duration, len(windows)) # 5. per-window inference half_frame = 0.5 / ENGINE.fps all_segments = [] per_window = [] preprocess_total = 0.0 generate_total = 0.0 for idx, (w_start, w_end) in enumerate(windows): progress( 0.25 + 0.7 * idx / max(len(windows), 1), desc=f"Window {idx + 1}/{len(windows)} — {w_start:.1f}–{w_end:.1f} s", ) requested_end = min(w_end - half_frame, duration) result = ENGINE.infer( model_rgb, model_flow, video_start=w_start, video_end=requested_end, rebase_frame_indices=True, ) preprocess_total += result.preprocess_sec generate_total += result.generate_sec per_window.append( { "idx": idx, "start": round(w_start, 4), "end": round(w_end, 4), "requested_end": round(requested_end, 4), "n_frames": result.n_frames, "height": result.height, "width": result.width, "n_input_tokens": result.n_input_tokens, "n_output_tokens": result.n_output_tokens, "raw_output": result.raw_text, "preprocess_sec": round(result.preprocess_sec, 3), "generate_sec": round(result.generate_sec, 3), } ) all_segments.extend(parse_model_output(result.raw_text, window_offset=w_start, window_index=idx)) # 6. merge overlapping predictions across windows progress(0.97, desc="Merging predictions…") merged = merge_shots(all_segments, eps=merge_eps) segments = [ { "start_time": round(float(s.start_time), 3), "end_time": round(float(s.end_time), 3), "kind": _kind(float(s.start_time), float(s.end_time)), "source_window": s.source_window, } for s in merged ] timings["preprocess_sec_total"] = round(preprocess_total, 3) timings["generate_sec_total"] = round(generate_total, 3) timings["wall_sec"] = round(time.perf_counter() - wall_started, 3) n_cuts = sum(1 for s in segments if s["kind"] == "cut") n_grad = len(segments) - n_cuts summary = ( f"### {len(segments)} transition{'s' if len(segments) != 1 else ''} " f"— {n_cuts} hard cut{'s' if n_cuts != 1 else ''}, " f"{n_grad} gradual\n" f"`{duration:.2f}s analysed` · `{len(windows)} window(s)` · " f"`{src_hw[1]}×{src_hw[0]} → {target_hw[1]}×{target_hw[0]}` · " f"`{timings['wall_sec']:.1f}s wall` " f"(flow {timings['flow_sec']:.1f}s, generate {timings['generate_sec_total']:.1f}s)" ) table = [ [ i, f"{s['start_time']:.2f}", f"{s['end_time']:.2f}", f"{s['end_time'] - s['start_time']:.2f}", s["kind"], ] for i, s in enumerate(segments, 1) ] if not table: table = [["—", "—", "—", "—", "no transition detected"]] record = { "model": MODEL_ID, "duration_sec": round(duration, 3), "resampled_from": resampled_from, "source_hw": {"height": src_hw[0], "width": src_hw[1]}, "model_hw": {"height": target_hw[0], "width": target_hw[1]}, "n_windows": len(windows), "n_raw_segments": len(all_segments), "n_segments": len(segments), "segments": segments, "time_base": "absolute_video", "windows": per_window, "timings": timings, "config": { "fps": DEFAULT_FPS, "window_size": window_size, "stride": stride, "strict_tail": False, "merge_eps": merge_eps, "max_pixels_override": MAX_PIXELS_OVERRIDE, "max_new_tokens": max_new_tokens, "timestamp_format": ENGINE.timestamp_format, "flow_codec": FLOW_CODEC, "flow_viz_on_gpu": True, "video_decoder": selected_video_backend(), "backend": "hf", }, } logger.info("[done] %s", json.dumps(record["timings"])) return ( str(_faststart(model_flow)), summary, _timeline_html(segments, duration), table, record, ) # --------------------------------------------------------------------------- # UI # --------------------------------------------------------------------------- EXAMPLES_DIR = Path(__file__).resolve().parent / "examples" # The authors' own showcase clips, taken from the TransVLM project page # (docs/assets/media/samples/ in github.com/heygen-com/TransVLM, Apache-2.0). # The repo publishes ground truth AND the paper's TransVLM predictions for each # one, so these examples are verifiable rather than decorative. # abrupt_cuts.mp4 = synth_fast_b8d0c0de/original.mp4 (STD-Synth, "abrupt cuts") # GT cuts at 5.10 / 7.26 / 12.50 s # gradual_transitions.mp4 = synth_base_3a2d5b40/s03.mp4 (STD-Synth, "normal transitions") # GT 2.52-8.20 s and 9.12-9.92 s # long_dissolve.mp4 = synth_long_78cdf2e9/s01.mp4 (STD-Synth, "long transitions") # GT 2.35-10.00 s and 11.98-12.36 s EXAMPLE_ROWS = [ [str(EXAMPLES_DIR / "abrupt_cuts.mp4"), 14.0], [str(EXAMPLES_DIR / "gradual_transitions.mp4"), 11.0], [str(EXAMPLES_DIR / "long_dissolve.mp4"), 14.0], ] EXAMPLE_ROWS = [row for row in EXAMPLE_ROWS if Path(row[0]).is_file()] HEADER = f""" # 🎬 TransVLM — Detect Any Shot Transition Hard cuts are easy. **Dissolves, fades, wipes and whip-pans are not** — they are gradual, and classic shot-boundary detectors either miss them or fire on camera motion. [TransVLM]({PAPER_URL}) fuses **optical flow into the vision tower at the input stage** (a 6-channel patch embed over RGB ⊕ flow) so a Qwen3-VL backbone can see *motion* as well as *appearance*, and asks it to write out the start and end time of every transition it finds. This Space runs the authors' full pipeline: 25 fps resample → `smart_resize` → whole-video **NeuFlow v2** optical flow → 10 s sliding windows → JSON parse → merge. [model]({f"https://huggingface.co/{MODEL_ID}"}) · [code]({CODE_URL}) · [paper]({PAPER_URL}) """ with gr.Blocks(theme=gr.themes.Citrus(), title="TransVLM — Shot Transition Detection") as demo: gr.Markdown(HEADER) with gr.Row(): with gr.Column(scale=1): video_in = gr.Video(label="Video", sources=["upload"], height=340) run_btn = gr.Button("Detect transitions", variant="primary") max_seconds_in = gr.Slider( 4, 30, value=DEFAULT_MAX_SECONDS, step=1, label="Seconds to analyse", info="Only the first N seconds are analysed, so one request fits a GPU slot.", ) with gr.Accordion("Advanced (paper protocol defaults)", open=False): window_size_in = gr.Slider( 4, 12, value=DEFAULT_WINDOW_SIZE, step=0.5, label="Sliding-window size (s)" ) stride_in = gr.Slider(2, 12, value=DEFAULT_STRIDE, step=0.5, label="Window stride (s)") merge_eps_in = gr.Slider( 0.0, 0.5, value=DEFAULT_MERGE_EPS, step=0.01, label="Merge tolerance (s)" ) max_new_tokens_in = gr.Slider( 128, 2048, value=DEFAULT_MAX_NEW_TOKENS, step=64, label="Max new tokens per window" ) gr.Markdown( "The model was trained and evaluated at **25 fps**, a **524,288-pixel** " "per-frame budget and **10 s / 9 s** windows. Changing these changes the " "signal the model sees." ) with gr.Column(scale=1): summary_out = gr.Markdown() timeline_out = gr.HTML() table_out = gr.Dataframe( headers=["#", "start (s)", "end (s)", "duration (s)", "type"], datatype=["str", "str", "str", "str", "str"], label="Detected transitions", wrap=True, ) flow_out = gr.Video( label="NeuFlow v2 optical flow (the model's second input stream)", height=280, autoplay=False, ) with gr.Accordion("Full run record (raw model output, timings, config)", open=False): json_out = gr.JSON() inputs = [video_in, max_seconds_in, window_size_in, stride_in, merge_eps_in, max_new_tokens_in] outputs = [flow_out, summary_out, timeline_out, table_out, json_out] run_btn.click(fn=detect, inputs=inputs, outputs=outputs, api_name="detect") if EXAMPLE_ROWS: gr.Examples( examples=EXAMPLE_ROWS, inputs=[video_in, max_seconds_in], outputs=outputs, fn=detect, cache_examples=True, cache_mode="lazy", label="Examples — the authors' own STD-Synth showcase clips", ) gr.Markdown( "Ground truth published with the paper, for the three examples above:\n\n" "| clip | ground-truth transitions |\n|---|---|\n" "| `abrupt_cuts` | three hard cuts at **5.10 s**, **7.26 s**, **12.50 s** |\n" "| `gradual_transitions` | **2.52–8.20 s** and **9.12–9.92 s** |\n" "| `long_dissolve` | **2.35–10.00 s** (7.6 s squeeze) and **11.98–12.36 s** |\n" ) if __name__ == "__main__": demo.queue(max_size=12).launch(mcp_server=True, show_error=True)