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()