tools/score_v3_m10_fast.py

"""Block A — Patched + multi-benchmark v3-M10 scorer.

Applies the Qwen3VLVisionPatchEmbed Conv3d → Linear monkey-patch
(from tools/run_qwen3_cache_fast.py) BEFORE any Qwen3 model is
loaded. On Blackwell + bf16, this gives ~64× speedup on the patch-
embed layer, bringing per-tick Qwen3 forward from ~16 s to ~0.26 s.

Supports four benchmarks via --benchmark:
  adas_to          - ADAS-TO Critic 285 clips
  sim_dataset      - CARLA Sim-to-Real 250 clips
  longdrive        - LongDrive 2.5 h continuous mp4
  kaggle_accident  - Kaggle accident competition 2,027 clips (zero-shot)

Output: appends "m10_v3" field to each existing *.qwen_scores.json,
or creates new JSONs for kaggle_accident.

Usage:
    python3 tools/score_v3_m10_fast.py --benchmark adas_to --skip_existing
    python3 tools/score_v3_m10_fast.py --benchmark sim_dataset --skip_existing
    python3 tools/score_v3_m10_fast.py --benchmark longdrive --skip_existing
    python3 tools/score_v3_m10_fast.py --benchmark kaggle_accident
"""
from __future__ import annotations

import sys
sys.path.insert(0, ".")

# ─── Apply Qwen3 fast-patch BEFORE loading any model ──────────────────────
import torch
import torch.nn as nn
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionPatchEmbed

_PATCH_APPLIED = {}


def _fast_patch_embed_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """Conv3d → Linear lazy replacement (math equivalent, ~64× faster on
    Blackwell + bf16)."""
    target_dtype = self.proj.weight.dtype

    if isinstance(self.proj, nn.Conv3d):
        conv = self.proj
        out_dim = conv.out_channels
        in_dim = (conv.in_channels * conv.kernel_size[0]
                  * conv.kernel_size[1] * conv.kernel_size[2])
        w_flat = conv.weight.detach().reshape(out_dim, in_dim).contiguous()
        bias = conv.bias.detach().clone() if conv.bias is not None else None
        new_proj = nn.Linear(in_dim, out_dim, bias=bias is not None)
        new_proj.weight.data.copy_(w_flat)
        if bias is not None:
            new_proj.bias.data.copy_(bias)
        new_proj.to(device=conv.weight.device, dtype=conv.weight.dtype)
        self.proj = new_proj
        if id(self) not in _PATCH_APPLIED:
            _PATCH_APPLIED[id(self)] = True
            print(f"[fast_patch] patched Qwen3VLVisionPatchEmbed @ id={id(self)}: "
                   f"Conv3d({in_dim}→{out_dim}) → Linear({in_dim}→{out_dim})",
                   flush=True)

    if hidden_states.dim() > 2 or hidden_states.shape[-1] != self.proj.in_features:
        hidden_states = hidden_states.reshape(-1, self.proj.in_features)
    hidden_states = hidden_states.to(dtype=target_dtype)
    return self.proj(hidden_states)


Qwen3VLVisionPatchEmbed.forward = _fast_patch_embed_forward
print("[fast_patch] Qwen3VLVisionPatchEmbed.forward replaced (lazy Conv3d → Linear)",
      flush=True)

# ─── Now imports that may load Qwen3 models ───────────────────────────────
import argparse
import csv
import json
import time
from pathlib import Path
from typing import List, Optional

import cv2
import numpy as np
import torch.nn.functional as F

ROOT = Path(__file__).resolve().parents[1]

# Reuse helpers from existing scorer
from tools import qwen_alert_demo as qad                            # noqa: E402
from tools.score_adasto_v3_m10 import load_v3_m10, score_one_clip   # noqa: E402

DEFAULT_QWEN3_BASE = ROOT / "models/Qwen3-VL-4B-Instruct"
DEFAULT_QWEN3_LORA = ROOT / "checkpoints/VLA/qwen3vl4b_cot_belief_perframe/best"
DEFAULT_M10_V3_HEAD = ROOT / "checkpoints/Policy/m10_qwen3vl4b_seed0/best/policy_head.pt"


# ─── Benchmark configs ────────────────────────────────────────────────────

def get_benchmark_config(name: str, args) -> dict:
    """Return paths + iter_clips function for the chosen benchmark."""
    if name == "adas_to":
        videos_dir = ROOT / "ADAS-TO-TEST"
        results_dir = ROOT / "ADAS-TO-TEST/results_qwen"
        json_files = sorted(results_dir.glob("*.qwen_scores.json"))
        clips = []
        for jp in json_files:
            cid = jp.name.replace(".qwen_scores.json", "")
            video = videos_dir / f"{cid}.mp4"
            if video.exists():
                clips.append((cid, video, jp))
        return dict(name=name, clips=clips, append_field="m10_v3",
                    create_jsons=False)

    if name == "sim_dataset":
        # Drive from the full takeover_manifest.csv (2,211 CARLA clips),
        # not the b50 stratified subset (250 clips). With --skip_existing,
        # already-scored clips are skipped, so this is incremental.
        videos_root = ROOT / "accident/sim_dataset/videos"
        results_dir = ROOT / "accident/results_qwen"
        results_dir.mkdir(parents=True, exist_ok=True)
        manifest_csv = ROOT / "accident/takeover_manifest.csv"
        clips = []
        all_videos = list(videos_root.rglob("*.mp4"))
        videos_by_id = {p.stem: p for p in all_videos}
        with manifest_csv.open() as fh:
            for row in csv.DictReader(fh):
                cid = row["clip"]
                video = videos_by_id.get(cid)
                if video is None or not video.exists():
                    continue
                jp = results_dir / f"{cid}.qwen_scores.json"
                clips.append((cid, video, jp))
        return dict(name=name, clips=clips, append_field="m10_v3",
                    create_jsons=True)

    if name == "longdrive":
        videos_dir = ROOT / "LongDrive"
        results_dir = ROOT / "LongDrive/results_qwen_smoke_44"
        # LongDrive: single mp4 → single JSON
        clips = []
        for video in sorted(videos_dir.glob("*.mp4")):
            cid = video.stem
            jp = results_dir / f"{cid}.qwen_scores.json"
            if jp.exists():
                clips.append((cid, video, jp))
        return dict(name=name, clips=clips, append_field="m10_v3",
                    create_jsons=False)

    if name == "kaggle_accident":
        videos_dir = ROOT / "accident/videos"
        metadata_csv = ROOT / "accident/test_metadata.csv"
        results_dir = ROOT / "accident/kaggle_zero_shot/results_v3_m10"
        results_dir.mkdir(parents=True, exist_ok=True)
        clips = []
        with metadata_csv.open() as fh:
            for row in csv.DictReader(fh):
                video = ROOT / "accident" / row["path"]   # path = "videos/xxx.mp4"
                if not video.exists():
                    continue
                cid = video.stem
                jp = results_dir / f"{cid}.qwen_scores.json"
                clips.append((cid, video, jp))
        return dict(name=name, clips=clips, append_field="m10_v3",
                    create_jsons=True)

    raise ValueError(f"Unknown benchmark: {name}")


# ─── Main scoring loop ────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--benchmark", required=True,
                     choices=["adas_to", "sim_dataset", "longdrive",
                                "kaggle_accident"])
    ap.add_argument("--skip_existing", action="store_true",
                     help="skip clips whose JSON already has m10_v3 field")
    ap.add_argument("--qwen3_base", type=Path, default=DEFAULT_QWEN3_BASE)
    ap.add_argument("--qwen3_lora", type=Path, default=DEFAULT_QWEN3_LORA)
    ap.add_argument("--m10_v3_head", type=Path, default=DEFAULT_M10_V3_HEAD)
    ap.add_argument("--frame_size", type=int, default=448)
    ap.add_argument("--tick_seconds", type=float, default=1.0)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--limit", type=int, default=0,
                     help="smoke-test: only score first N clips")
    args = ap.parse_args()

    cfg = get_benchmark_config(args.benchmark, args)
    clips = cfg["clips"]
    if args.limit:
        clips = clips[:args.limit]
    print(f"[score] benchmark={args.benchmark}  n_clips={len(clips)}")
    if not clips:
        print(f"[error] no clips found for {args.benchmark}", file=sys.stderr)
        return 2

    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    if device.type != "cuda":
        print("[warn] CUDA unavailable; will be slow", file=sys.stderr)

    model = load_v3_m10(device, args.qwen3_base, args.qwen3_lora,
                          args.m10_v3_head)

    n_total = len(clips)
    n_done = 0
    n_skipped = 0
    t_start = time.time()

    for cid, video, jp in clips:
        # Load existing JSON or create fresh
        if jp.exists():
            scores_data = json.loads(jp.read_text())
        elif cfg["create_jsons"]:
            scores_data = {}
        else:
            print(f"  [skip] no JSON at {jp}")
            continue

        # Skip if already has m10_v3 field
        if args.skip_existing and scores_data.get(cfg["append_field"]):
            n_skipped += 1
            continue

        # Determine ticks: reuse existing m10_v2 / pomdp_v3 ticks if present;
        # else build from video metadata
        ticks = []
        for src_field in ("m10_v2", "pomdp_v3"):
            if scores_data.get(src_field):
                ticks = [t["frame"] for t in scores_data[src_field]]
                break
        cap = cv2.VideoCapture(str(video))
        fps = cap.get(cv2.CAP_PROP_FPS) or 20.0
        n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        cap.release()
        if not ticks:
            tick_frames = max(1, int(round(args.tick_seconds * fps)))
            ticks = list(range(tick_frames - 1, n_frames, tick_frames))
            scores_data["fps"] = fps
            scores_data["n_total"] = n_frames
            scores_data["tick_frames"] = tick_frames

        window_frames = max(8, int(round(4.0 * fps)))

        t0 = time.time()
        scores_m10v3 = score_one_clip(model, video, ticks, window_frames,
                                          n_sample=8, frame_size=args.frame_size)
        scores_data[cfg["append_field"]] = scores_m10v3
        jp.parent.mkdir(parents=True, exist_ok=True)
        jp.write_text(json.dumps(scores_data))
        n_done += 1
        elapsed = time.time() - t0
        total = time.time() - t_start
        eta_min = (total / n_done) * (n_total - n_done - n_skipped) / 60.0
        print(f"  [{n_done + n_skipped:>4}/{n_total}] {cid[:50]:<50}  "
              f"ticks={len(ticks)}  {elapsed:.1f}s  ETA {eta_min:.1f}min",
              flush=True)

    wall = (time.time() - t_start) / 60.0
    print(f"\n[done] benchmark={args.benchmark}  scored={n_done}  "
          f"skipped={n_skipped}  total_time={wall:.1f}min")
    return 0


if __name__ == "__main__":
    sys.exit(main())