274 lines
10 KiB
Python
274 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Pre-quantized MLX checkpoint save/load for the FastWan DiT.
|
|
|
|
Loading the Diffusers fp32/fp16 checkpoint and quantizing at startup costs
|
|
both download size and load time on every run. This module persists an already
|
|
cast (and optionally already quantized) ``MLXWanDiT`` so 16 GB users download
|
|
and load roughly half the bytes and skip requantization entirely:
|
|
|
|
dit = mlx_dit_from_diffusers_safetensors(ckpt, cfg, quantization="int8")
|
|
save_mlx_dit_checkpoint(dit, "FastWan2.1-T2V-1.3B-mlx-int8")
|
|
...
|
|
dit = load_mlx_dit_checkpoint("FastWan2.1-T2V-1.3B-mlx-int8")
|
|
|
|
Format (one directory):
|
|
|
|
- ``mlx_dit.safetensors`` — every array, saved with ``mx.save_safetensors``.
|
|
Plain weights keep their key; a quantized weight ``K`` is stored as the
|
|
packed ``K`` plus ``K.scales`` (and ``K.biases`` for affine modes).
|
|
- ``mlx_dit.json`` — format version, the model config, the quantization spec,
|
|
and which keys are quantized, so the loader can rebuild ``QuantizedMatrix``
|
|
objects without guessing.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import shutil
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from fastvideo.logger import init_logger
|
|
from fastvideo.mlx_runtime.fastwan import (
|
|
MLXQuantizationSpec,
|
|
MLXWanDiT,
|
|
MLXWanTransformerBlock,
|
|
QuantizedMatrix,
|
|
ensure_quantization_supported,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
FORMAT_VERSION = 1
|
|
WEIGHTS_FILENAME = "mlx_dit.safetensors"
|
|
MANIFEST_FILENAME = "mlx_dit.json"
|
|
|
|
_BLOCK_PREFIX = "blocks"
|
|
|
|
_DTYPE_TO_NAME = {"float16": "fp16", "bfloat16": "bf16", "float32": "fp32"}
|
|
|
|
|
|
def _dtype_name(dtype) -> str:
|
|
"""Return the manifest name for a supported MLX data type.
|
|
|
|
Parameters:
|
|
dtype: The MLX data type to convert.
|
|
|
|
Returns:
|
|
str: The manifest name corresponding to the data type.
|
|
|
|
Raises:
|
|
ValueError: If the data type is not supported for checkpointing.
|
|
"""
|
|
import mlx.core as mx
|
|
|
|
for raw, name in _DTYPE_TO_NAME.items():
|
|
if dtype == getattr(mx, raw):
|
|
return name
|
|
raise ValueError(f"Unsupported MLX dtype for checkpointing: {dtype}")
|
|
|
|
|
|
def _name_to_dtype(name: str):
|
|
"""Convert a manifest dtype name to its corresponding MLX dtype.
|
|
|
|
Parameters:
|
|
name (str): Manifest name, such as ``"fp16"``, ``"bf16"``, or ``"fp32"``.
|
|
|
|
Returns:
|
|
The corresponding MLX dtype.
|
|
"""
|
|
import mlx.core as mx
|
|
|
|
return {"fp16": mx.float16, "bf16": mx.bfloat16, "fp32": mx.float32}[name]
|
|
|
|
|
|
def _flatten_weights(dit: MLXWanDiT) -> dict[str, Any]:
|
|
"""Combine model and transformer-block weights into a single flattened mapping.
|
|
|
|
Parameters:
|
|
dit (MLXWanDiT): Model whose weights should be flattened.
|
|
|
|
Returns:
|
|
dict[str, Any]: Mapping containing top-level weights and indexed transformer-block weights.
|
|
"""
|
|
flat: dict[str, Any] = dict(dit.weights)
|
|
for index, block in enumerate(dit.blocks):
|
|
for name, value in block.weights.items():
|
|
flat[f"{_BLOCK_PREFIX}.{index}.{name}"] = value
|
|
return flat
|
|
|
|
|
|
def save_mlx_dit_checkpoint(dit: MLXWanDiT, checkpoint_dir: str | Path) -> Path:
|
|
"""Save a plain or quantized MLX Wan DiT checkpoint to a directory.
|
|
|
|
Parameters:
|
|
dit (MLXWanDiT): Model whose weights and configuration will be saved.
|
|
checkpoint_dir (str | Path): Destination directory for the checkpoint.
|
|
|
|
Returns:
|
|
Path: Path to the checkpoint directory.
|
|
"""
|
|
import mlx.core as mx
|
|
|
|
checkpoint_dir = Path(checkpoint_dir)
|
|
arrays: dict[str, Any] = {}
|
|
quantized: dict[str, dict[str, Any]] = {}
|
|
spec: MLXQuantizationSpec | None = None
|
|
for key, value in _flatten_weights(dit).items():
|
|
if isinstance(value, QuantizedMatrix):
|
|
if spec is not None and value.spec != spec:
|
|
raise ValueError(f"Mixed quantization specs in one checkpoint ({spec} vs {value.spec} at '{key}') "
|
|
"are not supported.")
|
|
spec = value.spec
|
|
arrays[key] = value.weight
|
|
arrays[f"{key}.scales"] = value.scales
|
|
if value.biases is not None:
|
|
arrays[f"{key}.biases"] = value.biases
|
|
quantized[key] = {
|
|
"dequantized_dtype": _dtype_name(value.dequantized_dtype),
|
|
"has_biases": value.biases is not None,
|
|
}
|
|
else:
|
|
arrays[key] = value
|
|
|
|
manifest = {
|
|
"format_version": FORMAT_VERSION,
|
|
"config": dit.config,
|
|
"num_blocks": len(dit.blocks),
|
|
"quantization": None if spec is None else {
|
|
"mode": spec.mode,
|
|
"bits": spec.bits,
|
|
"group_size": spec.group_size,
|
|
},
|
|
"quantized_keys": quantized,
|
|
}
|
|
|
|
manifest_json = json.dumps(manifest, indent=2)
|
|
checkpoint_dir.parent.mkdir(parents=True, exist_ok=True)
|
|
staging_dir = Path(tempfile.mkdtemp(dir=checkpoint_dir.parent, prefix=f".{checkpoint_dir.name}.staging-"))
|
|
backup_root: Path | None = None
|
|
try:
|
|
staged_weights = staging_dir / WEIGHTS_FILENAME
|
|
staged_manifest = staging_dir / MANIFEST_FILENAME
|
|
mx.save_safetensors(str(staged_weights), arrays)
|
|
staged_manifest.write_text(manifest_json)
|
|
if checkpoint_dir.exists():
|
|
backup_root = Path(tempfile.mkdtemp(dir=checkpoint_dir.parent, prefix=f".{checkpoint_dir.name}.backup-"))
|
|
try:
|
|
checkpoint_dir.replace(backup_root / checkpoint_dir.name)
|
|
except Exception:
|
|
shutil.rmtree(backup_root, ignore_errors=True)
|
|
raise
|
|
try:
|
|
staging_dir.replace(checkpoint_dir)
|
|
except Exception:
|
|
if backup_root is not None:
|
|
(backup_root / checkpoint_dir.name).replace(checkpoint_dir)
|
|
shutil.rmtree(backup_root, ignore_errors=True)
|
|
raise
|
|
if backup_root is not None:
|
|
shutil.rmtree(backup_root, ignore_errors=True)
|
|
finally:
|
|
shutil.rmtree(staging_dir, ignore_errors=True)
|
|
logger.info("Saved MLX DiT checkpoint (%d arrays, quantization=%s) to %s", len(arrays),
|
|
spec.label if spec else "none", checkpoint_dir)
|
|
return checkpoint_dir
|
|
|
|
|
|
def load_mlx_dit_checkpoint(checkpoint_dir: str | Path, *, compile: bool = False) -> MLXWanDiT:
|
|
"""
|
|
Reconstruct an MLXWanDiT model from a versioned checkpoint.
|
|
|
|
Parameters:
|
|
checkpoint_dir (str | Path): Directory containing the checkpoint manifest and weights.
|
|
compile (bool): Whether to configure the reconstructed model for compilation.
|
|
|
|
Returns:
|
|
MLXWanDiT: The reconstructed model.
|
|
|
|
Raises:
|
|
FileNotFoundError: If the checkpoint manifest or weights file is missing.
|
|
ValueError: If the checkpoint format is unsupported or block weights are incomplete.
|
|
"""
|
|
import mlx.core as mx
|
|
|
|
checkpoint_dir = Path(checkpoint_dir)
|
|
manifest_path = checkpoint_dir / MANIFEST_FILENAME
|
|
weights_path = checkpoint_dir / WEIGHTS_FILENAME
|
|
if not manifest_path.exists() or not weights_path.exists():
|
|
from fastvideo.mlx_runtime.checkpoint_compat import mlx_checkpoint_missing_hint
|
|
raise FileNotFoundError(mlx_checkpoint_missing_hint(checkpoint_dir))
|
|
|
|
manifest = json.loads(manifest_path.read_text())
|
|
version = manifest.get("format_version")
|
|
if version != FORMAT_VERSION:
|
|
raise ValueError(f"MLX DiT checkpoint {checkpoint_dir} has format_version={version}; "
|
|
f"this FastVideo build reads version {FORMAT_VERSION}. Re-export the checkpoint.")
|
|
|
|
spec = None
|
|
if manifest["quantization"] is not None:
|
|
spec = MLXQuantizationSpec(**manifest["quantization"])
|
|
# The packed layout of mx.quantize output is mode-specific, so a build
|
|
# that cannot run the mode cannot use these arrays at all.
|
|
ensure_quantization_supported(spec)
|
|
|
|
arrays = mx.load(str(weights_path))
|
|
quantized_keys: dict[str, dict[str, Any]] = manifest["quantized_keys"]
|
|
|
|
def rebuild(key: str):
|
|
"""
|
|
Reconstructs a weight array or quantized matrix from checkpoint data.
|
|
|
|
Parameters:
|
|
key (str): The weight key to rebuild.
|
|
|
|
Returns:
|
|
The stored array for an unquantized weight or a reconstructed quantized matrix.
|
|
"""
|
|
if key not in quantized_keys:
|
|
return arrays[key]
|
|
info = quantized_keys[key]
|
|
assert spec is not None, f"Quantized key '{key}' in a checkpoint without a quantization spec"
|
|
return QuantizedMatrix(
|
|
weight=arrays[key],
|
|
scales=arrays[f"{key}.scales"],
|
|
biases=arrays[f"{key}.biases"] if info["has_biases"] else None,
|
|
spec=spec,
|
|
dequantized_dtype=_name_to_dtype(info["dequantized_dtype"]),
|
|
)
|
|
|
|
config = manifest["config"]
|
|
block_keys: dict[int, list[str]] = {}
|
|
top_level_keys: list[str] = []
|
|
for key in arrays:
|
|
if key.endswith(".scales") or key.endswith(".biases"):
|
|
continue
|
|
if key.startswith(f"{_BLOCK_PREFIX}."):
|
|
index_str, _, _ = key[len(_BLOCK_PREFIX) + 1:].partition(".")
|
|
block_keys.setdefault(int(index_str), []).append(key)
|
|
else:
|
|
top_level_keys.append(key)
|
|
|
|
weights = {key: rebuild(key) for key in top_level_keys}
|
|
|
|
num_blocks = int(manifest["num_blocks"])
|
|
if sorted(block_keys) != list(range(num_blocks)):
|
|
raise ValueError(f"MLX DiT checkpoint {checkpoint_dir} is missing block weights: "
|
|
f"manifest says {num_blocks} blocks, found indices {sorted(block_keys)}.")
|
|
|
|
inner_dim = int(config["num_attention_heads"]) * int(config["attention_head_dim"])
|
|
blocks = []
|
|
for index in range(num_blocks):
|
|
prefix = f"{_BLOCK_PREFIX}.{index}."
|
|
block_weights = {key[len(prefix):]: rebuild(key) for key in block_keys[index]}
|
|
blocks.append(
|
|
MLXWanTransformerBlock(
|
|
block_weights,
|
|
dim=inner_dim,
|
|
ffn_dim=int(config["ffn_dim"]),
|
|
num_heads=int(config["num_attention_heads"]),
|
|
eps=float(config["eps"]),
|
|
))
|
|
return MLXWanDiT(weights, blocks, config, compile=compile)
|