hy_generate_eval.py

#!/usr/bin/env python
"""Generate the VBench-I2V-100 evaluation videos on the HY-WorldPlay 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 24, like the DEV project's validation videos).

    CUDA_VISIBLE_DEVICES=5 python hy_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 hycache import harness as H
H.enter_repo()
import torch
from hycache.measure import make_ctrl
from hy_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_hy")
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 hy_ffff's MP4")
ap.add_argument("--tag", default="",
                help="Video-length tag for long videos, e.g. _16c with --frames 253: the reference "
                     "becomes hy_16c_ffff and every strategy gets the tag inserted before its "
                     "_c0FFFF suffix (hy_naive2step_16c_c0FFFF), so nothing collides with the "
                     "125-frame rows")
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]
if args.tag:
    for s in strategies:
        if s["is_reference"]:
            s["name"] = f"hy{args.tag}_ffff"
        elif s["name"].endswith("_c0FFFF"):
            s["name"] = s["name"][:-len("_c0FFFF")] + args.tag + "_c0FFFF"
        else:
            s["name"] += args.tag
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"))
from hyvideo.generate import 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, decode=True,
                                  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 hy_ffff run's MP4: FFFF is not regenerated, so there is
        # no same-process FFFF timing (hy_aggregate falls back to hy_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": "HY-WorldPlay VBench-I2V-100 (validation25 x 4 actions)",
            "strategy": s["name"], "base_model": "hy_worldplay", "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": 24,
            "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"[hy 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)