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

"""Builds the int8 ConvRot Gemma 4 encoder for LTX-2.5's sidecar; keys use the loaded model's names."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
from pathlib import Path
from typing import Optional

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


def main(argv = None) -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--base", default = "Lightricks/LTX-2.5-Diffusers")
    p.add_argument(
        "--repo", default = None, help = "local pack folder or Hub id to read (default: --base)"
    )
    p.add_argument("--out", required = True)
    p.add_argument(
        "--device", default = "cuda:0", help = "the at-load path quantizes on the card when it has one"
    )
    p.add_argument(
        "--venv-dir", default = None, help = "the transformers 5.10 sidecar (default: Unsloth's)"
    )
    p.add_argument("--hf-token", default = None)
    args = p.parse_args(argv)
    sys.path.insert(0, str(BACKEND))
    venv_dir = args.venv_dir
    if venv_dir is None:
        from utils.transformers_version import _VENV_T5_510_DIR
        venv_dir = _VENV_T5_510_DIR
    result = build(
        args.repo or args.base,
        Path(args.out),
        base = args.base,
        device = args.device,
        venv_dir = venv_dir,
        hf_token = args.hf_token,
    )
    print(json.dumps(result))
    return 0


def build(
    repo: str,
    out: Path,
    *,
    base: str,
    device: str,
    venv_dir: Optional[str] = None,
    hf_token: Optional[str] = None,
) -> dict:
    """Load ``repo``'s encoder the worker's way, quantized, and write it to ``out``; its size, sha256 and counts."""
    sys.path.insert(0, str(BACKEND))
    sys.path.insert(0, str(ROOT / "scripts"))
    from core.inference import video_ltx25_te_worker as worker

    if venv_dir:
        worker._activate_sidecar(venv_dir)
    import safetensors
    import torch

    from build_te_int8_convrot_checkpoint import _write_safetensors
    from core.inference.prequant_safetensors import UNSLOTH_FORMAT_KEY, UNSLOTH_METADATA_KEY
    from core.inference.video_ltx25 import LTX25_TE_HOSTED_FORMAT
    from core.inference.video_minimax_h3_te import _int8_convrot_linear_class

    device = _available_device(torch, device)
    on_card = device.startswith(("cuda", "xpu"))
    config = {
        "repo": repo,
        "hf_token": hf_token,
        "device": device,
        "dtype": "bfloat16",
        "resident": on_card,
        "te_quant": worker.TE_QUANT_INT8_CONVROT,
    }
    _tokenizer, model = worker._load(config)
    if config.get("_te_source") != "quantized_at_load":
        raise SystemExit(f"expected the at-load build, got {config.get('_te_source')!r}")
    quantized = sum(1 for m in model.modules() if isinstance(m, _int8_convrot_linear_class()))
    state = model.state_dict()
    # Tied lm_head is the embedding again: the loader re-ties it, so drop it before shipping.
    if _head_is_tied(state):
        state.pop("lm_head.weight")
    tensors = {k: v.detach().to("cpu").contiguous() for k, v in state.items()}
    del model
    names = sorted(tensors)
    try:
        import transformers

        cfg = transformers.AutoConfig.from_pretrained(
            repo, subfolder = "text_encoder", token = hf_token
        )
        te_class = (getattr(cfg, "architectures", None) or [None])[0]
        transformers_version = transformers.__version__
    except Exception:  # noqa: BLE001 -- forensic metadata only
        te_class, transformers_version = None, None
    metadata = {
        "base_model_id": base,
        "family": "ltx-2.5",
        "scheme": "int8",
        "component": "text_encoder",
        "te_class": te_class or "Gemma4UnifiedForConditionalGeneration",
        "torch_dtype": "bfloat16",
        "quant": {
            "format": "int8_convrot",
            "group_size": worker.CONVROT_GROUP,
            "weight_only": True,
            "scale": "per_output_channel_absmax_over_127",
            "linears": quantized,
            "quantized_on": torch.device(device).type if on_card else "cpu",
        },
        "lm_head": "tied to the input embedding, no key",
        "torch_version": torch.__version__,
        "transformers_version": transformers_version,
        "safetensors_version": safetensors.__version__,
    }
    header = {
        UNSLOTH_FORMAT_KEY: LTX25_TE_HOSTED_FORMAT,
        UNSLOTH_METADATA_KEY: json.dumps(metadata, sort_keys = True),
        "tensor_names": json.dumps(names),
    }
    # Plain tensors only, so the shared reader (prequant_safetensors.load_plain_prequant_safetensors) accepts it.
    header.update({name: json.dumps({"_type": "Tensor"}) for name in names})
    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)
    return {
        "out": str(out),
        "bytes": out.stat().st_size,
        "sha256": digest.hexdigest(),
        "quantized_linears": quantized,
        "tensors": len(names),
    }


def _available_device(torch, device: str) -> str:
    """``device`` when its backend is usable here, else ``cpu``: the worker moves a resident build onto the device."""
    kind = torch.device(device).type
    backend = getattr(torch, kind, None) if kind in ("cuda", "xpu") else None
    if backend is None:
        return device
    is_available = getattr(backend, "is_available", None)
    if callable(is_available) and is_available():
        return device
    print(f"{device} is not available on this host; building on the CPU", file = sys.stderr)
    return "cpu"


def _head_is_tied(state: dict) -> bool:
    """Whether ``lm_head.weight`` is the same storage as another tensor in ``state`` (the input embedding)."""
    head = state.get("lm_head.weight")
    if head is None:
        return False
    return any(
        name != "lm_head.weight" and t.shape == head.shape and t.data_ptr() == head.data_ptr()
        for name, t in state.items()
    )


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