From 20446c0bd8177c1baa6a3880e986699f8146a2ef Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Mon, 24 Aug 2026 09:52:56 +0800 Subject: [PATCH] Add overlap working-memory ablation and LoRA epoch2 record --- README.md | 171 ++++++++++ .../README.md | 77 +++++ .../metrics.csv | 21 ++ .../summary.json | 70 ++++ streaming_turn_model.py | 300 +++++++++++++++++- streaming_turn_model_directly.py | 220 ++++++++++++- 6 files changed, 851 insertions(+), 8 deletions(-) create mode 100644 results/easycom_omni_lora_sparse_epoch2_2026-08-24/README.md create mode 100644 results/easycom_omni_lora_sparse_epoch2_2026-08-24/metrics.csv create mode 100644 results/easycom_omni_lora_sparse_epoch2_2026-08-24/summary.json diff --git a/README.md b/README.md index acd8b73..2b26a60 100644 --- a/README.md +++ b/README.md @@ -1104,6 +1104,24 @@ This control should approach the previously reported frozen 8-second score of difference between dense-cache reconstruction and the earlier open-stream feature files; a close match isolates the regression to LoRA/head training. +### Recorded sparse-LoRA results + +The current single-run EasyCom results confirm negative transfer from this +sparse LoRA recipe: + +| Checkpoint | Full-timeline mean mAP | Strict causal mean mAP | +|---|---:|---:| +| Frozen Thinker dense-cache control | 58.16091 | 58.17146 | +| Sparse LoRA epoch1 | 56.52480 | 56.53450 | +| Sparse LoRA weights-only epoch2 | 54.19850 | 54.20711 | + +Epoch2 is a weights-only continuation: adapter/head weights resume, but AdamW +and the scheduler restart. It is 2.32630 mAP points below epoch1 and 3.96241 +below the frozen dense-cache control in the full-timeline protocol. The largest +regression is camera-wearer speaking. Do not continue the same recipe to a +third run. The complete snapshot and per-class metrics are archived in +[`results/easycom_omni_lora_sparse_epoch2_2026-08-24/`](results/easycom_omni_lora_sparse_epoch2_2026-08-24/). + For a speed-only alternative, `train_omni_turn_lora.py` supports a two-GPU producer-consumer benchmark. The producer GPU runs the frozen AV towers and places current-tail embeddings in a bounded CPU queue while the consumer GPU @@ -1233,6 +1251,159 @@ per-class AP over identical seeds. A useful result requires the fused model to beat both controls consistently; otherwise a gain could merely come from the extra fusion parameters or the stronger recent-window shortcut. +### Frozen serialized-recap ablation (implemented) + +Before adding a learned two-branch fusion module, test a stricter frozen-Omni +variant with exactly the same 2048-dimensional readout and prediction head as +the current 8-second baseline. Mode `overlap-working-memory` serializes the two +timescales inside one still-open Thinker media turn: + +```text +total AV presentation budget = 8 s +fixed history budget = 6 s = three cached 2 s blocks +working recap = 2 s = [t-2,t] + +t = 20.4 s: +fixed: [14,16] [16,18] [18,20] +working: [18.4,20.4] + -> one frozen Thinker state [2048] +``` + +The overlap is intentional. It presents recent evidence twice but does not add +AV-token duration beyond the matched 8-second budget. Thus this ablation trades +older unique context for stronger recent evidence instead of silently giving +the candidate more input tokens. Fixed blocks retain the global half-open 2 FPS +grid. The working recap uses an endpoint-aligned 2 FPS grid and includes the +current causal frame: + +```text +fixed [18,20): 18.0 18.5 19.0 19.5 +working [18.4,20.4]: 18.9 19.4 19.9 20.4 +``` + +Keeping four frames means the oldest visual tick is replaced by the current +frame rather than adding a fifth frame; this isolates endpoint recency without +increasing visual tokens. Audio remains the full, exactly half-open +`[start,end)` slice. Consequently, call this an endpoint-aligned recent-video +intervention, not an exact timestamp-preserving resampling of both modalities. +The recap is serialized as the next TMRoPE chunk; it is deliberately a +recent-evidence recap, not a claim that duplicated samples occur later in +physical time. This mode is an ablation, not native Qwen streaming semantics, +and the Omni backbone stays frozen. + +One deliberate side effect is a 2-second phase cycle: overlap falls toward +0.2 seconds just before a fixed boundary and jumps to 2.0 seconds at the +boundary. Every output stores this value in +`overlap_working_memory_policy.layout`. If aggregate mAP changes, also inspect +performance by `step_index % 10`; a boundary-only gain or failure is a chunk +phase artifact rather than general evidence for better recent semantics. + +Run a 24-second L20 smoke test first: + +```bash +mkdir -p features/easycom_overlap_working_8s_smoke + +CUDA_VISIBLE_DEVICES=3 \ +OMP_NUM_THREADS=4 \ +MKL_NUM_THREADS=4 \ +OPENBLAS_NUM_THREADS=4 \ +PYTORCH_ALLOC_CONF=expandable_segments:True \ +python streaming_turn_model_directly.py \ + --model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \ + --easycom-session-dir /data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed/Session_7 \ + --end 24 \ + --mode overlap-working-memory \ + --memory-seconds 8 \ + --working-memory-seconds 2 \ + --video-fps 2 \ + --frame-cache-size 128 \ + --decoder decord \ + --decoder-threads 1 \ + --max-open-readers 1 \ + --dtype bfloat16 \ + --flash-attention-2 \ + --device cuda:0 \ + --verify-every 10 \ + --output features/easycom_overlap_working_8s_smoke/Session_7_24s.pt +``` + +Expected shape is `(120, 2048)`. `overlap_cache_replay` should be numerically +exact or effectively exact. Inspect the saved timing policy with: + +```bash +python - <<'PY' +import torch + +p = torch.load( + "features/easycom_overlap_working_8s_smoke/Session_7_24s.pt", + map_location="cpu", + weights_only=False, +) +print(p["mode"], tuple(p["states"].shape)) +print(p["overlap_working_memory_policy"]["working_media"]) +for row in p["overlap_working_memory_policy"]["layout"][98:102]: + print(row) +PY +``` + +For one complete session, replace `--end 24` with its target and use a +session-level output filename: + +```bash +mkdir -p /data1/yfzx/wangx/Datasets/easycom/omni_features/overlap_working_2fps_8s + +CUDA_VISIBLE_DEVICES=3 \ +OMP_NUM_THREADS=4 \ +MKL_NUM_THREADS=4 \ +OPENBLAS_NUM_THREADS=4 \ +PYTORCH_ALLOC_CONF=expandable_segments:True \ +python streaming_turn_model_directly.py \ + --model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \ + --easycom-session-dir /data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed/Session_7 \ + --target /data1/yfzx/wangx/Datasets/easycom/easycom_target_frame/Session_7.npy \ + --mode overlap-working-memory \ + --memory-seconds 8 \ + --working-memory-seconds 2 \ + --video-fps 2 \ + --frame-cache-size 128 \ + --decoder decord \ + --decoder-threads 1 \ + --max-open-readers 1 \ + --dtype bfloat16 \ + --flash-attention-2 \ + --device cuda:0 \ + --output /data1/yfzx/wangx/Datasets/easycom/omni_features/overlap_working_2fps_8s/Session_7.pt +``` + +After all train/test sessions have been extracted, train the unchanged frozen +head by setting only the expected feature mode: + +```bash +python train_turn_head.py \ + --feature-root /data1/yfzx/wangx/Datasets/easycom/omni_features/overlap_working_2fps_8s \ + --target-root /data1/yfzx/wangx/Datasets/easycom/easycom_target_frame \ + --output checkpoints/easycom_overlap_working_frozen_8s.pt \ + --expected-mode overlap-working-memory \ + --epochs 50 \ + --batch-size 512 \ + --device cuda:0 + +python infer_turn_head.py \ + --feature-root /data1/yfzx/wangx/Datasets/easycom/omni_features/overlap_working_2fps_8s \ + --target-root /data1/yfzx/wangx/Datasets/easycom/easycom_target_frame \ + --checkpoint checkpoints/easycom_overlap_working_frozen_8s.pt \ + --output-dir results/easycom_overlap_working_frozen_8s \ + --expected-mode overlap-working-memory \ + --device cuda:0 +``` + +The decisive comparison is against the frozen single-state 8-second baseline +using the same split, head, optimization, and seeds. Report the full +per-horizon curve: the hypothesis predicts the clearest gain at 0.2--1.0 +seconds; degradation at 1.2--2.0 seconds would indicate that repeated local +cues displaced too much useful long context. Only if this parameter-matched +test is positive should the learned two-branch fusion experiment be added. + ## Single-window feature inspection `qwen_omni_features.py` is useful for inspecting a short window or layer, but diff --git a/results/easycom_omni_lora_sparse_epoch2_2026-08-24/README.md b/results/easycom_omni_lora_sparse_epoch2_2026-08-24/README.md new file mode 100644 index 0000000..f1262ec --- /dev/null +++ b/results/easycom_omni_lora_sparse_epoch2_2026-08-24/README.md @@ -0,0 +1,77 @@ +# EasyCom: sparse Thinker-LoRA weights-only continuation (epoch2 snapshot) + +Recorded on **2026-08-24** from the remote EasyCom evaluation console. This is +a single-run experiment snapshot, not a multi-seed paper result. + +## Experiment contract + +- Dataset: EasyCom with EgoSpeak's official split. +- Train sessions: 7, 11, 10, 6, 9, 4, 8, 5, and 12. +- Test sessions: 3, 2, and 1 (`n=23048` full-timeline rows). +- Backbone: Qwen2.5-Omni-3B Thinker with LoRA plus the `10 x 3` turn head. +- Context: 8 seconds; 2 FPS video; predictions and labels at 5 Hz. +- Training sampling: transition-sparse anchors (ordinary 1 Hz plus dense 5 Hz + anchors near target transitions). +- Evaluation: dense 5 Hz AV cache on every test timestamp. +- Continuation: adapter and head weights loaded from the epoch1 checkpoint, + but optimizer moments and the learning-rate scheduler were reinitialized. + Therefore, "epoch2" here means a weights-only second training run, not an + exact optimizer-state resume. +- Reported checkpoint directory: + `checkpoints/omni_turn_lora_easycom_sparse_8s_epoch2_ddp4/`. + +## Summary + +All values are percentages. + +| Checkpoint | Full-timeline mean mAP | Strict causal mean mAP | +|---|---:|---:| +| Frozen Thinker dense-cache control | 58.16091 | 58.17146 | +| Sparse LoRA epoch1 | 56.52480 | 56.53450 | +| Sparse LoRA weights-only epoch2 | **54.19850** | **54.20711** | +| Epoch2 minus epoch1 | -2.32630 | -2.32739 | +| Epoch2 minus frozen control | -3.96241 | -3.96435 | + +The second run is materially worse than both epoch1 and the frozen control. +This is not consistent with normal evaluator noise. + +## Per-horizon results + +| Horizon | Full mAP | Strict mAP | Background AP | Wearer AP | Others AP | +|---:|---:|---:|---:|---:|---:| +| 0.2 s | 65.55214 | 65.55385 | 49.03788 | 65.76312 | 81.85542 | +| 0.4 s | 60.96517 | 60.96812 | 43.06922 | 61.14251 | 78.68378 | +| 0.6 s | 56.79766 | 56.80301 | 39.37427 | 55.10528 | 75.91344 | +| 0.8 s | 54.11741 | 54.12457 | 36.94546 | 51.56720 | 73.83955 | +| 1.0 s | 52.09538 | 52.10400 | 35.40706 | 48.38835 | 72.49075 | +| 1.2 s | 51.35516 | 51.36504 | 35.30874 | 46.86261 | 71.89412 | +| 1.4 s | 50.99774 | 51.00914 | 35.29891 | 45.48818 | 72.20613 | +| 1.6 s | 50.47830 | 50.49067 | 34.98109 | 44.74277 | 71.71105 | +| 1.8 s | 50.00999 | 50.02296 | 34.64341 | 43.86161 | 71.52496 | +| 2.0 s | 49.61605 | 49.62970 | 34.26290 | 43.14791 | 71.43735 | + +The AP columns above are from the full-timeline protocol. Complete AP values +for both protocols are stored in [`metrics.csv`](metrics.csv), and +[`summary.json`](summary.json) stores the machine-readable aggregate result. + +## Interpretation and decision + +The largest class-specific regression is camera-wearer speaking. Relative to +epoch1, wearer AP falls from 68.77118 to 65.76312 at 0.2 seconds and from +47.53875 to 43.14791 at 2.0 seconds. Background and other-speaker AP decline +less at the long horizon. The pattern indicates negative transfer in the +wearer-speaking representation rather than a uniform metric offset. + +Likely contributors are sparse-train/dense-evaluation distribution mismatch, +too much LoRA capacity for EasyCom, and restarting the learning-rate schedule +during weights-only continuation. Do not run a third continuation epoch with +the same recipe. Keep the frozen 8-second result as the main representation +baseline and next test the parameter-matched frozen overlap-working-memory +ablation. + +## Provenance + +- Source: user-reported `eval_omni_turn_lora.py` console output. +- Raw remote `metrics.json` archived here: no; values were transcribed. +- Multi-seed: no. +- Publication status: diagnostic negative result only. diff --git a/results/easycom_omni_lora_sparse_epoch2_2026-08-24/metrics.csv b/results/easycom_omni_lora_sparse_epoch2_2026-08-24/metrics.csv new file mode 100644 index 0000000..6bb29b9 --- /dev/null +++ b/results/easycom_omni_lora_sparse_epoch2_2026-08-24/metrics.csv @@ -0,0 +1,21 @@ +protocol,horizon,seconds,background_ap_percent,camera_wearer_ap_percent,other_people_ap_percent,map_percent,num_steps +egospeak_full_timeline,1,0.2,49.03788,65.76312,81.85542,65.55214,23048 +egospeak_full_timeline,2,0.4,43.06922,61.14251,78.68378,60.96517,23048 +egospeak_full_timeline,3,0.6,39.37427,55.10528,75.91344,56.79766,23048 +egospeak_full_timeline,4,0.8,36.94546,51.56720,73.83955,54.11741,23048 +egospeak_full_timeline,5,1.0,35.40706,48.38835,72.49075,52.09538,23048 +egospeak_full_timeline,6,1.2,35.30874,46.86261,71.89412,51.35516,23048 +egospeak_full_timeline,7,1.4,35.29891,45.48818,72.20613,50.99774,23048 +egospeak_full_timeline,8,1.6,34.98109,44.74277,71.71105,50.47830,23048 +egospeak_full_timeline,9,1.8,34.64341,43.86161,71.52496,50.00999,23048 +egospeak_full_timeline,10,2.0,34.26290,43.14791,71.43735,49.61605,23048 +causal_valid_only,1,0.2,49.03788,65.76312,81.86055,65.55385,23045 +causal_valid_only,2,0.4,43.06922,61.14251,78.69264,60.96812,23042 +causal_valid_only,3,0.6,39.37427,55.10988,75.92489,56.80301,23039 +causal_valid_only,4,0.8,36.94546,51.57526,73.85298,54.12457,23036 +causal_valid_only,5,1.0,35.40706,48.39889,72.50604,52.10400,23033 +causal_valid_only,6,1.2,35.30874,46.87568,71.91070,51.36504,23030 +causal_valid_only,7,1.4,35.29891,45.50340,72.22511,51.00914,23027 +causal_valid_only,8,1.6,34.98337,44.75739,71.73126,50.49067,23024 +causal_valid_only,9,1.8,34.64563,43.87551,71.54773,50.02296,23021 +causal_valid_only,10,2.0,34.26505,43.16123,71.46280,49.62970,23018 diff --git a/results/easycom_omni_lora_sparse_epoch2_2026-08-24/summary.json b/results/easycom_omni_lora_sparse_epoch2_2026-08-24/summary.json new file mode 100644 index 0000000..24b54e2 --- /dev/null +++ b/results/easycom_omni_lora_sparse_epoch2_2026-08-24/summary.json @@ -0,0 +1,70 @@ +{ + "schema_version": 1, + "recorded_date": "2026-08-24", + "status": "single_run_snapshot", + "dataset": "EasyCom", + "split": { + "train_sessions": [7, 11, 10, 6, 9, 4, 8, 5, 12], + "test_sessions": [3, 2, 1] + }, + "model": "Qwen2.5-Omni-3B Thinker LoRA", + "checkpoint": "checkpoints/omni_turn_lora_easycom_sparse_8s_epoch2_ddp4", + "continuation": "adapter_and_head_weights_only_optimizer_and_scheduler_reinitialized", + "training_anchor_policy": "transition_sparse_ordinary_1hz_plus_transition_5hz", + "evaluation_anchor_policy": "dense_5hz", + "prediction_step_seconds": 0.2, + "video_fps": 2, + "memory_seconds": 8, + "metrics_percent": { + "epoch2": { + "egospeak_full_timeline_mean_map": 54.19850, + "causal_valid_only_mean_map": 54.20711, + "egospeak_full_timeline_horizon_map": [ + 65.55214, + 60.96517, + 56.79766, + 54.11741, + 52.09538, + 51.35516, + 50.99774, + 50.47830, + 50.00999, + 49.61605 + ], + "causal_valid_only_horizon_map": [ + 65.55385, + 60.96812, + 56.80301, + 54.12457, + 52.10400, + 51.36504, + 51.00914, + 50.49067, + 50.02296, + 49.62970 + ] + }, + "epoch1_reference": { + "egospeak_full_timeline_mean_map": 56.52480, + "causal_valid_only_mean_map": 56.53450 + }, + "frozen_dense_cache_control": { + "egospeak_full_timeline_mean_map": 58.16091, + "causal_valid_only_mean_map": 58.17146 + }, + "epoch2_minus_epoch1": { + "egospeak_full_timeline_mean_map": -2.32630, + "causal_valid_only_mean_map": -2.32739 + }, + "epoch2_minus_frozen_dense_cache_control": { + "egospeak_full_timeline_mean_map": -3.96241, + "causal_valid_only_mean_map": -3.96435 + } + }, + "decision": "stop_same_recipe_continuation_and_test_frozen_overlap_working_memory", + "provenance": { + "source": "User-reported remote eval_omni_turn_lora.py console tables", + "raw_metrics_json_archived_here": false, + "multi_seed": false + } +} diff --git a/streaming_turn_model.py b/streaming_turn_model.py index 669666e..5e7df2d 100644 --- a/streaming_turn_model.py +++ b/streaming_turn_model.py @@ -1,11 +1,13 @@ """Omni-only causal turn-taking features at 5 Hz. -This module implements two inference paths for Qwen2.5-Omni-3B: +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. @@ -54,7 +56,11 @@ from transformers import ( ) from transformers.cache_utils import DynamicCache -Mode = Literal["causal-prefix", "open-stream-replay"] +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", @@ -1642,6 +1648,296 @@ class OpenStreamReplayExtractor: ) +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, + while the recap contains 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, + ) -> 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) + # 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, + ) + + 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, + ) + 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, + ) + 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.""" diff --git a/streaming_turn_model_directly.py b/streaming_turn_model_directly.py index e0cbdd4..db4ee8e 100644 --- a/streaming_turn_model_directly.py +++ b/streaming_turn_model_directly.py @@ -496,6 +496,28 @@ class CausalSourceAV: """Map one global output tick to a source frame at or before it.""" global_time = grid_index / self.fps + clip_index, frame_index, source_time = self._source_frame_for_time( + global_time, + count_tail_reuse=count_tail_reuse, + ) + return clip_index, frame_index, global_time, source_time + + def _source_frame_for_time( + self, + global_time: float, + *, + count_tail_reuse: bool = True, + ) -> tuple[int, int, float]: + """Map an arbitrary causal timestamp to the latest observed frame.""" + + if global_time < -EPSILON: + raise ValueError(f"Negative video timestamp: {global_time}") + if global_time > self.required_duration + EPSILON: + raise ValueError( + f"Video timestamp {global_time:g}s exceeds required duration " + f"{self.required_duration:g}s" + ) + global_time = min(max(0.0, float(global_time)), self.required_duration) clip_index, clip, local_time = self._clip_for_time(global_time) frame_index = int( np.searchsorted(clip.source_frame_starts, local_time + EPSILON, side="right") @@ -510,7 +532,14 @@ class CausalSourceAV: ) if count_tail_reuse and local_time > float(clip.source_frame_starts[-1]) + EPSILON: self._tail_frame_reuses += 1 - return clip_index, frame_index, global_time, source_time + return clip_index, frame_index, source_time + + def frame_at_or_before(self, global_time: float) -> Any: + """Decode the latest source frame whose PTS is not later than ``time``.""" + + clip_index, frame_index, _ = self._source_frame_for_time(global_time) + self._request_count += 1 + return self._load_frame(clip_index, frame_index) def _load_frame(self, clip_index: int, frame_index: int) -> Any: from PIL import Image @@ -797,6 +826,89 @@ class CausalSourceAV: } +class EndpointAlignedAV: + """Working-window view whose final visual sample is exactly current time. + + Complete cached blocks use the source reader's fixed global half-open grid, + for example ``[18.0, 18.5, 19.0, 19.5]`` for ``[18,20)`` at 2 FPS. A + two-second working interval ending at ``t`` instead uses + ``[t-1.5, t-1.0, t-0.5, t]``. The number of visual samples therefore + stays matched while the working branch explicitly observes the current + causal frame. Audio remains the exact half-open slice ``[start,end)``. + """ + + def __init__( + self, + source: CausalSourceAV, + *, + max_interval_seconds: float, + ) -> None: + if max_interval_seconds <= 0: + raise ValueError("max_interval_seconds must be positive") + self.source = source + self.fps = source.fps + self.sample_rate = source.sample_rate + self.max_interval_seconds = float(max_interval_seconds) + max_frames = self.max_interval_seconds * self.fps + if not math.isclose(max_frames, round(max_frames), abs_tol=1e-6): + raise ValueError( + "working-memory seconds * video FPS must be integral, got " + f"{max_frames}" + ) + self.max_frames = int(round(max_frames)) + + def sample_times(self, start: float, end: float) -> list[float]: + if end <= start: + raise ValueError(f"Empty video interval [{start}, {end}]") + duration = end - start + if duration > self.max_interval_seconds + EPSILON: + raise ValueError( + f"Working interval is {duration:g}s, exceeds configured " + f"{self.max_interval_seconds:g}s" + ) + count = max(1, int(math.ceil(duration * self.fps - EPSILON))) + count = min(count, self.max_frames) + first = end - (count - 1) / self.fps + if first < start - EPSILON: + raise RuntimeError( + f"Endpoint-aligned grid starts before its interval: " + f"{first:g} < {start:g}" + ) + return [first + index / self.fps for index in range(count)] + + def frames(self, start: float, end: float) -> list[Any]: + frames = [ + self.source.frame_at_or_before(timestamp) + for timestamp in self.sample_times(start, end) + ] + if len(frames) % 2: + frames.append(frames[-1].copy()) + self.source._odd_frame_duplications += 1 + return frames + + def audio_slice(self, start: float, end: float) -> np.ndarray: + return self.source.audio_slice(start, end) + + def metadata(self) -> dict[str, Any]: + return { + "visual_sampling": "endpoint_aligned_latest_source_pts_not_after_tick", + "audio_interval": "half_open_[start,end)", + "working_memory_seconds": self.max_interval_seconds, + "output_video_fps": self.fps, + "full_window_frame_count": self.max_frames, + "full_window_relative_video_times": [ + -(self.max_frames - 1 - index) / self.fps + for index in range(self.max_frames) + ], + "includes_current_video_frame": True, + "oldest_visual_tick_replaced_by_current": True, + "timestamp_note": ( + "video is endpoint-aligned while audio retains the full " + "half-open working interval" + ), + } + + def _manifest_session_specs( manifest_path: Path, session_name: str, @@ -901,7 +1013,11 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument("--output", type=Path) parser.add_argument( "--mode", - choices=("causal-prefix", "open-stream-replay"), + choices=( + "causal-prefix", + "open-stream-replay", + "overlap-working-memory", + ), default="open-stream-replay", ) parser.add_argument("--model-id", default=core.MODEL_ID) @@ -909,6 +1025,15 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument("--end", type=float) parser.add_argument("--step-seconds", type=float, default=0.2) parser.add_argument("--native-block-seconds", type=float, default=2.0) + parser.add_argument( + "--working-memory-seconds", + type=float, + default=2.0, + help=( + "Sliding recent recap used by overlap-working-memory; its duration " + "is included in --memory-seconds rather than added on top" + ), + ) parser.add_argument( "--memory-seconds", type=float, @@ -1046,6 +1171,16 @@ def main( raise ValueError("model and processor must be supplied together") if args.verify_every < 0 or args.verify_continuous_every < 0: raise SystemExit("verification intervals must be >= 0") + if args.working_memory_seconds <= 0: + raise SystemExit("--working-memory-seconds must be positive") + if ( + args.mode == "overlap-working-memory" + and args.verify_continuous_every > 0 + ): + raise SystemExit( + "--verify-continuous-every is not defined for the intentional " + "overlap representation; use --verify-every" + ) if args.frame_cache_size < 0: raise SystemExit("--frame-cache-size must be >= 0") if args.output is None and not args.media_only: @@ -1131,6 +1266,11 @@ def main( return duration = end - args.start + if args.mode == "overlap-working-memory" and args.memory_seconds <= 0: + raise SystemExit( + "overlap-working-memory uses a matched finite presentation budget; " + "pass a positive --memory-seconds (recommended: 8)" + ) effective_memory = _effective_memory_seconds( args.memory_seconds, duration=duration, @@ -1167,6 +1307,8 @@ def main( model.requires_grad_(False) media_path_for_api = reader.clips[0].path + overlap_extractor = None + working_grid_av = None if args.mode == "causal-prefix": extractor = core.CausalPrefixExtractor( model, processor, config, fixed_grid_av=reader @@ -1179,7 +1321,7 @@ def main( use_audio_in_video=True, show_progress=not args.no_progress, ) - else: + elif args.mode == "open-stream-replay": extractor = core.OpenStreamReplayExtractor( model, processor, config, fixed_grid_av=reader ) @@ -1193,20 +1335,51 @@ def main( verify_continuous_every=args.verify_continuous_every, show_progress=not args.no_progress, ) + else: + working_grid_av = EndpointAlignedAV( + reader, + max_interval_seconds=args.working_memory_seconds, + ) + overlap_extractor = core.OverlappingWorkingMemoryExtractor( + model, + processor, + config, + working_memory_seconds=args.working_memory_seconds, + fixed_grid_av=reader, + working_grid_av=working_grid_av, + ) + result = overlap_extractor.extract( + video_path=media_path_for_api, + audio_path=None, + start=args.start, + end=end, + use_audio_in_video=True, + verify_every=args.verify_every, + show_progress=not args.no_progress, + ) if target_steps is not None and len(result.times) != target_steps: raise RuntimeError( f"Extractor returned {len(result.times)} states, expected {target_steps}" ) - persistent_memory = args.memory_seconds == 0 + persistent_memory = ( + args.mode == "open-stream-replay" and args.memory_seconds == 0 + ) payload: dict[str, Any] = { "model_id": args.model_id, "mode": args.mode, "stream_representation": ( "single_causal_prefix" if args.mode == "causal-prefix" - else "single_open_media_turn_with_2s_payload_commits" + else ( + "single_open_media_turn_with_2s_payload_commits" + if args.mode == "open-stream-replay" + else ( + "serialized_cached_complete_blocks_then_overlapping_" + "endpoint_aligned_recent_recap" + ) + ) ), "config": asdict(config), "requested_memory_seconds": args.memory_seconds, @@ -1221,7 +1394,9 @@ def main( "source_manifest_record": manifest_record, "direct_media": reader.metadata(), "video_sampling_grid": ( - "direct_pts_causal_global_grid_half_open_with_tail_duplication" + "fixed_global_half_open_plus_endpoint_aligned_working_recap" + if args.mode == "overlap-working-memory" + else "direct_pts_causal_global_grid_half_open_with_tail_duplication" ), "open_stream_policy": ( { @@ -1242,6 +1417,39 @@ def main( if args.mode == "open-stream-replay" else None ), + "overlap_working_memory_policy": ( + { + "total_presentation_seconds": config.memory_seconds, + "fixed_history_seconds": ( + config.memory_seconds - args.working_memory_seconds + ), + "fixed_block_seconds": config.native_block_seconds, + "working_memory_seconds": args.working_memory_seconds, + "fixed_blocks": ( + overlap_extractor.fixed_block_budget + if overlap_extractor is not None + else None + ), + "serialization": ( + "fixed complete AV payloads followed by overlapping " + "working AV payload in the same open media turn" + ), + "overlap_is_intentional": True, + "working_media": ( + working_grid_av.metadata() + if working_grid_av is not None + else None + ), + "layout": ( + overlap_extractor.layout_records + if overlap_extractor is not None + else None + ), + "backbone": "frozen", + } + if args.mode == "overlap-working-memory" + else None + ), "times": result.times, "states": result.states, "verification": result.verification,