Patch get_torch_device in setup to enforce CPU fallback

comfy_integration/setup.py

@@ -1,98 +1,135 @@
-import os
-import sys
-import shutil
-
-from core.settings import *
-
-def move_and_overwrite(src, dst):
-    if os.path.isdir(src):
-        if os.path.exists(dst):
-            shutil.rmtree(dst)
-        shutil.move(src, dst)
-    elif os.path.isfile(src):
-        if os.path.exists(dst):
-            os.remove(dst)
-        shutil.move(src, dst)
-
-def initialize_comfyui():
-    APP_DIR = sys.path[0]
-    COMFYUI_TEMP_DIR = "ComfyUI_temp"
-
-    print("--- Cloning ComfyUI Repository ---")
-    if not os.path.exists(COMFYUI_TEMP_DIR):
-        os.system(f"git clone https://github.com/comfy-Org/ComfyUI {COMFYUI_TEMP_DIR}")
-        print("✅ ComfyUI repository cloned.")
-    else:
-        print("✅ ComfyUI repository already exists.")
-
-    print(f"--- Merging ComfyUI from '{COMFYUI_TEMP_DIR}' to '{APP_DIR}' ---")
-    for item in os.listdir(COMFYUI_TEMP_DIR):
-        src_path = os.path.join(COMFYUI_TEMP_DIR, item)
-        dst_path = os.path.join(APP_DIR, item)
-        if item == '.git':
-            continue
-        move_and_overwrite(src_path, dst_path)
-    
-    try:
-        shutil.rmtree(COMFYUI_TEMP_DIR)
-        print("✅ ComfyUI merged and temporary directory removed.")
-    except OSError as e:
-        print(f"⚠️ Could not remove temporary directory '{COMFYUI_TEMP_DIR}': {e}")
-
-
-    print("--- Cloning third-party extensions for ComfyUI ---")
-
-    # 1. ComfyUI_IPAdapter_plus
-    ipadapter_plus_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI_IPAdapter_plus")
-    if not os.path.exists(ipadapter_plus_path):
-        os.system(f"git clone https://github.com/cubiq/ComfyUI_IPAdapter_plus.git {ipadapter_plus_path}")
-        print("✅ ComfyUI_IPAdapter_plus extension cloned.")
-    else:
-        print("✅ ComfyUI_IPAdapter_plus extension already exists.")
-
-    # 2. ComfyUI-InstantX-IPAdapter-SD3
-    ipadapter_plus_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI-InstantX-IPAdapter-SD3")
-    if not os.path.exists(ipadapter_plus_path):
-        os.system(f"git clone https://github.com/Slickytail/ComfyUI-InstantX-IPAdapter-SD3.git {ipadapter_plus_path}")
-        print("✅ ComfyUI-InstantX-IPAdapter-SD3 extension cloned.")
-    else:
-        print("✅ ComfyUI-InstantX-IPAdapter-SD3 extension already exists.")
-
-    # 3. ComfyUI-IPAdapter-Flux
-    ipadapter_flux_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI-IPAdapter-Flux")
-    if not os.path.exists(ipadapter_flux_path):
-        os.system(f"git clone https://github.com/Shakker-Labs/ComfyUI-IPAdapter-Flux.git {ipadapter_flux_path}")
-        print("✅ ComfyUI-IPAdapter-Flux extension cloned.")
-    else:
-        print("✅ ComfyUI-IPAdapter-Flux extension already exists.")
-
-    # 4. ComfyUI-Newbie-Nodes
-    newbie_nodes_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI-Newbie-Nodes")
-    if not os.path.exists(newbie_nodes_path):
-        os.system(f"git clone https://github.com/NewBieAI-Lab/ComfyUI-Newbie-Nodes.git {newbie_nodes_path}")
-        print("✅ ComfyUI-Newbie-Nodes extension cloned.")
-    else:
-        print("✅ ComfyUI-Newbie-Nodes extension already exists.")
-
-    # 5. ComfyUI-Anima-LLLite
-    anima_controlnet_lllite_nodes_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI-Anima-LLLite")
-    if not os.path.exists(anima_controlnet_lllite_nodes_path):
-        os.system(f"git clone https://github.com/kohya-ss/ComfyUI-Anima-LLLite.git {anima_controlnet_lllite_nodes_path}")
-        print("✅ ComfyUI-Anima-LLLite extension cloned.")
-    else:
-        print("✅ ComfyUI-Anima-LLLite extension already exists.")
-
-    print(f"✅ Current working directory is: {os.getcwd()}")
-    
-    import comfy.model_management
-    print("--- Environment Ready ---")
-
-    print("✅ ComfyUI initialized with default attention mechanism.")
-
-    for dir_path in CATEGORY_TO_DIR_MAP.values():
-        os.makedirs(os.path.join(APP_DIR, dir_path), exist_ok=True)
-    
-    os.makedirs(os.path.join(APP_DIR, INPUT_DIR), exist_ok=True)
-    os.makedirs(os.path.join(APP_DIR, OUTPUT_DIR), exist_ok=True)
-    
+import os
+import sys
+import shutil
+
+from core.settings import *
+
+def move_and_overwrite(src, dst):
+    if os.path.isdir(src):
+        if os.path.exists(dst):
+            shutil.rmtree(dst)
+        shutil.move(src, dst)
+    elif os.path.isfile(src):
+        if os.path.exists(dst):
+            os.remove(dst)
+        shutil.move(src, dst)
+
+def initialize_comfyui():
+    APP_DIR = sys.path[0]
+    COMFYUI_TEMP_DIR = "ComfyUI_temp"
+
+    print("--- Cloning ComfyUI Repository ---")
+    if not os.path.exists(COMFYUI_TEMP_DIR):
+        os.system(f"git clone https://github.com/comfy-Org/ComfyUI {COMFYUI_TEMP_DIR}")
+        print("✅ ComfyUI repository cloned.")
+    else:
+        print("✅ ComfyUI repository already exists.")
+
+    print(f"--- Merging ComfyUI from '{COMFYUI_TEMP_DIR}' to '{APP_DIR}' ---")
+    for item in os.listdir(COMFYUI_TEMP_DIR):
+        src_path = os.path.join(COMFYUI_TEMP_DIR, item)
+        dst_path = os.path.join(APP_DIR, item)
+        if item == '.git':
+            continue
+        move_and_overwrite(src_path, dst_path)
+    
+    try:
+        shutil.rmtree(COMFYUI_TEMP_DIR)
+        print("✅ ComfyUI merged and temporary directory removed.")
+    except OSError as e:
+        print(f"⚠️ Could not remove temporary directory '{COMFYUI_TEMP_DIR}': {e}")
+
+
+    print("--- Cloning third-party extensions for ComfyUI ---")
+
+    # 1. ComfyUI_IPAdapter_plus
+    ipadapter_plus_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI_IPAdapter_plus")
+    if not os.path.exists(ipadapter_plus_path):
+        os.system(f"git clone https://github.com/cubiq/ComfyUI_IPAdapter_plus.git {ipadapter_plus_path}")
+        print("✅ ComfyUI_IPAdapter_plus extension cloned.")
+    else:
+        print("✅ ComfyUI_IPAdapter_plus extension already exists.")
+
+    # 2. ComfyUI-InstantX-IPAdapter-SD3
+    ipadapter_plus_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI-InstantX-IPAdapter-SD3")
+    if not os.path.exists(ipadapter_plus_path):
+        os.system(f"git clone https://github.com/Slickytail/ComfyUI-InstantX-IPAdapter-SD3.git {ipadapter_plus_path}")
+        print("✅ ComfyUI-InstantX-IPAdapter-SD3 extension cloned.")
+    else:
+        print("✅ ComfyUI-InstantX-IPAdapter-SD3 extension already exists.")
+
+    # 3. ComfyUI-IPAdapter-Flux
+    ipadapter_flux_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI-IPAdapter-Flux")
+    if not os.path.exists(ipadapter_flux_path):
+        os.system(f"git clone https://github.com/Shakker-Labs/ComfyUI-IPAdapter-Flux.git {ipadapter_flux_path}")
+        print("✅ ComfyUI-IPAdapter-Flux extension cloned.")
+    else:
+        print("✅ ComfyUI-IPAdapter-Flux extension already exists.")
+
+    # 4. ComfyUI-Newbie-Nodes
+    newbie_nodes_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI-Newbie-Nodes")
+    if not os.path.exists(newbie_nodes_path):
+        os.system(f"git clone https://github.com/NewBieAI-Lab/ComfyUI-Newbie-Nodes.git {newbie_nodes_path}")
+        print("✅ ComfyUI-Newbie-Nodes extension cloned.")
+    else:
+        print("✅ ComfyUI-Newbie-Nodes extension already exists.")
+
+    # 5. ComfyUI-Anima-LLLite
+    anima_controlnet_lllite_nodes_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI-Anima-LLLite")
+    if not os.path.exists(anima_controlnet_lllite_nodes_path):
+        os.system(f"git clone https://github.com/kohya-ss/ComfyUI-Anima-LLLite.git {anima_controlnet_lllite_nodes_path}")
+        print("✅ ComfyUI-Anima-LLLite extension cloned.")
+    else:
+        print("✅ ComfyUI-Anima-LLLite extension already exists.")
+
+    print(f"✅ Current working directory is: {os.getcwd()}")
+    
+    import comfy.model_management as model_mgmt
+    # Patch get_torch_device to ensure CPU fallback on systems without a GPU.
+    import torch
+    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.
+        if torch.cuda.is_available():
+            return torch.device(torch.cuda.current_device())
+        # No GPU detected; default to CPU.
+        return torch.device("cpu")
+    model_mgmt.get_torch_device = _cpu_fallback_get_torch_device
+    
+    # Ensure the patched function is used by subsequent imports.
+    import importlib
+    importlib.reload(model_mgmt)
+    
+    print("--- Environment Ready ---")
+
+    print("✅ ComfyUI initialized with default attention mechanism.")
+
+    for dir_path in CATEGORY_TO_DIR_MAP.values():
+        os.makedirs(os.path.join(APP_DIR, dir_path), exist_ok=True)
+    
+    os.makedirs(os.path.join(APP_DIR, INPUT_DIR), exist_ok=True)
+    os.makedirs(os.path.join(APP_DIR, OUTPUT_DIR), exist_ok=True)
+    
     print("✅ All required model directories are present.")