#!/usr/bin/env python
"""
Re-save the in-memory quantized Cosmos3-Super transformer in the ROUND-TRIPPABLE
ModelOpt HF format and assemble a complete drop-in diffusers repo around it.

WHY THIS EXISTS
---------------
`export_hf_checkpoint()` writes a *deployment* checkpoint (hf_quant_config.json +
packed weights) meant for TRT-LLM / vLLM, which load the packed format with their
own kernels. diffusers' `from_pretrained` does NOT reconstruct the quantized
modules from that. Instead, on load, ModelOpt restores the quantized module
structure from a file named `modelopt_state.pth` in the checkpoint dir (via
`restore_from_modelopt_state`). That file is written by `save_pretrained()` once
`enable_huggingface_checkpointing()` has been called -- and is NOT written by
`export_hf_checkpoint()`. So for a diffusers drop-in repo, re-save with
`save_pretrained`. This script does that and then copies the non-transformer
pipeline components (VAE, tokenizers, model_index.json) from an existing assembled
dir, producing a self-contained repo you can publish.

Run in the ModelOpt venv, from the directory containing serve_cosmos3_diffusers.py:
    CUDA_VISIBLE_DEVICES=0 python repackage_for_hf.py --format nvfp4 \
        --serve-dir ./cosmos3-super-nvfp4-serve \
        --out-dir   ./cosmos3-super-nvfp4-hf \
        [--cache ./cosmos3-cache]        # optional: faster rebuild if you have it

Then verify the result loads + renders:
    python -i load_cosmos3_modelopt.py ./cosmos3-super-nvfp4-hf
"""
import argparse
import json
import os
import pathlib
import shutil

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch  # noqa: F401  (ensures CUDA init / dtype availability)
from modelopt.torch.opt import enable_huggingface_checkpointing

# Reuse the proven in-memory build (empty-on-meta -> quantize -> compress -> stream).
from serve_cosmos3_diffusers import build_quantized_transformer, try_restore_quantized

QUANT_TYPE = {"fp8": "FP8", "nvfp4": "NVFP4"}


def ensure_loadable_config(transformer_dir: str, fmt: str) -> None:
    """If save_pretrained wrote a quantization_config, make sure diffusers can construct
    it: NVIDIAModelOptConfig needs `quant_type`, and a truthy `modelopt_config` avoids the
    buggy get_config_from_quant_type() builder. (Structure restore itself comes from
    modelopt_state.pth; this just keeps config parsing from crashing on load.)"""
    cfg_path = pathlib.Path(transformer_dir) / "config.json"
    cfg = json.loads(cfg_path.read_text())
    qc = cfg.get("quantization_config")
    if isinstance(qc, dict):
        qc["quant_type"] = QUANT_TYPE[fmt]
        qc.setdefault("weight_only", True)
        if not qc.get("modelopt_config"):
            qc["modelopt_config"] = {"quant_cfg": {}, "algorithm": "max"}
        cfg["quantization_config"] = qc
        cfg_path.write_text(json.dumps(cfg, indent=2))
        print(f"[patch] quantization_config made loadable (quant_type={qc['quant_type']})")
    else:
        print("[patch] no embedded quantization_config; relying on modelopt_state.pth for restore")


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--format", choices=["fp8", "nvfp4"], default="nvfp4")
    ap.add_argument("--serve-dir", required=True,
                    help="existing assembled pipeline dir (source of VAE / tokenizers / model_index.json)")
    ap.add_argument("--out-dir", required=True, help="new drop-in repo dir to create")
    ap.add_argument("--cache", default=None, help="optional mto cache dir for a faster rebuild")
    ap.add_argument("--gpu-mem-fraction", type=float, default=0.85)
    args = ap.parse_args()

    # 1. quantized transformer in memory (restore from cache if available, else rebuild)
    model = try_restore_quantized(args.format, args.cache) if args.cache else None
    if model is None:
        model = build_quantized_transformer(args.format, args.gpu_mem_fraction)

    # 2. enable round-trippable HF checkpointing, then save the transformer
    enable_huggingface_checkpointing()
    tdir = os.path.join(args.out_dir, "transformer")
    os.makedirs(tdir, exist_ok=True)
    print(f"[save] writing round-trippable transformer (+ modelopt_state.pth) -> {tdir}")
    model.save_pretrained(tdir)
    state_file = os.path.join(tdir, "modelopt_state.pth")
    assert os.path.isfile(state_file), (
        f"expected {state_file} to exist -- enable_huggingface_checkpointing() must run "
        "before save_pretrained(); without modelopt_state.pth the repo won't load in diffusers"
    )
    ensure_loadable_config(tdir, args.format)

    # 3. copy the rest of the pipeline (everything except transformer/) from the serve dir
    print(f"[assemble] copying non-transformer components from {args.serve_dir}")
    for name in os.listdir(args.serve_dir):
        if name == "transformer":
            continue
        src = os.path.join(args.serve_dir, name)
        dst = os.path.join(args.out_dir, name)
        if os.path.isdir(src):
            shutil.copytree(src, dst, dirs_exist_ok=True)
        else:
            shutil.copy2(src, dst)

    print(f"[done] drop-in repo -> {args.out_dir}")
    print(f"       verify:  python -i load_cosmos3_modelopt.py {args.out_dir}")


if __name__ == "__main__":
    main()