Add reproducible EasyCom causal head pipeline

streaming_turn_model.py

@@ -4,24 +4,22 @@ This module implements two 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. ``commit-replay``: keep the Thinker KV cache for completed 2-second
-   multimodal blocks, and replay only the causal prefix of the current block.
-3. ``open-stream-replay``: cache the same completed AV blocks inside one
+2. ``open-stream-replay``: cache completed AV blocks inside one
    still-open user/media turn; closure and readout tokens remain temporary.
 
 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 commit/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.
+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 commit-replay \
+        --mode open-stream-replay \
         --memory-seconds 20 \
         --verify-every 25 \
         --output features/session.pt
@@ -32,6 +30,7 @@ from __future__ import annotations
 import argparse
 import copy
 import json
+import logging
 import math
 import subprocess
 from dataclasses import asdict, dataclass
@@ -39,16 +38,58 @@ from pathlib import Path
 from typing import Any, Literal
 
 import torch
+from easycom_head import (
+    TurnPredictionHead,
+    egospeak_multihorizon_loss,
+    miniroad_multihorizon_loss,
+    multihorizon_cross_entropy,
+)
 from qwen_omni_utils import process_mm_info
 from torch import nn
+from tqdm.auto import tqdm
 from transformers import (
     Qwen2_5OmniProcessor,
     Qwen2_5OmniThinkerForConditionalGeneration,
 )
 from transformers.cache_utils import DynamicCache
 
-Mode = Literal["causal-prefix", "commit-replay", "open-stream-replay"]
+Mode = Literal["causal-prefix", "open-stream-replay"]
 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(
@@ -132,6 +173,7 @@ class OpenMediaPieces:
     payload_embeddings: torch.Tensor
     payload_position_ids: torch.LongTensor
     closure_ids: torch.LongTensor
+    closure_position_ids: torch.LongTensor
 
 
 class FixedGridAV:
@@ -167,19 +209,22 @@ class FixedGridAV:
         )
 
     def frames(self, start: float, end: float) -> list[Any]:
-        """Return fixed-grid frames whose timestamps lie in [start, end]."""
+        """Return fixed-grid frames whose timestamps lie in [start, end)."""
 
         import numpy as np
         from PIL import Image
 
         first = max(0, int(math.ceil(start * self.fps - 1e-7)))
-        last = min(
-            self.num_frames - 1,
-            int(math.floor(end * 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:
-            # This can occur only when a caller starts between global visual
-            # ticks. Reuse the most recent past frame, never ceil into future.
+            # 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()
@@ -198,76 +243,6 @@ class FixedGridAV:
         return self.audio[first:last]
 
 
-class TurnPredictionHead(nn.Module):
-    """Map one Omni state to EgoSpeak-compatible [10, 3] logits."""
-
-    def __init__(
-        self,
-        hidden_size: int = 2048,
-        num_horizons: int = 10,
-        num_classes: int = 3,
-        dropout: float = 0.1,
-    ) -> None:
-        super().__init__()
-        self.num_horizons = num_horizons
-        self.num_classes = num_classes
-        self.net = nn.Sequential(
-            nn.LayerNorm(hidden_size),
-            nn.Linear(hidden_size, hidden_size // 2),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(hidden_size // 2, num_horizons * num_classes),
-        )
-
-    def forward(self, states: torch.Tensor) -> torch.Tensor:
-        logits = self.net(states)
-        return logits.view(*states.shape[:-1], self.num_horizons, self.num_classes)
-
-
-def egospeak_multihorizon_loss(
-    logits: torch.Tensor,
-    targets: torch.Tensor,
-    *,
-    reduction: Literal["mean", "sum", "none"] = "mean",
-) -> torch.Tensor:
-    """Reproduce EgoSpeak's anticipation loss for multi-positive targets.
-
-    ``logits`` and ``targets`` both have shape ``[..., 10, 3]``. EasyCom can
-    label camera-wearer and other speech simultaneously as ``[0, 1, 1]``.
-    MiniROAD L2-normalizes that row and applies a log-softmax loss.
-    """
-
-    if logits.shape != targets.shape:
-        raise ValueError(
-            f"logits/targets shape mismatch: {tuple(logits.shape)} vs "
-            f"{tuple(targets.shape)}"
-        )
-    if logits.shape[-1] != 3:
-        raise ValueError(f"Expected three classes, got {logits.shape[-1]}")
-    targets = nn.functional.normalize(
-        targets.to(device=logits.device, dtype=logits.dtype),
-        p=2,
-        dim=-1,
-    )
-    loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1)
-    if reduction == "mean":
-        return loss.mean()
-    if reduction == "sum":
-        return loss.sum()
-    if reduction == "none":
-        return loss
-    raise ValueError(f"Unknown reduction: {reduction}")
-
-
-def multihorizon_cross_entropy(
-    logits: torch.Tensor,
-    targets: torch.Tensor,
-) -> torch.Tensor:
-    """Backward-compatible name for :func:`egospeak_multihorizon_loss`."""
-
-    return egospeak_multihorizon_loss(logits, targets)
-
-
 def _require_integer_ratio(numerator: float, denominator: float, name: str) -> int:
     ratio = numerator / denominator
     rounded = round(ratio)
@@ -437,16 +412,23 @@ def _prepare_conversation(
         # 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
-    batch = processor(
-        text=[text],
-        audio=audios,
-        images=images,
-        videos=videos,
-        padding=True,
-        return_tensors="pt",
-        use_audio_in_video=use_audio_in_video,
+    processor_inputs: dict[str, Any] = {
+        "text": [text],
+        "audio": audios,
+        "videos": videos,
+        "padding": True,
+        "use_audio_in_video": use_audio_in_video,
         **processor_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 = batch.convert_to_tensors("pt")
     # BatchFeature.to(dtype) changes floating tensors only; IDs stay integral.
     batch = batch.to(model.device).to(model.dtype)
     return dict(batch)
@@ -662,18 +644,10 @@ def _open_media_pieces(
         payload_embeddings=embeddings[:, first : last + 1, :],
         payload_position_ids=positions[:, :, first : last + 1],
         closure_ids=closure_ids,
+        closure_position_ids=positions[:, :, last + 1 :],
     )
 
 
-def _positions_at_offset(
-    position_ids: torch.LongTensor,
-    offset: int,
-) -> torch.LongTensor:
-    """Preserve local 3D geometry while moving its minimum to ``offset``."""
-
-    return position_ids - position_ids.min() + offset
-
-
 def _tmrope_positions_per_chunk(
     model: Qwen2_5OmniThinkerForConditionalGeneration,
 ) -> int:
@@ -823,13 +797,22 @@ def _read_open_state(
     model: Qwen2_5OmniThinkerForConditionalGeneration,
     open_media_cache: Any,
     *,
-    next_rope_offset: int,
     closure_ids: torch.LongTensor,
+    closure_position_ids: torch.LongTensor,
     query_ids: torch.LongTensor,
 ) -> 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,
@@ -837,13 +820,24 @@ def _read_open_state(
     past_length = _cache_length(branch_cache)
     readout_length = readout_ids.shape[1]
     embeddings = model.get_input_embeddings()(readout_ids)
-    scalar_positions = torch.arange(
-        next_rope_offset,
-        next_rope_offset + readout_length,
+
+    # 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,
     )
-    position_ids = scalar_positions.view(1, 1, -1).expand(
-        3, readout_ids.shape[0], -1
+    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,
@@ -876,32 +870,41 @@ def _append_open_pieces(
     past_key_values: Any | None,
     rope_offset: int,
     include_prefix: bool,
-) -> tuple[Any, int]:
-    """Append one AV payload, adding the chat/media prefix exactly once."""
+) -> 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_ids = _positions_at_offset(
-            pieces.payload_position_ids,
-            rope_offset,
+        position_shift = rope_offset - int(
+            pieces.payload_position_ids.min().item()
         )
-    return _append_embeddings(
+        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:
@@ -982,18 +985,34 @@ class CausalPrefixExtractor:
         start: float,
         end: float,
         use_audio_in_video: bool,
+        show_progress: bool = True,
     ) -> ExtractionResult:
         times = _time_grid(start, end, self.config.step_seconds)
         states = []
-        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())
+        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),
@@ -1001,202 +1020,6 @@ class CausalPrefixExtractor:
         )
 
 
-class CommitReplayExtractor:
-    """Cache complete 2 s blocks and replay only the current causal prefix."""
-
-    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)
-        self.reference = CausalPrefixExtractor(
-            model, processor, config, fixed_grid_av=fixed_grid_av
-        )
-
-    def _prepare_one(
-        self,
-        *,
-        video_path: str | Path | None,
-        audio_path: str | Path | None,
-        interval: tuple[float, float],
-        use_audio_in_video: bool,
-        turn_index: int,
-    ) -> dict[str, torch.Tensor]:
-        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,
-        )
-        return _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,
-        )
-
-    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,
-    ) -> tuple[list[float], list[torch.Tensor], list[dict[str, Any]]]:
-        times = _time_grid(segment_start, segment_end, self.config.step_seconds)
-        committed_cache = None
-        committed_rope_offset = 0
-        block_start = segment_start
-        states: list[torch.Tensor] = []
-        checks: list[dict[str, Any]] = []
-
-        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
-
-            # DynamicCache is mutated by forward.  Every unfinished prefix must
-            # branch from the identical committed history.
-            branch_cache = clone_cache(committed_cache)
-            inputs = self._prepare_one(
-                video_path=video_path,
-                audio_path=audio_path,
-                interval=(block_start, current_time),
-                use_audio_in_video=use_audio_in_video,
-                turn_index=round(
-                    (block_start - segment_start)
-                    / self.config.native_block_seconds
-                ),
-            )
-            media_cache, next_rope_offset = _append_media_turn(
-                self.model,
-                inputs,
-                past_key_values=branch_cache,
-                rope_offset=committed_rope_offset,
-                use_audio_in_video=use_audio_in_video,
-            )
-            state = _read_state(
-                self.model,
-                media_cache,
-                next_rope_offset=next_rope_offset,
-                query_ids=self.query_ids,
-            )
-            states.append(state.cpu())
-
-            global_step = global_step_offset + local_step
-            if verify_every > 0 and global_step % verify_every == 0:
-                intervals = _block_intervals(
-                    segment_start,
-                    current_time,
-                    self.config.native_block_seconds,
-                )
-                exact = self.reference.state_for_intervals(
-                    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.reference.state_for_intervals(
-                    video_path=video_path,
-                    audio_path=audio_path,
-                    intervals=[(segment_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):
-                # media_cache excludes the readout query, so it is safe to keep.
-                committed_cache = media_cache
-                committed_rope_offset = next_rope_offset
-                block_start = block_end
-
-        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,
-    ) -> ExtractionResult:
-        all_times: list[float] = []
-        all_states: list[torch.Tensor] = []
-        all_checks: list[dict[str, Any]] = []
-        segment_start = start
-        global_steps = 0
-
-        # A bounded cache is implemented as non-overlapping memory segments.
-        # Predictions are still emitted for the entire session, but history is
-        # reset at each boundary.  Rolling KV eviction is intentionally not
-        # attempted because it also requires rebasing multimodal RoPE.
-        while segment_start < end - 1e-7:
-            segment_end = min(segment_start + self.config.memory_seconds, end)
-            times, states, checks = self._extract_segment(
-                video_path=video_path,
-                audio_path=audio_path,
-                segment_start=segment_start,
-                segment_end=segment_end,
-                use_audio_in_video=use_audio_in_video,
-                verify_every=verify_every,
-                verify_continuous_every=verify_continuous_every,
-                global_step_offset=global_steps,
-            )
-            all_times.extend(times)
-            all_states.extend(states)
-            all_checks.extend(checks)
-            global_steps += len(times)
-            segment_start = segment_end
-
-        return ExtractionResult(
-            times=torch.tensor(all_times, dtype=torch.float64),
-            states=torch.cat(all_states, dim=0),
-            verification=all_checks,
-        )
-
-
 class OpenStreamReplayExtractor:
     """Cache completed AV payloads inside one still-open multimodal turn."""
 
@@ -1271,7 +1094,13 @@ class OpenStreamReplayExtractor:
     ) -> torch.Tensor:
         """Replay the same open stream from an empty cache for verification."""
 
-        media_cache, next_offset, _, closure_ids = self._cache_from_intervals(
+        (
+            media_cache,
+            _,
+            _,
+            closure_ids,
+            closure_position_ids,
+        ) = self._cache_from_intervals(
             video_path=video_path,
             audio_path=audio_path,
             intervals=intervals,
@@ -1280,8 +1109,8 @@ class OpenStreamReplayExtractor:
         return _read_open_state(
             self.model,
             media_cache,
-            next_rope_offset=next_offset,
             closure_ids=closure_ids,
+            closure_position_ids=closure_position_ids,
             query_ids=self.query_ids,
         )
 
@@ -1292,7 +1121,7 @@ class OpenStreamReplayExtractor:
         audio_path: str | Path | None,
         intervals: list[tuple[float, float]],
         use_audio_in_video: bool,
-    ) -> tuple[Any, int, int, torch.LongTensor]:
+    ) -> tuple[Any, int, int, torch.LongTensor, torch.LongTensor]:
         """Build an open payload cache and expose its local TMRoPE base."""
 
         if not intervals:
@@ -1300,6 +1129,7 @@ class OpenStreamReplayExtractor:
         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):
@@ -1312,7 +1142,11 @@ class OpenStreamReplayExtractor:
             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 = _append_open_pieces(
+            (
+                media_cache,
+                next_offset,
+                closure_position_ids,
+            ) = _append_open_pieces(
                 self.model,
                 pieces,
                 past_key_values=media_cache,
@@ -1322,9 +1156,17 @@ class OpenStreamReplayExtractor:
             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
+        return (
+            media_cache,
+            next_offset,
+            payload_base,
+            closure_ids,
+            closure_position_ids,
+        )
 
     def _extract_segment(
         self,
@@ -1337,6 +1179,7 @@ class OpenStreamReplayExtractor:
         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
@@ -1373,6 +1216,7 @@ class OpenStreamReplayExtractor:
                     _,
                     payload_base,
                     _,
+                    _,
                 ) = self._cache_from_intervals(
                     video_path=video_path,
                     audio_path=audio_path,
@@ -1394,7 +1238,11 @@ class OpenStreamReplayExtractor:
             )
             payload_offset = payload_base + block_index * positions_per_chunk
             branch_cache = clone_cache(committed_cache)
-            media_cache, next_offset = _append_open_pieces(
+            (
+                media_cache,
+                _,
+                closure_position_ids,
+            ) = _append_open_pieces(
                 self.model,
                 pieces,
                 past_key_values=branch_cache,
@@ -1404,8 +1252,8 @@ class OpenStreamReplayExtractor:
             state = _read_open_state(
                 self.model,
                 media_cache,
-                next_rope_offset=next_offset,
                 closure_ids=pieces.closure_ids,
+                closure_position_ids=closure_position_ids,
                 query_ids=self.query_ids,
             )
             states.append(state.cpu())
@@ -1454,6 +1302,11 @@ class OpenStreamReplayExtractor:
                 committed_cache = media_cache
                 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(
@@ -1466,17 +1319,29 @@ class OpenStreamReplayExtractor:
         use_audio_in_video: bool,
         verify_every: int = 0,
         verify_continuous_every: int = 0,
+        show_progress: bool = True,
     ) -> ExtractionResult:
-        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,
-        )
+        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),
@@ -1485,7 +1350,7 @@ class OpenStreamReplayExtractor:
 
 
 def _block_intervals(start: float, end: float, block_seconds: float) -> list[tuple[float, float]]:
-    """Exact media-turn sequence represented by a commit/replay cache."""
+    """Exact AV payload intervals represented by an open-stream cache."""
 
     intervals = []
     cursor = start
@@ -1637,7 +1502,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--output", type=Path, required=True)
     parser.add_argument(
         "--mode",
-        choices=("causal-prefix", "commit-replay", "open-stream-replay"),
+        choices=("causal-prefix", "open-stream-replay"),
         default="open-stream-replay",
     )
     parser.add_argument("--model-id", default=MODEL_ID)
@@ -1661,6 +1526,11 @@ def parse_args() -> argparse.Namespace:
         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",
@@ -1674,6 +1544,7 @@ def parse_args() -> argparse.Namespace:
 
 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")
@@ -1683,7 +1554,6 @@ def main() -> None:
         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:
@@ -1758,19 +1628,7 @@ def main() -> None:
             start=args.start,
             end=end,
             use_audio_in_video=args.use_audio_in_video,
-        )
-    elif args.mode == "commit-replay":
-        extractor = CommitReplayExtractor(
-            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,
         )
     else:
         extractor = OpenStreamReplayExtractor(
@@ -1784,6 +1642,7 @@ def main() -> None:
             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] = {
@@ -1792,11 +1651,7 @@ def main() -> None:
         "stream_representation": (
             "single_causal_window"
             if args.mode == "causal-prefix"
-            else (
-                "sequence_of_2s_media_turns"
-                if args.mode == "commit-replay"
-                else "single_open_media_turn_with_2s_payload_commits"
-            )
+            else "single_open_media_turn_with_2s_payload_commits"
         ),
         "config": asdict(config),
         "start": args.start,
@@ -1808,7 +1663,7 @@ def main() -> None:
         "session_metadata": session_metadata,
         "input_video_fps": input_video_fps,
         "video_sampling_grid": (
-            "explicit_fixed_past_frames_with_duplication"
+            "explicit_fixed_half_open_past_frames_with_tail_duplication"
             if fixed_grid_av is not None
             else "qwen_utility_adaptive_sampling"
         ),