"""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")