From 080c45f049d31a104042c4929363ce3afa0fcec2 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Sun, 16 Aug 2026 22:19:43 +0800 Subject: [PATCH 1/3] Add reproducible EasyCom causal head pipeline --- README.md | 165 +++++----- easycom_head.py | 299 +++++++++++++++++ infer_turn_head.py | 289 +++++++++++++++++ prepare_easycom_session.py | 84 +++-- requirements.txt | 2 + run_easycom_head_experiment.sh | 35 ++ streaming_turn_model.py | 571 ++++++++++++--------------------- train_turn_head.py | 476 +++++++++++++++++++++++++++ 8 files changed, 1451 insertions(+), 470 deletions(-) create mode 100644 easycom_head.py create mode 100644 infer_turn_head.py create mode 100755 run_easycom_head_experiment.sh create mode 100644 train_turn_head.py diff --git a/README.md b/README.md index 5872967..167f585 100644 --- a/README.md +++ b/README.md @@ -32,9 +32,10 @@ classification head. | `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 | +| `streaming_turn_model.py` | Implemented | Produce one causal Omni state every 200 ms, with causal-prefix and open-stream 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 | +| `train_turn_head.py` | Implemented | Train the frozen-Omni anticipation head with the official EasyCom split | +| `infer_turn_head.py` | Implemented | Save EgoSpeak-format predictions and reproduce per-horizon AP/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 @@ -284,13 +285,17 @@ 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]`. +The target length must equal the number of complete 5 Hz intervals, +`floor(sum(source_frames) / 4)`. If a session ends with one to three residual +20 Hz frames, the target-derived duration discards only that incomplete final +interval; the discarded frame count is recorded in the sidecar. 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. +audio/video drift. The session-global visual grid is padded defensively at EOF +and then trimmed by frame index, so it contains exactly +`ceil(target_steps / 5 * video_fps)` frames on every FFmpeg version. On the remote L20 machine, use its own dataset paths. For example: @@ -336,7 +341,7 @@ That early state can contain audio from 0.2--2.0 seconds. At prediction time ## Causal 5 Hz extraction modes -`streaming_turn_model.py` implements three modes. +`streaming_turn_model.py` implements two modes. ### 1. `causal-prefix`: correctness and semantic baseline @@ -368,57 +373,7 @@ python streaming_turn_model.py \ --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 +### 2. `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; @@ -468,20 +423,6 @@ 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. @@ -515,6 +456,10 @@ 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. +Long runs use `tqdm` and update the progress bar after every completed state, +showing percentage, media time, elapsed time, ETA, and throughput. Pass +`--no-progress` to disable it. The Thinker-only path also suppresses Qwen's +irrelevant speech-output system prompt warning. Before a full session, run only the first four seconds and inspect both check types: @@ -536,7 +481,7 @@ python streaming_turn_model.py \ 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. +lower because blockwise media preprocessing changes 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 @@ -552,7 +497,7 @@ Each `.pt` file contains: - `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`; +- `mode`: `causal-prefix` 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; @@ -575,7 +520,8 @@ LayerNorm(2048) -> reshape [10, 3] ``` -`streaming_turn_model.py` exposes: +`easycom_head.py` defines the head and loss; they remain re-exported by +`streaming_turn_model.py` for compatibility: ```python from streaming_turn_model import ( @@ -584,16 +530,37 @@ from streaming_turn_model import ( ) ``` -Recommended initial procedure: +The supplied LSTR config uses Adam, 50 epochs, learning rate `7e-5`, weight +decay `5e-5`, dropout `0.2`, and 10 warmup epochs followed by cosine decay. +`train_turn_head.py` uses those defaults. Its batch size is 512 causal states, +not LSTR's 16 temporal windows, because the two model inputs have different +shapes. + +Train once on the official nine sessions with a fixed 50-epoch protocol: + +```bash +python train_turn_head.py \ + --feature-root features/causal_prefix_2fps \ + --target-root target_frame \ + --output checkpoints/causal_prefix_2fps_head.pth \ + --device cuda:0 +``` + +The default loss is `lstr-mce`, exactly matching LSTR's MCE treatment of a +multi-positive `[0,1,1]` row: it sums both active negative-log-softmax terms +without normalizing the target. The previous L2-normalized objective belongs +to MiniROAD and is available only as the explicit +`--loss miniroad-normalized` ablation. + +Recommended experimental 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. +4. For comparison with the supplied LSTR shell, train with the default + `egospeak_multihorizon_loss`/`lstr-mce` objective. 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. @@ -605,7 +572,22 @@ ablations. They should not replace the frozen-Omni linear/MLP-head baseline. ## EgoSpeak-compatible evaluation -For every valid current timestep, produce probabilities shaped: +Run inference on the official test sessions (`Session_3`, `Session_2`, and +`Session_1`): + +```bash +python infer_turn_head.py \ + --feature-root features/causal_prefix_2fps \ + --target-root target_frame \ + --checkpoint checkpoints/causal_prefix_2fps_head.pth \ + --output-dir results/causal_prefix_2fps_head \ + --device cuda:0 +``` + +This writes `pred/Session_X.npy` with the exact LSTR/EgoSpeak layout +`[target_time, class, future_horizon]`, i.e. `[T,3,10]`, plus `metrics.json`. +For every valid current timestep, the head itself produces probabilities +shaped: ```text [number of timesteps, 10 future horizons, 3 classes] @@ -621,8 +603,22 @@ Then reproduce EgoSpeak's Table 2 procedure: 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. +class index zero, so background is included. The script reports two protocols: + +- `egospeak_full_timeline`: the same full `[T,3,10]` array and evaluator as + LSTR, with zeros where a prediction would require pre-session context; +- `causal_valid_only`: excludes only those unavailable first `h` rows for + horizon `h`. This is the primary strictly causal score. + +The difference involves at most the first ten rows of each session, but both +are saved so the comparison is explicit. + +The convenience wrapper mirrors EgoSpeak's train/inference shells: + +```bash +chmod +x run_easycom_head_experiment.sh +CUDA_VISIBLE_DEVICES=0 ./run_easycom_head_experiment.sh all +``` Reference Table 2 targets from the EgoSpeak paper: @@ -701,7 +697,7 @@ 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, +open-stream 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`. @@ -732,8 +728,9 @@ Do not launch full EasyCom/Ego4D extraction until all of these pass: 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; +- open-stream cache replay is exactly equivalent to its from-scratch replay on + short clips; +- open-stream 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; diff --git a/easycom_head.py b/easycom_head.py new file mode 100644 index 0000000..ff5c138 --- /dev/null +++ b/easycom_head.py @@ -0,0 +1,299 @@ +"""Shared head, labels, and metrics for the EasyCom EgoSpeak experiment.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal, Sequence + +import numpy as np +import torch +from torch import nn + + +OFFICIAL_TRAIN_SESSIONS = ( + "Session_7", + "Session_11", + "Session_10", + "Session_6", + "Session_9", + "Session_4", + "Session_8", + "Session_5", + "Session_12", +) +OFFICIAL_TEST_SESSIONS = ("Session_3", "Session_2", "Session_1") +CLASS_NAMES = ( + "background", + "camera_wearer_speaking", + "other_people_speaking", +) +STEP_SECONDS = 0.2 +NUM_HORIZONS = 10 +NUM_CLASSES = 3 + + +class TurnPredictionHead(nn.Module): + """Map one Omni state to EgoSpeak-compatible ``[10, 3]`` logits.""" + + def __init__( + self, + hidden_size: int = 2048, + num_horizons: int = NUM_HORIZONS, + num_classes: int = NUM_CLASSES, + dropout: float = 0.1, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.num_horizons = num_horizons + self.num_classes = num_classes + self.dropout = dropout + 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 _validate_loss_inputs(logits: torch.Tensor, targets: torch.Tensor) -> None: + if logits.shape != targets.shape: + raise ValueError( + f"logits/targets shape mismatch: {tuple(logits.shape)} vs " + f"{tuple(targets.shape)}" + ) + if logits.shape[-1] != NUM_CLASSES: + raise ValueError( + f"Expected {NUM_CLASSES} classes, got {logits.shape[-1]}" + ) + + +def egospeak_multihorizon_loss( + logits: torch.Tensor, + targets: torch.Tensor, + *, + reduction: Literal["mean", "sum", "none"] = "mean", +) -> torch.Tensor: + """Reproduce the LSTR ``MCE`` loss used by the supplied shell/config. + + EasyCom targets may contain both active speaker classes as ``[0, 1, 1]``. + LSTR preserves those weights directly and sums their negative log-softmax + terms. It does not normalize the target row. + """ + + _validate_loss_inputs(logits, targets) + targets = targets.to(device=logits.device, dtype=logits.dtype) + 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 miniroad_multihorizon_loss( + logits: torch.Tensor, + targets: torch.Tensor, + *, + reduction: Literal["mean", "sum", "none"] = "mean", +) -> torch.Tensor: + """Reproduce MiniROAD's L2-normalized multi-positive anticipation loss.""" + + _validate_loss_inputs(logits, targets) + 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 alias for the LSTR/EgoSpeak loss.""" + + return egospeak_multihorizon_loss(logits, targets) + + +def normalize_session_name(value: str | int) -> str: + text = str(value) + if text.startswith("Session_"): + suffix = text.removeprefix("Session_") + else: + suffix = text + if not suffix.isdigit(): + raise ValueError(f"Invalid EasyCom session: {value}") + return f"Session_{int(suffix)}" + + +@dataclass +class SessionData: + name: str + states: torch.Tensor + targets: torch.Tensor + label_indices: torch.LongTensor + feature_metadata: dict[str, Any] + + @property + def hidden_size(self) -> int: + return int(self.states.shape[-1]) + + def future_examples( + self, + num_horizons: int = NUM_HORIZONS, + ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: + """Return states, future labels, and target indices for valid rows.""" + + offsets = torch.arange(1, num_horizons + 1, dtype=torch.long) + future_indices = self.label_indices[:, None] + offsets[None, :] + valid = (future_indices >= 0).all(dim=1) & ( + future_indices < len(self.targets) + ).all(dim=1) + rows = future_indices[valid] + return self.states[valid], self.targets[rows], rows + + +def load_session_data( + feature_root: str | Path, + target_root: str | Path, + session: str | int, + *, + expected_mode: str | None = "causal-prefix", +) -> SessionData: + """Load and validate one saved Omni feature payload and target array.""" + + name = normalize_session_name(session) + feature_path = Path(feature_root) / f"{name}.pt" + target_path = Path(target_root) / f"{name}.npy" + if not feature_path.is_file(): + raise FileNotFoundError(feature_path) + if not target_path.is_file(): + raise FileNotFoundError(target_path) + + payload = torch.load(feature_path, map_location="cpu", weights_only=False) + if not isinstance(payload, dict) or "states" not in payload: + raise ValueError(f"Invalid feature payload: {feature_path}") + states = payload["states"] + if not isinstance(states, torch.Tensor) or states.ndim != 2: + raise ValueError(f"states must be [T,D] in {feature_path}") + if not torch.is_floating_point(states): + raise ValueError(f"states must be floating point in {feature_path}") + if not torch.isfinite(states).all(): + raise ValueError(f"states contain NaN/Inf in {feature_path}") + + mode = payload.get("mode") + if expected_mode is not None and mode != expected_mode: + raise ValueError( + f"{feature_path} mode is {mode!r}, expected {expected_mode!r}" + ) + label_indices = payload.get("label_indices") + if label_indices is None: + raise ValueError(f"Missing label_indices in {feature_path}") + label_indices = torch.as_tensor(label_indices, dtype=torch.long) + if label_indices.ndim != 1 or len(label_indices) != len(states): + raise ValueError(f"Invalid label_indices in {feature_path}") + if len(label_indices) and not torch.equal( + label_indices, + torch.arange(len(label_indices), dtype=torch.long), + ): + raise ValueError( + f"{feature_path} does not use the verified zero-based dense alignment" + ) + + target_array = np.load(target_path, allow_pickle=False) + if target_array.ndim != 2 or target_array.shape[1] != NUM_CLASSES: + raise ValueError( + f"Expected target [T,{NUM_CLASSES}], got {target_array.shape}: " + f"{target_path}" + ) + targets = torch.from_numpy(np.asarray(target_array, dtype=np.float32)) + if len(label_indices) and int(label_indices[-1]) >= len(targets): + raise ValueError( + f"Last label index {int(label_indices[-1])} exceeds target length " + f"{len(targets)} in {name}" + ) + if len(states) != len(targets): + raise ValueError( + f"{name}: feature/target length mismatch: {len(states)} vs " + f"{len(targets)}" + ) + + metadata = { + key: payload.get(key) + for key in ( + "model_id", + "mode", + "config", + "start", + "end", + "num_steps", + "video_sampling_grid", + ) + } + return SessionData(name, states, targets, label_indices, metadata) + + +def average_precision_metrics( + scores: np.ndarray, + targets: np.ndarray, +) -> dict[str, Any]: + """Compute EgoSpeak frame AP for all three classes, including background.""" + + from sklearn.metrics import average_precision_score + + scores = np.asarray(scores, dtype=np.float64) + targets = np.asarray(targets, dtype=np.float64) + if scores.shape != targets.shape or scores.ndim != 2: + raise ValueError( + f"scores/targets must have equal [N,C] shape, got " + f"{scores.shape} and {targets.shape}" + ) + if scores.shape[1] != NUM_CLASSES: + raise ValueError(f"Expected {NUM_CLASSES} classes, got {scores.shape[1]}") + + aps = average_precision_score(targets, scores, average=None) + return { + "num_steps": int(len(scores)), + "class_ap": { + name: float(value) for name, value in zip(CLASS_NAMES, aps) + }, + "map": float(np.mean(aps)), + } + + +def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]: + """Return a validated head config, including legacy-checkpoint defaults.""" + + config = dict(payload.get("head_config", {})) + config.setdefault("hidden_size", 2048) + config.setdefault("num_horizons", NUM_HORIZONS) + config.setdefault("num_classes", NUM_CLASSES) + config.setdefault("dropout", 0.1) + return config + + +def validate_disjoint_sessions( + first: Sequence[str], + second: Sequence[str], +) -> None: + overlap = set(first) & set(second) + if overlap: + raise ValueError(f"Session sets overlap: {sorted(overlap)}") diff --git a/infer_turn_head.py b/infer_turn_head.py new file mode 100644 index 0000000..bcd1da6 --- /dev/null +++ b/infer_turn_head.py @@ -0,0 +1,289 @@ +"""Run the Omni anticipation head and evaluate it with EgoSpeak's EasyCom metric.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from tqdm.auto import tqdm + +from easycom_head import ( + CLASS_NAMES, + NUM_CLASSES, + NUM_HORIZONS, + OFFICIAL_TEST_SESSIONS, + STEP_SECONDS, + TurnPredictionHead, + average_precision_metrics, + checkpoint_head_config, + load_session_data, + normalize_session_name, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--feature-root", type=Path, required=True) + parser.add_argument("--target-root", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument( + "--sessions", + nargs="+", + default=list(OFFICIAL_TEST_SESSIONS), + help="Evaluation sessions; default is the official Session 3/2/1 test split.", + ) + parser.add_argument("--batch-size", type=int, default=1024) + parser.add_argument("--device", default="cuda:0") + parser.add_argument( + "--expected-mode", + default="causal-prefix", + help="Expected feature mode, or 'any' to disable the check.", + ) + parser.add_argument("--no-progress", action="store_true") + return parser.parse_args() + + +def _resolve_device(name: str) -> torch.device: + device = torch.device(name) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError(f"CUDA requested but unavailable: {name}") + return device + + +def _load_head( + checkpoint: Path, + device: torch.device, +) -> tuple[TurnPredictionHead, dict[str, Any]]: + payload = torch.load(checkpoint, map_location="cpu", weights_only=False) + if not isinstance(payload, dict): + raise ValueError(f"Invalid checkpoint: {checkpoint}") + state_dict = payload.get("head") + if not isinstance(state_dict, dict): + raise ValueError(f"Checkpoint has no head state_dict: {checkpoint}") + config = checkpoint_head_config(payload) + if int(config["num_horizons"]) != NUM_HORIZONS: + raise ValueError( + f"Checkpoint has {config['num_horizons']} horizons, expected {NUM_HORIZONS}" + ) + if int(config["num_classes"]) != NUM_CLASSES: + raise ValueError( + f"Checkpoint has {config['num_classes']} classes, expected {NUM_CLASSES}" + ) + head = TurnPredictionHead(**config) + head.load_state_dict(state_dict) + return head.to(device).eval(), payload + + +@torch.inference_mode() +def _predict_states( + head: TurnPredictionHead, + states: torch.Tensor, + device: torch.device, + batch_size: int, + *, + description: str, + show_progress: bool, +) -> np.ndarray: + chunks: list[torch.Tensor] = [] + starts = range(0, len(states), batch_size) + progress = tqdm( + starts, + total=(len(states) + batch_size - 1) // batch_size, + desc=description, + unit="batch", + disable=not show_progress, + ) + for start in progress: + batch = states[start : start + batch_size].to( + device=device, + dtype=torch.float32, + non_blocking=True, + ) + chunks.append(head(batch).softmax(dim=-1).cpu()) + return torch.cat(chunks, dim=0).numpy() + + +def _align_to_egospeak_timeline( + probabilities: np.ndarray, + label_indices: torch.LongTensor, + num_target_steps: int, +) -> tuple[np.ndarray, np.ndarray]: + """Return EgoSpeak ``[T,3,10]`` scores and their causal-valid mask.""" + + predictions = np.zeros( + (num_target_steps, NUM_CLASSES, NUM_HORIZONS), + dtype=np.float32, + ) + valid = np.zeros((num_target_steps, NUM_HORIZONS), dtype=bool) + current = label_indices.numpy() + for horizon in range(NUM_HORIZONS): + target_indices = current + horizon + 1 + usable = (target_indices >= 0) & (target_indices < num_target_steps) + rows = target_indices[usable] + predictions[rows, :, horizon] = probabilities[usable, horizon, :] + valid[rows, horizon] = True + return predictions, valid + + +def _evaluate( + session_predictions: dict[str, np.ndarray], + session_valid: dict[str, np.ndarray], + session_targets: dict[str, np.ndarray], +) -> dict[str, Any]: + full_results = [] + valid_results = [] + for horizon in range(NUM_HORIZONS): + full_scores = np.concatenate( + [scores[:, :, horizon] for scores in session_predictions.values()], + axis=0, + ) + full_targets = np.concatenate(list(session_targets.values()), axis=0) + full = average_precision_metrics(full_scores, full_targets) + full["horizon"] = horizon + 1 + full["seconds"] = (horizon + 1) * STEP_SECONDS + full_results.append(full) + + valid_scores = np.concatenate( + [ + session_predictions[name][session_valid[name][:, horizon], :, horizon] + for name in session_predictions + ], + axis=0, + ) + valid_targets = np.concatenate( + [ + session_targets[name][session_valid[name][:, horizon]] + for name in session_predictions + ], + axis=0, + ) + causal = average_precision_metrics(valid_scores, valid_targets) + causal["horizon"] = horizon + 1 + causal["seconds"] = (horizon + 1) * STEP_SECONDS + valid_results.append(causal) + + return { + "egospeak_full_timeline": { + "description": ( + "Exact EgoSpeak array/evaluator layout. Unavailable pre-session " + "causal predictions remain zero in the preallocated arrays." + ), + "horizons": full_results, + "mean_map": float(np.mean([item["map"] for item in full_results])), + }, + "causal_valid_only": { + "description": ( + "Primary causal metric; excludes only the first h rows of each " + "session for horizon h because no pre-session context exists." + ), + "horizons": valid_results, + "mean_map": float(np.mean([item["map"] for item in valid_results])), + }, + } + + +def _print_metrics(metrics: dict[str, Any]) -> None: + for protocol_name in ("egospeak_full_timeline", "causal_valid_only"): + protocol = metrics[protocol_name] + print(f"\n{protocol_name}:") + for result in protocol["horizons"]: + print( + f"Action anticipation ({result['seconds']:.2f}s) " + f"perframe mAP: {result['map'] * 100:.5f}" + ) + class_text = ", ".join( + f"{name}={value * 100:.5f}" + for name, value in result["class_ap"].items() + ) + print(f" AP: {class_text}; n={result['num_steps']}") + print( + f"Action anticipation (mean) perframe mAP: " + f"{protocol['mean_map'] * 100:.5f}" + ) + + +def main() -> None: + args = parse_args() + if args.batch_size <= 0: + raise ValueError("--batch-size must be positive") + device = _resolve_device(args.device) + expected_mode = None if args.expected_mode == "any" else args.expected_mode + sessions = [normalize_session_name(item) for item in args.sessions] + if len(sessions) != len(set(sessions)): + raise ValueError("Duplicate sessions were supplied") + head, checkpoint_payload = _load_head(args.checkpoint, device) + head_config = checkpoint_head_config(checkpoint_payload) + + args.output_dir.mkdir(parents=True, exist_ok=True) + prediction_dir = args.output_dir / "pred" + valid_dir = args.output_dir / "valid" + prediction_dir.mkdir(parents=True, exist_ok=True) + valid_dir.mkdir(parents=True, exist_ok=True) + session_predictions: dict[str, np.ndarray] = {} + session_valid: dict[str, np.ndarray] = {} + session_targets: dict[str, np.ndarray] = {} + + print(f"device: {device}") + print(f"sessions: {', '.join(sessions)}") + for name in sessions: + data = load_session_data( + args.feature_root, + args.target_root, + name, + expected_mode=expected_mode, + ) + if data.hidden_size != int(head_config["hidden_size"]): + raise ValueError( + f"{name} feature dim is {data.hidden_size}, checkpoint expects " + f"{head_config['hidden_size']}" + ) + probabilities = _predict_states( + head, + data.states, + device, + args.batch_size, + description=f"infer {name}", + show_progress=not args.no_progress, + ) + predictions, valid = _align_to_egospeak_timeline( + probabilities, + data.label_indices, + len(data.targets), + ) + target_array = data.targets.numpy() + session_predictions[name] = predictions + session_valid[name] = valid + session_targets[name] = target_array + np.save(prediction_dir / f"{name}.npy", predictions) + np.save(valid_dir / f"{name}.npy", valid) + print( + f"{name}: saved pred {predictions.shape}, " + f"causal-valid counts={valid.sum(axis=0).tolist()}" + ) + + metrics = _evaluate(session_predictions, session_valid, session_targets) + metrics.update( + { + "sessions": sessions, + "classes": list(CLASS_NAMES), + "prediction_layout": "[target_time, class, future_horizon]", + "checkpoint": str(args.checkpoint.resolve()), + "checkpoint_epoch": checkpoint_payload.get("epoch"), + "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"), + "checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"), + } + ) + metrics_path = args.output_dir / "metrics.json" + metrics_path.write_text(json.dumps(metrics, indent=2) + "\n") + _print_metrics(metrics) + print(f"\nmetrics: {metrics_path}") + print(f"EgoSpeak-compatible predictions: {prediction_dir}") + + +if __name__ == "__main__": + main() diff --git a/prepare_easycom_session.py b/prepare_easycom_session.py index e9df04d..ac5734e 100644 --- a/prepare_easycom_session.py +++ b/prepare_easycom_session.py @@ -42,10 +42,22 @@ class ClipInfo: def _run_json(command: list[str]) -> dict: - completed = subprocess.run( - command, check=True, capture_output=True, text=True - ) - return json.loads(completed.stdout) + try: + completed = subprocess.run( + command, check=True, capture_output=True, text=True + ) + except subprocess.CalledProcessError as exc: + stderr = (exc.stderr or "").strip() + detail = f"\n{stderr}" if stderr else "" + raise RuntimeError( + f"Command failed ({exc.returncode}): {shlex.join(command)}{detail}" + ) from exc + try: + return json.loads(completed.stdout) + except json.JSONDecodeError as exc: + raise RuntimeError( + f"Command returned invalid JSON: {shlex.join(command)}" + ) from exc def probe_clip(path: Path) -> ClipInfo: @@ -108,16 +120,22 @@ def validate_inventory(clips: list[ClipInfo], target_steps: int | None) -> int: 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) + frames_per_step = SOURCE_FPS // TASK_FPS + media_steps = total_frames // frames_per_step 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}" + f"Dense media gives {media_steps} complete 5 Hz steps from " + f"{total_frames} frames, but target has {target_steps}" ) - return media_steps + # A session may end with 1--3 source frames, which cannot form another + # complete 5 Hz target interval. The target length is authoritative; the + # final video/audio trim below discards only this sub-step tail. + return target_steps if target_steps is not None else media_steps + + +def fixed_grid_frame_count(*, fps: float, dense_duration: float) -> int: + """Number of samples on the grid ``n / fps < dense_duration``.""" + return int(math.ceil(dense_duration * fps - 1e-7)) def build_ffmpeg_command( @@ -127,6 +145,9 @@ def build_ffmpeg_command( fps: float, dense_duration: float, ) -> list[str]: + expected_frames = fixed_grid_frame_count( + fps=fps, dense_duration=dense_duration + ) # 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. @@ -158,11 +179,19 @@ def build_ffmpeg_command( "".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. + # Establish one global visual lattice after the clips have become a dense + # timeline. Some FFmpeg versions omit the last grid point when the input + # ends exactly on an integer-second boundary. Add two disposable clones, + # then trim by frame index, so EOF timestamp rounding cannot change the + # number of samples. The clones are retained only if a real final grid + # sample was missing. filters.append( - f"[vcat]fps=fps={fps}:start_time=0:round=down,format=yuv420p[vout]" + f"[vcat]trim=start=0:end={dense_duration:.9f}," + f"setpts=PTS-STARTPTS," + f"fps=fps={fps}:start_time=0:round=down," + f"tpad=stop_mode=clone:stop=2," + f"trim=start_frame=0:end_frame={expected_frames}," + f"setpts=N/({fps}*TB),format=yuv420p[vout]" ) filters.append( f"[acat]atrim=start=0:end={dense_duration:.9f}," @@ -220,20 +249,15 @@ def validate_output(path: Path, *, fps: float, dense_duration: float) -> dict: 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)), - } + expected_frames = fixed_grid_frame_count( + fps=fps, dense_duration=dense_duration + ) 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: + if actual_frames != expected_frames: raise RuntimeError( - f"Prepared frame count is {actual_frames}, expected one of " - f"{sorted(expected_frames)}: {path}" + f"Prepared frame count is {actual_frames}, expected " + f"{expected_frames}: {path}" ) audio_duration = float(audio["duration"]) # AAC frame padding can shift the container duration by a few milliseconds. @@ -245,7 +269,7 @@ def validate_output(path: Path, *, fps: float, dense_duration: float) -> dict: return { "video_fps": actual_fps, "video_frames": actual_frames, - "allowed_video_frames": sorted(expected_frames), + "expected_video_frames": expected_frames, "video_duration_seconds": float(video["duration"]), "audio_duration_seconds": audio_duration, "audio_sample_rate": int(audio["sample_rate"]), @@ -322,12 +346,16 @@ def main() -> None: output, fps=fps, dense_duration=duration ) + source_frames = sum(clip.frames for clip in clips) + used_source_frames = num_steps * (SOURCE_FPS // TASK_FPS) 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), + "source_frames_20hz": source_frames, + "used_source_frames_20hz": used_source_frames, + "discarded_tail_frames_20hz": source_frames - used_source_frames, "task_fps": TASK_FPS, "num_steps": num_steps, "dense_duration_seconds": duration, diff --git a/requirements.txt b/requirements.txt index 10df5a8..0690fb7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,5 +5,7 @@ accelerate qwen-omni-utils[decord] librosa Pillow +tqdm +scikit-learn # Optional on compatible NVIDIA GPUs: flash-attn diff --git a/run_easycom_head_experiment.sh b/run_easycom_head_experiment.sh new file mode 100755 index 0000000..bbf248d --- /dev/null +++ b/run_easycom_head_experiment.sh @@ -0,0 +1,35 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Frozen-Qwen head equivalent of EgoSpeak's run_LSTR_{train,inference}.sh. +# Override any path or device by setting the corresponding environment value. +ACTION="${1:-all}" +if [[ "${ACTION}" != "train" && "${ACTION}" != "infer" && "${ACTION}" != "all" ]]; then + echo "Usage: $0 [train|infer|all]" >&2 + exit 2 +fi +PROJECT_ROOT="${PROJECT_ROOT:-/data1/yfzx/wangx/Code/test_qwen_omni}" +FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/causal_prefix_2fps}" +TARGET_ROOT="${TARGET_ROOT:-${PROJECT_ROOT}/target_frame}" +CHECKPOINT="${CHECKPOINT:-${PROJECT_ROOT}/checkpoints/causal_prefix_2fps_head.pth}" +RESULT_ROOT="${RESULT_ROOT:-${PROJECT_ROOT}/results/causal_prefix_2fps_head}" +DEVICE="${DEVICE:-cuda:0}" + +cd "${PROJECT_ROOT}" + +if [[ "${ACTION}" == "train" || "${ACTION}" == "all" ]]; then + python train_turn_head.py \ + --feature-root "${FEATURE_ROOT}" \ + --target-root "${TARGET_ROOT}" \ + --output "${CHECKPOINT}" \ + --device "${DEVICE}" +fi + +if [[ "${ACTION}" == "infer" || "${ACTION}" == "all" ]]; then + python infer_turn_head.py \ + --feature-root "${FEATURE_ROOT}" \ + --target-root "${TARGET_ROOT}" \ + --checkpoint "${CHECKPOINT}" \ + --output-dir "${RESULT_ROOT}" \ + --device "${DEVICE}" +fi diff --git a/streaming_turn_model.py b/streaming_turn_model.py index a822db1..befea2b 100644 --- a/streaming_turn_model.py +++ b/streaming_turn_model.py @@ -4,24 +4,22 @@ This module implements two inference paths for Qwen2.5-Omni-3B: 1. ``causal-prefix``: recompute a causal AV window ending at every 200 ms. This is deliberately slow, but is the simplest correctness baseline. -2. ``commit-replay``: keep the Thinker KV cache for completed 2-second - multimodal blocks, and replay only the causal prefix of the current block. -3. ``open-stream-replay``: cache the same completed AV blocks inside one +2. ``open-stream-replay``: cache completed AV blocks inside one still-open user/media turn; closure and readout tokens remain temporary. Neither path uses an RNN. A small head maps one Thinker state to the same ``10 future horizons x 3 classes`` output used by EgoSpeak. -The commit/replay path touches version-specific Hugging Face internals. It is -therefore pinned to Transformers 4.52.3 and includes an optional exact-replay -check. Run that check before extracting a full dataset. +The open-stream replay path touches version-specific Hugging Face internals. +It is therefore pinned to Transformers 4.52.3 and includes an optional +exact-replay check. Run that check before extracting a full dataset. Example (no command is run merely by importing this file): python streaming_turn_model.py \ --video /path/to/session.mp4 \ --use-audio-in-video \ - --mode commit-replay \ + --mode open-stream-replay \ --memory-seconds 20 \ --verify-every 25 \ --output features/session.pt @@ -32,6 +30,7 @@ from __future__ import annotations import argparse import copy import json +import logging import math import subprocess from dataclasses import asdict, dataclass @@ -39,16 +38,58 @@ from pathlib import Path from typing import Any, Literal import torch +from easycom_head import ( + TurnPredictionHead, + egospeak_multihorizon_loss, + miniroad_multihorizon_loss, + multihorizon_cross_entropy, +) from qwen_omni_utils import process_mm_info from torch import nn +from tqdm.auto import tqdm from transformers import ( Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration, ) from transformers.cache_utils import DynamicCache -Mode = Literal["causal-prefix", "commit-replay", "open-stream-replay"] +Mode = Literal["causal-prefix", "open-stream-replay"] MODEL_ID = "Qwen/Qwen2.5-Omni-3B" +_KNOWN_QWEN_WARNING_TEXT = ( + "System prompt modified, audio output may not work", + "Unused or unrecognized kwargs: images.", + "Unused or unrecognized kwargs: return_tensors, images.", + "Unused or unrecognized kwargs: images, return_tensors.", +) + + +class _KnownQwenWarningFilter(logging.Filter): + """Hide Qwen's speech-output warning in Thinker-only extraction.""" + + def filter(self, record: logging.LogRecord) -> bool: + message = record.getMessage() + return not any(text in message for text in _KNOWN_QWEN_WARNING_TEXT) + + +def _configure_qwen_logging() -> None: + # Qwen2.5-Omni 4.52.3 emits this through Python's root logger whenever the + # conversation does not use its speech-generation system prompt. This + # extractor uses only the Thinker and never requests audio output. + def add_filter(target: logging.Filterer) -> None: + if not any( + isinstance(item, _KnownQwenWarningFilter) + for item in target.filters + ): + target.addFilter(_KnownQwenWarningFilter()) + + # The system-prompt warning is logged directly on root. The unused-kwargs + # warning is emitted by a Transformers child logger and handled by the + # library's top-level StreamHandler, so cover both paths. + for logger_name in (None, "transformers"): + logger = logging.getLogger(logger_name) + add_filter(logger) + for handler in logger.handlers: + add_filter(handler) def _choose_dtype( @@ -132,6 +173,7 @@ class OpenMediaPieces: payload_embeddings: torch.Tensor payload_position_ids: torch.LongTensor closure_ids: torch.LongTensor + closure_position_ids: torch.LongTensor class FixedGridAV: @@ -167,19 +209,22 @@ class FixedGridAV: ) def frames(self, start: float, end: float) -> list[Any]: - """Return fixed-grid frames whose timestamps lie in [start, end].""" + """Return fixed-grid frames whose timestamps lie in [start, end).""" import numpy as np from PIL import Image first = max(0, int(math.ceil(start * self.fps - 1e-7))) - last = min( - self.num_frames - 1, - int(math.floor(end * self.fps + 1e-7)), + last_exclusive = min( + self.num_frames, + int(math.ceil(end * self.fps - 1e-7)), ) + last = last_exclusive - 1 if first > last: - # This can occur only when a caller starts between global visual - # ticks. Reuse the most recent past frame, never ceil into future. + # A short interval can fall entirely between global visual ticks. + # Reuse the latest frame strictly before ``end``, never a future + # frame. Native 2-second blocks start on-grid and do not take this + # fallback path. first = last = max(0, min(self.num_frames - 1, last)) indices = list(range(first, last + 1)) arrays = self.reader.get_batch(indices).asnumpy() @@ -198,76 +243,6 @@ class FixedGridAV: return self.audio[first:last] -class TurnPredictionHead(nn.Module): - """Map one Omni state to EgoSpeak-compatible [10, 3] logits.""" - - def __init__( - self, - hidden_size: int = 2048, - num_horizons: int = 10, - num_classes: int = 3, - dropout: float = 0.1, - ) -> None: - super().__init__() - self.num_horizons = num_horizons - self.num_classes = num_classes - self.net = nn.Sequential( - nn.LayerNorm(hidden_size), - nn.Linear(hidden_size, hidden_size // 2), - nn.GELU(), - nn.Dropout(dropout), - nn.Linear(hidden_size // 2, num_horizons * num_classes), - ) - - def forward(self, states: torch.Tensor) -> torch.Tensor: - logits = self.net(states) - return logits.view(*states.shape[:-1], self.num_horizons, self.num_classes) - - -def egospeak_multihorizon_loss( - logits: torch.Tensor, - targets: torch.Tensor, - *, - reduction: Literal["mean", "sum", "none"] = "mean", -) -> torch.Tensor: - """Reproduce EgoSpeak's anticipation loss for multi-positive targets. - - ``logits`` and ``targets`` both have shape ``[..., 10, 3]``. EasyCom can - label camera-wearer and other speech simultaneously as ``[0, 1, 1]``. - MiniROAD L2-normalizes that row and applies a log-softmax loss. - """ - - if logits.shape != targets.shape: - raise ValueError( - f"logits/targets shape mismatch: {tuple(logits.shape)} vs " - f"{tuple(targets.shape)}" - ) - if logits.shape[-1] != 3: - raise ValueError(f"Expected three classes, got {logits.shape[-1]}") - targets = nn.functional.normalize( - targets.to(device=logits.device, dtype=logits.dtype), - p=2, - dim=-1, - ) - loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1) - if reduction == "mean": - return loss.mean() - if reduction == "sum": - return loss.sum() - if reduction == "none": - return loss - raise ValueError(f"Unknown reduction: {reduction}") - - -def multihorizon_cross_entropy( - logits: torch.Tensor, - targets: torch.Tensor, -) -> torch.Tensor: - """Backward-compatible name for :func:`egospeak_multihorizon_loss`.""" - - return egospeak_multihorizon_loss(logits, targets) - - def _require_integer_ratio(numerator: float, denominator: float, name: str) -> int: ratio = numerator / denominator rounded = round(ratio) @@ -437,16 +412,23 @@ def _prepare_conversation( # Qwen2.5-Omni uses this value to compute video_second_per_grid for # TMRoPE. This is essential for the 5 FPS ablation (default is 2 FPS). processor_kwargs["fps"] = video_fps - batch = processor( - text=[text], - audio=audios, - images=images, - videos=videos, - padding=True, - return_tensors="pt", - use_audio_in_video=use_audio_in_video, + processor_inputs: dict[str, Any] = { + "text": [text], + "audio": audios, + "videos": videos, + "padding": True, + "use_audio_in_video": use_audio_in_video, **processor_kwargs, - ) + } + # Passing images=None and return_tensors through Qwen's 4.52.3 processor + # produces a warning on every streaming step. Omit the absent modality and + # perform the same BatchFeature tensor conversion explicitly afterward. + if images is not None and ( + not isinstance(images, (list, tuple)) or len(images) > 0 + ): + processor_inputs["images"] = images + batch = processor(**processor_inputs) + batch = batch.convert_to_tensors("pt") # BatchFeature.to(dtype) changes floating tensors only; IDs stay integral. batch = batch.to(model.device).to(model.dtype) return dict(batch) @@ -662,18 +644,10 @@ def _open_media_pieces( payload_embeddings=embeddings[:, first : last + 1, :], payload_position_ids=positions[:, :, first : last + 1], closure_ids=closure_ids, + closure_position_ids=positions[:, :, last + 1 :], ) -def _positions_at_offset( - position_ids: torch.LongTensor, - offset: int, -) -> torch.LongTensor: - """Preserve local 3D geometry while moving its minimum to ``offset``.""" - - return position_ids - position_ids.min() + offset - - def _tmrope_positions_per_chunk( model: Qwen2_5OmniThinkerForConditionalGeneration, ) -> int: @@ -823,13 +797,22 @@ def _read_open_state( model: Qwen2_5OmniThinkerForConditionalGeneration, open_media_cache: Any, *, - next_rope_offset: int, closure_ids: torch.LongTensor, + closure_position_ids: torch.LongTensor, query_ids: torch.LongTensor, ) -> torch.Tensor: """Temporarily close an open media turn, read q_t, and discard the branch.""" branch_cache = clone_cache(open_media_cache) + if closure_position_ids.shape != ( + 3, + closure_ids.shape[0], + closure_ids.shape[1], + ): + raise ValueError( + "closure position/ID mismatch: " + f"{tuple(closure_position_ids.shape)} vs {tuple(closure_ids.shape)}" + ) readout_ids = torch.cat( [closure_ids.to(query_ids.device), query_ids], dim=1, @@ -837,13 +820,24 @@ def _read_open_state( past_length = _cache_length(branch_cache) readout_length = readout_ids.shape[1] embeddings = model.get_input_embeddings()(readout_ids) - scalar_positions = torch.arange( - next_rope_offset, - next_rope_offset + readout_length, + + # Qwen's AV TMRoPE gives the paired video/audio closing tokens the same + # position. Preserve those model-derived positions instead of replacing + # the whole closure with a scalar arange. Only the artificial assistant + # query follows the closure's maximum position. + closure_position_ids = closure_position_ids.to(query_ids.device) + query_start = int(closure_position_ids.max().item()) + 1 + query_scalar_positions = torch.arange( + query_start, + query_start + query_ids.shape[1], device=readout_ids.device, ) - position_ids = scalar_positions.view(1, 1, -1).expand( - 3, readout_ids.shape[0], -1 + query_position_ids = query_scalar_positions.view(1, 1, -1).expand( + 3, query_ids.shape[0], -1 + ) + position_ids = torch.cat( + [closure_position_ids, query_position_ids], + dim=2, ) cache_position = torch.arange( past_length, @@ -876,32 +870,41 @@ def _append_open_pieces( past_key_values: Any | None, rope_offset: int, include_prefix: bool, -) -> tuple[Any, int]: - """Append one AV payload, adding the chat/media prefix exactly once.""" +) -> tuple[Any, int, torch.LongTensor]: + """Append one AV payload and return its shifted closure positions.""" if include_prefix: if past_key_values is not None: raise ValueError("An open-media prefix can only start an empty cache") + payload_base = int(pieces.payload_position_ids.min().item()) + if rope_offset != payload_base: + raise ValueError( + "The first open payload must retain its local TMRoPE base: " + f"got {rope_offset}, expected {payload_base}" + ) embeddings = torch.cat( [pieces.prefix_embeddings, pieces.payload_embeddings], dim=1 ) position_ids = torch.cat( [pieces.prefix_position_ids, pieces.payload_position_ids], dim=2 ) + closure_position_ids = pieces.closure_position_ids else: if past_key_values is None: raise ValueError("Payload-only append requires an existing open cache") embeddings = pieces.payload_embeddings - position_ids = _positions_at_offset( - pieces.payload_position_ids, - rope_offset, + position_shift = rope_offset - int( + pieces.payload_position_ids.min().item() ) - return _append_embeddings( + position_ids = pieces.payload_position_ids + position_shift + closure_position_ids = pieces.closure_position_ids + position_shift + next_cache, next_rope_offset = _append_embeddings( model, embeddings, position_ids, past_key_values=past_key_values, ) + return next_cache, next_rope_offset, closure_position_ids class CausalPrefixExtractor: @@ -982,18 +985,34 @@ class CausalPrefixExtractor: start: float, end: float, use_audio_in_video: bool, + show_progress: bool = True, ) -> ExtractionResult: times = _time_grid(start, end, self.config.step_seconds) states = [] - for current_time in times: - context_start = max(start, current_time - self.config.memory_seconds) - state = self.state_for_intervals( - video_path=video_path, - audio_path=audio_path, - intervals=[(context_start, current_time)], - use_audio_in_video=use_audio_in_video, - ) - states.append(state.cpu()) + with tqdm( + total=len(times), + desc="causal-prefix", + unit="state", + dynamic_ncols=True, + mininterval=0.0, + miniters=1, + disable=not show_progress, + ) as progress: + for current_time in times: + context_start = max( + start, current_time - self.config.memory_seconds + ) + state = self.state_for_intervals( + video_path=video_path, + audio_path=audio_path, + intervals=[(context_start, current_time)], + use_audio_in_video=use_audio_in_video, + ) + states.append(state.cpu()) + progress.set_postfix_str( + f"media={current_time:.1f}s", refresh=False + ) + progress.update(1) return ExtractionResult( times=torch.tensor(times, dtype=torch.float64), states=torch.cat(states, dim=0), @@ -1001,202 +1020,6 @@ class CausalPrefixExtractor: ) -class CommitReplayExtractor: - """Cache complete 2 s blocks and replay only the current causal prefix.""" - - def __init__( - self, - model: Qwen2_5OmniThinkerForConditionalGeneration, - processor: Qwen2_5OmniProcessor, - config: StreamConfig, - fixed_grid_av: FixedGridAV | None = None, - ) -> None: - config.validate() - self.model = model - self.processor = processor - self.config = config - self.fixed_grid_av = fixed_grid_av - self.query_ids = _readout_query_ids(processor, model.device) - self.reference = CausalPrefixExtractor( - model, processor, config, fixed_grid_av=fixed_grid_av - ) - - def _prepare_one( - self, - *, - video_path: str | Path | None, - audio_path: str | Path | None, - interval: tuple[float, float], - use_audio_in_video: bool, - turn_index: int, - ) -> dict[str, torch.Tensor]: - conversation, manual_audios = build_chunk_conversation( - video_path=video_path, - audio_path=audio_path, - intervals=[interval], - use_audio_in_video=use_audio_in_video, - config=self.config, - fixed_grid_av=self.fixed_grid_av, - ) - text_override = _media_turn_template_suffix( - self.processor, - conversation[0], - turn_index=turn_index, - ) - return _prepare_conversation( - self.processor, - conversation, - use_audio_in_video=use_audio_in_video, - model=self.model, - manual_audios=manual_audios, - video_fps=self.config.video_fps, - text_override=text_override, - ) - - def _extract_segment( - self, - *, - video_path: str | Path | None, - audio_path: str | Path | None, - segment_start: float, - segment_end: float, - use_audio_in_video: bool, - verify_every: int, - verify_continuous_every: int, - global_step_offset: int, - ) -> tuple[list[float], list[torch.Tensor], list[dict[str, Any]]]: - times = _time_grid(segment_start, segment_end, self.config.step_seconds) - committed_cache = None - committed_rope_offset = 0 - block_start = segment_start - states: list[torch.Tensor] = [] - checks: list[dict[str, Any]] = [] - - for local_step, current_time in enumerate(times, start=1): - while current_time > block_start + self.config.native_block_seconds + 1e-6: - block_start += self.config.native_block_seconds - - # DynamicCache is mutated by forward. Every unfinished prefix must - # branch from the identical committed history. - branch_cache = clone_cache(committed_cache) - inputs = self._prepare_one( - video_path=video_path, - audio_path=audio_path, - interval=(block_start, current_time), - use_audio_in_video=use_audio_in_video, - turn_index=round( - (block_start - segment_start) - / self.config.native_block_seconds - ), - ) - media_cache, next_rope_offset = _append_media_turn( - self.model, - inputs, - past_key_values=branch_cache, - rope_offset=committed_rope_offset, - use_audio_in_video=use_audio_in_video, - ) - state = _read_state( - self.model, - media_cache, - next_rope_offset=next_rope_offset, - query_ids=self.query_ids, - ) - states.append(state.cpu()) - - global_step = global_step_offset + local_step - if verify_every > 0 and global_step % verify_every == 0: - intervals = _block_intervals( - segment_start, - current_time, - self.config.native_block_seconds, - ) - exact = self.reference.state_for_intervals( - video_path=video_path, - audio_path=audio_path, - intervals=intervals, - use_audio_in_video=use_audio_in_video, - ) - checks.append( - _compare_states(current_time, state, exact, kind="cache_replay") - ) - if ( - verify_continuous_every > 0 - and global_step % verify_continuous_every == 0 - ): - continuous = self.reference.state_for_intervals( - video_path=video_path, - audio_path=audio_path, - intervals=[(segment_start, current_time)], - use_audio_in_video=use_audio_in_video, - ) - checks.append( - _compare_states( - current_time, - state, - continuous, - kind="continuous_window_semantics", - ) - ) - - block_end = min( - block_start + self.config.native_block_seconds, - segment_end, - ) - if math.isclose(current_time, block_end, rel_tol=0.0, abs_tol=1e-6): - # media_cache excludes the readout query, so it is safe to keep. - committed_cache = media_cache - committed_rope_offset = next_rope_offset - block_start = block_end - - return times, states, checks - - def extract( - self, - *, - video_path: str | Path | None, - audio_path: str | Path | None, - start: float, - end: float, - use_audio_in_video: bool, - verify_every: int = 0, - verify_continuous_every: int = 0, - ) -> ExtractionResult: - all_times: list[float] = [] - all_states: list[torch.Tensor] = [] - all_checks: list[dict[str, Any]] = [] - segment_start = start - global_steps = 0 - - # A bounded cache is implemented as non-overlapping memory segments. - # Predictions are still emitted for the entire session, but history is - # reset at each boundary. Rolling KV eviction is intentionally not - # attempted because it also requires rebasing multimodal RoPE. - while segment_start < end - 1e-7: - segment_end = min(segment_start + self.config.memory_seconds, end) - times, states, checks = self._extract_segment( - video_path=video_path, - audio_path=audio_path, - segment_start=segment_start, - segment_end=segment_end, - use_audio_in_video=use_audio_in_video, - verify_every=verify_every, - verify_continuous_every=verify_continuous_every, - global_step_offset=global_steps, - ) - all_times.extend(times) - all_states.extend(states) - all_checks.extend(checks) - global_steps += len(times) - segment_start = segment_end - - return ExtractionResult( - times=torch.tensor(all_times, dtype=torch.float64), - states=torch.cat(all_states, dim=0), - verification=all_checks, - ) - - class OpenStreamReplayExtractor: """Cache completed AV payloads inside one still-open multimodal turn.""" @@ -1271,7 +1094,13 @@ class OpenStreamReplayExtractor: ) -> torch.Tensor: """Replay the same open stream from an empty cache for verification.""" - media_cache, next_offset, _, closure_ids = self._cache_from_intervals( + ( + media_cache, + _, + _, + closure_ids, + closure_position_ids, + ) = self._cache_from_intervals( video_path=video_path, audio_path=audio_path, intervals=intervals, @@ -1280,8 +1109,8 @@ class OpenStreamReplayExtractor: return _read_open_state( self.model, media_cache, - next_rope_offset=next_offset, closure_ids=closure_ids, + closure_position_ids=closure_position_ids, query_ids=self.query_ids, ) @@ -1292,7 +1121,7 @@ class OpenStreamReplayExtractor: audio_path: str | Path | None, intervals: list[tuple[float, float]], use_audio_in_video: bool, - ) -> tuple[Any, int, int, torch.LongTensor]: + ) -> tuple[Any, int, int, torch.LongTensor, torch.LongTensor]: """Build an open payload cache and expose its local TMRoPE base.""" if not intervals: @@ -1300,6 +1129,7 @@ class OpenStreamReplayExtractor: media_cache = None next_offset = 0 closure_ids = None + closure_position_ids = None payload_base = None positions_per_chunk = _tmrope_positions_per_chunk(self.model) for index, interval in enumerate(intervals): @@ -1312,7 +1142,11 @@ class OpenStreamReplayExtractor: if payload_base is None: payload_base = int(pieces.payload_position_ids.min().item()) payload_offset = payload_base + index * positions_per_chunk - media_cache, next_offset = _append_open_pieces( + ( + media_cache, + next_offset, + closure_position_ids, + ) = _append_open_pieces( self.model, pieces, past_key_values=media_cache, @@ -1322,9 +1156,17 @@ class OpenStreamReplayExtractor: closure_ids = pieces.closure_ids if closure_ids is None: raise RuntimeError("Open-stream replay did not produce closure tokens") + if closure_position_ids is None: + raise RuntimeError("Open-stream replay did not produce closure positions") if payload_base is None: raise RuntimeError("Open-stream replay did not establish a payload base") - return media_cache, next_offset, payload_base, closure_ids + return ( + media_cache, + next_offset, + payload_base, + closure_ids, + closure_position_ids, + ) def _extract_segment( self, @@ -1337,6 +1179,7 @@ class OpenStreamReplayExtractor: verify_every: int, verify_continuous_every: int, global_step_offset: int, + progress: Any, ) -> tuple[list[float], list[torch.Tensor], list[dict[str, Any]]]: times = _time_grid(segment_start, segment_end, self.config.step_seconds) committed_cache = None @@ -1373,6 +1216,7 @@ class OpenStreamReplayExtractor: _, payload_base, _, + _, ) = self._cache_from_intervals( video_path=video_path, audio_path=audio_path, @@ -1394,7 +1238,11 @@ class OpenStreamReplayExtractor: ) payload_offset = payload_base + block_index * positions_per_chunk branch_cache = clone_cache(committed_cache) - media_cache, next_offset = _append_open_pieces( + ( + media_cache, + _, + closure_position_ids, + ) = _append_open_pieces( self.model, pieces, past_key_values=branch_cache, @@ -1404,8 +1252,8 @@ class OpenStreamReplayExtractor: state = _read_open_state( self.model, media_cache, - next_rope_offset=next_offset, closure_ids=pieces.closure_ids, + closure_position_ids=closure_position_ids, query_ids=self.query_ids, ) states.append(state.cpu()) @@ -1454,6 +1302,11 @@ class OpenStreamReplayExtractor: committed_cache = media_cache block_start = block_end + progress.set_postfix_str( + f"media={current_time:.1f}s", refresh=False + ) + progress.update(1) + return times, states, checks def extract( @@ -1466,17 +1319,29 @@ class OpenStreamReplayExtractor: use_audio_in_video: bool, verify_every: int = 0, verify_continuous_every: int = 0, + show_progress: bool = True, ) -> ExtractionResult: - all_times, all_states, all_checks = self._extract_segment( - video_path=video_path, - audio_path=audio_path, - segment_start=start, - segment_end=end, - use_audio_in_video=use_audio_in_video, - verify_every=verify_every, - verify_continuous_every=verify_continuous_every, - global_step_offset=0, - ) + total_steps = len(_time_grid(start, end, self.config.step_seconds)) + with tqdm( + total=total_steps, + desc="open-stream-replay", + unit="state", + dynamic_ncols=True, + mininterval=0.0, + miniters=1, + disable=not show_progress, + ) as progress: + all_times, all_states, all_checks = self._extract_segment( + video_path=video_path, + audio_path=audio_path, + segment_start=start, + segment_end=end, + use_audio_in_video=use_audio_in_video, + verify_every=verify_every, + verify_continuous_every=verify_continuous_every, + global_step_offset=0, + progress=progress, + ) return ExtractionResult( times=torch.tensor(all_times, dtype=torch.float64), states=torch.cat(all_states, dim=0), @@ -1485,7 +1350,7 @@ class OpenStreamReplayExtractor: def _block_intervals(start: float, end: float, block_seconds: float) -> list[tuple[float, float]]: - """Exact media-turn sequence represented by a commit/replay cache.""" + """Exact AV payload intervals represented by an open-stream cache.""" intervals = [] cursor = start @@ -1637,7 +1502,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--output", type=Path, required=True) parser.add_argument( "--mode", - choices=("causal-prefix", "commit-replay", "open-stream-replay"), + choices=("causal-prefix", "open-stream-replay"), default="open-stream-replay", ) parser.add_argument("--model-id", default=MODEL_ID) @@ -1661,6 +1526,11 @@ def parse_args() -> argparse.Namespace: default=0, help="Also compare block replay with one continuous causal media window.", ) + parser.add_argument( + "--no-progress", + action="store_true", + help="Disable the per-state tqdm progress bar.", + ) parser.add_argument("--head-checkpoint", type=Path) parser.add_argument( "--dtype", @@ -1674,6 +1544,7 @@ def parse_args() -> argparse.Namespace: def main() -> None: args = parse_args() + _configure_qwen_logging() session_metadata = _apply_session_metadata(args) if args.video is None and args.audio is None: raise SystemExit("Pass --video, --audio, or both") @@ -1683,7 +1554,6 @@ def main() -> None: raise SystemExit("--verify-every must be >= 0") if args.verify_continuous_every < 0: raise SystemExit("--verify-continuous-every must be >= 0") - if args.num_steps is not None and args.num_steps <= 0: raise SystemExit("--num-steps must be positive") if args.target is not None: @@ -1758,19 +1628,7 @@ def main() -> None: start=args.start, end=end, use_audio_in_video=args.use_audio_in_video, - ) - elif args.mode == "commit-replay": - extractor = CommitReplayExtractor( - model, processor, config, fixed_grid_av=fixed_grid_av - ) - result = extractor.extract( - video_path=args.video, - audio_path=args.audio, - start=args.start, - end=end, - use_audio_in_video=args.use_audio_in_video, - verify_every=args.verify_every, - verify_continuous_every=args.verify_continuous_every, + show_progress=not args.no_progress, ) else: extractor = OpenStreamReplayExtractor( @@ -1784,6 +1642,7 @@ def main() -> None: use_audio_in_video=args.use_audio_in_video, verify_every=args.verify_every, verify_continuous_every=args.verify_continuous_every, + show_progress=not args.no_progress, ) payload: dict[str, Any] = { @@ -1792,11 +1651,7 @@ def main() -> None: "stream_representation": ( "single_causal_window" if args.mode == "causal-prefix" - else ( - "sequence_of_2s_media_turns" - if args.mode == "commit-replay" - else "single_open_media_turn_with_2s_payload_commits" - ) + else "single_open_media_turn_with_2s_payload_commits" ), "config": asdict(config), "start": args.start, @@ -1808,7 +1663,7 @@ def main() -> None: "session_metadata": session_metadata, "input_video_fps": input_video_fps, "video_sampling_grid": ( - "explicit_fixed_past_frames_with_duplication" + "explicit_fixed_half_open_past_frames_with_tail_duplication" if fixed_grid_av is not None else "qwen_utility_adaptive_sampling" ), diff --git a/train_turn_head.py b/train_turn_head.py new file mode 100644 index 0000000..cacb3d8 --- /dev/null +++ b/train_turn_head.py @@ -0,0 +1,476 @@ +"""Train the frozen-Omni EasyCom anticipation head. + +The defaults follow EgoSpeak's EasyCom LSTR optimization config where the +architectures are comparable: 50 epochs, Adam, 7e-5 learning rate, 5e-5 weight +decay, linear warmup for 10 epochs, cosine decay, dropout 0.2, and MCE loss. +""" + +from __future__ import annotations + +import argparse +import json +import math +import random +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch.optim.lr_scheduler import LambdaLR +from torch.utils.data import DataLoader, TensorDataset +from tqdm.auto import tqdm + +from easycom_head import ( + CLASS_NAMES, + NUM_CLASSES, + NUM_HORIZONS, + OFFICIAL_TRAIN_SESSIONS, + STEP_SECONDS, + TurnPredictionHead, + average_precision_metrics, + egospeak_multihorizon_loss, + load_session_data, + miniroad_multihorizon_loss, + normalize_session_name, + validate_disjoint_sessions, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--feature-root", type=Path, required=True) + parser.add_argument("--target-root", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--sessions", + nargs="+", + default=list(OFFICIAL_TRAIN_SESSIONS), + help="Training sessions; default is the official nine-session split.", + ) + parser.add_argument( + "--dev-sessions", + nargs="*", + default=[], + help=( + "Optional held-out training sessions used for checkpoint selection. " + "They are automatically removed from --sessions. Never use 1/2/3." + ), + ) + parser.add_argument("--epochs", type=int, default=50) + parser.add_argument("--batch-size", type=int, default=512) + parser.add_argument("--learning-rate", type=float, default=7e-5) + parser.add_argument("--weight-decay", type=float, default=5e-5) + parser.add_argument("--warmup-epochs", type=float, default=10.0) + parser.add_argument("--warmup-factor", type=float, default=0.3) + parser.add_argument("--dropout", type=float, default=0.2) + parser.add_argument( + "--loss", + choices=("lstr-mce", "miniroad-normalized"), + default="lstr-mce", + ) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--num-workers", type=int, default=0) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--expected-mode", + default="causal-prefix", + help="Expected feature mode, or 'any' to disable the check.", + ) + parser.add_argument("--no-progress", action="store_true") + return parser.parse_args() + + +def _validate_args(args: argparse.Namespace) -> None: + if args.epochs <= 0: + raise ValueError("--epochs must be positive") + if args.batch_size <= 0: + raise ValueError("--batch-size must be positive") + if args.learning_rate <= 0: + raise ValueError("--learning-rate must be positive") + if args.weight_decay < 0: + raise ValueError("--weight-decay cannot be negative") + if not 0 <= args.warmup_epochs <= args.epochs: + raise ValueError("--warmup-epochs must be in [0, epochs]") + if not 0 < args.warmup_factor <= 1: + raise ValueError("--warmup-factor must be in (0, 1]") + if not 0 <= args.dropout < 1: + raise ValueError("--dropout must be in [0, 1)") + if args.num_workers < 0: + raise ValueError("--num-workers cannot be negative") + + +def _resolve_device(name: str) -> torch.device: + device = torch.device(name) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError(f"CUDA requested but unavailable: {name}") + return device + + +def _seed_everything(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def _load_examples( + feature_root: Path, + target_root: Path, + sessions: list[str], + expected_mode: str | None, +) -> tuple[torch.Tensor, torch.Tensor, int, list[dict[str, Any]]]: + all_states: list[torch.Tensor] = [] + all_targets: list[torch.Tensor] = [] + metadata: list[dict[str, Any]] = [] + hidden_size: int | None = None + for name in sessions: + data = load_session_data( + feature_root, + target_root, + name, + expected_mode=expected_mode, + ) + states, targets, _ = data.future_examples(NUM_HORIZONS) + if hidden_size is None: + hidden_size = data.hidden_size + elif data.hidden_size != hidden_size: + raise ValueError( + f"Feature dimensions differ: {name} has {data.hidden_size}, " + f"expected {hidden_size}" + ) + all_states.append(states) + all_targets.append(targets) + metadata.append( + { + "session": name, + "num_states": len(data.states), + "num_training_examples": len(states), + **data.feature_metadata, + } + ) + print( + f"{name}: states={len(data.states)}, " + f"valid_10_horizon_examples={len(states)}" + ) + if hidden_size is None: + raise ValueError("No sessions selected") + return ( + torch.cat(all_states, dim=0), + torch.cat(all_targets, dim=0), + hidden_size, + metadata, + ) + + +def _scheduler( + optimizer: torch.optim.Optimizer, + *, + steps_per_epoch: int, + epochs: int, + warmup_epochs: float, + warmup_factor: float, +) -> LambdaLR: + total_steps = max(1, steps_per_epoch * epochs) + warmup_steps = steps_per_epoch * warmup_epochs + + def multiplier(step: int) -> float: + cosine = 0.5 * (1.0 + math.cos(math.pi * step / total_steps)) + if warmup_steps > 0 and step < warmup_steps: + alpha = step / warmup_steps + warmup = warmup_factor * (1.0 - alpha) + alpha + else: + warmup = 1.0 + return warmup * cosine + + return LambdaLR(optimizer, multiplier) + + +@torch.inference_mode() +def _evaluate_valid_horizons( + head: TurnPredictionHead, + states: torch.Tensor, + targets: torch.Tensor, + device: torch.device, + batch_size: int, +) -> dict[str, Any]: + head.eval() + chunks: list[torch.Tensor] = [] + for start in range(0, len(states), batch_size): + x = states[start : start + batch_size].to( + device=device, + dtype=torch.float32, + non_blocking=True, + ) + chunks.append(head(x).softmax(dim=-1).cpu()) + probabilities = torch.cat(chunks, dim=0).numpy() + target_array = targets.numpy() + + results = [] + for horizon in range(NUM_HORIZONS): + result = average_precision_metrics( + probabilities[:, horizon], + target_array[:, horizon], + ) + result["horizon"] = horizon + 1 + result["seconds"] = (horizon + 1) * STEP_SECONDS + results.append(result) + return { + "horizons": results, + "mean_map": float(np.mean([item["map"] for item in results])), + } + + +def _checkpoint( + head: TurnPredictionHead, + *, + epoch: int, + best_epoch: int, + best_dev_map: float | None, + train_loss: float, + args: argparse.Namespace, + train_sessions: list[str], + dev_sessions: list[str], + hidden_size: int, + feature_metadata: list[dict[str, Any]], + history: list[dict[str, Any]], +) -> dict[str, Any]: + return { + "format_version": 1, + "head": head.state_dict(), + "head_config": { + "hidden_size": hidden_size, + "num_horizons": NUM_HORIZONS, + "num_classes": NUM_CLASSES, + "dropout": args.dropout, + }, + "alignment": { + "state_label_offset": 0, + "future_offsets": list(range(1, NUM_HORIZONS + 1)), + "future_seconds": [ + (index + 1) * STEP_SECONDS for index in range(NUM_HORIZONS) + ], + }, + "classes": list(CLASS_NAMES), + "epoch": epoch, + "best_epoch": best_epoch, + "best_dev_map": best_dev_map, + "train_loss": train_loss, + "train_sessions": train_sessions, + "dev_sessions": dev_sessions, + "feature_metadata": feature_metadata, + "training_config": { + "epochs": args.epochs, + "batch_size": args.batch_size, + "learning_rate": args.learning_rate, + "weight_decay": args.weight_decay, + "warmup_epochs": args.warmup_epochs, + "warmup_factor": args.warmup_factor, + "loss": args.loss, + "optimizer": "Adam", + "scheduler": "linear_warmup_times_cosine", + "seed": args.seed, + "expected_mode": args.expected_mode, + }, + "history": history, + } + + +def main() -> None: + args = parse_args() + _validate_args(args) + _seed_everything(args.seed) + device = _resolve_device(args.device) + expected_mode = None if args.expected_mode == "any" else args.expected_mode + + requested_train = [normalize_session_name(item) for item in args.sessions] + dev_sessions = [normalize_session_name(item) for item in args.dev_sessions] + train_sessions = [item for item in requested_train if item not in dev_sessions] + validate_disjoint_sessions(train_sessions, dev_sessions) + if not train_sessions: + raise ValueError("No training sessions remain after removing dev sessions") + forbidden_test = {"Session_1", "Session_2", "Session_3"} + leaked = forbidden_test & (set(train_sessions) | set(dev_sessions)) + if leaked: + raise ValueError( + f"Official test sessions cannot be used for training/selection: " + f"{sorted(leaked)}" + ) + + print(f"device: {device}") + print(f"train sessions: {', '.join(train_sessions)}") + print(f"dev sessions: {', '.join(dev_sessions) if dev_sessions else 'none'}") + states, targets, hidden_size, feature_metadata = _load_examples( + args.feature_root, + args.target_root, + train_sessions, + expected_mode, + ) + print(f"training examples: {len(states)}, feature_dim={hidden_size}") + dev_states: torch.Tensor | None = None + dev_targets: torch.Tensor | None = None + if dev_sessions: + dev_states, dev_targets, dev_hidden_size, _ = _load_examples( + args.feature_root, + args.target_root, + dev_sessions, + expected_mode, + ) + if dev_hidden_size != hidden_size: + raise ValueError( + f"Dev feature dimension {dev_hidden_size} differs from " + f"training dimension {hidden_size}" + ) + print(f"development examples: {len(dev_states)}") + + dataset = TensorDataset(states, targets) + generator = torch.Generator().manual_seed(args.seed) + loader = DataLoader( + dataset, + batch_size=args.batch_size, + shuffle=True, + num_workers=args.num_workers, + pin_memory=device.type == "cuda", + persistent_workers=args.num_workers > 0, + generator=generator, + ) + head = TurnPredictionHead( + hidden_size=hidden_size, + num_horizons=NUM_HORIZONS, + num_classes=NUM_CLASSES, + dropout=args.dropout, + ).to(device) + optimizer = torch.optim.Adam( + head.parameters(), + lr=args.learning_rate, + weight_decay=args.weight_decay, + ) + scheduler = _scheduler( + optimizer, + steps_per_epoch=len(loader), + epochs=args.epochs, + warmup_epochs=args.warmup_epochs, + warmup_factor=args.warmup_factor, + ) + loss_function = ( + egospeak_multihorizon_loss + if args.loss == "lstr-mce" + else miniroad_multihorizon_loss + ) + + args.output.parent.mkdir(parents=True, exist_ok=True) + history: list[dict[str, Any]] = [] + best_dev_map = -math.inf if dev_sessions else None + best_epoch = args.epochs if not dev_sessions else 0 + for epoch in range(1, args.epochs + 1): + head.train() + loss_sum = 0.0 + example_count = 0 + progress = tqdm( + loader, + desc=f"head epoch {epoch}/{args.epochs}", + unit="batch", + disable=args.no_progress, + ) + for state_batch, target_batch in progress: + state_batch = state_batch.to( + device=device, + dtype=torch.float32, + non_blocking=True, + ) + target_batch = target_batch.to( + device=device, + dtype=torch.float32, + non_blocking=True, + ) + optimizer.zero_grad(set_to_none=True) + logits = head(state_batch) + loss = loss_function(logits, target_batch) + loss.backward() + optimizer.step() + scheduler.step() + batch_count = len(state_batch) + loss_sum += float(loss.detach()) * batch_count + example_count += batch_count + progress.set_postfix( + loss=f"{float(loss.detach()):.5f}", + lr=f"{scheduler.get_last_lr()[0]:.3g}", + ) + train_loss = loss_sum / example_count + record: dict[str, Any] = { + "epoch": epoch, + "train_loss": train_loss, + "learning_rate": scheduler.get_last_lr()[0], + } + + if dev_sessions: + assert dev_states is not None and dev_targets is not None + dev_metrics = _evaluate_valid_horizons( + head, + dev_states, + dev_targets, + device, + args.batch_size, + ) + dev_map = float(dev_metrics["mean_map"]) + record["dev_valid_mean_map"] = dev_map + print( + f"epoch {epoch:02d}: train_loss={train_loss:.6f}, " + f"dev_valid_mean_mAP={dev_map * 100:.3f}" + ) + if dev_map > float(best_dev_map): + best_dev_map = dev_map + best_epoch = epoch + torch.save( + _checkpoint( + head, + epoch=epoch, + best_epoch=best_epoch, + best_dev_map=best_dev_map, + train_loss=train_loss, + args=args, + train_sessions=train_sessions, + dev_sessions=dev_sessions, + hidden_size=hidden_size, + feature_metadata=feature_metadata, + history=history + [record], + ), + args.output, + ) + print(f"saved new best checkpoint: {args.output}") + else: + print(f"epoch {epoch:02d}: train_loss={train_loss:.6f}") + history.append(record) + + if not dev_sessions: + torch.save( + _checkpoint( + head, + epoch=args.epochs, + best_epoch=best_epoch, + best_dev_map=None, + train_loss=history[-1]["train_loss"], + args=args, + train_sessions=train_sessions, + dev_sessions=dev_sessions, + hidden_size=hidden_size, + feature_metadata=feature_metadata, + history=history, + ), + args.output, + ) + print(f"checkpoint: {args.output}") + if dev_sessions: + print( + f"best epoch: {best_epoch}, " + f"dev valid-only mean mAP: {float(best_dev_map) * 100:.3f}" + ) + summary_path = args.output.with_suffix(args.output.suffix + ".json") + payload = torch.load(args.output, map_location="cpu", weights_only=False) + summary = {key: value for key, value in payload.items() if key != "head"} + summary_path.write_text(json.dumps(summary, indent=2) + "\n") + print(f"training summary: {summary_path}") + + +if __name__ == "__main__": + main() From de29e7573405df7d78b818ed49d7d65104b8e316 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Mon, 17 Aug 2026 15:17:11 +0800 Subject: [PATCH 2/3] Add scripts for preparing and extracting 5 FPS audio-video clips, including `prepare_ego4d_2fps.py`, `run_extract_5fps_2pergpu.sh`, and `run_prepare_all_5fps.sh`. Introduce `streaming_turn_model_directly.py` for direct extraction from original media. Update `README.md` to reflect changes in Qwen's streaming design and causal extraction modes. Add new `AGENTS.md` file for documentation. --- AGENTS.md | 2 + README.md | 216 +++++++- prepare_ego4d_2fps.py | 421 ++++++++++++++ run_extract_5fps_2pergpu.sh | 94 ++++ run_prepare_all_5fps.sh | 63 +++ streaming_turn_model.py | 42 +- streaming_turn_model_directly.py | 913 +++++++++++++++++++++++++++++++ 7 files changed, 1722 insertions(+), 29 deletions(-) create mode 100644 AGENTS.md create mode 100755 prepare_ego4d_2fps.py create mode 100755 run_extract_5fps_2pergpu.sh create mode 100755 run_prepare_all_5fps.sh create mode 100644 streaming_turn_model_directly.py diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..2221f8b --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,2 @@ +作为一名资深音视频，深度学习国际知名导师。 +对于我的想法，引导我以思辨的方式继续，指导我完成icassp2027论文实验和编写。 \ No newline at end of file diff --git a/README.md b/README.md index 167f585..abf2715 100644 --- a/README.md +++ b/README.md @@ -317,33 +317,112 @@ script and reuse the existing MP4 without encoding it again: python prepare_easycom_session.py ... --fps 2 --reuse-existing ``` -## What Qwen's native 2-second chunk means +## What Qwen's native 2-second streaming design 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: +Qwen2.5-Omni was designed for **block-wise multimodal prefill**. Its +[technical report](https://arxiv.org/html/2503.20215) states that the audio +and visual encoders are modified for block-wise processing and that audio +self-attention is performed in independent two-second blocks. The released +[3B configuration](https://huggingface.co/Qwen/Qwen2.5-Omni-3B/raw/main/config.json) +uses `seconds_per_chunk=2`, `position_id_per_seconds=25`, and audio +`n_window=100`. -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: +Conceptually, a native online input is organized as follows: ```text -encode the complete [0, 2.0 s] block --> read a purported state at 0.2 s +open one user/media turn + encode/append AV [0,2] -> commit Thinker KV + encode/append AV [2,4] -> commit Thinker KV + encode/append AV [4,6] -> commit Thinker KV + ... +close the media turn -> generate text and/or speech ``` -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`. +Within each temporal group, video tokens precede the aligned audio tokens. +TMRoPE assigns 25 temporal positions per second, or 50 positions per native +two-second group, so audio and video share a physical-time coordinate system. +Completed blocks are appended to one causal Thinker sequence; the long-range +history is modeled by the Thinker rather than by cross-block full attention in +the modality encoders. + +This architecture is block-causal, not frame-causal. Attention inside an +audio block is non-causal. It is therefore invalid to encode the complete +`[0,2]` block and claim that its representation at 0.2 seconds saw no future: + +```text +encode complete [0,2] AV -> read state at 0.2 s # future leakage +``` + +For this project's 5 Hz targets, the current incomplete block must instead be +encoded only through the prediction time. These 0.2-second partial-block +readouts are an extension made for EgoSpeak; the Qwen paper does not define a +canonical hidden state every 200 ms. + +The public Transformers and local Qwen web-demo APIs normally receive a +completed media object and call the processor/model once. They retain the +model's block-wise masks and token organization, but do not expose the full +persistent live-session implementation used by the hosted real-time product. + +## Native Qwen streaming versus `causal-prefix` + +The two extraction modes are both causal with respect to raw observations, +but they define different model inputs. `causal-prefix` is a sliding-window +reference invented for this experiment; it is not Qwen's native streaming +execution model. + +| Property | `causal-prefix` | Native Qwen streaming idea | +|---|---|---| +| Session | New inference at every 0.2 s | One open media turn | +| Media organization | One complete `[max(0,t-M),t]` item | Fixed source-aligned 2 s blocks | +| Modality encoding | Re-encode the complete window | Encode/commit each completed block once | +| Thinker KV | Start empty for every state | Append to persistent KV | +| TMRoPE | Recomputed from each window origin | Advances with the stream | +| History | Exact rolling `M` seconds | Accumulates until a serving/context policy intervenes | +| 5 Hz readout | Close/read a new turn each time | Not specified by the original model | + +The difference becomes especially important after a sliding window fills. At +`t=20.8` with `M=20`, `causal-prefix` constructs one new media item covering: + +```text +[0.8,20.8] + -> encoder blocks [0.8,2.8], [2.8,4.8], ..., [18.8,20.8] +``` + +At `t=21.0`, every encoder boundary moves again: + +```text +[1.0,21.0] + -> encoder blocks [1.0,3.0], [3.0,5.0], ..., [19.0,21.0] +``` + +Native streaming keeps the source-aligned partition instead: + +```text +[0,2], [2,4], ..., [18,20], current partial [20,20.8] +``` + +Because Qwen's audio encoder uses block-local non-causal attention, moving a +block boundary changes the representation even when most raw samples are the +same. Window-relative TMRoPE, convolution/padding boundaries, chat closure, +and BF16/attention execution add further differences. This is why exact KV +replay can be numerically identical to its own from-scratch replay while still +differing from `causal-prefix`; the modes are evaluating different token and +encoder histories. + +Use the modes according to the scientific question: + +- use `open-stream-replay` for the primary online experiment, training, and + deployment-aligned evaluation; +- use `causal-prefix` as an expensive continuous-window reference and + short-clip ablation; +- never train a head on one representation and deploy it on the other without + retraining, because their feature distributions are not interchangeable. ## Causal 5 Hz extraction modes `streaming_turn_model.py` implements two modes. -### 1. `causal-prefix`: correctness and semantic baseline +### 1. `causal-prefix`: continuous-window reference At every 200 ms, independently process the causal context ending at the current time: @@ -358,7 +437,8 @@ Properties: - one continuous media window per prediction; - no cross-forward cache; - very expensive because overlapping context is recomputed; -- preferred reference when checking semantics and causality. +- useful for checking continuous sliding-window semantics and causality; +- not the native Qwen block-streaming execution model. Example: @@ -390,6 +470,18 @@ open AV prefix AV/chat closure -> never committed; readout branch only ``` +This matches the core Qwen design choices: fixed two-second AV groups, one +open media turn, completed-block commits, and incremental Thinker KV. It makes +two practical changes for this dataset: + +1. the current 0--2 second partial block is evaluated speculatively every + 0.2 seconds and discarded unless it reaches a complete block boundary; +2. the KV history is bounded and rebuilt from recent complete blocks rather + than accumulating for an entire multi-minute session. + +It should therefore be described as **bounded Qwen-style open streaming**, not +as bit-for-bit reproduction of an unbounded official serving session. + 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. @@ -420,15 +512,26 @@ 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 +maximum video duration and not a limit specified by Qwen. The original +append-only idea eventually reaches both linear KV-memory growth and the +Thinker's 32,768-token context limit, so an unbounded multi-minute stream is +not practical. The complete session is processed sequentially and a prediction is emitted every 0.2 seconds. -`open-stream-replay` instead uses blockwise rolling memory. With a 20-second +`open-stream-replay` 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 a short native-stream fidelity check, set `--memory-seconds` to at least +the tested clip duration. For example, the options +`--start 0 --end 40 --memory-seconds 40` prevent a rolling rebuild during that +40-second clip. +Do not extend this policy to complete EasyCom/Ego4D sessions: KV memory grows +linearly, and this implementation clones the cache for speculative partial +blocks and readout branches. + For the paper experiment, drive the output length from the preparation sidecar rather than a floating MP4 duration: @@ -488,6 +591,68 @@ 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. +## Direct extraction without a prepared 2 FPS MP4 + +`streaming_turn_model_directly.py` removes the offline video-transcoding +requirement. It reads original frame PTS with Decord, constructs a global +2 FPS lattice, maps every lattice point to the latest source frame whose PTS +is not later than that point, and passes explicit frame lists to Qwen. Audio is +decoded once to 16-kHz mono and sliced at exact causal boundaries. Therefore +no FFmpeg `fps` filter or second H.264/AAC encode can alter the experiment. + +Validate one Ego4D clip without loading Qwen: + +```bash +python streaming_turn_model_directly.py \ + --video /data/wangx/Datasets/ego4d/v1/clips/CLIP_UID.mp4 \ + --target /data/wangx/Datasets/ego4d/ego4d_target_frame/CLIP_UID.npy \ + --video-fps 2 \ + --media-only +``` + +Run a short native persistent-KV check: + +```bash +python streaming_turn_model_directly.py \ + --model-id /data/wangx/Models/Qwen2.5-Omni-3B \ + --video /data/wangx/Datasets/ego4d/v1/clips/CLIP_UID.mp4 \ + --end 40 \ + --mode open-stream-replay \ + --memory-seconds 0 \ + --video-fps 2 \ + --flash-attention-2 \ + --device cuda:0 \ + --output features/direct/CLIP_UID_first_40s.pt +``` + +Here `--memory-seconds 0` means that completed-block Thinker KV persists for +the entire requested interval. Do not use this setting blindly for complete +multi-minute recordings: it has linear KV growth and remains subject to the +Thinker's context limit. A full bounded-memory dataset run should explicitly +use the experimental budget, for example `--memory-seconds 20`. + +EasyCom can be read directly as the dense concatenation of retained clips; +filename timecode gaps are ignored and the rightmost 1920x1080 region is kept: + +```bash +python streaming_turn_model_directly.py \ + --model-id /path/to/Qwen2.5-Omni-3B \ + --easycom-session-dir /path/to/Main/Video_Compressed/Session_1 \ + --target target_frame/Session_1.npy \ + --mode open-stream-replay \ + --memory-seconds 20 \ + --video-fps 2 \ + --flash-attention-2 \ + --device cuda:0 \ + --output features/direct/Session_1.pt +``` + +Matching `states.shape[0]` to the target is necessary but not sufficient. +The direct reader also validates EasyCom's 20-Hz source inventory, source PTS +monotonicity, dense duration, crop bounds, audio sample count, and causal +frame-to-grid mapping. The original prepare scripts remain useful only as an +optional decode-speed/reproducibility cache. + ## Saved streaming tensors Each `.pt` file contains: @@ -540,9 +705,9 @@ Train once on the official nine sessions with a fixed 50-epoch protocol: ```bash python train_turn_head.py \ - --feature-root features/causal_prefix_2fps \ + --feature-root features/open_2fps \ --target-root target_frame \ - --output checkpoints/causal_prefix_2fps_head.pth \ + --output checkpoints/open_2fps_head.pth \ --device cuda:0 ``` @@ -577,10 +742,10 @@ Run inference on the official test sessions (`Session_3`, `Session_2`, and ```bash python infer_turn_head.py \ - --feature-root features/causal_prefix_2fps \ + --feature-root features/open_2fps \ --target-root target_frame \ - --checkpoint checkpoints/causal_prefix_2fps_head.pth \ - --output-dir results/causal_prefix_2fps_head \ + --checkpoint checkpoints/open_2fps_head.pth \ + --output-dir results/open_2fps_head \ --device cuda:0 ``` @@ -743,7 +908,8 @@ Once these checks pass, the first publishable experiment is: ```text frozen Qwen2.5-Omni-3B Thinker + trained 10x3 head -+ 20-second causal memory ++ bounded Qwen-style 2-second open streaming ++ 20-second rolling causal memory + full-session 5 Hz output + exact EgoSpeak split and Table 2 metric ``` diff --git a/prepare_ego4d_2fps.py b/prepare_ego4d_2fps.py new file mode 100755 index 0000000..875b576 --- /dev/null +++ b/prepare_ego4d_2fps.py @@ -0,0 +1,421 @@ +#!/usr/bin/env python3 +"""Prepare EgoSpeak's Ego4D split as exact 2 FPS audio-video clips. + +The 5 Hz target arrays are the timing authority. For a target with T rows, +the prepared clip has duration T / 5 seconds and exactly ceil(2T / 5) video +frames. Short source audio is padded with silence and all audio is trimmed to +the target duration. + +Processing is intentionally sequential. A single libx264 process uses all +available CPU cores on the intended extraction machine, so launching several +encoders normally creates CPU and disk contention instead of improving total +throughput. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +import shutil +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + + +TARGET_FPS = 5 +VIDEO_FPS = 2 +EXPECTED_TRAIN_CLIPS = 346 +EXPECTED_TEST_CLIPS = 87 + + +@dataclass(frozen=True) +class ClipSpec: + split: str + uid: str + steps: int + duration: float + video_frames: int + source: Path + target: Path + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Prepare the 433 clips in EgoSpeak's EGO4DSHUFFLE split as " + "strict fixed-grid 2 FPS MP4 files." + ) + ) + parser.add_argument( + "--data-root", + type=Path, + default=Path("/data/wangx/Datasets/ego4d"), + help="Ego4D dataset root (default: %(default)s)", + ) + parser.add_argument( + "--split-json", + type=Path, + help="EgoSpeak LSTR data_info.json; defaults to DATA_ROOT/egospeak_data_info.json", + ) + parser.add_argument( + "--video-root", + type=Path, + help="Raw clip directory; defaults to DATA_ROOT/v1/clips", + ) + parser.add_argument( + "--target-root", + type=Path, + help="Target directory; defaults to DATA_ROOT/ego4d_target_frame", + ) + parser.add_argument( + "--output-root", + type=Path, + help="Prepared video directory; defaults to DATA_ROOT/omni_prepared_2fps", + ) + parser.add_argument( + "--manifest", + type=Path, + help="Output TSV; defaults to DATA_ROOT/ego4d_egospeak.tsv", + ) + parser.add_argument( + "--split", + choices=("all", "train", "test"), + default="all", + help="Prepare all, train, or test clips (default: %(default)s)", + ) + parser.add_argument( + "--uid", + action="append", + help="Prepare only this split UID; repeat to select several clips", + ) + parser.add_argument( + "--limit", + type=int, + help="Prepare only the first N selected clips, for testing", + ) + parser.add_argument( + "--crf", + type=int, + default=18, + help="libx264 constant-rate-factor value (default: %(default)s)", + ) + parser.add_argument( + "--preset", + default="medium", + help="libx264 preset (default: %(default)s)", + ) + parser.add_argument( + "--force", + action="store_true", + help="Rebuild outputs even when their frame counts are already correct", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Validate inputs and print work without running ffmpeg", + ) + args = parser.parse_args() + + if args.limit is not None and args.limit <= 0: + parser.error("--limit must be positive") + if not 0 <= args.crf <= 51: + parser.error("--crf must be between 0 and 51") + return args + + +def resolve_paths(args: argparse.Namespace) -> None: + args.data_root = args.data_root.expanduser().resolve() + args.split_json = ( + args.split_json or args.data_root / "egospeak_data_info.json" + ).expanduser().resolve() + args.video_root = ( + args.video_root or args.data_root / "v1" / "clips" + ).expanduser().resolve() + args.target_root = ( + args.target_root or args.data_root / "ego4d_target_frame" + ).expanduser().resolve() + args.output_root = ( + args.output_root or args.data_root / "omni_prepared_2fps" + ).expanduser().resolve() + args.manifest = ( + args.manifest or args.data_root / "ego4d_egospeak.tsv" + ).expanduser().resolve() + + +def require_file(path: Path, description: str) -> None: + if not path.is_file(): + raise FileNotFoundError(f"Missing {description}: {path}") + + +def load_specs(args: argparse.Namespace) -> list[ClipSpec]: + require_file(args.split_json, "EgoSpeak split JSON") + if not args.video_root.is_dir(): + raise FileNotFoundError(f"Missing raw-video directory: {args.video_root}") + if not args.target_root.is_dir(): + raise FileNotFoundError(f"Missing target directory: {args.target_root}") + + with args.split_json.open() as file: + payload = json.load(file) + try: + split_info = payload["EGO4DSHUFFLE"] + train_uids = split_info["train_session_set"] + test_uids = split_info["test_session_set"] + except KeyError as error: + raise RuntimeError( + f"{args.split_json} is not EgoSpeak's LSTR data_info.json: " + f"missing {error}" + ) from error + + if len(train_uids) != EXPECTED_TRAIN_CLIPS: + raise RuntimeError( + f"Expected {EXPECTED_TRAIN_CLIPS} EGO4DSHUFFLE train clips, " + f"found {len(train_uids)}" + ) + if len(test_uids) != EXPECTED_TEST_CLIPS: + raise RuntimeError( + f"Expected {EXPECTED_TEST_CLIPS} EGO4DSHUFFLE test clips, " + f"found {len(test_uids)}" + ) + + specs: list[ClipSpec] = [] + seen: set[str] = set() + for split, uids in (("train", train_uids), ("test", test_uids)): + for uid in uids: + if uid in seen: + raise RuntimeError(f"Duplicate UID in EGO4DSHUFFLE: {uid}") + seen.add(uid) + + source = args.video_root / f"{uid}.mp4" + target_path = args.target_root / f"{uid}.npy" + require_file(source, "source video") + require_file(target_path, "target array") + + target = np.load(target_path, mmap_mode="r", allow_pickle=False) + if target.ndim != 2 or target.shape[1] != 3: + raise RuntimeError( + f"Expected target shape [T, 3], found {target.shape}: " + f"{target_path}" + ) + + steps = int(target.shape[0]) + duration = steps / TARGET_FPS + video_frames = math.ceil(steps * VIDEO_FPS / TARGET_FPS) + specs.append( + ClipSpec( + split=split, + uid=uid, + steps=steps, + duration=duration, + video_frames=video_frames, + source=source, + target=target_path, + ) + ) + return specs + + +def write_manifest(specs: list[ClipSpec], path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".partial") + with temporary.open("w", newline="") as file: + writer = csv.writer(file, delimiter="\t") + writer.writerow( + ["split", "uid", "steps", "duration", "video_frames_2fps"] + ) + for spec in specs: + writer.writerow( + [ + spec.split, + spec.uid, + spec.steps, + f"{spec.duration:.6f}", + spec.video_frames, + ] + ) + temporary.replace(path) + + +def select_specs(args: argparse.Namespace, specs: list[ClipSpec]) -> list[ClipSpec]: + selected = specs + if args.split != "all": + selected = [spec for spec in selected if spec.split == args.split] + if args.uid: + requested = set(args.uid) + known = {spec.uid for spec in specs} + missing = sorted(requested - known) + if missing: + raise RuntimeError( + "Requested UID(s) are not in EGO4DSHUFFLE: " + ", ".join(missing) + ) + selected = [spec for spec in selected if spec.uid in requested] + if args.limit is not None: + selected = selected[: args.limit] + if not selected: + raise RuntimeError("No clips selected") + return selected + + +def count_video_frames(path: Path) -> int | None: + if not path.is_file() or path.stat().st_size == 0: + return None + command = [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-count_frames", + "-show_entries", + "stream=nb_read_frames", + "-of", + "default=noprint_wrappers=1:nokey=1", + str(path), + ] + completed = subprocess.run(command, text=True, capture_output=True) + if completed.returncode != 0: + return None + try: + return int(completed.stdout.strip()) + except ValueError: + return None + + +def prepare_one( + args: argparse.Namespace, + spec: ClipSpec, + index: int, + total: int, +) -> str: + output = args.output_root / f"{spec.uid}_2fps.mp4" + actual = count_video_frames(output) + prefix = f"[{index}/{total}] [{spec.split}] {spec.uid}" + + if not args.force and actual == spec.video_frames: + print(f"{prefix}: reuse ({actual} frames)", flush=True) + return "reused" + + if actual is None: + reason = "missing" + else: + reason = f"actual={actual}, expected={spec.video_frames}" + print( + f"{prefix}: prepare ({reason}; {spec.steps} states, " + f"{spec.duration:.1f}s, {spec.video_frames} frames)", + flush=True, + ) + if args.dry_run: + return "dry-run" + + temporary = args.output_root / f".{spec.uid}.partial.mp4" + duration = f"{spec.duration:.6f}" + video_filter = ( + f"[0:v:0]trim=start=0:end={duration},setpts=PTS-STARTPTS," + f"fps=fps={VIDEO_FPS}:start_time=0:round=down," + "tpad=stop_mode=clone:stop=2," + f"trim=start_frame=0:end_frame={spec.video_frames}," + f"setpts=N/({VIDEO_FPS}*TB),format=yuv420p[vout]" + ) + audio_filter = ( + f"[0:a:0]apad,atrim=start=0:end={duration}," + "asetpts=PTS-STARTPTS[aout]" + ) + command = [ + "ffmpeg", + "-hide_banner", + "-loglevel", + "warning", + "-stats", + "-y", + "-i", + str(spec.source), + "-filter_complex", + f"{video_filter};{audio_filter}", + "-map", + "[vout]", + "-map", + "[aout]", + "-map_metadata", + "-1", + "-c:v", + "libx264", + "-preset", + args.preset, + "-crf", + str(args.crf), + "-c:a", + "aac", + "-ar", + "16000", + "-ac", + "1", + "-b:a", + "64k", + "-movflags", + "+faststart", + str(temporary), + ] + completed = subprocess.run(command) + if completed.returncode != 0: + raise RuntimeError( + f"ffmpeg failed for {spec.uid} with exit code {completed.returncode}; " + f"partial output was left at {temporary}" + ) + + actual = count_video_frames(temporary) + if actual != spec.video_frames: + raise RuntimeError( + f"Prepared frame count is {actual}, expected {spec.video_frames}: " + f"{temporary}" + ) + temporary.replace(output) + print(f"{prefix}: OK ({actual} frames) -> {output}", flush=True) + return "prepared" + + +def main() -> None: + args = parse_args() + resolve_paths(args) + for executable in ("ffmpeg", "ffprobe"): + if shutil.which(executable) is None: + raise RuntimeError(f"Required executable is not on PATH: {executable}") + + specs = load_specs(args) + write_manifest(specs, args.manifest) + selected = select_specs(args, specs) + args.output_root.mkdir(parents=True, exist_ok=True) + + total_steps = sum(spec.steps for spec in specs) + print( + f"Validated EGO4DSHUFFLE: {len(specs)} clips " + f"({EXPECTED_TRAIN_CLIPS} train, {EXPECTED_TEST_CLIPS} test), " + f"{total_steps} target states" + ) + print(f"Manifest: {args.manifest}") + print(f"Output: {args.output_root}") + print(f"Selected: {len(selected)} clips; sequential ffmpeg execution") + + counts = {"prepared": 0, "reused": 0, "dry-run": 0} + for index, spec in enumerate(selected, start=1): + result = prepare_one(args, spec, index, len(selected)) + counts[result] += 1 + + print( + "Complete: " + + ", ".join(f"{key}={value}" for key, value in counts.items() if value) + ) + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + print("\nInterrupted; rerun the same command to resume.", file=sys.stderr) + raise SystemExit(130) + except Exception as error: + print(f"ERROR: {error}", file=sys.stderr) + raise SystemExit(1) diff --git a/run_extract_5fps_2pergpu.sh b/run_extract_5fps_2pergpu.sh new file mode 100755 index 0000000..03be97e --- /dev/null +++ b/run_extract_5fps_2pergpu.sh @@ -0,0 +1,94 @@ +#!/usr/bin/env bash +set -euo pipefail + +MODEL="${MODEL:-/data1/yfzx/wangx/Models/Qwen2.5-Omni-3B}" +PROJECT_ROOT="${PROJECT_ROOT:-/data1/yfzx/wangx/Code/icassp2027-experiment}" +TARGET_ROOT="${TARGET_ROOT:-${PROJECT_ROOT}/target_frame}" +PREP_ROOT="${PREP_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_prepared}" +FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/causal_prefix_5fps}" +LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/causal_prefix_5fps}" + +mkdir -p "${FEATURE_ROOT}" "${LOG_ROOT}" +cd "${PROJECT_ROOT}" + +run_worker() { + local gpu="$1" + shift + + local sid session metadata target output log + for sid in "$@"; do + session="Session_${sid}" + metadata="${PREP_ROOT}/${session}/${session}.json" + target="${TARGET_ROOT}/${session}.npy" + output="${FEATURE_ROOT}/${session}.pt" + log="${LOG_ROOT}/${session}.log" + + if [[ ! -s "${metadata}" ]]; then + echo "Missing metadata: ${metadata}" >&2 + return 1 + fi + if [[ ! -s "${target}" ]]; then + echo "Missing target: ${target}" >&2 + return 1 + fi + if [[ -s "${output}" ]]; then + echo "Skipping existing output: ${output}" + continue + fi + + { + echo "$(date): GPU ${gpu} starting ${session}" + + CUDA_VISIBLE_DEVICES="${gpu}" PYTHONUNBUFFERED=1 \ + python streaming_turn_model.py \ + --model-id "${MODEL}" \ + --session-metadata "${metadata}" \ + --target "${target}" \ + --use-audio-in-video \ + --mode causal-prefix \ + --memory-seconds 20 \ + --video-fps 5 \ + --strict-fixed-video-grid \ + --flash-attention-2 \ + --output "${output}" + + echo "$(date): GPU ${gpu} completed ${session}" + } >"${log}" 2>&1 + done +} + +pids=() + +# Two independent extraction workers per GPU. Each GPU receives approximately +# the same total number of EasyCom timesteps. +run_worker 0 12 11 &2 + exit 1 +fi + +echo "ALL 12 CAUSAL 5 FPS EXTRACTIONS COMPLETED" diff --git a/run_prepare_all_5fps.sh b/run_prepare_all_5fps.sh new file mode 100755 index 0000000..cf18825 --- /dev/null +++ b/run_prepare_all_5fps.sh @@ -0,0 +1,63 @@ +#!/usr/bin/env bash +set -euo pipefail + +PROJECT_ROOT="${PROJECT_ROOT:-/data1/yfzx/wangx/Code/icassp2027-experiment}" +VIDEO_ROOT="${VIDEO_ROOT:-/data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed}" +TARGET_ROOT="${TARGET_ROOT:-${PROJECT_ROOT}/target_frame}" +PREP_ROOT="${PREP_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_prepared}" +LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/prepare_5fps}" + +mkdir -p "${LOG_ROOT}" +cd "${PROJECT_ROOT}" + +prepare_worker() { + local sid session session_dir target output_dir log + for sid in "$@"; do + session="Session_${sid}" + session_dir="${VIDEO_ROOT}/${session}" + target="${TARGET_ROOT}/${session}.npy" + output_dir="${PREP_ROOT}/${session}" + log="${LOG_ROOT}/${session}.log" + + if [[ ! -d "${session_dir}" ]]; then + echo "Missing source session: ${session_dir}" >&2 + return 1 + fi + if [[ ! -s "${target}" ]]; then + echo "Missing target: ${target}" >&2 + return 1 + fi + + { + echo "$(date): preparing ${session}" + + python prepare_easycom_session.py \ + --session-dir "${session_dir}" \ + --output-dir "${output_dir}" \ + --target "${target}" \ + --fps 2 5 \ + --reuse-existing + + echo "$(date): completed ${session}" + } >"${log}" 2>&1 + done +} + +# FFmpeg/libx264 is already multithreaded. Two session workers generally make +# better use of a many-core host without overwhelming storage or the CPU. +prepare_worker 1 3 5 7 9 11 &2 + exit 1 +fi + +echo "ALL 12 SESSIONS PREPARED AT 2 FPS AND 5 FPS" diff --git a/streaming_turn_model.py b/streaming_turn_model.py index befea2b..2d3f5ae 100644 --- a/streaming_turn_model.py +++ b/streaming_turn_model.py @@ -190,10 +190,10 @@ class FixedGridAV: def __init__(self, path: str | Path, expected_fps: float) -> None: try: import decord - import librosa + import numpy as np except ImportError as exc: raise RuntimeError( - "Fixed-grid AV extraction requires decord and librosa" + "Fixed-grid AV extraction requires decord and numpy" ) from exc self.path = Path(path).resolve() @@ -204,9 +204,43 @@ class FixedGridAV: 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 + # SoundFile cannot decode the AAC track in MP4, causing librosa to + # fall back to its deprecated and much slower audioread path. Decode + # directly to the exact float32/16-kHz/mono representation consumed by + # Qwen. Keeping the complete track in memory also makes repeated + # 200-ms slicing deterministic and inexpensive. + completed = subprocess.run( + [ + "ffmpeg", + "-nostdin", + "-v", + "error", + "-i", + str(self.path), + "-map", + "0:a:0", + "-vn", + "-ac", + "1", + "-ar", + str(self.sample_rate), + "-c:a", + "pcm_f32le", + "-f", + "f32le", + "pipe:1", + ], + check=False, + capture_output=True, ) + if completed.returncode != 0: + message = completed.stderr.decode("utf-8", errors="replace").strip() + raise RuntimeError( + f"Failed to decode embedded audio from {self.path}: {message}" + ) + self.audio = np.frombuffer(completed.stdout, dtype=" list[Any]: """Return fixed-grid frames whose timestamps lie in [start, end).""" diff --git a/streaming_turn_model_directly.py b/streaming_turn_model_directly.py new file mode 100644 index 0000000..54e6351 --- /dev/null +++ b/streaming_turn_model_directly.py @@ -0,0 +1,913 @@ +#!/usr/bin/env python3 +"""Extract causal Qwen2.5-Omni states directly from original media. + +Unlike ``prepare_easycom_session.py`` and ``prepare_ego4d_2fps.py``, this +entry point does not create a re-encoded fixed-frame-rate MP4. It builds a +global visual sampling lattice in Python, selects the newest decoded source +frame whose presentation timestamp is not later than each lattice time, and +passes explicit PIL frame lists plus exact 16-kHz audio slices to Qwen. + +The default mode follows the released Qwen2.5-Omni streaming design as closely +as the public Transformers implementation permits: + +* 2 FPS causal visual sampling (configurable for ablations); +* two-frame temporal patches, with the last observed frame duplicated when an + incomplete input contains an odd number of frames; +* a single open media turn with completed 2-second AV payload commits; +* persistent Thinker KV for the complete extraction (``--memory-seconds 0``); +* one temporary partial-block branch and readout every 200 ms. + +Examples +-------- +Ego4D, directly from one original clip:: + + python streaming_turn_model_directly.py \ + --model-id /data/wangx/Models/Qwen2.5-Omni-3B \ + --video /data/wangx/Datasets/ego4d/v1/clips/UID.mp4 \ + --target /data/wangx/Datasets/ego4d/ego4d_target_frame/UID.npy \ + --mode open-stream-replay \ + --memory-seconds 20 \ + --video-fps 2 \ + --flash-attention-2 \ + --device cuda:0 \ + --output features/ego4d_direct/UID.pt + +EasyCom, directly from the retained one-minute source clips:: + + python streaming_turn_model_directly.py \ + --model-id /path/to/Qwen2.5-Omni-3B \ + --easycom-session-dir /path/to/Main/Video_Compressed/Session_1 \ + --target target_frame/Session_1.npy \ + --mode open-stream-replay \ + --memory-seconds 20 \ + --video-fps 2 \ + --flash-attention-2 \ + --device cuda:0 \ + --output features/easycom_direct/Session_1.pt + +``--media-only`` performs source/timestamp/audio validation without loading +Qwen. This is useful before launching a full extraction on a remote machine. +""" + +from __future__ import annotations + +import argparse +import json +import math +import shutil +import subprocess +from collections import OrderedDict +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +import streaming_turn_model as core +from easycom_manifest import parse_timecode + + +EPSILON = 1e-7 +SAMPLE_RATE = 16_000 +EASYCOM_WIDTH = 1920 +EASYCOM_HEIGHT = 1080 + + +@dataclass(frozen=True) +class SourceClipSpec: + """One source file in a dense, gap-free media timeline.""" + + path: Path + duration: float | None = None + crop: tuple[int, int, int, int] | None = None + crop_rightmost_width: int | None = None + audio_path: Path | None = None + + +@dataclass +class IndexedSourceClip: + """Decoded indexing information for one source clip.""" + + path: Path + audio_path: Path + dense_start: float + dense_end: float + duration: float + crop: tuple[int, int, int, int] | None + reader: Any + source_frame_starts: np.ndarray + source_frame_ends: np.ndarray + source_fps: float + width: int + height: int + + @property + def num_frames(self) -> int: + return int(len(self.source_frame_starts)) + + +def _grid_indices(start: float, end: float, fps: float) -> list[int]: + """Indices of global samples ``n/fps`` in the half-open interval.""" + + first = max(0, int(math.ceil(start * fps - EPSILON))) + last_exclusive = int(math.ceil(end * fps - EPSILON)) + if first >= last_exclusive: + # Qwen's visual tower needs at least one observed image before its + # temporal-patch duplication. Reuse the latest global tick strictly + # before ``end``; this is identical to FixedGridAV's short-tail rule. + return [max(0, last_exclusive - 1)] + return list(range(first, last_exclusive)) + + +def _timestamp_array(reader: Any) -> tuple[np.ndarray, np.ndarray]: + """Return monotonically increasing, zero-origin frame PTS intervals.""" + + # Decord releases differ here: some accept an ndarray, while older ones + # require a Python list. A list works with both APIs. + indices = list(range(len(reader))) + raw = reader.get_frame_timestamp(indices) + if hasattr(raw, "asnumpy"): + raw = raw.asnumpy() + timestamps = np.asarray(raw, dtype=np.float64) + if timestamps.ndim == 1: + starts = timestamps + nominal = 1.0 / float(reader.get_avg_fps()) + ends = np.concatenate((starts[1:], starts[-1:] + nominal)) + elif timestamps.ndim == 2 and timestamps.shape[1] >= 2: + starts = timestamps[:, 0] + ends = timestamps[:, 1] + else: + raise RuntimeError( + f"Unexpected decord frame-timestamp shape: {timestamps.shape}" + ) + if not len(starts): + raise RuntimeError("Video contains no decoded frames") + origin = float(starts[0]) + starts = starts - origin + ends = ends - origin + if np.any(np.diff(starts) < -1e-9): + raise RuntimeError("Decoded video frame PTS values are not monotonic") + return starts, ends + + +def _decode_audio_exact(path: Path, duration: float) -> np.ndarray: + """Decode, then pad/trim one clip to its dense video duration.""" + + completed = subprocess.run( + [ + "ffmpeg", + "-nostdin", + "-v", + "error", + "-i", + str(path), + "-map", + "0:a:0", + "-vn", + "-ac", + "1", + "-ar", + str(SAMPLE_RATE), + "-c:a", + "pcm_f32le", + "-f", + "f32le", + "pipe:1", + ], + check=False, + capture_output=True, + ) + if completed.returncode != 0: + message = completed.stderr.decode("utf-8", errors="replace").strip() + raise RuntimeError(f"Failed to decode audio from {path}: {message}") + + decoded = np.frombuffer(completed.stdout, dtype=" None: + if not specs: + raise ValueError("At least one source clip is required") + if output_fps <= 0: + raise ValueError("output_fps must be positive") + if required_duration <= 0: + raise ValueError("required_duration must be positive") + if frame_cache_size < 0: + raise ValueError("frame_cache_size must be non-negative") + if shutil.which("ffmpeg") is None: + raise RuntimeError("ffmpeg is required to decode source audio") + + try: + import decord + except ImportError as exc: + raise RuntimeError( + "Direct extraction requires decord; install " + "qwen-omni-utils[decord]" + ) from exc + + self.fps = float(output_fps) + self.required_duration = float(required_duration) + self.frame_cache_size = int(frame_cache_size) + self._frame_cache: OrderedDict[tuple[int, int], Any] = OrderedDict() + self._request_count = 0 + self._cache_hits = 0 + self._odd_frame_duplications = 0 + self._short_interval_fallbacks = 0 + self._tail_frame_reuses = 0 + + indexed: list[IndexedSourceClip] = [] + audio_parts: list[np.ndarray] = [] + dense_cursor = 0.0 + + for spec_index, spec in enumerate(specs): + path = spec.path.expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(f"Source video does not exist: {path}") + print( + f"direct-media: indexing [{spec_index + 1}/{len(specs)}] " + f"{path.name}", + flush=True, + ) + reader = decord.VideoReader(str(path)) + if not len(reader): + raise RuntimeError(f"No video frames decoded from {path}") + starts, ends = _timestamp_array(reader) + source_fps = float(reader.get_avg_fps()) + first_shape = tuple(int(value) for value in reader[0].shape) + if len(first_shape) != 3: + raise RuntimeError( + f"Unexpected decoded frame shape {first_shape}: {path}" + ) + height, width = first_shape[:2] + + inferred_duration = len(reader) / source_fps + duration = ( + float(spec.duration) + if spec.duration is not None + else float(inferred_duration) + ) + if duration <= 0: + raise RuntimeError(f"Non-positive source duration for {path}") + + crop = spec.crop + if crop is None and spec.crop_rightmost_width is not None: + keep_width = int(spec.crop_rightmost_width) + if width < keep_width: + raise RuntimeError( + f"Cannot keep rightmost {keep_width}px from {width}px: {path}" + ) + crop = (width - keep_width, 0, keep_width, height) + if crop is not None: + x, y, crop_width, crop_height = crop + if ( + x < 0 + or y < 0 + or crop_width <= 0 + or crop_height <= 0 + or x + crop_width > width + or y + crop_height > height + ): + raise RuntimeError( + f"Invalid crop {crop} for {width}x{height}: {path}" + ) + + audio_path = (spec.audio_path or path).expanduser().resolve() + if not audio_path.is_file(): + raise FileNotFoundError(f"Audio source does not exist: {audio_path}") + + clip = IndexedSourceClip( + path=path, + audio_path=audio_path, + dense_start=dense_cursor, + dense_end=dense_cursor + duration, + duration=duration, + crop=crop, + reader=reader, + source_frame_starts=starts, + source_frame_ends=ends, + source_fps=source_fps, + width=width, + height=height, + ) + indexed.append(clip) + print( + f"direct-media: decoding audio [{spec_index + 1}/{len(specs)}] " + f"{audio_path.name}", + flush=True, + ) + audio_parts.append(_decode_audio_exact(audio_path, duration)) + dense_cursor = clip.dense_end + + # No later source can be observed before the requested endpoint. + if dense_cursor >= required_duration - EPSILON: + break + + self.clips = indexed + self.dense_starts = np.asarray( + [clip.dense_start for clip in indexed], dtype=np.float64 + ) + self.dense_duration = float(dense_cursor) + if self.dense_duration + EPSILON < self.required_duration: + raise RuntimeError( + f"Dense source duration is {self.dense_duration:.6f}s, shorter " + f"than required {self.required_duration:.6f}s" + ) + self.audio = np.concatenate(audio_parts) + required_samples = int(round(required_duration * SAMPLE_RATE)) + if len(self.audio) < required_samples: + self.audio = np.pad(self.audio, (0, required_samples - len(self.audio))) + elif len(self.audio) > required_samples: + self.audio = self.audio[:required_samples] + + max_grid_index = int( + math.ceil(required_duration * self.fps - EPSILON) + ) - 1 + self.expected_grid_frames = max(0, max_grid_index + 1) + if max_grid_index >= 0: + self._source_frame_for_grid(max_grid_index, count_tail_reuse=False) + + def _clip_for_time(self, global_time: float) -> tuple[int, IndexedSourceClip, float]: + clip_index = int(np.searchsorted(self.dense_starts, global_time, side="right") - 1) + clip_index = max(0, min(clip_index, len(self.clips) - 1)) + clip = self.clips[clip_index] + local_time = global_time - clip.dense_start + if local_time < 0: + local_time = 0.0 + return clip_index, clip, local_time + + def _source_frame_for_grid( + self, + grid_index: int, + *, + count_tail_reuse: bool = True, + ) -> tuple[int, int, float, float]: + """Map one global output tick to a source frame at or before it.""" + + global_time = grid_index / self.fps + clip_index, clip, local_time = self._clip_for_time(global_time) + frame_index = int( + np.searchsorted(clip.source_frame_starts, local_time + EPSILON, side="right") + - 1 + ) + frame_index = max(0, min(frame_index, clip.num_frames - 1)) + source_time = float(clip.source_frame_starts[frame_index]) + if source_time > local_time + 2 * EPSILON: + raise RuntimeError( + "Causal frame-selection violation: " + f"source PTS {source_time:.9f}s > grid-local time {local_time:.9f}s" + ) + if count_tail_reuse and local_time > float(clip.source_frame_starts[-1]) + EPSILON: + self._tail_frame_reuses += 1 + return clip_index, frame_index, global_time, source_time + + def _load_frame(self, clip_index: int, frame_index: int) -> Any: + from PIL import Image + + key = (clip_index, frame_index) + cached = self._frame_cache.get(key) + if cached is not None: + self._cache_hits += 1 + self._frame_cache.move_to_end(key) + return cached.copy() + + clip = self.clips[clip_index] + array = clip.reader[frame_index].asnumpy() + if clip.crop is not None: + x, y, width, height = clip.crop + array = array[y : y + height, x : x + width] + image = Image.fromarray(np.asarray(array)).copy() + if self.frame_cache_size: + self._frame_cache[key] = image + self._frame_cache.move_to_end(key) + while len(self._frame_cache) > self.frame_cache_size: + self._frame_cache.popitem(last=False) + return image.copy() + + def frames(self, start: float, end: float) -> list[Any]: + """Return causal global-grid frames in ``[start, end)``.""" + + if end <= start: + raise ValueError(f"Empty video interval [{start}, {end}]") + raw_first = int(math.ceil(start * self.fps - EPSILON)) + raw_last_exclusive = int(math.ceil(end * self.fps - EPSILON)) + grid_indices = _grid_indices(start, end, self.fps) + if raw_first >= raw_last_exclusive: + self._short_interval_fallbacks += 1 + + output = [] + for grid_index in grid_indices: + clip_index, frame_index, _, _ = self._source_frame_for_grid(grid_index) + output.append(self._load_frame(clip_index, frame_index)) + self._request_count += 1 + if len(output) % 2: + output.append(output[-1].copy()) + self._odd_frame_duplications += 1 + return output + + def audio_slice(self, start: float, end: float) -> np.ndarray: + """Return the exact 16-kHz mono samples in ``[start, end)``.""" + + first = max(0, int(math.floor(start * SAMPLE_RATE + EPSILON))) + last = min( + len(self.audio), int(math.floor(end * SAMPLE_RATE + EPSILON)) + ) + if last <= first: + raise ValueError(f"Empty audio interval [{start}, {end}]") + return self.audio[first:last] + + def sample_mapping(self, grid_index: int) -> dict[str, Any]: + """Describe one visual-grid mapping without decoding its pixels.""" + + clip_index, frame_index, grid_time, source_time = self._source_frame_for_grid( + grid_index, count_tail_reuse=False + ) + clip = self.clips[clip_index] + return { + "grid_index": int(grid_index), + "grid_time": float(grid_time), + "source_clip": str(clip.path), + "source_frame_index": int(frame_index), + "source_frame_pts_seconds": float(source_time), + "dense_source_frame_time": float(clip.dense_start + source_time), + } + + def metadata(self) -> dict[str, Any]: + return { + "decoder": "decord_with_frame_pts_and_ffmpeg_pcm_audio", + "visual_sampling": "global_grid_latest_source_pts_not_after_tick", + "interval_convention": "half_open_[start,end)", + "output_video_fps": self.fps, + "expected_grid_frames": self.expected_grid_frames, + "required_duration_seconds": self.required_duration, + "dense_source_duration_seconds": self.dense_duration, + "audio_sample_rate": SAMPLE_RATE, + "audio_samples": int(len(self.audio)), + "temporal_patch_padding": "duplicate_last_observed_frame_when_odd", + "source_clips": [ + { + "path": str(clip.path), + "audio_path": str(clip.audio_path), + "dense_start": clip.dense_start, + "dense_end": clip.dense_end, + "duration": clip.duration, + "source_fps": clip.source_fps, + "source_frames": clip.num_frames, + "source_size": [clip.width, clip.height], + "crop": list(clip.crop) if clip.crop is not None else None, + "first_source_pts": float(clip.source_frame_starts[0]), + "last_source_pts": float(clip.source_frame_starts[-1]), + } + for clip in self.clips + ], + "runtime_counters": { + "frame_requests": self._request_count, + "frame_cache_hits": self._cache_hits, + "odd_frame_duplications": self._odd_frame_duplications, + "short_interval_fallbacks": self._short_interval_fallbacks, + "tail_frame_reuses": self._tail_frame_reuses, + }, + } + + +def _manifest_session_specs( + manifest_path: Path, + session_name: str, +) -> tuple[list[SourceClipSpec], dict[str, Any], Path | None]: + payload = json.loads(manifest_path.read_text()) + matches = [ + session + for session in payload.get("sessions", []) + if session.get("session") == session_name + ] + if len(matches) != 1: + raise RuntimeError( + f"Expected exactly one {session_name!r} in {manifest_path}, " + f"found {len(matches)}" + ) + session = matches[0] + specs = [] + for record in session.get("clips", []): + duration = record.get("video_duration_seconds") + crop_record = record.get("egospeak_crop") + crop = None + if crop_record is not None: + crop = ( + int(crop_record["x"]), + int(crop_record["y"]), + int(crop_record["width"]), + int(crop_record["height"]), + ) + specs.append( + SourceClipSpec( + path=Path(record["path"]), + duration=float(duration) if duration is not None else None, + crop=crop, + ) + ) + if not specs: + raise RuntimeError(f"No source clips recorded for {session_name}") + target_record = session.get("target") or {} + target_path = Path(target_record["path"]) if target_record.get("path") else None + return specs, session, target_path + + +def _easycom_directory_specs(session_dir: Path) -> list[SourceClipSpec]: + session_dir = session_dir.expanduser().resolve() + if not session_dir.is_dir(): + raise FileNotFoundError(f"EasyCom session directory does not exist: {session_dir}") + paths = sorted(session_dir.glob("*.mp4"), key=lambda path: parse_timecode(path.stem)) + if not paths: + raise RuntimeError(f"No MP4 clips found in {session_dir}") + return [ + SourceClipSpec(path=path, crop_rightmost_width=EASYCOM_WIDTH) + for path in paths + ] + + +def _parse_crop(value: str | None) -> tuple[int, int, int, int] | None: + if value is None: + return None + try: + values = tuple(int(item) for item in value.split(",")) + except ValueError as exc: + raise argparse.ArgumentTypeError("crop must be x,y,width,height") from exc + if len(values) != 4: + raise argparse.ArgumentTypeError("crop must be x,y,width,height") + return values + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--video", type=Path, help="One original video clip") + source.add_argument( + "--easycom-session-dir", + type=Path, + help="Original EasyCom Session_N directory containing retained MP4 clips", + ) + source.add_argument( + "--easycom-manifest", + type=Path, + help="easycom_manifest.py JSON; also requires --session", + ) + parser.add_argument("--session", help="Session_N selected from --easycom-manifest") + parser.add_argument( + "--audio", + type=Path, + help="Optional separate audio for --video; defaults to its embedded track", + ) + parser.add_argument( + "--crop", + help="Optional crop for --video as x,y,width,height", + ) + parser.add_argument( + "--use-audio-in-video", + action="store_true", + help="Accepted for command compatibility; direct extraction always uses AV interleaving", + ) + parser.add_argument( + "--target", + type=Path, + help="target_perframe [T,3]; T is the authoritative full duration", + ) + parser.add_argument("--output", type=Path) + parser.add_argument( + "--mode", + choices=("causal-prefix", "open-stream-replay"), + default="open-stream-replay", + ) + parser.add_argument("--model-id", default=core.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=0.0, + help="0 keeps Thinker KV for the full extraction; positive values use a rolling window", + ) + parser.add_argument("--video-fps", type=float, default=2.0) + 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("--frame-cache-size", type=int, default=8) + parser.add_argument("--verify-every", type=int, default=0) + parser.add_argument("--verify-continuous-every", type=int, default=0) + parser.add_argument("--no-progress", action="store_true") + 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") + parser.add_argument( + "--media-only", + action="store_true", + help="Validate/index direct media and exit before loading Qwen", + ) + return parser.parse_args() + + +def _resolve_duration_and_target( + args: argparse.Namespace, + manifest_target: Path | None, +) -> tuple[Path | None, int | None, float]: + target_path = args.target + if target_path is None and manifest_target is not None and manifest_target.is_file(): + target_path = manifest_target + target_steps = None + if target_path is not None: + target_path = target_path.expanduser().resolve() + target_steps = core.target_num_steps(target_path) + if not math.isclose(args.start, 0.0, rel_tol=0.0, abs_tol=1e-9): + raise SystemExit("Full target extraction requires --start 0") + target_end = target_steps * args.step_seconds + if args.end is not None and not math.isclose( + args.end, target_end, rel_tol=0.0, abs_tol=1e-6 + ): + raise SystemExit( + f"--end {args.end} conflicts with target length {target_steps} " + f"at {args.step_seconds}s ({target_end}s)" + ) + end = target_end + else: + if args.end is None: + raise SystemExit("Pass --target for a full run or --end for a bounded test") + end = float(args.end) + if end <= args.start: + raise SystemExit("--end must be greater than --start") + return target_path, target_steps, end + + +def _effective_memory_seconds( + requested: float, + *, + duration: float, + block_seconds: float, +) -> float: + if requested < 0: + raise SystemExit("--memory-seconds must be >= 0") + if requested > 0: + return requested + blocks = max(1, int(math.ceil(duration / block_seconds - EPSILON))) + return blocks * block_seconds + + +def _print_media_summary(reader: CausalSourceAV, start: float, end: float) -> None: + print( + f"Direct media valid: {len(reader.clips)} source clip(s), " + f"required={end - start:.3f}s, grid={reader.expected_grid_frames} " + f"frames at {reader.fps:g} FPS, audio={len(reader.audio)} samples" + ) + candidates = [0] + if reader.expected_grid_frames > 1: + candidates.append(reader.expected_grid_frames - 1) + for grid_index in candidates: + mapping = reader.sample_mapping(grid_index) + print( + " grid " + f"{mapping['grid_time']:.3f}s -> " + f"{Path(mapping['source_clip']).name} " + f"frame={mapping['source_frame_index']} " + f"pts={mapping['source_frame_pts_seconds']:.6f}s" + ) + + +def main() -> None: + args = parse_args() + core._configure_qwen_logging() + if args.verify_every < 0 or args.verify_continuous_every < 0: + raise SystemExit("verification intervals must be >= 0") + if args.frame_cache_size < 0: + raise SystemExit("--frame-cache-size must be >= 0") + if args.output is None and not args.media_only: + raise SystemExit("--output is required unless --media-only is used") + if args.audio is not None and args.video is None: + raise SystemExit("--audio is supported only together with --video") + if args.session is not None and args.easycom_manifest is None: + raise SystemExit("--session is used only with --easycom-manifest") + if args.easycom_manifest is not None and args.session is None: + raise SystemExit("--easycom-manifest requires --session Session_N") + + manifest_record = None + manifest_target = None + if args.video is not None: + specs = [ + SourceClipSpec( + path=args.video, + crop=_parse_crop(args.crop), + audio_path=args.audio, + ) + ] + source_kind = "single_original_video" + elif args.easycom_session_dir is not None: + if args.crop is not None: + raise SystemExit("EasyCom direct extraction uses the fixed rightmost-1920 crop") + specs = _easycom_directory_specs(args.easycom_session_dir) + source_kind = "easycom_dense_original_clips" + else: + if args.crop is not None: + raise SystemExit("EasyCom manifest supplies its crop metadata") + specs, manifest_record, manifest_target = _manifest_session_specs( + args.easycom_manifest.expanduser().resolve(), args.session + ) + source_kind = "easycom_manifest_original_clips" + + target_path, target_steps, end = _resolve_duration_and_target( + args, manifest_target + ) + required_duration = end + reader = CausalSourceAV( + specs, + output_fps=args.video_fps, + required_duration=required_duration, + frame_cache_size=args.frame_cache_size, + ) + is_easycom = ( + args.easycom_session_dir is not None + or args.easycom_manifest is not None + ) + if is_easycom: + for clip in reader.clips: + if clip.height != EASYCOM_HEIGHT: + raise RuntimeError( + f"Expected EasyCom height {EASYCOM_HEIGHT}, got {clip.height}: " + f"{clip.path}" + ) + if not math.isclose( + clip.source_fps, 20.0, rel_tol=0.0, abs_tol=1e-6 + ): + raise RuntimeError( + f"Expected EasyCom source FPS 20, got {clip.source_fps:g}: " + f"{clip.path}" + ) + if target_steps is not None: + source_frames = sum(clip.num_frames for clip in reader.clips) + media_steps = source_frames // 4 + if media_steps != target_steps: + raise RuntimeError( + f"EasyCom dense source gives {media_steps} complete 5-Hz " + f"steps from {source_frames} frames, but target has " + f"{target_steps}" + ) + _print_media_summary(reader, args.start, end) + if args.media_only: + return + + duration = end - args.start + effective_memory = _effective_memory_seconds( + args.memory_seconds, + duration=duration, + block_seconds=args.native_block_seconds, + ) + if args.memory_seconds == 0 and duration > 120: + print( + "WARNING: persistent Thinker KV grows for the entire media " + f"duration ({duration:.1f}s). Long EasyCom sessions can exceed " + "the model context or GPU memory; a positive --memory-seconds is " + "a bounded-memory experiment rather than native persistent-KV " + "semantics.", + flush=True, + ) + config = core.StreamConfig( + step_seconds=args.step_seconds, + native_block_seconds=args.native_block_seconds, + memory_seconds=effective_memory, + video_fps=args.video_fps, + video_min_pixels=args.video_min_pixels, + video_max_pixels=args.video_max_pixels, + ) + config.validate() + + model, processor = core.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) + media_path_for_api = reader.clips[0].path + + if args.mode == "causal-prefix": + extractor = core.CausalPrefixExtractor( + model, processor, config, fixed_grid_av=reader + ) + result = extractor.extract( + video_path=media_path_for_api, + audio_path=None, + start=args.start, + end=end, + use_audio_in_video=True, + show_progress=not args.no_progress, + ) + else: + extractor = core.OpenStreamReplayExtractor( + model, processor, config, fixed_grid_av=reader + ) + result = extractor.extract( + video_path=media_path_for_api, + audio_path=None, + start=args.start, + end=end, + use_audio_in_video=True, + verify_every=args.verify_every, + verify_continuous_every=args.verify_continuous_every, + show_progress=not args.no_progress, + ) + + if target_steps is not None and len(result.times) != target_steps: + raise RuntimeError( + f"Extractor returned {len(result.times)} states, expected {target_steps}" + ) + + persistent_memory = args.memory_seconds == 0 + payload: dict[str, Any] = { + "model_id": args.model_id, + "mode": args.mode, + "stream_representation": ( + "single_causal_prefix" + if args.mode == "causal-prefix" + else "single_open_media_turn_with_2s_payload_commits" + ), + "config": asdict(config), + "requested_memory_seconds": args.memory_seconds, + "persistent_thinker_kv": persistent_memory, + "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(target_path) if target_path is not None else None, + "source_kind": source_kind, + "source_manifest_record": manifest_record, + "direct_media": reader.metadata(), + "video_sampling_grid": ( + "direct_pts_causal_global_grid_half_open_with_tail_duplication" + ), + "open_stream_policy": ( + { + "commit": "completed_2s_av_payload_only", + "temporary": "partial_payload_plus_media_closure_plus_readout", + "tmrope_positions_per_chunk": core._tmrope_positions_per_chunk(model), + "memory": ( + "persistent_full_extraction_thinker_kv" + if persistent_memory + else "rolling_complete_blocks_plus_current_partial_block" + ), + } + 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 = core._load_head( + args.head_checkpoint, core._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) + + 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({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 item.get("kind") == kind] + print( + f" {kind}: n={len(checks)}, " + f"worst max_abs={max(item['max_abs'] for item in checks):.6g}, " + f"lowest cosine={min(item['cosine'] for item in checks):.8f}" + ) + + +if __name__ == "__main__": + main() From 8145be410866dd0675038d92347aea34d24b84d5 Mon Sep 17 00:00:00 2001 From: WangXin93 Date: Mon, 17 Aug 2026 22:16:27 +0800 Subject: [PATCH 3/3] Add scripts for extracting EasyCom session features: `run_extract_easycom_direct.sh` for full session extraction and `run_smoke_easycom_direct.sh` for smoke testing on a single GPU. Both scripts utilize `streaming_turn_model_directly.py` for processing and include necessary configurations and logging. This enhances the EasyCom pipeline for better session management and feature extraction. --- run_extract_easycom_direct.sh | 134 ++++++++++++++++++++++++++++++++++ run_smoke_easycom_direct.sh | 66 +++++++++++++++++ 2 files changed, 200 insertions(+) create mode 100644 run_extract_easycom_direct.sh create mode 100644 run_smoke_easycom_direct.sh diff --git a/run_extract_easycom_direct.sh b/run_extract_easycom_direct.sh new file mode 100644 index 0000000..c95d420 --- /dev/null +++ b/run_extract_easycom_direct.sh @@ -0,0 +1,134 @@ +#!/usr/bin/env bash +# Extract open-stream Thinker states for all 12 EasyCom sessions via +# streaming_turn_model_directly.py (original clips, no prepared 2 FPS MP4). +# +# Recommended for full sessions: --memory-seconds 20 (bounded KV). +# Do not default to --memory-seconds 0 on multi-minute sessions. +set -euo pipefail + +PROJECT_ROOT="${PROJECT_ROOT:-/mnt/58462F4C462F29E8/Code/icassp2027-experiment}" +MODEL="${MODEL:-/mnt/58462F4C462F29E8/Models/Qwen2.5-Omni-3B}" +VIDEO_ROOT="${VIDEO_ROOT:-/mnt/58462F4C462F29E8/Datasets/easycom/Main/Video_Compressed}" +TARGET_ROOT="${TARGET_ROOT:-/mnt/58462F4C462F29E8/Code/EgoSpeak/dataset/EasyCom/target_perframe}" +FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/easycom_direct}" +LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/easycom_direct}" + +MEMORY_SECONDS="${MEMORY_SECONDS:-20}" +NATIVE_BLOCK_SECONDS="${NATIVE_BLOCK_SECONDS:-2}" +VIDEO_FPS="${VIDEO_FPS:-2}" +MODE="${MODE:-open-stream-replay}" +DTYPE="${DTYPE:-bfloat16}" +# Space-separated host GPU ids, e.g. "0" or "0 1 2 3". One session per GPU at a time. +GPUS="${GPUS:-0}" +# Optional: only run selected session ids, e.g. SESSIONS="1 2 3" +SESSIONS="${SESSIONS:-1 2 3 4 5 6 7 8 9 10 11 12}" + +# StreamConfig requires memory / native_block to be an integer (e.g. 20/2=10). +# Values like 15 yield 7.5 and crash before loading the model. +python - "${MEMORY_SECONDS}" "${NATIVE_BLOCK_SECONDS}" <<'PY' +import sys +memory, block = float(sys.argv[1]), float(sys.argv[2]) +ratio = memory / block +if abs(ratio - round(ratio)) > 1e-6: + raise SystemExit( + f"MEMORY_SECONDS={memory} is invalid: memory/native_block={ratio} " + f"(native_block={block}). Use a multiple of {block:g}, e.g. 16, 18, 20." + ) +print(f"memory_seconds={memory:g} native_block={block:g} blocks={int(round(ratio))}") +PY + +mkdir -p "${FEATURE_ROOT}" "${LOG_ROOT}" +cd "${PROJECT_ROOT}" + +run_one() { + local gpu="$1" + local sid="$2" + local session="Session_${sid}" + local session_dir="${VIDEO_ROOT}/${session}" + local target="${TARGET_ROOT}/${session}.npy" + local output="${FEATURE_ROOT}/${session}.pt" + local log="${LOG_ROOT}/${session}.log" + + if [[ ! -d "${session_dir}" ]]; then + echo "Missing session dir: ${session_dir}" >&2 + return 1 + fi + if [[ ! -s "${target}" ]]; then + echo "Missing target: ${target}" >&2 + return 1 + fi + if [[ -s "${output}" ]]; then + echo "Skipping existing output: ${output}" + return 0 + fi + + { + echo "$(date): GPU ${gpu} starting ${session}" + echo "model=${MODEL}" + echo "session_dir=${session_dir}" + echo "target=${target}" + echo "output=${output}" + echo "memory_seconds=${MEMORY_SECONDS} video_fps=${VIDEO_FPS} mode=${MODE}" + + if ! CUDA_VISIBLE_DEVICES="${gpu}" PYTHONUNBUFFERED=1 \ + python streaming_turn_model_directly.py \ + --model-id "${MODEL}" \ + --easycom-session-dir "${session_dir}" \ + --target "${target}" \ + --use-audio-in-video \ + --mode "${MODE}" \ + --memory-seconds "${MEMORY_SECONDS}" \ + --native-block-seconds "${NATIVE_BLOCK_SECONDS}" \ + --video-fps "${VIDEO_FPS}" \ + --dtype "${DTYPE}" \ + --flash-attention-2 \ + --device cuda:0 \ + --output "${output}" + then + echo "$(date): GPU ${gpu} FAILED ${session}" + exit 1 + fi + + echo "$(date): GPU ${gpu} completed ${session}" + } >"${log}" 2>&1 +} + +# Round-robin sessions across GPUs. +mapfile -t GPU_ARR < <(echo "${GPUS}" | tr ' ' '\n' | sed '/^$/d') +mapfile -t SID_ARR < <(echo "${SESSIONS}" | tr ' ' '\n' | sed '/^$/d') + +declare -A WORKERS=() +for gpu in "${GPU_ARR[@]}"; do + WORKERS["${gpu}"]="" +done + +i=0 +for sid in "${SID_ARR[@]}"; do + gpu="${GPU_ARR[$((i % ${#GPU_ARR[@]}))]}" + WORKERS["${gpu}"]+="${sid} " + i=$((i + 1)) +done + +pids=() +status=0 +for gpu in "${GPU_ARR[@]}"; do + sids="${WORKERS[${gpu}]}" + [[ -z "${sids// }" ]] && continue + ( + for sid in ${sids}; do + run_one "${gpu}" "${sid}" || exit 1 + done + ) & + pids+=("$!") +done + +for pid in "${pids[@]}"; do + wait "${pid}" || status=1 +done + +if [[ "${status}" -ne 0 ]]; then + echo "At least one EasyCom direct extraction failed. Check ${LOG_ROOT}." >&2 + exit 1 +fi + +echo "ALL EasyCom direct extractions completed -> ${FEATURE_ROOT}" diff --git a/run_smoke_easycom_direct.sh b/run_smoke_easycom_direct.sh new file mode 100644 index 0000000..b140835 --- /dev/null +++ b/run_smoke_easycom_direct.sh @@ -0,0 +1,66 @@ +#!/usr/bin/env bash +# Smoke: one EasyCom session on one GPU via streaming_turn_model_directly.py +# +# Short smoke (default): --end only, NO --target +# (passing both requires --end == full target duration) +# Full session: END= (empty) and --target is used +set -euo pipefail + +# Remote defaults (override as needed) +PROJECT_ROOT="${PROJECT_ROOT:-/data1/yfzx/wangx/Code/icassp2027-experiment}" +MODEL="${MODEL:-/data1/yfzx/wangx/Models/Qwen2.5-Omni-3B}" +VIDEO_ROOT="${VIDEO_ROOT:-/data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed}" +TARGET_ROOT="${TARGET_ROOT:-/data1/yfzx/wangx/Datasets/easycom/target_frame}" +FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/easycom_direct_smoke}" +LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/easycom_direct_smoke}" + +SESSION_ID="${SESSION_ID:-1}" +GPU="${GPU:-0}" +# Short smoke: END=4. Full session smoke: END="" (empty). +END="${END:-4}" +MEMORY_SECONDS="${MEMORY_SECONDS:-20}" +VIDEO_FPS="${VIDEO_FPS:-2}" + +session="Session_${SESSION_ID}" +session_dir="${VIDEO_ROOT}/${session}" +target="${TARGET_ROOT}/${session}.npy" +output="${FEATURE_ROOT}/${session}_gpu${GPU}.pt" +log="${LOG_ROOT}/${session}_gpu${GPU}.log" + +mkdir -p "${FEATURE_ROOT}" "${LOG_ROOT}" +cd "${PROJECT_ROOT}" + +[[ -d "${session_dir}" ]] || { echo "Missing: ${session_dir}" >&2; exit 1; } + +extra_args=() +if [[ -n "${END}" ]]; then + # Bounded smoke: do not pass --target (length would conflict with --end). + extra_args+=(--end "${END}") + MEMORY_SECONDS=0 +else + [[ -s "${target}" ]] || { echo "Missing: ${target}" >&2; exit 1; } + extra_args+=(--target "${target}") +fi + +echo "GPU=${GPU} session=${session} end=${END:-full} memory=${MEMORY_SECONDS}" +echo "log -> ${log}" + +{ + echo "$(date): smoke start ${session} on GPU ${GPU}" + CUDA_VISIBLE_DEVICES="${GPU}" PYTHONUNBUFFERED=1 \ + python streaming_turn_model_directly.py \ + --model-id "${MODEL}" \ + --easycom-session-dir "${session_dir}" \ + --use-audio-in-video \ + --mode open-stream-replay \ + --memory-seconds "${MEMORY_SECONDS}" \ + --video-fps "${VIDEO_FPS}" \ + --dtype bfloat16 \ + --flash-attention-2 \ + --device cuda:0 \ + --verify-every 5 \ + --verify-continuous-every 5 \ + "${extra_args[@]}" \ + --output "${output}" + echo "$(date): smoke done -> ${output}" +} 2>&1 | tee "${log}"