Upload model_management.py

comfy/model_management.py

@@ -177,14 +177,17 @@ def get_torch_device():
     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())
+        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())
+            else:
+                return torch.device(torch.cuda.current_device())
         else:
-            return torch.device(torch.cuda.current_device())
+            return torch.device("cpu")  # Fallback на CPU, если CUDA недоступен
 
 def get_total_memory(dev=None, torch_total_too=False):
     global directml_enabled