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

"""Convert a published ``.pt`` pre-quant checkpoint (denoiser or text encoder) to ``.safetensors``.

    python scripts/convert_prequant_to_safetensors.py SRC.pt [SRC2.pt ...] --out-dir DIR \
        [--repo owner/name --revision SHA]

Writes ``DIR/<owner>/<name>/<same stem>.safetensors`` (``DIR/<stem>.safetensors`` without ``--repo``), the
layout ``UNSLOTH_DIFFUSION_PREQUANT_MIRROR`` reads. Nothing is uploaded.

Denoisers use the container Qwen-Image-2.1's published artifacts already use (``prequant_safetensors``:
torchao's flattened tensors plus ``unsloth_format`` / ``unsloth_metadata``), so a released Studio that
resolves ``<Model>-<SCHEME>.safetensors`` reads it too. Two header keys are added: ``unsloth_quant_layout``
(how to rebuild every weight from its plain tensors without torchao, and for int8 that the weights came
from torchao v1, so torchao <= 0.17 rebuilds the v1 class the ``.pt`` held) and ``unsloth_source`` (the
file, its sha256 and the repo revision it was converted from). ``unsloth_metadata`` is copied unchanged.

Text encoders are plain fp8 / bf16 tensors and are written in the plain layout
``load_plain_prequant_safetensors`` reads, with no torchao involved.

Every weight is checked after writing: the file is read back through Studio's loader and each tensor
compared bit for bit with the ``.pt`` as the same install reads it.
"""

from __future__ import annotations

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

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


def sha256_of(path: str, chunk: int = 64 << 20) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as fh:
        while True:
            block = fh.read(chunk)
            if not block:
                break
            digest.update(block)
    return digest.hexdigest()


def _record_v1_facts(facts: list):
    """Wrap the legacy decoder so each v1 weight's facts are recorded before it is rebuilt (torchao >= 0.18).

    Patches a module global for the duration of one read, so conversions must not run concurrently in one
    process (this is a single-threaded CLI)."""
    from core.inference import prequant_legacy_int8 as legacy

    original = legacy._rebuild_weight

    def recording(name, w, standins, api):
        facts.append(_v1_facts_of(w, standins))
        return original(name, w, standins, api)

    legacy._rebuild_weight = recording
    return lambda: setattr(legacy, "_rebuild_weight", original)


def _v1_facts_of(w, standins = None) -> dict:
    """The INT8_V1_FACTS view of one v1 weight (a real torchao object or the decoder's stand-in)."""
    from core.inference.prequant_legacy_int8 import _ACT_QUANT, _LAYOUT

    state = getattr(w, "__dict__", {})
    aqt = state.get("original_weight_tensor")
    impl = getattr(aqt, "tensor_impl", None)
    act = state.get("input_quant_func")
    scale = getattr(impl, "scale", None)
    zero_point = getattr(impl, "zero_point", None)
    domain = getattr(aqt, "zero_point_domain", None)
    layout = getattr(impl, "_layout", None)
    if standins is not None:
        # torchao >= 0.18 read the pickle into stand-ins registered under the v1 names: report those names.
        if act is standins.get(_ACT_QUANT):
            act = _ACT_QUANT
        if isinstance(layout, standins.get(_LAYOUT, ())):
            layout = _LAYOUT
    return {
        "act_quant": (act if isinstance(act, str) else getattr(act, "__name__", str(act))).rsplit(
            ".", 1
        )[-1],
        "quant_kwargs": dict(state.get("quant_kwargs") or {}),
        "zero_point": None if zero_point is None else "tensor",
        "zero_point_domain": getattr(domain, "name", str(domain)),
        "quant_min": getattr(aqt, "quant_min", "?"),
        "quant_max": getattr(aqt, "quant_max", "?"),
        "scale_shape": "out" if scale is not None and scale.dim() == 1 else "other",
        "layout": layout.rsplit(".", 1)[-1] if isinstance(layout, str) else type(layout).__name__,
    }


def _to_int8_tensor_dict(state_dict: dict, facts: list) -> dict:
    """torchao <= 0.17 hands back v1 objects: record their facts and rebuild them as Int8Tensor to flatten."""
    from core.inference.prequant_legacy_int8 import (
        LEGACY_INT8_CLASS_NAMES,
        _ACT_QUANT,
        _int8_tensor_api,
        _rebuild_weight,
        _resolve,
    )

    classes = {n: _resolve(n) for n in (*LEGACY_INT8_CLASS_NAMES, _ACT_QUANT)}
    laqt = classes[LEGACY_INT8_CLASS_NAMES[0]]
    if laqt is None:
        return state_dict
    api = _int8_tensor_api()
    out = {}
    for key, value in state_dict.items():
        if isinstance(value, laqt):
            facts.append(_v1_facts_of(value))
            out[key] = _rebuild_weight(key, value, classes, api)
        else:
            out[key] = value
    return out


def _canonical(value):
    """``(class name, {attr: tensor or value})`` of a weight, recursing into torchao subclasses."""
    import torch

    if type(value) is torch.Tensor:
        return ("Tensor", value)
    flatten = getattr(value, "__tensor_flatten__", None)
    if flatten is None:
        return (type(value).__name__, value)
    names, ctx = flatten()
    return (type(value).__name__, {n: _canonical(getattr(value, n)) for n in names}, repr(ctx))


def _same(a, b) -> bool:
    import torch

    if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor):
        if a.dtype != b.dtype or tuple(a.shape) != tuple(b.shape):
            return False
        # Raw bytes, not torch.equal: -0.0 == 0.0 and NaN != NaN would otherwise decide. Flattened first:
        # a 0-dim tensor (nvfp4's per-tensor scale) cannot be viewed as bytes of another element size.
        return a.numel() == 0 or torch.equal(
            a.detach().contiguous().cpu().reshape(-1).view(torch.uint8),
            b.detach().contiguous().cpu().reshape(-1).view(torch.uint8),
        )
    if isinstance(a, tuple) and isinstance(b, tuple):
        return len(a) == len(b) and all(_same(x, y) for x, y in zip(a, b, strict = True))
    if isinstance(a, dict) and isinstance(b, dict):
        return a.keys() == b.keys() and all(_same(a[k], b[k]) for k in a)
    return a == b


def compare_state_dicts(reference: dict, candidate: dict) -> list:
    """Keys whose class, attributes or tensor bytes differ (empty = bit-identical)."""
    bad = sorted(set(reference) ^ set(candidate))
    for key in reference:
        if key in candidate and not _same(_canonical(reference[key]), _canonical(candidate[key])):
            bad.append(key)
    return bad


def _untie(flat: dict) -> dict:
    """safetensors refuses tensors that share storage (tied embeddings): write each one as its own copy.

    Studio's text-encoder loader re-ties after ``load_state_dict`` (``tie_weights``), exactly as it does for the
    published safetensors encoders, so the copy costs disk only."""
    seen: set = set()
    out = {}
    for key, value in flat.items():
        value = value.contiguous()
        ptr = value.untyped_storage().data_ptr() if value.numel() else None
        if ptr is not None and ptr in seen:
            value = value.clone()
        elif ptr is not None:
            seen.add(ptr)
        out[key] = value
    return out


def tied_copies(state_dict: dict) -> dict:
    """``{alias: source}`` for keys that share storage, shape and stride with an earlier key (a tied embedding)."""
    seen: dict = {}
    aliases = {}
    for key, value in state_dict.items():
        if not getattr(value, "numel", lambda: 0)():
            continue
        ident = (
            value.untyped_storage().data_ptr(),
            value.storage_offset(),
            tuple(value.shape),
            tuple(value.stride()),
            value.dtype,
        )
        if ident in seen:
            aliases[key] = seen[ident]
        else:
            seen[ident] = key
    return aliases


def write_plain_safetensors(
    path: str, *, fmt: str, state_dict: dict, metadata: dict, extra: dict
) -> None:
    """The torchao-free plain layout ``load_plain_prequant_safetensors`` reads (text encoders)."""
    import torch
    from safetensors.torch import save_file

    from core.inference.prequant_safetensors import (
        UNSLOTH_FORMAT_KEY,
        UNSLOTH_METADATA_KEY,
        UNSLOTH_ROOT_KEYS_KEY,
        UNSLOTH_ROOT_PREFIX,
    )

    names, roots, flat = [], [], {}
    for key, value in state_dict.items():
        if type(value) is not torch.Tensor:
            raise ValueError(f"{key!r} is a {type(value).__name__}, not a plain tensor")
        if "." in key:
            names.append(key)
            flat[key] = value.contiguous()
        else:
            roots.append(key)
            flat[f"{UNSLOTH_ROOT_PREFIX}{key}"] = value.contiguous()
    header = {name: json.dumps({"_type": "Tensor"}) for name in names}
    header["tensor_names"] = json.dumps(names)
    header[UNSLOTH_FORMAT_KEY] = fmt
    header[UNSLOTH_METADATA_KEY] = json.dumps(metadata, default = str)
    if roots:
        header[UNSLOTH_ROOT_KEYS_KEY] = json.dumps(roots)
    header.update(extra)
    save_file(_untie(flat), path, metadata = header)


def convert(
    src: str,
    out_dir: str,
    *,
    repo: str | None,
    revision: str | None,
    sha: str | None,
    drop_tied: bool = False,
) -> dict:
    import torch

    from core.inference import diffusion_prequant as dp
    from core.inference import prequant_safetensors as ps
    from core.inference.diffusion_te_prequant import TE_PREQUANT_FORMAT
    from core.inference.prequant_native import (
        INT8_SOURCE_V1,
        INT8_V1_FACTS,
        QUANT_LAYOUT_KEY,
        SOURCE_KEY,
        quant_layout_header,
    )

    t0 = time.time()
    name = os.path.basename(src)
    stem = os.path.splitext(name)[0] if name.endswith((".pt", ".pth")) else name
    dest_dir = os.path.join(out_dir, *repo.split("/")) if repo else out_dir
    os.makedirs(dest_dir, exist_ok = True)
    dst = os.path.join(dest_dir, stem + ".safetensors")
    source = {"file": name, "sha256": sha or sha256_of(src), "bytes": os.path.getsize(src)}
    if repo:
        source["repo"] = repo
    if revision:
        source["revision"] = revision
    try:
        import torchao
        source["converter_torchao"] = torchao.__version__
    except Exception:  # noqa: BLE001
        pass
    source["converter_torch"] = torch.__version__

    # Read as Studio does (allowlisted weights_only); on torchao >= 0.18 v1 int8 loads as Int8Tensor.
    facts: list = []
    restore = _record_v1_facts(facts)
    try:
        if str(src).lower().endswith(".safetensors"):
            ckpt = dp._load_prequant_checkpoint(src, map_location = "cpu", mmap = True)
        else:
            # Bypasses Studio's nvfp4 pickle refusal: the only way a local nvfp4 .pt becomes .safetensors.
            ckpt = dp._torch_load_prequant(src, map_location = "cpu", mmap = True)
    finally:
        restore()
    fmt = ckpt.get("format")
    is_te = fmt == TE_PREQUANT_FORMAT
    metadata = dict(ckpt.get("metadata") or {})
    state_dict = ckpt["state_dict"]
    extra = {SOURCE_KEY: json.dumps(source)}
    dropped: dict = {}
    if is_te:
        written = state_dict
        if drop_tied:
            # Store a tied embedding once (UMT5's encoder.embed_tokens.weight is shared.weight): the TE loader fills
            # the alias from the model's own tie (diffusion_te_prequant._fill_tied_aliases). Older readers refuse such
            # a file and load the dense encoder, so only use it for an encoder no older release was pointed at.
            # A tied lm_head stays: the loader never fills a head alias (whether a head is tied is per config).
            dropped = {
                alias: key
                for alias, key in tied_copies(state_dict).items()
                if not alias.endswith("lm_head.weight")
            }
            written = {k: v for k, v in state_dict.items() if k not in dropped}
        write_plain_safetensors(dst, fmt = fmt, state_dict = written, metadata = metadata, extra = extra)
        back = ps.load_plain_prequant_safetensors(dst)
        for alias, source_key in dropped.items():
            back["state_dict"][alias] = back["state_dict"][source_key]
    else:
        if fmt not in dp.PREQUANT_FORMATS:
            raise ValueError(f"{src}: unrecognised format {fmt!r}")
        flat_sd = _to_int8_tensor_dict(state_dict, facts)
        int8_source = None
        if facts:
            int8 = sum(type(v).__name__ == "Int8Tensor" for v in flat_sd.values())
            if len(facts) != int8:
                raise ValueError(
                    f"{src}: {int8} int8 weights but {len(facts)} came from torchao v1"
                )
            distinct = {json.dumps(f, sort_keys = True, default = str) for f in facts}
            if distinct != {json.dumps(INT8_V1_FACTS, sort_keys = True)}:
                raise ValueError(
                    f"{src}: v1 int8 weights do not match the supported v1 layout: {sorted(distinct)[:3]}"
                )
            int8_source = INT8_SOURCE_V1
        layout = quant_layout_header(flat_sd, int8_source = int8_source)
        from safetensors.torch import save_file

        helpers = ps._torchao_helpers()
        if helpers is None:
            raise RuntimeError("torchao >= 0.16 is required to write the flattened layout")
        undotted = ps.unsupported_state_dict_keys(flat_sd)
        if undotted:
            raise ValueError(
                f"{src}: root-level quantized weights cannot be flattened: {undotted[:8]}"
            )
        roots = ps._root_level_keys(flat_sd)
        quantizable = {k: v for k, v in flat_sd.items() if k not in set(roots)}
        flat, torchao_metadata = helpers[0](quantizable)
        header = dict(torchao_metadata or {})
        header[ps.UNSLOTH_FORMAT_KEY] = str(fmt)
        header[ps.UNSLOTH_METADATA_KEY] = json.dumps(metadata, default = str)
        header[QUANT_LAYOUT_KEY] = json.dumps(layout)
        header.update(extra)
        if roots:
            flat = dict(flat)
            for key in roots:
                flat[f"{ps.UNSLOTH_ROOT_PREFIX}{key}"] = flat_sd[key].contiguous()
            header[ps.UNSLOTH_ROOT_KEYS_KEY] = json.dumps([str(k) for k in roots])
        save_file(_untie(dict(flat)), dst, metadata = header)
        back = ps.load_prequant_safetensors(dst, mmap = True)
    # Bit-identity against the .pt as this install reads it.
    bad = compare_state_dicts(state_dict, back["state_dict"])
    if bad:
        raise AssertionError(f"{dst}: {len(bad)} tensors differ from {src}, e.g. {bad[:3]}")
    if back["format"] != fmt or back["metadata"] != metadata:
        raise AssertionError(f"{dst}: format or metadata changed in the round trip")
    rec = {
        "src": src,
        "dst": dst,
        "kind": "text_encoder" if is_te else "transformer",
        "format": fmt,
        "scheme": metadata.get("scheme"),
        "tensors": len(state_dict),
        "quantized": sum(1 for v in state_dict.values() if type(v) is not torch.Tensor),
        "int8_source": None if is_te else layout.get("int8_source"),
        "reader": back.get("reader"),
        "src_bytes": source["bytes"],
        "src_sha256": source["sha256"],
        "dst_bytes": os.path.getsize(dst),
        "dst_sha256": sha256_of(dst),
        "seconds": round(time.time() - t0, 1),
        "verified_bit_identical": True,
        "tied_copies_dropped": dropped,
    }
    return rec


def main(argv = None) -> int:
    p = argparse.ArgumentParser(description = __doc__.split("\n\n")[0])
    p.add_argument("src", nargs = "+")
    p.add_argument("--out-dir", required = True)
    p.add_argument(
        "--repo", default = None, help = "owner/name the .pt came from (mirror layout + provenance)"
    )
    p.add_argument("--revision", default = None)
    p.add_argument("--sha256", default = None, help = "known sha256 of a single SRC (skips hashing it)")
    p.add_argument("--report", default = None, help = "append one JSON line per file here")
    p.add_argument(
        "--drop-tied-copies",
        action = "store_true",
        help = "text encoders: store a tied tensor once (the loader re-fills the alias from the model's tie)",
    )
    args = p.parse_args(argv)
    sys.path.insert(0, str(BACKEND))
    for src in args.src:
        rec = convert(
            src,
            args.out_dir,
            repo = args.repo,
            revision = args.revision,
            sha = args.sha256 if len(args.src) == 1 else None,
            drop_tied = args.drop_tied_copies,
        )
        line = json.dumps(rec)
        print(line, flush = True)
        if args.report:
            with open(args.report, "a") as fh:
                fh.write(line + "\n")
    return 0


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