# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""Pre-casts the TE to layerwise fp8 storage once; Ideogram 4's vendor fp8 needs --loader ideogram4."""

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

BACKEND = Path(__file__).resolve().parent.parent / "studio" / "backend"


def main(argv = None) -> int:
    p = argparse.ArgumentParser()
    p.add_argument(
        "--base", required = True, help = "diffusers base repo (carries the component subfolder)"
    )
    p.add_argument("--family", required = True, help = "diffusion/video family name or alias")
    p.add_argument(
        "--component",
        default = "text_encoder",
        help = "pipeline component attribute (also the repo subfolder)",
    )
    p.add_argument(
        "--config-subfolder",
        default = None,
        help = "where the encoder lives inside --base (default: the component name; "
        "pass '' for a standalone encoder repo whose config sits at the root, "
        "e.g. HiDream's Llama text_encoder_4)",
    )
    p.add_argument("--scheme", default = "fp8", choices = ["fp8"])
    p.add_argument("--out", required = True, help = "output .pt path for the checkpoint")
    p.add_argument("--dtype", default = "bfloat16", choices = ["bfloat16"])
    p.add_argument("--hf-token", default = None)
    p.add_argument(
        "--base-model-id",
        default = None,
        help = "base_model_id to record (default: --base); the upstream id when --base is an ungated mirror",
    )
    p.add_argument(
        "--loader",
        default = "auto",
        choices = ["auto", "ideogram4"],
        help = "how to load the dense encoder: 'auto' = transformers from_pretrained on the config's "
        "architecture; 'ideogram4' = Unsloth's Ideogram 4 loader, which dequantizes the vendor "
        "fp8 + per-channel weight_scale Qwen3-VL encoder to bf16 (from_pretrained cannot read it)",
    )
    args = p.parse_args(argv)

    sys.path.insert(0, str(BACKEND))
    import torch
    import transformers

    from core.inference.diffusion_precision import _cast_fp8
    from core.inference.diffusion_te_prequant import TE_PREQUANT_FORMAT

    family = args.family.strip().lower()

    subfolder = args.component if args.config_subfolder is None else args.config_subfolder
    from_pretrained_kwargs = {"token": args.hf_token}
    if subfolder:
        from_pretrained_kwargs["subfolder"] = subfolder

    print(f"== build TE prequant ({family}/{args.component}/{args.scheme}) ==", flush = True)
    print(f"  loading dense encoder from {args.base} (subfolder={subfolder!r}) ...", flush = True)
    t0 = time.time()
    if args.loader == "ideogram4":
        if args.component != "text_encoder" or subfolder != "text_encoder":
            p.error("--loader ideogram4 reads <base>/text_encoder only")
        from core.inference.diffusion_ideogram4 import load_ideogram4_text_encoder

        # Same module Studio builds, so the cast below matches its dense-load-then-cast.
        encoder = load_ideogram4_text_encoder(args.base, torch.bfloat16, hf_token = args.hf_token)
        encoder_cls_name = type(encoder).__name__
    else:
        config = transformers.AutoConfig.from_pretrained(args.base, **from_pretrained_kwargs)
        # Prefer the checkpoint's own architecture; AutoModel.from_config gives an unusable bare base class.
        arch = (getattr(config, "architectures", None) or [None])[0]
        if arch and hasattr(transformers, arch):
            encoder_cls_name = arch
        else:
            encoder = transformers.AutoModel.from_config(config)
            encoder_cls_name = type(encoder).__name__
            del encoder
        encoder = getattr(transformers, encoder_cls_name).from_pretrained(
            args.base,
            torch_dtype = torch.bfloat16,
            **from_pretrained_kwargs,
        )
    print(f"  casting in place (layerwise {args.scheme}) ...", flush = True)

    class _Target:
        dtype = torch.bfloat16

    _cast_fp8(encoder, _Target())

    state_dict = {
        k: (v.detach().to("cpu") if hasattr(v, "detach") else v)
        for k, v in encoder.state_dict().items()
    }
    metadata = {
        "base_model_id": args.base_model_id or args.base,
        "family": family,
        "scheme": args.scheme,
        "component": args.component,
        "te_class": encoder_cls_name,
        "torch_dtype": args.dtype,
        "cast_backend": "diffusers_layerwise",
        # str(): a pickled TorchVersion makes torch.load(weights_only=True) reject the artifact.
        "torch_version": str(torch.__version__),
        "transformers_version": str(transformers.__version__),
    }
    ckpt = {
        "format": TE_PREQUANT_FORMAT,
        "metadata": metadata,
        "state_dict": state_dict,
    }
    out = Path(args.out)
    out.parent.mkdir(parents = True, exist_ok = True)
    torch.save(ckpt, out)
    size_gb = out.stat().st_size / 1e9
    print(f"  saved {out}  ({size_gb:.2f} GB) in {time.time() - t0:.0f}s", flush = True)
    print(f"  metadata: {metadata}", flush = True)
    print("BUILD-TE-PREQUANT-DONE", flush = True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())