From bde72487bb70cc03dd0932c9ebeebac055469e17 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Sun, 16 Aug 2026 00:02:58 +0800 Subject: [PATCH] Initial commit: Omni-only turn-taking experiment scaffold. Add EasyCom prep, Qwen2.5-Omni feature extraction, streaming turn model prototype, and docs for the ICASSP 2027 ego-speak study. Co-authored-by: Cursor --- README.md | 752 ++++++++ audit_egospeak_features.py | 169 ++ easycom_manifest.py | 321 ++++ prepare_easycom_session.py | 357 ++++ qwen_omni_features.py | 270 +++ reports/egospeak_easycom_feature_audit.json | 458 +++++ requirements.txt | 9 + smoke_test.py | 147 ++ streaming_turn_model.py | 1866 +++++++++++++++++++ 9 files changed, 4349 insertions(+) create mode 100644 README.md create mode 100644 audit_egospeak_features.py create mode 100644 easycom_manifest.py create mode 100644 prepare_easycom_session.py create mode 100644 qwen_omni_features.py create mode 100644 reports/egospeak_easycom_feature_audit.json create mode 100644 requirements.txt create mode 100644 smoke_test.py create mode 100644 streaming_turn_model.py diff --git a/README.md b/README.md new file mode 100644 index 0000000..5872967 --- /dev/null +++ b/README.md @@ -0,0 +1,752 @@ +# Omni-only Turn-Taking Experiment + +This project tests whether **Qwen2.5-Omni-3B can replace both the handcrafted +audio/visual feature extractors and the downstream RNN/Transformer used by +EgoSpeak** for egocentric turn-taking anticipation. + +The main experiment is deliberately simple: + +```text +causal audio + video up to time t + | + v + Qwen2.5-Omni-3B Thinker + | + fused state q_t [2048] + | + v + trainable prediction head + | + v + future logits [10 horizons, 3 classes] +``` + +There is **no RNN after Qwen**. The Omni Thinker is the temporal model. The +only supervised component in the first experiment is a small `10 x 3` +classification head. + +## Current status + +| File | Status | Purpose | +|---|---|---| +| `smoke_test.py` | Verified on one NVIDIA L20 | Text and 12-second EasyCom audio-video inference | +| `easycom_manifest.py` | Implemented, awaiting dataset-side run | Discover one-minute clips, preserve session timecodes, and check 5 Hz targets | +| `qwen_omni_features.py` | Implemented, not executed here | Load only the 3B Thinker and extract fused states from one AV window | +| `streaming_turn_model.py` | Prototype implemented, not executed here | Produce one causal Omni state every 200 ms, with reference and KV-cache modes | +| `TurnPredictionHead` | Defined, untrained | Map `[2048]` to `[10, 3]` logits | +| Head training and EgoSpeak evaluation script | Not yet implemented | Align saved states with `target_perframe`, train the head, and reproduce Table 2 mAP | + +The official full model, Thinker input path, embedded audio, video decoding, +and FlashAttention-2 have passed a 12-second EasyCom smoke test on one NVIDIA +L20. Streaming extraction, cache equivalence, target alignment, training, and +paper metrics remain unverified. + +## Why Omni-only instead of Omni + RNN + +EgoSpeak separately extracts visual and audio features and then uses GRU, +LSTR, or Mamba for temporal modeling: + +```text +ResNet50 video feature + wav2vec2 audio feature -> temporal model -> prediction +``` + +Qwen2.5-Omni already contains: + +- audio and visual encoders; +- joint audio-video temporal positions; +- a causal decoder-only Thinker with KV cache; +- pretrained acoustic, visual, and semantic knowledge. + +Adding a GRU after a continuously processed Omni sequence would model history +twice and would make it harder to attribute an improvement to Omni. The primary +comparison is therefore: + +```text +EgoSpeak: ResNet50 + wav2vec2 + GRU +Ours: Qwen2.5-Omni Thinker + small prediction head +``` + +An Omni-feature-plus-GRU model can be kept as an ablation, but it is not the +main architecture. + +## EgoSpeak contract that this project reproduces + +The inspected reference repository is `/Volumes/760P/Code/EgoSpeak`. +It contains no Qwen or Omni implementation and no raw-media feature extraction +scripts. It distributes precomputed arrays instead: + +```text +DATASET/ +|- rgb_kinetics_resnet50/ # [T, 2048] +|- audio_wav2vec2/ # [T, 5120] +`- target_perframe/ # [T, 3] +``` + +The useful part for this project is the exact task interface. + +### Time and classes + +Both EasyCom and the EgoSpeak `EGO4DSHUFFLE` split use **5 FPS**, so one task +timestep is **0.2 seconds**. The three output columns are: + +```text +0: background +1: camera wearer speaking +2: other people speaking +``` + +They are not strictly one-hot. Background is exclusive, but simultaneous +camera-wearer and other-person speech is stored as `[0, 1, 1]`. The inspected +EasyCom features contain 7,711 such overlap steps. Preserve the full +multi-positive target rather than applying `argmax`. + +Relevant EgoSpeak files: + +- `LSTR/data/data_info.json`: FPS, classes, and session splits; +- `MiniROAD/configs/EasyCom/miniroad_aa_easycom.yaml`: 128-step training + windows, stride 4, three classes, and anticipation length 10; +- `MiniROAD/datasets/dataset.py`: feature loading and future-label indexing; +- `MiniROAD/criterions/loss.py`: anticipation loss; +- `MiniROAD/trainer/eval.py`: per-horizon AP and mean anticipation mAP. + +### Future target alignment + +For current task index `i`, the prediction target is: + +```python +target[i + 1 : i + 11] +``` + +Thus the ten output rows represent: + +```text +t+0.2, t+0.4, t+0.6, ..., t+2.0 seconds +``` + +If `states[i]` is verified to correspond to EgoSpeak task index `i`, the +offline head-training tensors should be constructed as follows: + +```python +target = target_perframe # [T, 3], multi-positive +usable = min(len(states), len(target)) - 10 + +x = states[:usable] # [N, 2048] +y = stack( + [target[h : h + usable] for h in range(1, 11)], + axis=1, +) # [N, 10, 3] +``` + +Do not silently assume that the first saved Omni state matches label index +zero. EgoSpeak does not publish the script that generated its original +ResNet/wav2vec2 features, so the exact zero-time convention is not documented. +Before training, inspect a short clip with known speech boundaries and test the +candidate mappings `states[i] -> target[i]` and `states[i] -> target[i+1]`. +Store the selected offset in every feature file/checkpoint. + +### Splits + +Use the published split, not a random frame split. + +EasyCom: + +- train, preserving EgoSpeak's list: `Session_7`, `Session_11`, `Session_10`, + `Session_6`, `Session_9`, `Session_4`, `Session_8`, `Session_5`, `Session_12`; +- test, preserving EgoSpeak's list: `Session_3`, `Session_2`, `Session_1`. + +This is the **EgoSpeak comparison split**. The EasyCom dataset repository +itself recommends a different train/validation/test split; do not substitute +that split when reproducing EgoSpeak Table 2. + +Ego4D Audio-Visual Diarization: + +- use the `EGO4DSHUFFLE` key; +- 346 training clips; +- 87 test clips; +- each clip is approximately five minutes. + +Never mix timesteps from the same session across training and validation. + +## EasyCom raw-media layout and manifest + +The extracted release used here has this relevant structure: + +```text +extracted_easycom/ +`- Main/ + |- Video_Compressed/ + | |- Session_1/ + | | |- 00-00-000.mp4 + | | |- 01-00-000.mp4 + | | `- ... + | `- Session_2 ... Session_12/ + |- Glasses_Microphone_Array_Audio/ + |- Close_Microphone_Audio/ + |- Speech_Transcriptions/ + `- Metadata/ +``` + +`Video_Compressed` is the AR-glasses wearer's egocentric stream with embedded +stereo AAC audio. Each session is divided into approximately one-minute clips; +the `mm-ss-SSS` filename is the clip's start time on the original recording +timeline. Some intervals are absent because EasyCom redacted one-minute +segments. These raw time gaps are expected. + +EgoSpeak does not preserve those empty intervals. Its appendix states that it +crops every EasyCom clip to retain only the 1920-pixel-wide video region and +then merges the retained clips into one dense video per session. Therefore the +EgoSpeak-compatible label time for a clip is its cumulative position in the +concatenation, not the `mm-ss-SSS` source time. The manifest stores both. + +Build a raw-video manifest and, when the EgoSpeak labels are available, check +their duration at 5 Hz: + +```bash +python easycom_manifest.py \ + --dataset-root /data1/yfzx/wangx/Datasets/easycom/extracted_easycom \ + --target-root /path/to/EasyCom/target_perframe \ + --output manifests/easycom.json +``` + +Omit `--target-root` for the first media-only inventory. The manifest reports +raw-source discontinuities as informational and computes each clip's dense +EgoSpeak concatenation offset and corresponding 5 Hz target index. For the +audio-only ablation, use the same embedded glasses audio as A+V (extract it to +mono WAV if needed); do not switch to close microphones, because that changes +the sensor and makes the modality ablation unfair. + +Audit downloaded EgoSpeak arrays, and optionally compare every session length +against the raw manifest: + +```bash +python audit_egospeak_features.py \ + --feature-root /path/to/EgoSpeak/dataset/EasyCom \ + --raw-manifest manifests/easycom.json \ + --output reports/egospeak_easycom_audit.json +``` + +The inspected release contains 95,453 synchronized steps across all 12 +sessions: 5.302944 hours at 5 Hz. Every session has identical RGB, flow, audio, +and target length. Shapes are `[T,2048]`, `[T,2048]`, `[T,5120]`, and `[T,3]`. + +## Recommended comparison matrix + +| Experiment | Visual sampling | Audio | Output rate | +|---|---:|---|---:| +| EgoSpeak official A+V | 5 Hz ResNet features | 5 Hz wav2vec2 features | 5 Hz | +| Omni audio-only | none | raw glasses-wearer waveform | 5 Hz | +| Omni A+V primary | 2 FPS | raw embedded waveform | 5 Hz | +| Omni A+V FPS ablation | 5 FPS | raw embedded waveform | 5 Hz | + +Two FPS is the primary Omni setting because it matches Qwen's default video +sampling and permits longer context. Five FPS is the visual-rate ablation. +Input feature rates need not match EgoSpeak; the split, causal prediction +times, three classes, ten target horizons, and evaluator must match exactly. + +## Prepare an EgoSpeak-compatible EasyCom session + +Do not pass the raw one-minute files independently to Omni and do not preserve +the gaps implied by their filenames. EgoSpeak concatenates the retained clips +into a dense session timeline and crops the participant-ID banner on the left. + +`prepare_easycom_session.py` reproduces this policy and creates videos whose +stored frames already lie on a fixed session-global 2 FPS or 5 FPS grid. This +prevents Qwen's per-window video loader from choosing a new adaptive set of +frames for every overlapping prefix. + +Session 1 has been checked against the local release: + +```text +24 retained clips +27,768 source frames at 20 FPS +/ 4 frames per task step += 6,942 steps at 5 Hz += target_perframe/Session_1.npy length exactly +``` + +Preview the commands and validate frames against the target without encoding: + +```bash +python prepare_easycom_session.py \ + --session-dir /Volumes/MX500B/Code/easycom/Main/Video_Compressed/Session_1 \ + --output-dir prepared_easycom/Session_1 \ + --target /Volumes/MX500B/Code/EgoSpeak/dataset/EasyCom/target_perframe/Session_1.npy \ + --fps 2 5 \ + --dry-run +``` + +Remove `--dry-run` to run FFmpeg. The outputs are: + +```text +Session_1_2fps.mp4 primary Omni A+V input +Session_1_5fps.mp4 visual-rate ablation +Session_1.json clip order, exact T, crop and sampling provenance +Session_1.ffconcat reproducible dense concatenation order +``` + +The script fails rather than silently padding or truncating when +`sum(source_frames) / 4 != target.shape[0]`. + +Each clip's decoded video and audio are trimmed to the exact duration implied +by its video frame count before concatenation. This prevents AAC priming/tail +padding from accumulating at the 24 clip boundaries and causing progressive +audio/video drift. + +On the remote L20 machine, use its own dataset paths. For example: + +```bash +python prepare_easycom_session.py \ + --session-dir /data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed/Session_1 \ + --output-dir /data1/yfzx/wangx/Datasets/easycom/omni_prepared/Session_1 \ + --target /path/to/EgoSpeak/dataset/EasyCom/target_perframe/Session_1.npy \ + --fps 2 5 +``` + +This is the command that performs the potentially long FFmpeg processing; it +has not been run by the code authoring workflow. + +If encoding completed but a later validation/metadata step failed, update the +script and reuse the existing MP4 without encoding it again: + +```bash +python prepare_easycom_session.py ... --fps 2 --reuse-existing +``` + +## What Qwen's native 2-second chunk means + +Qwen2.5-Omni is **not restricted to receiving audio once every two seconds**. +It can receive an arbitrary-duration AV input. The two-second value describes +internal processing: + +1. the audio encoder divides acoustic features into approximately two-second + attention blocks; +2. audio and video tokens are interleaved in two-second temporal groups before + entering the Thinker. + +The audio encoder attention inside a block is non-causal. Therefore this is +invalid for online evaluation: + +```text +encode the complete [0, 2.0 s] block +-> read a purported state at 0.2 s +``` + +That early state can contain audio from 0.2--2.0 seconds. At prediction time +`t`, the model must receive only audio and video with timestamps `<= t`. + +## Causal 5 Hz extraction modes + +`streaming_turn_model.py` implements three modes. + +### 1. `causal-prefix`: correctness and semantic baseline + +At every 200 ms, independently process the causal context ending at the +current time: + +```text +[max(0, t-20 s), t] AV -> Omni -> q_t +``` + +Properties: + +- strict no-future input; +- one continuous media window per prediction; +- no cross-forward cache; +- very expensive because overlapping context is recomputed; +- preferred reference when checking semantics and causality. + +Example: + +```bash +python streaming_turn_model.py \ + --video /path/to/easycom_session.mp4 \ + --use-audio-in-video \ + --mode causal-prefix \ + --start 0 \ + --end 20 \ + --memory-seconds 20 \ + --output features/easycom_20s_prefix.pt +``` + +### 2. `commit-replay`: closed-turn compatibility baseline + +The extractor retains Thinker KV for completed two-second media blocks. During +the unfinished block it branches from the same committed cache and evaluates +only the available prefix: + +```text +committed historical KV + current block [0.0, 0.2] -> q_t +committed historical KV + current block [0.0, 0.4] -> q_t +... +committed historical KV + current block [0.0, 2.0] -> q_t + -> commit media KV +``` + +The temporary readout/assistant query is never committed. Only the completed +media block enters long-term cache. This avoids recomputing the full history, +but still re-encodes at most the current two-second prefix. + +Example with numerical checking: + +```bash +python streaming_turn_model.py \ + --video /path/to/easycom_session.mp4 \ + --use-audio-in-video \ + --mode commit-replay \ + --start 0 \ + --end 20 \ + --memory-seconds 20 \ + --verify-every 5 \ + --verify-continuous-every 5 \ + --output features/easycom_20s_stream.pt +``` + +`--verify-every 5` compares cached processing with a from-scratch replay of the +same sequence of two-second media turns. The saved cosine similarity should be +extremely close to 1. If it is not, the TMRoPE/cache path is invalid and must +not be used for a paper result. + +`--verify-continuous-every 5` is a different check: it compares the cached +sequence of two-second media turns with a single continuous causal window. +This comparison need not be numerically identical, but it measures whether the +input organization materially changes the representation. Save and report the +two check types separately (`cache_replay` and `continuous_window_semantics`). + +There is an important experimental distinction: `commit-replay` represents +history as a sequence of two-second media turns. Its numerical check proves +cache equivalence to that same turn sequence, not equivalence to one continuous +20-second video placeholder. Compare it against `causal-prefix` on identical +clips and report the representation if it changes mAP. + +### 3. `open-stream-replay`: recommended efficient mode + +This mode keeps one user/media turn open across the whole bounded-memory +segment. The first block contributes the chat and AV opening markers once; +later completed two-second blocks contribute only interleaved audio/video +payload tokens: + +```text +open AV prefix + completed [0,2] payload -> commit KV + current [2,2.2] payload -> temporary branch -> close/read q_t + current [2,2.4] payload -> temporary branch -> close/read q_t + ... + completed [2,4] payload -> commit KV +AV/chat closure -> never committed; readout branch only +``` + +The audio encoder accepts a shorter final chunk with padding and an +actual-length mask. For the two-frame visual patch requirement, strict-grid +mode duplicates only the most recent past frame when necessary. + +Run a four-second validation before using it for extraction: + +```bash +python streaming_turn_model.py \ + --model-id /path/to/Qwen2.5-Omni-3B \ + --video prepared_easycom/Session_1/Session_1_2fps.mp4 \ + --use-audio-in-video \ + --mode open-stream-replay \ + --start 0 --end 4 \ + --video-fps 2 \ + --strict-fixed-video-grid \ + --verify-every 5 \ + --verify-continuous-every 5 \ + --flash-attention-2 \ + --output features/Session_1_open_first_4s.pt +``` + +`cache_replay` must be identical to replaying the same open stream from an +empty cache. `continuous_window_semantics` compares it with one normal +continuous Qwen AV input. Do not assume equivalence: the modality encoders are +still invoked per tail block in the efficient path, so this numerical check is +required even though artificial user-turn boundaries have been removed. + +## Whole-session prediction and the 20-second limit + +Twenty seconds is a practical memory/context setting for the 3090, not the +maximum video duration. The complete session is processed sequentially and a +prediction is emitted every 0.2 seconds. + +The legacy `commit-replay` prototype resets Thinker KV at each non-overlapping +20-second boundary: + +```text +session: 0 -------- 20 -------- 40 -------- 60 ... +cache: [segment 1] [segment 2] [segment 3] ... +output: continuous predictions across the entire session +``` + +This preserves complete output coverage but limits available history after a +reset. Report the memory length and additionally evaluate the first two seconds +after each reset versus all other timesteps. Rolling KV eviction is a possible +future optimization, but requires correct multimodal position rebasing. + +`open-stream-replay` instead uses blockwise rolling memory. With a 20-second +budget it rebuilds the open cache from the latest nine completed blocks +(18 seconds) when a new block starts, then adds the current 0--2 second tail. +Its effective history therefore varies from 18 to 20 seconds and never drops +to an empty 0.2-second context at a 20-second boundary. + +For the paper experiment, drive the output length from the preparation sidecar +rather than a floating MP4 duration: + +```bash +python streaming_turn_model.py \ + --session-metadata prepared_easycom/Session_1/Session_1.json \ + --use-audio-in-video \ + --mode open-stream-replay \ + --memory-seconds 20 \ + --video-fps 2 \ + --strict-fixed-video-grid \ + --verify-every 25 \ + --verify-continuous-every 25 \ + --output features/Session_1_omni_2fps.pt +``` + +The sidecar selects the matching fixed-FPS MP4 and sets `N` to its exact +`num_steps`. Alternatively pass `--video ... --target ...` explicitly. A raw +20 FPS MP4 with `--video-fps 2` is rejected when +`--strict-fixed-video-grid` is enabled. + +In strict A+V mode the extractor does not delegate temporal frame selection to +the generic video utility. It reads the already-fixed MP4 grid explicitly, +uses only frames with timestamps `<= t`, and duplicates the most recent past +frame when the visual tower needs an even frame count. Audio is loaded once at +16 kHz and sliced exactly at each causal boundary. The processor still receives +`use_audio_in_video=True`, so Qwen's audio/video token interleaving is retained. + +Before a full session, run only the first four seconds and inspect both check +types: + +```bash +python streaming_turn_model.py \ + --model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \ + --video /data1/yfzx/wangx/Datasets/easycom/omni_prepared/Session_1/Session_1_2fps.mp4 \ + --use-audio-in-video \ + --mode open-stream-replay \ + --start 0 --end 4 \ + --video-fps 2 \ + --strict-fixed-video-grid \ + --verify-every 5 \ + --verify-continuous-every 5 \ + --flash-attention-2 \ + --output features/Session_1_first_4s_check.pt +``` + +Do not launch all sessions unless `cache_replay` cosine is nearly 1. The +`continuous_window_semantics` cosine answers a different question and may be +lower because media-turn boundaries change the representation. + +Video can be sampled at 2 FPS while predictions remain at 5 Hz. Audio and the +causal boundary update every 200 ms; consecutive prediction times may use the +same most-recent visual frame. A frame with timestamp later than `t` must never +be included. + +## Saved streaming tensors + +Each `.pt` file contains: + +- `times`: `[N]`, prediction times in seconds; +- `states`: `[N, 2048]`, causal fused Thinker states; +- `label_indices`: `[0, ..., N-1]`, direct `target_perframe` row alignment; +- `valid_10_horizon_steps`: `N-10`, because the last ten current rows do not + have all future labels; +- `mode`: `causal-prefix`, `commit-replay`, or `open-stream-replay`; +- `stream_representation`: explicit input-organization metadata; +- `config`: step, native block, memory, FPS, and pixel settings; +- `verification`: cached-versus-exact numerical checks; +- `input_video_fps` and `video_sampling_grid`: sampling provenance; +- `logits`: optional `[N, 10, 3]` output if a trained head is supplied; +- `probabilities`: optional softmax probabilities. + +## Training the `10 x 3` head + +The head must be trained. It is not part of pretrained Qwen. + +The first experiment should freeze all Qwen parameters and train only: + +```text +LayerNorm(2048) +-> Linear(2048, 1024) +-> GELU +-> Dropout +-> Linear(1024, 30) +-> reshape [10, 3] +``` + +`streaming_turn_model.py` exposes: + +```python +from streaming_turn_model import ( + TurnPredictionHead, + egospeak_multihorizon_loss, +) +``` + +Recommended initial procedure: + +1. Freeze Qwen and extract all training-session states offline. +2. Verify state/label time offset and truncate only the final ten invalid + current positions. +3. Keep EgoSpeak targets as float multi-positive tensors `[N, 10, 3]`; never + convert them with `argmax`. +4. Train with `egospeak_multihorizon_loss`. It exactly mirrors MiniROAD's + anticipation objective: L2-normalize each three-class target, then compute + its weighted negative log-softmax. +5. Select hyperparameters using only training sessions. For EasyCom, reserve a + training session for development or use session-level cross-validation; + never tune on Sessions 1--3. +6. After choosing the epoch/hyperparameters, optionally retrain on all nine + EasyCom training sessions and evaluate the fixed model once on the test set. + +Class-balanced/focal loss, a learned turn-query token, and LoRA are later +ablations. They should not replace the frozen-Omni linear/MLP-head baseline. + +## EgoSpeak-compatible evaluation + +For every valid current timestep, produce probabilities shaped: + +```text +[number of timesteps, 10 future horizons, 3 classes] +``` + +Then reproduce EgoSpeak's Table 2 procedure: + +1. for each horizon independently, concatenate test-session predictions and + ground truth; +2. compute per-class frame-level average precision; +3. average the three class AP values to obtain horizon mAP; +4. report all ten horizon mAP values; +5. average the ten horizon mAP values for the final score. + +With EgoSpeak's `ignore_index=-100`, the released evaluator does not remove +class index zero, so background is included. Reuse the same implementation or +verify a reimplementation numerically. + +Reference Table 2 targets from the EgoSpeak paper: + +- EasyCom GRU A+V: **60.6 mean mAP**; +- EasyCom Transformer A+V: **58.7 mean mAP**; +- Ego4D Transformer audio: **69.2 mean mAP**; +- Ego4D Transformer A+V: **69.0 mean mAP**; +- Ego4D GRU audio: **69.2 mean mAP**; +- Ego4D GRU A+V: **68.2 mean mAP**. + +Do not compare these numbers with EgoSpeak Table 3 target-speaker-only AP; +Table 3 uses a different metric. + +## Experiments that demonstrate Omni's semantic advantage + +The headline comparison alone cannot establish that improvements come from +semantics. Run the following controlled experiments: + +1. **Omni audio vs Omni audio+video**: measure whether visual grounding helps. +2. **Per-horizon curve**: semantic completion should be especially visible at + longer horizons such as 1.2--2.0 seconds. +3. **Context length**: compare 4, 10, 15, and 20 seconds without adding an RNN. +4. **Layer ablation**: compare audio-encoder output and several Thinker layers, + because an intermediate layer may be more discriminative than the final + generation-oriented layer. +5. **Semantic-completion subset**: contrast completed utterances with + hesitation/incomplete pauses while matching pause duration and basic + prosody as closely as possible. +6. **Question/statement/backchannel subsets**: report target-speaker AP and + premature-interruption rate on linguistically meaningful events. +7. **Causal transcript baseline**: compare Omni audio with causal ASR text plus + a language model to isolate lexical/semantic contribution. + +The frozen-Omni head is the cleanest representation test. LoRA can be reported +as a second-stage ceiling, not as the only result. + +## Single-window feature inspection + +`qwen_omni_features.py` is useful for inspecting a short window or layer, but +it is not the main whole-session experiment. + +If audio is embedded in the video: + +```bash +python qwen_omni_features.py \ + --video /path/to/window.mp4 \ + --use-audio-in-video \ + --output features/window.pt +``` + +If audio and video are separate: + +```bash +python qwen_omni_features.py \ + --video /path/to/window.mp4 \ + --audio /path/to/window.wav \ + --output features/window.pt +``` + +Do not provide the same audio both as a separate file and via +`--use-audio-in-video`. + +Saved single-window fields are: + +- `token_features`: `[B, token_length, 2048]`; +- `last_token`: `[B, 2048]`; +- `tail_mean`: `[B, 2048]`; +- `attention_mask`: valid-token mask. + +## Installation + +Install `ffmpeg`, then: + +```bash +pip install -r requirements.txt +``` + +The implementation is pinned to `transformers==4.52.3` because the +commit/replay path uses version-specific Thinker internals. On Linux, +`qwen-omni-utils[decord]` is recommended for faster video decoding. On a +compatible NVIDIA GPU, optionally install `flash-attn` and pass +`--flash-attention-2`. + +## Ego4D Audio-Visual Diarization download + +Do not download the complete Ego4D release for this experiment. Download only +the compressed Audio-Visual Diarization benchmark clips and annotations: + +```bash +ego4d \ + --output_directory=/path/to/ego4d \ + --datasets clips annotations \ + --benchmarks AV +``` + +The approximately 1 TB figure on the Ego4D download page refers to all +benchmark clips combined, not only the AV subset. The annotated AVD subset has +about 572 five-minute clips (roughly 50 hours). EgoSpeak uses a custom 346/87 +train/test split, about 36.1 hours total. Reproducing its Table 2 requires the +exact IDs in EgoSpeak's `EGO4DSHUFFLE` split. + +## Required validation before a full run + +Do not launch full EasyCom/Ego4D extraction until all of these pass: + +- no input audio sample or video frame has timestamp later than prediction + time `t`; +- `states` and `target_perframe` lengths/time origins are reconciled and the + chosen label offset is saved; +- commit/replay cosine equivalence is near 1 on short clips; +- commit/replay and causal-prefix predictions are compared on identical clips; +- a 20-second run fits the 3090 with the selected pixel budget and attention + implementation; +- saved outputs contain exactly ten future horizons in the order 0.2--2.0 s; +- the local metric reproduces EgoSpeak's evaluator on known predictions; +- test session IDs never enter training, validation, normalization statistics, + or hyperparameter selection. + +Once these checks pass, the first publishable experiment is: + +```text +frozen Qwen2.5-Omni-3B Thinker ++ trained 10x3 head ++ 20-second causal memory ++ full-session 5 Hz output ++ exact EgoSpeak split and Table 2 metric +``` diff --git a/audit_egospeak_features.py b/audit_egospeak_features.py new file mode 100644 index 0000000..a65bba4 --- /dev/null +++ b/audit_egospeak_features.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python3 +"""Audit EgoSpeak EasyCom feature rates, shapes, labels, and raw alignment. + +Example: + python audit_egospeak_features.py \ + --feature-root /path/to/EgoSpeak/dataset/EasyCom \ + --raw-manifest manifests/easycom.json \ + --output reports/egospeak_easycom_audit.json +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import numpy as np + +from easycom_manifest import ALL_SESSIONS, TASK_FPS, TEST_SESSIONS, TRAIN_SESSIONS + + +FEATURE_SPECS = { + "rgb_kinetics_resnet50": 2048, + "flow_kinetics_resnet50": 2048, + "audio_wav2vec2": 5120, + "target_perframe": 3, +} + + +def load_raw_steps(path: Path | None) -> dict[str, int]: + if path is None: + return {} + payload = json.loads(path.read_text()) + return { + session["session"]: int(session["expected_5hz_steps"]) + for session in payload["sessions"] + if "expected_5hz_steps" in session + } + + +def audit( + feature_root: Path, + *, + raw_manifest: Path | None, +) -> dict[str, Any]: + feature_root = feature_root.resolve() + raw_steps = load_raw_steps(raw_manifest) + sessions: list[dict[str, Any]] = [] + issues: list[str] = [] + all_targets: list[np.ndarray] = [] + + for session in ALL_SESSIONS: + arrays: dict[str, np.ndarray] = {} + shapes: dict[str, list[int]] = {} + dtypes: dict[str, str] = {} + for modality, expected_dim in FEATURE_SPECS.items(): + path = feature_root / modality / f"{session}.npy" + if not path.is_file(): + issues.append(f"{session}: missing {path}") + continue + array = np.load(path, mmap_mode="r") + arrays[modality] = array + shapes[modality] = list(array.shape) + dtypes[modality] = str(array.dtype) + if array.ndim != 2 or array.shape[1] != expected_dim: + issues.append( + f"{session}: {modality} expected [T,{expected_dim}], got {array.shape}" + ) + + if len(arrays) != len(FEATURE_SPECS): + continue + lengths = {name: int(array.shape[0]) for name, array in arrays.items()} + if len(set(lengths.values())) != 1: + issues.append(f"{session}: unequal modality lengths {lengths}") + + target = np.asarray(arrays["target_perframe"]) + if not np.all((target == 0) | (target == 1)): + issues.append(f"{session}: target contains values outside {{0,1}}") + allowed_rows = {(1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0), (0.0, 1.0, 1.0)} + observed_rows = {tuple(row) for row in np.unique(target, axis=0)} + unexpected = observed_rows - allowed_rows + if unexpected: + issues.append(f"{session}: unexpected target rows {sorted(unexpected)}") + all_targets.append(target) + + num_steps = lengths["target_perframe"] + record: dict[str, Any] = { + "session": session, + "split": "train" if session in TRAIN_SESSIONS else "test", + "num_steps": num_steps, + "duration_seconds_at_5hz": num_steps / TASK_FPS, + "shapes": shapes, + "dtypes": dtypes, + "class_positive_steps": target.sum(axis=0).astype(int).tolist(), + "overlap_wearer_other_steps": int( + np.sum((target[:, 1] == 1) & (target[:, 2] == 1)) + ), + } + if session in raw_steps: + record["raw_manifest_expected_steps"] = raw_steps[session] + record["feature_minus_raw_steps"] = num_steps - raw_steps[session] + if abs(record["feature_minus_raw_steps"]) > 1: + issues.append( + f"{session}: feature/raw mismatch " + f"{record['feature_minus_raw_steps']:+d} steps" + ) + sessions.append(record) + + combined = np.concatenate(all_targets, axis=0) + unique_rows, row_counts = np.unique(combined, axis=0, return_counts=True) + split_steps = { + split: sum(item["num_steps"] for item in sessions if item["split"] == split) + for split in ("train", "test") + } + return { + "feature_root": str(feature_root), + "raw_manifest": str(raw_manifest.resolve()) if raw_manifest is not None else None, + "task_fps": TASK_FPS, + "step_seconds": 1.0 / TASK_FPS, + "sessions": sessions, + "total_steps": int(combined.shape[0]), + "total_duration_hours": float(combined.shape[0] / TASK_FPS / 3600), + "split_steps": split_steps, + "class_positive_steps": combined.sum(axis=0).astype(int).tolist(), + "target_row_counts": { + str(row.astype(int).tolist()): int(count) + for row, count in zip(unique_rows, row_counts) + }, + "issues": issues, + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--feature-root", type=Path, required=True) + parser.add_argument("--raw-manifest", type=Path) + parser.add_argument("--output", type=Path) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + report = audit(args.feature_root, raw_manifest=args.raw_manifest) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(f"Saved report to {args.output}") + + print( + f"{report['total_steps']} steps at {report['task_fps']} Hz = " + f"{report['total_duration_hours']:.6f} hours" + ) + print(f"split steps: {report['split_steps']}") + print(f"class-positive steps [background, wearer, other]: {report['class_positive_steps']}") + print(f"target rows: {report['target_row_counts']}") + if report["issues"]: + print(f"Found {len(report['issues'])} issue(s):") + for issue in report["issues"]: + print(f" - {issue}") + else: + if report["raw_manifest"] is not None: + print("All modality lengths, dimensions, labels, and raw steps match.") + else: + print("All feature-modality lengths, dimensions, and labels are internally consistent.") + + +if __name__ == "__main__": + main() diff --git a/easycom_manifest.py b/easycom_manifest.py new file mode 100644 index 0000000..4101a35 --- /dev/null +++ b/easycom_manifest.py @@ -0,0 +1,321 @@ +#!/usr/bin/env python3 +"""Build and validate an EgoSpeak-compatible EasyCom raw-media manifest. + +EasyCom stores each session as approximately one-minute MP4 segments named by +their session-relative start time (``mm-ss-SSS.mp4``). This script preserves +that global time origin and optionally checks EgoSpeak ``target_perframe`` +arrays at 5 Hz. + +Example: + python easycom_manifest.py \ + --dataset-root /data1/yfzx/wangx/Datasets/easycom/extracted_easycom \ + --target-root /path/to/EasyCom/target_perframe \ + --output easycom_manifest.json +""" + +from __future__ import annotations + +import argparse +import json +import re +import subprocess +from fractions import Fraction +from pathlib import Path +from typing import Any + + +TASK_FPS = 5 +TRAIN_SESSIONS = ( + "Session_7", + "Session_11", + "Session_10", + "Session_6", + "Session_9", + "Session_4", + "Session_8", + "Session_5", + "Session_12", +) +TEST_SESSIONS = ("Session_3", "Session_2", "Session_1") +ALL_SESSIONS = TRAIN_SESSIONS + TEST_SESSIONS +TIMECODE_RE = re.compile(r"^(\d{2})-(\d{2})-(\d{3})$") + + +def parse_timecode(stem: str) -> float: + match = TIMECODE_RE.fullmatch(stem) + if match is None: + raise ValueError(f"Expected mm-ss-SSS filename, got {stem!r}") + minutes, seconds, milliseconds = (int(value) for value in match.groups()) + if seconds >= 60: + raise ValueError(f"Invalid seconds in timecode {stem!r}") + return minutes * 60.0 + seconds + milliseconds / 1000.0 + + +def probe_video(path: Path) -> dict[str, Any]: + completed = subprocess.run( + [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=width,height,avg_frame_rate,nb_frames,duration:format=duration", + "-of", + "json", + str(path), + ], + check=True, + capture_output=True, + text=True, + ) + payload = json.loads(completed.stdout) + streams = payload.get("streams", []) + if not streams: + raise RuntimeError(f"No video stream found in {path}") + stream = streams[0] + rate_text = stream.get("avg_frame_rate", "0/0") + if rate_text in ("0/0", "N/A"): + raise RuntimeError(f"No usable video frame rate in {path}") + fps = float(Fraction(rate_text)) + + frame_text = stream.get("nb_frames") + if frame_text not in (None, "N/A"): + num_frames = int(frame_text) + else: + duration_text = stream.get("duration") or payload.get("format", {}).get("duration") + if duration_text is None: + raise RuntimeError(f"No frame count or duration in {path}") + num_frames = round(float(duration_text) * fps) + + # EgoSpeak processes EasyCom at 20 video frames/s and groups four frames + # per 5 Hz task step. The released MP4s are nominally already 20 FPS. + duration = num_frames / fps + frames_at_20hz = round(duration * 20.0) + width = int(stream["width"]) + height = int(stream["height"]) + return { + "width": width, + "height": height, + "fps": fps, + "num_frames": num_frames, + "video_duration_seconds": duration, + "frames_at_20hz": frames_at_20hz, + # EasyCom videos may contain an ID banner to the left of the 1920px + # video. EgoSpeak crops it and retains the rightmost video region. + "egospeak_crop": { + "x": max(0, width - 1920), + "y": 0, + "width": min(1920, width), + "height": height, + }, + } + + +def find_target(target_root: Path, session: str) -> Path | None: + direct = target_root / f"{session}.npy" + if direct.is_file(): + return direct + matches = sorted(target_root.rglob(f"{session}.npy")) + if len(matches) > 1: + raise RuntimeError( + f"Multiple targets found for {session}: " + + ", ".join(str(path) for path in matches) + ) + return matches[0] if matches else None + + +def target_metadata(path: Path) -> dict[str, Any]: + try: + import numpy as np + except ImportError as exc: + raise RuntimeError("Install numpy to inspect target_perframe arrays") from exc + + target = np.load(path, mmap_mode="r") + if target.ndim not in (1, 2): + raise ValueError(f"Unexpected target shape in {path}: {target.shape}") + if target.ndim == 2 and target.shape[1] != 3: + raise ValueError(f"Expected three classes in {path}, got {target.shape}") + return { + "path": str(path.resolve()), + "shape": list(target.shape), + "num_steps": int(target.shape[0]), + "duration_seconds_at_5hz": float(target.shape[0] / TASK_FPS), + } + + +def build_manifest( + dataset_root: Path, + *, + target_root: Path | None, + probe_media: bool, +) -> dict[str, Any]: + dataset_root = dataset_root.resolve() + video_root = dataset_root / "Main" / "Video_Compressed" + if not video_root.is_dir(): + raise FileNotFoundError(f"EasyCom video directory not found: {video_root}") + + sessions: list[dict[str, Any]] = [] + all_issues: list[str] = [] + all_discontinuities: list[str] = [] + for session in ALL_SESSIONS: + session_dir = video_root / session + if not session_dir.is_dir(): + all_issues.append(f"missing session directory: {session_dir}") + continue + + clip_paths = sorted( + session_dir.glob("*.mp4"), key=lambda path: parse_timecode(path.stem) + ) + if not clip_paths: + all_issues.append(f"no MP4 clips in {session_dir}") + continue + + clips: list[dict[str, Any]] = [] + session_issues: list[str] = [] + session_discontinuities: list[str] = [] + previous_source_end: float | None = None + concatenated_20hz_frames = 0 + for path in clip_paths: + source_start = parse_timecode(path.stem) + media = probe_video(path) if probe_media else None + duration = media["video_duration_seconds"] if media is not None else None + source_end = source_start + duration if duration is not None else None + + if previous_source_end is not None: + delta = source_start - previous_source_end + if abs(delta) > 0.05: + kind = "gap" if delta > 0 else "overlap" + session_discontinuities.append( + f"{kind} of {abs(delta):.3f}s before {path.name}" + ) + if source_end is not None: + previous_source_end = source_end + + record: dict[str, Any] = { + "path": str(path.resolve()), + "timecode": path.stem, + "source_session_start_seconds": source_start, + } + if media is not None and source_end is not None: + concat_start_frame = concatenated_20hz_frames + concatenated_20hz_frames += int(media["frames_at_20hz"]) + record.update( + { + **media, + "source_session_end_seconds": source_end, + "egospeak_concat_start_seconds": concat_start_frame / 20.0, + "egospeak_concat_end_seconds": concatenated_20hz_frames / 20.0, + "target_start_index": concat_start_frame // 4, + "target_end_index_exclusive": concatenated_20hz_frames // 4, + } + ) + clips.append(record) + + source_session_end = ( + max(float(clip["source_session_end_seconds"]) for clip in clips) + if probe_media + else None + ) + split = "train" if session in TRAIN_SESSIONS else "test" + session_record: dict[str, Any] = { + "session": session, + "split": split, + "clips": clips, + "issues": session_issues, + "source_timeline_discontinuities": session_discontinuities, + } + if source_session_end is not None: + session_record["source_session_end_seconds"] = source_session_end + session_record["egospeak_concat_duration_seconds"] = ( + concatenated_20hz_frames / 20.0 + ) + session_record["expected_5hz_steps"] = concatenated_20hz_frames // 4 + + if target_root is not None: + target_path = find_target(target_root.resolve(), session) + if target_path is None: + session_issues.append(f"target_perframe missing for {session}") + else: + target = target_metadata(target_path) + session_record["target"] = target + if source_session_end is not None: + difference = target["num_steps"] - (concatenated_20hz_frames // 4) + session_record["target_minus_media_steps"] = difference + # Small encoder/container rounding differences are normal; + # anything beyond one task step needs manual inspection. + if abs(difference) > 1: + session_issues.append( + f"target/media length mismatch: {difference:+d} steps" + ) + + all_issues.extend(f"{session}: {issue}" for issue in session_issues) + all_discontinuities.extend( + f"{session}: {item}" for item in session_discontinuities + ) + sessions.append(session_record) + + return { + "dataset": "EasyCom", + "dataset_root": str(dataset_root), + "video_root": str(video_root), + "task_fps": TASK_FPS, + "prediction_step_seconds": 1.0 / TASK_FPS, + "class_names": [ + "background", + "camera wearer speaking", + "other people speaking", + ], + "future_target_offsets": list(range(1, 11)), + "train_sessions": list(TRAIN_SESSIONS), + "test_sessions": list(TEST_SESSIONS), + "sessions": sessions, + "source_timeline_discontinuities": all_discontinuities, + "issues": all_issues, + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset-root", type=Path, required=True) + parser.add_argument("--target-root", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--no-probe", + action="store_true", + help="Skip ffprobe duration and continuity checks.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + manifest = build_manifest( + args.dataset_root, + target_root=args.target_root, + probe_media=not args.no_probe, + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(manifest, indent=2) + "\n") + + num_clips = sum(len(session["clips"]) for session in manifest["sessions"]) + print( + f"Saved {len(manifest['sessions'])} sessions / {num_clips} clips " + f"to {args.output}" + ) + if manifest["issues"]: + print(f"Found {len(manifest['issues'])} issue(s):") + for issue in manifest["issues"]: + print(f" - {issue}") + else: + print("No missing-session or target-length issues found.") + discontinuities = manifest["source_timeline_discontinuities"] + if discontinuities: + print( + f"Observed {len(discontinuities)} source-timeline discontinuity(ies); " + "EgoSpeak removes these when concatenating retained clips." + ) + + +if __name__ == "__main__": + main() diff --git a/prepare_easycom_session.py b/prepare_easycom_session.py new file mode 100644 index 0000000..e9df04d --- /dev/null +++ b/prepare_easycom_session.py @@ -0,0 +1,357 @@ +#!/usr/bin/env python3 +"""Prepare one EasyCom session for causal Qwen2.5-Omni extraction. + +The released EasyCom session consists of retained, approximately one-minute +MP4 clips. EgoSpeak removes source-timeline gaps by concatenating those clips +and crops the left participant-ID banner. This script reproduces that dense +timeline and writes fixed-global-rate 2 FPS and/or 5 FPS AV files. + +It intentionally never fills gaps inferred from clip filenames. The number +of 5 Hz task steps is derived from ``target_perframe`` (preferred) or from the +sum of the released 20 FPS video frames. +""" + +from __future__ import annotations + +import argparse +import json +import math +import shlex +import subprocess +from dataclasses import asdict, dataclass +from fractions import Fraction +from pathlib import Path + +from easycom_manifest import parse_timecode + + +SOURCE_FPS = 20 +TASK_FPS = 5 +KEEP_WIDTH = 1920 +KEEP_HEIGHT = 1080 + + +@dataclass(frozen=True) +class ClipInfo: + path: str + timecode_seconds: float + width: int + height: int + fps: float + frames: int + + +def _run_json(command: list[str]) -> dict: + completed = subprocess.run( + command, check=True, capture_output=True, text=True + ) + return json.loads(completed.stdout) + + +def probe_clip(path: Path) -> ClipInfo: + payload = _run_json( + [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=width,height,avg_frame_rate,nb_frames", + "-of", + "json", + str(path), + ] + ) + streams = payload.get("streams", []) + if not streams: + raise RuntimeError(f"No video stream in {path}") + stream = streams[0] + if stream.get("nb_frames") in (None, "N/A"): + raise RuntimeError(f"Exact nb_frames is required for {path}") + fps = float(Fraction(stream["avg_frame_rate"])) + return ClipInfo( + path=str(path.resolve()), + timecode_seconds=parse_timecode(path.stem), + width=int(stream["width"]), + height=int(stream["height"]), + fps=fps, + frames=int(stream["nb_frames"]), + ) + + +def load_target_steps(path: Path) -> int: + try: + import numpy as np + except ImportError as exc: + raise RuntimeError("numpy is required when --target is used") from exc + target = np.load(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 ffconcat_quote(path: str) -> str: + # ffconcat uses single-quoted paths; escape an embedded quote as '\''. + return "'" + path.replace("'", "'\\''") + "'" + + +def validate_inventory(clips: list[ClipInfo], target_steps: int | None) -> int: + if not clips: + raise ValueError("No clips found") + for clip in clips: + if not abs(clip.fps - SOURCE_FPS) < 1e-6: + raise ValueError(f"Expected 20 FPS, got {clip.fps} in {clip.path}") + if clip.width < KEEP_WIDTH or clip.height != KEEP_HEIGHT: + raise ValueError( + f"Cannot take rightmost {KEEP_WIDTH}x{KEEP_HEIGHT} from " + f"{clip.width}x{clip.height}: {clip.path}" + ) + total_frames = sum(clip.frames for clip in clips) + if total_frames % (SOURCE_FPS // TASK_FPS): + raise ValueError( + f"Total source frames {total_frames} are not divisible into 5 Hz steps" + ) + media_steps = total_frames // (SOURCE_FPS // TASK_FPS) + if target_steps is not None and media_steps != target_steps: + raise ValueError( + f"Dense media gives {media_steps} steps but target has {target_steps}" + ) + return media_steps + + +def build_ffmpeg_command( + clips: list[ClipInfo], + output: Path, + *, + fps: float, + dense_duration: float, +) -> list[str]: + # Trim every decoded stream independently before concatenation. Released + # AAC clips can contain roughly one codec-padding frame beyond their video + # duration; concatenating them first would accumulate AV drift. + command = [ + "ffmpeg", + "-hide_banner", + "-loglevel", + "warning", + "-y", + ] + for clip in clips: + command.extend(["-i", clip.path]) + + filters: list[str] = [] + concat_inputs: list[str] = [] + for index, clip in enumerate(clips): + duration = clip.frames / SOURCE_FPS + filters.append( + f"[{index}:v:0]trim=start=0:end={duration:.9f}," + f"setpts=PTS-STARTPTS," + f"crop={KEEP_WIDTH}:{KEEP_HEIGHT}:iw-{KEEP_WIDTH}:0[v{index}]" + ) + filters.append( + f"[{index}:a:0]atrim=start=0:end={duration:.9f}," + f"asetpts=PTS-STARTPTS[a{index}]" + ) + concat_inputs.extend([f"[v{index}]", f"[a{index}]"]) + filters.append( + "".join(concat_inputs) + + f"concat=n={len(clips)}:v=1:a=1[vcat][acat]" + ) + # fps=start_time=0 establishes one global visual lattice after the clips + # have become a dense timeline. The final atrim is a defensive guard + # against decoder rounding at the last sample. + filters.append( + f"[vcat]fps=fps={fps}:start_time=0:round=down,format=yuv420p[vout]" + ) + filters.append( + f"[acat]atrim=start=0:end={dense_duration:.9f}," + "asetpts=PTS-STARTPTS[aout]" + ) + + command.extend( + [ + "-filter_complex", + ";".join(filters), + "-map", + "[vout]", + "-map", + "[aout]", + "-c:v", + "libx264", + "-preset", + "medium", + "-crf", + "18", + "-c:a", + "aac", + "-b:a", + "128k", + "-movflags", + "+faststart", + str(output), + ] + ) + return command + + +def validate_output(path: Path, *, fps: float, dense_duration: float) -> dict: + payload = _run_json( + [ + "ffprobe", + "-v", + "error", + "-show_entries", + "stream=index,codec_type,avg_frame_rate,nb_frames,duration,sample_rate,channels", + "-of", + "json", + str(path), + ] + ) + video = next( + (stream for stream in payload.get("streams", []) if stream["codec_type"] == "video"), + None, + ) + audio = next( + (stream for stream in payload.get("streams", []) if stream["codec_type"] == "audio"), + None, + ) + if video is None or audio is None: + raise RuntimeError(f"Prepared file must contain video and audio: {path}") + actual_fps = float(Fraction(video["avg_frame_rate"])) + actual_frames = int(video["nb_frames"]) + # FFmpeg may either retain or discard the final partial visual interval, + # depending on timestamps delivered by the concat filter. Both are valid + # fixed grids; for 1388.4 s at 2 FPS this permits 2776 or 2777 frames. + frame_product = dense_duration * fps + expected_frames = { + int(math.floor(frame_product + 1e-7)), + int(math.ceil(frame_product - 1e-7)), + } + if not math.isclose(actual_fps, fps, rel_tol=0.0, abs_tol=1e-6): + raise RuntimeError(f"Prepared FPS is {actual_fps}, expected {fps}: {path}") + if actual_frames not in expected_frames: + raise RuntimeError( + f"Prepared frame count is {actual_frames}, expected one of " + f"{sorted(expected_frames)}: {path}" + ) + audio_duration = float(audio["duration"]) + # AAC frame padding can shift the container duration by a few milliseconds. + if abs(audio_duration - dense_duration) > 0.05: + raise RuntimeError( + f"Prepared audio duration {audio_duration:.6f}s differs from dense " + f"timeline {dense_duration:.6f}s: {path}" + ) + return { + "video_fps": actual_fps, + "video_frames": actual_frames, + "allowed_video_frames": sorted(expected_frames), + "video_duration_seconds": float(video["duration"]), + "audio_duration_seconds": audio_duration, + "audio_sample_rate": int(audio["sample_rate"]), + "audio_channels": int(audio["channels"]), + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--session-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--target", type=Path) + parser.add_argument( + "--fps", + type=float, + nargs="+", + default=[2.0, 5.0], + help="Fixed global visual rates to generate (default: 2 5).", + ) + parser.add_argument("--dry-run", action="store_true") + parser.add_argument( + "--reuse-existing", + action="store_true", + help="Validate and reuse an existing output MP4 instead of re-encoding it.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + clip_paths = sorted( + args.session_dir.glob("*.mp4"), key=lambda path: parse_timecode(path.stem) + ) + clips = [probe_clip(path) for path in clip_paths] + target_steps = load_target_steps(args.target) if args.target else None + num_steps = validate_inventory(clips, target_steps) + session = args.session_dir.name + duration = num_steps / TASK_FPS + + args.output_dir.mkdir(parents=True, exist_ok=True) + concat_path = args.output_dir / f"{session}.ffconcat" + concat_text = "ffconcat version 1.0\n" + "".join( + f"file {ffconcat_quote(clip.path)}\n" for clip in clips + ) + concat_path.write_text(concat_text) + + outputs: dict[str, str] = {} + output_media: dict[str, dict] = {} + commands: list[list[str]] = [] + for fps in args.fps: + if fps <= 0: + raise ValueError("Every --fps value must be positive") + fps_label = f"{fps:g}".replace(".", "p") + output = args.output_dir / f"{session}_{fps_label}fps.mp4" + outputs[f"{fps:g}"] = str(output.resolve()) + command = build_ffmpeg_command( + clips, + output, + fps=fps, + dense_duration=duration, + ) + commands.append(command) + if args.dry_run: + print(shlex.join(command)) + elif args.reuse_existing and output.is_file(): + print(f"Reusing {output} ...", flush=True) + output_media[f"{fps:g}"] = validate_output( + output, fps=fps, dense_duration=duration + ) + else: + print(f"Preparing {output} ...", flush=True) + subprocess.run(command, check=True) + output_media[f"{fps:g}"] = validate_output( + output, fps=fps, dense_duration=duration + ) + + metadata = { + "session": session, + "source_session_dir": str(args.session_dir.resolve()), + "source_clip_order": [asdict(clip) for clip in clips], + "source_num_clips": len(clips), + "source_frames_20hz": sum(clip.frames for clip in clips), + "task_fps": TASK_FPS, + "num_steps": num_steps, + "dense_duration_seconds": duration, + "crop": { + "policy": "rightmost_1920x1080", + "width": KEEP_WIDTH, + "height": KEEP_HEIGHT, + }, + "visual_sampling": "fixed_session_global_grid", + "outputs": outputs, + "output_media": output_media, + "target": str(args.target.resolve()) if args.target else None, + "dry_run": bool(args.dry_run), + "reuse_existing": bool(args.reuse_existing), + "commands": [shlex.join(command) for command in commands], + } + metadata_path = args.output_dir / f"{session}.json" + metadata_path.write_text(json.dumps(metadata, indent=2) + "\n") + print( + f"{session}: {len(clips)} clips, {num_steps} steps, " + f"{duration:.3f}s dense timeline" + ) + print(f"Saved metadata to {metadata_path}") + + +if __name__ == "__main__": + main() diff --git a/qwen_omni_features.py b/qwen_omni_features.py new file mode 100644 index 0000000..9daebcf --- /dev/null +++ b/qwen_omni_features.py @@ -0,0 +1,270 @@ +"""Load Qwen2.5-Omni-3B and extract causal audio-video features. + +This research example does not generate text or speech: only the Thinker +(multimodal understanding model) is loaded. + +Example: + python qwen_omni_features.py --video clip.mp4 \ + --use-audio-in-video --output clip_features.pt +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal, Mapping + +import torch +from qwen_omni_utils import process_mm_info +from transformers import ( + Qwen2_5OmniProcessor, + Qwen2_5OmniThinkerForConditionalGeneration, +) + + +MODEL_ID = "Qwen/Qwen2.5-Omni-3B" + + +@dataclass +class OmniFeatures: + """Features returned for one or more audio-video examples.""" + + token_features: torch.Tensor # [B, T, D] + attention_mask: torch.Tensor # [B, T] + last_token: torch.Tensor # [B, D] + tail_mean: torch.Tensor # [B, D] + hidden_states: tuple[torch.Tensor, ...] | None = None + + +def choose_dtype(dtype: Literal["auto", "bfloat16", "float16", "float32"]): + 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: str | Mapping[str, str | int] = "auto", + use_flash_attention_2: bool = False, +) -> tuple[Qwen2_5OmniThinkerForConditionalGeneration, Qwen2_5OmniProcessor]: + """Load the multimodal Thinker and processor without the Talker.""" + + 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 + + +def build_conversation( + *, + video_path: str | Path | None = None, + audio_path: str | Path | None = None, + prompt: str = ( + "Observe the conversation up to the current time and represent the " + "multimodal evidence useful for predicting the next speaking turn." + ), +) -> list[dict[str, Any]]: + """Build one Qwen chat containing video, optional audio, and instruction.""" + + if video_path is None and audio_path is None: + raise ValueError("At least one of video_path and audio_path is required") + + content: list[dict[str, str]] = [] + if video_path is not None: + content.append({"type": "video", "video": str(Path(video_path).resolve())}) + if audio_path is not None: + content.append({"type": "audio", "audio": str(Path(audio_path).resolve())}) + content.append({"type": "text", "text": prompt}) + + return [ + { + "role": "system", + "content": [{ + "type": "text", + "text": ( + "You are a causal multimodal conversation observer. " + "Use only information available up to the current time." + ), + }], + }, + {"role": "user", "content": content}, + ] + + +def prepare_inputs( + processor: Qwen2_5OmniProcessor, + conversations: list[list[dict[str, Any]]], + *, + use_audio_in_video: bool, + model_device: torch.device, + model_dtype: torch.dtype, +): + """Convert a batch of conversations into model tensors.""" + + texts = processor.apply_chat_template( + conversations, add_generation_prompt=True, tokenize=False + ) + audios, images, videos = process_mm_info( + conversations, use_audio_in_video=use_audio_in_video + ) + inputs = processor( + text=texts, + audio=audios, + images=images, + videos=videos, + padding=True, + return_tensors="pt", + use_audio_in_video=use_audio_in_video, + ) + # BatchFeature.to(dtype) casts floating tensors only; IDs/masks stay integer. + return inputs.to(model_device).to(model_dtype) + + +def _last_valid_token(states: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + positions = torch.arange(states.shape[1], device=states.device).unsqueeze(0) + indices = positions.masked_fill(~mask.bool(), -1).max(dim=1).values.clamp_min(0) + batch_indices = torch.arange(states.shape[0], device=states.device) + return states[batch_indices, indices] + + +def _tail_mean( + states: torch.Tensor, mask: torch.Tensor, tail_tokens: int +) -> torch.Tensor: + if tail_tokens <= 0: + raise ValueError("tail_tokens must be positive") + # Ranking valid tokens from the right supports both left and right padding. + reverse_ranks = mask.long().flip(1).cumsum(dim=1).flip(1) + tail_mask = mask.bool() & (reverse_ranks <= tail_tokens) + weights = tail_mask.unsqueeze(-1).to(states.dtype) + return (states * weights).sum(dim=1) / weights.sum(dim=1).clamp_min(1) + + +@torch.inference_mode() +def extract_features( + model: Qwen2_5OmniThinkerForConditionalGeneration, + processor: Qwen2_5OmniProcessor, + conversations: list[list[dict[str, Any]]], + *, + use_audio_in_video: bool = False, + layer: int = -1, + tail_tokens: int = 8, + return_all_layers: bool = False, + output_device: str | torch.device = "cpu", +) -> OmniFeatures: + """Extract fused Thinker states without generation. + + ``last_token`` is the recommended one-vector-per-window baseline. The + Thinker is causal, so it summarizes all earlier prompt and AV tokens. + """ + + inputs = prepare_inputs( + processor, + conversations, + use_audio_in_video=use_audio_in_video, + model_device=model.device, + model_dtype=model.dtype, + ) + # A direct forward ends at the observation boundary; generate() would append + # answer tokens and change which representation is being extracted. + outputs = model( + **inputs, + use_audio_in_video=use_audio_in_video, + output_hidden_states=True, + use_cache=False, + return_dict=True, + ) + if outputs.hidden_states is None: + raise RuntimeError("The model did not return hidden states") + + states = outputs.hidden_states[layer] + mask = inputs["attention_mask"] + return OmniFeatures( + token_features=states.to(output_device), + attention_mask=mask.to(output_device), + last_token=_last_valid_token(states, mask).to(output_device), + tail_mean=_tail_mean(states, mask, tail_tokens).to(output_device), + hidden_states=( + tuple(x.to(output_device) for x in outputs.hidden_states) + if return_all_layers else None + ), + ) + + +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("--output", type=Path, required=True) + parser.add_argument("--model-id", default=MODEL_ID) + parser.add_argument("--prompt", default=None) + parser.add_argument("--use-audio-in-video", action="store_true") + parser.add_argument("--tail-tokens", type=int, default=8) + parser.add_argument("--layer", type=int, default=-1) + parser.add_argument( + "--dtype", + choices=("auto", "bfloat16", "float16", "float32"), + default="auto", + ) + 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() + 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.video is None: + raise SystemExit("--use-audio-in-video requires --video") + + model, processor = load_qwen_omni( + args.model_id, + dtype=args.dtype, + device_map={"": args.device}, + use_flash_attention_2=args.flash_attention_2, + ) + conversation_kwargs: dict[str, Any] = { + "video_path": args.video, + "audio_path": args.audio, + } + if args.prompt is not None: + conversation_kwargs["prompt"] = args.prompt + conversation = build_conversation(**conversation_kwargs) + features = extract_features( + model, + processor, + [conversation], + use_audio_in_video=args.use_audio_in_video, + layer=args.layer, + tail_tokens=args.tail_tokens, + ) + + args.output.parent.mkdir(parents=True, exist_ok=True) + torch.save({ + "model_id": args.model_id, + "layer": args.layer, + "token_features": features.token_features, + "attention_mask": features.attention_mask, + "last_token": features.last_token, + "tail_mean": features.tail_mean, + }, args.output) + print(f"Saved features to {args.output}") + print(f"token_features: {tuple(features.token_features.shape)}") + print(f"last_token: {tuple(features.last_token.shape)}") + print(f"tail_mean: {tuple(features.tail_mean.shape)}") + + +if __name__ == "__main__": + main() diff --git a/reports/egospeak_easycom_feature_audit.json b/reports/egospeak_easycom_feature_audit.json new file mode 100644 index 0000000..8b365ea --- /dev/null +++ b/reports/egospeak_easycom_feature_audit.json @@ -0,0 +1,458 @@ +{ + "feature_root": "/Users/wangx/.mounty/MX500B/Code/EgoSpeak/dataset/EasyCom", + "raw_manifest": null, + "task_fps": 5, + "step_seconds": 0.2, + "sessions": [ + { + "session": "Session_7", + "split": "train", + "num_steps": 7270, + "duration_seconds_at_5hz": 1454.0, + "shapes": { + "rgb_kinetics_resnet50": [ + 7270, + 2048 + ], + "flow_kinetics_resnet50": [ + 7270, + 2048 + ], + "audio_wav2vec2": [ + 7270, + 5120 + ], + "target_perframe": [ + 7270, + 3 + ] + }, + "dtypes": { + "rgb_kinetics_resnet50": "float32", + "flow_kinetics_resnet50": "float32", + "audio_wav2vec2": "float32", + "target_perframe": "float64" + }, + "class_positive_steps": [ + 1529, + 1994, + 4198 + ], + "overlap_wearer_other_steps": 451 + }, + { + "session": "Session_11", + "split": "train", + "num_steps": 6837, + "duration_seconds_at_5hz": 1367.4, + "shapes": { + "rgb_kinetics_resnet50": [ + 6837, + 2048 + ], + "flow_kinetics_resnet50": [ + 6837, + 2048 + ], + "audio_wav2vec2": [ + 6837, + 5120 + ], + "target_perframe": [ + 6837, + 3 + ] + }, + "dtypes": { + "rgb_kinetics_resnet50": "float32", + "flow_kinetics_resnet50": "float32", + "audio_wav2vec2": "float32", + "target_perframe": "float64" + }, + "class_positive_steps": [ + 1662, + 965, + 4676 + ], + "overlap_wearer_other_steps": 466 + }, + { + "session": "Session_10", + "split": "train", + "num_steps": 7199, + "duration_seconds_at_5hz": 1439.8, + "shapes": { + "rgb_kinetics_resnet50": [ + 7199, + 2048 + ], + "flow_kinetics_resnet50": [ + 7199, + 2048 + ], + "audio_wav2vec2": [ + 7199, + 5120 + ], + "target_perframe": [ + 7199, + 3 + ] + }, + "dtypes": { + "rgb_kinetics_resnet50": "float32", + "flow_kinetics_resnet50": "float32", + "audio_wav2vec2": "float32", + "target_perframe": "float64" + }, + "class_positive_steps": [ + 1920, + 2014, + 3757 + ], + "overlap_wearer_other_steps": 492 + }, + { + "session": "Session_6", + "split": "train", + "num_steps": 8819, + "duration_seconds_at_5hz": 1763.8, + "shapes": { + "rgb_kinetics_resnet50": [ + 8819, + 2048 + ], + "flow_kinetics_resnet50": [ + 8819, + 2048 + ], + "audio_wav2vec2": [ + 8819, + 5120 + ], + "target_perframe": [ + 8819, + 3 + ] + }, + "dtypes": { + "rgb_kinetics_resnet50": "float32", + "flow_kinetics_resnet50": "float32", + "audio_wav2vec2": "float32", + "target_perframe": "float64" + }, + "class_positive_steps": [ + 1520, + 3232, + 4967 + ], + "overlap_wearer_other_steps": 900 + }, + { + "session": "Session_9", + "split": "train", + "num_steps": 8528, + "duration_seconds_at_5hz": 1705.6, + "shapes": { + "rgb_kinetics_resnet50": [ + 8528, + 2048 + ], + "flow_kinetics_resnet50": [ + 8528, + 2048 + ], + "audio_wav2vec2": [ + 8528, + 5120 + ], + "target_perframe": [ + 8528, + 3 + ] + }, + "dtypes": { + "rgb_kinetics_resnet50": "float32", + "flow_kinetics_resnet50": "float32", + "audio_wav2vec2": "float32", + "target_perframe": "float64" + }, + "class_positive_steps": [ + 2957, + 884, + 4808 + ], + "overlap_wearer_other_steps": 121 + }, + { + "session": "Session_4", + "split": "train", + "num_steps": 8658, + "duration_seconds_at_5hz": 1731.6, + "shapes": { + "rgb_kinetics_resnet50": [ + 8658, + 2048 + ], + "flow_kinetics_resnet50": [ + 8658, + 2048 + ], + "audio_wav2vec2": [ + 8658, + 5120 + ], + "target_perframe": [ + 8658, + 3 + ] + }, + "dtypes": { + "rgb_kinetics_resnet50": "float32", + "flow_kinetics_resnet50": "float32", + "audio_wav2vec2": "float32", + "target_perframe": "float64" + }, + "class_positive_steps": [ + 1386, + 1013, + 6681 + ], + "overlap_wearer_other_steps": 422 + }, + { + "session": "Session_8", + "split": "train", + "num_steps": 7642, + "duration_seconds_at_5hz": 1528.4, + "shapes": { + "rgb_kinetics_resnet50": [ + 7642, + 2048 + ], + "flow_kinetics_resnet50": [ + 7642, + 2048 + ], + "audio_wav2vec2": [ + 7642, + 5120 + ], + "target_perframe": [ + 7642, + 3 + ] + }, + "dtypes": { + "rgb_kinetics_resnet50": "float32", + "flow_kinetics_resnet50": "float32", + "audio_wav2vec2": "float32", + "target_perframe": "float64" + }, + "class_positive_steps": [ + 1859, + 1778, + 4533 + ], + "overlap_wearer_other_steps": 528 + }, + { + "session": "Session_5", + "split": "train", + "num_steps": 8558, + "duration_seconds_at_5hz": 1711.6, + "shapes": { + "rgb_kinetics_resnet50": [ + 8558, + 2048 + ], + "flow_kinetics_resnet50": [ + 8558, + 2048 + ], + "audio_wav2vec2": [ + 8558, + 5120 + ], + "target_perframe": [ + 8558, + 3 + ] + }, + "dtypes": { + "rgb_kinetics_resnet50": "float32", + "flow_kinetics_resnet50": "float32", + "audio_wav2vec2": "float32", + "target_perframe": "float64" + }, + "class_positive_steps": [ + 2256, + 2814, + 4083 + ], + "overlap_wearer_other_steps": 595 + }, + { + "session": "Session_12", + "split": "train", + "num_steps": 8894, + "duration_seconds_at_5hz": 1778.8, + "shapes": { + "rgb_kinetics_resnet50": [ + 8894, + 2048 + ], + "flow_kinetics_resnet50": [ + 8894, + 2048 + ], + "audio_wav2vec2": [ + 8894, + 5120 + ], + "target_perframe": [ + 8894, + 3 + ] + }, + "dtypes": { + "rgb_kinetics_resnet50": "float32", + "flow_kinetics_resnet50": "float32", + "audio_wav2vec2": "float32", + "target_perframe": "float64" + }, + "class_positive_steps": [ + 1252, + 3353, + 6099 + ], + "overlap_wearer_other_steps": 1810 + }, + { + "session": "Session_3", + "split": "test", + "num_steps": 8400, + "duration_seconds_at_5hz": 1680.0, + "shapes": { + "rgb_kinetics_resnet50": [ + 8400, + 2048 + ], + "flow_kinetics_resnet50": [ + 8400, + 2048 + ], + "audio_wav2vec2": [ + 8400, + 5120 + ], + "target_perframe": [ + 8400, + 3 + ] + }, + "dtypes": { + "rgb_kinetics_resnet50": "float32", + "flow_kinetics_resnet50": "float32", + "audio_wav2vec2": "float32", + "target_perframe": "float64" + }, + "class_positive_steps": [ + 2024, + 2547, + 4521 + ], + "overlap_wearer_other_steps": 692 + }, + { + "session": "Session_2", + "split": "test", + "num_steps": 7706, + "duration_seconds_at_5hz": 1541.2, + "shapes": { + "rgb_kinetics_resnet50": [ + 7706, + 2048 + ], + "flow_kinetics_resnet50": [ + 7706, + 2048 + ], + "audio_wav2vec2": [ + 7706, + 5120 + ], + "target_perframe": [ + 7706, + 3 + ] + }, + "dtypes": { + "rgb_kinetics_resnet50": "float32", + "flow_kinetics_resnet50": "float32", + "audio_wav2vec2": "float32", + "target_perframe": "float64" + }, + "class_positive_steps": [ + 1718, + 1350, + 5152 + ], + "overlap_wearer_other_steps": 514 + }, + { + "session": "Session_1", + "split": "test", + "num_steps": 6942, + "duration_seconds_at_5hz": 1388.4, + "shapes": { + "rgb_kinetics_resnet50": [ + 6942, + 2048 + ], + "flow_kinetics_resnet50": [ + 6942, + 2048 + ], + "audio_wav2vec2": [ + 6942, + 5120 + ], + "target_perframe": [ + 6942, + 3 + ] + }, + "dtypes": { + "rgb_kinetics_resnet50": "float32", + "flow_kinetics_resnet50": "float32", + "audio_wav2vec2": "float32", + "target_perframe": "float64" + }, + "class_positive_steps": [ + 1538, + 2266, + 3858 + ], + "overlap_wearer_other_steps": 720 + } + ], + "total_steps": 95453, + "total_duration_hours": 5.302944444444444, + "split_steps": { + "train": 72405, + "test": 23048 + }, + "class_positive_steps": [ + 21621, + 24210, + 57333 + ], + "target_row_counts": { + "[0, 0, 1]": 49622, + "[0, 1, 0]": 16499, + "[0, 1, 1]": 7711, + "[1, 0, 0]": 21621 + }, + "issues": [] +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..10df5a8 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,9 @@ +torch +numpy +transformers==4.52.3 +accelerate +qwen-omni-utils[decord] +librosa +Pillow + +# Optional on compatible NVIDIA GPUs: flash-attn diff --git a/smoke_test.py b/smoke_test.py new file mode 100644 index 0000000..b564a71 --- /dev/null +++ b/smoke_test.py @@ -0,0 +1,147 @@ +#!/usr/bin/env python3 +"""Minimal single-GPU smoke test for a local Qwen2.5-Omni checkpoint. + +Examples: + python smoke_test.py --model /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B + python smoke_test.py --model /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \ + --video /data1/example.mp4 + +Qwen2.5-Omni generation currently is more reliable on a single GPU for this +use case. In particular, do not replace the explicit device map below with +``device_map=\"auto\"``: it can place cache tensors on different GPUs. +""" + +import argparse +import time + +import torch +from qwen_omni_utils import process_mm_info +from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True, help="Local model directory or HF model ID") + parser.add_argument("--video", help="Optional local .mp4 file for an audio-video test") + parser.add_argument("--device", default="cuda:0", help="One CUDA device, e.g. cuda:0") + parser.add_argument("--max-new-tokens", type=int, default=64) + parser.add_argument( + "--video-fps", + type=float, + default=1.0, + help="Frames per second sampled by the Omni video processor (default: 1.0)", + ) + parser.add_argument( + "--video-max-pixels", + type=int, + default=256 * 28 * 28, + help="Maximum pixels per sampled video frame (default: 200704)", + ) + parser.add_argument( + "--flash-attn2", + action="store_true", + help="Use FlashAttention-2; requires the flash_attn package.", + ) + return parser.parse_args() + + +def move_tensors_to_device(batch: dict, device: str) -> dict: + """Move processor tensors without accidentally converting integer IDs.""" + return { + name: value.to(device) if torch.is_tensor(value) else value + for name, value in batch.items() + } + + +def main() -> None: + args = parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("A CUDA GPU is required for this smoke test.") + + print(f"GPU: {torch.cuda.get_device_name(args.device)}") + print(f"CUDA: {torch.version.cuda}") + print(f"PyTorch: {torch.__version__}") + print(f"Loading: {args.model}") + + started = time.perf_counter() + # A 3B model fits comfortably on one L20 (48 GB). Keep all weights and + # generation cache on the same device; device_map='auto' may split them. + model_load_kwargs = dict( + torch_dtype=torch.bfloat16, + device_map={"": args.device}, + ) + if args.flash_attn2: + model_load_kwargs["attn_implementation"] = "flash_attention_2" + + model = Qwen2_5OmniForConditionalGeneration.from_pretrained( + args.model, + **model_load_kwargs, + ) + model.disable_talker() # We need Thinker text states only; saves VRAM. + model.eval() + processor = Qwen2_5OmniProcessor.from_pretrained(args.model) + print(f"Model loaded in {time.perf_counter() - started:.1f}s") + + if args.video: + content = [ + { + "type": "video", + "video": args.video, + # These are consumed by qwen_omni_utils while sampling the + # video. They prevent quadratic vision attention from + # allocating tens of GB for a short, high-resolution clip. + "fps": args.video_fps, + "max_pixels": args.video_max_pixels, + }, + { + "type": "text", + "text": "Briefly describe the people, speech, and interaction in this video.", + }, + ] + print( + f"Video: {args.video} " + f"(sample_fps={args.video_fps}, max_pixels={args.video_max_pixels})" + ) + else: + content = [ + {"type": "text", "text": "Reply with exactly: Qwen Omni smoke test passed."}, + ] + + messages = [{"role": "user", "content": content}] + prompt = processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + audios, images, videos = process_mm_info( + messages, use_audio_in_video=bool(args.video) + ) + inputs = processor( + text=prompt, + audio=audios, + images=images, + videos=videos, + return_tensors="pt", + padding=True, + use_audio_in_video=bool(args.video), + ) + inputs = move_tensors_to_device(inputs, args.device) + + started = time.perf_counter() + with torch.inference_mode(): + output_ids = model.generate( + **inputs, + max_new_tokens=args.max_new_tokens, + do_sample=False, + return_audio=False, + ) + + # Remove the prompt tokens before decoding the generated completion. + generated_ids = output_ids[:, inputs["input_ids"].shape[1] :] + answer = processor.batch_decode( + generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False + )[0] + print(f"Generation completed in {time.perf_counter() - started:.1f}s") + print(f"\nModel response:\n{answer.strip()}") + + +if __name__ == "__main__": + main() diff --git a/streaming_turn_model.py b/streaming_turn_model.py new file mode 100644 index 0000000..a822db1 --- /dev/null +++ b/streaming_turn_model.py @@ -0,0 +1,1866 @@ +"""Omni-only causal turn-taking features at 5 Hz. + +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 + 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. + +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 \ + --memory-seconds 20 \ + --verify-every 25 \ + --output features/session.pt +""" + +from __future__ import annotations + +import argparse +import copy +import json +import math +import subprocess +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Literal + +import torch +from qwen_omni_utils import process_mm_info +from torch import nn +from transformers import ( + Qwen2_5OmniProcessor, + Qwen2_5OmniThinkerForConditionalGeneration, +) +from transformers.cache_utils import DynamicCache + +Mode = Literal["causal-prefix", "commit-replay", "open-stream-replay"] +MODEL_ID = "Qwen/Qwen2.5-Omni-3B" + + +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 + + 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") + _require_integer_ratio( + self.native_block_seconds, self.step_seconds, "native block / step" + ) + _require_integer_ratio( + self.memory_seconds, self.native_block_seconds, "memory / native block" + ) + + +@dataclass +class ExtractionResult: + """One fused Thinker state for every causal prediction time.""" + + times: torch.Tensor # [N], seconds in the original session + states: torch.Tensor # [N, D] + verification: list[dict[str, Any]] + + +@dataclass +class OpenMediaPieces: + """One prepared AV interval split around its interleaved media payload.""" + + prefix_embeddings: torch.Tensor + prefix_position_ids: torch.LongTensor + payload_embeddings: torch.Tensor + payload_position_ids: torch.LongTensor + closure_ids: torch.LongTensor + + +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 librosa + except ImportError as exc: + raise RuntimeError( + "Fixed-grid AV extraction requires decord and librosa" + ) 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) + self.audio, _ = librosa.load( + str(self.path), sr=self.sample_rate, mono=True + ) + + def frames(self, start: float, end: float) -> list[Any]: + """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)), + ) + 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. + 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] + + +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) + 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, +) -> 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 + batch = processor( + text=[text], + audio=audios, + images=images, + videos=videos, + padding=True, + return_tensors="pt", + use_audio_in_video=use_audio_in_video, + **processor_kwargs, + ) + # BatchFeature.to(dtype) changes floating tensors only; IDs stay integral. + batch = batch.to(model.device).to(model.dtype) + return dict(batch) + + +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, + ) + + +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: + 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, +) -> 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=False, + return_dict=True, + cache_position=cache_position, + ) + return outputs.last_hidden_state[:, -1, :] + + +@torch.inference_mode() +def _read_open_state( + model: Qwen2_5OmniThinkerForConditionalGeneration, + open_media_cache: Any, + *, + next_rope_offset: int, + closure_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) + 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) + scalar_positions = torch.arange( + next_rope_offset, + next_rope_offset + readout_length, + device=readout_ids.device, + ) + position_ids = scalar_positions.view(1, 1, -1).expand( + 3, readout_ids.shape[0], -1 + ) + 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=False, + return_dict=True, + cache_position=cache_position, + ) + return outputs.last_hidden_state[:, -1, :] + + +@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]: + """Append one AV payload, adding the chat/media prefix exactly once.""" + + if include_prefix: + if past_key_values is not None: + raise ValueError("An open-media prefix can only start an empty cache") + 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 + ) + 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, + ) + return _append_embeddings( + model, + embeddings, + position_ids, + past_key_values=past_key_values, + ) + + +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, + ) + + def extract( + self, + *, + video_path: str | Path | None, + audio_path: str | Path | None, + start: float, + end: float, + use_audio_in_video: bool, + ) -> 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()) + return ExtractionResult( + times=torch.tensor(times, dtype=torch.float64), + states=torch.cat(states, dim=0), + verification=[], + ) + + +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.""" + + 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, + ) -> 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, + ) + return _open_media_pieces( + self.model, + inputs, + use_audio_in_video=use_audio_in_video, + ) + + 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, next_offset, _, closure_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, + next_rope_offset=next_offset, + closure_ids=closure_ids, + query_ids=self.query_ids, + ) + + 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, + ) -> tuple[Any, int, int, torch.LongTensor]: + """Build an open payload cache and expose its local TMRoPE base.""" + + if not intervals: + raise ValueError("At least one interval is required") + media_cache = None + next_offset = 0 + closure_ids = None + payload_base = None + positions_per_chunk = _tmrope_positions_per_chunk(self.model) + for index, interval in enumerate(intervals): + pieces = self._prepare_pieces( + video_path=video_path, + audio_path=audio_path, + interval=interval, + use_audio_in_video=use_audio_in_video, + ) + 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( + 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 payload_base is None: + raise RuntimeError("Open-stream replay did not establish a payload base") + return media_cache, next_offset, payload_base, closure_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, + ) -> 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]] = [] + + 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, + ) + ( + 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, + ) + + 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, next_offset = _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, + next_rope_offset=next_offset, + closure_ids=pieces.closure_ids, + 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( + 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 + 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, 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, + ) + return ExtractionResult( + times=torch.tensor(all_times, dtype=torch.float64), + states=torch.cat(all_states, dim=0), + verification=all_checks, + ) + + +def _block_intervals(start: float, end: float, block_seconds: float) -> list[tuple[float, float]]: + """Exact media-turn sequence represented by a commit/replay cache.""" + + intervals = [] + cursor = start + while cursor < end - 1e-7: + next_end = min(cursor + block_seconds, end) + intervals.append((cursor, next_end)) + cursor = next_end + return intervals + + +def _compare_states( + current_time: float, + cached: torch.Tensor, + exact: torch.Tensor, + *, + kind: str, +) -> dict[str, Any]: + cached_f = cached.float() + exact_f = exact.float() + return { + "kind": kind, + "time": float(current_time), + "max_abs": float((cached_f - exact_f).abs().max().item()), + "mean_abs": float((cached_f - exact_f).abs().mean().item()), + "cosine": float(nn.functional.cosine_similarity(cached_f, exact_f).item()), + } + + +def probe_duration(path: str | Path) -> float: + """Read media duration with ffprobe.""" + + completed = subprocess.run( + [ + "ffprobe", + "-v", + "error", + "-show_entries", + "format=duration", + "-of", + "default=noprint_wrappers=1:nokey=1", + str(Path(path).resolve()), + ], + check=True, + capture_output=True, + text=True, + ) + return float(completed.stdout.strip()) + + +def probe_video_fps(path: str | Path) -> float: + """Read the stored video frame rate with ffprobe.""" + + completed = subprocess.run( + [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=avg_frame_rate", + "-of", + "default=noprint_wrappers=1:nokey=1", + str(Path(path).resolve()), + ], + check=True, + capture_output=True, + text=True, + ) + numerator, denominator = completed.stdout.strip().split("/", maxsplit=1) + return float(numerator) / float(denominator) + + +def target_num_steps(path: str | Path) -> int: + """Read an EgoSpeak target_perframe ``[T,3]`` length.""" + + try: + import numpy as np + except ImportError as exc: + raise RuntimeError("numpy is required for --target") from exc + target = np.load(Path(path), mmap_mode="r") + if target.ndim != 2 or target.shape[1] != 3: + raise ValueError(f"Expected target shape [T,3], got {target.shape}") + return int(target.shape[0]) + + +def _apply_session_metadata(args: argparse.Namespace) -> dict[str, Any] | None: + """Fill the prepared video and task length from a preparation sidecar.""" + + if args.session_metadata is None: + return None + metadata = json.loads(args.session_metadata.read_text()) + fps_key = f"{args.video_fps:g}" + try: + metadata_video = Path(metadata["outputs"][fps_key]) + except KeyError as exc: + raise SystemExit( + f"Metadata has no fixed {fps_key} FPS output: {args.session_metadata}" + ) from exc + if args.video is not None and args.video.resolve() != metadata_video.resolve(): + raise SystemExit("--video conflicts with the selected --session-metadata output") + args.video = metadata_video + metadata_steps = int(metadata["num_steps"]) + if args.num_steps is not None and args.num_steps != metadata_steps: + raise SystemExit("--num-steps conflicts with --session-metadata") + args.num_steps = metadata_steps + return metadata + + +def _hidden_size(model: Qwen2_5OmniThinkerForConditionalGeneration) -> int: + text_config = getattr(model.config, "text_config", None) + if text_config is not None: + return int(text_config.hidden_size) + return int(model.config.hidden_size) + + +def _load_head( + checkpoint: Path, + hidden_size: int, + device: torch.device, +) -> TurnPredictionHead: + head = TurnPredictionHead(hidden_size=hidden_size) + payload = torch.load(checkpoint, map_location="cpu", weights_only=True) + state_dict = payload.get("head", payload) if isinstance(payload, dict) else payload + head.load_state_dict(state_dict) + return head.to(device).eval() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--video", type=Path) + parser.add_argument("--audio", type=Path) + parser.add_argument( + "--session-metadata", + type=Path, + help="JSON sidecar written by prepare_easycom_session.py.", + ) + parser.add_argument( + "--target", + type=Path, + help="target_perframe .npy; its T fixes the full-session output length.", + ) + parser.add_argument( + "--num-steps", + type=int, + help="Exact number of 5 Hz outputs (normally read from --target/metadata).", + ) + parser.add_argument("--use-audio-in-video", action="store_true") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--mode", + choices=("causal-prefix", "commit-replay", "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("--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("--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() + 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, + ) + 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, + ) + 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, + ) + 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, + ) + + payload: dict[str, Any] = { + "model_id": args.model_id, + "mode": args.mode, + "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" + ) + ), + "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_past_frames_with_duplication" + if fixed_grid_av is not None + else "qwen_utility_adaptive_sampling" + ), + "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", + "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()