From dd35763ad6ba8db441fdccbe05f66fba797b8931 Mon Sep 17 00:00:00 2001
From: Max LI <114852914+2001maxli@users.noreply.github.com>
Date: Sun, 27 Sep 2026 14:53:25 -0700
Subject: [PATCH] [bugfix] Fix MLX prompt enhancement sampler compatibility
 (#1891)

---
 .github/workflows/ci-macos-mlx.yml            |  2 ++
 fastvideo/mlx_runtime/prompt_enhance.py       |  3 +-
 .../tests/mlx/test_mlx_prompt_enhance.py      | 32 +++++++++++++++++++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/ci-macos-mlx.yml b/.github/workflows/ci-macos-mlx.yml
index 62054c839..3ef19a311 100644
--- a/.github/workflows/ci-macos-mlx.yml
+++ b/.github/workflows/ci-macos-mlx.yml
@@ -92,6 +92,7 @@ jobs:
             fastvideo/tests/mlx/test_frame_upsample.py \
             fastvideo/tests/mlx/test_mlx_fast_spatial.py \
             fastvideo/tests/mlx/test_mlx_refine.py \
+            fastvideo/tests/mlx/test_mlx_prompt_enhance.py \
             fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \
             fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \
             fastvideo/tests/mlx/test_wan22_sample.py \
@@ -156,6 +157,7 @@ jobs:
             fastvideo/tests/mlx/test_frame_upsample.py \
             fastvideo/tests/mlx/test_mlx_fast_spatial.py \
             fastvideo/tests/mlx/test_mlx_refine.py \
+            fastvideo/tests/mlx/test_mlx_prompt_enhance.py \
             fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \
             fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \
             fastvideo/tests/mlx/test_wan22_sample.py \
diff --git a/fastvideo/mlx_runtime/prompt_enhance.py b/fastvideo/mlx_runtime/prompt_enhance.py
index 60a4ed4bb..78e95e35d 100644
--- a/fastvideo/mlx_runtime/prompt_enhance.py
+++ b/fastvideo/mlx_runtime/prompt_enhance.py
@@ -194,6 +194,7 @@ def enhance_prompt_mlx_lm(
     """
     try:
         from mlx_lm import generate, load
+        from mlx_lm.sample_utils import make_sampler
     except ImportError as exc:  # pragma: no cover - optional dep
         raise RuntimeError("mlx-lm is not installed. `uv pip install mlx-lm` or use "
                            "--enhance-prompt-backend template.") from exc
@@ -226,7 +227,7 @@ def enhance_prompt_mlx_lm(
         tokenizer,
         prompt=chat,
         max_tokens=max_tokens,
-        temp=temp,
+        sampler=make_sampler(temp=temp),
         verbose=False,
     )
     enhanced = _clean_llm_output(raw, original=text)
diff --git a/fastvideo/tests/mlx/test_mlx_prompt_enhance.py b/fastvideo/tests/mlx/test_mlx_prompt_enhance.py
index 51bcd5992..1279573cd 100644
--- a/fastvideo/tests/mlx/test_mlx_prompt_enhance.py
+++ b/fastvideo/tests/mlx/test_mlx_prompt_enhance.py
@@ -4,6 +4,10 @@
 from __future__ import annotations
 
 import json
+import sys
+from types import ModuleType, SimpleNamespace
+
+import pytest
 
 from fastvideo.mlx_runtime.prompt_enhance import (
     DEFAULT_ENHANCE_SYSTEM_PROMPT,
@@ -14,6 +18,34 @@ from fastvideo.mlx_runtime.prompt_enhance import (
 )
 
 
+@pytest.mark.parametrize("backend", ["mlx-lm", "auto"])
+def test_enhancement_uses_available_mlx_lm(monkeypatch, backend):
+    """Both backends must reach generation through the sampler-based API."""
+    package = ModuleType("mlx_lm")
+    sampling = ModuleType("mlx_lm.sample_utils")
+    temperatures = []
+    expected_sampler = object()
+    enhanced = "A cat walking on grass in warm evening light."
+
+    def make_sampler(temp):
+        temperatures.append(temp)
+        return expected_sampler
+
+    def generate(model_arg, tokenizer_arg, *, prompt, max_tokens, sampler, verbose):
+        assert sampler is expected_sampler
+        return enhanced
+
+    package.load = lambda path: (object(), SimpleNamespace())
+    package.generate = generate
+    sampling.make_sampler = make_sampler
+    monkeypatch.setitem(sys.modules, "mlx_lm", package)
+    monkeypatch.setitem(sys.modules, "mlx_lm.sample_utils", sampling)
+    result = enhance_prompt("A cat walking on grass", backend=backend)
+    assert result.backend == "mlx-lm"
+    assert result.enhanced == enhanced
+    assert temperatures == [0.6]
+
+
 def test_template_adds_cinematic_cues_to_thin_prompt() -> None:
     raw = "A red fox in the snow"
     out = enhance_prompt_template(raw)