comfy-env 0.0.8__py3-none-any.whl

comfy_env/env/detection.py

@@ -0,0 +1,176 @@
+"""
+GPU detection for automatic CUDA version selection.
+
+Detects Blackwell GPUs (RTX 50xx, B100, B200) which require CUDA 12.8,
+vs older GPUs which use CUDA 12.4.
+
+This runs BEFORE PyTorch is installed, so we use nvidia-smi directly.
+"""
+
+import subprocess
+from typing import List, Dict, Optional
+
+
+def detect_gpu_info() -> List[Dict[str, str]]:
+    """
+    Detect GPU name and compute capability using nvidia-smi.
+
+    Returns:
+        List of dicts with 'name' and 'compute_cap' keys.
+        Empty list if detection fails.
+    """
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", "--query-gpu=name,compute_cap", "--format=csv,noheader"],
+            capture_output=True,
+            text=True,
+            timeout=10
+        )
+        if result.returncode == 0:
+            gpus = []
+            for line in result.stdout.strip().split('\n'):
+                if not line.strip():
+                    continue
+                parts = [p.strip() for p in line.split(',')]
+                name = parts[0] if parts else "Unknown"
+                cc = parts[1] if len(parts) > 1 else "0.0"
+                gpus.append({"name": name, "compute_cap": cc})
+            return gpus
+    except FileNotFoundError:
+        # nvidia-smi not found - no NVIDIA GPU or driver not installed
+        pass
+    except subprocess.TimeoutExpired:
+        pass
+    except Exception:
+        pass
+    return []
+
+
+def is_blackwell_gpu(name: str, compute_cap: str) -> bool:
+    """
+    Check if a GPU is Blackwell architecture.
+
+    Args:
+        name: GPU name from nvidia-smi
+        compute_cap: Compute capability string (e.g., "8.9", "12.0")
+
+    Returns:
+        True if Blackwell (requires CUDA 12.8)
+    """
+    name_upper = name.upper()
+
+    # Check by name patterns
+    blackwell_patterns = [
+        "RTX 50",      # RTX 5090, 5080, 5070, etc.
+        "RTX50",       # Without space
+        "B100",        # Datacenter Blackwell
+        "B200",        # Datacenter Blackwell
+        "GB202",       # Blackwell die
+        "GB203",
+        "GB205",
+        "GB206",
+        "GB207",
+    ]
+
+    if any(pattern in name_upper for pattern in blackwell_patterns):
+        return True
+
+    # Check by compute capability (10.0+ = Blackwell)
+    try:
+        cc = float(compute_cap)
+        if cc >= 10.0:
+            return True
+    except (ValueError, TypeError):
+        pass
+
+    return False
+
+
+def needs_cuda_128() -> bool:
+    """
+    Check if any detected GPU requires CUDA 12.8.
+
+    Returns:
+        True if Blackwell GPU detected, False otherwise.
+    """
+    gpus = detect_gpu_info()
+
+    for gpu in gpus:
+        if is_blackwell_gpu(gpu["name"], gpu["compute_cap"]):
+            return True
+
+    return False
+
+
+def is_legacy_gpu(compute_cap: str) -> bool:
+    """
+    Check if GPU is Pascal or older (requires legacy CUDA/PyTorch).
+
+    Args:
+        compute_cap: Compute capability string (e.g., "6.1", "7.5")
+
+    Returns:
+        True if Pascal or older (compute < 7.5)
+    """
+    try:
+        cc = float(compute_cap)
+        return cc < 7.5  # Turing starts at 7.5
+    except (ValueError, TypeError):
+        return False
+
+
+def detect_cuda_version() -> Optional[str]:
+    """
+    Get recommended CUDA version based on detected GPU.
+
+    Returns:
+        "12.4" for Pascal or older (compute < 7.5),
+        "12.8" for Turing or newer (compute >= 7.5),
+        None if no GPU detected.
+
+    GPU Architecture Reference:
+        - Pascal (GTX 10xx, P100): compute 6.0-6.1 → CUDA 12.4
+        - Turing (RTX 20xx, T4): compute 7.5 → CUDA 12.8
+        - Ampere (RTX 30xx, A100): compute 8.0-8.6 → CUDA 12.8
+        - Ada (RTX 40xx, L40): compute 8.9 → CUDA 12.8
+        - Hopper (H100): compute 9.0 → CUDA 12.8
+        - Blackwell (RTX 50xx, B100/B200): compute 10.0+ → CUDA 12.8
+    """
+    gpus = detect_gpu_info()
+    if not gpus:
+        return None
+
+    # Check if any GPU is legacy (Pascal or older)
+    for gpu in gpus:
+        if is_legacy_gpu(gpu.get("compute_cap", "0.0")):
+            return "12.4"
+
+    # Turing or newer - use modern stack
+    return "12.8"
+
+
+def get_gpu_summary() -> str:
+    """
+    Get a human-readable summary of detected GPUs.
+
+    Returns:
+        Summary string for logging.
+    """
+    gpus = detect_gpu_info()
+
+    if not gpus:
+        return "No NVIDIA GPU detected"
+
+    lines = []
+    for i, gpu in enumerate(gpus):
+        cc = gpu.get("compute_cap", "0.0")
+        is_legacy = is_legacy_gpu(cc)
+        if is_legacy:
+            tag = " [Pascal - CUDA 12.4]"
+        elif is_blackwell_gpu(gpu["name"], cc):
+            tag = " [Blackwell - CUDA 12.8]"
+        else:
+            tag = " [CUDA 12.8]"
+        lines.append(f"  GPU {i}: {gpu['name']} (sm_{cc.replace('.', '')}){tag}")
+
+    return "\n".join(lines)