runtime_loader.py

"""Runtime-only loader for Capicu's quantized Cellpose-SAM exports."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import torch
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from torch import nn

REPO_ID = "capicu-ai/cellpose-sam-wquant-w8a16"


def _digest(path: Path) -> str:
    value = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            value.update(block)
    return value.hexdigest()


def _restore(packed: torch.Tensor, scales: torch.Tensor, state: dict, device):
    bits = int(state["bits"])
    per_byte = 8 // bits
    packed = packed.to(device=device, dtype=torch.uint8).flatten()
    mask = (1 << bits) - 1
    codes = torch.stack(
        [(packed >> (index * bits)) & mask for index in range(per_byte)],
        dim=1,
    ).flatten()[: int(state["code_count"])]
    signed = codes.to(torch.int16) - int(state["qmax"])
    grouped = signed.reshape(
        int(state["rows"]),
        int(state["n_groups"]),
        int(state["group_size"]),
    )
    restored = grouped.float() * scales.to(device).unsqueeze(-1)
    restored = restored.reshape(int(state["rows"]), int(state["padded_inner"]))
    return restored[..., : int(state["inner"])].reshape(state["original_shape"])


def _activation(tensor: torch.Tensor, bits: int | None) -> torch.Tensor:
    if bits is None:
        return tensor
    qmax = (1 << (bits - 1)) - 1
    scale = tensor.detach().abs().amax().clamp_min(1e-12) / qmax
    return torch.round(tensor / scale).clamp(-qmax, qmax) * scale


class _Packed(nn.Module):
    def _setup(self, packed, scales, bias, state, activation_bits):
        self.register_buffer("packed_weight", packed.cpu())
        self.register_buffer("weight_scales", scales.cpu())
        self.register_buffer("bias", None if bias is None else bias.cpu())
        self.state = state
        self.activation_bits = activation_bits

    def _weight(self, tensor):
        return _restore(
            self.packed_weight,
            self.weight_scales,
            self.state,
            tensor.device,
        ).to(dtype=tensor.dtype)


class _Linear(_Packed):
    def __init__(self, source, **payload):
        super().__init__()
        self.in_features = source.in_features
        self.out_features = source.out_features
        self._setup(**payload)

    def forward(self, tensor):
        tensor = _activation(tensor, self.activation_bits)
        bias = None if self.bias is None else self.bias.to(tensor.device, tensor.dtype)
        return F.linear(tensor, self._weight(tensor), bias)


class _Conv2d(_Packed):
    def __init__(self, source, **payload):
        super().__init__()
        self.stride = source.stride
        self.padding = source.padding
        self.dilation = source.dilation
        self.groups = source.groups
        self._setup(**payload)

    def forward(self, tensor):
        tensor = _activation(tensor, self.activation_bits)
        bias = None if self.bias is None else self.bias.to(tensor.device, tensor.dtype)
        return F.conv2d(
            tensor,
            self._weight(tensor),
            bias,
            self.stride,
            self.padding,
            self.dilation,
            self.groups,
        )


def _replace(root: nn.Module, name: str, module: nn.Module) -> None:
    parent_name, _, child_name = name.rpartition(".")
    parent = root.get_submodule(parent_name) if parent_name else root
    if child_name.isdigit() and isinstance(parent, (nn.Sequential, nn.ModuleList)):
        parent[int(child_name)] = module
    else:
        setattr(parent, child_name, module)


def load_model(
    device: str | torch.device = "cpu",
    local_repo: str | Path | None = None,
):
    """Download and load the ready-to-run quantized Cellpose-SAM model."""
    root = (
        Path(local_repo)
        if local_repo is not None
        else Path(
            snapshot_download(
                REPO_ID,
                allow_patterns=[
                    "config.json",
                    "model.safetensors",
                ],
            )
        )
    )
    manifest = json.loads((root / "config.json").read_text())
    weights_path = root / "model.safetensors"
    if _digest(weights_path) != manifest["weights"]["sha256"]:
        raise ValueError("model checksum mismatch")

    from cellpose import models as cellpose_models

    target = torch.device(device)
    model = cellpose_models.CellposeModel(
        gpu=target.type == "cuda",
        pretrained_model=manifest["base_model"],
        device=target,
        use_bfloat16=False,
    )
    state = load_file(str(weights_path), device="cpu")
    network = model.net.cpu().eval()
    for item in manifest["replacements"]:
        name = item["name"]
        source = network.get_submodule(name)
        prefix = f"{name}."
        payload = {
            "packed": state[f"{prefix}packed_weight"],
            "scales": state[f"{prefix}weight_scales"],
            "bias": state.get(f"{prefix}bias"),
            "state": item["quant_state"],
            "activation_bits": item.get("activation_bits"),
        }
        replacement = (
            _Linear(source, **payload)
            if item["kind"] == "linear"
            else _Conv2d(source, **payload)
        )
        _replace(network, name, replacement)
    network.load_state_dict(state, strict=True)
    model.net = network.to(target).eval()
    return model