Spaces:
Running on Zero
Running on Zero
File size: 2,202 Bytes
440855b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """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
|