eval_predictor_rollout.py

#!/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()