From 5a016e492cdb5e39ba862fbb443355db1045fc39 Mon Sep 17 00:00:00 2001
From: Manuel Cartagena Herrera <mcartagenah@users.noreply.github.com>
Date: Tue, 11 Aug 2026 20:23:43 -0400
Subject: [PATCH] fix(gallery): handle MPS float64 mask inputs (#5903)

Co-authored-by: Alexandre Teixeira <111787685+alteixeira20@users.noreply.github.com>
---
 routes/gallery/gallery_routes.py         | 29 ++++++++++++++------
 tests/test_gallery_model_input_device.py | 34 ++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 8 deletions(-)
 create mode 100644 tests/test_gallery_model_input_device.py

diff --git a/routes/gallery/gallery_routes.py b/routes/gallery/gallery_routes.py
index 457df210d..e6b5e0713 100644
--- a/routes/gallery/gallery_routes.py
+++ b/routes/gallery/gallery_routes.py
@@ -127,6 +127,25 @@ def _load_grounding_backend():
     return cached
 
 
+def _model_input_to_device(value, device: str, torch):
+    if not hasattr(value, "to"):
+        return value
+    if (
+        device == "mps"
+        and hasattr(torch, "float64")
+        and getattr(value, "dtype", None) == torch.float64
+    ):
+        return value.to(device=device, dtype=torch.float32)
+    return value.to(device)
+
+
+def _model_inputs_to_device(inputs, device: str, torch) -> Dict[str, Any]:
+    return {
+        key: _model_input_to_device(value, device, torch)
+        for key, value in inputs.items()
+    }
+
+
 def _ground_text_to_box(image, text: str, *, threshold: float = 0.05):
     query = (text or "").strip()
     if not query:
@@ -142,10 +161,7 @@ def _ground_text_to_box(image, text: str, *, threshold: float = 0.05):
         labels.append(f"a photo of {query}")
     try:
         inputs = processor(text=[labels], images=image, return_tensors="pt")
-        model_inputs = {
-            k: (v.to(device) if hasattr(v, "to") else v)
-            for k, v in inputs.items()
-        }
+        model_inputs = _model_inputs_to_device(inputs, device, torch)
         with torch.no_grad():
             outputs = model(**model_inputs)
         target_sizes = torch.tensor([[image.height, image.width]])
@@ -1869,10 +1885,7 @@ def setup_gallery_routes() -> APIRouter:
 
         try:
             inputs = processor(image, **kwargs)
-            model_inputs = {
-                k: (v.to(device) if hasattr(v, "to") else v)
-                for k, v in inputs.items()
-            }
+            model_inputs = _model_inputs_to_device(inputs, device, torch)
             with torch.no_grad():
                 outputs = model(**model_inputs)
             masks = processor.image_processor.post_process_masks(
diff --git a/tests/test_gallery_model_input_device.py b/tests/test_gallery_model_input_device.py
new file mode 100644
index 000000000..7e63eb000
--- /dev/null
+++ b/tests/test_gallery_model_input_device.py
@@ -0,0 +1,34 @@
+import routes.gallery_routes as gallery_routes
+
+
+class _TorchSentinel:
+    float32 = object()
+    float64 = object()
+
+
+class _FakeTensor:
+    def __init__(self, dtype):
+        self.dtype = dtype
+        self.to_args = None
+
+    def to(self, *args, **kwargs):
+        self.to_args = (args, kwargs)
+        return self
+
+
+def test_model_inputs_to_device_casts_mps_float64_to_float32():
+    float_tensor = _FakeTensor(_TorchSentinel.float64)
+    int_tensor = _FakeTensor("int64")
+    plain_value = object()
+
+    result = gallery_routes._model_inputs_to_device(
+        {"points": float_tensor, "labels": int_tensor, "plain": plain_value},
+        "mps",
+        _TorchSentinel,
+    )
+
+    assert result["points"] is float_tensor
+    assert float_tensor.to_args == ((), {"device": "mps", "dtype": _TorchSentinel.float32})
+    assert int_tensor.to_args == (("mps",), {})
+    assert result["plain"] is plain_value
+