#!/usr/bin/env python """Generate the VBench-I2V-100 evaluation videos on the LingBot-World-V2 1.3B causal-fast base. 100 jobs = 25 validation cases x 4 camera actions, 125 frames, 480x832, seed 0. For every job the FFFF reference is generated first and every cache strategy is timed back to back with it in the same process (the protocol's matched-FFFF latency), scored against its decoded frames (PSNR/SSIM/LPIPS, all 125 frames), and written as mp4 (fps 16, the model's native rate). CUDA_VISIBLE_DEVICES=5 python lb_generate_eval.py --shard 0 --num-shards 2 """ import argparse, gc, json, os, sys, time sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from lbcache import harness as H H.enter_repo() import torch from lbcache.measure import make_ctrl from lb_strategies import load_strategies from eval.pixel_metrics import PixelMetrics ROOT = os.path.dirname(os.path.abspath(__file__)) ap = argparse.ArgumentParser() ap.add_argument("--shard", type=int, default=0) ap.add_argument("--num-shards", type=int, default=1) ap.add_argument("--out-root", default="eval_out_lb") ap.add_argument("--frames", type=int, default=125) ap.add_argument("--seed", type=int, default=0) ap.add_argument("--limit", type=int, default=None) ap.add_argument("--only-strategy", default=None) ap.add_argument("--strategies", default=None, help="Comma-separated subset (FFFF always kept)") ap.add_argument("--overwrite", action="store_true") ap.add_argument("--regenerate-ffff", action="store_true", help="Regenerate the FFFF reference instead of reading lb_ffff's MP4") args = ap.parse_args() out_root = args.out_root if os.path.isabs(args.out_root) else os.path.join(ROOT, args.out_root) cases = H.load_cases() jobs = [(c, ai) for c in cases for ai in range(len(H.ACTIONS))] if args.limit: jobs = jobs[:args.limit] shard_jobs = jobs[args.shard::args.num_shards] strategies = load_strategies() if args.only_strategy: strategies = [s for s in strategies if s["is_reference"] or s["name"] == args.only_strategy] if args.strategies: keep = set(args.strategies.split(",")) strategies = [s for s in strategies if s["is_reference"] or s["name"] in keep] ref_name = strategies[0]["name"] print(f"shard {args.shard}/{args.num_shards}: {len(shard_jobs)} jobs x {len(strategies)} strategies", flush=True) pipe = H.load_pipeline() metrics = PixelMetrics(torch.device("cuda")) save_video = H.save_video def paths(strategy, case, ai): stem = f"case_{case['split_case_id']:04d}_action_{ai:02d}" d = os.path.join(out_root, "per_prompt", strategy); os.makedirs(d, exist_ok=True) v = os.path.join(out_root, "generated_videos", strategy); os.makedirs(v, exist_ok=True) return os.path.join(d, stem + ".json"), os.path.join(v, stem + ".mp4") def is_done(strategy, case, ai): rp, vp = paths(strategy, case, ai) if args.overwrite or not os.path.exists(rp): return False try: rec = json.load(open(rp)) return rec.get("status") == "complete" and os.path.exists(rec.get("video", "")) except Exception: return False def write_json(path, payload): tmp = path + ".tmp" json.dump(payload, open(tmp, "w"), indent=2) os.replace(tmp, path) def run_one(strategy, case, ai): kw = {strategy["param"]: strategy["value"]} if strategy.get("param") else {} n_steps = strategy.get("num_inference_steps") or H.NUM_STEPS first = strategy.get("first_chunk_schedule") ctrl = make_ctrl(pipe, strategy["method"], schedule=strategy["schedule"], coefficients=strategy["coefficients"], num_steps=n_steps, first_chunk_schedule=first, **kw) gc.collect(); torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats() video, _, timing = H.generate(pipe, case, H.ACTIONS[ai], controller=ctrl, num_frames=args.frames, seed=args.seed, num_steps=n_steps, first_chunk_steps=(len(first) if first and len(first) != n_steps else None)) return video, timing, ctrl.summary() # warm-up (TaylorSeer's compiled forecast, allocator) -- not recorded if shard_jobs: for s in strategies: run_one(s, *shard_jobs[0]) torch.cuda.empty_cache() t_start = time.time() for n, (case, ai) in enumerate(shard_jobs): todo = [s for s in strategies if not is_done(s["name"], case, ai)] if not todo: continue ref_from_mp4 = False ref_rp, ref_vp = paths(ref_name, case, ai) if (not args.regenerate_ffff and os.path.exists(ref_vp) and not any(s["is_reference"] for s in todo)): # Reference from the lb_ffff run's MP4: FFFF is not regenerated, so there is # no same-process FFFF timing (lb_aggregate falls back to lb_ffff's records). from torchvision.io import read_video vid, _, _ = read_video(ref_vp, pts_unit="sec", output_format="TCHW") ref_frames = (vid.to(torch.device("cuda"), torch.float16) / 255.0) ref_video, ref_timing, ref_summary, ref_from_mp4 = None, None, None, True else: ref_video, ref_timing, ref_summary = run_one(strategies[0], case, ai) ref_frames = ref_video[0].permute(1, 0, 2, 3).contiguous() # [T, C, H, W] in [0,1] for s in strategies: if s not in todo: continue if s["is_reference"]: video, timing, summary = ref_video, ref_timing, ref_summary pix = PixelMetrics.identity(ref_frames.shape[0]) else: video, timing, summary = run_one(s, case, ai) frames = video[0].permute(1, 0, 2, 3).contiguous() pix = metrics.compute(frames, ref_frames.to(torch.float16)) rp, vp = paths(s["name"], case, ai) save_video(video, vp) write_json(rp, { "status": "complete", "protocol": "LingBot-World-V2 VBench-I2V-100 (validation25 x 4 actions)", "strategy": s["name"], "base_model": "lingbot_world_v2_1.3b_causal_fast", "method": s["method"], "param": s["param"], "param_value": s["value"], "target_speedup": s["target"], "schedule": s["schedule"], "num_inference_steps": s.get("num_inference_steps") or 4, "first_chunk_schedule": s.get("first_chunk_schedule"), "case_id": case["case_id"], "split_case_id": case["split_case_id"], "action_id": ai, "action_name": H.ACTIONS[ai][0], "prompt": case["caption"], "image_path": case["image_path"], "seed": args.seed, "video": vp, "num_frames": int(video.shape[2]), "height": int(video.shape[3]), "width": int(video.shape[4]), "fps": H.FPS, "policy_latency_ms": timing["denoise_dit_ms"], "excluded_context_kv_latency_ms": timing["context_kv_dit_ms"], "excluded_text_cache_ms": timing["text_cache_ms"], "matched_ffff_policy_latency_ms": ref_timing["denoise_dit_ms"] if ref_timing else None, "latency_source": ("generation run (FFFF from its own records)" if ref_from_mp4 else "generation run, paired with FFFF in the same process"), "reference_source": "ffff_mp4" if ref_from_mp4 else "ffff_frames", "wall_generation_s": timing["wall_s"], "peak_mem_gib": timing["peak_mem_gib"], "reference_strategy": ref_name, "pixel_metrics_vs_ffff": pix, "cache_diagnostics": dict(summary), }) del video del ref_video, ref_frames done = n + 1 rate = (time.time() - t_start) / done print(f"[lb shard {args.shard}] {done}/{len(shard_jobs)} jobs ({case['split_case_id']:04d}/{ai}) " f"{rate:.0f}s/job eta {(len(shard_jobs) - done) * rate / 60:.0f} min", flush=True) print(f"shard {args.shard} done in {(time.time() - t_start) / 60:.1f} min", flush=True)