#!/usr/bin/env python
"""
Re-save the in-memory quantized Cosmos3-Super-Image2Video-4Step transformer in the
ROUND-TRIPPABLE ModelOpt HF format and assemble a complete drop-in diffusers repo
around it. Sibling of repackage_for_hf.py -- identical logic, just imports the
build helpers from serve_cosmos3_i2v4step_diffusers instead (which targets
nvidia/Cosmos3-Super-Image2Video-4Step and drops the scheduler swap this
checkpoint doesn't want -- see that file's docstring).

Run in the ModelOpt venv, from the directory containing serve_cosmos3_i2v4step_diffusers.py:
    CUDA_VISIBLE_DEVICES=0 python repackage_for_hf_i2v4step.py --format fp8 \
        --serve-dir ./cosmos3-i2v4step-fp8-serve \
        --out-dir   ./cosmos3-i2v4step-fp8-hf \
        [--cache ./cosmos3-i2v4step-cache]

Then verify the result loads + renders:
    python -i load_cosmos3_modelopt.py ./cosmos3-i2v4step-fp8-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

from serve_cosmos3_i2v4step_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="fp8")
    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()

    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)

    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)

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