Add extract_ego4d_batch.py for efficient EgoSpeak clip extraction. This script loads the Qwen model once and processes selected manifest rows sequentially, ensuring each clip is handled independently. Update README with usage instructions and details on the extraction process, including handling of completed outputs and memory management.

README.md

@@ -685,6 +685,38 @@ conversion can differ slightly; validate a short Decord/PyNvVideoCodec pair
 before extracting every session, and do not mix decoder backends across
 train/test splits.
 
+### Persistent-model Ego4D batch extraction
+
+`extract_ego4d_batch.py` loads Qwen and its processor once, then extracts the
+selected EgoSpeak manifest rows sequentially. Each clip receives a new media
+reader, Thinker KV cache, TMRoPE origin, and rolling-memory state; only the
+immutable model weights and processor are reused. Manifest indices are
+one-based data-row indices and do not count the TSV header.
+
+For example, process manifest samples 81 through the end on visible GPU 0:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+OMP_NUM_THREADS=4 \
+MKL_NUM_THREADS=4 \
+OPENBLAS_NUM_THREADS=4 \
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+python extract_ego4d_batch.py \
+  --data-root /data1/yfzx/wangx/Datasets/ego4d \
+  --model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \
+  --start-index 81 \
+  --memory-seconds 20 \
+  --video-fps 2 \
+  --decoder decord \
+  --dtype bfloat16 \
+  --flash-attention-2 \
+  --device cuda:0
+```
+
+Completed non-empty `.pt` files are skipped. Each new output is first written
+to a process-specific `.partial` path and atomically renamed after extraction,
+so rerunning the same range resumes safely.
+
 For CPU Decord, the direct reader intentionally indexes each clip and releases
 that temporary reader, then keeps an LRU of at most two live readers during
 extraction. Each reader defaults to one decode thread. These limits avoid one

extract_ego4d_batch.py

@@ -0,0 +1,454 @@
+#!/usr/bin/env python3
+"""Extract EgoSpeak's Ego4D clips while loading Qwen only once.
+
+Manifest sample indices are one-based and refer to data rows, excluding the
+header.  For example, ``--start-index 31 --end-index 60`` selects exactly the
+same clips as ``sed -n '32,61p' ego4d_egospeak.tsv``.
+
+Only immutable model weights and the processor are reused.  Every clip invokes
+the direct extractor independently, which creates a new media reader, rolling
+KV cache, TMRoPE origin, and readout sequence.  Consequently no temporal state
+is carried between dataset samples.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import gc
+import math
+import os
+import sys
+import time
+import traceback
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Sequence
+
+import torch
+
+import streaming_turn_model as core
+import streaming_turn_model_directly as direct
+
+
+TARGET_HZ = 5
+DEFAULT_OUTPUT_NAME = "open_stream_2fps_20s_decord"
+EXPECTED_MANIFEST_COLUMNS = (
+    "split",
+    "uid",
+    "steps",
+    "duration",
+    "video_frames_2fps",
+)
+
+
+@dataclass(frozen=True)
+class ClipRecord:
+    index: int
+    split: str
+    uid: str
+    steps: int
+    duration: float
+    video_frames_2fps: int
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--data-root",
+        type=Path,
+        default=Path("/data/wangx/Datasets/ego4d"),
+        help="Ego4D experiment root (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--manifest",
+        type=Path,
+        help="Defaults to DATA_ROOT/ego4d_egospeak.tsv",
+    )
+    parser.add_argument(
+        "--video-root",
+        type=Path,
+        help="Defaults to DATA_ROOT/v1/clips",
+    )
+    parser.add_argument(
+        "--target-root",
+        type=Path,
+        help="Defaults to DATA_ROOT/ego4d_target_frame",
+    )
+    parser.add_argument(
+        "--output-root",
+        type=Path,
+        help=(
+            "Defaults to DATA_ROOT/omni_features/"
+            f"{DEFAULT_OUTPUT_NAME}"
+        ),
+    )
+    parser.add_argument(
+        "--model-id",
+        default="/data/wangx/Models/Qwen2.5-Omni-3B",
+    )
+    parser.add_argument(
+        "--start-index",
+        type=int,
+        default=1,
+        help="First one-based manifest data row to process (default: 1)",
+    )
+    parser.add_argument(
+        "--end-index",
+        type=int,
+        help="Last one-based manifest data row to process, inclusive",
+    )
+    parser.add_argument(
+        "--split",
+        choices=("all", "train", "test"),
+        default="all",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Replace completed outputs instead of resuming around them",
+    )
+    parser.add_argument(
+        "--continue-on-error",
+        action="store_true",
+        help="Record a failed clip and continue with the next selected clip",
+    )
+    parser.add_argument("--memory-seconds", type=float, default=20.0)
+    parser.add_argument("--step-seconds", type=float, default=0.2)
+    parser.add_argument("--native-block-seconds", type=float, default=2.0)
+    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=64)
+    parser.add_argument(
+        "--decoder",
+        choices=("decord", "pynvcodec"),
+        default="decord",
+    )
+    parser.add_argument("--decoder-gpu-id", type=int, default=0)
+    parser.add_argument("--decoder-threads", type=int, default=1)
+    parser.add_argument("--max-open-readers", type=int, default=1)
+    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(
+        "--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(
+        "--empty-cache-every",
+        type=int,
+        default=0,
+        help="Call torch.cuda.empty_cache every N processed clips; 0 disables",
+    )
+    args = parser.parse_args(argv)
+
+    if args.start_index <= 0:
+        parser.error("--start-index must be positive")
+    if args.end_index is not None and args.end_index < args.start_index:
+        parser.error("--end-index must be >= --start-index")
+    if args.empty_cache_every < 0:
+        parser.error("--empty-cache-every must be >= 0")
+    if args.step_seconds <= 0:
+        parser.error("--step-seconds must be positive")
+    return args
+
+
+def resolve_paths(args: argparse.Namespace) -> None:
+    args.data_root = args.data_root.expanduser().resolve()
+    args.manifest = (
+        args.manifest or args.data_root / "ego4d_egospeak.tsv"
+    ).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_features" / DEFAULT_OUTPUT_NAME
+    ).expanduser().resolve()
+
+
+def load_manifest(path: Path) -> list[ClipRecord]:
+    if not path.is_file():
+        raise FileNotFoundError(f"Missing EgoSpeak manifest: {path}")
+    records: list[ClipRecord] = []
+    seen: set[str] = set()
+    with path.open(newline="") as file:
+        reader = csv.DictReader(file, delimiter="\t")
+        if reader.fieldnames != list(EXPECTED_MANIFEST_COLUMNS):
+            raise RuntimeError(
+                f"Unexpected manifest columns in {path}: {reader.fieldnames}"
+            )
+        for index, row in enumerate(reader, start=1):
+            split = row["split"].strip()
+            uid = row["uid"].strip()
+            if split not in {"train", "test"}:
+                raise RuntimeError(f"Invalid split at manifest row {index}: {split}")
+            if not uid or uid in seen:
+                raise RuntimeError(
+                    f"Missing or duplicate UID at manifest row {index}: {uid!r}"
+                )
+            seen.add(uid)
+            steps = int(row["steps"])
+            duration = float(row["duration"])
+            video_frames = int(row["video_frames_2fps"])
+            if steps <= 0:
+                raise RuntimeError(f"Non-positive steps for {uid}: {steps}")
+            expected_duration = steps / TARGET_HZ
+            if not math.isclose(
+                duration, expected_duration, rel_tol=0.0, abs_tol=1e-6
+            ):
+                raise RuntimeError(
+                    f"Manifest duration mismatch for {uid}: {duration} vs "
+                    f"{expected_duration} from {steps} target steps"
+                )
+            records.append(
+                ClipRecord(
+                    index=index,
+                    split=split,
+                    uid=uid,
+                    steps=steps,
+                    duration=duration,
+                    video_frames_2fps=video_frames,
+                )
+            )
+    if not records:
+        raise RuntimeError(f"No clips in {path}")
+    return records
+
+
+def select_records(
+    records: list[ClipRecord], args: argparse.Namespace
+) -> list[ClipRecord]:
+    end_index = args.end_index if args.end_index is not None else len(records)
+    selected = [
+        record
+        for record in records
+        if args.start_index <= record.index <= end_index
+        and (args.split == "all" or record.split == args.split)
+    ]
+    if not selected:
+        raise RuntimeError(
+            f"No manifest rows selected by {args.start_index=} "
+            f"{args.end_index=} {args.split=}"
+        )
+    return selected
+
+
+def preflight(selected: list[ClipRecord], args: argparse.Namespace) -> None:
+    if not args.video_root.is_dir():
+        raise FileNotFoundError(f"Missing video root: {args.video_root}")
+    if not args.target_root.is_dir():
+        raise FileNotFoundError(f"Missing target root: {args.target_root}")
+    model_path = Path(args.model_id).expanduser()
+    if model_path.is_absolute() and not model_path.is_dir():
+        raise FileNotFoundError(f"Missing local model: {model_path}")
+
+    problems: list[str] = []
+    for record in selected:
+        video = args.video_root / f"{record.uid}.mp4"
+        target = args.target_root / f"{record.uid}.npy"
+        if not video.is_file() or video.stat().st_size == 0:
+            problems.append(f"missing video: {video}")
+        if not target.is_file() or target.stat().st_size == 0:
+            problems.append(f"missing target: {target}")
+            continue
+        try:
+            target_steps = core.target_num_steps(target)
+        except Exception as error:  # report all bad inputs before model loading
+            problems.append(f"invalid target {target}: {error}")
+            continue
+        if target_steps != record.steps:
+            problems.append(
+                f"target length mismatch {record.uid}: manifest={record.steps}, "
+                f"npy={target_steps}"
+            )
+    if problems:
+        preview = "\n".join(f"- {problem}" for problem in problems[:20])
+        suffix = "" if len(problems) <= 20 else f"\n... and {len(problems)-20} more"
+        raise RuntimeError(f"Ego4D preflight failed:\n{preview}{suffix}")
+
+
+def clip_argv(
+    record: ClipRecord,
+    args: argparse.Namespace,
+    output: Path,
+) -> list[str]:
+    values = [
+        "--model-id",
+        str(args.model_id),
+        "--video",
+        str(args.video_root / f"{record.uid}.mp4"),
+        "--target",
+        str(args.target_root / f"{record.uid}.npy"),
+        "--mode",
+        "open-stream-replay",
+        "--step-seconds",
+        str(args.step_seconds),
+        "--native-block-seconds",
+        str(args.native_block_seconds),
+        "--memory-seconds",
+        str(args.memory_seconds),
+        "--video-fps",
+        str(args.video_fps),
+        "--video-min-pixels",
+        str(args.video_min_pixels),
+        "--video-max-pixels",
+        str(args.video_max_pixels),
+        "--frame-cache-size",
+        str(args.frame_cache_size),
+        "--decoder",
+        args.decoder,
+        "--decoder-gpu-id",
+        str(args.decoder_gpu_id),
+        "--decoder-threads",
+        str(args.decoder_threads),
+        "--max-open-readers",
+        str(args.max_open_readers),
+        "--verify-every",
+        str(args.verify_every),
+        "--verify-continuous-every",
+        str(args.verify_continuous_every),
+        "--dtype",
+        args.dtype,
+        "--device",
+        args.device,
+        "--output",
+        str(output),
+    ]
+    if args.flash_attention_2:
+        values.append("--flash-attention-2")
+    if args.no_progress:
+        values.append("--no-progress")
+    return values
+
+
+def format_eta(seconds: float) -> str:
+    if not math.isfinite(seconds) or seconds < 0:
+        return "unknown"
+    hours, remainder = divmod(int(seconds), 3600)
+    minutes, secs = divmod(remainder, 60)
+    return f"{hours:02d}:{minutes:02d}:{secs:02d}"
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    args = parse_args(argv)
+    resolve_paths(args)
+    core._configure_qwen_logging()
+
+    records = load_manifest(args.manifest)
+    selected = select_records(records, args)
+    preflight(selected, args)
+    (args.output_root / "train").mkdir(parents=True, exist_ok=True)
+    (args.output_root / "test").mkdir(parents=True, exist_ok=True)
+
+    pending = [
+        record
+        for record in selected
+        if args.overwrite
+        or not (args.output_root / record.split / f"{record.uid}.pt").is_file()
+        or (args.output_root / record.split / f"{record.uid}.pt").stat().st_size == 0
+    ]
+    skipped = len(selected) - len(pending)
+    print(
+        f"Selected manifest samples {selected[0].index}-{selected[-1].index}: "
+        f"{len(selected)} clip(s), pending={len(pending)}, skipped={skipped}",
+        flush=True,
+    )
+    print(f"Output root: {args.output_root}", flush=True)
+    if not pending:
+        print("All selected clips already have non-empty outputs.")
+        return
+
+    print(f"Loading Qwen once from {args.model_id} ...", flush=True)
+    load_started = time.perf_counter()
+    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)
+    print(
+        f"Model loaded once in {time.perf_counter() - load_started:.1f}s; "
+        "per-clip KV/media state will reset.",
+        flush=True,
+    )
+
+    run_started = time.perf_counter()
+    completed = 0
+    failures: list[tuple[ClipRecord, str]] = []
+    for pending_index, record in enumerate(pending, start=1):
+        output = args.output_root / record.split / f"{record.uid}.pt"
+        temporary = output.with_name(f"{output.name}.partial.{os.getpid()}")
+        print(
+            f"\n[{pending_index}/{len(pending)} | manifest {record.index}] "
+            f"{record.split}/{record.uid}: {record.steps} states, "
+            f"{record.duration:.1f}s",
+            flush=True,
+        )
+        clip_started = time.perf_counter()
+        try:
+            direct.main(
+                clip_argv(record, args, temporary),
+                model=model,
+                processor=processor,
+            )
+            os.replace(temporary, output)
+            completed += 1
+            elapsed = time.perf_counter() - clip_started
+            average = (time.perf_counter() - run_started) / completed
+            remaining = len(pending) - pending_index
+            print(
+                f"Completed {output} in {elapsed:.1f}s; "
+                f"ETA {format_eta(average * remaining)}",
+                flush=True,
+            )
+        except KeyboardInterrupt:
+            print(f"Interrupted; partial path: {temporary}", file=sys.stderr)
+            raise
+        except Exception as error:
+            failures.append((record, str(error)))
+            print(
+                f"FAILED manifest sample {record.index} ({record.uid}): {error}\n"
+                f"Partial path, if created: {temporary}",
+                file=sys.stderr,
+                flush=True,
+            )
+            traceback.print_exc()
+            if not args.continue_on_error:
+                raise
+        finally:
+            # The model and processor stay referenced above.  Collection here
+            # releases the completed clip's reader, PIL frames, branch caches,
+            # and Python cache wrappers before opening the next sample.
+            gc.collect()
+            processed = completed + len(failures)
+            if (
+                args.empty_cache_every > 0
+                and processed > 0
+                and processed % args.empty_cache_every == 0
+                and torch.cuda.is_available()
+            ):
+                torch.cuda.empty_cache()
+
+    total = time.perf_counter() - run_started
+    print(
+        f"\nBatch finished in {format_eta(total)}: completed={completed}, "
+        f"skipped={skipped}, failed={len(failures)}",
+        flush=True,
+    )
+    if failures:
+        for record, message in failures:
+            print(f"- manifest {record.index} {record.uid}: {message}", file=sys.stderr)
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    main()

streaming_turn_model_directly.py

@@ -59,7 +59,7 @@ import subprocess
 from collections import OrderedDict
 from dataclasses import asdict, dataclass
 from pathlib import Path
-from typing import Any
+from typing import Any, Sequence
 
 import numpy as np
 import torch
@@ -778,7 +778,7 @@ def _parse_crop(value: str | None) -> tuple[int, int, int, int] | None:
     return values
 
 
-def parse_args() -> argparse.Namespace:
+def parse_args(argv: Sequence[str] | None = None) -> 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")
@@ -873,7 +873,7 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Validate/index direct media and exit before loading Qwen",
     )
-    return parser.parse_args()
+    return parser.parse_args(argv)
 
 
 def _resolve_duration_and_target(
@@ -941,9 +941,23 @@ def _print_media_summary(reader: CausalSourceAV, start: float, end: float) -> No
         )
 
 
-def main() -> None:
+def main(
-    args = parse_args()
+    argv: Sequence[str] | None = None,
+    *,
+    model: Any | None = None,
+    processor: Any | None = None,
+) -> None:
+    """Run one direct extraction, optionally reusing loaded Qwen objects.
+
+    ``model`` and ``processor`` must be supplied together.  Reusing them is
+    safe across independent clips because this function constructs a new media
+    reader, extractor, and temporal cache on every call.
+    """
+
+    args = parse_args(argv)
     core._configure_qwen_logging()
+    if (model is None) != (processor is None):
+        raise ValueError("model and processor must be supplied together")
     if args.verify_every < 0 or args.verify_continuous_every < 0:
         raise SystemExit("verification intervals must be >= 0")
     if args.frame_cache_size < 0:
@@ -1055,12 +1069,15 @@ def main() -> None:
     )
     config.validate()
 
-    model, processor = core.load_qwen_omni(
+    if model is None:
-        args.model_id,
+        model, processor = core.load_qwen_omni(
-        dtype=args.dtype,
+            args.model_id,
-        device_map={"": args.device},
+            dtype=args.dtype,
-        use_flash_attention_2=args.flash_attention_2,
+            device_map={"": args.device},
-    )
+            use_flash_attention_2=args.flash_attention_2,
+        )
+    else:
+        model.eval()
     model.requires_grad_(False)
     media_path_for_api = reader.clips[0].path