diff --git a/fusion_mlx/engine/gguf_guard.py b/fusion_mlx/engine/gguf_guard.py new file mode 100644 index 00000000..8a969b68 --- /dev/null +++ b/fusion_mlx/engine/gguf_guard.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GGUF load guard for MLX engines. + +mlx-lm / mlx-vlm have no GGUF load path (mx.save_gguf is one-way export). +Loading a .gguf file or a GGUF-only directory crashes inside mlx_lm.load +with an opaque error. This guard detects GGUF targets up front and raises +a clear, actionable ValueError pointing the user at the MLX-native path. +""" + +from __future__ import annotations + +import logging +from pathlib import Path + +logger = logging.getLogger(__name__) + +GGUF_SUFFIX = ".gguf" +_MLX_WEIGHT_SUFFIXES = (".safetensors", ".npz") + + +def _local_exists(model_name: str) -> bool: + try: + return Path(model_name).exists() + except OSError: + return False + + +def is_gguf_model(model_name: str) -> bool: + """Return True if model_name points at a GGUF-only local target.""" + if not model_name: + return False + name = model_name.strip() + if name.lower().endswith(GGUF_SUFFIX): + return _local_exists(name) + p = Path(name) + try: + if not p.is_dir(): + return False + except OSError: + return False + children = list(p.iterdir()) + has_gguf = any(c.name.lower().endswith(GGUF_SUFFIX) for c in children) + if not has_gguf: + return False + has_mlx_weights = any( + c.name.lower().endswith(_MLX_WEIGHT_SUFFIXES) for c in children + ) + has_config = (p / "config.json").exists() + return not has_config and not has_mlx_weights + + +class GGUFLoadError(ValueError): + """Raised when a GGUF target is given to an MLX engine.""" + + +def assert_not_gguf(model_name: str, engine_kind: str = "MLX") -> None: + """Raise GGUFLoadError if model_name is a GGUF-only target. + + Call this right before mlx_lm.load / mlx_vlm.load so the error is + raised before any weight download or parse attempt. + """ + if not is_gguf_model(model_name): + return + logger.warning("GGUF target rejected by %s engine: %s", engine_kind, model_name) + raise GGUFLoadError( + f"{model_name} is a GGUF model; {engine_kind} engines cannot load " + f"GGUF (mlx-lm/mlx-vlm have no GGUF load path, only export). " + f"Use an MLX-native checkpoint instead: download an " + f"'mlx-community/-mlx' repo, or convert via the " + f"POST /v1/convert endpoint (pytorch|safetensors -> MLX). " + f"GGUF is a terminal format and cannot be converted back to MLX." + ) diff --git a/fusion_mlx/engines/batched.py b/fusion_mlx/engines/batched.py index a7651973..8c6bbf5a 100644 --- a/fusion_mlx/engines/batched.py +++ b/fusion_mlx/engines/batched.py @@ -303,6 +303,9 @@ def _load_model_sync(): if self._lora_path: load_kwargs["adapter_path"] = self._lora_path logger.info("Applying LoRA adapter: %s", self._lora_path) + from ..engine.gguf_guard import assert_not_gguf + + assert_not_gguf(self._model_name, engine_kind="LLM") model, tokenizer = load(self._model_name, **load_kwargs) elapsed = time.monotonic() - start # Estimate model size from loaded weights diff --git a/fusion_mlx/engines/embedding.py b/fusion_mlx/engines/embedding.py index 1c0149f9..af2b4c67 100644 --- a/fusion_mlx/engines/embedding.py +++ b/fusion_mlx/engines/embedding.py @@ -172,6 +172,9 @@ def load(self): logger.info( "Loading embedding model via mlx-embeddings: %s", self._model_name ) + from ..engine.gguf_guard import assert_not_gguf + + assert_not_gguf(self._model_name, engine_kind="Embedding") self._model, self._processor = load( self._model_name, tokenizer_config={"trust_remote_code": self._trust_remote_code}, diff --git a/fusion_mlx/engines/reranker.py b/fusion_mlx/engines/reranker.py index b7ed7de2..7117f7db 100644 --- a/fusion_mlx/engines/reranker.py +++ b/fusion_mlx/engines/reranker.py @@ -183,6 +183,9 @@ def _load_causal_lm(self) -> tuple[Any, Any]: model_path = str(self._model_name) tokenizer_config = {"trust_remote_code": self._trust_remote_code} + from ..engine.gguf_guard import assert_not_gguf + + assert_not_gguf(model_path, engine_kind="Reranker") custom_loaded = maybe_load_custom_quantization( model_path, is_vlm=False, @@ -239,6 +242,9 @@ def _load_jina_reranker(self) -> tuple[Any, Any]: model_path = str(self._model_name) tokenizer_config = {"trust_remote_code": self._trust_remote_code} + from ..engine.gguf_guard import assert_not_gguf + + assert_not_gguf(model_path, engine_kind="Reranker") custom_loaded = maybe_load_custom_quantization( model_path, is_vlm=False, @@ -494,6 +500,9 @@ def load(self) -> None: patch_qwen3_vl_processor_for_torch_free_image_loading() from mlx_embeddings import load + from ..engine.gguf_guard import assert_not_gguf + + assert_not_gguf(self._model_name, engine_kind="Reranker") self._model, self._processor = load( self._model_name, tokenizer_config={"trust_remote_code": self._trust_remote_code}, diff --git a/fusion_mlx/engines/vlm.py b/fusion_mlx/engines/vlm.py index de3a6f70..000f162e 100644 --- a/fusion_mlx/engines/vlm.py +++ b/fusion_mlx/engines/vlm.py @@ -346,6 +346,9 @@ async def start(self) -> None: def _load_vlm_sync(): start = time.monotonic() logger.info("Loading VLM model: %s", self._model_name) + from ..engine.gguf_guard import assert_not_gguf + + assert_not_gguf(self._model_name, engine_kind="VLM") model, processor = vlm_load( self._model_name, trust_remote_code=self._trust_remote_code ) diff --git a/tests/unit/test_gguf_guard.py b/tests/unit/test_gguf_guard.py new file mode 100644 index 00000000..c3dcdbc7 --- /dev/null +++ b/tests/unit/test_gguf_guard.py @@ -0,0 +1,69 @@ +import pytest + +from fusion_mlx.engine.gguf_guard import ( + GGUFLoadError, + assert_not_gguf, + is_gguf_model, +) + + +def test_direct_gguf_file_detected(tmp_path): + f = tmp_path / "model.gguf" + f.write_bytes(b"\x00") + assert is_gguf_model(str(f)) is True + + +def test_gguf_file_guard_raises(tmp_path): + f = tmp_path / "model.gguf" + f.write_bytes(b"\x00") + with pytest.raises(GGUFLoadError) as exc: + assert_not_gguf(str(f), engine_kind="LLM") + assert "GGUF" in str(exc.value) + assert "mlx-community" in str(exc.value) + + +def test_mlx_dir_not_flagged(tmp_path): + (tmp_path / "config.json").write_text("{}") + (tmp_path / "model.safetensors").write_bytes(b"\x00") + assert is_gguf_model(str(tmp_path)) is False + assert_not_gguf(str(tmp_path), engine_kind="LLM") + + +def test_gguf_only_dir_flagged(tmp_path): + (tmp_path / "model.gguf").write_bytes(b"\x00") + assert is_gguf_model(str(tmp_path)) is True + + +def test_gguf_dir_with_config_not_flagged(tmp_path): + (tmp_path / "config.json").write_text("{}") + (tmp_path / "model.gguf").write_bytes(b"\x00") + assert is_gguf_model(str(tmp_path)) is False + + +def test_empty_name_not_flagged(): + assert is_gguf_model("") is False + assert_not_gguf("", engine_kind="LLM") + + +def test_nonexistent_path_not_flagged(tmp_path): + assert is_gguf_model(str(tmp_path / "nope.gguf")) is False + assert_not_gguf(str(tmp_path / "nope.gguf"), engine_kind="LLM") + + +def test_hf_repo_id_not_flagged(): + assert is_gguf_model("mlx-community/Qwen2.5-7B-Instruct-4bit") is False + assert is_gguf_model("Qwen2.5-7B-Instruct") is False + + +def test_error_message_mentions_convert_endpoint(tmp_path): + f = tmp_path / "m.gguf" + f.write_bytes(b"\x00") + with pytest.raises(GGUFLoadError) as exc: + assert_not_gguf(str(f)) + assert "/v1/convert" in str(exc.value) + + +def test_guard_no_op_for_normal_path(tmp_path): + (tmp_path / "config.json").write_text("{}") + (tmp_path / "w.safetensors").write_bytes(b"\x00") + assert_not_gguf(str(tmp_path), engine_kind="VLM")