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