Add shared-AV multiwindow Qwen feature extraction

streaming_turn_model_directly.py

@@ -1061,6 +1061,14 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         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(
+        "--memory-seconds-list", type=float, nargs="+",
+        help="Export independent open-stream banks in context_Cs/ with shared AV encoding",
+    )
+    parser.add_argument(
+        "--overwrite", action="store_true",
+        help="Allow replacing existing multiwindow banks (single-window behavior is unchanged)",
+    )
     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(
@@ -1222,6 +1230,21 @@ def main(
 
     args = parse_args(argv)
     core._configure_qwen_logging()
+    multi_caps = None
+    if args.memory_seconds_list is not None:
+        from streaming_turn_multiwindow import output_paths, validate_memory_seconds
+
+        if args.mode != "open-stream-replay" or args.head_checkpoint is not None:
+            raise SystemExit("multiwindow export requires open-stream-replay without --head-checkpoint")
+        if args.memory_seconds != 0:
+            raise SystemExit("use --memory-seconds-list instead of --memory-seconds, not both")
+        if not math.isclose(args.start, 0.0, abs_tol=1e-9) or not math.isclose(args.step_seconds, 0.2, abs_tol=1e-9):
+            raise SystemExit("multiwindow dense banks require --start 0 and --step-seconds 0.2")
+        multi_caps = validate_memory_seconds(args.memory_seconds_list, args.native_block_seconds)
+        if args.output is not None and not args.media_only and not args.overwrite:
+            existing = [str(path) for path in output_paths(args.output, multi_caps).values() if path.exists()]
+            if existing:
+                raise FileExistsError(f"multiwindow outputs already exist; use --overwrite explicitly: {existing}")
     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:
@@ -1340,7 +1363,7 @@ def main(
             "pass a positive --memory-seconds (recommended: 8)"
         )
     effective_memory = _effective_memory_seconds(
-        args.memory_seconds,
+        max(multi_caps) if multi_caps else args.memory_seconds,
         duration=duration,
         block_seconds=args.native_block_seconds,
     )
@@ -1404,7 +1427,20 @@ def main(
 
     overlap_extractor = None
     working_grid_av = None
-    if args.mode == "causal-prefix":
+    multi_results = None
+    if multi_caps:
+        from streaming_turn_multiwindow import MultiWindowOpenStreamReplayExtractor
+
+        extractor = MultiWindowOpenStreamReplayExtractor(
+            model, processor, config, memory_seconds_list=multi_caps, fixed_grid_av=reader,
+        )
+        multi_results = 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,
+        )
+        result = multi_results[max(multi_caps)]
+    elif args.mode == "causal-prefix":
         extractor = core.CausalPrefixExtractor(
             model, processor, config, fixed_grid_av=reader
         )
@@ -1616,6 +1652,9 @@ def main(
         "verification": result.verification,
         "extraction_diagnostics": result.diagnostics,
     }
+    if multi_results is not None:
+        save_multiwindow_banks(payload, multi_results, args.output, overwrite=args.overwrite)
+        return
     if args.head_checkpoint is not None:
         head, head_loss_name = core._load_head(
             args.head_checkpoint, core._hidden_size(model), model.device
@@ -1647,5 +1686,56 @@ def main(
             )
 
 
+def save_multiwindow_banks(
+    payload: dict[str, Any], results: dict[float, core.ExtractionResult], output: Path,
+    *, overwrite: bool = False,
+) -> dict[float, Path]:
+    """Write ordinary head-compatible payloads; never overwrite implicitly."""
+    from local_turn_features import reference_alignment
+    from streaming_turn_multiwindow import output_paths
+
+    paths = output_paths(output, sorted(results))
+    banks = {}
+    for cap, result in results.items():
+        if len(result.times) != payload["num_steps"]:
+            raise ValueError(f"{cap:g}s bank has a different number of dense rows")
+        bank = dict(payload)
+        bank.update(
+            config={**payload["config"], "memory_seconds": cap},
+            requested_memory_seconds=cap, persistent_thinker_kv=False,
+            times=result.times, states=result.states, verification=result.verification,
+            extraction_diagnostics=result.diagnostics,
+            partial_extraction=payload["target_path"] is None,
+            open_stream_policy={**payload["open_stream_policy"],
+                                "memory": "rolling_complete_blocks_plus_current_partial_block",
+                                "rolling_replay_inputs": "reuse_cached_complete_block_av_embeddings"},
+            multiwindow_policy={
+                "memory_seconds_list": sorted(results), "context_cap_seconds": cap,
+                "shared": "media_reader_preprocessing_AV_block_and_tail_embeddings",
+                "independent": "per_context_Thinker_KV_replay_and_readout",
+                "long_context_KV_cropping": False,
+                "equivalence_reference": "single_window_open_stream_replay_same_cap",
+                "runtime_scope": "joint_multi_bank_extraction_not_single_bank_online_latency",
+            },
+        )
+        reference_alignment(bank)
+        if not torch.isfinite(result.states).all():
+            raise ValueError(f"{cap:g}s bank contains non-finite states")
+        if paths[cap].exists() and not overwrite:
+            raise FileExistsError(paths[cap])
+        banks[cap] = bank
+    for cap, result in results.items():
+        paths[cap].parent.mkdir(parents=True, exist_ok=True)
+        temporary = paths[cap].with_suffix(".pt.tmp")
+        torch.save(banks[cap], temporary)
+        temporary.replace(paths[cap])
+        print(f"Saved {cap:g}s bank: {paths[cap]} states={tuple(result.states.shape)}")
+        checks = [item for item in result.verification if item.get("kind") == "cache_replay"]
+        if checks:
+            print(f"  independent replay n={len(checks)} max_abs={max(c['max_abs'] for c in checks):.6g} "
+                  f"min_cosine={min(c['cosine'] for c in checks):.8f}")
+    return paths
+
+
 if __name__ == "__main__":
     main()