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 <cursoragent@cursor.com>
This commit is contained in:
Xin Wang
@@ -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
|
||||
```
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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": []
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
torch
|
||||
numpy
|
||||
transformers==4.52.3
|
||||
accelerate
|
||||
qwen-omni-utils[decord]
|
||||
librosa
|
||||
Pillow
|
||||
|
||||
# Optional on compatible NVIDIA GPUs: flash-attn
|
||||
+147
@@ -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()
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user