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

"""Int8 ConvRot TE checkpoint, group-256 Hadamard; keys must match the strict load_state_dict the
loader expects; lm_head omitted, pipelines read hidden states only"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import sys
from pathlib import Path

BACKEND = Path(__file__).resolve().parent.parent / "studio" / "backend"
DECODER_LINEAR = re.compile(
    r"^model\.(language_model\.)?layers\.\d+\.(self_attn\.[qkvo]_proj|mlp\.(gate|up|down)_proj)\.weight$"
)
# UMT5 / T5 encoder blocks (Wan's UMT5-XXL): self-attention q/k/v/o and the MLP, gated (wi_0 / wi_1) or not (wi), + wo.
T5_ENCODER_LINEAR = re.compile(
    r"^encoder\.block\.\d+\.layer\.(0\.SelfAttention\.[qkvo]|1\.DenseReluDense\.(wi_0|wi_1|wi|wo))\.weight$"
)
_T5_CLASSES = ("UMT5EncoderModel", "T5EncoderModel")
# The encoder classes' _keep_in_fp32_modules: a projection kept dense is stored float32, as from_pretrained holds it.
_T5_KEEP_FP32 = ("wo",)


def _linear_pattern(te_class):
    return T5_ENCODER_LINEAR if te_class in _T5_CLASSES else DECODER_LINEAR


_ST_DTYPES = {"float32": "F32", "bfloat16": "BF16", "float16": "F16", "int8": "I8", "uint8": "U8"}


def _write_safetensors(path: Path, tensors: dict, metadata: dict) -> None:
    """safetensors with a canonical header (sorted keys) so a rebuild is byte-identical; ``save_file`` orders the
    metadata map differently per run. Widest dtype first, so every tensor starts at its own alignment."""
    import torch

    order = sorted(tensors, key = lambda n: (-tensors[n].element_size(), n))
    header: dict = {"__metadata__": dict(sorted(metadata.items()))}
    offset = 0
    for name in order:
        t = tensors[name]
        size = t.numel() * t.element_size()
        header[name] = {
            "dtype": _ST_DTYPES[str(t.dtype).replace("torch.", "")],
            "shape": list(t.shape),
            "data_offsets": [offset, offset + size],
        }
        offset += size
    blob = json.dumps(header, sort_keys = True, separators = (",", ":")).encode("utf-8")
    blob += b" " * (-len(blob) % 8)
    with open(path, "wb") as fh:
        fh.write(len(blob).to_bytes(8, "little"))
        fh.write(blob)
        for name in order:
            fh.write(tensors[name].contiguous().view(torch.uint8).numpy().tobytes())


def _download_component(base: str, component: str, token) -> Path:
    """The component folder of ``base``, data files only: the build reads JSON and safetensors (safe_open), never
    anything executable, so nothing else is fetched."""
    from huggingface_hub import snapshot_download

    root = snapshot_download(
        base,
        allow_patterns = [f"{component}/*.json", f"{component}/*.safetensors"],
        token = token,
    )
    return Path(root) / component


def main(argv = None) -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--base", required = True, help = "diffusers base repo carrying text_encoder/")
    p.add_argument("--family", required = True)
    p.add_argument("--component", default = "text_encoder")
    p.add_argument("--src", default = None, help = "local dense encoder folder (skips the download)")
    p.add_argument("--out", required = True, help = "output .safetensors path")
    p.add_argument("--group", type = int, default = 256)
    p.add_argument(
        "--device", default = "cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") != "" else "cpu"
    )
    p.add_argument("--hf-token", default = None)
    p.add_argument(
        "--load",
        choices = ("shards", "transformers", "ideogram4"),
        default = "shards",
        help = "ideogram4: build the encoder with Unsloth's Ideogram 4 loader (vendor fp8 + weight_scale "
        "dequantized to a bf16 Qwen3VLModel), then quantize it like --load transformers",
    )
    p.add_argument(
        "--te-class", default = None, help = "transformers class to load (--load transformers)"
    )
    p.add_argument(
        "--base-model-id", default = None, help = "base_model_id to record (default: --base)"
    )
    p.add_argument(
        "--keep-dense",
        default = "",
        help = "comma-separated Linear leaf names left dense (e.g. 'wo' for UMT5); fp32-kept modules stay float32",
    )
    args = p.parse_args(argv)
    keep_dense = tuple(t for t in (x.strip() for x in args.keep_dense.split(",")) if t)
    if args.load in ("transformers", "ideogram4"):
        return _main_transformers(args)

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

    from core.inference.diffusion_te_prequant import (
        TE_PREQUANT_FORMAT_INT8_CONVROT,
        quantize_int8_convrot_weight,
    )
    from core.inference.prequant_safetensors import UNSLOTH_FORMAT_KEY, UNSLOTH_METADATA_KEY

    src = (
        Path(args.src)
        if args.src
        else _download_component(args.base, args.component, args.hf_token)
    )
    config = json.loads((src / "config.json").read_text(encoding = "utf-8"))
    index_path = src / "model.safetensors.index.json"
    if index_path.is_file():
        weight_map = json.loads(index_path.read_text(encoding = "utf-8"))["weight_map"]
    else:
        with safe_open(str(src / "model.safetensors"), framework = "pt") as handle:
            weight_map = {name: "model.safetensors" for name in handle.keys()}
    te_class = (config.get("architectures") or [None])[0]
    pattern = _linear_pattern(te_class)
    keep_fp32 = _T5_KEEP_FP32 if te_class in _T5_CLASSES else ()
    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    quant = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": args.group}
    blob = torch.tensor(list(json.dumps(quant).encode("utf-8")), dtype = torch.uint8)
    tensors: dict = {}
    quantized = 0
    for shard in sorted(set(weight_map.values())):
        with safe_open(str(src / shard), framework = "pt", device = "cpu") as handle:
            for name in sorted(n for n, s in weight_map.items() if s == shard):
                if name == "lm_head.weight":
                    continue
                tensor = handle.get_tensor(name)
                leaf = (
                    name[: -len(".weight")].rsplit(".", 1)[-1] if name.endswith(".weight") else ""
                )
                if pattern.match(name) and leaf not in keep_dense:
                    q, scale = quantize_int8_convrot_weight(tensor.to(device), args.group)
                    prefix = name[: -len(".weight")]
                    tensors[name] = q.cpu().contiguous()
                    tensors[prefix + ".weight_scale"] = scale.cpu().contiguous()
                    tensors[prefix + ".comfy_quant"] = blob.clone()
                    quantized += 1
                elif pattern.match(name) and leaf in keep_fp32:
                    tensors[name] = tensor.to(torch.float32).contiguous()
                else:
                    tensors[name] = tensor.to(torch.bfloat16).contiguous()
    if not quantized:
        raise SystemExit("no projections matched; is this a Qwen3 / Qwen3-VL or UMT5 text encoder?")
    metadata = {
        "base_model_id": args.base_model_id or args.base,
        "family": args.family,
        "scheme": "int8",
        "component": args.component,
        "te_class": te_class,
        "torch_dtype": "bfloat16",
        "quant": dict(
            quant,
            weight_only = True,
            scale = "per_output_channel_absmax_over_127",
            linears = quantized,
            pattern = pattern.pattern,
            kept_dense = list(keep_dense),
        ),
        "lm_head": "dropped" if te_class not in _T5_CLASSES else "absent",
        "torch_version": torch.__version__,
        "transformers_version": transformers.__version__,
        "safetensors_version": safetensors.__version__,
    }
    names = sorted(tensors)
    header = {
        UNSLOTH_FORMAT_KEY: TE_PREQUANT_FORMAT_INT8_CONVROT,
        UNSLOTH_METADATA_KEY: json.dumps(metadata, sort_keys = True),
        "tensor_names": json.dumps(names),
    }
    header.update({name: json.dumps({"_type": "Tensor"}) for name in names})
    out = Path(args.out)
    out.parent.mkdir(parents = True, exist_ok = True)
    tmp = out.with_name(out.name + ".tmp")
    _write_safetensors(tmp, tensors, header)
    os.replace(tmp, out)
    digest = hashlib.sha256()
    with open(out, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 24), b""):
            digest.update(chunk)
    print(
        json.dumps(
            {
                "out": str(out),
                "bytes": out.stat().st_size,
                "sha256": digest.hexdigest(),
                "quantized_linears": quantized,
                "tensors": len(names),
            }
        )
    )
    return 0


def _main_transformers(args) -> int:
    """``--load transformers``: quantize the runtime state dict of the encoder class the pipeline builds."""
    sys.path.insert(0, str(BACKEND))
    import safetensors
    import torch
    import transformers

    from core.inference.diffusion_te_prequant import (
        TE_INT8_CONVROT_LINEAR,
        TE_PREQUANT_FORMAT_INT8_CONVROT,
        quantize_int8_convrot_weight,
    )
    from core.inference.prequant_safetensors import UNSLOTH_FORMAT_KEY, UNSLOTH_METADATA_KEY

    src = args.src or args.base
    if args.load == "ideogram4":
        if args.component != "text_encoder":
            raise SystemExit("--load ideogram4 reads <base>/text_encoder only")
        from core.inference.diffusion_ideogram4 import load_ideogram4_text_encoder

        # The exact module the runtime builds (from_pretrained cannot read the vendor fp8 layout).
        encoder = load_ideogram4_text_encoder(src, torch.bfloat16, hf_token = args.hf_token).eval()
        te_class = type(encoder).__name__
    else:
        kwargs = {"token": args.hf_token, "torch_dtype": torch.bfloat16}
        if not args.src:
            kwargs["subfolder"] = args.component
        config = transformers.AutoConfig.from_pretrained(
            src, **{k: v for k, v in kwargs.items() if k != "torch_dtype"}
        )
        te_class = args.te_class or (getattr(config, "architectures", None) or [None])[0]
        encoder = getattr(transformers, te_class).from_pretrained(src, **kwargs).eval()
    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    quant = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": args.group}
    blob = torch.tensor(list(json.dumps(quant).encode("utf-8")), dtype = torch.uint8)
    keep_dense = tuple(t for t in (x.strip() for x in args.keep_dense.split(",")) if t)
    linears = {
        name
        for name, module in encoder.named_modules()
        if isinstance(module, torch.nn.Linear)
        and TE_INT8_CONVROT_LINEAR.search(name)
        and name.rsplit(".", 1)[-1] not in keep_dense
    }
    # The encoder's own fp32-kept modules (T5's ``wo``) are stored float32 when left dense, as the shard path does.
    keep_fp32 = tuple(str(m) for m in (getattr(encoder, "_keep_in_fp32_modules", None) or ()))
    tensors: dict = {}
    for name, tensor in encoder.state_dict().items():
        if name == "lm_head.weight":
            continue
        prefix = name[: -len(".weight")] if name.endswith(".weight") else None
        if prefix in linears:
            q, scale = quantize_int8_convrot_weight(tensor.to(device), args.group)
            tensors[name] = q.cpu().contiguous()
            tensors[prefix + ".weight_scale"] = scale.cpu().contiguous()
            tensors[prefix + ".comfy_quant"] = blob.clone()
        else:
            if not tensor.is_floating_point():
                dtype = tensor.dtype
            elif any(t in name.split(".") for t in keep_fp32):
                dtype = torch.float32
            else:
                dtype = torch.bfloat16
            tensors[name] = tensor.detach().to("cpu", dtype)
            tensors[name] = tensors[name].contiguous().clone()
    if not linears:
        raise SystemExit(f"no Linear in {te_class} matched TE_INT8_CONVROT_LINEAR")
    metadata = {
        "base_model_id": args.base_model_id or args.base,
        "family": args.family,
        "scheme": "int8",
        "component": args.component,
        "te_class": te_class,
        "torch_dtype": "bfloat16",
        "quant": dict(
            quant,
            weight_only = True,
            scale = "per_output_channel_absmax_over_127",
            linears = len(linears),
        ),
        "lm_head": "dropped",
        "torch_version": torch.__version__,
        "transformers_version": transformers.__version__,
        "safetensors_version": safetensors.__version__,
    }
    names = sorted(tensors)
    header = {
        UNSLOTH_FORMAT_KEY: TE_PREQUANT_FORMAT_INT8_CONVROT,
        UNSLOTH_METADATA_KEY: json.dumps(metadata, sort_keys = True),
        "tensor_names": json.dumps(names),
    }
    header.update({name: json.dumps({"_type": "Tensor"}) for name in names})
    out = Path(args.out)
    out.parent.mkdir(parents = True, exist_ok = True)
    tmp = out.with_name(out.name + ".tmp")
    _write_safetensors(tmp, tensors, header)
    os.replace(tmp, out)
    digest = hashlib.sha256()
    with open(out, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 24), b""):
            digest.update(chunk)
    print(
        json.dumps(
            {
                "out": str(out),
                "bytes": out.stat().st_size,
                "sha256": digest.hexdigest(),
                "te_class": te_class,
                "quantized_linears": len(linears),
                "tensors": len(names),
            }
        )
    )
    return 0


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