Patch external ComfyUI model_management for CPU fallback before import

comfy_integration/setup.py

@@ -86,43 +86,80 @@ def initialize_comfyui():
     # Ensure torch treats CUDA as unavailable on CPU‑only environment.
     import torch
     torch.cuda.is_available = lambda: False
-
-    import comfy.model_management as model_mgmt
-    # Patch get_torch_device to enforce CPU fallback.
-    def _cpu_fallback_get_torch_device():
-        """Return appropriate torch device, falling back to CPU when needed.
++
++    # Apply patch to the external ComfyUI model_management before import.
++    external_model_mgmt_path = os.path.join(APP_DIR, "comfy", "model_management.py")
++    if os.path.exists(external_model_mgmt_path):
++        try:
++            with open(external_model_mgmt_path, "r", encoding="utf-8") as f:
++                content = f.read()
++            start = content.find("def get_torch_device():")
++            if start != -1:
++                # Find the end of the function (next def or end of file)
++                end = content.find("\ndef ", start + 1)
++                if end == -1:
++                    end = len(content)
++                new_func = """def get_torch_device():
++    \"\"\"Return appropriate torch device, falling back to CPU when needed.\"\"\"
++    global directml_enabled, cpu_state
++    if directml_enabled:
++        return directml_device
++    if cpu_state == CPUState.MPS:
++        return torch.device(\"mps\")
++    if cpu_state == CPUState.CPU:
++        return torch.device(\"cpu\")
++    if is_intel_xpu():
++        return torch.device(\"xpu\", torch.xpu.current_device())
++    if is_ascend_npu():
++        return torch.device(\"npu\", torch.npu.current_device())
++    if is_mlu():
++        return torch.device(\"mlu\", torch.mlu.current_device())
++    if torch.cuda.is_available():
++        return torch.device(torch.cuda.current_device())
++    return torch.device(\"cpu\")
++"""
++                patched_content = content[:start] + new_func + content[end:]
++                with open(external_model_mgmt_path, "w", encoding="utf-8") as f:
++                    f.write(patched_content)
++        except Exception as e:
++            print(f"⚠️ Failed to patch external model_management: {e}")
++
++    import comfy.model_management as model_mgmt
++    # Patch get_torch_device to enforce CPU fallback.
++    def _cpu_fallback_get_torch_device():
++        """Return appropriate torch device, falling back to CPU when needed.
 +
 +        This patch replaces the original implementation that assumed a CUDA GPU
 +        and raised a RuntimeError on CPU‑only systems. It respects the existing
 +        ``directml_enabled`` and ``cpu_state`` flags, and falls back to CPU when
 +        no GPU backend is available.
 +        """
-        # DirectML path – unchanged.
-        if model_mgmt.directml_enabled:
-            return model_mgmt.directml_device
-        # Apple Silicon / MPS backend.
-        if model_mgmt.cpu_state == model_mgmt.CPUState.MPS:
-            return torch.device("mps")
-        # Explicit CPU request.
-        if model_mgmt.cpu_state == model_mgmt.CPUState.CPU:
-            return torch.device("cpu")
-        # Intel XPU, Ascend NPU, MLU – retain original handling.
-        if model_mgmt.is_intel_xpu():
-            return torch.device("xpu", torch.xpu.current_device())
-        if model_mgmt.is_ascend_npu():
-            return torch.device("npu", torch.npu.current_device())
-        if model_mgmt.is_mlu():
-            return torch.device("mlu", torch.mlu.current_device())
-        # CUDA – use if available (will be False due to monkey‑patch).
-        if torch.cuda.is_available():
-            return torch.device(torch.cuda.current_device())
-        # Default to CPU.
-        return torch.device("cpu")
-    model_mgmt.get_torch_device = _cpu_fallback_get_torch_device
-    import importlib
-    importlib.reload(model_mgmt)
-    
-    print("--- Environment Ready ---")
++        # DirectML path – unchanged.
++        if model_mgmt.directml_enabled:
++            return model_mgmt.directml_device
++        # Apple Silicon / MPS backend.
++        if model_mgmt.cpu_state == model_mgmt.CPUState.MPS:
++            return torch.device("mps")
++        # Explicit CPU request.
++        if model_mgmt.cpu_state == model_mgmt.CPUState.CPU:
++            return torch.device("cpu")
++        # Intel XPU, Ascend NPU, MLU – retain original handling.
++        if model_mgmt.is_intel_xpu():
++            return torch.device("xpu", torch.xpu.current_device())
++        if model_mgmt.is_ascend_npu():
++            return torch.device("npu", torch.npu.current_device())
++        if model_mgmt.is_mlu():
++            return torch.device("mlu", torch.mlu.current_device())
++        # CUDA – use if available (will be False due to monkey‑patch).
++        if torch.cuda.is_available():
++            return torch.device(torch.cuda.current_device())
++        # Default to CPU.
++        return torch.device("cpu")
++    model_mgmt.get_torch_device = _cpu_fallback_get_torch_device
++    import importlib
++    importlib.reload(model_mgmt)
++    
++    print("--- Environment Ready ---")
     
     print("✅ ComfyUI initialized with default attention mechanism.")