From 5701c79fab40927ae94b30fcdf79337bb0ad32e8 Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Tue, 30 Jul 2024 23:04:15 +0200 Subject: [PATCH] Prevent Grounding DINO and Segment Anything from being moved to MPS - they don't work on MPS devices. --- invokeai/backend/grounded_sam/grounding_dino_pipeline.py | 4 ++++ invokeai/backend/grounded_sam/segment_anything_model.py | 3 +++ 2 files changed, 7 insertions(+) diff --git a/invokeai/backend/grounded_sam/grounding_dino_pipeline.py b/invokeai/backend/grounded_sam/grounding_dino_pipeline.py index 028aeb283d..1fc92b5e12 100644 --- a/invokeai/backend/grounded_sam/grounding_dino_pipeline.py +++ b/invokeai/backend/grounded_sam/grounding_dino_pipeline.py @@ -21,6 +21,10 @@ class GroundingDinoPipeline: return results def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> "GroundingDinoPipeline": + # HACK(ryand): The GroundingDinoPipeline does not work on MPS devices. We only allow it to be moved to CPU or + # CUDA. + if device is not None and device.type not in {"cpu", "cuda"}: + device = None self._pipeline.model.to(device=device, dtype=dtype) self._pipeline.device = self._pipeline.model.device return self diff --git a/invokeai/backend/grounded_sam/segment_anything_model.py b/invokeai/backend/grounded_sam/segment_anything_model.py index 86106b869c..1cc105c5fd 100644 --- a/invokeai/backend/grounded_sam/segment_anything_model.py +++ b/invokeai/backend/grounded_sam/segment_anything_model.py @@ -16,6 +16,9 @@ class SegmentAnythingModel: self._sam_processor = sam_processor def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> "SegmentAnythingModel": + # HACK(ryand): The SAM pipeline does not work on MPS devices. We only allow it to be moved to CPU or CUDA. + if device is not None and device.type not in {"cpu", "cuda"}: + device = None self._sam_model.to(device=device, dtype=dtype) return self