#!/usr/bin/env python """Generate the VBench-I2V-100 videos for the LingBot Layer-17 Predictor strategies (the analogue of eval/generate_predictor_eval.py for Causal-Forcing-a). 100 jobs = 25 validation cases x 4 camera actions, 125 frames = 8 chunks, seed 0. Chunk 0 always runs the full 4-step grid; on chunks 1..7 the pattern decides which steps a Predictor serves (FPFF / FPPF / FPPP) or which steps are visited at all (FPF: steps 0, 1, 3 with one Predictor step at t=937 anchored on t=999). A Predictor step runs 1 of the 30 Teacher blocks and is booked as 1/30 of a full forward. FFFF is generated in the same process for every job (matched-FFFF latency, pre-MP4 reference frames) but lb_ffff's own records are never rewritten. CUDA_VISIBLE_DEVICES=0 python lb_generate_predictor_eval.py --models atc_chunk --patterns FPPF,FPF,FPPP --shard 0 --num-shards 4 """ import argparse, gc, json, os, sys, time ROOT = os.path.dirname(os.path.abspath(__file__)) LBA = "/local/zoubin/cz/projects/lingbot-world-v2-a" for p in (ROOT, LBA): if p not in sys.path: sys.path.insert(0, p) from lbcache import harness as H H.enter_repo() import torch from PIL import Image from safetensors.torch import load_file from lbcache.pipeline import DEFAULT_TIMESTEPS_INDEX, prepare from eval.pixel_metrics import PixelMetrics from predictor_training.lb_data import FRAMES_PER_CHUNK, NUM_CHUNKS, NUM_STEPS, TOKENS_PER_CHUNK from predictor_training.metadata import build_single_block_predictor, read_predictor_config from predictor_training.online import online_predictor_step MODELS = { "atc_chunk": "output/layer17_stage1_atc_chunk_100c4a_125f_seed0_4gpu_b8_acc2_2000steps", "disca": "output/layer17_stage1_disca_100c4a_125f_seed0_4gpu_b8_acc2_2000steps", } PATTERNS = {"FPFF": {1}, "FPPF": {1, 2}, "FPPP": {1, 2, 3}, "FPF": {1}} RUN_STEPS = {"FPF": [0, 1, 3]} PREDICTOR_COMPUTE_EQUIV = 1.0 / 30.0 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("--models", default="atc_chunk,disca") ap.add_argument("--patterns", default="FPPF,FPF,FPPP") ap.add_argument("--checkpoint-step", type=int, default=2000) ap.add_argument("--tag", default="", help="suffix for another checkpoint of the same model, e.g. _s1000") 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("--overwrite", action="store_true") 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 = [{"name": "lb_ffff", "model": None, "pattern": "FFFF", "is_reference": True}] for model in args.models.split(","): for pattern in args.patterns.split(","): strategies.append({"name": f"lb_{model}{args.tag}_{pattern.lower()}", "model": model, "pattern": pattern.upper(), "is_reference": False}) ref_name = strategies[0]["name"] print(f"shard {args.shard}/{args.num_shards}: {len(shard_jobs)} jobs x {len(strategies)} strategies", flush=True) device = torch.device("cuda") pipe = H.load_pipeline() teacher = pipe.model metrics = PixelMetrics(device) predictors, predictor_configs = {}, {} for model in args.models.split(","): path = MODELS[model] path = path if os.path.isabs(path) else os.path.join(LBA, path) weights = path if path.endswith(".safetensors") else os.path.join(path, f"checkpoint_step_{args.checkpoint_step:04d}", "predictor.safetensors") config = read_predictor_config(weights) predictor = build_single_block_predictor(teacher, config, gradient_checkpointing=False) predictor.load_state_dict(load_file(weights, device="cpu"), strict=True) predictor.to(device).eval().requires_grad_(False) predictors[model], predictor_configs[model] = predictor, config print(f"[predictor] {model}: {weights} ({config})", flush=True) class FinalHiddenCapture: def __init__(self, model): self.enabled, self.value = False, None model.head.register_forward_pre_hook(self._hook) def _hook(self, _m, inputs): if self.enabled: self.value = inputs[0].detach() def start(self): self.value, self.enabled = None, True def finish(self): self.enabled = False if self.value is None: raise RuntimeError("Teacher final hidden was not captured") v, self.value = self.value, None return v capture = FinalHiddenCapture(teacher) @torch.no_grad() def rollout(case, action, predictor=None, predictor_steps=frozenset(), run_steps=None): img = Image.open(case["image_path"]).convert("RGB") torch.manual_seed(args.seed); torch.cuda.manual_seed_all(args.seed) prep = prepare(pipe, case["caption"], img, H.action_path(action, args.frames), chunk_size=FRAMES_PER_CHUNK, max_area=480 * 832 + 1, frame_num=args.frames, shift=5.0, seed=args.seed) scheduler = pipe.scheduler ts = scheduler.timesteps[list(DEFAULT_TIMESTEPS_INDEX)] pipe._cross_attn_initialized = False denoise_events, context_events = [], [] full_calls = predictor_calls = 0 previous_hidden = None preds = [] with torch.amp.autocast('cuda', dtype=pipe.param_dtype): lat = prep["noise"].split(FRAMES_PER_CHUNK, dim=1); cond = prep["y"].split(FRAMES_PER_CHUNK, dim=1) plk = prep["plucker"].split(FRAMES_PER_CHUNK, dim=2) for chunk in range(len(lat)): cur = lat[chunk] kw = {'context': [prep["context"][0]], 'seq_len': prep["max_seq_len"], 'y': [cond[chunk]], 'dit_cond_dict': {"c2ws_plucker_emb": plk[chunk].chunk(1, dim=0)}, 'kv_cache': prep["self_kv_cache"], 'crossattn_cache': prep["cross_kv_cache"], 'current_start': chunk * TOKENS_PER_CHUNK, 'max_attention_size': prep["kv_size"], 'frame_seqlen': prep["frame_seqlen"]} current_hidden = [None] * NUM_STEPS steps = list(range(NUM_STEPS)) if (run_steps is None or chunk == 0) else list(run_steps) for si, step in enumerate(steps): t = torch.stack([ts[step]]).to(device) s_ev = torch.cuda.Event(enable_timing=True); e_ev = torch.cuda.Event(enable_timing=True) s_ev.record() if predictor is not None and chunk > 0 and step in predictor_steps: anchor_step = steps[si - 1] hidden, flow = online_predictor_step( predictor=predictor, teacher=teacher, noisy_input=cur[None], y=cond[chunk][None], plucker=plk[chunk], timestep=t.float(), anchor_timestep=torch.stack([ts[anchor_step]]).to(device).float(), anchor_hidden=current_hidden[anchor_step], previous_hidden=previous_hidden[step], history_cache=prep["self_kv_cache"][predictor.source_layer], cross_cache=prep["cross_kv_cache"][predictor.source_layer], chunk=chunk) noise_pred = flow[0].float() current_hidden[step] = hidden predictor_calls += 1 else: capture.start() noise_pred = pipe.model(x=[cur.to(device)], t=t, cross_attn_first_call=not pipe._cross_attn_initialized, **kw)[0] current_hidden[step] = capture.finish() pipe._cross_attn_initialized = True full_calls += 1 e_ev.record() denoise_events.append((s_ev, e_ev)) x0 = pipe._convert_flow_pred_to_x0(flow_pred=noise_pred, xt=cur, timestep=ts[step], scheduler=scheduler) if si < len(steps) - 1: cur = scheduler.add_noise(x0, torch.randn(x0.shape, generator=prep["seed_g"], device=x0.device, dtype=x0.dtype), ts[steps[si + 1]]) preds.append(x0) s_ev = torch.cuda.Event(enable_timing=True); e_ev = torch.cuda.Event(enable_timing=True) s_ev.record() pipe.model(x=[x0], t=torch.stack([scheduler.timesteps[-1] * 0.0]).to(device), cross_attn_first_call=False, **kw) e_ev.record() context_events.append((s_ev, e_ev)) previous_hidden = current_hidden out = pipe.vae.decode([torch.cat(preds, dim=1)])[0] torch.cuda.synchronize() timing = {"denoise_dit_ms": sum(s.elapsed_time(e) for s, e in denoise_events), "context_kv_dit_ms": sum(s.elapsed_time(e) for s, e in context_events), "num_denoise_forwards": len(denoise_events), "text_cache_ms": 0.0} video = out.float().clamp(-1, 1).add(1).div(2).unsqueeze(0) # [1, C, T, H, W] in [0, 1] return video, timing, {"full_forwards": full_calls, "predictor_forwards": predictor_calls} 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, _ = 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): gc.collect(); torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats() wall0 = time.time() if strategy["is_reference"]: video, timing, counts = rollout(case, H.ACTIONS[ai]) else: video, timing, counts = rollout(case, H.ACTIONS[ai], predictors[strategy["model"]], PATTERNS[strategy["pattern"]], RUN_STEPS.get(strategy["pattern"])) timing["wall_s"] = time.time() - wall0 timing["peak_mem_gib"] = torch.cuda.max_memory_allocated() / 1024 ** 3 return video, timing, counts if shard_jobs: for s in strategies: run_one(s, *shard_jobs[0]) torch.cuda.empty_cache() t_start = time.time() n_chunks = ((args.frames - 1) // 4 + 1) // FRAMES_PER_CHUNK for n, (case, ai) in enumerate(shard_jobs): todo = [s for s in strategies if not s["is_reference"] and not is_done(s["name"], case, ai)] if not todo: continue ref_video, ref_timing, _ = run_one(strategies[0], case, ai) ref_frames = ref_video[0].permute(1, 0, 2, 3).contiguous() for s in todo: video, timing, counts = 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) H.save_video(video, vp) compute_equiv = counts["full_forwards"] + PREDICTOR_COMPUTE_EQUIV * counts["predictor_forwards"] 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": f"predictor_{s['model']}", "param": "pattern", "param_value": s["pattern"], "target_speedup": (n_chunks * NUM_STEPS) / compute_equiv, "schedule": s["pattern"], "num_inference_steps": 4, "first_chunk_schedule": "FFFF", "predictor_config": predictor_configs[s["model"]], "checkpoint_step": args.checkpoint_step, "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": 0.0, "matched_ffff_policy_latency_ms": ref_timing["denoise_dit_ms"], "latency_source": "generation run, paired with FFFF in the same process", "reference_source": "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": {"denoise_forwards": timing["num_denoise_forwards"], "full_forwards": counts["full_forwards"], "predictor_forwards": counts["predictor_forwards"], "predictor_compute_equivalent": PREDICTOR_COMPUTE_EQUIV, "compute_equivalent_forwards": compute_equiv, "middle_steps": (n_chunks - 1) * 2, "middle_compute_equivalent": compute_equiv - 4 - (n_chunks - 1) * 2.0}, }) del video, frames del ref_video, ref_frames done = n + 1 rate = (time.time() - t_start) / done print(f"[lb-pred 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)