2372 lines
86 KiB
Python
2372 lines
86 KiB
Python
"""Omni-only causal turn-taking features at 5 Hz.
|
|
|
|
This module implements three inference paths for Qwen2.5-Omni-3B:
|
|
|
|
1. ``causal-prefix``: recompute a causal AV window ending at every 200 ms.
|
|
This is deliberately slow, but is the simplest correctness baseline.
|
|
2. ``open-stream-replay``: cache completed AV blocks inside one
|
|
still-open user/media turn; closure and readout tokens remain temporary.
|
|
3. ``overlap-working-memory``: a frozen ablation that serializes cached
|
|
complete blocks and an overlapping ``[t-2,t]`` recent recap.
|
|
|
|
Neither path uses an RNN. A small head maps one Thinker state to the same
|
|
``10 future horizons x 3 classes`` output used by EgoSpeak.
|
|
|
|
The open-stream replay path touches version-specific Hugging Face internals.
|
|
It is therefore pinned to Transformers 4.52.3 and includes an optional
|
|
exact-replay check. Run that check before extracting a full dataset.
|
|
|
|
Example (no command is run merely by importing this file):
|
|
|
|
python streaming_turn_model.py \
|
|
--video /path/to/session.mp4 \
|
|
--use-audio-in-video \
|
|
--mode open-stream-replay \
|
|
--memory-seconds 20 \
|
|
--verify-every 25 \
|
|
--output features/session.pt
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import copy
|
|
import json
|
|
import logging
|
|
import math
|
|
import subprocess
|
|
from collections import OrderedDict
|
|
from dataclasses import asdict, dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Literal
|
|
|
|
import torch
|
|
from easycom_head import (
|
|
TurnPredictionHead,
|
|
egospeak_multihorizon_loss,
|
|
miniroad_multihorizon_loss,
|
|
multihorizon_cross_entropy,
|
|
)
|
|
from qwen_omni_utils import process_mm_info
|
|
from torch import nn
|
|
from tqdm.auto import tqdm
|
|
from transformers import (
|
|
Qwen2_5OmniProcessor,
|
|
Qwen2_5OmniThinkerForConditionalGeneration,
|
|
)
|
|
from transformers.cache_utils import DynamicCache
|
|
|
|
Mode = Literal[
|
|
"causal-prefix",
|
|
"open-stream-replay",
|
|
"overlap-working-memory",
|
|
]
|
|
MODEL_ID = "Qwen/Qwen2.5-Omni-3B"
|
|
_KNOWN_QWEN_WARNING_TEXT = (
|
|
"System prompt modified, audio output may not work",
|
|
"Unused or unrecognized kwargs: images.",
|
|
"Unused or unrecognized kwargs: return_tensors, images.",
|
|
"Unused or unrecognized kwargs: images, return_tensors.",
|
|
)
|
|
|
|
|
|
class _KnownQwenWarningFilter(logging.Filter):
|
|
"""Hide Qwen's speech-output warning in Thinker-only extraction."""
|
|
|
|
def filter(self, record: logging.LogRecord) -> bool:
|
|
message = record.getMessage()
|
|
return not any(text in message for text in _KNOWN_QWEN_WARNING_TEXT)
|
|
|
|
|
|
def _configure_qwen_logging() -> None:
|
|
# Qwen2.5-Omni 4.52.3 emits this through Python's root logger whenever the
|
|
# conversation does not use its speech-generation system prompt. This
|
|
# extractor uses only the Thinker and never requests audio output.
|
|
def add_filter(target: logging.Filterer) -> None:
|
|
if not any(
|
|
isinstance(item, _KnownQwenWarningFilter)
|
|
for item in target.filters
|
|
):
|
|
target.addFilter(_KnownQwenWarningFilter())
|
|
|
|
# The system-prompt warning is logged directly on root. The unused-kwargs
|
|
# warning is emitted by a Transformers child logger and handled by the
|
|
# library's top-level StreamHandler, so cover both paths.
|
|
for logger_name in (None, "transformers"):
|
|
logger = logging.getLogger(logger_name)
|
|
add_filter(logger)
|
|
for handler in logger.handlers:
|
|
add_filter(handler)
|
|
|
|
|
|
def _choose_dtype(
|
|
dtype: Literal["auto", "bfloat16", "float16", "float32"],
|
|
) -> str | torch.dtype:
|
|
if dtype == "auto":
|
|
return "auto"
|
|
return getattr(torch, dtype)
|
|
|
|
|
|
def load_qwen_omni(
|
|
model_id: str = MODEL_ID,
|
|
*,
|
|
dtype: Literal["auto", "bfloat16", "float16", "float32"] = "auto",
|
|
device_map: Any = "auto",
|
|
use_flash_attention_2: bool = False,
|
|
) -> tuple[Qwen2_5OmniThinkerForConditionalGeneration, Qwen2_5OmniProcessor]:
|
|
"""Load only Qwen's multimodal Thinker and its processor."""
|
|
|
|
model_kwargs: dict[str, Any] = {
|
|
"torch_dtype": _choose_dtype(dtype),
|
|
"device_map": device_map,
|
|
}
|
|
if use_flash_attention_2:
|
|
model_kwargs["attn_implementation"] = "flash_attention_2"
|
|
model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
|
|
model_id, **model_kwargs
|
|
)
|
|
processor = Qwen2_5OmniProcessor.from_pretrained(model_id)
|
|
model.eval()
|
|
return model, processor
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StreamConfig:
|
|
"""Temporal and visual settings for one extraction run."""
|
|
|
|
step_seconds: float = 0.2
|
|
native_block_seconds: float = 2.0
|
|
memory_seconds: float = 20.0
|
|
video_fps: float = 2.0
|
|
video_min_pixels: int = 64 * 28 * 28
|
|
video_max_pixels: int = 128 * 28 * 28
|
|
readout_last_n_layers: int = 1
|
|
|
|
def validate(self) -> None:
|
|
if self.step_seconds <= 0:
|
|
raise ValueError("step_seconds must be positive")
|
|
if self.native_block_seconds <= 0:
|
|
raise ValueError("native_block_seconds must be positive")
|
|
if self.memory_seconds < self.native_block_seconds:
|
|
raise ValueError("memory_seconds must be at least one native block")
|
|
if self.video_fps <= 0:
|
|
raise ValueError("video_fps must be positive")
|
|
if self.video_min_pixels <= 0:
|
|
raise ValueError("video_min_pixels must be positive")
|
|
if self.video_max_pixels < self.video_min_pixels:
|
|
raise ValueError("video_max_pixels must be >= video_min_pixels")
|
|
if self.readout_last_n_layers <= 0:
|
|
raise ValueError("readout_last_n_layers must be positive")
|
|
_require_integer_ratio(
|
|
self.native_block_seconds, self.step_seconds, "native block / step"
|
|
)
|
|
_require_integer_ratio(
|
|
self.memory_seconds, self.native_block_seconds, "memory / native block"
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class ExtractionResult:
|
|
"""One fused Thinker state for every causal prediction time."""
|
|
|
|
times: torch.Tensor # [N], seconds in the original session
|
|
states: torch.Tensor # [N, D]
|
|
verification: list[dict[str, Any]]
|
|
|
|
|
|
@dataclass
|
|
class OpenMediaPieces:
|
|
"""One prepared AV interval split around its interleaved media payload."""
|
|
|
|
prefix_embeddings: torch.Tensor
|
|
prefix_position_ids: torch.LongTensor
|
|
payload_embeddings: torch.Tensor
|
|
payload_position_ids: torch.LongTensor
|
|
closure_ids: torch.LongTensor
|
|
closure_position_ids: torch.LongTensor
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CachedOpenMediaBlock:
|
|
"""Prepared model inputs for one complete source-aligned AV block."""
|
|
|
|
start: float
|
|
end: float
|
|
pieces: OpenMediaPieces
|
|
|
|
|
|
class CompletedBlockPieceCache:
|
|
"""Bounded ring of complete blocks reusable during strict-window replay.
|
|
|
|
The modality encoders operate on each native block independently of the
|
|
Thinker history. Their embeddings and local TMRoPE coordinates can
|
|
therefore be reused when the strict rolling window rebuilds the Thinker
|
|
KV. Only the Thinker replay must be recomputed.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
origin: float,
|
|
block_seconds: float,
|
|
max_blocks: int,
|
|
) -> None:
|
|
if block_seconds <= 0:
|
|
raise ValueError("block_seconds must be positive")
|
|
if max_blocks < 0:
|
|
raise ValueError("max_blocks must be >= 0")
|
|
self.origin = origin
|
|
self.block_seconds = block_seconds
|
|
self.max_blocks = max_blocks
|
|
self._blocks: OrderedDict[int, CachedOpenMediaBlock] = OrderedDict()
|
|
|
|
def _index(self, start: float) -> int:
|
|
ratio = (start - self.origin) / self.block_seconds
|
|
index = round(ratio)
|
|
if not math.isclose(ratio, index, rel_tol=0.0, abs_tol=1e-6):
|
|
raise ValueError(
|
|
f"Block start {start:g} is not aligned to origin "
|
|
f"{self.origin:g} at {self.block_seconds:g}s"
|
|
)
|
|
return index
|
|
|
|
def put(self, block: CachedOpenMediaBlock) -> None:
|
|
duration = block.end - block.start
|
|
if not math.isclose(
|
|
duration,
|
|
self.block_seconds,
|
|
rel_tol=0.0,
|
|
abs_tol=1e-6,
|
|
):
|
|
raise ValueError(
|
|
f"Cached block duration is {duration:g}s, expected "
|
|
f"{self.block_seconds:g}s"
|
|
)
|
|
if self.max_blocks == 0:
|
|
return
|
|
index = self._index(block.start)
|
|
self._blocks[index] = block
|
|
self._blocks.move_to_end(index)
|
|
while len(self._blocks) > self.max_blocks:
|
|
self._blocks.popitem(last=False)
|
|
|
|
def get_many(
|
|
self,
|
|
intervals: list[tuple[float, float]],
|
|
) -> list[CachedOpenMediaBlock]:
|
|
blocks = []
|
|
for start, end in intervals:
|
|
index = self._index(start)
|
|
try:
|
|
block = self._blocks[index]
|
|
except KeyError as exc:
|
|
available = list(self._blocks)
|
|
raise RuntimeError(
|
|
"Completed-block embedding cache miss for interval "
|
|
f"[{start:g}, {end:g}); available block indices="
|
|
f"{available}"
|
|
) from exc
|
|
if not (
|
|
math.isclose(block.start, start, rel_tol=0.0, abs_tol=1e-6)
|
|
and math.isclose(block.end, end, rel_tol=0.0, abs_tol=1e-6)
|
|
):
|
|
raise RuntimeError(
|
|
"Completed-block embedding cache interval mismatch: "
|
|
f"requested [{start:g}, {end:g}), cached "
|
|
f"[{block.start:g}, {block.end:g})"
|
|
)
|
|
blocks.append(block)
|
|
return blocks
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._blocks)
|
|
|
|
|
|
class FixedGridAV:
|
|
"""Decode only frames on a prepared MP4's fixed global grid.
|
|
|
|
Qwen's normal path cannot represent a causal 2 FPS prefix shorter than
|
|
0.5 seconds because its video tower requires an even number of frames.
|
|
Explicit frame lists let the utility duplicate the latest past frame,
|
|
never a future frame. Audio is loaded once and sliced exactly at 16 kHz.
|
|
"""
|
|
|
|
sample_rate = 16_000
|
|
|
|
def __init__(self, path: str | Path, expected_fps: float) -> None:
|
|
try:
|
|
import decord
|
|
import numpy as np
|
|
except ImportError as exc:
|
|
raise RuntimeError(
|
|
"Fixed-grid AV extraction requires decord and numpy"
|
|
) from exc
|
|
|
|
self.path = Path(path).resolve()
|
|
self.reader = decord.VideoReader(str(self.path))
|
|
self.fps = float(self.reader.get_avg_fps())
|
|
if not math.isclose(self.fps, expected_fps, rel_tol=0.0, abs_tol=1e-6):
|
|
raise ValueError(
|
|
f"Prepared video is {self.fps:g} FPS, expected {expected_fps:g}"
|
|
)
|
|
self.num_frames = len(self.reader)
|
|
# SoundFile cannot decode the AAC track in MP4, causing librosa to
|
|
# fall back to its deprecated and much slower audioread path. Decode
|
|
# directly to the exact float32/16-kHz/mono representation consumed by
|
|
# Qwen. Keeping the complete track in memory also makes repeated
|
|
# 200-ms slicing deterministic and inexpensive.
|
|
completed = subprocess.run(
|
|
[
|
|
"ffmpeg",
|
|
"-nostdin",
|
|
"-v",
|
|
"error",
|
|
"-i",
|
|
str(self.path),
|
|
"-map",
|
|
"0:a:0",
|
|
"-vn",
|
|
"-ac",
|
|
"1",
|
|
"-ar",
|
|
str(self.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 embedded audio from {self.path}: {message}"
|
|
)
|
|
self.audio = np.frombuffer(completed.stdout, dtype="<f4").copy()
|
|
if not len(self.audio):
|
|
raise RuntimeError(f"Decoded audio is empty: {self.path}")
|
|
|
|
def frames(self, start: float, end: float) -> list[Any]:
|
|
"""Return fixed-grid frames whose timestamps lie in ``[start, end)``.
|
|
|
|
Half-open ownership is essential at native block boundaries. For
|
|
example, the frame at exactly 2.0 s belongs to ``[2.0, 4.0)`` and must
|
|
not also be committed in ``[0.0, 2.0)``. Audio slicing below follows
|
|
the same convention.
|
|
"""
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
first = max(0, int(math.ceil(start * self.fps - 1e-7)))
|
|
last_exclusive = min(
|
|
self.num_frames,
|
|
int(math.ceil(end * self.fps - 1e-7)),
|
|
)
|
|
last = last_exclusive - 1
|
|
if first > last:
|
|
# A short interval can fall entirely between global visual ticks.
|
|
# Reuse the latest frame strictly before ``end``, never a future
|
|
# frame. Native 2-second blocks start on-grid and do not take this
|
|
# fallback path.
|
|
first = last = max(0, min(self.num_frames - 1, last))
|
|
indices = list(range(first, last + 1))
|
|
arrays = self.reader.get_batch(indices).asnumpy()
|
|
frames = [Image.fromarray(np.asarray(array)) for array in arrays]
|
|
if len(frames) % 2:
|
|
frames.append(frames[-1].copy())
|
|
return frames
|
|
|
|
def audio_slice(self, start: float, end: float) -> Any:
|
|
first = max(0, int(math.floor(start * self.sample_rate + 1e-7)))
|
|
last = min(
|
|
len(self.audio), int(math.floor(end * self.sample_rate + 1e-7))
|
|
)
|
|
if last <= first:
|
|
raise ValueError(f"Empty audio interval [{start}, {end}]")
|
|
return self.audio[first:last]
|
|
|
|
|
|
def _require_integer_ratio(numerator: float, denominator: float, name: str) -> int:
|
|
ratio = numerator / denominator
|
|
rounded = round(ratio)
|
|
if not math.isclose(ratio, rounded, rel_tol=0.0, abs_tol=1e-6):
|
|
raise ValueError(f"{name} must be an integer, got {ratio}")
|
|
return rounded
|
|
|
|
|
|
def _time_grid(start: float, end: float, step: float) -> list[float]:
|
|
"""Return start+step ... <=end without cumulative floating-point drift."""
|
|
|
|
count = int(math.floor((end - start) / step + 1e-7))
|
|
return [start + i * step for i in range(1, count + 1)]
|
|
|
|
|
|
def _media_content(
|
|
*,
|
|
video_path: str | Path | None,
|
|
audio_path: str | Path | None,
|
|
start: float,
|
|
end: float,
|
|
use_audio_in_video: bool,
|
|
config: StreamConfig,
|
|
) -> list[dict[str, Any]]:
|
|
if end <= start:
|
|
raise ValueError(f"Empty media interval [{start}, {end}]")
|
|
if video_path is None and audio_path is None:
|
|
raise ValueError("At least one of video_path and audio_path is required")
|
|
if use_audio_in_video and video_path is None:
|
|
raise ValueError("use_audio_in_video=True requires video_path")
|
|
if use_audio_in_video and audio_path is not None:
|
|
raise ValueError("Do not pass separate audio when use_audio_in_video=True")
|
|
|
|
content: list[dict[str, Any]] = []
|
|
if video_path is not None:
|
|
content.append(
|
|
{
|
|
"type": "video",
|
|
"video": str(Path(video_path).resolve()),
|
|
"video_start": start,
|
|
"video_end": end,
|
|
"fps": config.video_fps,
|
|
# Two frames are the minimum required by temporal_patch_size=2.
|
|
# For a short prefix both frames are sampled from <= end.
|
|
"min_frames": 2,
|
|
"min_pixels": config.video_min_pixels,
|
|
"max_pixels": config.video_max_pixels,
|
|
}
|
|
)
|
|
if audio_path is not None:
|
|
content.append(
|
|
{
|
|
"type": "audio",
|
|
"audio": str(Path(audio_path).resolve()),
|
|
"audio_start": start,
|
|
"audio_end": end,
|
|
}
|
|
)
|
|
return content
|
|
|
|
|
|
def build_chunk_conversation(
|
|
*,
|
|
video_path: str | Path | None,
|
|
audio_path: str | Path | None,
|
|
intervals: list[tuple[float, float]],
|
|
use_audio_in_video: bool,
|
|
config: StreamConfig,
|
|
fixed_grid_av: FixedGridAV | None = None,
|
|
) -> tuple[list[dict[str, Any]], list[Any] | None]:
|
|
"""Build one user turn per AV block, without an assistant query."""
|
|
|
|
if fixed_grid_av is not None:
|
|
if not use_audio_in_video or audio_path is not None:
|
|
raise ValueError("FixedGridAV requires embedded-audio AV mode")
|
|
conversation = []
|
|
audios = []
|
|
for start, end in intervals:
|
|
conversation.append(
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "video",
|
|
"video": fixed_grid_av.frames(start, end),
|
|
"fps": config.video_fps,
|
|
"min_pixels": config.video_min_pixels,
|
|
"max_pixels": config.video_max_pixels,
|
|
}
|
|
],
|
|
}
|
|
)
|
|
audios.append(fixed_grid_av.audio_slice(start, end))
|
|
return conversation, audios
|
|
|
|
conversation = [
|
|
{
|
|
"role": "user",
|
|
"content": _media_content(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
start=start,
|
|
end=end,
|
|
use_audio_in_video=use_audio_in_video,
|
|
config=config,
|
|
),
|
|
}
|
|
for start, end in intervals
|
|
]
|
|
return conversation, None
|
|
|
|
|
|
def _single_template_text(value: Any) -> str:
|
|
"""Normalize processor chat-template output across Transformers versions."""
|
|
|
|
if isinstance(value, str):
|
|
return value
|
|
if (
|
|
isinstance(value, (list, tuple))
|
|
and len(value) == 1
|
|
and isinstance(value[0], str)
|
|
):
|
|
return value[0]
|
|
raise TypeError(
|
|
"Expected one chat-template string, got "
|
|
f"{type(value).__name__}: {value!r}"
|
|
)
|
|
|
|
|
|
def _prepare_conversation(
|
|
processor: Qwen2_5OmniProcessor,
|
|
conversation: list[dict[str, Any]],
|
|
*,
|
|
use_audio_in_video: bool,
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
|
manual_audios: list[Any] | None = None,
|
|
video_fps: float | None = None,
|
|
text_override: str | None = None,
|
|
audio_feature_device: str | torch.device | None = None,
|
|
audio_right_padding_samples: int | None = None,
|
|
) -> dict[str, torch.Tensor]:
|
|
"""Decode/crop media and produce tensors for media turns only."""
|
|
|
|
text = text_override
|
|
if text is None:
|
|
text = _single_template_text(
|
|
processor.apply_chat_template(
|
|
conversation,
|
|
add_generation_prompt=False,
|
|
tokenize=False,
|
|
)
|
|
)
|
|
if manual_audios is None:
|
|
audios, images, videos = process_mm_info(
|
|
conversation,
|
|
use_audio_in_video=use_audio_in_video,
|
|
)
|
|
else:
|
|
# Avoid asking audio_process to open a list of explicit video frames as
|
|
# a media file. The processor still receives use_audio_in_video=True,
|
|
# preserving Qwen's AV token interleaving behavior.
|
|
_, images, videos = process_mm_info(
|
|
conversation,
|
|
use_audio_in_video=False,
|
|
)
|
|
audios = manual_audios
|
|
processor_kwargs: dict[str, Any] = {}
|
|
if videos is not None and video_fps is not None:
|
|
# Qwen2.5-Omni uses this value to compute video_second_per_grid for
|
|
# TMRoPE. This is essential for the 5 FPS ablation (default is 2 FPS).
|
|
processor_kwargs["fps"] = video_fps
|
|
processor_inputs: dict[str, Any] = {
|
|
"text": [text],
|
|
"audio": audios,
|
|
"videos": videos,
|
|
"padding": True,
|
|
"use_audio_in_video": use_audio_in_video,
|
|
**processor_kwargs,
|
|
}
|
|
audio_kwargs: dict[str, Any] = {}
|
|
if audio_feature_device is not None:
|
|
audio_kwargs["device"] = str(audio_feature_device)
|
|
if audio_right_padding_samples is not None:
|
|
if audio_right_padding_samples < 0:
|
|
raise ValueError("audio_right_padding_samples must be non-negative")
|
|
if audios is None or len(audios) != 1:
|
|
raise ValueError(
|
|
"Bounded audio padding currently requires exactly one audio input"
|
|
)
|
|
actual_samples = len(audios[0])
|
|
feature_extractor = processor.feature_extractor
|
|
hop_length = int(getattr(feature_extractor, "hop_length", 1))
|
|
if hop_length <= 0:
|
|
raise ValueError(f"Invalid audio hop_length: {hop_length}")
|
|
requested_length = actual_samples + audio_right_padding_samples
|
|
max_length = math.ceil(requested_length / hop_length) * hop_length
|
|
model_max_length = getattr(feature_extractor, "n_samples", None)
|
|
if model_max_length is not None:
|
|
max_length = min(max_length, int(model_max_length))
|
|
if max_length < actual_samples:
|
|
raise ValueError(
|
|
f"Bounded audio max_length {max_length} truncates "
|
|
f"{actual_samples} input samples"
|
|
)
|
|
audio_kwargs["max_length"] = max_length
|
|
if audio_kwargs:
|
|
# ProcessorMixin forwards nested modality kwargs to the underlying
|
|
# WhisperFeatureExtractor. This preserves Qwen's token expansion and
|
|
# feature_attention_mask construction while allowing CUDA STFT and a
|
|
# bounded amount of right-side silence.
|
|
processor_inputs["audio_kwargs"] = audio_kwargs
|
|
# Passing images=None and return_tensors through Qwen's 4.52.3 processor
|
|
# produces a warning on every streaming step. Omit the absent modality and
|
|
# perform the same BatchFeature tensor conversion explicitly afterward.
|
|
if images is not None and (
|
|
not isinstance(images, (list, tuple)) or len(images) > 0
|
|
):
|
|
processor_inputs["images"] = images
|
|
batch = processor(**processor_inputs)
|
|
batch = _processor_batch_convert_to_tensors(batch)
|
|
# BatchFeature.to(dtype) changes floating tensors only; IDs stay integral.
|
|
batch = _processor_batch_to_device(batch, model.device)
|
|
batch = _processor_batch_to_dtype(batch, model.dtype)
|
|
return dict(batch)
|
|
|
|
|
|
def _processor_batch_convert_to_tensors(batch: Any) -> Any:
|
|
"""Materialize the processor's NumPy/list outputs as PyTorch tensors."""
|
|
|
|
return batch.convert_to_tensors("pt")
|
|
|
|
|
|
def _processor_batch_to_device(batch: Any, device: torch.device) -> Any:
|
|
"""Move one processor batch to the model device."""
|
|
|
|
return batch.to(device)
|
|
|
|
|
|
def _processor_batch_to_dtype(batch: Any, dtype: torch.dtype) -> Any:
|
|
"""Cast only floating processor tensors to the model dtype."""
|
|
|
|
return batch.to(dtype)
|
|
|
|
|
|
def _readout_query_ids(
|
|
processor: Qwen2_5OmniProcessor,
|
|
device: torch.device,
|
|
) -> torch.LongTensor:
|
|
"""Derive the assistant-generation suffix from the installed chat template."""
|
|
|
|
probe = [{"role": "user", "content": [{"type": "text", "text": "."}]}]
|
|
without_query = _single_template_text(
|
|
processor.apply_chat_template(
|
|
probe, add_generation_prompt=False, tokenize=False
|
|
)
|
|
)
|
|
with_query = _single_template_text(
|
|
processor.apply_chat_template(
|
|
probe, add_generation_prompt=True, tokenize=False
|
|
)
|
|
)
|
|
if not with_query.startswith(without_query):
|
|
raise RuntimeError("The chat template does not expose an append-only assistant prompt")
|
|
suffix = with_query[len(without_query) :]
|
|
if not suffix:
|
|
raise RuntimeError("The chat template returned an empty assistant prompt")
|
|
return processor.tokenizer(
|
|
suffix,
|
|
add_special_tokens=False,
|
|
return_tensors="pt",
|
|
)["input_ids"].to(device)
|
|
|
|
|
|
def _media_turn_template_suffix(
|
|
processor: Qwen2_5OmniProcessor,
|
|
current_turn: dict[str, Any],
|
|
*,
|
|
turn_index: int,
|
|
) -> str:
|
|
"""Return only the template text appended by media turn ``turn_index``."""
|
|
|
|
if turn_index < 0:
|
|
raise ValueError("turn_index must be non-negative")
|
|
if turn_index == 0:
|
|
return _single_template_text(
|
|
processor.apply_chat_template(
|
|
[current_turn], add_generation_prompt=False, tokenize=False
|
|
)
|
|
)
|
|
dummy_turn = {
|
|
"role": "user",
|
|
"content": [{"type": "video", "video": "template-placeholder.mp4"}],
|
|
}
|
|
prefix_conversation = [dummy_turn for _ in range(turn_index)]
|
|
prefix_text = _single_template_text(
|
|
processor.apply_chat_template(
|
|
prefix_conversation,
|
|
add_generation_prompt=False,
|
|
tokenize=False,
|
|
)
|
|
)
|
|
full_text = _single_template_text(
|
|
processor.apply_chat_template(
|
|
[*prefix_conversation, current_turn],
|
|
add_generation_prompt=False,
|
|
tokenize=False,
|
|
)
|
|
)
|
|
if not full_text.startswith(prefix_text):
|
|
raise RuntimeError("Chat template is not append-only across media turns")
|
|
suffix = full_text[len(prefix_text) :]
|
|
if not suffix:
|
|
raise RuntimeError("Chat template produced an empty media-turn suffix")
|
|
return suffix
|
|
|
|
|
|
def _config_token_id(config: Any, name: str) -> int:
|
|
"""Support both *_token_id and *_token_index config spellings."""
|
|
|
|
for key in (f"{name}_token_id", f"{name}_token_index"):
|
|
value = getattr(config, key, None)
|
|
if value is not None:
|
|
return int(value)
|
|
raise AttributeError(f"Model config has no token ID for {name}")
|
|
|
|
|
|
def _masked_scatter_features(
|
|
embeddings: torch.Tensor,
|
|
input_ids: torch.LongTensor,
|
|
token_id: int,
|
|
features: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
mask = (input_ids == token_id).unsqueeze(-1).expand_as(embeddings)
|
|
expected = int(mask.sum().item())
|
|
actual = features.numel()
|
|
if expected != actual:
|
|
raise ValueError(
|
|
f"Multimodal placeholder mismatch: mask has {expected} values, "
|
|
f"features have {actual}"
|
|
)
|
|
return embeddings.masked_scatter(
|
|
mask,
|
|
features.to(device=embeddings.device, dtype=embeddings.dtype),
|
|
)
|
|
|
|
|
|
def _multimodal_embeddings(
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
|
inputs: dict[str, torch.Tensor],
|
|
) -> torch.Tensor:
|
|
"""Reproduce the Thinker merge but skip its large vocabulary LM head."""
|
|
|
|
input_ids = inputs["input_ids"]
|
|
embeddings = model.get_input_embeddings()(input_ids)
|
|
|
|
if inputs.get("input_features") is not None:
|
|
audio_features = model.get_audio_features(
|
|
inputs["input_features"],
|
|
feature_attention_mask=inputs.get("feature_attention_mask"),
|
|
audio_feature_lengths=None,
|
|
)
|
|
embeddings = _masked_scatter_features(
|
|
embeddings,
|
|
input_ids,
|
|
_config_token_id(model.config, "audio"),
|
|
audio_features,
|
|
)
|
|
|
|
if inputs.get("pixel_values") is not None:
|
|
image_features = model.get_image_features(
|
|
inputs["pixel_values"], inputs.get("image_grid_thw")
|
|
)
|
|
embeddings = _masked_scatter_features(
|
|
embeddings,
|
|
input_ids,
|
|
_config_token_id(model.config, "image"),
|
|
image_features,
|
|
)
|
|
|
|
if inputs.get("pixel_values_videos") is not None:
|
|
video_features = model.get_video_features(
|
|
inputs["pixel_values_videos"], inputs.get("video_grid_thw")
|
|
)
|
|
embeddings = _masked_scatter_features(
|
|
embeddings,
|
|
input_ids,
|
|
_config_token_id(model.config, "video"),
|
|
video_features,
|
|
)
|
|
return embeddings
|
|
|
|
|
|
def _local_position_ids(
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
|
inputs: dict[str, torch.Tensor],
|
|
*,
|
|
use_audio_in_video: bool,
|
|
) -> torch.LongTensor:
|
|
feature_mask = inputs.get("feature_attention_mask")
|
|
audio_feature_lengths = feature_mask.sum(dim=1) if feature_mask is not None else None
|
|
# Positional arguments deliberately mirror Transformers 4.52.3. Passing
|
|
# position_ids explicitly avoids the model-global mutable rope_deltas path.
|
|
position_ids, _ = model.get_rope_index(
|
|
inputs["input_ids"],
|
|
inputs.get("image_grid_thw"),
|
|
inputs.get("video_grid_thw"),
|
|
inputs.get("attention_mask"),
|
|
use_audio_in_video,
|
|
audio_feature_lengths,
|
|
inputs.get("video_second_per_grid"),
|
|
)
|
|
return position_ids
|
|
|
|
|
|
def _open_media_pieces(
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
|
inputs: dict[str, torch.Tensor],
|
|
*,
|
|
use_audio_in_video: bool,
|
|
) -> OpenMediaPieces:
|
|
"""Remove only the outer AV/chat closure while retaining media payload."""
|
|
|
|
input_ids = inputs["input_ids"]
|
|
if input_ids.shape[0] != 1:
|
|
raise ValueError("Open-stream extraction supports batch size 1 only")
|
|
audio_id = _config_token_id(model.config, "audio")
|
|
video_id = _config_token_id(model.config, "video")
|
|
payload_mask = (input_ids == audio_id) | (input_ids == video_id)
|
|
indices = payload_mask[0].nonzero(as_tuple=False).flatten()
|
|
if indices.numel() == 0:
|
|
raise ValueError("Prepared AV turn contains no audio/video payload tokens")
|
|
first = int(indices[0].item())
|
|
last = int(indices[-1].item())
|
|
if not bool(payload_mask[0, first : last + 1].all().item()):
|
|
raise RuntimeError(
|
|
"Audio/video payload is not contiguous; cannot open-stream this "
|
|
"processor template safely"
|
|
)
|
|
|
|
embeddings = _multimodal_embeddings(model, inputs)
|
|
positions = _local_position_ids(
|
|
model,
|
|
inputs,
|
|
use_audio_in_video=use_audio_in_video,
|
|
)
|
|
closure_ids = input_ids[:, last + 1 :]
|
|
if closure_ids.shape[1] == 0:
|
|
raise RuntimeError("Prepared AV turn has no closing/template tokens")
|
|
return OpenMediaPieces(
|
|
prefix_embeddings=embeddings[:, :first, :],
|
|
prefix_position_ids=positions[:, :, :first],
|
|
payload_embeddings=embeddings[:, first : last + 1, :],
|
|
payload_position_ids=positions[:, :, first : last + 1],
|
|
closure_ids=closure_ids,
|
|
closure_position_ids=positions[:, :, last + 1 :],
|
|
)
|
|
|
|
|
|
def _tmrope_positions_per_chunk(
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
|
) -> int:
|
|
positions_per_second = float(model.config.position_id_per_seconds)
|
|
seconds_per_chunk = float(model.config.seconds_per_chunk)
|
|
value = positions_per_second * seconds_per_chunk
|
|
rounded = round(value)
|
|
if not math.isclose(value, rounded, rel_tol=0.0, abs_tol=1e-6):
|
|
raise ValueError(f"Non-integral TMRoPE positions per chunk: {value}")
|
|
return int(rounded)
|
|
|
|
|
|
def _cache_length(cache: Any | None) -> int:
|
|
if cache is None:
|
|
return 0
|
|
if hasattr(cache, "get_seq_length"):
|
|
return int(cache.get_seq_length())
|
|
if isinstance(cache, (tuple, list)) and cache:
|
|
return int(cache[0][0].shape[-2])
|
|
raise TypeError(f"Unsupported cache type: {type(cache)!r}")
|
|
|
|
|
|
def clone_cache(cache: Any | None) -> Any | None:
|
|
"""Clone a cache because DynamicCache.update mutates its argument in place."""
|
|
|
|
if cache is None:
|
|
return None
|
|
if hasattr(cache, "to_legacy_cache"):
|
|
legacy = cache.to_legacy_cache()
|
|
cloned = tuple((key.clone(), value.clone()) for key, value in legacy)
|
|
return DynamicCache.from_legacy_cache(cloned)
|
|
if isinstance(cache, (tuple, list)):
|
|
return tuple((key.clone(), value.clone()) for key, value in cache)
|
|
return copy.deepcopy(cache)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def _append_embeddings(
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
|
embeddings: torch.Tensor,
|
|
position_ids: torch.LongTensor,
|
|
*,
|
|
past_key_values: Any | None,
|
|
) -> tuple[Any, int]:
|
|
"""Append precomputed embeddings/3D positions to a Thinker KV cache."""
|
|
|
|
current_length = embeddings.shape[1]
|
|
if position_ids.shape != (3, embeddings.shape[0], current_length):
|
|
raise ValueError(
|
|
"position_ids/embeddings mismatch: "
|
|
f"{tuple(position_ids.shape)} vs {tuple(embeddings.shape)}"
|
|
)
|
|
past_length = _cache_length(past_key_values)
|
|
cache_position = torch.arange(
|
|
past_length,
|
|
past_length + current_length,
|
|
device=embeddings.device,
|
|
)
|
|
attention_mask = torch.ones(
|
|
(embeddings.shape[0], past_length + current_length),
|
|
dtype=torch.long,
|
|
device=embeddings.device,
|
|
)
|
|
outputs = model.model(
|
|
attention_mask=attention_mask,
|
|
position_ids=position_ids,
|
|
past_key_values=past_key_values,
|
|
inputs_embeds=embeddings,
|
|
use_cache=True,
|
|
output_hidden_states=False,
|
|
return_dict=True,
|
|
cache_position=cache_position,
|
|
)
|
|
next_rope_offset = int(position_ids.max().item()) + 1
|
|
return outputs.past_key_values, next_rope_offset
|
|
|
|
|
|
@torch.inference_mode()
|
|
def _append_media_turn(
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
|
inputs: dict[str, torch.Tensor],
|
|
*,
|
|
past_key_values: Any | None,
|
|
rope_offset: int,
|
|
use_audio_in_video: bool,
|
|
) -> tuple[Any, int]:
|
|
"""Append one media turn and return its cache and next 3D-RoPE offset."""
|
|
|
|
input_ids = inputs["input_ids"]
|
|
embeddings = _multimodal_embeddings(model, inputs)
|
|
position_ids = _local_position_ids(
|
|
model, inputs, use_audio_in_video=use_audio_in_video
|
|
) + rope_offset
|
|
return _append_embeddings(
|
|
model,
|
|
embeddings,
|
|
position_ids,
|
|
past_key_values=past_key_values,
|
|
)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def _read_state(
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
|
media_cache: Any,
|
|
*,
|
|
next_rope_offset: int,
|
|
query_ids: torch.LongTensor,
|
|
last_n_layers: int = 1,
|
|
) -> torch.Tensor:
|
|
"""Read one q_t without committing the artificial assistant query."""
|
|
|
|
query_cache = clone_cache(media_cache)
|
|
past_length = _cache_length(query_cache)
|
|
query_length = query_ids.shape[1]
|
|
embeddings = model.get_input_embeddings()(query_ids)
|
|
scalar_positions = torch.arange(
|
|
next_rope_offset,
|
|
next_rope_offset + query_length,
|
|
device=query_ids.device,
|
|
)
|
|
position_ids = scalar_positions.view(1, 1, -1).expand(3, query_ids.shape[0], -1)
|
|
cache_position = torch.arange(
|
|
past_length,
|
|
past_length + query_length,
|
|
device=query_ids.device,
|
|
)
|
|
attention_mask = torch.ones(
|
|
(query_ids.shape[0], past_length + query_length),
|
|
dtype=torch.long,
|
|
device=query_ids.device,
|
|
)
|
|
outputs = model.model(
|
|
attention_mask=attention_mask,
|
|
position_ids=position_ids,
|
|
past_key_values=query_cache,
|
|
inputs_embeds=embeddings,
|
|
use_cache=True,
|
|
output_hidden_states=last_n_layers > 1,
|
|
return_dict=True,
|
|
cache_position=cache_position,
|
|
)
|
|
return _readout_hidden_state(outputs, last_n_layers=last_n_layers)
|
|
|
|
|
|
def _readout_hidden_state(outputs: Any, *, last_n_layers: int) -> torch.Tensor:
|
|
"""Return the query state, optionally averaging the final decoder layers."""
|
|
|
|
if last_n_layers == 1:
|
|
return outputs.last_hidden_state[:, -1, :]
|
|
hidden_states = outputs.hidden_states
|
|
if hidden_states is None or len(hidden_states) < last_n_layers:
|
|
available = 0 if hidden_states is None else len(hidden_states)
|
|
raise RuntimeError(
|
|
f"Requested the last {last_n_layers} hidden states, but the model "
|
|
f"returned {available}"
|
|
)
|
|
return torch.stack(
|
|
[state[:, -1, :] for state in hidden_states[-last_n_layers:]],
|
|
dim=0,
|
|
).mean(dim=0)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def _read_open_state(
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
|
open_media_cache: Any,
|
|
*,
|
|
closure_ids: torch.LongTensor,
|
|
closure_position_ids: torch.LongTensor,
|
|
query_ids: torch.LongTensor,
|
|
last_n_layers: int = 1,
|
|
) -> torch.Tensor:
|
|
"""Temporarily close an open media turn, read q_t, and discard the branch."""
|
|
|
|
branch_cache = clone_cache(open_media_cache)
|
|
if closure_position_ids.shape != (
|
|
3,
|
|
closure_ids.shape[0],
|
|
closure_ids.shape[1],
|
|
):
|
|
raise ValueError(
|
|
"closure position/ID mismatch: "
|
|
f"{tuple(closure_position_ids.shape)} vs {tuple(closure_ids.shape)}"
|
|
)
|
|
readout_ids = torch.cat(
|
|
[closure_ids.to(query_ids.device), query_ids],
|
|
dim=1,
|
|
)
|
|
past_length = _cache_length(branch_cache)
|
|
readout_length = readout_ids.shape[1]
|
|
embeddings = model.get_input_embeddings()(readout_ids)
|
|
|
|
# Qwen's AV TMRoPE gives the paired video/audio closing tokens the same
|
|
# position. Preserve those model-derived positions instead of replacing
|
|
# the whole closure with a scalar arange. Only the artificial assistant
|
|
# query follows the closure's maximum position.
|
|
closure_position_ids = closure_position_ids.to(query_ids.device)
|
|
query_start = int(closure_position_ids.max().item()) + 1
|
|
query_scalar_positions = torch.arange(
|
|
query_start,
|
|
query_start + query_ids.shape[1],
|
|
device=readout_ids.device,
|
|
)
|
|
query_position_ids = query_scalar_positions.view(1, 1, -1).expand(
|
|
3, query_ids.shape[0], -1
|
|
)
|
|
position_ids = torch.cat(
|
|
[closure_position_ids, query_position_ids],
|
|
dim=2,
|
|
)
|
|
cache_position = torch.arange(
|
|
past_length,
|
|
past_length + readout_length,
|
|
device=readout_ids.device,
|
|
)
|
|
attention_mask = torch.ones(
|
|
(readout_ids.shape[0], past_length + readout_length),
|
|
dtype=torch.long,
|
|
device=readout_ids.device,
|
|
)
|
|
outputs = model.model(
|
|
attention_mask=attention_mask,
|
|
position_ids=position_ids,
|
|
past_key_values=branch_cache,
|
|
inputs_embeds=embeddings,
|
|
use_cache=True,
|
|
output_hidden_states=last_n_layers > 1,
|
|
return_dict=True,
|
|
cache_position=cache_position,
|
|
)
|
|
return _readout_hidden_state(outputs, last_n_layers=last_n_layers)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def _append_open_pieces(
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
|
pieces: OpenMediaPieces,
|
|
*,
|
|
past_key_values: Any | None,
|
|
rope_offset: int,
|
|
include_prefix: bool,
|
|
) -> tuple[Any, int, torch.LongTensor]:
|
|
"""Append one AV payload and return its shifted closure positions."""
|
|
|
|
if include_prefix:
|
|
if past_key_values is not None:
|
|
raise ValueError("An open-media prefix can only start an empty cache")
|
|
payload_base = int(pieces.payload_position_ids.min().item())
|
|
if rope_offset != payload_base:
|
|
raise ValueError(
|
|
"The first open payload must retain its local TMRoPE base: "
|
|
f"got {rope_offset}, expected {payload_base}"
|
|
)
|
|
embeddings = torch.cat(
|
|
[pieces.prefix_embeddings, pieces.payload_embeddings], dim=1
|
|
)
|
|
position_ids = torch.cat(
|
|
[pieces.prefix_position_ids, pieces.payload_position_ids], dim=2
|
|
)
|
|
closure_position_ids = pieces.closure_position_ids
|
|
else:
|
|
if past_key_values is None:
|
|
raise ValueError("Payload-only append requires an existing open cache")
|
|
embeddings = pieces.payload_embeddings
|
|
position_shift = rope_offset - int(
|
|
pieces.payload_position_ids.min().item()
|
|
)
|
|
position_ids = pieces.payload_position_ids + position_shift
|
|
closure_position_ids = pieces.closure_position_ids + position_shift
|
|
next_cache, next_rope_offset = _append_embeddings(
|
|
model,
|
|
embeddings,
|
|
position_ids,
|
|
past_key_values=past_key_values,
|
|
)
|
|
return next_cache, next_rope_offset, closure_position_ids
|
|
|
|
|
|
class CausalPrefixExtractor:
|
|
"""Slow reference: independently encode the last M seconds at every step."""
|
|
|
|
def __init__(
|
|
self,
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
|
processor: Qwen2_5OmniProcessor,
|
|
config: StreamConfig,
|
|
fixed_grid_av: FixedGridAV | None = None,
|
|
) -> None:
|
|
config.validate()
|
|
self.model = model
|
|
self.processor = processor
|
|
self.config = config
|
|
self.fixed_grid_av = fixed_grid_av
|
|
self.query_ids = _readout_query_ids(processor, model.device)
|
|
|
|
def state_for_intervals(
|
|
self,
|
|
*,
|
|
video_path: str | Path | None,
|
|
audio_path: str | Path | None,
|
|
intervals: list[tuple[float, float]],
|
|
use_audio_in_video: bool,
|
|
) -> torch.Tensor:
|
|
if not intervals:
|
|
raise ValueError("At least one interval is required")
|
|
|
|
# Replay media turns sequentially from an empty cache. Different turns
|
|
# contain different frame counts, so batching them would make the HF
|
|
# video processor attempt torch.stack([T1,C,H,W], [T2,C,H,W], ...).
|
|
media_cache = None
|
|
next_offset = 0
|
|
for turn_index, interval in enumerate(intervals):
|
|
conversation, manual_audios = build_chunk_conversation(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
intervals=[interval],
|
|
use_audio_in_video=use_audio_in_video,
|
|
config=self.config,
|
|
fixed_grid_av=self.fixed_grid_av,
|
|
)
|
|
text_override = _media_turn_template_suffix(
|
|
self.processor,
|
|
conversation[0],
|
|
turn_index=turn_index,
|
|
)
|
|
inputs = _prepare_conversation(
|
|
self.processor,
|
|
conversation,
|
|
use_audio_in_video=use_audio_in_video,
|
|
model=self.model,
|
|
manual_audios=manual_audios,
|
|
video_fps=self.config.video_fps,
|
|
text_override=text_override,
|
|
)
|
|
media_cache, next_offset = _append_media_turn(
|
|
self.model,
|
|
inputs,
|
|
past_key_values=media_cache,
|
|
rope_offset=next_offset,
|
|
use_audio_in_video=use_audio_in_video,
|
|
)
|
|
return _read_state(
|
|
self.model,
|
|
media_cache,
|
|
next_rope_offset=next_offset,
|
|
query_ids=self.query_ids,
|
|
last_n_layers=self.config.readout_last_n_layers,
|
|
)
|
|
|
|
def extract(
|
|
self,
|
|
*,
|
|
video_path: str | Path | None,
|
|
audio_path: str | Path | None,
|
|
start: float,
|
|
end: float,
|
|
use_audio_in_video: bool,
|
|
show_progress: bool = True,
|
|
) -> ExtractionResult:
|
|
times = _time_grid(start, end, self.config.step_seconds)
|
|
states = []
|
|
with tqdm(
|
|
total=len(times),
|
|
desc="causal-prefix",
|
|
unit="state",
|
|
dynamic_ncols=True,
|
|
mininterval=0.0,
|
|
miniters=1,
|
|
disable=not show_progress,
|
|
) as progress:
|
|
for current_time in times:
|
|
context_start = max(
|
|
start, current_time - self.config.memory_seconds
|
|
)
|
|
state = self.state_for_intervals(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
intervals=[(context_start, current_time)],
|
|
use_audio_in_video=use_audio_in_video,
|
|
)
|
|
states.append(state.cpu())
|
|
progress.set_postfix_str(
|
|
f"media={current_time:.1f}s", refresh=False
|
|
)
|
|
progress.update(1)
|
|
return ExtractionResult(
|
|
times=torch.tensor(times, dtype=torch.float64),
|
|
states=torch.cat(states, dim=0),
|
|
verification=[],
|
|
)
|
|
|
|
|
|
class OpenStreamReplayExtractor:
|
|
"""Cache completed AV payloads inside one still-open multimodal turn."""
|
|
|
|
def __init__(
|
|
self,
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
|
processor: Qwen2_5OmniProcessor,
|
|
config: StreamConfig,
|
|
fixed_grid_av: FixedGridAV | None = None,
|
|
) -> None:
|
|
config.validate()
|
|
model_chunk_seconds = float(model.config.seconds_per_chunk)
|
|
if not math.isclose(
|
|
config.native_block_seconds,
|
|
model_chunk_seconds,
|
|
rel_tol=0.0,
|
|
abs_tol=1e-6,
|
|
):
|
|
raise ValueError(
|
|
"open-stream-replay requires native_block_seconds to match "
|
|
f"model.config.seconds_per_chunk ({model_chunk_seconds:g})"
|
|
)
|
|
self.model = model
|
|
self.processor = processor
|
|
self.config = config
|
|
self.fixed_grid_av = fixed_grid_av
|
|
self.query_ids = _readout_query_ids(processor, model.device)
|
|
self.continuous_reference = CausalPrefixExtractor(
|
|
model,
|
|
processor,
|
|
config,
|
|
fixed_grid_av=fixed_grid_av,
|
|
)
|
|
|
|
def _prepare_pieces(
|
|
self,
|
|
*,
|
|
video_path: str | Path | None,
|
|
audio_path: str | Path | None,
|
|
interval: tuple[float, float],
|
|
use_audio_in_video: bool,
|
|
audio_feature_device: str | torch.device | None = None,
|
|
audio_right_padding_samples: int | None = None,
|
|
) -> OpenMediaPieces:
|
|
conversation, manual_audios = build_chunk_conversation(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
intervals=[interval],
|
|
use_audio_in_video=use_audio_in_video,
|
|
config=self.config,
|
|
fixed_grid_av=self.fixed_grid_av,
|
|
)
|
|
inputs = _prepare_conversation(
|
|
self.processor,
|
|
conversation,
|
|
use_audio_in_video=use_audio_in_video,
|
|
model=self.model,
|
|
manual_audios=manual_audios,
|
|
video_fps=self.config.video_fps,
|
|
audio_feature_device=audio_feature_device,
|
|
audio_right_padding_samples=audio_right_padding_samples,
|
|
)
|
|
pieces = _open_media_pieces(
|
|
self.model,
|
|
inputs,
|
|
use_audio_in_video=use_audio_in_video,
|
|
)
|
|
# A standalone tail is at most one native 2 s block. Its temporal
|
|
# TMRoPE coordinates must therefore stay inside one 50-position chunk.
|
|
# This catches endpoint-inclusive frame selection (or another
|
|
# preprocessing regression) before overlapping coordinates are
|
|
# committed to the next block.
|
|
duration = interval[1] - interval[0]
|
|
if duration <= self.config.native_block_seconds + 1e-6:
|
|
temporal_positions = pieces.payload_position_ids[0]
|
|
temporal_span = int(
|
|
temporal_positions.max().item()
|
|
- temporal_positions.min().item()
|
|
)
|
|
positions_per_chunk = _tmrope_positions_per_chunk(self.model)
|
|
if temporal_span >= positions_per_chunk:
|
|
raise RuntimeError(
|
|
"One open-stream AV interval spans multiple native TMRoPE "
|
|
f"chunks: interval={interval}, temporal_span={temporal_span}, "
|
|
f"chunk_size={positions_per_chunk}. Check frame-boundary "
|
|
"ownership and audio preprocessing."
|
|
)
|
|
return pieces
|
|
|
|
def state_from_scratch(
|
|
self,
|
|
*,
|
|
video_path: str | Path | None,
|
|
audio_path: str | Path | None,
|
|
intervals: list[tuple[float, float]],
|
|
use_audio_in_video: bool,
|
|
) -> torch.Tensor:
|
|
"""Replay the same open stream from an empty cache for verification."""
|
|
|
|
(
|
|
media_cache,
|
|
_,
|
|
_,
|
|
closure_ids,
|
|
closure_position_ids,
|
|
) = self._cache_from_intervals(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
intervals=intervals,
|
|
use_audio_in_video=use_audio_in_video,
|
|
)
|
|
return _read_open_state(
|
|
self.model,
|
|
media_cache,
|
|
closure_ids=closure_ids,
|
|
closure_position_ids=closure_position_ids,
|
|
query_ids=self.query_ids,
|
|
last_n_layers=self.config.readout_last_n_layers,
|
|
)
|
|
|
|
def _cache_from_intervals(
|
|
self,
|
|
*,
|
|
video_path: str | Path | None,
|
|
audio_path: str | Path | None,
|
|
intervals: list[tuple[float, float]],
|
|
use_audio_in_video: bool,
|
|
prepared_blocks: list[CachedOpenMediaBlock] | None = None,
|
|
) -> tuple[Any, int, int, torch.LongTensor, torch.LongTensor]:
|
|
"""Build an open payload cache and expose its local TMRoPE base.
|
|
|
|
``prepared_blocks`` is used only by the production rolling-window
|
|
path. Correctness/reference calls omit it and independently rerun the
|
|
processor and modality encoders from the source media.
|
|
"""
|
|
|
|
if not intervals:
|
|
raise ValueError("At least one interval is required")
|
|
if (
|
|
prepared_blocks is not None
|
|
and len(prepared_blocks) != len(intervals)
|
|
):
|
|
raise ValueError(
|
|
"prepared block count does not match replay intervals: "
|
|
f"{len(prepared_blocks)} vs {len(intervals)}"
|
|
)
|
|
media_cache = None
|
|
next_offset = 0
|
|
closure_ids = None
|
|
closure_position_ids = None
|
|
payload_base = None
|
|
positions_per_chunk = _tmrope_positions_per_chunk(self.model)
|
|
for index, interval in enumerate(intervals):
|
|
if prepared_blocks is None:
|
|
pieces = self._prepare_pieces(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
interval=interval,
|
|
use_audio_in_video=use_audio_in_video,
|
|
)
|
|
else:
|
|
prepared = prepared_blocks[index]
|
|
if not (
|
|
math.isclose(
|
|
prepared.start,
|
|
interval[0],
|
|
rel_tol=0.0,
|
|
abs_tol=1e-6,
|
|
)
|
|
and math.isclose(
|
|
prepared.end,
|
|
interval[1],
|
|
rel_tol=0.0,
|
|
abs_tol=1e-6,
|
|
)
|
|
):
|
|
raise ValueError(
|
|
"prepared block/replay interval mismatch: "
|
|
f"prepared [{prepared.start:g}, {prepared.end:g}), "
|
|
f"requested [{interval[0]:g}, {interval[1]:g})"
|
|
)
|
|
pieces = prepared.pieces
|
|
if payload_base is None:
|
|
payload_base = int(pieces.payload_position_ids.min().item())
|
|
payload_offset = payload_base + index * positions_per_chunk
|
|
(
|
|
media_cache,
|
|
next_offset,
|
|
closure_position_ids,
|
|
) = _append_open_pieces(
|
|
self.model,
|
|
pieces,
|
|
past_key_values=media_cache,
|
|
rope_offset=payload_offset,
|
|
include_prefix=index == 0,
|
|
)
|
|
closure_ids = pieces.closure_ids
|
|
if closure_ids is None:
|
|
raise RuntimeError("Open-stream replay did not produce closure tokens")
|
|
if closure_position_ids is None:
|
|
raise RuntimeError("Open-stream replay did not produce closure positions")
|
|
if payload_base is None:
|
|
raise RuntimeError("Open-stream replay did not establish a payload base")
|
|
return (
|
|
media_cache,
|
|
next_offset,
|
|
payload_base,
|
|
closure_ids,
|
|
closure_position_ids,
|
|
)
|
|
|
|
def _extract_segment(
|
|
self,
|
|
*,
|
|
video_path: str | Path | None,
|
|
audio_path: str | Path | None,
|
|
segment_start: float,
|
|
segment_end: float,
|
|
use_audio_in_video: bool,
|
|
verify_every: int,
|
|
verify_continuous_every: int,
|
|
global_step_offset: int,
|
|
progress: Any,
|
|
) -> tuple[list[float], list[torch.Tensor], list[dict[str, Any]]]:
|
|
times = _time_grid(segment_start, segment_end, self.config.step_seconds)
|
|
committed_cache = None
|
|
block_start = segment_start
|
|
history_start = segment_start
|
|
payload_base = None
|
|
positions_per_chunk = _tmrope_positions_per_chunk(self.model)
|
|
states: list[torch.Tensor] = []
|
|
checks: list[dict[str, Any]] = []
|
|
replay_block_count = max(
|
|
0,
|
|
_require_integer_ratio(
|
|
self.config.memory_seconds,
|
|
self.config.native_block_seconds,
|
|
"memory / native block",
|
|
)
|
|
- 1,
|
|
)
|
|
needs_rolling_rebuild = (
|
|
self.config.memory_seconds
|
|
< segment_end - segment_start - 1e-6
|
|
)
|
|
completed_blocks = CompletedBlockPieceCache(
|
|
origin=segment_start,
|
|
block_seconds=self.config.native_block_seconds,
|
|
max_blocks=replay_block_count if needs_rolling_rebuild else 0,
|
|
)
|
|
|
|
for local_step, current_time in enumerate(times, start=1):
|
|
while current_time > block_start + self.config.native_block_seconds + 1e-6:
|
|
block_start += self.config.native_block_seconds
|
|
|
|
# At the first prediction in a new block beyond the memory budget,
|
|
# rebuild the open cache from the most recent complete blocks. This
|
|
# avoids an empty-history reset while keeping 18s completed history
|
|
# + up to 2s current tail for the default 20s configuration.
|
|
if (
|
|
committed_cache is not None
|
|
and block_start - history_start
|
|
>= self.config.memory_seconds - 1e-6
|
|
):
|
|
history_start = block_start - (
|
|
self.config.memory_seconds - self.config.native_block_seconds
|
|
)
|
|
history_intervals = _block_intervals(
|
|
history_start,
|
|
block_start,
|
|
self.config.native_block_seconds,
|
|
)
|
|
if history_intervals:
|
|
prepared_history = completed_blocks.get_many(
|
|
history_intervals
|
|
)
|
|
(
|
|
committed_cache,
|
|
_,
|
|
payload_base,
|
|
_,
|
|
_,
|
|
) = self._cache_from_intervals(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
intervals=history_intervals,
|
|
use_audio_in_video=use_audio_in_video,
|
|
prepared_blocks=prepared_history,
|
|
)
|
|
else:
|
|
# A one-block memory budget intentionally retains no
|
|
# completed history before the current partial block.
|
|
committed_cache = None
|
|
payload_base = None
|
|
|
|
pieces = self._prepare_pieces(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
interval=(block_start, current_time),
|
|
use_audio_in_video=use_audio_in_video,
|
|
)
|
|
if payload_base is None:
|
|
payload_base = int(pieces.payload_position_ids.min().item())
|
|
block_index = round(
|
|
(block_start - history_start)
|
|
/ self.config.native_block_seconds
|
|
)
|
|
payload_offset = payload_base + block_index * positions_per_chunk
|
|
branch_cache = clone_cache(committed_cache)
|
|
(
|
|
media_cache,
|
|
_,
|
|
closure_position_ids,
|
|
) = _append_open_pieces(
|
|
self.model,
|
|
pieces,
|
|
past_key_values=branch_cache,
|
|
rope_offset=payload_offset,
|
|
include_prefix=committed_cache is None,
|
|
)
|
|
state = _read_open_state(
|
|
self.model,
|
|
media_cache,
|
|
closure_ids=pieces.closure_ids,
|
|
closure_position_ids=closure_position_ids,
|
|
query_ids=self.query_ids,
|
|
last_n_layers=self.config.readout_last_n_layers,
|
|
)
|
|
states.append(state.cpu())
|
|
|
|
global_step = global_step_offset + local_step
|
|
if verify_every > 0 and global_step % verify_every == 0:
|
|
intervals = _block_intervals(
|
|
history_start,
|
|
current_time,
|
|
self.config.native_block_seconds,
|
|
)
|
|
exact = self.state_from_scratch(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
intervals=intervals,
|
|
use_audio_in_video=use_audio_in_video,
|
|
)
|
|
checks.append(
|
|
_compare_states(current_time, state, exact, kind="cache_replay")
|
|
)
|
|
if (
|
|
verify_continuous_every > 0
|
|
and global_step % verify_continuous_every == 0
|
|
):
|
|
continuous = self.continuous_reference.state_for_intervals(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
intervals=[(history_start, current_time)],
|
|
use_audio_in_video=use_audio_in_video,
|
|
)
|
|
checks.append(
|
|
_compare_states(
|
|
current_time,
|
|
state,
|
|
continuous,
|
|
kind="continuous_window_semantics",
|
|
)
|
|
)
|
|
|
|
block_end = min(
|
|
block_start + self.config.native_block_seconds,
|
|
segment_end,
|
|
)
|
|
if math.isclose(current_time, block_end, rel_tol=0.0, abs_tol=1e-6):
|
|
# The cache ends at AV payload, before media/chat closure.
|
|
committed_cache = media_cache
|
|
if math.isclose(
|
|
block_end - block_start,
|
|
self.config.native_block_seconds,
|
|
rel_tol=0.0,
|
|
abs_tol=1e-6,
|
|
):
|
|
completed_blocks.put(
|
|
CachedOpenMediaBlock(
|
|
start=block_start,
|
|
end=block_end,
|
|
pieces=pieces,
|
|
)
|
|
)
|
|
block_start = block_end
|
|
|
|
progress.set_postfix_str(
|
|
f"media={current_time:.1f}s", refresh=False
|
|
)
|
|
progress.update(1)
|
|
|
|
return times, states, checks
|
|
|
|
def extract(
|
|
self,
|
|
*,
|
|
video_path: str | Path | None,
|
|
audio_path: str | Path | None,
|
|
start: float,
|
|
end: float,
|
|
use_audio_in_video: bool,
|
|
verify_every: int = 0,
|
|
verify_continuous_every: int = 0,
|
|
show_progress: bool = True,
|
|
) -> ExtractionResult:
|
|
total_steps = len(_time_grid(start, end, self.config.step_seconds))
|
|
with tqdm(
|
|
total=total_steps,
|
|
desc="open-stream-replay",
|
|
unit="state",
|
|
dynamic_ncols=True,
|
|
mininterval=0.0,
|
|
miniters=1,
|
|
disable=not show_progress,
|
|
) as progress:
|
|
all_times, all_states, all_checks = self._extract_segment(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
segment_start=start,
|
|
segment_end=end,
|
|
use_audio_in_video=use_audio_in_video,
|
|
verify_every=verify_every,
|
|
verify_continuous_every=verify_continuous_every,
|
|
global_step_offset=0,
|
|
progress=progress,
|
|
)
|
|
return ExtractionResult(
|
|
times=torch.tensor(all_times, dtype=torch.float64),
|
|
states=torch.cat(all_states, dim=0),
|
|
verification=all_checks,
|
|
)
|
|
|
|
|
|
class OverlappingWorkingMemoryExtractor:
|
|
"""Serialize cached long memory followed by an overlapping recent recap.
|
|
|
|
This is a frozen-representation ablation, not the native Qwen streaming
|
|
path. At each prediction time, the total AV presentation budget is split
|
|
between source-aligned complete blocks and one sliding working interval.
|
|
For an 8-second budget and a 2-second working interval, the serialized
|
|
Thinker input is three cached 2-second blocks followed by ``[t-2,t]``.
|
|
The final fixed block and the recap intentionally overlap.
|
|
|
|
Fixed and working media readers may implement different visual grids. In
|
|
the direct-media experiment the fixed blocks remain half-open/global-grid.
|
|
The recap reuses global-grid historical frames but replaces its oldest
|
|
visual slot with a frame sampled exactly at current time ``t``.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
|
processor: Qwen2_5OmniProcessor,
|
|
config: StreamConfig,
|
|
*,
|
|
working_memory_seconds: float,
|
|
fixed_grid_av: FixedGridAV,
|
|
working_grid_av: FixedGridAV,
|
|
audio_feature_device: str | torch.device | None = None,
|
|
audio_right_padding_samples: int | None = None,
|
|
) -> None:
|
|
config.validate()
|
|
if working_memory_seconds <= 0:
|
|
raise ValueError("working_memory_seconds must be positive")
|
|
if not math.isclose(
|
|
working_memory_seconds,
|
|
config.native_block_seconds,
|
|
rel_tol=0.0,
|
|
abs_tol=1e-6,
|
|
):
|
|
raise ValueError(
|
|
"The first overlap-working-memory experiment requires one "
|
|
"native block of working memory: "
|
|
f"{working_memory_seconds:g}s vs "
|
|
f"{config.native_block_seconds:g}s"
|
|
)
|
|
total_blocks = _require_integer_ratio(
|
|
config.memory_seconds,
|
|
config.native_block_seconds,
|
|
"total presentation memory / native block",
|
|
)
|
|
if total_blocks < 1:
|
|
raise ValueError("The presentation budget must contain one block")
|
|
|
|
self.model = model
|
|
self.processor = processor
|
|
self.config = config
|
|
self.working_memory_seconds = float(working_memory_seconds)
|
|
if (
|
|
audio_right_padding_samples is not None
|
|
and audio_right_padding_samples < 0
|
|
):
|
|
raise ValueError("audio_right_padding_samples must be non-negative")
|
|
self.audio_feature_device = audio_feature_device
|
|
self.audio_right_padding_samples = audio_right_padding_samples
|
|
# The working recap consumes one block from the matched token budget.
|
|
self.fixed_block_budget = total_blocks - 1
|
|
self.fixed_extractor = OpenStreamReplayExtractor(
|
|
model,
|
|
processor,
|
|
config,
|
|
fixed_grid_av=fixed_grid_av,
|
|
)
|
|
self.working_extractor = OpenStreamReplayExtractor(
|
|
model,
|
|
processor,
|
|
config,
|
|
fixed_grid_av=working_grid_av,
|
|
)
|
|
self.query_ids = self.fixed_extractor.query_ids
|
|
self.layout_records: list[dict[str, Any]] = []
|
|
|
|
def _state_from_prepared(
|
|
self,
|
|
*,
|
|
video_path: str | Path | None,
|
|
audio_path: str | Path | None,
|
|
fixed_intervals: list[tuple[float, float]],
|
|
fixed_blocks: list[CachedOpenMediaBlock],
|
|
working_pieces: OpenMediaPieces,
|
|
use_audio_in_video: bool,
|
|
cached_fixed_payload: tuple[Any, int] | None = None,
|
|
) -> torch.Tensor:
|
|
positions_per_chunk = _tmrope_positions_per_chunk(self.model)
|
|
if fixed_intervals:
|
|
if cached_fixed_payload is None:
|
|
(
|
|
fixed_cache,
|
|
_,
|
|
payload_base,
|
|
_,
|
|
_,
|
|
) = self.fixed_extractor._cache_from_intervals(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
intervals=fixed_intervals,
|
|
use_audio_in_video=use_audio_in_video,
|
|
prepared_blocks=fixed_blocks,
|
|
)
|
|
else:
|
|
fixed_cache, payload_base = cached_fixed_payload
|
|
branch_cache = clone_cache(fixed_cache)
|
|
working_offset = payload_base + len(fixed_intervals) * positions_per_chunk
|
|
include_prefix = False
|
|
else:
|
|
branch_cache = None
|
|
working_offset = int(working_pieces.payload_position_ids.min().item())
|
|
include_prefix = True
|
|
|
|
(
|
|
media_cache,
|
|
_,
|
|
closure_position_ids,
|
|
) = _append_open_pieces(
|
|
self.model,
|
|
working_pieces,
|
|
past_key_values=branch_cache,
|
|
rope_offset=working_offset,
|
|
include_prefix=include_prefix,
|
|
)
|
|
return _read_open_state(
|
|
self.model,
|
|
media_cache,
|
|
closure_ids=working_pieces.closure_ids,
|
|
closure_position_ids=closure_position_ids,
|
|
query_ids=self.query_ids,
|
|
last_n_layers=self.config.readout_last_n_layers,
|
|
)
|
|
|
|
def extract(
|
|
self,
|
|
*,
|
|
video_path: str | Path | None,
|
|
audio_path: str | Path | None,
|
|
start: float,
|
|
end: float,
|
|
use_audio_in_video: bool,
|
|
verify_every: int = 0,
|
|
show_progress: bool = True,
|
|
) -> ExtractionResult:
|
|
if verify_every < 0:
|
|
raise ValueError("verify_every must be >= 0")
|
|
times = _time_grid(start, end, self.config.step_seconds)
|
|
block_seconds = self.config.native_block_seconds
|
|
completed_blocks = CompletedBlockPieceCache(
|
|
origin=start,
|
|
block_seconds=block_seconds,
|
|
max_blocks=max(1, self.fixed_block_budget),
|
|
)
|
|
next_fixed_index = 0
|
|
cached_interval_key: tuple[tuple[float, float], ...] | None = None
|
|
cached_fixed_payload: tuple[Any, int] | None = None
|
|
states: list[torch.Tensor] = []
|
|
checks: list[dict[str, Any]] = []
|
|
self.layout_records = []
|
|
|
|
with tqdm(
|
|
total=len(times),
|
|
desc="overlap-working-memory",
|
|
unit="state",
|
|
dynamic_ncols=True,
|
|
mininterval=0.0,
|
|
miniters=1,
|
|
disable=not show_progress,
|
|
) as progress:
|
|
for step_index, current_time in enumerate(times, start=1):
|
|
completed_count = int(
|
|
math.floor((current_time - start) / block_seconds + 1e-7)
|
|
)
|
|
if self.fixed_block_budget:
|
|
while next_fixed_index < completed_count:
|
|
block_start = start + next_fixed_index * block_seconds
|
|
block_end = block_start + block_seconds
|
|
pieces = self.fixed_extractor._prepare_pieces(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
interval=(block_start, block_end),
|
|
use_audio_in_video=use_audio_in_video,
|
|
audio_feature_device=self.audio_feature_device,
|
|
audio_right_padding_samples=(
|
|
self.audio_right_padding_samples
|
|
),
|
|
)
|
|
completed_blocks.put(
|
|
CachedOpenMediaBlock(
|
|
start=block_start,
|
|
end=block_end,
|
|
pieces=pieces,
|
|
)
|
|
)
|
|
next_fixed_index += 1
|
|
|
|
first_fixed_index = max(
|
|
0, completed_count - self.fixed_block_budget
|
|
)
|
|
fixed_intervals = [
|
|
(
|
|
start + index * block_seconds,
|
|
start + (index + 1) * block_seconds,
|
|
)
|
|
for index in range(first_fixed_index, completed_count)
|
|
]
|
|
fixed_blocks = (
|
|
completed_blocks.get_many(fixed_intervals)
|
|
if fixed_intervals
|
|
else []
|
|
)
|
|
interval_key = tuple(fixed_intervals)
|
|
if interval_key != cached_interval_key:
|
|
if fixed_intervals:
|
|
(
|
|
fixed_cache,
|
|
_,
|
|
payload_base,
|
|
_,
|
|
_,
|
|
) = self.fixed_extractor._cache_from_intervals(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
intervals=fixed_intervals,
|
|
use_audio_in_video=use_audio_in_video,
|
|
prepared_blocks=fixed_blocks,
|
|
)
|
|
cached_fixed_payload = (fixed_cache, payload_base)
|
|
else:
|
|
cached_fixed_payload = None
|
|
cached_interval_key = interval_key
|
|
|
|
working_start = max(start, current_time - self.working_memory_seconds)
|
|
working_pieces = self.working_extractor._prepare_pieces(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
interval=(working_start, current_time),
|
|
use_audio_in_video=use_audio_in_video,
|
|
audio_feature_device=self.audio_feature_device,
|
|
audio_right_padding_samples=self.audio_right_padding_samples,
|
|
)
|
|
state = self._state_from_prepared(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
fixed_intervals=fixed_intervals,
|
|
fixed_blocks=fixed_blocks,
|
|
working_pieces=working_pieces,
|
|
use_audio_in_video=use_audio_in_video,
|
|
cached_fixed_payload=cached_fixed_payload,
|
|
)
|
|
states.append(state.cpu())
|
|
|
|
fixed_start = (
|
|
fixed_intervals[0][0] if fixed_intervals else current_time
|
|
)
|
|
fixed_end = fixed_intervals[-1][1] if fixed_intervals else current_time
|
|
overlap_seconds = (
|
|
max(0.0, min(fixed_end, current_time) - max(fixed_start, working_start))
|
|
if fixed_intervals
|
|
else 0.0
|
|
)
|
|
self.layout_records.append(
|
|
{
|
|
"time": float(current_time),
|
|
"fixed_start": float(fixed_start),
|
|
"fixed_end": float(fixed_end),
|
|
"fixed_blocks": len(fixed_intervals),
|
|
"working_start": float(working_start),
|
|
"working_end": float(current_time),
|
|
"overlap_seconds": float(overlap_seconds),
|
|
}
|
|
)
|
|
|
|
if verify_every and step_index % verify_every == 0:
|
|
exact = self._state_from_prepared(
|
|
video_path=video_path,
|
|
audio_path=audio_path,
|
|
fixed_intervals=fixed_intervals,
|
|
fixed_blocks=fixed_blocks,
|
|
working_pieces=working_pieces,
|
|
use_audio_in_video=use_audio_in_video,
|
|
)
|
|
checks.append(
|
|
_compare_states(
|
|
current_time,
|
|
state,
|
|
exact,
|
|
kind="overlap_cache_replay",
|
|
)
|
|
)
|
|
|
|
progress.set_postfix_str(
|
|
f"media={current_time:.1f}s overlap={overlap_seconds:.1f}s",
|
|
refresh=False,
|
|
)
|
|
progress.update(1)
|
|
|
|
return ExtractionResult(
|
|
times=torch.tensor(times, dtype=torch.float64),
|
|
states=torch.cat(states, dim=0),
|
|
verification=checks,
|
|
)
|
|
|
|
|
|
def _block_intervals(start: float, end: float, block_seconds: float) -> list[tuple[float, float]]:
|
|
"""Exact AV payload intervals represented by an open-stream cache."""
|
|
|
|
intervals = []
|
|
cursor = start
|
|
while cursor < end - 1e-7:
|
|
next_end = min(cursor + block_seconds, end)
|
|
intervals.append((cursor, next_end))
|
|
cursor = next_end
|
|
return intervals
|
|
|
|
|
|
def _compare_states(
|
|
current_time: float,
|
|
cached: torch.Tensor,
|
|
exact: torch.Tensor,
|
|
*,
|
|
kind: str,
|
|
) -> dict[str, Any]:
|
|
cached_f = cached.float()
|
|
exact_f = exact.float()
|
|
return {
|
|
"kind": kind,
|
|
"time": float(current_time),
|
|
"max_abs": float((cached_f - exact_f).abs().max().item()),
|
|
"mean_abs": float((cached_f - exact_f).abs().mean().item()),
|
|
"cosine": float(nn.functional.cosine_similarity(cached_f, exact_f).item()),
|
|
}
|
|
|
|
|
|
def probe_duration(path: str | Path) -> float:
|
|
"""Read media duration with ffprobe."""
|
|
|
|
completed = subprocess.run(
|
|
[
|
|
"ffprobe",
|
|
"-v",
|
|
"error",
|
|
"-show_entries",
|
|
"format=duration",
|
|
"-of",
|
|
"default=noprint_wrappers=1:nokey=1",
|
|
str(Path(path).resolve()),
|
|
],
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
return float(completed.stdout.strip())
|
|
|
|
|
|
def probe_video_fps(path: str | Path) -> float:
|
|
"""Read the stored video frame rate with ffprobe."""
|
|
|
|
completed = subprocess.run(
|
|
[
|
|
"ffprobe",
|
|
"-v",
|
|
"error",
|
|
"-select_streams",
|
|
"v:0",
|
|
"-show_entries",
|
|
"stream=avg_frame_rate",
|
|
"-of",
|
|
"default=noprint_wrappers=1:nokey=1",
|
|
str(Path(path).resolve()),
|
|
],
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
numerator, denominator = completed.stdout.strip().split("/", maxsplit=1)
|
|
return float(numerator) / float(denominator)
|
|
|
|
|
|
def target_num_steps(path: str | Path) -> int:
|
|
"""Read an EgoSpeak target_perframe ``[T,3]`` length."""
|
|
|
|
try:
|
|
import numpy as np
|
|
except ImportError as exc:
|
|
raise RuntimeError("numpy is required for --target") from exc
|
|
target = np.load(Path(path), mmap_mode="r")
|
|
if target.ndim != 2 or target.shape[1] != 3:
|
|
raise ValueError(f"Expected target shape [T,3], got {target.shape}")
|
|
return int(target.shape[0])
|
|
|
|
|
|
def _apply_session_metadata(args: argparse.Namespace) -> dict[str, Any] | None:
|
|
"""Fill the prepared video and task length from a preparation sidecar."""
|
|
|
|
if args.session_metadata is None:
|
|
return None
|
|
metadata = json.loads(args.session_metadata.read_text())
|
|
fps_key = f"{args.video_fps:g}"
|
|
try:
|
|
metadata_video = Path(metadata["outputs"][fps_key])
|
|
except KeyError as exc:
|
|
raise SystemExit(
|
|
f"Metadata has no fixed {fps_key} FPS output: {args.session_metadata}"
|
|
) from exc
|
|
if args.video is not None and args.video.resolve() != metadata_video.resolve():
|
|
raise SystemExit("--video conflicts with the selected --session-metadata output")
|
|
args.video = metadata_video
|
|
metadata_steps = int(metadata["num_steps"])
|
|
if args.num_steps is not None and args.num_steps != metadata_steps:
|
|
raise SystemExit("--num-steps conflicts with --session-metadata")
|
|
args.num_steps = metadata_steps
|
|
return metadata
|
|
|
|
|
|
def _hidden_size(model: Qwen2_5OmniThinkerForConditionalGeneration) -> int:
|
|
text_config = getattr(model.config, "text_config", None)
|
|
if text_config is not None:
|
|
return int(text_config.hidden_size)
|
|
return int(model.config.hidden_size)
|
|
|
|
|
|
def _load_head(
|
|
checkpoint: Path,
|
|
hidden_size: int,
|
|
device: torch.device,
|
|
) -> TurnPredictionHead:
|
|
head = TurnPredictionHead(hidden_size=hidden_size)
|
|
payload = torch.load(checkpoint, map_location="cpu", weights_only=True)
|
|
state_dict = payload.get("head", payload) if isinstance(payload, dict) else payload
|
|
head.load_state_dict(state_dict)
|
|
return head.to(device).eval()
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--video", type=Path)
|
|
parser.add_argument("--audio", type=Path)
|
|
parser.add_argument(
|
|
"--session-metadata",
|
|
type=Path,
|
|
help="JSON sidecar written by prepare_easycom_session.py.",
|
|
)
|
|
parser.add_argument(
|
|
"--target",
|
|
type=Path,
|
|
help="target_perframe .npy; its T fixes the full-session output length.",
|
|
)
|
|
parser.add_argument(
|
|
"--num-steps",
|
|
type=int,
|
|
help="Exact number of 5 Hz outputs (normally read from --target/metadata).",
|
|
)
|
|
parser.add_argument("--use-audio-in-video", action="store_true")
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument(
|
|
"--mode",
|
|
choices=("causal-prefix", "open-stream-replay"),
|
|
default="open-stream-replay",
|
|
)
|
|
parser.add_argument("--model-id", default=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("--memory-seconds", type=float, default=20.0)
|
|
parser.add_argument("--video-fps", type=float, default=2.0)
|
|
parser.add_argument(
|
|
"--strict-fixed-video-grid",
|
|
action="store_true",
|
|
help="Require stored MP4 FPS to equal --video-fps.",
|
|
)
|
|
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("--verify-every", type=int, default=0)
|
|
parser.add_argument(
|
|
"--verify-continuous-every",
|
|
type=int,
|
|
default=0,
|
|
help="Also compare block replay with one continuous causal media window.",
|
|
)
|
|
parser.add_argument(
|
|
"--no-progress",
|
|
action="store_true",
|
|
help="Disable the per-state tqdm progress bar.",
|
|
)
|
|
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")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
_configure_qwen_logging()
|
|
session_metadata = _apply_session_metadata(args)
|
|
if args.video is None and args.audio is None:
|
|
raise SystemExit("Pass --video, --audio, or both")
|
|
if args.use_audio_in_video and args.audio is not None:
|
|
raise SystemExit("Do not combine --use-audio-in-video with --audio")
|
|
if args.verify_every < 0:
|
|
raise SystemExit("--verify-every must be >= 0")
|
|
if args.verify_continuous_every < 0:
|
|
raise SystemExit("--verify-continuous-every must be >= 0")
|
|
if args.num_steps is not None and args.num_steps <= 0:
|
|
raise SystemExit("--num-steps must be positive")
|
|
if args.target is not None:
|
|
target_steps = target_num_steps(args.target)
|
|
if args.num_steps is not None and args.num_steps != target_steps:
|
|
raise SystemExit("--num-steps conflicts with --target length")
|
|
args.num_steps = target_steps
|
|
if args.num_steps is not None:
|
|
if not math.isclose(args.start, 0.0, rel_tol=0.0, abs_tol=1e-9):
|
|
raise SystemExit("Exact full-session --num-steps extraction requires --start 0")
|
|
derived_end = args.num_steps * args.step_seconds
|
|
if args.end is not None and not math.isclose(
|
|
args.end, derived_end, rel_tol=0.0, abs_tol=1e-6
|
|
):
|
|
raise SystemExit(
|
|
f"--end {args.end} conflicts with {args.num_steps} steps "
|
|
f"at {args.step_seconds}s ({derived_end}s)"
|
|
)
|
|
end = derived_end
|
|
else:
|
|
media_for_duration = args.video if args.video is not None else args.audio
|
|
end = args.end if args.end is not None else probe_duration(media_for_duration)
|
|
if end <= args.start:
|
|
raise SystemExit("--end must be greater than --start")
|
|
|
|
config = StreamConfig(
|
|
step_seconds=args.step_seconds,
|
|
native_block_seconds=args.native_block_seconds,
|
|
memory_seconds=args.memory_seconds,
|
|
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,
|
|
)
|
|
config.validate()
|
|
|
|
input_video_fps = None
|
|
fixed_grid_av = None
|
|
if args.video is not None:
|
|
input_video_fps = probe_video_fps(args.video)
|
|
fixed_grid = math.isclose(
|
|
input_video_fps, config.video_fps, rel_tol=0.0, abs_tol=1e-6
|
|
)
|
|
if args.strict_fixed_video_grid and not fixed_grid:
|
|
raise SystemExit(
|
|
f"Input is {input_video_fps:g} FPS, expected a prepared fixed "
|
|
f"{config.video_fps:g} FPS video"
|
|
)
|
|
if not fixed_grid:
|
|
print(
|
|
f"WARNING: input is {input_video_fps:g} FPS while Qwen sampling is "
|
|
f"{config.video_fps:g} FPS. Use prepare_easycom_session.py and "
|
|
"--strict-fixed-video-grid for the paper experiment."
|
|
)
|
|
elif args.strict_fixed_video_grid and args.use_audio_in_video:
|
|
fixed_grid_av = FixedGridAV(args.video, config.video_fps)
|
|
|
|
model, processor = load_qwen_omni(
|
|
args.model_id,
|
|
dtype=args.dtype,
|
|
device_map={"": args.device},
|
|
use_flash_attention_2=args.flash_attention_2,
|
|
)
|
|
model.requires_grad_(False)
|
|
|
|
if args.mode == "causal-prefix":
|
|
extractor = CausalPrefixExtractor(
|
|
model, processor, config, fixed_grid_av=fixed_grid_av
|
|
)
|
|
result = extractor.extract(
|
|
video_path=args.video,
|
|
audio_path=args.audio,
|
|
start=args.start,
|
|
end=end,
|
|
use_audio_in_video=args.use_audio_in_video,
|
|
show_progress=not args.no_progress,
|
|
)
|
|
else:
|
|
extractor = OpenStreamReplayExtractor(
|
|
model, processor, config, fixed_grid_av=fixed_grid_av
|
|
)
|
|
result = extractor.extract(
|
|
video_path=args.video,
|
|
audio_path=args.audio,
|
|
start=args.start,
|
|
end=end,
|
|
use_audio_in_video=args.use_audio_in_video,
|
|
verify_every=args.verify_every,
|
|
verify_continuous_every=args.verify_continuous_every,
|
|
show_progress=not args.no_progress,
|
|
)
|
|
|
|
payload: dict[str, Any] = {
|
|
"model_id": args.model_id,
|
|
"mode": args.mode,
|
|
"stream_representation": (
|
|
"single_causal_window"
|
|
if args.mode == "causal-prefix"
|
|
else "single_open_media_turn_with_2s_payload_commits"
|
|
),
|
|
"config": asdict(config),
|
|
"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(args.target.resolve()) if args.target else None,
|
|
"session_metadata": session_metadata,
|
|
"input_video_fps": input_video_fps,
|
|
"video_sampling_grid": (
|
|
"explicit_fixed_half_open_past_frames_with_tail_duplication"
|
|
if fixed_grid_av is not None
|
|
else "qwen_utility_adaptive_sampling"
|
|
),
|
|
"media_interval_policy": (
|
|
"half_open_[start,end)"
|
|
if fixed_grid_av is not None
|
|
else "qwen_utility_defined"
|
|
),
|
|
"open_stream_policy": (
|
|
{
|
|
"commit": "completed_2s_av_payload_only",
|
|
"temporary": "partial_payload_plus_media_closure_plus_readout",
|
|
"tmrope_positions_per_chunk": _tmrope_positions_per_chunk(model),
|
|
"memory": "rolling_complete_blocks_plus_current_partial_block",
|
|
"rolling_replay_inputs": (
|
|
"reuse_cached_complete_block_av_embeddings"
|
|
),
|
|
"memory_seconds_range": [
|
|
config.memory_seconds - config.native_block_seconds,
|
|
config.memory_seconds,
|
|
],
|
|
}
|
|
if args.mode == "open-stream-replay"
|
|
else None
|
|
),
|
|
"times": result.times,
|
|
"states": result.states,
|
|
"verification": result.verification,
|
|
}
|
|
if args.head_checkpoint is not None:
|
|
head = _load_head(args.head_checkpoint, _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)
|
|
|
|
if args.num_steps is not None and len(result.times) != args.num_steps:
|
|
raise RuntimeError(
|
|
f"Extractor returned {len(result.times)} states, expected {args.num_steps}"
|
|
)
|
|
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({str(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 str(item.get("kind", "unspecified")) == kind
|
|
]
|
|
worst = max(item["max_abs"] for item in checks)
|
|
lowest_cosine = min(item["cosine"] for item in checks)
|
|
print(
|
|
f" {kind}: n={len(checks)}, worst max_abs={worst:.6g}, "
|
|
f"lowest cosine={lowest_cosine:.8f}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|