#!/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()