| |
| """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() |
|
|
