Fix remaining CUDA device capability errors in wgp.py and sage2_core.py

wan/modules/sage2_core.py

@@ -54,19 +54,29 @@ import warnings
 import os
 
 def is_sage2_supported():
-    device_count = torch.cuda.device_count()
-    for i in range(device_count):
-        major, minor = torch.cuda.get_device_capability(i)
-        if major < 8:
-            return False
-    return True
+    if not torch.cuda.is_available():
+        return False
+    try:
+        device_count = torch.cuda.device_count()
+        for i in range(device_count):
+            major, minor = torch.cuda.get_device_capability(i)
+            if major < 8:
+                return False
+        return True
+    except (RuntimeError, AssertionError):
+        return False
 
 def get_cuda_arch_versions():
-    cuda_archs = []
-    for i in range(torch.cuda.device_count()):
-        major, minor = torch.cuda.get_device_capability(i)
-        cuda_archs.append(f"sm{major}{minor}")
-    return cuda_archs
+    if not torch.cuda.is_available():
+        return ["sm75"]  # Fallback architecture
+    try:
+        cuda_archs = []
+        for i in range(torch.cuda.device_count()):
+            major, minor = torch.cuda.get_device_capability(i)
+            cuda_archs.append(f"sm{major}{minor}")
+        return cuda_archs
+    except (RuntimeError, AssertionError):
+        return ["sm75"]  # Fallback architecture
 
 def sageattn(
     qkv_list,

wgp.py

@@ -1496,12 +1496,25 @@ attention_modes_installed = get_attention_modes()
 attention_modes_supported = get_supported_attention_modes()
 args = _parse_args()
 
-major, minor = torch.cuda.get_device_capability(args.gpu if len(args.gpu) > 0 else None)
-if  major < 8:
-    print("Switching to FP16 models when possible as GPU architecture doesn't support optimed BF16 Kernels")
-    bfloat16_supported = False
+# Check CUDA availability and device capability safely
+if torch.cuda.is_available():
+    try:
+        device_id = args.gpu[0] if len(args.gpu) > 0 else None
+        major, minor = torch.cuda.get_device_capability(device_id)
+        if major < 8:
+            print("Switching to FP16 models when possible as GPU architecture doesn't support optimized BF16 Kernels")
+            bfloat16_supported = False
+        else:
+            bfloat16_supported = True
+    except (RuntimeError, AssertionError, IndexError):
+        # Fallback if CUDA device is not accessible or invalid device ID
+        print("CUDA device not accessible or invalid device ID, using fallback settings")
+        major, minor = 0, 0
+        bfloat16_supported = False
 else:
-    bfloat16_supported = True
+    print("CUDA not available, using CPU mode")
+    major, minor = 0, 0
+    bfloat16_supported = False
 
 args.flow_reverse = True
 processing_device = args.gpu