File size: 3,909 Bytes
b30771b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | """Load this quantized Cosmos3-Nano (NVFP4-AWQ, safetensors). Self-contained — no project `src/` needed.
Requires: diffusers (git main / >=0.39), nvidia-modelopt, torch (cu128), safetensors, and a
**Blackwell GPU (sm_120 / RTX 5090)** — NVFP4 restore re-enters ModelOpt's FP4 path, which queries
`torch.cuda.get_device_capability` (FP4 needs sm ≥ 10.0), so a CUDA device must be present at load.
from load_quantized import load
pipe = load() # uses this dir, or pass a repo id / local dir
import torch
with torch.autocast("cuda", torch.bfloat16):
img = pipe("a corgi astronaut", num_frames=1, height=480, width=480).video[0][0]
Format (Path B; see ../docs/reports/session_4.md): the NVFP4 transformer is stored as **safetensors**
(`transformer/diffusion_pytorch_model.safetensors`: 505 weight-only NVFP4 weights — packed FP4 E2M1
`weight` (uint8) + per-block-16 E4M3 `weight_quantizer._scale` + per-tensor FP32
`weight_quantizer._double_scale` global scale + `_amax` + the AWQ `input_quantizer._pre_quant_scale`)
plus a tiny tensor-free `transformer/modelopt_state.pt` structural sidecar. The original
`transformer/modelopt_quantized.pt` is **retained** as a fallback (loadable via
`modelopt.torch.opt.restore`); this loader does NOT use it.
NVFP4 restore is **device-order-sensitive**: build the skeleton + restore on CPU, THEN move the whole
pipeline to the GPU (this loader does exactly that). Restoring directly onto a CUDA module splits the
NVFP4 sub-tensors across devices and trips dequantize.
SECURITY: `modelopt_state.pt` is loaded with `torch.load(weights_only=False)`, which executes
arbitrary pickle. Load this checkpoint ONLY from a source you trust (a tampered sidecar = remote
code execution). The safetensors weights themselves are safe; only the structural sidecar is pickle.
"""
import glob
import os
import torch
from diffusers import Cosmos3OmniPipeline, Cosmos3OmniTransformer, UniPCMultistepScheduler
import modelopt.torch.opt as mto
from safetensors.torch import load_file
def load_transformer(local):
"""Materialize the quantized transformer from safetensors + the structural sidecar (no `.pt`).
Built + restored + state-loaded on CPU; the caller moves the pipeline to the GPU afterward
(NVFP4 restore directly onto CUDA splits the packed-data/block-scale/global-scale across devices).
"""
cfg = {**Cosmos3OmniTransformer.load_config(f"{local}/transformer/config.json"), "action_gen": False}
tf = Cosmos3OmniTransformer.from_config(cfg).to(torch.bfloat16)
state = torch.load(f"{local}/transformer/modelopt_state.pt", weights_only=False)
restored = mto.restore_from_modelopt_state(tf, state)
if restored is not None:
tf = restored
tensors = {}
for shard in sorted(glob.glob(f"{local}/transformer/*.safetensors")):
tensors.update(load_file(shard))
tf.load_state_dict(tensors, strict=True)
return tf
def load(repo_or_dir=".", device="cuda"):
if os.path.isdir(repo_or_dir):
local = repo_or_dir
else:
from huggingface_hub import snapshot_download
local = snapshot_download(repo_or_dir)
tf = load_transformer(local) # on CPU
pipe = Cosmos3OmniPipeline.from_pretrained(
local, transformer=tf, torch_dtype=torch.bfloat16, enable_safety_checker=False
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=10.0)
return pipe.to(device) # move the whole pipeline (incl. the restored NVFP4 transformer) together
if __name__ == "__main__":
pipe = load()
with torch.autocast("cuda", dtype=torch.bfloat16): # required: float32 rotary tensors -> bf16 linears
img = pipe("A red panda astronaut floating in a nebula, highly detailed",
num_frames=1, height=480, width=480).video[0][0]
img.save("out.png")
print("saved out.png")
|