#!/usr/bin/env python3 """Extract causal Qwen2.5-Omni states directly from original media. Unlike ``prepare_easycom_session.py`` and ``prepare_ego4d_2fps.py``, this entry point does not create a re-encoded fixed-frame-rate MP4. It builds a global visual sampling lattice in Python, selects the newest decoded source frame whose presentation timestamp is not later than each lattice time, and passes explicit PIL frame lists plus exact 16-kHz audio slices to Qwen. The default mode follows the released Qwen2.5-Omni streaming design as closely as the public Transformers implementation permits: * 2 FPS causal visual sampling (configurable for ablations); * two-frame temporal patches, with the last observed frame duplicated when an incomplete input contains an odd number of frames; * a single open media turn with completed 2-second AV payload commits; * persistent Thinker KV for the complete extraction (``--memory-seconds 0``); * one temporary partial-block branch and readout every 200 ms. Examples -------- Ego4D, directly from one original clip:: python streaming_turn_model_directly.py \ --model-id /data/wangx/Models/Qwen2.5-Omni-3B \ --video /data/wangx/Datasets/ego4d/v1/clips/UID.mp4 \ --target /data/wangx/Datasets/ego4d/ego4d_target_frame/UID.npy \ --mode open-stream-replay \ --memory-seconds 20 \ --video-fps 2 \ --flash-attention-2 \ --device cuda:0 \ --output features/ego4d_direct/UID.pt EasyCom, directly from the retained one-minute source clips:: python streaming_turn_model_directly.py \ --model-id /path/to/Qwen2.5-Omni-3B \ --easycom-session-dir /path/to/Main/Video_Compressed/Session_1 \ --target target_frame/Session_1.npy \ --mode open-stream-replay \ --memory-seconds 20 \ --video-fps 2 \ --flash-attention-2 \ --device cuda:0 \ --output features/easycom_direct/Session_1.pt ``--media-only`` performs source/timestamp/audio validation without loading Qwen. This is useful before launching a full extraction on a remote machine. """ from __future__ import annotations import argparse import json import math import shutil import subprocess from collections import OrderedDict from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Sequence import numpy as np import torch import streaming_turn_model as core from easycom_manifest import parse_timecode EPSILON = 1e-7 SAMPLE_RATE = 16_000 EASYCOM_WIDTH = 1920 EASYCOM_HEIGHT = 1080 @dataclass(frozen=True) class SourceClipSpec: """One source file in a dense, gap-free media timeline.""" path: Path duration: float | None = None crop: tuple[int, int, int, int] | None = None crop_rightmost_width: int | None = None audio_path: Path | None = None @dataclass class IndexedSourceClip: """Decoded indexing information for one source clip.""" path: Path audio_path: Path dense_start: float dense_end: float duration: float crop: tuple[int, int, int, int] | None reader: Any | None source_frame_starts: np.ndarray source_frame_ends: np.ndarray source_fps: float width: int height: int @property def num_frames(self) -> int: return int(len(self.source_frame_starts)) def _grid_indices(start: float, end: float, fps: float) -> list[int]: """Indices of global samples ``n/fps`` in the half-open interval.""" first = max(0, int(math.ceil(start * fps - EPSILON))) last_exclusive = int(math.ceil(end * fps - EPSILON)) if first >= last_exclusive: # Qwen's visual tower needs at least one observed image before its # temporal-patch duplication. Reuse the latest global tick strictly # before ``end``; this is identical to FixedGridAV's short-tail rule. return [max(0, last_exclusive - 1)] return list(range(first, last_exclusive)) def _timestamp_array(reader: Any) -> tuple[np.ndarray, np.ndarray]: """Return monotonically increasing, zero-origin frame PTS intervals.""" # Decord releases differ here: some accept an ndarray, while older ones # require a Python list. A list works with both APIs. indices = list(range(len(reader))) raw = reader.get_frame_timestamp(indices) if hasattr(raw, "asnumpy"): raw = raw.asnumpy() timestamps = np.asarray(raw, dtype=np.float64) if timestamps.ndim == 1: starts = timestamps nominal = 1.0 / float(reader.get_avg_fps()) ends = np.concatenate((starts[1:], starts[-1:] + nominal)) elif timestamps.ndim == 2 and timestamps.shape[1] >= 2: starts = timestamps[:, 0] ends = timestamps[:, 1] else: raise RuntimeError( f"Unexpected decord frame-timestamp shape: {timestamps.shape}" ) if not len(starts): raise RuntimeError("Video contains no decoded frames") origin = float(starts[0]) starts = starts - origin ends = ends - origin if np.any(np.diff(starts) < -1e-9): raise RuntimeError("Decoded video frame PTS values are not monotonic") return starts, ends def _parse_ffprobe_rate(value: str, path: Path) -> float: """Parse an ffprobe rational such as ``20/1``.""" try: numerator, denominator = value.split("/", maxsplit=1) rate = float(numerator) / float(denominator) except (AttributeError, ValueError, ZeroDivisionError) as exc: raise RuntimeError(f"Invalid ffprobe frame rate {value!r}: {path}") from exc if not math.isfinite(rate) or rate <= 0: raise RuntimeError(f"Non-positive ffprobe frame rate {rate}: {path}") return rate def _ffprobe_video_index( path: Path, ) -> tuple[np.ndarray, np.ndarray, float, int, int]: """Index display-order frame PTS without keeping a software decoder open.""" if shutil.which("ffprobe") is None: raise RuntimeError("ffprobe is required by non-Decord video backends") completed = subprocess.run( [ "ffprobe", "-v", "error", "-select_streams", "v:0", "-show_frames", "-show_entries", ( "stream=width,height,avg_frame_rate:" "frame=best_effort_timestamp_time,pkt_duration_time" ), "-of", "json", str(path), ], check=False, capture_output=True, text=True, ) if completed.returncode != 0: raise RuntimeError( f"Failed to index video with ffprobe: {path}: " f"{completed.stderr.strip()}" ) try: payload = json.loads(completed.stdout) stream = payload["streams"][0] width = int(stream["width"]) height = int(stream["height"]) source_fps = _parse_ffprobe_rate(stream["avg_frame_rate"], path) raw_starts = [ float(frame["best_effort_timestamp_time"]) for frame in payload["frames"] ] except (IndexError, KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: raise RuntimeError(f"Incomplete ffprobe frame index for {path}") from exc if not raw_starts: raise RuntimeError(f"Video contains no indexed frames: {path}") starts = np.asarray(raw_starts, dtype=np.float64) starts -= float(starts[0]) if np.any(np.diff(starts) < -1e-9): raise RuntimeError(f"ffprobe video frame PTS values are not monotonic: {path}") nominal = 1.0 / source_fps ends = np.concatenate((starts[1:], starts[-1:] + nominal)) return starts, ends, source_fps, width, height def _decode_audio_exact(path: Path, duration: float) -> np.ndarray: """Decode, then pad/trim one clip to its dense video duration.""" completed = subprocess.run( [ "ffmpeg", "-nostdin", "-v", "error", "-threads", "1", "-i", str(path), "-map", "0:a:0", "-vn", "-ac", "1", "-ar", str(SAMPLE_RATE), "-c:a", "pcm_f32le", "-f", "f32le", "pipe:1", ], check=False, capture_output=True, ) if completed.returncode != 0: message = completed.stderr.decode("utf-8", errors="replace").strip() raise RuntimeError(f"Failed to decode audio from {path}: {message}") decoded = np.frombuffer(completed.stdout, dtype=" None: if not specs: raise ValueError("At least one source clip is required") if output_fps <= 0: raise ValueError("output_fps must be positive") if required_duration <= 0: raise ValueError("required_duration must be positive") if frame_cache_size < 0: raise ValueError("frame_cache_size must be non-negative") if decoder_backend not in {"decord", "torchcodec", "pynvcodec"}: raise ValueError(f"Unsupported decoder backend: {decoder_backend}") if decoder_gpu_id < 0: raise ValueError("decoder_gpu_id must be non-negative") if decoder_threads <= 0: raise ValueError("decoder_threads must be positive") if max_open_readers <= 0: raise ValueError("max_open_readers must be positive") if shutil.which("ffmpeg") is None: raise RuntimeError("ffmpeg is required to decode source audio") decord = None self._decord = None self._torchcodec = None self._torchcodec_decoder_cls = None self._nvc = None if decoder_backend == "decord": try: import decord except ImportError as exc: raise RuntimeError( "The decord backend requires qwen-omni-utils[decord]" ) from exc self._decord = decord elif decoder_backend == "torchcodec": try: import torchcodec from torchcodec.decoders import VideoDecoder except ImportError as exc: raise RuntimeError( "The torchcodec backend requires TorchCodec; install it " "with `python -m pip install torchcodec`" ) from exc self._torchcodec = torchcodec self._torchcodec_decoder_cls = VideoDecoder else: try: import PyNvVideoCodec as nvc except ImportError as exc: raise RuntimeError( "The pynvcodec backend requires PyNvVideoCodec; install it " "with `python -m pip install PyNvVideoCodec`" ) from exc self._nvc = nvc self.fps = float(output_fps) self.required_duration = float(required_duration) self.frame_cache_size = int(frame_cache_size) self.decoder_backend = decoder_backend self.decoder_gpu_id = int(decoder_gpu_id) self.decoder_threads = int(decoder_threads) self.max_open_readers = int(max_open_readers) self._decord_readers: OrderedDict[int, Any] = OrderedDict() self._decord_reader_opens = 0 self._decord_reader_evictions = 0 self._torchcodec_readers: OrderedDict[int, Any] = OrderedDict() self._torchcodec_reader_opens = 0 self._torchcodec_reader_evictions = 0 self._nvc_decoder = None self._nvc_source_path: Path | None = None self._nvc_decoder_reconfigurations = 0 self._frame_cache: OrderedDict[tuple[int, int], Any] = OrderedDict() self._request_count = 0 self._cache_hits = 0 self._odd_frame_duplications = 0 self._short_interval_fallbacks = 0 self._tail_frame_reuses = 0 indexed: list[IndexedSourceClip] = [] audio_parts: list[np.ndarray] = [] dense_cursor = 0.0 for spec_index, spec in enumerate(specs): path = spec.path.expanduser().resolve() if not path.is_file(): raise FileNotFoundError(f"Source video does not exist: {path}") print( f"direct-media: indexing [{spec_index + 1}/{len(specs)}] " f"{path.name}", flush=True, ) if decoder_backend == "decord": assert decord is not None index_reader = decord.VideoReader( str(path), num_threads=self.decoder_threads ) if not len(index_reader): raise RuntimeError(f"No video frames decoded from {path}") starts, ends = _timestamp_array(index_reader) source_fps = float(index_reader.get_avg_fps()) first_shape = tuple(int(value) for value in index_reader[0].shape) if len(first_shape) != 3: raise RuntimeError( f"Unexpected decoded frame shape {first_shape}: {path}" ) height, width = first_shape[:2] # VideoReader.__del__ releases its native handle. Do not retain # one reader/thread pool for every one-minute EasyCom clip. del index_reader reader = None else: reader = None starts, ends, source_fps, width, height = _ffprobe_video_index(path) inferred_duration = len(starts) / source_fps duration = ( float(spec.duration) if spec.duration is not None else float(inferred_duration) ) if duration <= 0: raise RuntimeError(f"Non-positive source duration for {path}") crop = spec.crop if crop is None and spec.crop_rightmost_width is not None: keep_width = int(spec.crop_rightmost_width) if width < keep_width: raise RuntimeError( f"Cannot keep rightmost {keep_width}px from {width}px: {path}" ) crop = (width - keep_width, 0, keep_width, height) if crop is not None: x, y, crop_width, crop_height = crop if ( x < 0 or y < 0 or crop_width <= 0 or crop_height <= 0 or x + crop_width > width or y + crop_height > height ): raise RuntimeError( f"Invalid crop {crop} for {width}x{height}: {path}" ) audio_path = (spec.audio_path or path).expanduser().resolve() if not audio_path.is_file(): raise FileNotFoundError(f"Audio source does not exist: {audio_path}") clip = IndexedSourceClip( path=path, audio_path=audio_path, dense_start=dense_cursor, dense_end=dense_cursor + duration, duration=duration, crop=crop, reader=reader, source_frame_starts=starts, source_frame_ends=ends, source_fps=source_fps, width=width, height=height, ) indexed.append(clip) print( f"direct-media: decoding audio [{spec_index + 1}/{len(specs)}] " f"{audio_path.name}", flush=True, ) audio_parts.append(_decode_audio_exact(audio_path, duration)) dense_cursor = clip.dense_end # No later source can be observed before the requested endpoint. if dense_cursor >= required_duration - EPSILON: break self.clips = indexed self.dense_starts = np.asarray( [clip.dense_start for clip in indexed], dtype=np.float64 ) self.dense_duration = float(dense_cursor) if self.dense_duration + EPSILON < self.required_duration: raise RuntimeError( f"Dense source duration is {self.dense_duration:.6f}s, shorter " f"than required {self.required_duration:.6f}s" ) self.audio = np.concatenate(audio_parts) required_samples = int(round(required_duration * SAMPLE_RATE)) if len(self.audio) < required_samples: self.audio = np.pad(self.audio, (0, required_samples - len(self.audio))) elif len(self.audio) > required_samples: self.audio = self.audio[:required_samples] max_grid_index = int( math.ceil(required_duration * self.fps - EPSILON) ) - 1 self.expected_grid_frames = max(0, max_grid_index + 1) if max_grid_index >= 0: self._source_frame_for_grid(max_grid_index, count_tail_reuse=False) def _clip_for_time(self, global_time: float) -> tuple[int, IndexedSourceClip, float]: clip_index = int(np.searchsorted(self.dense_starts, global_time, side="right") - 1) clip_index = max(0, min(clip_index, len(self.clips) - 1)) clip = self.clips[clip_index] local_time = global_time - clip.dense_start if local_time < 0: local_time = 0.0 return clip_index, clip, local_time def _source_frame_for_grid( self, grid_index: int, *, count_tail_reuse: bool = True, ) -> tuple[int, int, float, float]: """Map one global output tick to a source frame at or before it.""" global_time = grid_index / self.fps clip_index, frame_index, source_time = self._source_frame_for_time( global_time, count_tail_reuse=count_tail_reuse, ) return clip_index, frame_index, global_time, source_time def _source_frame_for_time( self, global_time: float, *, count_tail_reuse: bool = True, ) -> tuple[int, int, float]: """Map an arbitrary causal timestamp to the latest observed frame.""" if global_time < -EPSILON: raise ValueError(f"Negative video timestamp: {global_time}") if global_time > self.required_duration + EPSILON: raise ValueError( f"Video timestamp {global_time:g}s exceeds required duration " f"{self.required_duration:g}s" ) global_time = min(max(0.0, float(global_time)), self.required_duration) clip_index, clip, local_time = self._clip_for_time(global_time) frame_index = int( np.searchsorted(clip.source_frame_starts, local_time + EPSILON, side="right") - 1 ) frame_index = max(0, min(frame_index, clip.num_frames - 1)) source_time = float(clip.source_frame_starts[frame_index]) if source_time > local_time + 2 * EPSILON: raise RuntimeError( "Causal frame-selection violation: " f"source PTS {source_time:.9f}s > grid-local time {local_time:.9f}s" ) if count_tail_reuse and local_time > float(clip.source_frame_starts[-1]) + EPSILON: self._tail_frame_reuses += 1 return clip_index, frame_index, source_time def frame_at_or_before(self, global_time: float) -> Any: """Decode the latest source frame whose PTS is not later than ``time``.""" clip_index, frame_index, _ = self._source_frame_for_time(global_time) self._request_count += 1 return self._load_frame(clip_index, frame_index) def _load_frame(self, clip_index: int, frame_index: int) -> Any: from PIL import Image key = (clip_index, frame_index) cached = self._frame_cache.get(key) if cached is not None: self._cache_hits += 1 self._frame_cache.move_to_end(key) return cached.copy() clip = self.clips[clip_index] if self.decoder_backend == "decord": reader = self._decord_reader_for_clip(clip_index) array = reader[frame_index].asnumpy() elif self.decoder_backend == "torchcodec": decoder = self._torchcodec_reader_for_clip(clip_index) decoded = decoder.get_frames_at(indices=[frame_index]) frame_tensor = decoded.data if hasattr(decoded, "data") else decoded if frame_tensor.ndim == 4 and frame_tensor.shape[0] == 1: frame_tensor = frame_tensor[0] if frame_tensor.ndim != 3: raise RuntimeError( "TorchCodec output must be one CHW or HWC frame, got " f"{tuple(frame_tensor.shape)} for {clip.path}" ) if frame_tensor.shape[0] in (3, 4): frame_tensor = frame_tensor[:3].permute(1, 2, 0) elif frame_tensor.shape[-1] == 4: frame_tensor = frame_tensor[..., :3] elif frame_tensor.shape[-1] != 3: raise RuntimeError( "TorchCodec output has no RGB channel axis: " f"{tuple(frame_tensor.shape)} for {clip.path}" ) if frame_tensor.dtype != torch.uint8: raise RuntimeError( "TorchCodec RGB output must be uint8, got " f"{frame_tensor.dtype} for {clip.path}" ) array = frame_tensor.cpu().contiguous().numpy() else: decoder = self._pynv_decoder_for_clip(clip) decoded_frame = decoder[frame_index] frame_tensor = torch.from_dlpack(decoded_frame) if frame_tensor.ndim != 3 or frame_tensor.shape[-1] != 3: raise RuntimeError( "PyNvVideoCodec RGB output must have HWC shape, got " f"{tuple(frame_tensor.shape)} for {clip.path}" ) if frame_tensor.dtype != torch.uint8: raise RuntimeError( "PyNvVideoCodec RGB output must be uint8, got " f"{frame_tensor.dtype} for {clip.path}" ) array = frame_tensor.cpu().contiguous().numpy() if clip.crop is not None: x, y, width, height = clip.crop array = array[y : y + height, x : x + width] image = Image.fromarray(np.asarray(array)).copy() if self.frame_cache_size: self._frame_cache[key] = image self._frame_cache.move_to_end(key) while len(self._frame_cache) > self.frame_cache_size: self._frame_cache.popitem(last=False) return image.copy() def _decord_reader_for_clip(self, clip_index: int) -> Any: """Return an LRU-managed CPU reader instead of retaining every clip.""" reader = self._decord_readers.get(clip_index) if reader is not None: self._decord_readers.move_to_end(clip_index) return reader if self._decord is None: raise RuntimeError("Decord was not initialized") clip = self.clips[clip_index] reader = self._decord.VideoReader( str(clip.path), num_threads=self.decoder_threads ) if len(reader) != clip.num_frames: raise RuntimeError( f"Decord frame count changed for {clip.path}: " f"runtime={len(reader)}, indexed={clip.num_frames}" ) self._decord_readers[clip_index] = reader self._decord_readers.move_to_end(clip_index) self._decord_reader_opens += 1 while len(self._decord_readers) > self.max_open_readers: _, evicted_reader = self._decord_readers.popitem(last=False) del evicted_reader self._decord_reader_evictions += 1 return reader def _torchcodec_reader_for_clip(self, clip_index: int) -> Any: """Return an LRU-managed, exact-seek TorchCodec CPU decoder.""" reader = self._torchcodec_readers.get(clip_index) if reader is not None: self._torchcodec_readers.move_to_end(clip_index) return reader if self._torchcodec_decoder_cls is None: raise RuntimeError("TorchCodec was not initialized") clip = self.clips[clip_index] reader = self._torchcodec_decoder_cls( str(clip.path), num_ffmpeg_threads=self.decoder_threads, ) decoded_length = getattr(reader.metadata, "num_frames", None) if decoded_length is not None and int(decoded_length) != clip.num_frames: raise RuntimeError( "TorchCodec/ffprobe frame-count mismatch for " f"{clip.path}: decoder={decoded_length}, ffprobe={clip.num_frames}" ) self._torchcodec_readers[clip_index] = reader self._torchcodec_readers.move_to_end(clip_index) self._torchcodec_reader_opens += 1 while len(self._torchcodec_readers) > self.max_open_readers: _, evicted_reader = self._torchcodec_readers.popitem(last=False) del evicted_reader self._torchcodec_reader_evictions += 1 return reader def _pynv_decoder_for_clip(self, clip: IndexedSourceClip) -> Any: """Return one reusable NVDEC decoder configured for ``clip``.""" if self._nvc is None: raise RuntimeError("PyNvVideoCodec was not initialized") source_path = clip.path if self._nvc_decoder is None: self._nvc_decoder = self._new_pynv_decoder(source_path) self._nvc_source_path = source_path elif self._nvc_source_path != source_path: torch.cuda.synchronize(self.decoder_gpu_id) try: self._nvc_decoder.reconfigure_decoder(str(source_path)) except Exception: # A rebuild is safer than silently using stale decoder state. self._nvc_decoder = self._new_pynv_decoder(source_path) self._nvc_source_path = source_path self._nvc_decoder_reconfigurations += 1 decoded_length = len(self._nvc_decoder) if decoded_length != clip.num_frames: raise RuntimeError( "PyNvVideoCodec/ffprobe frame-count mismatch for " f"{clip.path}: decoder={decoded_length}, ffprobe={clip.num_frames}" ) return self._nvc_decoder def _new_pynv_decoder(self, source_path: Path) -> Any: assert self._nvc is not None return self._nvc.SimpleDecoder( str(source_path), gpu_id=self.decoder_gpu_id, use_device_memory=True, need_scanned_stream_metadata=True, decoder_cache_size=4, output_color_type=self._nvc.OutputColorType.RGB, ) def frames(self, start: float, end: float) -> list[Any]: """Return causal global-grid frames in ``[start, end)``.""" if end <= start: raise ValueError(f"Empty video interval [{start}, {end}]") raw_first = int(math.ceil(start * self.fps - EPSILON)) raw_last_exclusive = int(math.ceil(end * self.fps - EPSILON)) grid_indices = _grid_indices(start, end, self.fps) if raw_first >= raw_last_exclusive: self._short_interval_fallbacks += 1 output = [] for grid_index in grid_indices: clip_index, frame_index, _, _ = self._source_frame_for_grid(grid_index) output.append(self._load_frame(clip_index, frame_index)) self._request_count += 1 if len(output) % 2: output.append(output[-1].copy()) self._odd_frame_duplications += 1 return output def audio_slice(self, start: float, end: float) -> np.ndarray: """Return the exact 16-kHz mono samples in ``[start, end)``.""" first = max(0, int(math.floor(start * SAMPLE_RATE + EPSILON))) last = min( len(self.audio), int(math.floor(end * SAMPLE_RATE + EPSILON)) ) if last <= first: raise ValueError(f"Empty audio interval [{start}, {end}]") return self.audio[first:last] def sample_mapping(self, grid_index: int) -> dict[str, Any]: """Describe one visual-grid mapping without decoding its pixels.""" clip_index, frame_index, grid_time, source_time = self._source_frame_for_grid( grid_index, count_tail_reuse=False ) clip = self.clips[clip_index] return { "grid_index": int(grid_index), "grid_time": float(grid_time), "source_clip": str(clip.path), "source_frame_index": int(frame_index), "source_frame_pts_seconds": float(source_time), "dense_source_frame_time": float(clip.dense_start + source_time), } def metadata(self) -> dict[str, Any]: decoder_names = { "decord": "decord_with_frame_pts_and_ffmpeg_pcm_audio", "torchcodec": "torchcodec_cpu_with_ffprobe_pts_and_ffmpeg_pcm_audio", "pynvcodec": ( "pynvvideocodec_nvdec_with_ffprobe_pts_and_ffmpeg_pcm_audio" ), } return { "decoder": decoder_names[self.decoder_backend], "decoder_backend": self.decoder_backend, "decoder_gpu_id": ( self.decoder_gpu_id if self.decoder_backend == "pynvcodec" else None ), "decoder_threads": ( self.decoder_threads if self.decoder_backend in {"decord", "torchcodec"} else None ), "max_open_readers": ( self.max_open_readers if self.decoder_backend in {"decord", "torchcodec"} else None ), "decoder_version": ( getattr(self._nvc, "__version__", None) if self.decoder_backend == "pynvcodec" else ( getattr(self._torchcodec, "__version__", None) if self.decoder_backend == "torchcodec" else None ) ), "visual_sampling": "global_grid_latest_source_pts_not_after_tick", "interval_convention": "half_open_[start,end)", "output_video_fps": self.fps, "expected_grid_frames": self.expected_grid_frames, "required_duration_seconds": self.required_duration, "dense_source_duration_seconds": self.dense_duration, "audio_sample_rate": SAMPLE_RATE, "audio_samples": int(len(self.audio)), "temporal_patch_padding": "duplicate_last_observed_frame_when_odd", "source_clips": [ { "path": str(clip.path), "audio_path": str(clip.audio_path), "dense_start": clip.dense_start, "dense_end": clip.dense_end, "duration": clip.duration, "source_fps": clip.source_fps, "source_frames": clip.num_frames, "source_size": [clip.width, clip.height], "crop": list(clip.crop) if clip.crop is not None else None, "first_source_pts": float(clip.source_frame_starts[0]), "last_source_pts": float(clip.source_frame_starts[-1]), } for clip in self.clips ], "runtime_counters": { "frame_requests": self._request_count, "frame_cache_hits": self._cache_hits, "odd_frame_duplications": self._odd_frame_duplications, "short_interval_fallbacks": self._short_interval_fallbacks, "tail_frame_reuses": self._tail_frame_reuses, "decoder_reconfigurations": self._nvc_decoder_reconfigurations, "decord_reader_opens": self._decord_reader_opens, "decord_reader_evictions": self._decord_reader_evictions, "torchcodec_reader_opens": self._torchcodec_reader_opens, "torchcodec_reader_evictions": self._torchcodec_reader_evictions, }, } class GlobalHistoryCurrentAV: """Working-window view with cached global history plus the current frame. Complete cached blocks use the source reader's fixed global half-open grid, for example ``[18.0, 18.5, 19.0, 19.5]`` for ``[18,20)`` at 2 FPS. A two-second working interval ending at ``t`` uses the latest global-grid samples at or before ``t-1.5``, ``t-1.0``, and ``t-0.5``, followed by the current causal frame at ``t``. At ``t=20.4`` this is ``[18.5, 19.0, 19.5, 20.4]``. Historical pixels therefore reuse the same source-frame cache as complete blocks and only the current frame advances at 5 Hz. Audio remains the exact half-open slice ``[start,end)``. """ def __init__( self, source: CausalSourceAV, *, max_interval_seconds: float, ) -> None: if max_interval_seconds <= 0: raise ValueError("max_interval_seconds must be positive") self.source = source self.fps = source.fps self.sample_rate = source.sample_rate self.max_interval_seconds = float(max_interval_seconds) max_frames = self.max_interval_seconds * self.fps if not math.isclose(max_frames, round(max_frames), abs_tol=1e-6): raise ValueError( "working-memory seconds * video FPS must be integral, got " f"{max_frames}" ) self.max_frames = int(round(max_frames)) def sample_times(self, start: float, end: float) -> list[float]: if end <= start: raise ValueError(f"Empty video interval [{start}, {end}]") duration = end - start if duration > self.max_interval_seconds + EPSILON: raise ValueError( f"Working interval is {duration:g}s, exceeds configured " f"{self.max_interval_seconds:g}s" ) count = max(1, int(math.ceil(duration * self.fps - EPSILON))) count = min(count, self.max_frames) history_times = [] for slots_before_current in range(count - 1, 0, -1): desired_time = end - slots_before_current / self.fps # Quantize only historical samples to the session-global grid. # During the short initial warm-up, clamp to the causal interval # start rather than selecting media from before the requested # working window. grid_index = int(math.floor(desired_time * self.fps + EPSILON)) history_times.append(max(start, grid_index / self.fps)) return [*history_times, end] def frames(self, start: float, end: float) -> list[Any]: frames = [ self.source.frame_at_or_before(timestamp) for timestamp in self.sample_times(start, end) ] if len(frames) % 2: frames.append(frames[-1].copy()) self.source._odd_frame_duplications += 1 return frames def audio_slice(self, start: float, end: float) -> np.ndarray: return self.source.audio_slice(start, end) def metadata(self) -> dict[str, Any]: return { "visual_sampling": "global_grid_history_plus_current_causal_frame", "audio_interval": "half_open_[start,end)", "working_memory_seconds": self.max_interval_seconds, "output_video_fps": self.fps, "full_window_frame_count": self.max_frames, "desired_history_offsets_seconds": [ -(self.max_frames - 1 - index) / self.fps for index in range(self.max_frames - 1) ], "current_frame_offset_seconds": 0.0, "history_timestamp_quantization": "floor_to_session_global_grid", "includes_current_video_frame": True, "oldest_visual_tick_replaced_by_current": True, "timestamp_note": ( "historical video targets are quantized causally to the " "global grid for decoder/cache reuse; the final video sample " "is exactly current time while audio retains the full " "half-open working interval" ), } def _manifest_session_specs( manifest_path: Path, session_name: str, ) -> tuple[list[SourceClipSpec], dict[str, Any], Path | None]: payload = json.loads(manifest_path.read_text()) matches = [ session for session in payload.get("sessions", []) if session.get("session") == session_name ] if len(matches) != 1: raise RuntimeError( f"Expected exactly one {session_name!r} in {manifest_path}, " f"found {len(matches)}" ) session = matches[0] specs = [] for record in session.get("clips", []): duration = record.get("video_duration_seconds") crop_record = record.get("egospeak_crop") crop = None if crop_record is not None: crop = ( int(crop_record["x"]), int(crop_record["y"]), int(crop_record["width"]), int(crop_record["height"]), ) specs.append( SourceClipSpec( path=Path(record["path"]), duration=float(duration) if duration is not None else None, crop=crop, ) ) if not specs: raise RuntimeError(f"No source clips recorded for {session_name}") target_record = session.get("target") or {} target_path = Path(target_record["path"]) if target_record.get("path") else None return specs, session, target_path def _easycom_directory_specs(session_dir: Path) -> list[SourceClipSpec]: session_dir = session_dir.expanduser().resolve() if not session_dir.is_dir(): raise FileNotFoundError(f"EasyCom session directory does not exist: {session_dir}") paths = sorted(session_dir.glob("*.mp4"), key=lambda path: parse_timecode(path.stem)) if not paths: raise RuntimeError(f"No MP4 clips found in {session_dir}") return [ SourceClipSpec(path=path, crop_rightmost_width=EASYCOM_WIDTH) for path in paths ] def _parse_crop(value: str | None) -> tuple[int, int, int, int] | None: if value is None: return None try: values = tuple(int(item) for item in value.split(",")) except ValueError as exc: raise argparse.ArgumentTypeError("crop must be x,y,width,height") from exc if len(values) != 4: raise argparse.ArgumentTypeError("crop must be x,y,width,height") return values def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) source = parser.add_mutually_exclusive_group(required=True) source.add_argument("--video", type=Path, help="One original video clip") source.add_argument( "--easycom-session-dir", type=Path, help="Original EasyCom Session_N directory containing retained MP4 clips", ) source.add_argument( "--easycom-manifest", type=Path, help="easycom_manifest.py JSON; also requires --session", ) parser.add_argument("--session", help="Session_N selected from --easycom-manifest") parser.add_argument( "--audio", type=Path, help="Optional separate audio for --video; defaults to its embedded track", ) parser.add_argument( "--crop", help="Optional crop for --video as x,y,width,height", ) parser.add_argument( "--use-audio-in-video", action="store_true", help="Accepted for command compatibility; direct extraction always uses AV interleaving", ) parser.add_argument( "--target", type=Path, help="target_perframe [T,3]; T is the authoritative full duration", ) parser.add_argument("--output", type=Path) parser.add_argument( "--mode", choices=( "causal-prefix", "open-stream-replay", "overlap-working-memory", ), default="open-stream-replay", ) parser.add_argument("--model-id", default=core.MODEL_ID) parser.add_argument("--start", type=float, default=0.0) parser.add_argument("--end", type=float) parser.add_argument("--step-seconds", type=float, default=0.2) parser.add_argument("--native-block-seconds", type=float, default=2.0) parser.add_argument( "--working-memory-seconds", type=float, default=2.0, help=( "Sliding recent recap used by overlap-working-memory; its duration " "is included in --memory-seconds rather than added on top" ), ) parser.add_argument( "--memory-seconds", type=float, default=0.0, help="0 keeps Thinker KV for the full extraction; positive values use a rolling window", ) parser.add_argument("--video-fps", type=float, default=2.0) parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28) parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28) parser.add_argument("--frame-cache-size", type=int, default=8) parser.add_argument( "--audio-feature-device", default=None, help=( "Whisper log-Mel device for overlap-working-memory, for example " "model, cuda:0, or cpu; omitted preserves processor defaults" ), ) parser.add_argument( "--audio-right-padding-samples", type=int, default=None, help=( "For overlap-working-memory, bound processor audio padding to " "actual samples plus this many zeros; omitted uses model maximum" ), ) parser.add_argument( "--decoder", choices=("decord", "torchcodec", "pynvcodec"), default="decord", help="Video pixel decoder; torchcodec is CPU and pynvcodec uses NVDEC", ) parser.add_argument( "--decoder-gpu-id", type=int, default=0, help="GPU visible to this process used by PyNvVideoCodec", ) parser.add_argument( "--decoder-threads", type=int, default=1, help="CPU decode threads per Decord or TorchCodec reader", ) parser.add_argument( "--max-open-readers", type=int, default=2, help="Maximum simultaneously open Decord or TorchCodec readers", ) parser.add_argument("--verify-every", type=int, default=0) parser.add_argument("--verify-continuous-every", type=int, default=0) parser.add_argument("--no-progress", action="store_true") parser.add_argument("--head-checkpoint", type=Path) parser.add_argument( "--dtype", choices=("auto", "bfloat16", "float16", "float32"), default="bfloat16", ) parser.add_argument("--flash-attention-2", action="store_true") parser.add_argument("--device", default="cuda:0") parser.add_argument( "--media-only", action="store_true", help="Validate/index direct media and exit before loading Qwen", ) return parser.parse_args(argv) def _resolve_duration_and_target( args: argparse.Namespace, manifest_target: Path | None, ) -> tuple[Path | None, int | None, float]: target_path = args.target if target_path is None and manifest_target is not None and manifest_target.is_file(): target_path = manifest_target target_steps = None if target_path is not None: target_path = target_path.expanduser().resolve() target_steps = core.target_num_steps(target_path) if not math.isclose(args.start, 0.0, rel_tol=0.0, abs_tol=1e-9): raise SystemExit("Full target extraction requires --start 0") target_end = target_steps * args.step_seconds if args.end is not None and not math.isclose( args.end, target_end, rel_tol=0.0, abs_tol=1e-6 ): raise SystemExit( f"--end {args.end} conflicts with target length {target_steps} " f"at {args.step_seconds}s ({target_end}s)" ) end = target_end else: if args.end is None: raise SystemExit("Pass --target for a full run or --end for a bounded test") end = float(args.end) if end <= args.start: raise SystemExit("--end must be greater than --start") return target_path, target_steps, end def _effective_memory_seconds( requested: float, *, duration: float, block_seconds: float, ) -> float: if requested < 0: raise SystemExit("--memory-seconds must be >= 0") if requested > 0: return requested blocks = max(1, int(math.ceil(duration / block_seconds - EPSILON))) return blocks * block_seconds def _print_media_summary(reader: CausalSourceAV, start: float, end: float) -> None: print( f"Direct media valid: {len(reader.clips)} source clip(s), " f"required={end - start:.3f}s, grid={reader.expected_grid_frames} " f"frames at {reader.fps:g} FPS, audio={len(reader.audio)} samples" ) candidates = [0] if reader.expected_grid_frames > 1: candidates.append(reader.expected_grid_frames - 1) for grid_index in candidates: mapping = reader.sample_mapping(grid_index) print( " grid " f"{mapping['grid_time']:.3f}s -> " f"{Path(mapping['source_clip']).name} " f"frame={mapping['source_frame_index']} " f"pts={mapping['source_frame_pts_seconds']:.6f}s" ) def main( argv: Sequence[str] | None = None, *, model: Any | None = None, processor: Any | None = None, ) -> None: """Run one direct extraction, optionally reusing loaded Qwen objects. ``model`` and ``processor`` must be supplied together. Reusing them is safe across independent clips because this function constructs a new media reader, extractor, and temporal cache on every call. """ args = parse_args(argv) core._configure_qwen_logging() if (model is None) != (processor is None): raise ValueError("model and processor must be supplied together") if args.verify_every < 0 or args.verify_continuous_every < 0: raise SystemExit("verification intervals must be >= 0") if args.working_memory_seconds <= 0: raise SystemExit("--working-memory-seconds must be positive") if ( args.mode == "overlap-working-memory" and args.verify_continuous_every > 0 ): raise SystemExit( "--verify-continuous-every is not defined for the intentional " "overlap representation; use --verify-every" ) if args.frame_cache_size < 0: raise SystemExit("--frame-cache-size must be >= 0") if ( args.audio_right_padding_samples is not None and args.audio_right_padding_samples < 0 ): raise SystemExit("--audio-right-padding-samples must be non-negative") if args.mode != "overlap-working-memory" and ( args.audio_feature_device is not None or args.audio_right_padding_samples is not None ): raise SystemExit( "Audio preprocessing flags are currently implemented only for " "--mode overlap-working-memory" ) if args.output is None and not args.media_only: raise SystemExit("--output is required unless --media-only is used") if args.audio is not None and args.video is None: raise SystemExit("--audio is supported only together with --video") if args.session is not None and args.easycom_manifest is None: raise SystemExit("--session is used only with --easycom-manifest") if args.easycom_manifest is not None and args.session is None: raise SystemExit("--easycom-manifest requires --session Session_N") manifest_record = None manifest_target = None if args.video is not None: specs = [ SourceClipSpec( path=args.video, crop=_parse_crop(args.crop), audio_path=args.audio, ) ] source_kind = "single_original_video" elif args.easycom_session_dir is not None: if args.crop is not None: raise SystemExit("EasyCom direct extraction uses the fixed rightmost-1920 crop") specs = _easycom_directory_specs(args.easycom_session_dir) source_kind = "easycom_dense_original_clips" else: if args.crop is not None: raise SystemExit("EasyCom manifest supplies its crop metadata") specs, manifest_record, manifest_target = _manifest_session_specs( args.easycom_manifest.expanduser().resolve(), args.session ) source_kind = "easycom_manifest_original_clips" target_path, target_steps, end = _resolve_duration_and_target( args, manifest_target ) required_duration = end reader = CausalSourceAV( specs, output_fps=args.video_fps, required_duration=required_duration, frame_cache_size=args.frame_cache_size, decoder_backend=args.decoder, decoder_gpu_id=args.decoder_gpu_id, decoder_threads=args.decoder_threads, max_open_readers=args.max_open_readers, ) is_easycom = ( args.easycom_session_dir is not None or args.easycom_manifest is not None ) if is_easycom: for clip in reader.clips: if clip.height != EASYCOM_HEIGHT: raise RuntimeError( f"Expected EasyCom height {EASYCOM_HEIGHT}, got {clip.height}: " f"{clip.path}" ) if not math.isclose( clip.source_fps, 20.0, rel_tol=0.0, abs_tol=1e-6 ): raise RuntimeError( f"Expected EasyCom source FPS 20, got {clip.source_fps:g}: " f"{clip.path}" ) if target_steps is not None: source_frames = sum(clip.num_frames for clip in reader.clips) media_steps = source_frames // 4 if media_steps != target_steps: raise RuntimeError( f"EasyCom dense source gives {media_steps} complete 5-Hz " f"steps from {source_frames} frames, but target has " f"{target_steps}" ) _print_media_summary(reader, args.start, end) if args.media_only: # Exercise the selected pixel decoder as well as its metadata path. probe_end = min(end, args.start + max(args.step_seconds, 1.0 / reader.fps)) reader.frames(args.start, probe_end) print(f"Decoded media probe successfully with {args.decoder}") return duration = end - args.start if args.mode == "overlap-working-memory" and args.memory_seconds <= 0: raise SystemExit( "overlap-working-memory uses a matched finite presentation budget; " "pass a positive --memory-seconds (recommended: 8)" ) effective_memory = _effective_memory_seconds( args.memory_seconds, duration=duration, block_seconds=args.native_block_seconds, ) if args.memory_seconds == 0 and duration > 120: print( "WARNING: persistent Thinker KV grows for the entire media " f"duration ({duration:.1f}s). Long EasyCom sessions can exceed " "the model context or GPU memory; a positive --memory-seconds is " "a bounded-memory experiment rather than native persistent-KV " "semantics.", flush=True, ) config = core.StreamConfig( step_seconds=args.step_seconds, native_block_seconds=args.native_block_seconds, memory_seconds=effective_memory, video_fps=args.video_fps, video_min_pixels=args.video_min_pixels, video_max_pixels=args.video_max_pixels, ) config.validate() if model is None: model, processor = core.load_qwen_omni( args.model_id, dtype=args.dtype, device_map={"": args.device}, use_flash_attention_2=args.flash_attention_2, ) else: model.eval() model.requires_grad_(False) media_path_for_api = reader.clips[0].path audio_feature_device: str | torch.device | None = args.audio_feature_device if audio_feature_device == "model": audio_feature_device = model.device if args.mode == "overlap-working-memory": print( "Overlap audio preprocessing: " f"feature_device={audio_feature_device or 'processor default'}, " "padding=" + ( "model maximum" if args.audio_right_padding_samples is None else ( "actual samples + >= " f"{args.audio_right_padding_samples} right-padding samples" ) ), flush=True, ) overlap_extractor = None working_grid_av = None if args.mode == "causal-prefix": extractor = core.CausalPrefixExtractor( model, processor, config, fixed_grid_av=reader ) result = extractor.extract( video_path=media_path_for_api, audio_path=None, start=args.start, end=end, use_audio_in_video=True, show_progress=not args.no_progress, ) elif args.mode == "open-stream-replay": extractor = core.OpenStreamReplayExtractor( model, processor, config, fixed_grid_av=reader ) result = extractor.extract( video_path=media_path_for_api, audio_path=None, start=args.start, end=end, use_audio_in_video=True, verify_every=args.verify_every, verify_continuous_every=args.verify_continuous_every, show_progress=not args.no_progress, ) else: working_grid_av = GlobalHistoryCurrentAV( reader, max_interval_seconds=args.working_memory_seconds, ) overlap_extractor = core.OverlappingWorkingMemoryExtractor( model, processor, config, working_memory_seconds=args.working_memory_seconds, fixed_grid_av=reader, working_grid_av=working_grid_av, audio_feature_device=audio_feature_device, audio_right_padding_samples=args.audio_right_padding_samples, ) result = overlap_extractor.extract( video_path=media_path_for_api, audio_path=None, start=args.start, end=end, use_audio_in_video=True, verify_every=args.verify_every, show_progress=not args.no_progress, ) if target_steps is not None and len(result.times) != target_steps: raise RuntimeError( f"Extractor returned {len(result.times)} states, expected {target_steps}" ) persistent_memory = ( args.mode == "open-stream-replay" and args.memory_seconds == 0 ) payload: dict[str, Any] = { "model_id": args.model_id, "mode": args.mode, "stream_representation": ( "single_causal_prefix" if args.mode == "causal-prefix" else ( "single_open_media_turn_with_2s_payload_commits" if args.mode == "open-stream-replay" else ( "serialized_cached_complete_blocks_then_overlapping_" "global_history_current_frame_recent_recap" ) ) ), "config": asdict(config), "requested_memory_seconds": args.memory_seconds, "persistent_thinker_kv": persistent_memory, "start": args.start, "end": end, "num_steps": len(result.times), "label_indices": torch.arange(len(result.times), dtype=torch.long), "valid_10_horizon_steps": max(0, len(result.times) - 10), "target_path": str(target_path) if target_path is not None else None, "source_kind": source_kind, "source_manifest_record": manifest_record, "direct_media": reader.metadata(), "video_sampling_grid": ( "fixed_global_half_open_plus_global_history_current_frame_recap" if args.mode == "overlap-working-memory" else "direct_pts_causal_global_grid_half_open_with_tail_duplication" ), "open_stream_policy": ( { "commit": "completed_2s_av_payload_only", "temporary": "partial_payload_plus_media_closure_plus_readout", "tmrope_positions_per_chunk": core._tmrope_positions_per_chunk(model), "memory": ( "persistent_full_extraction_thinker_kv" if persistent_memory else "rolling_complete_blocks_plus_current_partial_block" ), "rolling_replay_inputs": ( None if persistent_memory else "reuse_cached_complete_block_av_embeddings" ), } if args.mode == "open-stream-replay" else None ), "overlap_working_memory_policy": ( { "total_presentation_seconds": config.memory_seconds, "fixed_history_seconds": ( config.memory_seconds - args.working_memory_seconds ), "fixed_block_seconds": config.native_block_seconds, "working_memory_seconds": args.working_memory_seconds, "fixed_blocks": ( overlap_extractor.fixed_block_budget if overlap_extractor is not None else None ), "serialization": ( "fixed complete AV payloads followed by overlapping " "working AV payload in the same open media turn" ), "overlap_is_intentional": True, "working_media": ( working_grid_av.metadata() if working_grid_av is not None else None ), "audio_feature_device": ( str(audio_feature_device) if audio_feature_device is not None else "processor_default" ), "audio_right_padding_samples": args.audio_right_padding_samples, "layout": ( overlap_extractor.layout_records if overlap_extractor is not None else None ), "backbone": "frozen", } if args.mode == "overlap-working-memory" else None ), "times": result.times, "states": result.states, "verification": result.verification, } if args.head_checkpoint is not None: head = core._load_head( args.head_checkpoint, core._hidden_size(model), model.device ) with torch.inference_mode(): logits = head( result.states.to(model.device, dtype=model.dtype) ).float().cpu() payload["logits"] = logits payload["probabilities"] = logits.softmax(dim=-1) args.output.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, args.output) print(f"Saved {len(result.times)} causal states to {args.output}") print(f"states: {tuple(result.states.shape)}") if result.verification: kinds = sorted({item.get("kind", "unspecified") for item in result.verification}) print(f"verification checks: {len(result.verification)}") for kind in kinds: checks = [item for item in result.verification if item.get("kind") == kind] print( f" {kind}: n={len(checks)}, " f"worst max_abs={max(item['max_abs'] for item in checks):.6g}, " f"lowest cosine={min(item['cosine'] for item in checks):.8f}" ) if __name__ == "__main__": main()