unsloth/scripts/hunyuan_trim_e2e.py

#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""End-to-end pixel validation of the HunyuanVideo-1.5 attention trim: same seed, stock vs trim,
per-frame LPIPS + mean luma (black-frame guard), plus a full-resolution trim gen for real
wall-clock. Stock is only run at MODEST settings (a full 121-frame stock gen is ~19 min).

Run: CUDA_VISIBLE_DEVICES=3 python scripts/hunyuan_trim_e2e.py
"""

from __future__ import annotations

import argparse
import os
import sys
import time
from pathlib import Path

os.environ.setdefault("BITSANDBYTES_NOWELCOME", "1")

import numpy as np
import torch

_REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_REPO_ROOT / "studio" / "backend"))
OUT = _REPO_ROOT / "outputs" / "video_speedmem"
OUT.mkdir(parents = True, exist_ok = True)


def _import_diffusers():
    import diffusers.utils.import_utils as iu

    iu._bitsandbytes_available = False
    import diffusers

    return diffusers


def _gen(pipe, seed, frames, steps, w, h):
    g = torch.Generator(device = "cuda").manual_seed(seed)
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    out = pipe(
        prompt = "a cat playing piano on a stage, cinematic",
        num_frames = frames,
        width = w,
        height = h,
        num_inference_steps = steps,
        generator = g,
        output_type = "np",
    )
    torch.cuda.synchronize()
    wall = (time.perf_counter() - t0) * 1e3
    frames_np = out.frames[0]  # [F,H,W,C] in [0,1]
    return np.asarray(frames_np), wall


def _luma(frames):
    # BT.601 luma over [F,H,W,C] in [0,1]
    r, gg, b = frames[..., 0], frames[..., 1], frames[..., 2]
    return float((0.299 * r + 0.587 * gg + 0.114 * b).mean())


def _lpips_mean(
    loss_fn,
    a,
    b,
    stride = 4,
):
    vals = []
    for i in range(0, len(a), stride):
        ta = torch.from_numpy(a[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
        tb = torch.from_numpy(b[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
        with torch.no_grad():
            vals.append(loss_fn(ta.cuda(), tb.cuda()).item())
    return float(np.mean(vals)) if vals else None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
    ap.add_argument("--cmp-frames", type = int, default = 25)
    ap.add_argument("--cmp-steps", type = int, default = 50)
    ap.add_argument("--cmp-w", type = int, default = 832)
    ap.add_argument("--cmp-h", type = int, default = 480)
    ap.add_argument("--full-frames", type = int, default = 121)
    ap.add_argument("--full-steps", type = int, default = 50)
    ap.add_argument("--seed", type = int, default = 1234)
    args = ap.parse_args()

    diffusers = _import_diffusers()
    from core.inference.diffusion_attention import install_hunyuan_attention_trim
    from core.inference.video_families import detect_video_family
    import lpips

    print(f"loading {args.repo} ...", flush = True)
    pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype = torch.bfloat16).to(
        "cuda"
    )
    fam = detect_video_family(args.repo) or detect_video_family("hunyuanvideo-1.5")
    loss_fn = lpips.LPIPS(net = "alex").cuda()

    fr, st, w, hh = args.cmp_frames, args.cmp_steps, args.cmp_w, args.cmp_h
    # ---- STOCK x2 (nondeterminism floor) ----
    print(f"\n[stock#1] gen {fr}f/{st}steps {w}x{hh} ...", flush = True)
    stock1, w_stock = _gen(pipe, args.seed, fr, st, w, hh)
    print(
        f"[stock#1] wall={w_stock:.0f} ms  luma={_luma(stock1):.4f}  frames={stock1.shape}",
        flush = True,
    )
    print(f"[stock#2] gen (same seed, measures run-to-run nondeterminism) ...", flush = True)
    stock2, _ = _gen(pipe, args.seed, fr, st, w, hh)

    # ---- TRIM (same seed) ----
    engaged = install_hunyuan_attention_trim(pipe, fam, logger = None)
    print(f"\ninstall_hunyuan_attention_trim engaged={engaged}", flush = True)
    trim, w_trim = _gen(pipe, args.seed, fr, st, w, hh)
    print(
        f"[trim ] wall={w_trim:.0f} ms  luma={_luma(trim):.4f}  ({w_stock/w_trim:.2f}x vs stock)",
        flush = True,
    )

    floor = _lpips_mean(loss_fn, stock1, stock2)
    lp = _lpips_mean(loss_fn, stock1, trim)
    print(f"\nLPIPS(stock#1, stock#2) = {floor:.5f}   <- nondeterminism floor", flush = True)
    print(f"LPIPS(stock#1, trim  ) = {lp:.5f}   <- trim vs stock", flush = True)
    print(
        f"  => trim is {'WITHIN' if lp <= floor * 1.5 + 0.002 else 'ABOVE'} the nondeterminism floor",
        flush = True,
    )

    # ---- FULL-RES TRIM (wall-clock + black-frame guard; stock at full-res is ~19 min, skipped) ----
    print(
        f"\n[trim-full] gen {args.full_frames}f/{args.full_steps}steps {args.cmp_w}x{args.cmp_h} ...",
        flush = True,
    )
    full, w_full = _gen(pipe, args.seed, args.full_frames, args.full_steps, args.cmp_w, args.cmp_h)
    print(
        f"[trim-full] wall={w_full/1000:.1f} s  luma={_luma(full):.4f}  frames={full.shape}",
        flush = True,
    )

    try:
        from PIL import Image

        Image.fromarray((trim[0] * 255).astype("uint8")).save(OUT / "vid_hunyuan_trim_cmp.png")
        Image.fromarray((full[0] * 255).astype("uint8")).save(OUT / "vid_hunyuan_trim_full.png")
        print(f"\nsaved sample frames to {OUT}", flush = True)
    except Exception as exc:  # noqa: BLE001
        print(f"(png save skipped: {exc})", flush = True)


if __name__ == "__main__":
    main()