feat: add experimental macOS MPS inference support

boogu/models/transformers/transformer_boogu.py

@@ -206,7 +206,7 @@ def __init__(
         self.head_dim = dim // num_attention_heads
         self.modulation = modulation
 
-        if "cpu" in os.getenv("device", "cpu"):
+        if "cpu" in os.getenv("device", "cpu") or "mps" == os.getenv("device", ""):
             processor = BooguImageAttnProcessor()
 
         else:
@@ -416,15 +416,15 @@ def __init__(
         self.modulation = modulation
         self.hidden_size = dim
 
-        if "cpu" in os.getenv("device", "cpu"):
+        if "cpu" in os.getenv("device", "cpu") or "mps" == os.getenv("device", ""):
             processor = BooguImageAttnProcessor()
         else:
             try:
                 processor = BooguImageAttnProcessorFlash2Varlen()
             except ImportError:
                 processor = BooguImageAttnProcessor()
 
-        if "cpu" in os.getenv("device", "cpu"):
+        if "cpu" in os.getenv("device", "cpu") or "mps" == os.getenv("device", ""):
             double_stream_processor = BooguImageDoubleStreamSelfAttnProcessor(
                 head_dim=self.head_dim,
                 num_attention_heads=num_attention_heads,

boogu/pipelines/boogu/pipeline_boogu.py

@@ -271,8 +271,8 @@ def __init__(
 
     def _validate_device_format(
         self,
-        device: Literal[None, "cpu", "cuda", "cuda:x"] = "cpu",
-        rewriter_device: Literal[None, "cpu", "cuda", "cuda:x", "auto"] = "cpu",
+        device: Literal[None, "cpu", "cuda", "cuda:x", "mps"] = "cpu",
+        rewriter_device: Literal[None, "cpu", "cuda", "cuda:x", "auto", "mps"] = "cpu",
     ):
         device = device.lower() if isinstance(device, str) else device
         rewriter_device = (
@@ -294,8 +294,8 @@ def _check_device_strategy_validity(
         enable_model_cpu_offload_flag: bool = None,
         enable_sequential_cpu_offload_flag: bool = None,
         enable_group_offload_flag: bool = None,
-        rewriter_device: Literal[None, "cpu", "cuda", "cuda:x", "auto"] = None,
-        device: Literal[None, "cpu", "cuda", "cuda:x"] = None,
+        rewriter_device: Literal[None, "cpu", "cuda", "cuda:x", "auto", "mps"] = None,
+        device: Literal[None, "cpu", "cuda", "cuda:x", "mps"] = None,
         use_rewrite_text_instruction: bool = False,
         use_dashscope_remote_rewriting: bool = False,
         dashscope_api_key: str = None,
@@ -377,11 +377,11 @@ def _normalize_device_name(device_name):
 
     def devices_manager(
         self,
-        instant_device_2_use: Literal[None, "cpu", "cuda", "cuda:x"] = None,
-        instant_rewriter_device: Literal[None, "cpu", "cuda", "cuda:x", "auto"] = None,
-        user_set_pipe_device: Literal[None, "cpu", "cuda", "cuda:x"] = None,
-        user_set_rewriter_device: Literal[None, "cpu", "cuda", "cuda:x", "auto"] = None,
-        execution_device: Literal[None, "cpu", "cuda", "cuda:x"] = None,
+        instant_device_2_use: Literal[None, "cpu", "cuda", "cuda:x", "mps"] = None,
+        instant_rewriter_device: Literal[None, "cpu", "cuda", "cuda:x", "auto", "mps"] = None,
+        user_set_pipe_device: Literal[None, "cpu", "cuda", "cuda:x", "mps"] = None,
+        user_set_rewriter_device: Literal[None, "cpu", "cuda", "cuda:x", "auto", "mps"] = None,
+        execution_device: Literal[None, "cpu", "cuda", "cuda:x", "mps"] = None,
         unload_rewriter_level: Literal["keep", "cpu", "destroy"] = "destroy",
         enable_model_cpu_offload_flag: bool = None,
         enable_sequential_cpu_offload_flag: bool = None,
@@ -2716,12 +2716,19 @@ def __call__(
         return_dict: bool = True,
         verbose: bool = False,
         step_func=None,
-        device: Literal[None, "cpu", "cuda", "cuda:x"] = "cuda",
-        rewriter_device: Literal[None, "cpu", "cuda", "cuda:x", "auto"] = "cpu",
+        device: Literal[None, "cpu", "cuda", "cuda:x", "mps"] = None,
+        rewriter_device: Literal[None, "cpu", "cuda", "cuda:x", "auto", "mps"] = "cpu",
         unload_rewriter_level: Literal["keep", "cpu", "destroy"] = "destroy",
         enable_inner_devices_manager: bool = False,
     ):
 
+        if device is None:
+            device = str(self._execution_device)
+        if device.startswith("mps"):
+            device = "mps"
+        elif device == "cuda:0":
+            device = "cuda"
+
         if enable_inner_devices_manager is not None:
             self.enable_inner_devices_manager = enable_inner_devices_manager
 

boogu/pipelines/boogu/pipeline_boogu_turbo.py

@@ -158,6 +158,28 @@ def processing(self, *args, **kwargs):
         dtype = kwargs["dtype"]
         step_func = kwargs.get("step_func", None)
 
+        # Workaround for MPS backend: `@torch.no_grad()` triggers an MPS Graph
+        # assertion during transformer forward. Enabling grad bypasses the bug.
+        mps_no_grad_workaround = not torch.is_grad_enabled() and str(latents.device).startswith("mps")
+        if mps_no_grad_workaround:
+            torch.set_grad_enabled(True)
+
+        try:
+            return self._dmd_processing(
+                latents, ref_latents, instruction_embeds, freqs_cis,
+                instruction_attention_mask, num_inference_steps,
+                timesteps, device, dtype, step_func,
+            )
+        finally:
+            if mps_no_grad_workaround:
+                torch.set_grad_enabled(False)
+
+    def _dmd_processing(
+        self,
+        latents, ref_latents, instruction_embeds, freqs_cis,
+        instruction_attention_mask, num_inference_steps,
+        timesteps, device, dtype, step_func,
+    ):
         # --- DMD constraints (mirror the standalone turbo pipeline) ---
         task_type = self._get_task_type_by_ref_latents(ref_latents)
         if (

boogu/utils/validator_utils.py

@@ -7,7 +7,7 @@ def get_device_validator(additional_types: Optional[List[str]] = None):
     """
     Factory function that returns a validator for device arguments.
 
-    Base supported formats: 'cpu', 'cuda', or 'cuda:x' (where x is an integer).
+    Base supported formats: 'cpu', 'cuda', 'cuda:x' (where x is an integer), or 'mps'.
     Additional formats can be provided via `additional_types` (e.g., ['auto']).
     """
     # Initialize as an empty list if None is provided
@@ -27,7 +27,7 @@ def validate_device_format(value: str):
         # ^ and $ ensure the entire string is matched
         # (cpu|cuda) matches these exact words
         # |cuda:\d+ matches 'cuda:' followed by one or more digits (\d+)
-        if re.match(r"^(cpu|cuda|cuda:\d+)$", value):
+        if re.match(r"^(cpu|cuda|cuda:\d+|mps)$", value):
             return value
 
         # Check if the value is in the additionally allowed types (e.g., 'auto')
@@ -36,7 +36,7 @@ def validate_device_format(value: str):
 
         # If it doesn't match any allowed format, raise ArgumentTypeError.
         # argparse will automatically catch this and print a user-friendly error message.
-        allowed_msg = "'cpu', 'cuda', 'cuda:x' (where x is an integer like 'cuda:0')"
+        allowed_msg = "'cpu', 'cuda', 'cuda:x' (where x is an integer like 'cuda:0'), or 'mps'"
         if additional_types:
             allowed_msg += f", or one of {additional_types}"
 
@@ -87,7 +87,7 @@ def _normalize_bool_flag(value):
         return False
 
     device = str(device).strip().lower()
-    if not re.match(r"^(cpu|cuda|cuda:\d+)$", device):
+    if not re.match(r"^(cpu|cuda|cuda:\d+|mps)$", device):
         return False
 
     # CPU offload strategies need a non-CPU execution device to be meaningful.