180 lines
9.8 KiB
Python
180 lines
9.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Extract EasyCom 2/4/6/8/10/12s context banks with one shared AV pass.
|
|
|
|
One frozen model is loaded for all selected sessions. The current partial
|
|
native block is re-encoded at each 5-Hz cutoff, once rather than once per cap.
|
|
Existing direct-media validation and single-window CLI remain unchanged.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib.metadata
|
|
import json
|
|
import math
|
|
import re
|
|
import subprocess
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Sequence
|
|
|
|
import torch
|
|
|
|
import streaming_turn_model as core
|
|
import streaming_turn_model_directly as direct
|
|
from streaming_turn_multiwindow import DEFAULT_MEMORY_SECONDS, output_paths, validate_memory_seconds
|
|
|
|
|
|
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("--dataset-root", type=Path, help="EasyCom root containing Main/Video_Compressed")
|
|
source.add_argument("--easycom-video-root", type=Path, help="Directory directly containing Session_N")
|
|
parser.add_argument("--target-root", type=Path, help="Required for full extraction; target_perframe [T,3]")
|
|
parser.add_argument("--sessions", nargs="+", default=[f"Session_{i}" for i in range(1, 13)])
|
|
parser.add_argument("--memory-seconds-list", type=float, nargs="+", default=list(DEFAULT_MEMORY_SECONDS))
|
|
parser.add_argument("--output-root", type=Path, required=True)
|
|
parser.add_argument("--model-id", default=core.MODEL_ID)
|
|
parser.add_argument("--backbone-loader", choices=("qwen", "roma"), default="qwen")
|
|
parser.add_argument("--dtype", choices=("auto", "bfloat16", "float16", "float32"), default="bfloat16")
|
|
parser.add_argument("--device", default="cuda:0")
|
|
parser.add_argument("--flash-attention-2", action="store_true")
|
|
parser.add_argument("--end", type=float, help="Bounded zero-origin smoke only; does not attach full targets")
|
|
parser.add_argument("--readout-last-n-layers", type=int, default=4)
|
|
parser.add_argument("--readout-layer-fusion", choices=("mean", "stack"), default="stack")
|
|
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("--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=2)
|
|
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("--media-only", action="store_true")
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def session_commands(args: argparse.Namespace) -> list[tuple[str, list[str]]]:
|
|
"""Preflight all paths/settings before loading the backbone."""
|
|
caps = validate_memory_seconds(args.memory_seconds_list, 2.0)
|
|
core.StreamConfig(
|
|
memory_seconds=max(caps), video_fps=args.video_fps,
|
|
video_min_pixels=args.video_min_pixels, video_max_pixels=args.video_max_pixels,
|
|
readout_last_n_layers=args.readout_last_n_layers, readout_layer_fusion=args.readout_layer_fusion,
|
|
).validate()
|
|
if len(set(args.sessions)) != len(args.sessions) or not args.sessions:
|
|
raise ValueError("sessions must be nonempty and unique")
|
|
if args.end is None and args.target_root is None:
|
|
raise ValueError("full extraction requires --target-root; use --end only for smoke")
|
|
if args.end is not None:
|
|
if not math.isfinite(args.end) or args.end <= 0:
|
|
raise ValueError("smoke end must be finite and positive")
|
|
core._require_integer_ratio(args.end, 0.2, "duration / step")
|
|
if args.verify_every < 0 or args.verify_continuous_every < 0:
|
|
raise ValueError("verification intervals must be non-negative")
|
|
if args.decoder_threads <= 0 or args.max_open_readers <= 0 or args.frame_cache_size < 0:
|
|
raise ValueError("invalid reader/cache limits")
|
|
video_root = args.easycom_video_root or args.dataset_root / "Main" / "Video_Compressed"
|
|
commands = []
|
|
for session in args.sessions:
|
|
if re.fullmatch(r"Session_(?:[1-9]|1[0-2])", session) is None:
|
|
raise ValueError(f"invalid EasyCom session: {session}")
|
|
session_dir = video_root / session
|
|
if not session_dir.is_dir():
|
|
raise FileNotFoundError(session_dir)
|
|
direct._easycom_directory_specs(session_dir)
|
|
output = args.output_root / f"{session}.pt"
|
|
if not args.overwrite and not args.media_only:
|
|
for path in output_paths(output, caps).values():
|
|
if path.exists():
|
|
raise FileExistsError(f"{path}; use --overwrite explicitly or a fresh output root")
|
|
command = ["--easycom-session-dir", str(session_dir), "--output", str(output),
|
|
"--mode", "open-stream-replay", "--memory-seconds-list", *map(str, caps)]
|
|
for name in (
|
|
"model_id", "backbone_loader", "dtype", "device", "readout_last_n_layers",
|
|
"readout_layer_fusion", "video_fps", "video_min_pixels", "video_max_pixels",
|
|
"decoder", "decoder_gpu_id", "decoder_threads", "max_open_readers", "frame_cache_size",
|
|
"verify_every", "verify_continuous_every",
|
|
):
|
|
command.extend(["--" + name.replace("_", "-"), str(getattr(args, name))])
|
|
for name in ("flash_attention_2", "no_progress", "media_only", "overwrite"):
|
|
if getattr(args, name):
|
|
command.append("--" + name.replace("_", "-"))
|
|
if args.end is not None:
|
|
command.extend(["--end", str(args.end)])
|
|
else:
|
|
target = args.target_root / f"{session}.npy"
|
|
core.target_num_steps(target) # Missing/invalid targets fail before model loading.
|
|
command.extend(["--target", str(target)])
|
|
commands.append((session, command))
|
|
return commands
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> None:
|
|
args = parse_args(argv)
|
|
commands = session_commands(args)
|
|
if args.media_only:
|
|
for _, command in commands:
|
|
direct.main(command)
|
|
return
|
|
if args.end is not None:
|
|
print("SMOKE ONLY: partial banks are not a full-session training/evaluation dataset.", flush=True)
|
|
core._configure_qwen_logging()
|
|
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, backbone_loader=args.backbone_loader,
|
|
)
|
|
model_loading_seconds = time.perf_counter() - load_started
|
|
repository = Path(__file__).resolve().parent
|
|
revision = subprocess.run(["git", "rev-parse", "HEAD"], cwd=repository, capture_output=True, text=True, check=False)
|
|
dirty = subprocess.run(["git", "status", "--porcelain"], cwd=repository, capture_output=True, text=True, check=False)
|
|
summary = {
|
|
"format": "easycom-multiwindow-extraction-v1", "arguments": {
|
|
key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items()
|
|
},
|
|
"split": {"train": [f"Session_{i}" for i in range(4, 13)],
|
|
"test": ["Session_1", "Session_2", "Session_3"]},
|
|
"backbone": "frozen", "partial_extraction": args.end is not None,
|
|
"code_revision": revision.stdout.strip() if revision.returncode == 0 else None,
|
|
"worktree_dirty": bool(dirty.stdout.strip()) if dirty.returncode == 0 else None,
|
|
"versions": {name: importlib.metadata.version(name) for name in ("torch", "transformers", "qwen-omni-utils")},
|
|
"model_loading_seconds": model_loading_seconds,
|
|
"session_wall_scope": "direct_media_index_decode_extract_save_excluding_model_load_and_summary_scan",
|
|
"runtime_scope": "joint_multi_bank_with_shared_AV_not_single_window_online_runtime",
|
|
"sessions": {},
|
|
}
|
|
for session, command in commands:
|
|
started = time.perf_counter()
|
|
direct.main(command, model=model, processor=processor)
|
|
elapsed = time.perf_counter() - started
|
|
records = {}
|
|
for cap, path in output_paths(args.output_root / f"{session}.pt", sorted(args.memory_seconds_list)).items():
|
|
payload = torch.load(path, map_location="cpu", weights_only=False)
|
|
checks = [c for c in payload["verification"] if c.get("kind") == "cache_replay"]
|
|
spans = payload["extraction_diagnostics"]["actual_context_seconds"]
|
|
records[f"{cap:g}s"] = {
|
|
"path": str(path), "shape": list(payload["states"].shape),
|
|
"num_steps": payload["num_steps"], "max_actual_context_seconds": spans.max().item(),
|
|
"replay_checks": len(checks),
|
|
"worst_replay_max_abs": max((c["max_abs"] for c in checks), default=None),
|
|
"lowest_replay_cosine": min((c["cosine"] for c in checks), default=None),
|
|
"production_processor_calls_shared": payload["num_steps"],
|
|
}
|
|
del payload
|
|
summary["sessions"][session] = {"joint_wall_seconds": elapsed, "banks": records}
|
|
args.output_root.mkdir(parents=True, exist_ok=True)
|
|
summary_path = args.output_root / "extraction_summary.json"
|
|
temporary = summary_path.with_suffix(".json.tmp")
|
|
temporary.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
|
|
temporary.replace(summary_path)
|
|
print(f"Summary: {args.output_root / 'extraction_summary.json'}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|