Cccccz/HY

1#!/usr/bin/env python32"""Online Full-DiT vs direct-reuse vs F-P-F-P Predictor evaluation."""3 4from __future__ import annotations5 6import argparse7import csv8import gc9import json10import os11import time12from collections import defaultdict13from pathlib import Path14from types import SimpleNamespace15 16import torch17import torch.nn.functional as F18from safetensors.torch import load_file, save_file19 20from hyvideo.commons.infer_state import initialize_infer_state21from hyvideo.generate import pose_to_input22from hyvideo.pipelines.worldplay_video_pipeline import HunyuanVideo_1_5_Pipeline23from models import HYWorldPlayPredictor24from predictor_training.checkpoint import load_predictor_weights25from tools.build_predictor_dataset import configure_exact_teacher, parse_case_ids26 27 28DEFAULT_POSE = "w-7,s-8,a-8,d-8,left-8,right-8,up-8,down-8"29 30 31def parse_args() -> argparse.Namespace:32    parser = argparse.ArgumentParser(description=__doc__)33    parser.add_argument("--case_csv", default="assets/test_case.csv")34    parser.add_argument("--case_ids", default="1-10")35    parser.add_argument("--dataset_dir", default="datasets/predictor_v1")36    parser.add_argument("--weights", required=True)37    parser.add_argument("--output_dir", default="outputs/predictor_v1/online")38    parser.add_argument("--modes", default="full,reuse,predictor")39    parser.add_argument("--base_model", default="models-ms/HunyuanVideo-1.5")40    parser.add_argument(41        "--action_ckpt",42        default="models-ms/HY-WorldPlay/ar_distilled_action_model/diffusion_pytorch_model.safetensors",43    )44    parser.add_argument("--seed", type=int, default=0)45    parser.add_argument("--height", type=int, default=480)46    parser.add_argument("--width", type=int, default=832)47    parser.add_argument("--video_frames", type=int, default=253)48    parser.add_argument("--pose", default=DEFAULT_POSE)49    parser.add_argument("--save_latents", action=argparse.BooleanOptionalAction, default=True)50    return parser.parse_args()51 52 53def load_cases(csv_path: Path) -> dict[int, dict[str, str]]:54    with csv_path.open("r", encoding="utf-8", newline="") as handle:55        rows = list(csv.DictReader(handle))56    result = {}57    for case_id, row in enumerate(rows, start=1):58        image_path = (csv_path.parent.parent / row["image_name"]).resolve()59        if not image_path.is_file():60            image_path = Path(row["image_name"]).resolve()61        if not image_path.is_file():62            raise FileNotFoundError(f"Missing image for case {case_id}: {row['image_name']}")63        result[case_id] = {"image_path": str(image_path), "caption": row["caption"]}64    return result65 66 67def init_state() -> None:68    initialize_infer_state(69        SimpleNamespace(70            sage_blocks_range="0-53",71            use_sageattn=False,72            enable_torch_compile=False,73            use_fp8_gemm=False,74            quant_type="fp8-per-block",75            include_patterns="double_blocks",76            use_vae_parallel=False,77        )78    )79 80 81def aggregate_timing(records: list[dict]) -> dict[str, float]:82    totals: dict[str, float] = defaultdict(float)83    counts: dict[str, int] = defaultdict(int)84    for record in records:85        name = str(record["name"])86        totals[name] += float(record["elapsed_s"])87        counts[name] += 188    result = {}89    for name in sorted(totals):90        result[f"{name}_total_s"] = totals[name]91        result[f"{name}_count"] = counts[name]92        result[f"{name}_mean_s"] = totals[name] / counts[name]93    return result94 95 96def latent_metrics(candidate: torch.Tensor, baseline: torch.Tensor) -> dict[str, float]:97    candidate = candidate.float()98    baseline = baseline.float()99    mse = F.mse_loss(candidate, baseline)100    baseline_energy = baseline.square().mean().clamp_min(1e-12)101    return {102        "latent_mse_vs_full": float(mse),103        "latent_nrmse_vs_full": float(torch.sqrt(mse / baseline_energy)),104        "latent_mae_vs_full": float(F.l1_loss(candidate, baseline)),105        "latent_cosine_vs_full": float(106            F.cosine_similarity(candidate.flatten(1), baseline.flatten(1)).mean()107        ),108    }109 110 111def write_json_atomic(path: Path, payload: dict) -> None:112    path.parent.mkdir(parents=True, exist_ok=True)113    tmp = path.with_suffix(path.suffix + ".tmp")114    tmp.write_text(json.dumps(payload, indent=2, ensure_ascii=False, sort_keys=True) + "\n")115    os.replace(tmp, path)116 117 118def summarize(runs: list[dict]) -> dict[str, dict[str, float]]:119    by_mode: dict[str, dict[str, list[float]]] = defaultdict(lambda: defaultdict(list))120    for run in runs:121        mode = run["mode"]122        for name, value in run.items():123            if name not in {"case_id", "mode"} and isinstance(value, (int, float)):124                by_mode[mode][name].append(float(value))125        for name, value in run["timing"].items():126            if name.endswith("_s"):127                by_mode[mode][f"timing.{name}"].append(float(value))128    return {129        mode: {130            name: sum(values) / len(values)131            for name, values in sorted(metrics.items())132        }133        for mode, metrics in by_mode.items()134    }135 136 137def main() -> None:138    args = parse_args()139    modes = [item.strip() for item in args.modes.split(",") if item.strip()]140    if any(mode not in {"full", "reuse", "predictor"} for mode in modes):141        raise ValueError(f"Unsupported modes: {modes}")142    if not modes or modes[0] != "full":143        raise ValueError("Modes must start with full so alternatives use the same-run baseline")144    if args.video_frames != 13 and args.video_frames != 253:145        raise ValueError("Online evaluator currently supports the 13-frame smoke or 253-frame run")146 147    os.environ["HY_PROFILE_TIMING"] = "1"148    init_state()149    device = torch.device("cuda:0")150    dtype = torch.bfloat16151    case_ids = parse_case_ids(args.case_ids)152    cases = load_cases(Path(args.case_csv).resolve())153    latent_frames = (args.video_frames - 1) // 4 + 1154    viewmats, Ks, action = pose_to_input(args.pose, latent_frames)155 156    pipe = HunyuanVideo_1_5_Pipeline.create_pipeline(157        pretrained_model_name_or_path=str(Path(args.base_model).resolve()),158        transformer_version="480p_i2v",159        enable_offloading=True,160        enable_group_offloading=False,161        create_sr_pipeline=False,162        force_sparse_attn=False,163        transformer_dtype=dtype,164        action_ckpt=str(Path(args.action_ckpt).resolve()),165    )166    configure_exact_teacher(pipe.transformer)167    pipe.transformer.eval().requires_grad_(False)168 169    predictor = None170    if "predictor" in modes:171        predictor = HYWorldPlayPredictor()172        predictor.load_teacher_initialization(args.action_ckpt)173        load_predictor_weights(predictor, args.weights)174        predictor.to(device="cpu", dtype=dtype).eval()175 176    output_dir = Path(args.output_dir).resolve()177    latent_dir = output_dir / "latents"178    output_path = output_dir / "metrics.json"179    dataset_dir = Path(args.dataset_dir).resolve()180    runs: list[dict] = []181 182    for case_id in case_ids:183        case = cases[case_id]184        case_tensors = load_file(185            str(dataset_dir / "cases" / f"case_{case_id:02d}.safetensors"),186            device="cpu",187        )188        baseline = None189        for mode in modes:190            runtime = None191            if predictor is not None:192                predictor.to(device=device if mode == "predictor" else "cpu")193            if mode == "predictor":194                runtime = {195                    "model": predictor,196                    "current_txt": case_tensors["current_txt"].to(device=device, dtype=dtype),197                    "cached_txt": case_tensors["cached_txt"].to(device=device, dtype=dtype),198                    "vec_txt": case_tensors["vec_txt"].to(device=device, dtype=dtype),199                }200            torch.manual_seed(args.seed)201            torch.cuda.manual_seed_all(args.seed)202            gc.collect()203            torch.cuda.empty_cache()204            torch.cuda.reset_peak_memory_stats(device)205            torch.cuda.synchronize(device)206            started = time.perf_counter()207            output = pipe(208                enable_sr=False,209                prompt=case["caption"],210                aspect_ratio="16:9",211                num_inference_steps=4,212                video_length=args.video_frames,213                negative_prompt="",214                seed=args.seed,215                output_type="latent",216                prompt_rewrite=False,217                return_pre_sr_video=False,218                reference_image=case["image_path"],219                viewmats=viewmats.unsqueeze(0),220                Ks=Ks.unsqueeze(0),221                action=action.unsqueeze(0),222                few_step=True,223                chunk_latent_frames=4,224                model_type="ar",225                user_height=args.height,226                user_width=args.width,227                transformer_resident_ar_rollout=True,228                predictor_mode=mode,229                predictor_runtime=runtime if mode == "predictor" else None,230            )231            torch.cuda.synchronize(device)232            wall_time = time.perf_counter() - started233            latent = output.videos.detach().to(device="cpu", dtype=torch.float32).contiguous()234            run = {235                "case_id": case_id,236                "mode": mode,237                "wall_time_s": wall_time,238                "peak_memory_gib": torch.cuda.max_memory_allocated(device) / 1024**3,239                "timing": aggregate_timing(pipe._timing_records),240            }241            if mode == "full":242                baseline = latent243                run.update(244                    {245                        "latent_mse_vs_full": 0.0,246                        "latent_nrmse_vs_full": 0.0,247                        "latent_mae_vs_full": 0.0,248                        "latent_cosine_vs_full": 1.0,249                    }250                )251            else:252                if baseline is None:253                    raise RuntimeError("Full baseline is unavailable")254                run.update(latent_metrics(latent, baseline))255            runs.append(run)256 257            if args.save_latents:258                latent_dir.mkdir(parents=True, exist_ok=True)259                save_file(260                    {"latents": latent.to(torch.bfloat16)},261                    str(latent_dir / f"case_{case_id:02d}_{mode}.safetensors"),262                )263            payload = {264                "config": vars(args),265                "predictor_parameters": (266                    predictor.trainable_parameter_count() if predictor is not None else 0267                ),268                "runs": runs,269                "summary_mean": summarize(runs),270            }271            write_json_atomic(output_path, payload)272            print(json.dumps(run, indent=2, sort_keys=True), flush=True)273 274 275if __name__ == "__main__":276    main()277