checkpoint_compat ¶

def discover_mlx_checkpoint(*candidates: str | Path | None) -> Path | None:
    """Return the first candidate that is a packed MLX DiT directory."""
    for candidate in candidates:
        if candidate is None:
            continue
        path = Path(candidate)
        if is_mlx_dit_checkpoint(path):
            return path
    return None

def is_mlx_dit_checkpoint(path: str | Path) -> bool:
    """Return True if ``path`` is a packed FastMetal / MLX DiT directory."""
    checkpoint_dir = Path(path)
    return (checkpoint_dir / MLX_DIT_MANIFEST).is_file() and (checkpoint_dir / MLX_DIT_WEIGHTS).is_file()

def mlx_checkpoint_missing_hint(checkpoint_dir: str | Path) -> str:
    """Extra FileNotFoundError text when a directory is not a packed MLX DiT."""
    nvidia_reason = nvidia_fastwan_qad_reason(checkpoint_dir)
    prefix = (f"Not an MLX DiT checkpoint directory: {checkpoint_dir} "
              f"(expected {MLX_DIT_MANIFEST} and {MLX_DIT_WEIGHTS}).")
    if nvidia_reason is not None:
        return prefix + "\n\n" + _nvidia_fastwan_qad_message(Path(checkpoint_dir), nvidia_reason)
    return prefix + "\n\n" + _fastmetal_howto()

def nvidia_fastwan_qad_reason(path: str | Path) -> str | None:
    """Return a short reason if ``path`` looks like NVIDIA FastWan-QAD, else None."""
    checkpoint = Path(path)
    haystack = _path_haystack(checkpoint)
    if "int8" in haystack and is_mlx_dit_checkpoint(checkpoint):
        return None
    if _NVIDIA_FASTWAN_QAD_RE.search(haystack) and "int8" not in haystack:
        return "NVIDIA FastWan-QAD checkpoint name (NVFP4/FP8, not FastMetal INT8)"
    if (checkpoint / CUDA_QAD_OVERLAY_DIR).is_dir() or (checkpoint.parent / CUDA_QAD_OVERLAY_DIR).is_dir():
        return f"CUDA QAD overlay directory ({CUDA_QAD_OVERLAY_DIR}/)"
    for config_path in _config_candidates(checkpoint):
        marker = _quant_marker_in_file(config_path)
        if marker is not None:
            return f"{config_path.name} contains NVIDIA quantization marker {marker!r}"
    return None

def raise_if_unsupported_mlx_checkpoint(*paths: str | Path | None) -> None:
    """Raise if any path is an NVIDIA FastWan-QAD tree that must not run on MLX.

    Packed FastMetal / MLX DiT directories are always allowed, including the
    older FastWan-QAD-INT8 directory name, because those already contain
    ``mlx_dit.json``.
    """
    for path in paths:
        if path is None:
            continue
        checkpoint = Path(path)
        if is_mlx_dit_checkpoint(checkpoint):
            continue
        reason = nvidia_fastwan_qad_reason(checkpoint)
        if reason is None:
            continue
        raise UnsupportedMLXCheckpointError(_nvidia_fastwan_qad_message(checkpoint, reason))

def resolve_mlx_checkpoint(explicit: str | Path | None, *search_roots: str | Path | None) -> Path | None:
    """Prefer an explicit ``--mlx-checkpoint``, otherwise scan search roots."""
    if explicit is not None:
        return Path(explicit)
    return discover_mlx_checkpoint(*search_roots)