fix(studio): recover MLX VLM image prompts (#7094)
* fix(studio): recover mlx vlm image prompts * fix(studio): detect serialized vlm media items * studio: recover MLX VLM prompts when model_type only lives on _config _mlx_vlm_model_config only fell back to _config when config was entirely missing, so a model that exposes a config without a model_type (while _config carries it) skipped model-aware recovery. Prefer whichever of config / _config actually has a model_type. Adds a focused test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Daniel Han <danielhanchen@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Long Yixing
GitHub
parent
4beb0a3a5f
commit
14d0e853fa
2 changed files with 297 additions and 12 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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": "<image> 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] == "<image> 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"] = "<image> 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] == "<image> 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"<image>\n{media[0]}", quoted[:1])
|
||||
assert not _prompt_serializes_vlm_media(f"<image>\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'<image>\nExplain "this" {media}', quoted)
|
||||
json_media = [{"type": "image_url"}]
|
||||
json_repr = '{"type": "image_url"}'
|
||||
assert _prompt_serializes_vlm_media(f"<image>\n{json_repr}", [{"content": json_media}])
|
||||
assert not _prompt_serializes_vlm_media(
|
||||
f"<image>\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."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue