fix(gallery): handle MPS float64 mask inputs (#5903)

Co-authored-by: Alexandre Teixeira <111787685+alteixeira20@users.noreply.github.com>

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 = {
+        model_inputs = _model_inputs_to_device(inputs, device, torch)
-            k: (v.to(device) if hasattr(v, "to") else v)
-            for k, v in inputs.items()
-        }
         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 = {
+            model_inputs = _model_inputs_to_device(inputs, device, torch)
-                k: (v.to(device) if hasattr(v, "to") else v)
-                for k, v in inputs.items()
-            }
             with torch.no_grad():
                 outputs = model(**model_inputs)
             masks = processor.image_processor.post_process_masks(

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
+