Update ComfyUI/comfy/model_management.py

ComfyUI/comfy/model_management.py

@@ -163,30 +163,35 @@ def is_ixuca():
     return False
 
 def get_torch_device():
+    import torch
     global directml_enabled
     global cpu_state
+
     if directml_enabled:
         global directml_device
         return directml_device
+
     if cpu_state == CPUState.MPS:
         return torch.device("mps")
+
     if cpu_state == CPUState.CPU:
         return torch.device("cpu")
-    else:
-        if is_intel_xpu():
-            return torch.device("xpu", torch.xpu.current_device())
-        elif is_ascend_npu():
-            return torch.device("npu", torch.npu.current_device())
-        elif is_mlu():
-            return torch.device("mlu", torch.mlu.current_device())
-        else:
-            return torch.device(torch.cuda.current_device())
-            if torch.cuda.is_available():
+
+    if is_intel_xpu():
+        return torch.device("xpu", torch.xpu.current_device())
+    elif is_ascend_npu():
+        return torch.device("npu", torch.npu.current_device())
+    elif is_mlu():
+        return torch.device("mlu", torch.mlu.current_device())
+
+    # Fallback to CUDA if available
+    if torch.cuda.is_available():
         return torch.device(f"cuda:{torch.cuda.current_device()}")
     else:
-        print("⚠️ No GPU found. Falling back to CPU.")
+        print("⚠️ No compatible GPU found. Using CPU.")
         return torch.device("cpu")
 
+
 def get_total_memory(dev=None, torch_total_too=False):
     global directml_enabled
     if dev is None: