infer.py

#!/usr/bin/env python3
"""Generate one video with the pinned public HunyuanVideo inference pipeline."""

import argparse
import hashlib
import json
import os
import re
import sys
import time
from pathlib import Path


MODEL_ID = "vanzl/DiffusionART-HunyuanVideo-7-2-1"
BASE_ID = "hunyuanvideo-community/HunyuanVideo"
BASE_REVISION = "e8c2aaa66fe3742a32c11a6766aecbf07c56e773"
FRAMES = 53
WIDTH = HEIGHT = 640
FPS = 8
STEPS = 30
SHIFT = 7.0
GUIDANCE_SCALE = 6.018  # The public pipeline multiplies this by 1000.


def arguments():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--prompt", required=True, help="Text prompt for one video")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output", required=True, type=Path, help="Output MP4 path")
    parser.add_argument("--cache-dir", required=True, type=Path,
                        help="Writable directory for both pinned Hugging Face snapshots and runtime caches")
    parser.add_argument("--revision", required=True,
                        help="Exact 40-character commit SHA of the private model release")
    parser.add_argument("--cpu-offload", action="store_true",
                        help="Move inactive modules to CPU; the default H800 acceptance did not need this")
    args = parser.parse_args()
    if not args.prompt.strip():
        parser.error("--prompt must contain text")
    if not re.fullmatch(r"[0-9a-f]{40}", args.revision):
        parser.error("--revision must be a fixed 40-character lowercase commit SHA")
    if args.output.suffix.lower() != ".mp4":
        parser.error("--output must end in .mp4")
    if args.output.exists():
        parser.error(f"output already exists: {args.output}")
    return args


def configure_caches(root: Path):
    root.mkdir(parents=True, exist_ok=True)
    hf_home = root / "hf"
    hub = hf_home / "hub"
    for directory in (hf_home, hub, root / "xdg", root / "tmp", root / "triton"):
        directory.mkdir(parents=True, exist_ok=True)
    os.environ.update({
        "HF_HOME": str(hf_home),
        "HF_HUB_CACHE": str(hub),
        "HUGGINGFACE_HUB_CACHE": str(hub),
        "HF_ASSETS_CACHE": str(hf_home / "assets"),
        "XDG_CACHE_HOME": str(root / "xdg"),
        "TMPDIR": str(root / "tmp"),
        "TRITON_CACHE_DIR": str(root / "triton"),
        "TOKENIZERS_PARALLELISM": "false",
        "HF_HUB_DISABLE_TELEMETRY": "1",
        "PYTHONNOUSERSITE": "1",
    })
    return hub


def sha256(path: Path):
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for block in iter(lambda: source.read(16 * 1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def validate_video(path: Path):
    import av
    import numpy as np

    frame_count = 0
    pixel_sum = 0.0
    squared_sum = 0.0
    pixel_min = 255
    pixel_max = 0
    with av.open(str(path)) as container:
        streams = [stream for stream in container.streams if stream.type == "video"]
        if len(streams) != 1:
            raise RuntimeError(f"expected one video stream, got {len(streams)}")
        stream = streams[0]
        width, height = stream.width, stream.height
        fps = float(stream.average_rate)
        for frame in container.decode(stream):
            pixels = frame.to_ndarray(format="rgb24")
            if pixels.shape != (HEIGHT, WIDTH, 3):
                raise RuntimeError(f"unexpected frame shape: {pixels.shape}")
            values = pixels.astype(np.float64)
            pixel_sum += float(values.sum())
            squared_sum += float(np.square(values).sum())
            pixel_min = min(pixel_min, int(pixels.min()))
            pixel_max = max(pixel_max, int(pixels.max()))
            frame_count += 1
    if (frame_count, width, height) != (FRAMES, WIDTH, HEIGHT):
        raise RuntimeError(f"unexpected video geometry: {(frame_count, width, height)}")
    if abs(fps - FPS) > 0.01:
        raise RuntimeError(f"unexpected FPS: {fps}")
    count = frame_count * width * height * 3
    mean = pixel_sum / count
    std = max(0.0, squared_sum / count - mean * mean) ** 0.5
    if pixel_max <= pixel_min or pixel_max == 0 or std <= 1.0:
        raise RuntimeError(f"blank or nearly constant video: range={pixel_min}:{pixel_max}, std={std}")
    return {"decoded_frames": frame_count, "width": width, "height": height,
            "fps": fps, "pixel_min": pixel_min, "pixel_max": pixel_max,
            "pixel_mean": mean, "pixel_std": std}


def main():
    args = arguments()
    cache_root = args.cache_dir.expanduser().resolve()
    hub_cache = configure_caches(cache_root)
    output = args.output.expanduser().resolve()
    output.parent.mkdir(parents=True, exist_ok=True)

    # Import after setting cache variables, so neither library selects a home cache.
    import torch
    from diffusers import FlowMatchEulerDiscreteScheduler
    from diffusers.utils import export_to_video
    from fastvideo.models.hunyuan_hf.pipeline_hunyuan import HunyuanVideoPipeline
    from fastvideo.utils.load import load_transformer
    from huggingface_hub import snapshot_download

    if not torch.cuda.is_available():
        raise RuntimeError("A CUDA GPU is required")
    if torch.version.cuda != "12.1":
        raise RuntimeError(f"expected the verified PyTorch CUDA 12.1 build, got {torch.version.cuda}")
    model = Path(snapshot_download(MODEL_ID, revision=args.revision,
                                   cache_dir=str(hub_cache), token=True))
    base = Path(snapshot_download(BASE_ID, revision=BASE_REVISION,
                                  cache_dir=str(hub_cache)))
    if model.name != args.revision or base.name != BASE_REVISION:
        raise RuntimeError("Hugging Face snapshot revision differs from the requested commit")

    started = time.perf_counter()
    transformer = load_transformer("hunyuan_hf", str(model / "transformer"),
                                   str(base), torch.float32)
    if transformer.dtype != torch.float32:
        raise RuntimeError(f"unexpected transformer dtype: {transformer.dtype}")
    # The public base snapshot supplies both text tokenizers and encoders. This
    # preserves the pipeline's default prompt template and avoids saved embeds.
    pipe = HunyuanVideoPipeline.from_pretrained(
        str(base), transformer=transformer, torch_dtype=torch.bfloat16)
    if not isinstance(pipe.scheduler, FlowMatchEulerDiscreteScheduler):
        raise RuntimeError(f"unexpected scheduler: {type(pipe.scheduler).__name__}")
    if abs(float(pipe.scheduler.config.shift) - SHIFT) > 1e-6:
        raise RuntimeError(f"unexpected scheduler shift: {pipe.scheduler.config.shift}")
    pipe.enable_vae_tiling()
    if args.cpu_offload:
        pipe.enable_model_cpu_offload()
    else:
        pipe.to("cuda:0")
    load_seconds = time.perf_counter() - started

    generator = torch.Generator(device="cpu").manual_seed(args.seed)
    torch.cuda.reset_peak_memory_stats(0)
    started = time.perf_counter()
    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
        frames = pipe(prompt=[args.prompt], height=HEIGHT, width=WIDTH,
                      num_frames=FRAMES, num_inference_steps=STEPS,
                      guidance_scale=GUIDANCE_SCALE, generator=generator).frames[0]
    inference_seconds = time.perf_counter() - started
    peak_bytes = torch.cuda.max_memory_allocated(0)

    partial = output.with_name(f"{output.stem}.partial.{os.getpid()}.mp4")
    try:
        export_to_video(frames, str(partial), fps=FPS)
        del frames
        decoded = validate_video(partial)
        partial.replace(output)
    finally:
        if partial.exists():
            partial.unlink()
    result = {"output": str(output), "sha256": sha256(output),
              "model_revision": args.revision, "base_revision": BASE_REVISION,
              "prompt": args.prompt, "seed": args.seed,
              "sampling": {"frames": FRAMES, "width": WIDTH, "height": HEIGHT,
                           "fps": FPS, "steps": STEPS, "shift": SHIFT, "eta": 0,
                           "classic_cfg": False, "guidance_scale": GUIDANCE_SCALE,
                           "embedded_guidance": 6018},
              "load_seconds": round(load_seconds, 2),
              "inference_seconds": round(inference_seconds, 2),
              "peak_cuda_allocated_bytes": peak_bytes, "decoded": decoded}
    print(json.dumps(result, ensure_ascii=False, indent=2), flush=True)


if __name__ == "__main__":
    try:
        main()
    except Exception as error:
        print(f"inference failed: {error}", file=sys.stderr, flush=True)
        raise