lb_generate_eval.py

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