eval/generate_eval.py

#!/usr/bin/env python
"""Generate the Extended-251 videos for every strategy of one base model.

Work is sharded by prompt, not by strategy, because the protocol requires each
strategy's clip to be paired with the FFFF clip for the *same* prompt and seed
(section 8.1). A shard therefore generates FFFF first, keeps its decoded frames,
and immediately scores every cache strategy against them -- no 25 GB of reference
frames on disk, and the pairing cannot get mismatched.

    CUDA_VISIBLE_DEVICES=0 python eval/generate_eval.py \
        --base self_forcing --shard 0 --num-shards 4 --out-root eval_out

Re-running skips prompts whose per-strategy record is already complete.
"""

import argparse
import json
import os
import sys
import time

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

from eval.strategies import load_strategies  # noqa: E402
from harness import denoising_steps_for, BASES, enter_base, load_pipeline, use_denoising_steps  # noqa: E402

MAPPING = os.path.join(ROOT, "assets/vbench8_extended_subset_mapping.json")


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--base", choices=sorted(BASES), required=True)
    p.add_argument("--shard", type=int, default=0)
    p.add_argument("--num-shards", type=int, default=1)
    p.add_argument("--out-root", default="eval_out")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--num-latent-frames", type=int, default=21)
    p.add_argument("--fps", type=int, default=16)
    p.add_argument("--limit", type=int, default=None, help="Smoke-test: first N prompts")
    p.add_argument("--only-strategy", default=None, help="Smoke-test: one strategy name")
    p.add_argument("--strategies", default=None,
                   help="Comma-separated subset to generate (the FFFF reference is always kept)")
    p.add_argument("--overwrite", action="store_true")
    p.add_argument("--regenerate-ffff", action="store_true",
                   help="Regenerate the FFFF reference instead of reading its MP4")
    return p.parse_args()


def main():
    args = parse_args()
    out_root = (args.out_root if os.path.isabs(args.out_root)
                else os.path.join(ROOT, args.out_root))

    with open(MAPPING) as f:
        mapping = json.load(f)
    rows = mapping["rows"]
    if args.limit:
        rows = rows[:args.limit]
    shard_rows = rows[args.shard::args.num_shards]

    strategies = load_strategies(args.base)
    if args.only_strategy:
        keep = {strategies[0]["name"], args.only_strategy}
        strategies = [s for s in strategies if s["name"] in keep]
    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"]
    assert strategies[0]["is_reference"], "first strategy must be FFFF"

    print(f"base={args.base} shard {args.shard}/{args.num_shards}: "
          f"{len(shard_rows)} prompts x {len(strategies)} strategies", flush=True)

    base = enter_base(args.base)

    import torch
    from einops import rearrange
    from torchvision.io import read_video, write_video
    from utils.misc import set_seed

    from cachelib import (CacheController, build_method, cached_inference, install,
                          parse_schedule)
    from eval.pixel_metrics import PixelMetrics

    torch.set_grad_enabled(False)
    device = torch.device("cuda")
    pipeline = load_pipeline(base)
    num_steps = len(pipeline.denoising_step_list)
    metrics = PixelMetrics(device)

    def record_path(strategy, row):
        d = os.path.join(out_root, "per_prompt", strategy)
        os.makedirs(d, exist_ok=True)
        return os.path.join(d, f"{row['prompt_suite']}_{row['suite_index']:03d}.json")

    def video_path(strategy, row):
        d = os.path.join(out_root, "generated_videos", strategy, row["prompt_suite"])
        os.makedirs(d, exist_ok=True)
        return os.path.join(d, f"{row['suite_index']:03d}.mp4")

    def ref_cache_path(row):
        d = os.path.join(out_root, "ref_cache", args.base)
        os.makedirs(d, exist_ok=True)
        return os.path.join(d, f"{row['prompt_suite']}_{row['suite_index']:03d}.pt")

    def is_done(strategy, row):
        p = record_path(strategy, row)
        if args.overwrite or not os.path.exists(p):
            return False
        try:
            with open(p) as f:
                rec = json.load(f)
            return rec.get("status") == "complete" and os.path.exists(rec.get("video", ""))
        except Exception:
            return False

    def write_json(path, payload):
        tmp = path + ".tmp"
        with open(tmp, "w") as f:
            json.dump(payload, f, indent=2)
        os.replace(tmp, path)  # atomic: a partial file is never seen as complete

    def run_one(strategy, row):
        kw = dict(indicator=strategy["indicator"], coefficients=strategy["coefficients"])
        if strategy["method"] != "none" and strategy.get("param"):
            kw[strategy["param"]] = strategy["value"]
        n_steps = strategy.get("num_inference_steps") or num_steps
        original_steps = use_denoising_steps(pipeline, n_steps)
        sched = strategy.get("schedule") or ("F" * n_steps)
        first = strategy.get("first_chunk_schedule")
        # Chunk 0 may have a longer schedule than the rest (e.g. FFFF over naive 2-step).
        first_steps = (denoising_steps_for(original_steps, len(first))
                       if first and len(first) != n_steps else None)
        ctrl = CacheController(
            build_method(strategy["method"], **kw), num_steps=n_steps,
            forced_steps=parse_schedule(sched, n_steps),
            first_chunk_forced_steps=(parse_schedule(first, len(first)) if first else None))
        install(pipeline.generator.model, ctrl)

        # Identical noise for every strategy: same seed, same draw order, same shape.
        set_seed(args.seed)
        noise = torch.randn([1, args.num_latent_frames, 16, 60, 104],
                            device=device, dtype=torch.bfloat16)
        wall0 = time.time()
        video, latents, timing = cached_inference(
            pipeline, ctrl, noise, [row["extended_prompt"]], decode=True,
            first_chunk_steps=first_steps, sampler=strategy.get("sampler", "renoise"))
        wall = time.time() - wall0
        pipeline.vae.model.clear_cache()
        pipeline.denoising_step_list = original_steps
        return video, timing, ctrl.summary(), wall

    # Warm up every strategy once on a throwaway prompt and discard the result.
    # Without this the first recorded prompt carries one-off costs -- notably
    # TaylorSeer's torch.compile of the forecast, which made its first video look
    # slower than FFFF.
    if shard_rows:
        print("warm-up pass (not recorded)", flush=True)
        for strategy in strategies:
            run_one(strategy, shard_rows[0])
        torch.cuda.empty_cache()

    t_start = time.time()
    for n, row in enumerate(shard_rows):
        todo = [s for s in strategies if not is_done(s["name"], row)]
        if not todo:
            continue
        # The reference is needed for pixel metrics even when only some strategies
        # are outstanding.
        need_ref = any(not s["is_reference"] for s in todo)
        ref_frames = None
        # The FFFF reference frames (pre-MP4, fp16) are cached on disk after their
        # first generation, so later strategy runs do not regenerate FFFF at all.
        rc = ref_cache_path(row)
        ref_from_mp4 = False
        if need_ref and os.path.exists(rc):
            ref_frames = torch.load(rc, map_location=device)
            need_ref = False
        elif need_ref and os.path.exists(video_path(ref_name, row)) and not args.regenerate_ffff:
            # Reference from the FFFF run's MP4 (H.264, ~40+ dB above the 12-20 dB
            # strategy PSNRs, so the compression adds <1 % to the MSE); never
            # regenerate FFFF just to compare against it.
            vid, _, _ = read_video(video_path(ref_name, row), pts_unit="sec", output_format="TCHW")
            ref_frames = (vid.to(device, torch.float16) / 255.0)
            ref_from_mp4 = True
            need_ref = False

        for strategy in strategies:
            if strategy["is_reference"]:
                if strategy not in todo and not need_ref:
                    continue
            elif strategy not in todo:
                continue

            video, timing, summary, wall = run_one(strategy, row)
            frames = video[0]  # [T, C, H, W] in [0, 1]

            if strategy["is_reference"]:
                ref_frames = frames.detach().to(torch.float16)
                if not os.path.exists(rc):
                    torch.save(ref_frames.cpu(), rc + ".tmp")
                    os.replace(rc + ".tmp", rc)
                pix = PixelMetrics.identity(frames.shape[0])
            else:
                pix = metrics.compute(frames, ref_frames)

            if strategy in todo:
                vp = video_path(strategy["name"], row)
                arr = (255.0 * rearrange(frames, "t c h w -> t h w c")).clamp(0, 255)
                write_video(vp, arr.to(torch.uint8).cpu(), fps=args.fps)

                write_json(record_path(strategy["name"], row), {
                    "status": "complete",
                    "protocol": "Self-Forcing Extended-251 Full Evaluation",
                    "strategy": strategy["name"],
                    "base_model": args.base,
                    "method": strategy["method"],
                    "param": strategy["param"],
                    "param_value": strategy["value"],
                    "target_speedup": strategy["target"],
                    "schedule": strategy.get("schedule") or ("F" * (strategy.get("num_inference_steps") or 4)),
                    "num_inference_steps": strategy.get("num_inference_steps") or 4,
                    "first_chunk_schedule": strategy.get("first_chunk_schedule"),
                    "sampler": strategy.get("sampler", "renoise"),
                    "global_index": row["global_index"],
                    "prompt_suite": row["prompt_suite"],
                    "suite_index": row["suite_index"],
                    "prompt": row["extended_prompt"],
                    "original_prompt": row["original_prompt"],
                    "seed": args.seed,
                    "video": vp,
                    "num_frames": int(frames.shape[0]),
                    "height": int(frames.shape[2]),
                    "width": int(frames.shape[3]),
                    "fps": args.fps,
                    "policy_latency_ms": timing["denoise_dit_ms"],
                    "excluded_context_kv_latency_ms": timing["context_kv_dit_ms"],
                    "wall_generation_s": wall,
                    "reference_strategy": ref_name,
                    "pixel_metrics_vs_ffff": pix,
                    "reference_source": "ffff_mp4" if ref_from_mp4 else "ffff_frames",
                    "cache_diagnostics": dict(summary),
                })
            del frames, video

        ref_frames = None
        done = n + 1
        rate = (time.time() - t_start) / done
        print(f"[{args.base} shard {args.shard}] {done}/{len(shard_rows)} prompts "
              f"({row['prompt_suite']}/{row['suite_index']:03d})  "
              f"{rate:.1f}s/prompt  eta {(len(shard_rows) - done) * rate / 60:.0f} min",
              flush=True)

    print(f"shard {args.shard} done in {(time.time() - t_start) / 60:.1f} min", flush=True)


if __name__ == "__main__":
    main()