From 9d7c3dc22ae06edcddcd775a6bb351e91e21bafa Mon Sep 17 00:00:00 2001
From: MarcBarnaba <marcbarnaba@gmail.com>
Date: Thu, 30 Jul 2026 18:10:42 +0200
Subject: [PATCH 1/2] Add shared cuda->mps->cpu device-selection helper for
 generators

BaseGenerator subclasses currently hardcode torch.cuda.is_available(),
so the four bundled models only run on Windows/Linux with an NVIDIA
GPU. Add select_device()/select_dtype() to base.py so extensions can
detect Apple Silicon (MPS) and fall back to CPU in one shared place
instead of duplicating the check per extension. unload() now also
releases MPS memory via torch.mps.empty_cache().

Windows/CUDA behavior is unchanged.

Ref: https://github.com/lightningpixel/modly/issues/166
---
 api/services/generators/base.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/api/services/generators/base.py b/api/services/generators/base.py
index a23b01db..f8277df7 100644
--- a/api/services/generators/base.py
+++ b/api/services/generators/base.py
@@ -11,6 +11,33 @@ class GenerationCancelled(Exception):
     """Raised by generators when a cancel_event is set mid-generation."""
 
 
+def select_device() -> str:
+    """
+    Picks the best torch device available on this machine: CUDA (NVIDIA) ->
+    MPS (Apple Silicon) -> CPU. Extensions should call this instead of
+    hardcoding `.cuda()` / `torch.cuda.is_available()` so device logic lives
+    in one place and Windows/CUDA behavior is unaffected.
+    """
+    import torch
+
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+def select_dtype(device: str):
+    """
+    Picks a safe default dtype for `device`. MPS has incomplete fp16 kernel
+    coverage (attention, layer norm, some interpolation ops), so extensions
+    get fp32 there and on CPU; CUDA keeps fp16 for speed/memory.
+    """
+    import torch
+
+    return torch.float16 if device == "cuda" else torch.float32
+
+
 def smooth_progress(
     progress_cb: Callable[[int, str], None],
     start: int,
@@ -83,6 +110,8 @@ def unload(self) -> None:
             import torch
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
+            elif torch.backends.mps.is_available():
+                torch.mps.empty_cache()
         except ImportError:
             pass
         # Force the OS to reclaim unused memory from this process

From ba1d6943f85994ceae510d7f3f7c21d645adb1c3 Mon Sep 17 00:00:00 2001
From: MarcBarnaba <marcbarnaba@gmail.com>
Date: Fri, 31 Jul 2026 01:02:50 +0200
Subject: [PATCH 2/2] Enable PyTorch's MPS CPU-fallback for extension
 subprocesses on macOS
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

torch.backends.mps hard-crashes with NotImplementedError on ops that
have no Metal kernel yet (e.g. 3D pooling), instead of transparently
falling back to CPU like on other backends. The opt-in env var that
enables the fallback must be set before the process's first `import
torch` — setting it inside select_device() is too late for extensions
that import torch before calling it (confirmed by testing TripoSG end
to end). Set it once, for every extension subprocess, in
ExtensionProcess._build_env() on macOS.

Ref: https://github.com/lightningpixel/modly/issues/166
---
 api/services/extension_process.py |  6 ++++++
 api/services/generators/base.py   | 10 ++++++++++
 2 files changed, 16 insertions(+)

diff --git a/api/services/extension_process.py b/api/services/extension_process.py
index 60ce6787..d4cc2272 100644
--- a/api/services/extension_process.py
+++ b/api/services/extension_process.py
@@ -75,6 +75,12 @@ def _build_env(self) -> dict:
         env["MODLY_API_DIR"] = str(Path(__file__).parent.parent)
         if sys.platform == "darwin":
             env.setdefault("NUMBA_DISABLE_JIT", "1")
+            # Must be set before the subprocess's first `import torch` — PyTorch
+            # reads this once to decide whether MPS ops with no Metal kernel
+            # (e.g. 3D pooling) fall back to CPU or raise NotImplementedError.
+            # Setting it inside generator.py is too late, since generator.py
+            # itself imports torch before calling select_device().
+            env.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
         # Pass the exact model_dir so runner.py doesn't have to re-derive it
         # from manifest["id"] (which is the ext_id, not the composite node id).
         # runner.py extracts the node id from MODEL_DIR's trailing path component.
diff --git a/api/services/generators/base.py b/api/services/generators/base.py
index f8277df7..b739a1ff 100644
--- a/api/services/generators/base.py
+++ b/api/services/generators/base.py
@@ -17,12 +17,22 @@ def select_device() -> str:
     MPS (Apple Silicon) -> CPU. Extensions should call this instead of
     hardcoding `.cuda()` / `torch.cuda.is_available()` so device logic lives
     in one place and Windows/CUDA behavior is unaffected.
+
+    PyTorch only falls back to CPU for MPS ops with no Metal kernel (e.g. 3D
+    pooling) — instead of raising NotImplementedError — if
+    PYTORCH_ENABLE_MPS_FALLBACK=1 is set *before the process's first `import
+    torch`*. That's set for every extension subprocess in
+    ExtensionProcess._build_env(); the setdefault() below is just a
+    best-effort backstop for callers that reach select_device() before
+    importing torch themselves.
     """
+    import os
     import torch
 
     if torch.cuda.is_available():
         return "cuda"
     if torch.backends.mps.is_available():
+        os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
         return "mps"
     return "cpu"