File size: 8,493 Bytes
38a51ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | #!/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)
|