load_checkpoint.py

#!/usr/bin/env python3
"""Standalone loader for the Cosmos3-Nano blockwise FP8 mixed-precision checkpoint.

Loads the safetensors checkpoint (no .pt dependency) and optionally runs inference.

Dependencies: torch, diffusers, modelopt, safetensors
Environment: CUDA GPU with >= 25 GB VRAM (480p/57f) or >= 19 GB (480p/1f smoke)

Usage:
    # Load and verify (no inference)
    python load_checkpoint.py --verify

    # Run inference with a prompt
    python load_checkpoint.py --prompt "A cat sitting on a windowsill"

    # Smoke test (1 frame, 8 steps)
    python load_checkpoint.py --prompt "A cat" --steps 8 --frames 1

    # Full quality (57 frames, 35 steps)
    python load_checkpoint.py --prompt "A cat" --steps 35 --frames 57
"""

from __future__ import annotations

import argparse
import glob
import os
import random
import sys

import numpy as np
import torch

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
SIDECAR = os.path.join(SCRIPT_DIR, "transformer", "modelopt_state.pt")
CONFIG = os.path.join(SCRIPT_DIR, "transformer", "config.json")
SAFETENSORS_GLOB = os.path.join(SCRIPT_DIR, "transformer", "*.safetensors")


def seed_everything(seed: int) -> None:
    """Set deterministic generation (INV-5)."""
    os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    try:
        torch.use_deterministic_algorithms(True, warn_only=True)
    except Exception:
        pass


def load_transformer(ckpt_dir: str | None = None):
    """Load the quantized Cosmos3OmniTransformer from safetensors + sidecar.

    Never opens modelopt_quantized.pt. The structural sidecar (~670 KB) restores
    the quantizer wrappers; safetensors provides the actual weights and scales.
    """
    import modelopt.torch.opt as mto
    from diffusers import Cosmos3OmniTransformer
    from safetensors.torch import load_file

    ckpt_dir = ckpt_dir or SCRIPT_DIR
    config_path = os.path.join(ckpt_dir, "transformer", "config.json")
    sidecar_path = os.path.join(ckpt_dir, "transformer", "modelopt_state.pt")
    safe_pattern = os.path.join(ckpt_dir, "transformer", "*.safetensors")

    cfg = {**Cosmos3OmniTransformer.load_config(config_path), "action_gen": False}
    transformer = Cosmos3OmniTransformer.from_config(cfg).to(torch.bfloat16)

    state = torch.load(sidecar_path, weights_only=False)
    restored = mto.restore_from_modelopt_state(transformer, state)
    if restored is not None:
        transformer = restored

    shards = sorted(glob.glob(safe_pattern))
    if not shards:
        raise FileNotFoundError(f"no safetensors under {ckpt_dir}/transformer/")
    tensors: dict = {}
    for shard in shards:
        tensors.update(load_file(shard))
    transformer.load_state_dict(tensors, strict=True)

    return transformer


def load_pipeline(ckpt_dir: str | None = None, device: str = "cuda"):
    """Load the full pipeline with UniPC scheduler (flow_shift=10.0, INV-5)."""
    from diffusers import Cosmos3OmniPipeline, UniPCMultistepScheduler

    ckpt_dir = ckpt_dir or SCRIPT_DIR
    transformer = load_transformer(ckpt_dir)
    pipe = Cosmos3OmniPipeline.from_pretrained(
        ckpt_dir, transformer=transformer, torch_dtype=torch.bfloat16,
        enable_safety_checker=False,
    )
    pipe.scheduler = UniPCMultistepScheduler.from_config(
        pipe.scheduler.config, flow_shift=10.0,
    )
    return pipe.to(device)


def verify(ckpt_dir: str | None = None) -> None:
    """Load the checkpoint and print quantizer stats (no inference)."""
    ckpt_dir = ckpt_dir or SCRIPT_DIR
    print(f"Loading from {ckpt_dir}...")
    transformer = load_transformer(ckpt_dir)

    n_enabled = 0
    for _name, mod in transformer.named_modules():
        wq = getattr(mod, "weight_quantizer", None)
        if wq is not None and getattr(wq, "is_enabled", False):
            n_enabled += 1

    n_params = sum(p.numel() for p in transformer.parameters())
    print(f"Loaded: {n_enabled} quantized modules, {n_params / 1e9:.2f}B parameters")
    print("Verify OK" if n_enabled == 217 else f"UNEXPECTED quantizer count: {n_enabled}")


def generate(
    prompt: str,
    ckpt_dir: str | None = None,
    seed: int = 123,
    steps: int = 8,
    frames: int = 1,
    height: int = 480,
    width: int = 640,
    output_dir: str = "output",
) -> None:
    """Run inference and save output frames."""
    from PIL import Image

    seed_everything(seed)
    pipe = load_pipeline(ckpt_dir)

    print(f"Generating: {width}x{height}, {frames}f, {steps} steps, seed={seed}")
    with torch.autocast("cuda", torch.bfloat16):
        result = pipe(
            prompt=prompt,
            num_frames=frames,
            height=height,
            width=width,
            num_inference_steps=steps,
            generator=torch.Generator("cpu").manual_seed(seed),
        )

    video = result.video
    if isinstance(video, (list, tuple)) and video and isinstance(video[0], (list, tuple)):
        video = video[0]

    os.makedirs(output_dir, exist_ok=True)
    for i, frame in enumerate(video):
        if not isinstance(frame, Image.Image):
            frame = Image.fromarray(frame)
        frame.save(os.path.join(output_dir, f"frame_{i:04d}.png"))
    print(f"Saved {len(video)} frames to {output_dir}/")


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="Cosmos3-Nano blockwise FP8 checkpoint loader")
    p.add_argument("--ckpt-dir", default=None, help="Checkpoint directory (default: script dir)")
    p.add_argument("--verify", action="store_true", help="Load and verify only (no inference)")
    p.add_argument("--prompt", default=None, help="Generation prompt")
    p.add_argument("--seed", type=int, default=123, help="Random seed (default: 123)")
    p.add_argument("--steps", type=int, default=8, help="Denoising steps (default: 8)")
    p.add_argument("--frames", type=int, default=1, help="Number of frames (default: 1)")
    p.add_argument("--height", type=int, default=480, help="Frame height (default: 480)")
    p.add_argument("--width", type=int, default=640, help="Frame width (default: 640)")
    p.add_argument("--output-dir", default="output", help="Output directory (default: output)")
    return p


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)

    if args.verify:
        verify(args.ckpt_dir)
        return 0

    if args.prompt is None:
        print("Error: --prompt required (or use --verify)")
        return 1

    generate(
        prompt=args.prompt,
        ckpt_dir=args.ckpt_dir,
        seed=args.seed,
        steps=args.steps,
        frames=args.frames,
        height=args.height,
        width=args.width,
        output_dir=args.output_dir,
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())