icassp2027-experiment/streaming_turn_model.py

"""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,
    checkpoint_head_config,
    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
    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="<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,
    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,
) -> TurnPredictionHead:
    payload = torch.load(checkpoint, map_location="cpu", weights_only=True)
    if isinstance(payload, dict) and "head" in payload:
        config = checkpoint_head_config(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
    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(
        "--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,
    )
    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()