Support VLM serving with MLX (#10)

Author: liang2kl

README.md

@@ -47,6 +47,8 @@ python -m paroquant.cli.chat --model $MODEL
 python -m paroquant.cli.serve --model $MODEL --port 8000
 ```
 
+Add `--llm-only` if you do not wish to load the VLM components.
+
 ### Agent with Tool Calling
 
 Start the API server first, then install the agent dependencies and run:

assets/model_card.jinja

@@ -49,6 +49,8 @@ python -m paroquant.cli.chat --model {{ paro_model_path }}
 python -m paroquant.cli.serve --model {{ paro_model_path }} --port 8000
 ```
 
+Add `--llm-only` if you do not wish to load the VLM components.
+
 {% if supports_tool_call -%}
 ### Agent with Tool Calling
 

paroquant/cli/serve.py

@@ -20,20 +20,73 @@ def _serve_vllm():
 
 
 def _serve_mlx():
-    import mlx_lm.server
-    from mlx_lm.utils import load_tokenizer
+    import os
+    import sys
 
     from paroquant.inference.backends.mlx.load import load as paro_load
 
-    def _patched_load(path_or_hf_repo, tokenizer_config=None, adapter_path=None, **kwargs):
-        model, _, _ = paro_load(path_or_hf_repo, force_text=True)
-        tokenizer = load_tokenizer(path_or_hf_repo, tokenizer_config_extra=tokenizer_config)
-        tokenizer._tool_call_start = None
-        tokenizer._tool_call_end = None
-        return model, tokenizer
+    original_argv = list(sys.argv)
+    model_arg = None
+    llm_only = False
+    stripped_argv = [original_argv[0]]
+    i = 1
+    while i < len(original_argv):
+        arg = original_argv[i]
+        if arg == "--model":
+            if i + 1 >= len(original_argv):
+                raise ValueError("--model expects a value")
+            model_arg = original_argv[i + 1]
+            i += 2
+            continue
+        if arg.startswith("--model="):
+            model_arg = arg.split("=", 1)[1]
+            i += 1
+            continue
+        if arg == "--llm-only":
+            llm_only = True
+            i += 1
+            continue
+        stripped_argv.append(arg)
+        i += 1
+
+    if not model_arg:
+        model_arg = os.environ.get("MODEL")
+    if not model_arg:
+        raise ValueError("Model path is required (use --model or MODEL environment variable).")
+
+    model, processor, is_vlm = paro_load(model_arg, force_text=llm_only)
+
+    if is_vlm:
+        import mlx_vlm.server as mlx_server
+
+        os.environ["MODEL"] = model_arg
+        sys.argv = stripped_argv
+
+        def _patched_load(path_or_hf_repo, *args, **kwargs):
+            return model, processor
+
+        _uvicorn_run = mlx_server.uvicorn.run
+
+        def _run_no_reload(*args, **kwargs):
+            kwargs["reload"] = False
+            return _uvicorn_run(*args, **kwargs)
+
+        mlx_server.uvicorn.run = _run_no_reload
+    else:
+        import mlx_lm.server as mlx_server
+
+        tokenizer = getattr(processor, "tokenizer", processor)
+        if hasattr(tokenizer, "_tool_call_start"):
+            tokenizer._tool_call_start = None
+        if hasattr(tokenizer, "_tool_call_end"):
+            tokenizer._tool_call_end = None
+        sys.argv = stripped_argv
+
+        def _patched_load(path_or_hf_repo, tokenizer_config=None, adapter_path=None, **kwargs):
+            return model, tokenizer
 
-    mlx_lm.server.load = _patched_load
-    mlx_lm.server.main()
+    mlx_server.load = _patched_load
+    mlx_server.main()
 
 
 def main():

paroquant/inference/backends/mlx/load.py

@@ -176,6 +176,8 @@ def load(model_path: str, lazy: bool = False, force_text: bool = False) -> tuple
         weights = _convert_autoawq(weights, group_size)
     if hasattr(model, "sanitize"):
         weights = model.sanitize(weights)
+    if is_vlm and hasattr(model, "vision_tower") and hasattr(model.vision_tower, "sanitize"):
+        weights = model.vision_tower.sanitize(weights)
 
     _patch_rotation_layers(model, weights, bits, group_size)
     model.load_weights(list(weights.items()), strict=False)