| |
| """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) |
| 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) |
|
|
