1652 lines
64 KiB
Python
1652 lines
64 KiB
Python
#!/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="<f4")
|
|
expected = int(round(duration * SAMPLE_RATE))
|
|
output = np.zeros(expected, dtype=np.float32)
|
|
used = min(expected, len(decoded))
|
|
if used:
|
|
output[:used] = decoded[:used]
|
|
return output
|
|
|
|
|
|
class CausalSourceAV:
|
|
"""Direct PTS-aware reader implementing the existing extractor AV API."""
|
|
|
|
sample_rate = SAMPLE_RATE
|
|
|
|
def __init__(
|
|
self,
|
|
specs: list[SourceClipSpec],
|
|
*,
|
|
output_fps: float,
|
|
required_duration: float,
|
|
frame_cache_size: int = 8,
|
|
decoder_backend: str = "decord",
|
|
decoder_gpu_id: int = 0,
|
|
decoder_threads: int = 1,
|
|
max_open_readers: int = 2,
|
|
) -> 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",
|
|
"completed-block-hold",
|
|
"micro-chunk-200ms",
|
|
"overlap-working-memory",
|
|
),
|
|
default="open-stream-replay",
|
|
)
|
|
parser.add_argument("--model-id", default=core.MODEL_ID)
|
|
parser.add_argument(
|
|
"--backbone-loader",
|
|
choices=("qwen", "roma"),
|
|
default="qwen",
|
|
help=(
|
|
"qwen loads a Thinker checkpoint directly; roma loads the full "
|
|
"ROMA checkpoint and keeps only its frozen Thinker"
|
|
),
|
|
)
|
|
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(
|
|
"--readout-last-n-layers",
|
|
type=int,
|
|
default=1,
|
|
help=(
|
|
"Uniformly average the final N Thinker hidden layers at the "
|
|
"causal readout token (default: 1 preserves the original result)."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--readout-layer-fusion",
|
|
choices=("mean", "stack"),
|
|
default="mean",
|
|
help=(
|
|
"Fuse selected query layers by a fixed mean, or preserve [N,D] "
|
|
"for a trainable downstream layer mixture."
|
|
),
|
|
)
|
|
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,
|
|
readout_last_n_layers=args.readout_last_n_layers,
|
|
readout_layer_fusion=args.readout_layer_fusion,
|
|
)
|
|
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,
|
|
backbone_loader=args.backbone_loader,
|
|
)
|
|
else:
|
|
loaded_backbone = core.backbone_metadata(model).get("loader")
|
|
if loaded_backbone not in {"unknown", args.backbone_loader}:
|
|
raise RuntimeError(
|
|
"Preloaded backbone does not match CLI provenance: "
|
|
f"model={loaded_backbone}, requested={args.backbone_loader}"
|
|
)
|
|
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 in {
|
|
"open-stream-replay",
|
|
"completed-block-hold",
|
|
"micro-chunk-200ms",
|
|
}:
|
|
extractor_class = {
|
|
"open-stream-replay": core.OpenStreamReplayExtractor,
|
|
"completed-block-hold": core.CompletedBlockHoldExtractor,
|
|
"micro-chunk-200ms": core.MicroChunkReplayExtractor,
|
|
}[args.mode]
|
|
extractor = extractor_class(
|
|
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
|
|
in {
|
|
"open-stream-replay",
|
|
"completed-block-hold",
|
|
"micro-chunk-200ms",
|
|
}
|
|
and args.memory_seconds == 0
|
|
)
|
|
payload: dict[str, Any] = {
|
|
"model_id": args.model_id,
|
|
"backbone_loader": args.backbone_loader,
|
|
"backbone_metadata": core.backbone_metadata(model),
|
|
"mode": args.mode,
|
|
"stream_representation": {
|
|
"causal-prefix": "single_causal_prefix",
|
|
"open-stream-replay": (
|
|
"native_2s_completed_blocks_plus_temporary_partial_tail"
|
|
),
|
|
"completed-block-hold": (
|
|
"native_2s_completed_blocks_forward_held_to_5hz"
|
|
),
|
|
"micro-chunk-200ms": "non_native_200ms_payload_commits",
|
|
"overlap-working-memory": (
|
|
"serialized_cached_complete_blocks_then_overlapping_"
|
|
"global_history_current_frame_recent_recap"
|
|
),
|
|
}[args.mode],
|
|
"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
|
|
),
|
|
"chunk_ablation_policy": (
|
|
{
|
|
"kind": "native_completed_block_hold",
|
|
"native_chunk_seconds": config.native_block_seconds,
|
|
"output_step_seconds": config.step_seconds,
|
|
"tail_policy": "discard_incomplete_tail_and_hold_last_state",
|
|
"initial_state": "zeros_until_first_complete_native_block",
|
|
"fresh_at_every_5hz_step": False,
|
|
"memory": (
|
|
"persistent_full_extraction_thinker_kv"
|
|
if persistent_memory
|
|
else "rolling_native_completed_blocks"
|
|
),
|
|
"backbone": "frozen",
|
|
}
|
|
if args.mode == "completed-block-hold"
|
|
else (
|
|
{
|
|
"kind": "non_native_micro_chunk_stress_test",
|
|
"micro_chunk_seconds": config.step_seconds,
|
|
"native_chunk_seconds": config.native_block_seconds,
|
|
"tmrope_positions_per_micro_step": (
|
|
core._tmrope_positions_per_step(
|
|
model, config.step_seconds
|
|
)
|
|
),
|
|
"tail_policy": "commit_each_5hz_interval_as_a_payload",
|
|
"fresh_at_every_5hz_step": True,
|
|
"memory": (
|
|
"persistent_full_extraction_thinker_kv"
|
|
if persistent_memory
|
|
else "rolling_micro_payloads_on_native_boundaries"
|
|
),
|
|
"backbone": "frozen",
|
|
}
|
|
if args.mode == "micro-chunk-200ms"
|
|
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,
|
|
"extraction_diagnostics": result.diagnostics,
|
|
}
|
|
if args.head_checkpoint is not None:
|
|
head, head_loss_name = 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"] = core.marginal_probabilities_from_logits(
|
|
logits,
|
|
head_loss_name,
|
|
)
|
|
payload["head_loss"] = head_loss_name
|
|
|
|
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()
|