checkpoint ¶

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():
        raise FileNotFoundError(f"Not an MLX DiT checkpoint directory: {checkpoint_dir} "
                                f"(expected {MANIFEST_FILENAME} and {WEIGHTS_FILENAME}).")

    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)

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