Buckets:
| """Frozen Qwen3-VL-2B-Instruct feature extractor for the AIDE protocol frames (Driver-WM App. B.2). | |
| # smoke test (20 frames, timing + peak memory -> results/feats/qwen3vl2b_bench.json) | |
| flock /tmp/s13_gpu.lock python aide/features/extract_qwen3vl.py --smoke 20 | |
| # chunked full run (exit code 3 = time budget exhausted, work remains; 0 = all done) | |
| flock /tmp/s13_gpu.lock python aide/features/extract_qwen3vl.py --minutes 20 | |
| For every clip and each of the 4 views the 10 protocol frames (<view>frames/<frame_idx>.jpg, | |
| frame_idx = cache frame_idx10) are fed one image at a time (batched across images) with the | |
| verbatim Driver-WM prompt (protocol.QWEN_PROMPT), bf16, and the LAST layer hidden states are | |
| mean-pooled over the full token sequence -> 2048-d ("pooled") and over the vision tokens only | |
| ("vis_pooled"). Output per clip: results/feats/qwen3vl2b/<id>.npz | |
| pooled [10,4,2048] fp16, vis_pooled [10,4,2048] fp16, views [4], frame_idx [10], | |
| n_tokens [10,4], n_vis_tokens [10,4], meta json string. | |
| Implementation note: the last-layer hidden states are read as `model.model(...).last_hidden_state` | |
| (= output_hidden_states[-1] after the final RMSNorm) so the 151k-vocab logits are never | |
| materialised. Resumable (existing npz skipped); processing order test -> val -> train, by id. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import os | |
| import sys | |
| import time | |
| from concurrent.futures import ThreadPoolExecutor | |
| from pathlib import Path | |
| import numpy as np | |
| ROOT = Path(__file__).resolve().parents[2] | |
| sys.path.insert(0, str(ROOT)) | |
| from aide.data import protocol as P # noqa: E402 | |
| MODEL_ID = "Qwen/Qwen3-VL-2B-Instruct" | |
| OUT_DIR = P.RESULTS / "feats" / "qwen3vl2b" | |
| BENCH_JSON = P.RESULTS / "feats" / "qwen3vl2b_bench.json" | |
| EXIT_TIME_UP = 3 | |
| def clip_order(cache): | |
| order = [] | |
| for split in ("test", "val", "train"): | |
| ids = sorted(cache["clip_id"][cache["split"] == split].tolist()) | |
| order += ids | |
| return order | |
| def image_paths(cid: str, frame_idx): | |
| return [(fi, v, P.AIDE_ROOT / cid / f"{v}frames/{int(f)}.jpg") | |
| for fi, f in enumerate(frame_idx) for v in P.VIEWS] | |
| def load_images(paths): | |
| from PIL import Image | |
| ims = [] | |
| for _, _, p in paths: | |
| im = Image.open(p) | |
| im.load() | |
| ims.append(im.convert("RGB")) | |
| return ims | |
| class Extractor: | |
| def __init__(self, batch=8, max_pixels=None, min_pixels=None, add_generation_prompt=True): | |
| import torch | |
| from transformers import AutoProcessor, Qwen3VLForConditionalGeneration | |
| self.torch = torch | |
| self.batch = batch | |
| self.processor = AutoProcessor.from_pretrained(MODEL_ID) | |
| ip = self.processor.image_processor | |
| if max_pixels or min_pixels: | |
| size = dict(ip.size) | |
| if max_pixels: | |
| size["longest_edge"] = int(max_pixels) | |
| if min_pixels: | |
| size["shortest_edge"] = int(min_pixels) | |
| ip.size = size | |
| self.image_size_cfg = dict(ip.size) | |
| self.model = Qwen3VLForConditionalGeneration.from_pretrained( | |
| MODEL_ID, dtype=torch.bfloat16, attn_implementation="sdpa", device_map="cuda").eval() | |
| for p in self.model.parameters(): | |
| p.requires_grad_(False) | |
| self.image_token_id = self.model.config.image_token_id | |
| self.add_generation_prompt = add_generation_prompt | |
| msgs = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": P.QWEN_PROMPT}]}] | |
| self.text = self.processor.apply_chat_template(msgs, tokenize=False, | |
| add_generation_prompt=add_generation_prompt) | |
| self.hidden = self.model.config.text_config.hidden_size | |
| self.meta = { | |
| "model_id": MODEL_ID, "dtype": "bfloat16", "attn": "sdpa", "batch": batch, | |
| "image_size_cfg": self.image_size_cfg, "add_generation_prompt": add_generation_prompt, | |
| "prompt_sha256": hashlib.sha256(P.QWEN_PROMPT.encode()).hexdigest(), "prompt": P.QWEN_PROMPT, | |
| "pooling": "mean over attention_mask==1 tokens of model.model(...).last_hidden_state; " | |
| "vis_pooled = mean over input_ids==image_token_id positions", | |
| "hidden": self.hidden, "template_text": self.text, | |
| } | |
| def _masked_mean(h, mask): | |
| m = mask.unsqueeze(-1).to(h.dtype) | |
| return (h * m).sum(1) / m.sum(1).clamp(min=1) | |
| def encode(self, images): | |
| """images: list of PIL -> pooled [n,2048] f16, vis_pooled [n,2048] f16, n_tok [n], n_vis [n].""" | |
| torch = self.torch | |
| pooled, vis, ntok, nvis = [], [], [], [] | |
| for i in range(0, len(images), self.batch): | |
| ims = images[i:i + self.batch] | |
| inputs = self.processor(text=[self.text] * len(ims), images=ims, return_tensors="pt", padding=True) | |
| inputs = {k: v.to("cuda") for k, v in inputs.items()} | |
| with torch.inference_mode(): | |
| out = self.model.model(**inputs) | |
| h = out.last_hidden_state # [B,L,2048] | |
| am = inputs["attention_mask"].bool() | |
| vm = (inputs["input_ids"] == self.image_token_id) & am | |
| pooled.append(self._masked_mean(h, am).to(torch.float16).cpu().numpy()) | |
| vis.append(self._masked_mean(h, vm).to(torch.float16).cpu().numpy()) | |
| ntok.append(am.sum(1).cpu().numpy()) | |
| nvis.append(vm.sum(1).cpu().numpy()) | |
| return np.concatenate(pooled), np.concatenate(vis), np.concatenate(ntok), np.concatenate(nvis) | |
| def save_clip(out_path, cid, frame_idx, pooled, vis, ntok, nvis, meta): | |
| nf, nv = len(frame_idx), len(P.VIEWS) | |
| arr = dict( | |
| pooled=pooled.reshape(nf, nv, -1).astype(np.float16), | |
| vis_pooled=vis.reshape(nf, nv, -1).astype(np.float16), | |
| n_tokens=ntok.reshape(nf, nv).astype(np.int32), n_vis_tokens=nvis.reshape(nf, nv).astype(np.int32), | |
| views=np.array(P.VIEWS), frame_idx=np.asarray(frame_idx, dtype=np.int32), | |
| clip_id=np.array(cid), meta=np.array(json.dumps(meta)), | |
| ) | |
| tmp = out_path.with_suffix(".tmp.npz") | |
| np.savez(tmp, **arr) | |
| os.replace(tmp, out_path) | |
| def smoke(ex: Extractor, cache, n_frames: int): | |
| import torch | |
| order = clip_order(cache) | |
| cid = order[0] | |
| row = int(np.flatnonzero(cache["clip_id"] == cid)[0]) | |
| paths = image_paths(cid, cache["frame_idx10"][row])[:n_frames] | |
| ims = load_images(paths) | |
| # warm-up (first batch has compile/alloc overhead) | |
| torch.cuda.synchronize() | |
| ex.encode(ims[:ex.batch]) | |
| torch.cuda.synchronize() | |
| torch.cuda.reset_peak_memory_stats() | |
| t0 = time.time() | |
| pooled, vis, ntok, nvis = ex.encode(ims) | |
| torch.cuda.synchronize() | |
| dt = time.time() - t0 | |
| rec = { | |
| "utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "clip_id": cid, "n_frames": len(ims), | |
| "batch": ex.batch, "ms_per_frame": 1000 * dt / len(ims), "frames_per_s": len(ims) / dt, | |
| "peak_mem_allocated_GB": torch.cuda.max_memory_allocated() / 1e9, | |
| "peak_mem_reserved_GB": torch.cuda.max_memory_reserved() / 1e9, | |
| "n_tokens_per_image": {"min": int(ntok.min()), "max": int(ntok.max())}, | |
| "n_vis_tokens_per_image": {"min": int(nvis.min()), "max": int(nvis.max())}, | |
| "image_wh": list(ims[0].size), "pooled_norm_mean": float(np.linalg.norm(pooled.astype(np.float32), axis=1).mean()), | |
| "vis_pooled_norm_mean": float(np.linalg.norm(vis.astype(np.float32), axis=1).mean()), | |
| "pooled_finite": bool(np.isfinite(pooled).all() and np.isfinite(vis).all()), | |
| "est_total_frames": int(len(cache["clip_id"]) * P.N_SAMPLE * len(P.VIEWS)), | |
| "gpu": torch.cuda.get_device_name(0), "meta": ex.meta, | |
| } | |
| rec["est_total_gpu_hours"] = rec["est_total_frames"] * dt / len(ims) / 3600 | |
| BENCH_JSON.parent.mkdir(parents=True, exist_ok=True) | |
| recs = json.loads(BENCH_JSON.read_text()) if BENCH_JSON.exists() else [] | |
| recs.append(rec) | |
| BENCH_JSON.write_text(json.dumps(recs, indent=1)) | |
| np.savez(BENCH_JSON.parent / "qwen3vl2b_smoke.npz", pooled=pooled, vis_pooled=vis, | |
| paths=np.array([str(p) for _, _, p in paths])) | |
| print(json.dumps({k: v for k, v in rec.items() if k != "meta"}, indent=1)) | |
| def profile(ex: Extractor, cache, n_images=8, reps=3): | |
| """Per-image time breakdown (jpeg decode / processor / full forward / vision tower) -> bench json.""" | |
| import torch | |
| cid = clip_order(cache)[1] | |
| row = int(np.flatnonzero(cache["clip_id"] == cid)[0]) | |
| paths = image_paths(cid, cache["frame_idx10"][row])[:n_images] | |
| t = time.time() | |
| ims = load_images(paths) | |
| t_load = (time.time() - t) / len(ims) | |
| ex.encode(ims) # warm-up | |
| torch.cuda.synchronize() | |
| t = time.time() | |
| for _ in range(reps): | |
| inputs = ex.processor(text=[ex.text] * len(ims), images=ims, return_tensors="pt", padding=True) | |
| t_pre = (time.time() - t) / reps / len(ims) | |
| inputs = {k: v.to("cuda") for k, v in inputs.items()} | |
| with torch.inference_mode(): | |
| torch.cuda.synchronize(); t = time.time() | |
| for _ in range(reps): | |
| ex.model.model(**inputs) | |
| torch.cuda.synchronize(); t_fwd = (time.time() - t) / reps / len(ims) | |
| torch.cuda.synchronize(); t = time.time() | |
| for _ in range(reps): | |
| ex.model.model.visual(inputs["pixel_values"], grid_thw=inputs["image_grid_thw"]) | |
| torch.cuda.synchronize(); t_vis = (time.time() - t) / reps / len(ims) | |
| rec = {"utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "kind": "profile", "clip_id": cid, | |
| "n_images": len(ims), "reps": reps, "batch": ex.batch, | |
| "ms_jpeg_decode": 1000 * t_load, "ms_processor_cpu": 1000 * t_pre, | |
| "ms_forward_gpu": 1000 * t_fwd, "ms_vision_tower_gpu": 1000 * t_vis, | |
| "pixel_values_shape": list(inputs["pixel_values"].shape), "torch_threads": torch.get_num_threads()} | |
| recs = json.loads(BENCH_JSON.read_text()) if BENCH_JSON.exists() else [] | |
| recs.append(rec) | |
| BENCH_JSON.write_text(json.dumps(recs, indent=1)) | |
| print(json.dumps(rec, indent=1)) | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--minutes", type=float, default=20.0, help="wall-clock budget; exit 3 when exhausted") | |
| ap.add_argument("--batch", type=int, default=8) | |
| ap.add_argument("--smoke", type=int, default=0, help="run a timing smoke test on N frames and exit") | |
| ap.add_argument("--profile", action="store_true", help="per-image time breakdown -> bench json, then exit") | |
| ap.add_argument("--max-pixels", type=int, default=None) | |
| ap.add_argument("--min-pixels", type=int, default=None) | |
| ap.add_argument("--no-generation-prompt", action="store_true") | |
| ap.add_argument("--limit", type=int, default=0, help="stop after N clips (debug)") | |
| a = ap.parse_args() | |
| t0 = time.time() | |
| cache = np.load(P.CACHE_NPZ, allow_pickle=False) | |
| ex = Extractor(batch=a.batch, max_pixels=a.max_pixels, min_pixels=a.min_pixels, | |
| add_generation_prompt=not a.no_generation_prompt) | |
| print(f"model loaded in {time.time() - t0:.1f}s; image size cfg {ex.image_size_cfg}", flush=True) | |
| if a.smoke: | |
| smoke(ex, cache, a.smoke) | |
| return 0 | |
| if a.profile: | |
| profile(ex, cache, n_images=a.batch) | |
| return 0 | |
| OUT_DIR.mkdir(parents=True, exist_ok=True) | |
| order = clip_order(cache) | |
| pos = {c: i for i, c in enumerate(cache["clip_id"].tolist())} | |
| todo = [c for c in order if not (OUT_DIR / f"{c}.npz").exists()] | |
| print(f"{len(order) - len(todo)} done, {len(todo)} to do", flush=True) | |
| if not todo: | |
| return 0 | |
| pool = ThreadPoolExecutor(4) | |
| fut = pool.submit(load_images, image_paths(todo[0], cache["frame_idx10"][pos[todo[0]]])) | |
| n_done, t_loop = 0, time.time() | |
| for k, cid in enumerate(todo): | |
| ims = fut.result() | |
| if k + 1 < len(todo): | |
| nxt = todo[k + 1] | |
| fut = pool.submit(load_images, image_paths(nxt, cache["frame_idx10"][pos[nxt]])) | |
| fi = cache["frame_idx10"][pos[cid]] | |
| pooled, vis, ntok, nvis = ex.encode(ims) | |
| assert np.isfinite(pooled).all() and np.isfinite(vis).all(), cid | |
| save_clip(OUT_DIR / f"{cid}.npz", cid, fi, pooled, vis, ntok, nvis, ex.meta) | |
| n_done += 1 | |
| if n_done % 10 == 0: | |
| rate = (time.time() - t_loop) / n_done | |
| print(f"[{n_done}/{len(todo)}] {cid} {rate:.1f}s/clip, remaining ~{rate * (len(todo) - n_done) / 60:.0f} min", | |
| flush=True) | |
| if a.limit and n_done >= a.limit: | |
| break | |
| if time.time() - t0 > a.minutes * 60 and k + 1 < len(todo): | |
| print(f"time budget reached after {n_done} clips ({(time.time() - t0) / 60:.1f} min); exit 3", flush=True) | |
| pool.shutdown(wait=False, cancel_futures=True) | |
| return EXIT_TIME_UP | |
| pool.shutdown(wait=False, cancel_futures=True) | |
| remaining = len([c for c in order if not (OUT_DIR / f"{c}.npz").exists()]) | |
| print(f"done {n_done} clips this chunk; remaining {remaining}") | |
| return 0 if remaining == 0 else EXIT_TIME_UP | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |
Xet Storage Details
- Size:
- 13.2 kB
- Xet hash:
- 2c79843c84190260774a2fc01f18817d906fb49449350b15d58b9aa8c7e18493
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.