diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 6287b184a6..163ade10c4 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -5,15 +5,120 @@ Drop-in replacement for InferenceBackend — same interface, uses mlx-lm/mlx-vlm instead of torch/transformers for model loading and generation. """ +import json import os import threading from typing import Optional, Generator +from core.inference.message_content import content_to_text from core.inference.runtime_context import runtime_context_length from loggers import get_logger logger = get_logger(__name__) +def _mlx_vlm_model_config(model): + """Return the loaded MLX model config and its type, preferring whichever of + config / _config actually carries a model_type.""" + + def _model_type(cfg): + return cfg.get("model_type") if isinstance(cfg, dict) else getattr(cfg, "model_type", None) + + configs = [ + cfg + for cfg in (getattr(model, "config", None), getattr(model, "_config", None)) + if cfg is not None + ] + for cfg in configs: + model_type = _model_type(cfg) + if model_type is not None: + return cfg, model_type + return (configs[0] if configs else None), None + + +def _render_registered_vlm_prompt(processor, model, messages, num_images): + """Render through mlx-vlm when it declares a formatter for this model.""" + from mlx_vlm import prompt_utils + + config, model_type = _mlx_vlm_model_config(model) + if config is None: + return None + if model_type not in getattr(prompt_utils, "MODEL_CONFIG", {}): + return None + + rendered = prompt_utils.apply_chat_template( + processor, + config, + messages, + add_generation_prompt = True, + num_images = num_images, + ) + if isinstance(rendered, str) and rendered.strip(): + return rendered + raise RuntimeError("mlx-vlm's registered renderer returned an empty prompt.") + + +def _count_vlm_images(content): + if isinstance(content, list): + return sum(_count_vlm_images(item) for item in content) + if not isinstance(content, dict): + return 0 + if str(content.get("type", "")).lower() in ("image", "image_url", "input_image"): + return 1 + return _count_vlm_images(content.get("content")) + + +def _vlm_media_reprs(content): + if isinstance(content, list): + values = ( + {str(content), json.dumps(content, ensure_ascii = False)} + if _count_vlm_images(content) + else set() + ) + for item in content: + values.update(_vlm_media_reprs(item)) + return values + if not isinstance(content, dict): + return set() + if str(content.get("type", "")).lower() in ("image", "image_url", "input_image"): + return {str(content), json.dumps(content, ensure_ascii = False)} + return _vlm_media_reprs(content.get("content")) + + +def _prompt_serializes_vlm_media(prompt, messages): + """Detect templates that embed the exact structured media object repr.""" + media_reprs = set() + for message in messages: + if isinstance(message, dict): + media_reprs.update(_vlm_media_reprs(message.get("content"))) + text_content = [ + content_to_text(message.get("content")) for message in messages if isinstance(message, dict) + ] + return any( + prompt.count(media_repr) > sum(content.count(media_repr) for content in text_content) + for media_repr in media_reprs + ) + + +def _vlm_prompt_issue(prompt, messages): + if not isinstance(prompt, str) or not prompt.strip(): + return "an empty prompt" + if _prompt_serializes_vlm_media(prompt, messages): + return "serialized structured image content" + return None + + +def _vlm_messages_have_tool_history(messages): + return any( + isinstance(message, dict) + and ( + message.get("role") == "tool" + or message.get("tool_calls") + or message.get("tool_call_id") + ) + for message in messages + ) + + def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): """Map mlx stream stats onto the usage/timings shape llama-server emits.""" prompt_n = int(prompt_n or 0) @@ -422,9 +527,7 @@ class MLXInferenceBackend: {"type": "text", "text": content}, ] elif isinstance(content, list): - has_image = any( - p.get("type") == "image" for p in content if isinstance(p, dict) - ) + has_image = _count_vlm_images(content) > 0 if not has_image: content.insert(0, {"type": "image"}) break @@ -632,17 +735,87 @@ class MLXInferenceBackend: ): chat_target = getattr(self._processor, "tokenizer", self._processor) - prompt = apply_chat_template_for_generation( - chat_target, - messages, - tools = tools, - enable_thinking = enable_thinking, - reasoning_effort = reasoning_effort, - preserve_thinking = preserve_thinking, - ) - # mlx_vlm's stream_generate handles pixel_values (None for text-only) images = [image] if image is not None else None + attached_images = 0 if images is None else len(images) + structured_images = sum( + _count_vlm_images(message.get("content")) + for message in messages + if isinstance(message, dict) + ) + if structured_images != attached_images: + raise RuntimeError( + f"VLM conversation contains {structured_images} structured image " + f"item(s) for {attached_images} attached image(s)." + ) + prompt = None + has_tool_history = _vlm_messages_have_tool_history(messages) + prompt_error = None + try: + prompt = apply_chat_template_for_generation( + chat_target, + messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + except Exception as exc: + if images is None or has_tool_history: + raise + prompt_error = exc + prompt_issue = ( + _vlm_prompt_issue(prompt, messages) if prompt_error is None else "a rendering error" + ) + if prompt_issue and has_tool_history: + raise RuntimeError( + f"VLM chat template returned {prompt_issue} and cannot be recovered " + "without dropping tool-call history." + ) from prompt_error + + if images is not None and prompt_issue: + if tools or any( + value is not None + for value in (enable_thinking, reasoning_effort, preserve_thinking) + ): + if prompt_error is not None: + raise prompt_error + raise RuntimeError( + f"VLM chat template returned {prompt_issue} and cannot be recovered " + "without dropping requested tools or reasoning controls." + ) + try: + recovered_prompt = _render_registered_vlm_prompt( + self._processor, + self._model, + messages, + len(images), + ) + except Exception as recovery_error: + if prompt_error is not None: + raise prompt_error + raise RuntimeError( + f"VLM chat template returned {prompt_issue}; model-aware " + f"recovery failed: {recovery_error}" + ) from recovery_error + if recovered_prompt is None: + if prompt_error is not None: + raise prompt_error + raise RuntimeError( + f"VLM chat template returned {prompt_issue}, and no registered " + "MLX VLM renderer was available for this model." + ) + recovered_issue = _vlm_prompt_issue(recovered_prompt, messages) + if recovered_issue: + if prompt_error is not None: + raise prompt_error + raise RuntimeError( + f"Model-aware VLM rendering returned {recovered_issue} for " + f"{attached_images} attached image(s)." + ) + prompt = recovered_prompt + elif prompt_issue: + raise RuntimeError(f"VLM chat template returned {prompt_issue}.") from prompt_error from core.inference.chat_template_helpers import detect_think_prefill diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index 55a3198a6b..29fbb45158 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -333,6 +333,118 @@ def test_mlx_generate_chat_response_accepts_template_kwargs(): ), f"{name!r} must default to None so existing callers stay valid" +def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): + from core.inference.mlx_inference import MLXInferenceBackend + + calls = {"generic": [], "model": [], "stream": []} + state = {"generic": "serialized", "model": " model-aware"} + prompt_utils = SimpleNamespace( + MODEL_CONFIG = {"deepseek_vl_v2": object()}, + apply_chat_template = lambda *_args, **kwargs: ( + calls["model"].append(kwargs) or state["model"] + ), + ) + mlx_vlm = types.ModuleType("mlx_vlm") + mlx_vlm.prompt_utils = prompt_utils + mlx_vlm.stream_generate = lambda *_args, **kwargs: ( + calls["stream"].append((_args, kwargs)) + or iter([SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1)]) + ) + monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm) + + def generic(_target, _messages, **kwargs): + calls["generic"].append(kwargs) + if isinstance(state["generic"], Exception): + raise state["generic"] + if state["generic"] == "serialized": + return f"User: {_messages[0]['content']}" + return state["generic"] + + monkeypatch.setattr( + "core.inference.chat_template_helpers.apply_chat_template_for_generation", + generic, + ) + backend = MLXInferenceBackend() + backend._model = SimpleNamespace(config = {"model_type": "deepseek_vl_v2"}) + backend._processor = SimpleNamespace(tokenizer = SimpleNamespace()) + args = ([{"role": "user", "content": [{"type": "image"}]}], object(), 0, 1, 0, 0, 1, 1, None) + tools = [{"function": {"name": "search"}}] + assert list(backend._generate_vlm(*args)) == ["ok"] + assert calls["model"][0]["num_images"] == 1 + assert calls["stream"][0][0][2] == " model-aware" + with pytest.raises(RuntimeError, match = "dropping requested tools"): + list(backend._generate_vlm(*args, tools = tools)) + with pytest.raises(RuntimeError, match = "dropping requested tools or reasoning"): + list(backend._generate_vlm(*args, enable_thinking = False)) + backend._processor = SimpleNamespace(chat_template = "template") + state["generic"] = " healthy generic" + assert list(backend._generate_vlm(*args, tools = tools, enable_thinking = False)) == ["ok"] + assert calls["generic"][-1]["enable_thinking"] is False + assert calls["stream"][-1][0][2] == " healthy generic" + state["generic"] = "generic prompt" + text_messages = [{"role": "user", "content": "hello"}] + assert list(backend._generate_vlm(*((text_messages, None) + args[2:]), tools = tools)) == ["ok"] + assert calls["generic"][-1]["tools"] == tools + assert calls["stream"][-1][0][2] == "generic prompt" + two_images = [{"role": "user", "content": [{"type": "image"}, {"type": "image"}]}] + with pytest.raises(RuntimeError, match = "2 structured image item"): + list(backend._generate_vlm(*((two_images,) + args[1:]), tools = tools)) + state["generic"] = "serialized" + tool_history = args[0] + [{"role": "assistant", "tool_calls": [{"id": "call-1"}]}] + with pytest.raises(RuntimeError, match = "tool-call history"): + list(backend._generate_vlm(*((tool_history,) + args[1:]), tools = tools)) + state["generic"] = ValueError("generic rendering failed") + state["model"] = f"User: {args[0][0]['content']}" + with pytest.raises(ValueError, match = "generic rendering failed"): + list(backend._generate_vlm(*args)) + + +def test_mlx_vlm_image_injection_reuses_media_aliases(monkeypatch): + from core.inference.mlx_inference import MLXInferenceBackend, _prompt_serializes_vlm_media + + media = [{"type": "image"}] + quoted = [{"role": "user", "content": media}, {"role": "user", "content": f"Explain {media}"}] + assert _prompt_serializes_vlm_media(f"\n{media[0]}", quoted[:1]) + assert not _prompt_serializes_vlm_media(f"\nExplain {media}", quoted) + assert _prompt_serializes_vlm_media(f"User: {media}\nExplain {media}", quoted) + quoted[1]["content"] = [{"type": "text", "text": f'Explain "this" {media}'}] + assert not _prompt_serializes_vlm_media(f'\nExplain "this" {media}', quoted) + json_media = [{"type": "image_url"}] + json_repr = '{"type": "image_url"}' + assert _prompt_serializes_vlm_media(f"\n{json_repr}", [{"content": json_media}]) + assert not _prompt_serializes_vlm_media( + f"\nExplain {json_repr}", + [{"content": json_media}, {"content": f"Explain {json_repr}"}], + ) + + backend = MLXInferenceBackend() + backend._model = object() + backend._is_vlm = True + captured = [] + backend._generate_vlm = lambda messages, *_args, **_kwargs: ( + captured.append(messages) or iter(()) + ) + messages = [{"role": "user", "content": [{"type": "image_url"}]}] + list(backend.generate_chat_response(messages, image = object())) + assert captured[0][0]["content"] == [{"type": "image_url"}] + + +def test_mlx_vlm_model_config_prefers_config_with_model_type(): + from core.inference.mlx_inference import _mlx_vlm_model_config + + # config present but missing model_type must fall back to _config + m = SimpleNamespace(config = {}, _config = {"model_type": "deepseek_vl_v2"}) + assert _mlx_vlm_model_config(m) == ({"model_type": "deepseek_vl_v2"}, "deepseek_vl_v2") + # an object config whose model_type is None also falls back + m = SimpleNamespace(config = SimpleNamespace(model_type = None), _config = {"model_type": "qwen2_vl"}) + assert _mlx_vlm_model_config(m)[1] == "qwen2_vl" + # a config that already carries a model_type is preferred and returned unchanged + assert _mlx_vlm_model_config(SimpleNamespace(config = {"model_type": "gemma3"})) == ( + {"model_type": "gemma3"}, + "gemma3", + ) + + def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): """Mac text path must route through apply_chat_template_for_generation so reasoning / tool kwargs reach the tokenizer."""