load_decoder.py

"""Strictly load the 500K single-frame weights into a MiniMax H3 VAE."""

from __future__ import annotations

import hashlib
from pathlib import Path
from typing import Any

from safetensors.torch import load_file


EXPECTED_SHA256 = "6c5ff2caa8fade6769f4dd53ee244f77a06652c8cb66b2fedc93f75046d9f001"
EXPECTED_TENSOR_COUNT = 585


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


def _split_state(state: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
    decoder = {
        key.removeprefix("decoder."): value
        for key, value in state.items()
        if key.startswith("decoder.")
    }
    post_quant_conv = {
        key.removeprefix("post_quant_conv."): value
        for key, value in state.items()
        if key.startswith("post_quant_conv.")
    }
    unexpected = sorted(
        key
        for key in state
        if not key.startswith("decoder.")
        and not key.startswith("post_quant_conv.")
    )
    if unexpected or not decoder or not post_quant_conv:
        raise RuntimeError(
            "single-frame decoder key contract differs: "
            f"unexpected={unexpected[:5]}, decoder={len(decoder)}, "
            f"post_quant_conv={len(post_quant_conv)}"
        )
    return decoder, post_quant_conv


def load_single_frame_decoder(
    vae: Any,
    checkpoint_path: str | Path,
    *,
    verify_sha256: bool = True,
) -> Any:
    """Replace only ``decoder`` and ``post_quant_conv`` on an H3 VAE."""

    checkpoint = Path(checkpoint_path)
    if verify_sha256 and _sha256(checkpoint) != EXPECTED_SHA256:
        raise RuntimeError("single-frame decoder SHA-256 differs")
    state = load_file(str(checkpoint), device="cpu")
    if len(state) != EXPECTED_TENSOR_COUNT:
        raise RuntimeError(
            f"expected {EXPECTED_TENSOR_COUNT} tensors, found {len(state)}"
        )
    decoder, post_quant_conv = _split_state(state)
    vae.decoder.load_state_dict(decoder, strict=True)
    vae.post_quant_conv.load_state_dict(post_quant_conv, strict=True)
    return vae