455 lines
15 KiB
Python
455 lines
15 KiB
Python
#!/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", "torchcodec", "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()
|