#!/usr/bin/env python3 """Online Full-DiT vs direct-reuse vs F-P-F-P Predictor evaluation.""" from __future__ import annotations import argparse import csv import gc import json import os import time from collections import defaultdict from pathlib import Path from types import SimpleNamespace import torch import torch.nn.functional as F from safetensors.torch import load_file, save_file from hyvideo.commons.infer_state import initialize_infer_state from hyvideo.generate import pose_to_input from hyvideo.pipelines.worldplay_video_pipeline import HunyuanVideo_1_5_Pipeline from models import HYWorldPlayPredictor from predictor_training.checkpoint import load_predictor_weights from tools.build_predictor_dataset import configure_exact_teacher, parse_case_ids DEFAULT_POSE = "w-7,s-8,a-8,d-8,left-8,right-8,up-8,down-8" def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--case_csv", default="assets/test_case.csv") parser.add_argument("--case_ids", default="1-10") parser.add_argument("--dataset_dir", default="datasets/predictor_v1") parser.add_argument("--weights", required=True) parser.add_argument("--output_dir", default="outputs/predictor_v1/online") parser.add_argument("--modes", default="full,reuse,predictor") parser.add_argument("--base_model", default="models-ms/HunyuanVideo-1.5") parser.add_argument( "--action_ckpt", default="models-ms/HY-WorldPlay/ar_distilled_action_model/diffusion_pytorch_model.safetensors", ) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--height", type=int, default=480) parser.add_argument("--width", type=int, default=832) parser.add_argument("--video_frames", type=int, default=253) parser.add_argument("--pose", default=DEFAULT_POSE) parser.add_argument("--save_latents", action=argparse.BooleanOptionalAction, default=True) return parser.parse_args() def load_cases(csv_path: Path) -> dict[int, dict[str, str]]: with csv_path.open("r", encoding="utf-8", newline="") as handle: rows = list(csv.DictReader(handle)) result = {} for case_id, row in enumerate(rows, start=1): image_path = (csv_path.parent.parent / row["image_name"]).resolve() if not image_path.is_file(): image_path = Path(row["image_name"]).resolve() if not image_path.is_file(): raise FileNotFoundError(f"Missing image for case {case_id}: {row['image_name']}") result[case_id] = {"image_path": str(image_path), "caption": row["caption"]} return result def init_state() -> None: initialize_infer_state( SimpleNamespace( sage_blocks_range="0-53", use_sageattn=False, enable_torch_compile=False, use_fp8_gemm=False, quant_type="fp8-per-block", include_patterns="double_blocks", use_vae_parallel=False, ) ) def aggregate_timing(records: list[dict]) -> dict[str, float]: totals: dict[str, float] = defaultdict(float) counts: dict[str, int] = defaultdict(int) for record in records: name = str(record["name"]) totals[name] += float(record["elapsed_s"]) counts[name] += 1 result = {} for name in sorted(totals): result[f"{name}_total_s"] = totals[name] result[f"{name}_count"] = counts[name] result[f"{name}_mean_s"] = totals[name] / counts[name] return result def latent_metrics(candidate: torch.Tensor, baseline: torch.Tensor) -> dict[str, float]: candidate = candidate.float() baseline = baseline.float() mse = F.mse_loss(candidate, baseline) baseline_energy = baseline.square().mean().clamp_min(1e-12) return { "latent_mse_vs_full": float(mse), "latent_nrmse_vs_full": float(torch.sqrt(mse / baseline_energy)), "latent_mae_vs_full": float(F.l1_loss(candidate, baseline)), "latent_cosine_vs_full": float( F.cosine_similarity(candidate.flatten(1), baseline.flatten(1)).mean() ), } def write_json_atomic(path: Path, payload: dict) -> None: path.parent.mkdir(parents=True, exist_ok=True) tmp = path.with_suffix(path.suffix + ".tmp") tmp.write_text(json.dumps(payload, indent=2, ensure_ascii=False, sort_keys=True) + "\n") os.replace(tmp, path) def summarize(runs: list[dict]) -> dict[str, dict[str, float]]: by_mode: dict[str, dict[str, list[float]]] = defaultdict(lambda: defaultdict(list)) for run in runs: mode = run["mode"] for name, value in run.items(): if name not in {"case_id", "mode"} and isinstance(value, (int, float)): by_mode[mode][name].append(float(value)) for name, value in run["timing"].items(): if name.endswith("_s"): by_mode[mode][f"timing.{name}"].append(float(value)) return { mode: { name: sum(values) / len(values) for name, values in sorted(metrics.items()) } for mode, metrics in by_mode.items() } def main() -> None: args = parse_args() modes = [item.strip() for item in args.modes.split(",") if item.strip()] if any(mode not in {"full", "reuse", "predictor"} for mode in modes): raise ValueError(f"Unsupported modes: {modes}") if not modes or modes[0] != "full": raise ValueError("Modes must start with full so alternatives use the same-run baseline") if args.video_frames != 13 and args.video_frames != 253: raise ValueError("Online evaluator currently supports the 13-frame smoke or 253-frame run") os.environ["HY_PROFILE_TIMING"] = "1" init_state() device = torch.device("cuda:0") dtype = torch.bfloat16 case_ids = parse_case_ids(args.case_ids) cases = load_cases(Path(args.case_csv).resolve()) latent_frames = (args.video_frames - 1) // 4 + 1 viewmats, Ks, action = pose_to_input(args.pose, latent_frames) pipe = HunyuanVideo_1_5_Pipeline.create_pipeline( pretrained_model_name_or_path=str(Path(args.base_model).resolve()), transformer_version="480p_i2v", enable_offloading=True, enable_group_offloading=False, create_sr_pipeline=False, force_sparse_attn=False, transformer_dtype=dtype, action_ckpt=str(Path(args.action_ckpt).resolve()), ) configure_exact_teacher(pipe.transformer) pipe.transformer.eval().requires_grad_(False) predictor = None if "predictor" in modes: predictor = HYWorldPlayPredictor() predictor.load_teacher_initialization(args.action_ckpt) load_predictor_weights(predictor, args.weights) predictor.to(device="cpu", dtype=dtype).eval() output_dir = Path(args.output_dir).resolve() latent_dir = output_dir / "latents" output_path = output_dir / "metrics.json" dataset_dir = Path(args.dataset_dir).resolve() runs: list[dict] = [] for case_id in case_ids: case = cases[case_id] case_tensors = load_file( str(dataset_dir / "cases" / f"case_{case_id:02d}.safetensors"), device="cpu", ) baseline = None for mode in modes: runtime = None if predictor is not None: predictor.to(device=device if mode == "predictor" else "cpu") if mode == "predictor": runtime = { "model": predictor, "current_txt": case_tensors["current_txt"].to(device=device, dtype=dtype), "cached_txt": case_tensors["cached_txt"].to(device=device, dtype=dtype), "vec_txt": case_tensors["vec_txt"].to(device=device, dtype=dtype), } torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) gc.collect() torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats(device) torch.cuda.synchronize(device) started = time.perf_counter() output = pipe( enable_sr=False, prompt=case["caption"], aspect_ratio="16:9", num_inference_steps=4, video_length=args.video_frames, negative_prompt="", seed=args.seed, output_type="latent", prompt_rewrite=False, return_pre_sr_video=False, reference_image=case["image_path"], viewmats=viewmats.unsqueeze(0), Ks=Ks.unsqueeze(0), action=action.unsqueeze(0), few_step=True, chunk_latent_frames=4, model_type="ar", user_height=args.height, user_width=args.width, transformer_resident_ar_rollout=True, predictor_mode=mode, predictor_runtime=runtime if mode == "predictor" else None, ) torch.cuda.synchronize(device) wall_time = time.perf_counter() - started latent = output.videos.detach().to(device="cpu", dtype=torch.float32).contiguous() run = { "case_id": case_id, "mode": mode, "wall_time_s": wall_time, "peak_memory_gib": torch.cuda.max_memory_allocated(device) / 1024**3, "timing": aggregate_timing(pipe._timing_records), } if mode == "full": baseline = latent run.update( { "latent_mse_vs_full": 0.0, "latent_nrmse_vs_full": 0.0, "latent_mae_vs_full": 0.0, "latent_cosine_vs_full": 1.0, } ) else: if baseline is None: raise RuntimeError("Full baseline is unavailable") run.update(latent_metrics(latent, baseline)) runs.append(run) if args.save_latents: latent_dir.mkdir(parents=True, exist_ok=True) save_file( {"latents": latent.to(torch.bfloat16)}, str(latent_dir / f"case_{case_id:02d}_{mode}.safetensors"), ) payload = { "config": vars(args), "predictor_parameters": ( predictor.trainable_parameter_count() if predictor is not None else 0 ), "runs": runs, "summary_mean": summarize(runs), } write_json_atomic(output_path, payload) print(json.dumps(run, indent=2, sort_keys=True), flush=True) if __name__ == "__main__": main()