"""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 gc 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, checkpoint_head_config, checkpoint_loss_name, egospeak_multihorizon_loss, marginal_probabilities_from_logits, 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", ] BackboneLoader = Literal["qwen", "roma"] 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, backbone_loader: BackboneLoader = "qwen", ) -> tuple[Qwen2_5OmniThinkerForConditionalGeneration, Qwen2_5OmniProcessor]: """Load a multimodal Thinker and its processor. The normal Qwen checkpoints can be loaded directly as a Thinker. ROMA is released as a complete Qwen2.5-Omni model whose safetensors also contain custom ``thinker.gate_*`` modules. For a frozen-backbone experiment we deliberately ignore those binary proactive-speaking heads, load the full checkpoint so the ``thinker.*`` prefix is handled correctly, detach the Thinker, and release the speech-generation modules. """ 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" if backbone_loader == "qwen": model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained( model_id, **model_kwargs ) metadata: dict[str, Any] = { "loader": "qwen", "checkpoint_layout": "thinker", "original_speak_head": "not_applicable", } elif backbone_loader == "roma": try: from transformers import Qwen2_5OmniModel except ImportError as error: raise RuntimeError( "The ROMA loader requires transformers.Qwen2_5OmniModel. " "Use the ROMA/Qwen2.5-Omni environment rather than an older " "Transformers installation." ) from error loaded = Qwen2_5OmniModel.from_pretrained( model_id, output_loading_info=True, **model_kwargs, ) if not isinstance(loaded, tuple) or len(loaded) != 2: raise RuntimeError("ROMA loading did not return checkpoint audit data") full_model, loading_info = loaded missing = [str(key) for key in loading_info.get("missing_keys", [])] unexpected = [ str(key) for key in loading_info.get("unexpected_keys", []) ] mismatched = [ str(item) for item in loading_info.get("mismatched_keys", []) ] errors = [str(item) for item in loading_info.get("error_msgs", [])] non_gate_unexpected = [ key for key in unexpected if not key.startswith("thinker.gate_") ] if missing or mismatched or errors or non_gate_unexpected: raise RuntimeError( "ROMA checkpoint did not load cleanly as a Qwen2.5-Omni " "backbone: " f"missing={missing[:8]}, mismatched={mismatched[:8]}, " f"errors={errors[:4]}, unexpected_non_gate=" f"{non_gate_unexpected[:8]}" ) if not hasattr(full_model, "thinker"): raise RuntimeError("ROMA full checkpoint has no .thinker module") model = full_model.thinker gate_modules_present = bool( hasattr(model, "gate_mixer") and ( hasattr(model, "gate_head") or hasattr(model, "gate_head_pro_fc1") ) ) full_model.thinker = None del full_model gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() metadata = { "loader": "roma", "checkpoint_layout": "full_qwen2_5_omni", "ignored_checkpoint_keys": unexpected, "ignored_checkpoint_key_policy": "thinker.gate_*_only", "original_speak_head": ( "loaded_but_not_called; EasyCom head is trained anew" if gate_modules_present else ( "checkpoint_weights_ignored_by_standard_transformers; " "EasyCom head is trained anew" ) ), "original_speak_head_modules_present": gate_modules_present, } print( "ROMA backbone loaded cleanly: " f"gate_modules_present={gate_modules_present}, " f"ignored_gate_tensors={len(unexpected)}; " "Talker/Token2Wav released.", flush=True, ) else: raise ValueError(f"Unknown backbone_loader: {backbone_loader}") processor = Qwen2_5OmniProcessor.from_pretrained(model_id) model.eval() text_config = getattr(model.config, "text_config", None) if text_config is None and hasattr(model.config, "get_text_config"): text_config = model.config.get_text_config() metadata.update( { "model_id": str(model_id), "hidden_size": int(getattr(text_config, "hidden_size")), "num_hidden_layers": int( getattr(text_config, "num_hidden_layers") ), "seconds_per_chunk": float( getattr(model.config, "seconds_per_chunk") ), } ) setattr(model, "_turn_backbone_metadata", metadata) return model, processor def backbone_metadata( model: Qwen2_5OmniThinkerForConditionalGeneration, ) -> dict[str, Any]: """Return JSON/pickle-safe provenance attached by ``load_qwen_omni``.""" metadata = getattr(model, "_turn_backbone_metadata", None) if not isinstance(metadata, dict): return {"loader": "unknown"} return copy.deepcopy(metadata) @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 readout_layer_fusion: str = "mean" 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") if self.readout_layer_fusion not in {"mean", "stack"}: raise ValueError("readout_layer_fusion must be 'mean' or 'stack'") if ( self.readout_layer_fusion == "stack" and self.readout_last_n_layers < 2 ): raise ValueError( "stack readout requires readout_last_n_layers >= 2" ) _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] or preserved query layers [N,L,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=" 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, layer_fusion: str = "mean", ) -> 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, layer_fusion=layer_fusion, ) def _readout_hidden_state( outputs: Any, *, last_n_layers: int, layer_fusion: str, ) -> torch.Tensor: """Return one query vector or an ordered stack of final-layer vectors.""" if layer_fusion not in {"mean", "stack"}: raise ValueError(f"Unknown readout layer fusion: {layer_fusion}") if last_n_layers == 1: stacked = outputs.last_hidden_state[:, -1:, :] else: 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}" ) stacked = torch.stack( [state[:, -1, :] for state in hidden_states[-last_n_layers:]], dim=1, ) return stacked.mean(dim=1) if layer_fusion == "mean" else stacked @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, layer_fusion: str = "mean", ) -> 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, layer_fusion=layer_fusion, ) @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, layer_fusion=self.config.readout_layer_fusion, ) 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, layer_fusion=self.config.readout_layer_fusion, ) 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, layer_fusion=self.config.readout_layer_fusion, ) 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, layer_fusion=self.config.readout_layer_fusion, ) 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() cached_flat = cached_f.reshape(1, -1) exact_flat = exact_f.reshape(1, -1) 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_flat, exact_flat, dim=-1).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, ) -> tuple[TurnPredictionHead, str]: payload = torch.load(checkpoint, map_location="cpu", weights_only=True) if isinstance(payload, dict) and "head" in payload: config = checkpoint_head_config(payload) loss_name = checkpoint_loss_name(payload) if int(config["hidden_size"]) != hidden_size: raise ValueError( f"Head hidden size {config['hidden_size']} != model {hidden_size}" ) head = TurnPredictionHead(**config) state_dict = payload["head"] else: head = TurnPredictionHead(hidden_size=hidden_size) state_dict = payload loss_name = "lstr-mce" head.load_state_dict(state_dict) return head.to(device).eval(), loss_name 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( "--backbone-loader", choices=("qwen", "roma"), default="qwen", help=( "qwen loads a Thinker checkpoint directly; roma loads the full " "ROMA Qwen2.5-Omni checkpoint, keeps its Thinker, and ignores the " "original binary Speak Head" ), ) 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( "--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("--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, readout_layer_fusion=args.readout_layer_fusion, ) 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, backbone_loader=args.backbone_loader, ) 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, "backbone_loader": args.backbone_loader, "backbone_metadata": backbone_metadata(model), "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, head_loss_name = _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"] = marginal_probabilities_from_logits( logits, head_loss_name, ) payload["head_loss"] = head_loss_name 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()