anima_mlx/utils/weights.py

"""Weight loading helpers for local Anima safetensors artifacts.

The Anima checkpoints are BF16. ``safetensors`` can expose those tensors to
PyTorch reliably, while its NumPy path cannot represent BF16 on the system
Python used by this project. Keep the PyTorch dependency lazy so metadata-only
and source-contract tests can run without importing heavy runtimes.
"""

from __future__ import annotations

from pathlib import Path
from typing import Callable, Mapping


MLX_WEIGHT_ROOT_NAME = "mlx_weights"


def load_torch_safetensors_subset(
    path: str | Path,
    *,
    prefix: str | None = None,
    strip_prefix: str | None = None,
    key_filter: Callable[[str], bool] | None = None,
    dtype: str = "float32",
) -> dict[str, object]:
    """Load a filtered safetensors subset as CPU PyTorch tensors.

    Args:
        path: Safetensors file path.
        prefix: Optional source-key prefix filter.
        strip_prefix: Optional prefix to remove from returned keys.
        key_filter: Optional arbitrary source-key predicate.
        dtype: Either ``"float32"`` or ``"native"``.
    """

    try:
        import torch
        from safetensors import safe_open
    except ModuleNotFoundError as exc:  # pragma: no cover - depends on env
        raise RuntimeError("PyTorch and safetensors are required to load BF16 weights") from exc

    if dtype not in {"float32", "native"}:
        raise ValueError(f"unsupported dtype: {dtype}")

    weights: dict[str, object] = {}
    with safe_open(str(path), framework="pt", device="cpu") as handle:
        for source_key in handle.keys():
            if prefix is not None and not source_key.startswith(prefix):
                continue
            if key_filter is not None and not key_filter(source_key):
                continue
            target_key = source_key
            if strip_prefix is not None:
                if not target_key.startswith(strip_prefix):
                    continue
                target_key = target_key.removeprefix(strip_prefix)
            tensor = handle.get_tensor(source_key)
            if dtype == "float32":
                tensor = tensor.float()
            weights[target_key] = tensor
    return weights


def torch_to_mlx_weights(
    weights: Mapping[str, object],
    *,
    dtype: str = "float32",
) -> dict[str, object]:
    """Convert a mapping of CPU PyTorch tensors to MLX arrays.

    MLX is imported lazily because importing it inside the default sandbox can
    crash the local Metal runtime. Callers should do this only in MLX-enabled
    processes.
    """

    import numpy as np
    import mlx.core as mx

    dtype_map = {
        "float32": mx.float32,
        "float16": mx.float16,
        "bfloat16": mx.bfloat16,
    }
    if dtype not in dtype_map:
        raise ValueError(f"unsupported MLX dtype: {dtype}")

    converted: dict[str, object] = {}
    for key, tensor in weights.items():
        if hasattr(tensor, "detach"):
            array = tensor.detach().cpu().float().numpy()
        else:
            array = np.asarray(tensor, dtype=np.float32)
        converted[key] = mx.array(array, dtype=dtype_map[dtype])
    return converted


def load_native_mlx_safetensors(
    path: str | Path,
    *,
    prefix: str | None = None,
    strip_prefix: str | None = None,
    key_filter: Callable[[str], bool] | None = None,
    dtype: str = "native",
) -> dict[str, object]:
    """Load a safetensors file with MLX and optionally filter/cast its arrays."""

    import mlx.core as mx

    dtype_map = {
        "native": None,
        "float32": mx.float32,
        "float16": mx.float16,
        "bfloat16": mx.bfloat16,
    }
    if dtype not in dtype_map:
        raise ValueError(f"unsupported MLX dtype: {dtype}")

    arrays = mx.load(str(path), format="safetensors")
    target_dtype = dtype_map[dtype]
    filtered: dict[str, object] = {}
    for source_key, value in arrays.items():
        if prefix is not None and not source_key.startswith(prefix):
            continue
        if key_filter is not None and not key_filter(source_key):
            continue
        target_key = source_key
        if strip_prefix is not None:
            if not target_key.startswith(strip_prefix):
                continue
            target_key = target_key.removeprefix(strip_prefix)
        if target_dtype is not None and hasattr(value, "astype"):
            value = value.astype(target_dtype)
        filtered[target_key] = value
    return filtered


def load_mlx_safetensors_subset(
    path: str | Path,
    *,
    prefix: str | None = None,
    strip_prefix: str | None = None,
    key_filter: Callable[[str], bool] | None = None,
    dtype: str = "float32",
) -> dict[str, object]:
    """Load a safetensors subset into MLX arrays via the BF16-safe torch path."""

    if _is_converted_mlx_weight_path(path):
        return load_native_mlx_safetensors(
            path,
            prefix=prefix,
            strip_prefix=strip_prefix,
            key_filter=key_filter,
            dtype=dtype,
        )

    torch_weights = load_torch_safetensors_subset(
        path,
        prefix=prefix,
        strip_prefix=strip_prefix,
        key_filter=key_filter,
        dtype="float32",
    )
    return torch_to_mlx_weights(torch_weights, dtype=dtype)


def _is_converted_mlx_weight_path(path: str | Path) -> bool:
    resolved = Path(path)
    return MLX_WEIGHT_ROOT_NAME in resolved.parts or ".mlx." in resolved.name