"""Packed on-disk storage for LTX-2.5 encoder precision plans.

`ltx_fake_quantize` and `ltx_gptq` store quantized values expanded back to
BF16, so a 4.5-bit plan still costs 26 GB on disk. This module is the missing
second half: the same values in their native widths.

The format is one safetensors file. A raw tensor keeps its checkpoint name. A
quantized tensor is split into named parts:

* nvfp4 - ``{name}::nvfp4_codes`` (uint8, two 4-bit codes per byte, even
  column in the low nibble; a code is ``sign << 3 | magnitude`` indexing
  ``E2M1_LEVELS``), ``{name}::nvfp4_group_scale`` (float8_e4m3fn,
  [out, in/16]), ``{name}::nvfp4_global_scale`` (float32 scalar), and - when
  the GPTQ path smoothed the tensor - ``{name}::pre_quant_scale`` (float32,
  [in], the AWQ scale the stored weight was *divided* by).
* int8 - ``{name}::int8`` and ``{name}::int8_scale`` (float32, [out]).

Exactness is the design constraint, not an aspiration: unpacking replays the
producers' arithmetic - ``(sign * level) * effective`` then the pre-scale
division, in float32, cast to the stored dtype last - so a packed tensor
dequantizes to the same values the BF16 fake-quant file would have carried.
The one tolerated difference is that int8 cannot store a negative zero, so a
``-0.0`` produced by ``round()`` collapses to ``+0.0``; ``torch.equal`` treats
the two as equal and the packer counts them rather than hiding them.

Group scales are recorded as the exact float8 bytes the quantizer used, not
recomputed from the output. Recomputing them is impossible in general: the
GPTQ column loop derives each group's scale from weights that were already
compensated, and the AWQ division afterwards destroys the grid alignment that
recovery would need. That is why packing happens inside the build
(`ltx_gptq --packed-output`) instead of as a post-pass over the artifact.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parent))

# --- inlined from ltx_fake_quantize -----------------------------------------
#: The eight magnitudes E2M1 can represent, and the midpoints between them that
#: decide which one a value rounds to.
E2M1_LEVELS = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)
NVFP4_GROUP = 16
#: E4M3's largest finite value; the per-tensor scale is chosen so the per-group
#: scales divided by it land inside this range.
E4M3_MAX = 448.0
GPU_TENSOR_LIMIT = 1_000_000_000

# --- inlined from ltx_prompt_embedding_gate ---------------------------------
#: Lightricks re-exported the encoder under the naming Google used before the
#: `gemma4_unified` port landed in `transformers`. Without these rules
#: `from_pretrained` reports every checkpoint tensor UNEXPECTED and every
#: parameter MISSING, and hands back a randomly initialised model that looks
#: like it loaded.
LTX_RENAMES = [
    (r"^model\.", "language_model."),
    (r"^vision_model\.", "embed_vision."),
    (r"^audio_projector\.", "embed_audio."),
    (r"^multi_modal_projector\.embedding_projection",
     "embed_vision.multimodal_embedder.embedding_projection"),
]


def register_ltx_renames() -> None:
    from transformers.conversion_mapping import (
        get_checkpoint_conversion_mapping,
        register_checkpoint_conversion_mapping,
    )
    from transformers.core_model_loading import WeightRenaming

    existing = get_checkpoint_conversion_mapping("gemma4_unified") or []
    rules = [WeightRenaming(source_patterns=s, target_patterns=t)
             for s, t in LTX_RENAMES]
    register_checkpoint_conversion_mapping(
        "gemma4_unified", rules + existing, overwrite=True)
# ----------------------------------------------------------------------------

FORMAT_VERSION = "haverbex-packed-v1"

#: Part suffixes. `::` cannot appear in a checkpoint tensor name, so packed
#: parts can never collide with a raw tensor.
CODES = "::nvfp4_codes"
GROUP_SCALE = "::nvfp4_group_scale"
GLOBAL_SCALE = "::nvfp4_global_scale"
PRE_SCALE = "::pre_quant_scale"
INT8 = "::int8"
INT8_SCALE = "::int8_scale"

_SUFFIXES = (CODES, GROUP_SCALE, GLOBAL_SCALE, PRE_SCALE, INT8, INT8_SCALE)

#: Parts indexed by output row, so a chunked read can slice them. The other two
#: are not: `GLOBAL_SCALE` is a scalar and `PRE_SCALE` is per *input* channel.
_ROW_SLICEABLE = (CODES, GROUP_SCALE, INT8, INT8_SCALE)

#: Part suffix -> buffer name on a resident module. `::` is illegal in a
#: `register_buffer` name, so the mapping cannot be derived.
_BUFFER_NAMES = {CODES: "codes", GROUP_SCALE: "group_scale",
                 GLOBAL_SCALE: "global_scale", PRE_SCALE: "pre_scale",
                 INT8: "ints", INT8_SCALE: "int8_scale"}


def nibble_pack(codes: torch.Tensor) -> torch.Tensor:
    """[out, in] uint8 (values 0..15) -> [out, in/2], even column low nibble."""
    if codes.shape[-1] % 2:
        raise ValueError(f"odd inner width {codes.shape[-1]} cannot nibble-pack")
    return codes[..., 0::2] | (codes[..., 1::2] << 4)


def nibble_unpack(packed: torch.Tensor) -> torch.Tensor:
    out = torch.empty(*packed.shape[:-1], packed.shape[-1] * 2, dtype=torch.uint8,
                      device=packed.device)
    out[..., 0::2] = packed & 0x0F
    out[..., 1::2] = packed >> 4
    return out


def encode_nvfp4_column(normalized: torch.Tensor, codes: torch.Tensor) -> torch.Tensor:
    """4-bit code for one already-bucketized column: sign bit over magnitude.

    `torch.sign` maps an exact zero to 0, which multiplies out to ``+0.0``; the
    encoding gives it sign bit 0 so decode lands on the same ``+0.0``.
    """
    return codes.to(torch.uint8) | ((normalized < 0).to(torch.uint8) << 3)


def decode_nvfp4(codes: torch.Tensor, group_scale: torch.Tensor,
                 global_scale: torch.Tensor,
                 pre_scale: torch.Tensor | None) -> torch.Tensor:
    """Replay of `NVFP4Column.quantize` / `quantize_nvfp4`, in float32.

    `levels` is built on the codes' device: at load time everything is on the
    CPU, but a `PackedLinear` decodes wherever its weight lives.
    """
    levels = torch.tensor(E2M1_LEVELS, dtype=torch.float32, device=codes.device)
    sign = torch.where((codes & 0x8) != 0, -1.0, 1.0)
    magnitude = levels[(codes & 0x7).long()]

    effective = group_scale.float() * global_scale.float()
    effective = torch.where(effective > 0, effective, torch.ones_like(effective))
    out_features, in_features = codes.shape
    grid = (sign * magnitude).reshape(out_features, in_features // NVFP4_GROUP,
                                      NVFP4_GROUP)
    grid = (grid * effective.unsqueeze(-1)).reshape(out_features, in_features)
    if pre_scale is not None:
        grid = grid / pre_scale.float()
    return grid


def decode_int8(ints: torch.Tensor, scale: torch.Tensor,
                pre_scale: torch.Tensor | None) -> torch.Tensor:
    """Replay of `quantize_int8` / `Int8Column.quantize`, in float32."""
    grid = ints.float() * scale.float().unsqueeze(-1)
    if pre_scale is not None:
        grid = grid / pre_scale.float()
    return grid


def pack_rtn_nvfp4(w: torch.Tensor) -> tuple[torch.Tensor, dict]:
    """`ltx_fake_quantize.quantize_nvfp4` with the codes and scales kept.

    Returns the dequantized float32 tensor (identical to what the fake
    quantizer produces) and the packed parts.
    """
    out_features, in_features = w.shape
    if in_features % NVFP4_GROUP:
        raise ValueError(f"nvfp4 needs a multiple of {NVFP4_GROUP}, got {in_features}")
    amax = w.abs().amax()
    if amax == 0:
        raise ValueError("all-zero tensor should be stored raw, not packed")
    global_scale = amax / (E2M1_LEVELS[-1] * E4M3_MAX)

    groups = w.reshape(out_features, in_features // NVFP4_GROUP, NVFP4_GROUP)
    group_amax = groups.abs().amax(dim=-1, keepdim=True)
    scale = (group_amax / E2M1_LEVELS[-1] / global_scale).to(torch.float8_e4m3fn)
    effective = scale.float() * global_scale
    effective = torch.where(effective > 0, effective, torch.ones_like(effective))

    levels = torch.tensor(E2M1_LEVELS, device=w.device, dtype=w.dtype)
    midpoints = torch.tensor(E2M1_MIDPOINTS, device=w.device, dtype=w.dtype)
    normalized = groups / effective
    codes = torch.bucketize(normalized.abs(), midpoints, out_int32=True)
    dequant = (torch.sign(normalized) * levels[codes] * effective).reshape(
        out_features, in_features)

    packed_codes = (codes.to(torch.uint8)
                    | ((normalized < 0).to(torch.uint8) << 3)).reshape(
                        out_features, in_features)
    parts = {
        CODES: nibble_pack(packed_codes).cpu(),
        GROUP_SCALE: scale.squeeze(-1).cpu(),
        GLOBAL_SCALE: global_scale.detach().float().reshape(1).cpu(),
    }
    return dequant, parts


def pack_rtn(writer: PackWriter, name: str, tensor: torch.Tensor,
             width: float, device: str) -> torch.Tensor:
    """`ltx_fake_quantize.quantize_tensor`, but keeping the codes.

    Returns the fake-quantized tensor in the original dtype so the caller's
    BF16 write path is unchanged. Same device policy as the original: a tensor
    whose float32 form exceeds `GPU_TENSOR_LIMIT` is quantized on the CPU.
    """
    original = tensor.dtype
    where = "cpu" if tensor.numel() * 4 > GPU_TENSOR_LIMIT else device
    w = tensor.to(where, torch.float32)
    if width == 4.5:
        if w.abs().amax() == 0:
            # quantize_nvfp4 returns an all-zero tensor unchanged; store it raw
            writer.store_raw(name, tensor)
            return tensor
        dequant, parts = pack_rtn_nvfp4(w)
        result = dequant.to("cpu", original)
        writer.store_nvfp4(name, parts, result)
    else:
        dequant, parts = pack_rtn_int8(w)
        result = dequant.to("cpu", original)
        writer.store_int8(name, parts, result)
    return result


def pack_rtn_int8(w: torch.Tensor) -> tuple[torch.Tensor, dict]:
    """`ltx_fake_quantize.quantize_int8` with the integers kept."""
    scale = w.abs().amax(dim=-1, keepdim=True) / 127.0
    scale = torch.where(scale > 0, scale, torch.ones_like(scale))
    ints = (w / scale).round().clamp_(-127, 127)
    dequant = ints * scale
    parts = {
        INT8: ints.to(torch.int8).cpu(),
        INT8_SCALE: scale.squeeze(-1).float().cpu(),
    }
    return dequant, parts


class PackWriter:
    """Accumulates packed parts and raw tensors, then writes one file.

    Held in RAM rather than streamed: the whole point of the format is that
    the payload is ~8.5 GB, which fits beside the build. Every quantized store
    verifies round-trip equality against the reference tensor before
    accepting it - a packed file that does not reproduce its builder's values
    must not be creatable through this class.
    """

    def __init__(self) -> None:
        self.tensors: dict[str, torch.Tensor] = {}
        self.kinds: dict[str, str] = {}
        self.negative_zero_collapses = 0

    def store_raw(self, name: str, tensor: torch.Tensor) -> None:
        if any(s in name for s in _SUFFIXES):
            raise ValueError(f"raw name collides with a part suffix: {name}")
        self.tensors[name] = tensor.detach().contiguous().cpu()
        self.kinds[name] = "raw"

    def _verify(self, name: str, reference: torch.Tensor) -> None:
        got = unpack_tensor(self.tensors, name, self.kinds[name],
                            reference.dtype)
        if not torch.equal(got, reference.cpu()):
            diff = (got.float() - reference.float().cpu()).abs()
            raise SystemExit(
                f"packed round-trip mismatch on {name}: max {diff.max():.3e} at "
                f"{int(diff.argmax())} - refusing to write a lossy pack")
        signs = got.signbit() != reference.cpu().signbit()
        self.negative_zero_collapses += int(signs.sum())

    def store_nvfp4(self, name: str, parts: dict, reference: torch.Tensor,
                    pre_scale: torch.Tensor | None = None) -> None:
        for suffix, tensor in parts.items():
            self.tensors[name + suffix] = tensor.contiguous()
        if pre_scale is not None:
            self.tensors[name + PRE_SCALE] = pre_scale.detach().float().cpu()
        self.kinds[name] = "nvfp4"
        self._verify(name, reference)

    def store_int8(self, name: str, parts: dict, reference: torch.Tensor,
                   pre_scale: torch.Tensor | None = None) -> None:
        for suffix, tensor in parts.items():
            self.tensors[name + suffix] = tensor.contiguous()
        if pre_scale is not None:
            self.tensors[name + PRE_SCALE] = pre_scale.detach().float().cpu()
        self.kinds[name] = "int8"
        self._verify(name, reference)

    def save(self, path: Path, metadata: dict[str, str]) -> None:
        from safetensors.torch import save_file
        payload = {
            "format": FORMAT_VERSION,
            "kinds": json.dumps(self.kinds, separators=(",", ":")),
            "negative_zero_collapses": str(self.negative_zero_collapses),
            **metadata,
        }
        path.parent.mkdir(parents=True, exist_ok=True)
        save_file(self.tensors, str(path), metadata=payload)


def unpack_tensor(tensors: dict, name: str, kind: str,
                  dtype: torch.dtype) -> torch.Tensor:
    if kind == "raw":
        return tensors[name]
    pre = tensors.get(name + PRE_SCALE)
    if kind == "nvfp4":
        grid = decode_nvfp4(nibble_unpack(tensors[name + CODES]),
                            tensors[name + GROUP_SCALE],
                            tensors[name + GLOBAL_SCALE],
                            pre)
    elif kind == "int8":
        grid = decode_int8(tensors[name + INT8], tensors[name + INT8_SCALE], pre)
    else:
        raise ValueError(f"unknown kind {kind!r} for {name}")
    return grid.to(dtype)


class PackedCheckpoint:
    """Read-side API mirroring `safetensors.safe_open` for packed files."""

    def __init__(self, path: Path | str) -> None:
        from safetensors import safe_open
        self.path = Path(path)
        self._file = safe_open(str(self.path), framework="pt")
        self.metadata = self._file.metadata() or {}
        if self.metadata.get("format") != FORMAT_VERSION:
            raise SystemExit(
                f"{path}: format {self.metadata.get('format')!r}, "
                f"this reader speaks {FORMAT_VERSION}")
        self.kinds: dict[str, str] = json.loads(self.metadata["kinds"])
        self._present = set(self._file.keys())

    def __enter__(self) -> PackedCheckpoint:
        return self

    def __exit__(self, *exc) -> None:
        return None

    def keys(self) -> list[str]:
        return list(self.kinds)

    def get_tensor(self, name: str, dtype: torch.dtype = torch.bfloat16):
        kind = self.kinds[name]
        if kind == "raw":
            return self._file.get_tensor(name)
        parts = {name + s: self._file.get_tensor(name + s)
                 for s in _SUFFIXES if name + s in self._present}
        return unpack_tensor(parts, name, kind, dtype)

    def get_parts(self, name: str) -> dict:
        """The stored parts of a packed tensor, keyed by bare suffix, undecoded.

        What a resident module needs: the bytes as written, with no decode and
        no reassembly.
        """
        return {s: self._file.get_tensor(name + s)
                for s in _SUFFIXES if name + s in self._present}

    def row_count(self, name: str) -> int:
        """Output rows of a packed tensor, without decoding it."""
        kind = self.kinds[name]
        key = name if kind == "raw" else name + (CODES if kind == "nvfp4" else INT8)
        return self._file.get_slice(key).get_shape()[0]

    def get_tensor_rows(self, name: str, start: int, stop: int,
                        dtype: torch.dtype = torch.bfloat16):
        """Output rows `[start, stop)` of a tensor, decoding only those rows.

        Lets a caller work through a table that does not fit beside the model -
        the LTX aggregates are 3.08 GB each in float32. Row-sliced parts are
        the ones indexed by output channel; the global scale is a scalar and
        the AWQ pre-scale is per *input* channel, so both are read whole.
        """
        kind = self.kinds[name]
        if kind == "raw":
            return self._file.get_slice(name)[start:stop].to(dtype)
        parts = {}
        for suffix in _SUFFIXES:
            key = name + suffix
            if key not in self._present:
                continue
            parts[key] = (self._file.get_slice(key)[start:stop]
                          if suffix in _ROW_SLICEABLE
                          else self._file.get_tensor(key))
        return unpack_tensor(parts, name, kind, dtype)


def open_maybe_packed(path: Path | str):
    """`safe_open` for BF16 checkpoints, `PackedCheckpoint` for packed ones.

    Both expose `keys()` and `get_tensor(name)`, which is all the aggregate
    and asset readers use. Detection is by content, not extension: a packed
    file carries `format` metadata that a plain checkpoint does not.
    """
    from safetensors import safe_open
    handle = safe_open(str(path), framework="pt")
    if (handle.metadata() or {}).get("format") == FORMAT_VERSION:
        return PackedCheckpoint(path)
    return handle


def install_resident(model, reader: PackedCheckpoint, wanted: dict, rename,
                     device: str | None = None, fold_pre_scale: bool = False):
    """Replace the quantized modules of a meta-device model with packed ones.

    Two passes, because a module has to exist before its bias can be fed into
    it: first swap every module whose weight is packed, then feed the raw
    tensors - norms, layer scalars, and any bias - into whatever module now
    sits at that path.

    The coverage guard is the same one the dequantized path uses, and it is
    what makes a silent rename failure impossible: every parameter the model
    declared must be accounted for, either by a swap or by a raw feed.
    """
    from accelerate.utils import set_module_tensor_to_device

    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    swapped, raw_names = set(), []
    for source_name in reader.keys():  # noqa: SIM118 - PackedCheckpoint, not a dict
        param_name = rename(source_name)
        if param_name not in wanted:
            continue  # aggregates and asset blobs live outside the module tree
        if reader.kinds[source_name] == "raw":
            raw_names.append((source_name, param_name))
            continue
        if not param_name.endswith(".weight"):
            raise SystemExit(
                f"{source_name} is packed but renames to {param_name}, which is "
                "not a module weight; a resident build has nowhere to put it")
        path = param_name[: -len(".weight")]
        module = model.get_submodule(path)
        parts = {k: v.to(device) for k, v in reader.get_parts(source_name).items()}

        if isinstance(module, torch.nn.Embedding):
            # Subclasses carry forward behaviour a bare gather would lose -
            # Gemma's scales by ~sqrt(hidden). Anything else must stop the build.
            extra = set(dict(module.named_buffers()))- {"embed_scale"}
            if type(module) is not torch.nn.Embedding and extra:
                raise SystemExit(
                    f"{path} is a {type(module).__name__} carrying {sorted(extra)}; "
                    "PackedEmbedding only reproduces embed_scale")
            packed = PackedEmbedding(parts[INT8], parts[INT8_SCALE],
                                     embed_scale=getattr(module, "embed_scale", None),
                                     padding_idx=module.padding_idx)
        elif isinstance(module, torch.nn.Linear):
            packed = PackedLinear(reader.kinds[source_name], parts,
                                  module.out_features, module.in_features,
                                  fold_pre_scale=fold_pre_scale)
        else:
            raise SystemExit(f"{path} is a {type(module).__name__}, and this "
                             "build only knows how to pack Linear and Embedding")
        parent_path, _, attribute = path.rpartition(".")
        setattr(model.get_submodule(parent_path) if parent_path else model,
                attribute, packed)
        swapped.add(param_name)

    # Second pass, after every swap: a bias whose module was replaced lands on
    # the new module's buffer, which only exists once the swap has happened.
    for source_name, param_name in raw_names:
        value = reader.get_tensor(source_name, wanted[param_name].dtype).to(device)
        path, _, attribute = param_name.rpartition(".")
        owner = model.get_submodule(path) if path else model
        if isinstance(owner, (PackedLinear, PackedEmbedding)):
            # `set_module_tensor_to_device` reads the current value to find its
            # device, and a packed module declares `bias` as a None buffer.
            owner.register_buffer(attribute, value, persistent=False)
        else:
            set_module_tensor_to_device(model, param_name, device, value=value)

    assigned = swapped | {p for _, p in raw_names}
    missing = [n for n in wanted if n not in assigned and "rotary" not in n]
    if missing:
        raise SystemExit(
            f"{len(missing)} parameters were not fed from the packed file, "
            f"e.g. {missing[:5]}. The packed names do not match the model; "
            "extend LTX_RENAMES.")

    model.eval()
    print(f"resident packed model on {device}: "
          f"{resident_bytes(model) / 2**30:.3f} GiB", flush=True)
    return model


def resident_bytes(model) -> int:
    """Everything the model holds: packed buffers plus whatever stayed dense.

    Counted once. A packed module's parts are registered buffers, so walking
    `model.buffers()` after adding `packed_bytes()` counts them twice - which
    is what the first A7 run reported (13.4 GiB against an allocator peak of
    7.75 GiB). Packed buffers are collected by identity first and skipped in
    the dense pass.
    """
    seen, total = set(), 0
    for module in model.modules():
        if isinstance(module, (PackedLinear, PackedEmbedding)):
            total += module.packed_bytes()
            for name in (*_BUFFER_NAMES.values(), "bias"):
                tensor = getattr(module, name, None)
                if tensor is not None:
                    seen.add(id(tensor))
    for tensor in list(model.parameters()) + list(model.buffers()):
        if id(tensor) in seen or tensor.device.type == "meta":
            continue
        seen.add(id(tensor))
        total += tensor.numel() * tensor.element_size()
    return total


def check_gpu_kernels(device: str | None = None) -> None:
    """Refuse a wheel that has no kernels for this card, and say what to install.

    Nothing in this format needs anything unusual from a GPU - no fp8 units, no
    minimum compute capability, no bf16 tensor cores. What can be missing is
    PyTorch's own kernels: the current default wheel on PyPI is a cu130 build,
    and cu130 dropped Volta.

    Left alone, that surfaces as

        CUDA error: no kernel image is available for execution on the device

    raised from the first kernel launch - which is after an 8.46 GB load, inside
    somebody else's library, and long after `torch.cuda.is_available()` returned
    True. Checking `get_arch_list` costs nothing and moves the failure to the
    place where the fix makes sense.
    """
    if device is not None and not str(device).startswith("cuda"):
        return
    if not torch.cuda.is_available():
        return

    major, minor = torch.cuda.get_device_capability(0)
    arch = f"sm_{major}{minor}"
    compiled = torch.cuda.get_arch_list()
    # An empty list means a build that does not report them; do not guess.
    if not compiled or arch in compiled:
        return

    name = torch.cuda.get_device_name(0)
    raise SystemExit(
        f"this torch ({torch.__version__}) has no kernels for {name} ({arch}).\n"
        f"It was built for {', '.join(compiled)}, and the first CUDA op would "
        f"fail with 'no kernel image is available for execution on the device'.\n"
        f"The model is fine - it needs no custom kernels. Install a torch built "
        f"for your card, e.g. for {arch}:\n"
        f"  pip install torch --index-url https://download.pytorch.org/whl/cu128\n"
        f"or pass device='cpu' to load without touching the GPU.")


def load_packed_model(model_dir: str, packed_path: Path | str,
                      gpu_budget: str | None = None, resident: bool = False,
                      fold_pre_scale: bool = False, device: str | None = None):
    """Build the encoder from a packed file, split across GPU and CPU.

    Mirrors what `AutoModel.from_pretrained(..., device_map="auto")` does for
    the BF16 checkpoint: same conversion-mapping renames, same guard that
    every parameter was actually fed from the file. Dequantization happens
    tensor-by-tensor, so peak memory is one dequantized tensor above the
    final footprint - the 26 GB BF16 file never exists.

    Note what the default does and does not save. The weights land as BF16, so
    the *resident* footprint is the same 26 GB the BF16 checkpoint would take;
    what packing buys there is disk and the load-time peak. `gpu_budget`
    defaults to 13 GiB - what is left of a 16 GB card - and `LTX_PACKED_GPU_BUDGET`
    raises it on a larger card, where the whole encoder fits and the CPU
    offload that budget forces is pure slowdown.

    With `resident`, the quantized tensors stay in their native widths behind
    `PackedLinear` / `PackedEmbedding` and no device map is needed: the model
    costs what the file costs. `fold_pre_scale` is passed through to the
    linears and is not bit-identical - see `PackedLinear`.
    """
    import os
    import re

    check_gpu_kernels(device)

    if gpu_budget is None:
        gpu_budget = os.environ.get("LTX_PACKED_GPU_BUDGET", "13GiB")

    from accelerate import infer_auto_device_map, init_empty_weights
    from accelerate.utils import set_module_tensor_to_device
    from transformers import AutoConfig, AutoModel
    from transformers.conversion_mapping import get_checkpoint_conversion_mapping

    register_ltx_renames()
    rules = get_checkpoint_conversion_mapping("gemma4_unified") or []

    def rename(name: str) -> str:
        for rule in rules:
            sources = rule.source_patterns
            targets = rule.target_patterns
            if isinstance(sources, str):
                sources, targets = [sources], [targets]
            for source, target in zip(sources, targets, strict=True):
                new = re.sub(source, target, name)
                if new != name:
                    return new
        return name

    config = AutoConfig.from_pretrained(model_dir)
    # `from_config` does not resolve the attention implementation the way
    # `from_pretrained` does; left unset, a standalone build falls back to
    # eager and mis-applies the boolean SDPA mask (see ltx_gptq).
    config._attn_implementation = "sdpa"
    if hasattr(config, "text_config"):
        config.text_config._attn_implementation = "sdpa"
    with init_empty_weights():
        model = AutoModel.from_config(config)
    model = model.to(torch.bfloat16)

    reader = PackedCheckpoint(packed_path)
    wanted = dict(model.state_dict())
    if resident:
        return install_resident(model, reader, wanted, rename,
                                device=device, fold_pre_scale=fold_pre_scale)

    device_map = infer_auto_device_map(
        model, max_memory={0: gpu_budget, "cpu": "40GiB"},
        dtype=torch.bfloat16, no_split_module_classes=["Gemma4UnifiedTextDecoderLayer"])

    # `set_module_tensor_to_device` places tensors but installs no hooks, so a
    # split map produces a model that only fails once a forward crosses the
    # boundary - as a device mismatch deep inside a layernorm. Refuse it here
    # instead, and say what to do about it.
    placements = set(device_map.values())
    if len(placements) > 1:
        raise SystemExit(
            f"the packed model does not fit in {gpu_budget} and would be split "
            # str(): a device map mixes GPU ordinals with "cpu", and sorting
            # those against each other is a TypeError - which is how this
            # message first announced itself.
            f"across {sorted(map(str, placements))}. This loader dispatches no "
            "hooks, so a "
            "split model raises mid-forward. Raise LTX_PACKED_GPU_BUDGET, or use "
            "resident=True, which needs only what the file costs.")

    def target_device(param_name: str):
        candidate = param_name
        while candidate:
            if candidate in device_map:
                return device_map[candidate]
            candidate = candidate.rsplit(".", 1)[0] if "." in candidate else ""
        return device_map.get("", "cpu")

    assigned = set()
    for source_name in reader.keys():  # noqa: SIM118 - PackedCheckpoint, not a dict
        param_name = rename(source_name)
        if param_name not in wanted:
            continue  # aggregates and asset blobs live outside the module tree
        value = reader.get_tensor(source_name, wanted[param_name].dtype)
        set_module_tensor_to_device(model, param_name, target_device(param_name),
                                    value=value)
        assigned.add(param_name)

    missing = [name for name in wanted
               if name not in assigned and "rotary" not in name]
    if missing:
        raise SystemExit(
            f"{len(missing)} parameters were not fed from {packed_path}, "
            f"e.g. {missing[:5]}. The packed names do not match the model; "
            "extend LTX_RENAMES.")
    model.eval()
    return model


# ---------------------------------------------------------------------------
# Resident modules
#
# `load_packed_model` dequantizes to BF16, which saves the load-time peak and
# nothing else: the model that comes out is the same 26.264 GB it always was.
# These modules keep the packed widths in memory and dequantize inside
# `forward`, which is what turns 8.463 GB on disk into 8.463 GB resident.
#
# Affordable here because the encoder runs once per prompt - no KV cache, no
# autoregressive loop. Measured on this checkpoint's real shapes at 1024
# tokens, a full dequantize costs 9-12% of the matmul it feeds.
# ---------------------------------------------------------------------------


class PackedLinear(torch.nn.Module):
    """`nn.Linear` whose weight is stored in its native widths.

    `forward` rebuilds the weight with the same arithmetic `unpack_tensor` uses
    at load time, so its output is bit-identical to the dequantized model's.
    That equality is the point: without it a resident run cannot be compared
    against any figure recorded from the dequantized path.

    `fold_pre_scale` trades the equality for speed. AWQ smoothing is
    ``W' = W * s`` and the stored weight carries the closing ``/ s``, so
    ``x @ (Q/s).T == (x/s) @ Q.T`` - the division can move onto the activation
    and shrink an ``[out, in]`` elementwise pass to a ``[tokens, in]`` one. It
    changes float rounding order, so it is off by default and its difference is
    measured rather than assumed.
    """

    def __init__(self, kind: str, parts: dict, out_features: int,
                 in_features: int, bias: torch.Tensor | None = None,
                 dtype: torch.dtype = torch.bfloat16,
                 fold_pre_scale: bool = False) -> None:
        super().__init__()
        if kind not in ("nvfp4", "int8"):
            raise ValueError(f"{kind!r} is not a packed weight kind")
        self.kind = kind
        self.out_features = out_features
        self.in_features = in_features
        self.compute_dtype = dtype

        for suffix, buffer in _BUFFER_NAMES.items():
            tensor = parts.get(name_part(suffix, parts))
            self.register_buffer(buffer, tensor, persistent=False)
        self.register_buffer("bias", None if bias is None else bias.detach(),
                             persistent=False)
        self.fold_pre_scale = bool(fold_pre_scale) and self.pre_scale is not None

    def packed_bytes(self) -> int:
        """What this module actually costs, for a device map that cannot infer it."""
        total = 0
        for buffer in (*_BUFFER_NAMES.values(), "bias"):
            tensor = getattr(self, buffer, None)
            if tensor is not None:
                total += tensor.numel() * tensor.element_size()
        return total

    @property
    def weight(self) -> torch.Tensor:
        """The dequantized weight.

        A property rather than a buffer so that code reaching for
        `module.weight` - `transformers` does, in places - still works, at the
        cost of materializing it for that one call.
        """
        return self.dequantize()

    def dequantize(self, pre_scale: bool = True) -> torch.Tensor:
        pre = self.pre_scale if (pre_scale and not self.fold_pre_scale) else None
        if self.kind == "nvfp4":
            grid = decode_nvfp4(nibble_unpack(self.codes), self.group_scale,
                                self.global_scale, pre)
        else:
            grid = decode_int8(self.ints, self.int8_scale, pre)
        return grid.to(self.compute_dtype)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.fold_pre_scale:
            x = x / self.pre_scale.to(x.dtype)
        return torch.nn.functional.linear(x, self.dequantize(), self.bias)

    def extra_repr(self) -> str:
        return (f"in_features={self.in_features}, out_features={self.out_features}, "
                f"kind={self.kind}, packed={self.packed_bytes() / 2**20:.1f} MiB"
                + (", folded" if self.fold_pre_scale else ""))


class PackedEmbedding(torch.nn.Module):
    """`nn.Embedding` over int8 rows with a per-row scale.

    `model.embed_tokens` is [262144, 3840] - 1.008 GB packed against 2.016 GB
    BF16, the largest single tensor in the file, and the one DmitryDB's recipe
    keeps in BF16 as a precision island. Quantizing it is most of our disk
    advantage, and half of that is given back if it dequantizes at load.

    Gathering before scaling means only the rows a prompt touches are ever
    dequantized, which is a few thousand of 262144.

    `embed_scale` is not optional in practice. Gemma's embedding is a
    `Gemma4UnifiedTextScaledWordEmbedding`, whose forward is
    ``super().forward(ids) * embed_scale`` with `embed_scale` around sqrt(3840).
    Replacing the module without carrying that factor would drop every
    embedding by ~62x and still produce plausible-looking tensors, so the caller
    must read it off the module it is replacing.
    """

    def __init__(self, ints: torch.Tensor, scale: torch.Tensor,
                 dtype: torch.dtype = torch.bfloat16,
                 embed_scale: torch.Tensor | None = None,
                 padding_idx: int | None = None) -> None:
        super().__init__()
        self.num_embeddings, self.embedding_dim = ints.shape
        self.compute_dtype = dtype
        self.padding_idx = padding_idx
        self.register_buffer("ints", ints, persistent=False)
        self.register_buffer("int8_scale", scale, persistent=False)
        self.register_buffer("embed_scale", embed_scale, persistent=False)

    def packed_bytes(self) -> int:
        return (self.ints.numel() * self.ints.element_size()
                + self.int8_scale.numel() * self.int8_scale.element_size())

    @property
    def weight(self) -> torch.Tensor:
        return decode_int8(self.ints, self.int8_scale, None).to(self.compute_dtype)

    def forward(self, index: torch.Tensor) -> torch.Tensor:
        rows = self.ints[index].float()
        scale = self.int8_scale[index].unsqueeze(-1).float()
        out = (rows * scale).to(self.compute_dtype)
        if self.embed_scale is not None:
            # Cast then multiply, in that order, because that is what
            # `Gemma4UnifiedTextScaledWordEmbedding.forward` does.
            out = out * self.embed_scale.to(self.compute_dtype)
        return out

    def extra_repr(self) -> str:
        scaled = "" if self.embed_scale is None else f", embed_scale={float(self.embed_scale):.4g}"
        return (f"{self.num_embeddings}, {self.embedding_dim}, kind=int8, "
                f"packed={self.packed_bytes() / 2**20:.1f} MiB{scaled}")


def name_part(suffix: str, parts: dict) -> str:
    """The key in `parts` carrying `suffix`, or a miss the caller tolerates.

    `parts` comes either straight from a packer (keys are bare suffixes) or
    from a checkpoint read (keys are `name + suffix`), and both are worth
    supporting so a test does not have to fabricate tensor names.
    """
    if suffix in parts:
        return suffix
    for key in parts:
        if key.endswith(suffix):
            return key
    return suffix


def packed_linear_apply(reader, name: str, x: torch.Tensor,
                        bias: torch.Tensor | None = None, rows: int = 512,
                        dtype: torch.dtype = torch.float32) -> torch.Tensor:
    """`F.linear(x, W, bias)` for a packed W, `rows` output rows at a time.

    The two LTX aggregate tables are [4096, 188160] and [2048, 188160]. Reading
    one whole and casting it to float32 costs 3.08 GB, which is why
    `apply_aggregates` only runs after the encoder has been unloaded. Chunking
    is over *output* rows, so the contraction over `in_features` is untouched
    and the result is bit-identical to decoding the table in one piece.
    """
    out_features = reader.row_count(name)
    pieces = []
    for start in range(0, out_features, rows):
        stop = min(start + rows, out_features)
        weight = reader.get_tensor_rows(name, start, stop, dtype)
        piece = None if bias is None else bias[start:stop]
        pieces.append(torch.nn.functional.linear(x, weight, piece))
        del weight
    return torch.cat(pieces, dim=-1)


def cli_info(args) -> int:
    reader = PackedCheckpoint(args.packed)
    from collections import Counter
    counts = Counter(reader.kinds.values())
    size = Path(args.packed).stat().st_size
    print(f"{args.packed}: {size / 1e9:.3f} GB, {dict(counts)}")
    print(f"negative-zero collapses at pack time: "
          f"{reader.metadata.get('negative_zero_collapses')}")
    for key in sorted(set(reader.metadata) - {"kinds"}):
        print(f"  {key}: {reader.metadata[key][:100]}")
    return 0


def cli_verify(args) -> int:
    """Compare every tensor against a reference BF16 checkpoint, streaming."""
    from safetensors import safe_open
    reader = PackedCheckpoint(args.packed)
    reference = safe_open(str(args.reference), framework="pt")
    names = set(reference.keys())
    missing = sorted(set(reader.kinds) - names)
    extra = sorted(names - set(reader.kinds))
    if missing or extra:
        raise SystemExit(f"tensor sets differ: missing {missing[:5]}, extra {extra[:5]}")
    worst = 0
    for i, name in enumerate(sorted(names)):
        want = reference.get_tensor(name)
        got = reader.get_tensor(name, want.dtype)
        if not torch.equal(got, want):
            raise SystemExit(f"{name}: dequantized values differ from reference")
        worst += int((got.signbit() != want.signbit()).sum())
        if i % 100 == 0:
            print(f"  {i}/{len(names)} verified", flush=True)
    print(f"all {len(names)} tensors value-exact; {worst} zero-sign differences")
    return 0


def cli_emit_bf16(args) -> int:
    """Materialize a plain BF16 checkpoint a stock loader can read."""
    from safetensors.torch import save_file
    reader = PackedCheckpoint(args.packed)
    tensors = {}
    for name in sorted(reader.kinds):
        tensors[name] = reader.get_tensor(name, torch.bfloat16)
    save_file(tensors, str(args.output))
    print(f"wrote {args.output} ({Path(args.output).stat().st_size / 1e9:.2f} GB)")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="command", required=True)
    sub.add_parser("info").add_argument("packed")
    p = sub.add_parser("verify")
    p.add_argument("packed")
    p.add_argument("--reference", required=True)
    p = sub.add_parser("emit-bf16")
    p.add_argument("packed")
    p.add_argument("--output", required=True)
    args = parser.parse_args()
    return {"info": cli_info, "verify": cli_verify,
            "emit-bf16": cli_emit_bf16}[args.command](args)


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