lb_generate_predictor_eval.py

#!/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)