comfy-env 0.1.15__py3-none-any.whl → 0.1.16__py3-none-any.whl

comfy_env/config/parser.py

@@ -1,88 +1,42 @@
-"""
-Configuration parsing for comfy-env.
-
-Loads comfy-env.toml (a superset of pixi.toml) and provides typed config objects.
-"""
+"""Configuration parsing for comfy-env."""
 
 import copy
-import sys
-from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional, Dict, Any, List
-import tomli
-
-# --- Types&Constants ---
-CONFIG_FILE_NAME = "comfy-env.toml"
+from typing import Any, Dict, List, Optional
 
-@dataclass
-class NodeReq:
-    """A node dependency (another ComfyUI custom node)."""
-    name: str
-    repo: str  # GitHub repo, e.g., "owner/repo"
+import tomli
 
-@dataclass
-class ComfyEnvConfig:
-    """Configuration from comfy-env.toml."""
-    python: Optional[str] = None
-    cuda_packages: List[str] = field(default_factory=list)
-    apt_packages: List[str] = field(default_factory=list)
-    env_vars: Dict[str, str] = field(default_factory=dict)
-    node_reqs: List[NodeReq] = field(default_factory=list)
-    pixi_passthrough: Dict[str, Any] = field(default_factory=dict)
+from .types import ComfyEnvConfig, NodeDependency
 
-    @property
-    def has_cuda(self) -> bool:
-        return bool(self.cuda_packages)
-# --- Types&Constants ---
+CONFIG_FILE_NAME = "comfy-env.toml"
 
 
 def load_config(path: Path) -> ComfyEnvConfig:
-    """Load config from a TOML file."""
+    """Load and parse comfy-env.toml."""
     path = Path(path)
     if not path.exists():
         raise FileNotFoundError(f"Config file not found: {path}")
     with open(path, "rb") as f:
-        data = tomli.load(f)
-    return _parse_config(data)
+        return parse_config(tomli.load(f))
 
 
 def discover_config(node_dir: Path) -> Optional[ComfyEnvConfig]:
-    """Find and load comfy-env.toml from a directory."""
+    """Find and load comfy-env.toml from directory."""
     config_path = Path(node_dir) / CONFIG_FILE_NAME
-    if config_path.exists():
-        return load_config(config_path)
+    return load_config(config_path) if config_path.exists() else None
 
-    return None
 
-
-def _parse_config(data: Dict[str, Any]) -> ComfyEnvConfig:
+def parse_config(data: Dict[str, Any]) -> ComfyEnvConfig:
     """Parse TOML data into ComfyEnvConfig."""
-    # Make a copy so we can pop our custom sections
     data = copy.deepcopy(data)
 
-    # Extract python version (top-level key)
     python_version = data.pop("python", None)
-    if python_version is not None:
-        python_version = str(python_version)
-
-    # Extract [cuda] section
-    cuda_data = data.pop("cuda", {})
-    cuda_packages = _ensure_list(cuda_data.get("packages", []))
-
-    # Extract [apt] section
-    apt_data = data.pop("apt", {})
-    apt_packages = _ensure_list(apt_data.get("packages", []))
-
-    # Extract [env_vars] section
-    env_vars_data = data.pop("env_vars", {})
-    env_vars = {str(k): str(v) for k, v in env_vars_data.items()}
-
-    # Extract [node_reqs] section
-    node_reqs_data = data.pop("node_reqs", {})
-    node_reqs = _parse_node_reqs(node_reqs_data)
+    python_version = str(python_version) if python_version else None
 
-    # Everything else passes through to pixi.toml
-    pixi_passthrough = data
+    cuda_packages = _ensure_list(data.pop("cuda", {}).get("packages", []))
+    apt_packages = _ensure_list(data.pop("apt", {}).get("packages", []))
+    env_vars = {str(k): str(v) for k, v in data.pop("env_vars", {}).items()}
+    node_reqs = _parse_node_reqs(data.pop("node_reqs", {}))
 
     return ComfyEnvConfig(
         python=python_version,
@@ -90,25 +44,17 @@ def _parse_config(data: Dict[str, Any]) -> ComfyEnvConfig:
         apt_packages=apt_packages,
         env_vars=env_vars,
         node_reqs=node_reqs,
-        pixi_passthrough=pixi_passthrough,
+        pixi_passthrough=data,
     )
 
 
-def _parse_node_reqs(data: Dict[str, Any]) -> List[NodeReq]:
+def _parse_node_reqs(data: Dict[str, Any]) -> List[NodeDependency]:
     """Parse [node_reqs] section."""
-    node_reqs = []
-    for name, value in data.items():
-        if isinstance(value, str):
-            node_reqs.append(NodeReq(name=name, repo=value))
-        elif isinstance(value, dict):
-            node_reqs.append(NodeReq(name=name, repo=value.get("repo", "")))
-    return node_reqs
+    return [
+        NodeDependency(name=name, repo=value if isinstance(value, str) else value.get("repo", ""))
+        for name, value in data.items()
+    ]
 
 
 def _ensure_list(value) -> List:
-    """Ensure value is a list."""
-    if isinstance(value, list):
-        return value
-    if value:
-        return [value]
-    return []
+    return value if isinstance(value, list) else ([value] if value else [])

comfy_env/config/types.py

@@ -0,0 +1,37 @@
+"""Configuration types for comfy-env."""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+
+@dataclass
+class NodeDependency:
+    """A ComfyUI custom node dependency."""
+    name: str
+    repo: str  # "owner/repo" or full URL
+
+
+NodeReq = NodeDependency  # Backwards compat
+
+
+@dataclass
+class ComfyEnvConfig:
+    """Parsed comfy-env.toml configuration."""
+    python: Optional[str] = None
+    cuda_packages: List[str] = field(default_factory=list)
+    apt_packages: List[str] = field(default_factory=list)
+    env_vars: Dict[str, str] = field(default_factory=dict)
+    node_reqs: List[NodeDependency] = field(default_factory=list)
+    pixi_passthrough: Dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def has_cuda(self) -> bool:
+        return bool(self.cuda_packages)
+
+    @property
+    def has_dependencies(self) -> bool:
+        return bool(
+            self.cuda_packages or self.apt_packages or self.node_reqs
+            or self.pixi_passthrough.get("dependencies")
+            or self.pixi_passthrough.get("pypi-dependencies")
+        )

comfy_env/detection/__init__.py

@@ -0,0 +1,77 @@
+"""
+Detection layer - Pure functions for system detection.
+
+No side effects. These functions gather information about the runtime environment.
+"""
+
+from .cuda import (
+    CUDA_VERSION_ENV_VAR,
+    detect_cuda_version,
+    get_cuda_from_torch,
+    get_cuda_from_nvml,
+    get_cuda_from_nvcc,
+    get_cuda_from_env,
+)
+from .gpu import (
+    GPUInfo,
+    CUDAEnvironment,
+    COMPUTE_TO_ARCH,
+    detect_gpu,
+    detect_gpus,
+    detect_cuda_environment,
+    get_compute_capability,
+    compute_capability_to_architecture,
+    get_recommended_cuda_version,
+    get_gpu_summary,
+)
+from .platform import (
+    PlatformInfo,
+    detect_platform,
+    get_platform_tag,
+    get_pixi_platform,
+    get_library_extension,
+    get_executable_suffix,
+    is_linux,
+    is_windows,
+    is_macos,
+)
+from .runtime import (
+    RuntimeEnv,
+    detect_runtime,
+    parse_wheel_requirement,
+)
+
+__all__ = [
+    # CUDA detection
+    "CUDA_VERSION_ENV_VAR",
+    "detect_cuda_version",
+    "get_cuda_from_torch",
+    "get_cuda_from_nvml",
+    "get_cuda_from_nvcc",
+    "get_cuda_from_env",
+    # GPU detection
+    "GPUInfo",
+    "CUDAEnvironment",
+    "COMPUTE_TO_ARCH",
+    "detect_gpu",
+    "detect_gpus",
+    "detect_cuda_environment",
+    "get_compute_capability",
+    "compute_capability_to_architecture",
+    "get_recommended_cuda_version",
+    "get_gpu_summary",
+    # Platform detection
+    "PlatformInfo",
+    "detect_platform",
+    "get_platform_tag",
+    "get_pixi_platform",
+    "get_library_extension",
+    "get_executable_suffix",
+    "is_linux",
+    "is_windows",
+    "is_macos",
+    # Runtime detection
+    "RuntimeEnv",
+    "detect_runtime",
+    "parse_wheel_requirement",
+]

comfy_env/detection/cuda.py

@@ -0,0 +1,61 @@
+"""CUDA version detection. Priority: env -> torch -> nvcc"""
+
+from __future__ import annotations
+
+import os
+import re
+import subprocess
+
+CUDA_VERSION_ENV_VAR = "COMFY_ENV_CUDA_VERSION"
+
+
+def detect_cuda_version() -> str | None:
+    """Detect CUDA version from available sources."""
+    return get_cuda_from_env() or get_cuda_from_torch() or get_cuda_from_nvcc()
+
+
+def get_cuda_from_env() -> str | None:
+    """Get CUDA version from environment variable override."""
+    override = os.environ.get(CUDA_VERSION_ENV_VAR, "").strip()
+    if not override:
+        return None
+    if "." not in override and len(override) >= 2:
+        return f"{override[:-1]}.{override[-1]}"
+    return override
+
+
+def get_cuda_from_torch() -> str | None:
+    """Get CUDA version from PyTorch."""
+    try:
+        import torch
+        if torch.cuda.is_available() and torch.version.cuda:
+            return torch.version.cuda
+    except Exception:
+        pass
+    return None
+
+
+def get_cuda_from_nvml() -> str | None:
+    """Get CUDA version from NVML."""
+    try:
+        import pynvml
+        pynvml.nvmlInit()
+        try:
+            cuda_version = pynvml.nvmlSystemGetCudaDriverVersion_v2()
+            return f"{cuda_version // 1000}.{(cuda_version % 1000) // 10}"
+        finally:
+            pynvml.nvmlShutdown()
+    except Exception:
+        pass
+    return None
+
+
+def get_cuda_from_nvcc() -> str | None:
+    """Get CUDA version from nvcc compiler."""
+    try:
+        r = subprocess.run(["nvcc", "--version"], capture_output=True, text=True, timeout=5)
+        if r.returncode == 0 and (m := re.search(r"release (\d+\.\d+)", r.stdout)):
+            return m.group(1)
+    except Exception:
+        pass
+    return None

comfy_env/detection/gpu.py

@@ -0,0 +1,230 @@
+"""GPU detection. Methods: NVML -> PyTorch -> nvidia-smi -> sysfs"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from .cuda import CUDA_VERSION_ENV_VAR
+
+COMPUTE_TO_ARCH = {
+    (5, 0): "Maxwell", (5, 2): "Maxwell", (5, 3): "Maxwell",
+    (6, 0): "Pascal", (6, 1): "Pascal", (6, 2): "Pascal",
+    (7, 0): "Volta", (7, 2): "Volta", (7, 5): "Turing",
+    (8, 0): "Ampere", (8, 6): "Ampere", (8, 7): "Ampere", (8, 9): "Ada",
+    (9, 0): "Hopper",
+    (10, 0): "Blackwell", (10, 1): "Blackwell", (10, 2): "Blackwell",
+}
+
+_cache: tuple[float, "CUDAEnvironment | None"] = (0, None)
+CACHE_TTL = 60
+
+
+@dataclass
+class GPUInfo:
+    index: int
+    name: str
+    compute_capability: tuple[int, int]
+    architecture: str
+    vram_total_mb: int = 0
+    vram_free_mb: int = 0
+    uuid: str = ""
+    pci_bus_id: str = ""
+    driver_version: str = ""
+
+    def cc_str(self) -> str: return f"{self.compute_capability[0]}.{self.compute_capability[1]}"
+    def sm_version(self) -> str: return f"sm_{self.compute_capability[0]}{self.compute_capability[1]}"
+
+
+@dataclass
+class CUDAEnvironment:
+    gpus: list[GPUInfo] = field(default_factory=list)
+    driver_version: str = ""
+    cuda_runtime_version: str = ""
+    recommended_cuda: str = ""
+    detection_method: str = ""
+
+
+def compute_capability_to_architecture(major: int, minor: int) -> str:
+    if arch := COMPUTE_TO_ARCH.get((major, minor)):
+        return arch
+    if major >= 10: return "Blackwell"
+    if major == 9: return "Hopper"
+    if major == 8: return "Ada" if minor >= 9 else "Ampere"
+    if major == 7: return "Turing" if minor >= 5 else "Volta"
+    if major == 6: return "Pascal"
+    return "Maxwell" if major == 5 else "Unknown"
+
+
+def get_compute_capability(gpu_index: int = 0) -> tuple[int, int] | None:
+    env = detect_cuda_environment()
+    return env.gpus[gpu_index].compute_capability if env.gpus and gpu_index < len(env.gpus) else None
+
+
+def get_recommended_cuda_version(gpus: list[GPUInfo] | None = None) -> str:
+    """Blackwell: 12.8, Pascal: 12.4, others: 12.8"""
+    if override := os.environ.get(CUDA_VERSION_ENV_VAR, "").strip():
+        return f"{override[:-1]}.{override[-1]}" if "." not in override and len(override) >= 2 else override
+
+    gpus = gpus if gpus is not None else detect_cuda_environment().gpus
+    if not gpus: return ""
+
+    for gpu in gpus:
+        if gpu.compute_capability[0] >= 10: return "12.8"
+    for gpu in gpus:
+        cc = gpu.compute_capability
+        if cc[0] < 7 or (cc[0] == 7 and cc[1] < 5): return "12.4"
+    return "12.8"
+
+
+def detect_gpu() -> GPUInfo | None:
+    env = detect_cuda_environment()
+    return env.gpus[0] if env.gpus else None
+
+
+def detect_gpus() -> list[GPUInfo]:
+    return detect_cuda_environment().gpus
+
+
+def detect_cuda_environment(force_refresh: bool = False) -> CUDAEnvironment:
+    global _cache
+    if not force_refresh and _cache[1] and time.time() - _cache[0] < CACHE_TTL:
+        return _cache[1]
+
+    gpus, method = None, "none"
+    for name, fn in [("nvml", _detect_nvml), ("torch", _detect_torch), ("smi", _detect_smi), ("sysfs", _detect_sysfs)]:
+        if result := fn():
+            gpus, method = result, name
+            break
+
+    env = CUDAEnvironment(
+        gpus=gpus or [], driver_version=_get_driver_version(),
+        cuda_runtime_version=_get_cuda_version(),
+        recommended_cuda=get_recommended_cuda_version(gpus or []), detection_method=method,
+    )
+    _cache = (time.time(), env)
+    return env
+
+
+def get_gpu_summary() -> str:
+    env = detect_cuda_environment()
+    if not env.gpus:
+        override = os.environ.get(CUDA_VERSION_ENV_VAR)
+        return f"No GPU detected (using {CUDA_VERSION_ENV_VAR}={override})" if override else f"No GPU (set {CUDA_VERSION_ENV_VAR} to override)"
+
+    lines = [f"Detection: {env.detection_method}"]
+    if env.driver_version: lines.append(f"Driver: {env.driver_version}")
+    if env.cuda_runtime_version: lines.append(f"CUDA: {env.cuda_runtime_version}")
+    lines.append(f"Recommended: CUDA {env.recommended_cuda}")
+    lines.append("")
+    for gpu in env.gpus:
+        vram = f"{gpu.vram_total_mb}MB" if gpu.vram_total_mb else "?"
+        lines.append(f"  GPU {gpu.index}: {gpu.name} ({gpu.sm_version()}) [{gpu.architecture}] {vram}")
+    return "\n".join(lines)
+
+
+def _parse_cc(s: str) -> tuple[int, int]:
+    try:
+        if "." in s: p = s.split("."); return (int(p[0]), int(p[1]))
+        if len(s) >= 2: return (int(s[:-1]), int(s[-1]))
+    except (ValueError, IndexError): pass
+    return (0, 0)
+
+
+def _detect_nvml() -> list[GPUInfo] | None:
+    try:
+        import pynvml
+        pynvml.nvmlInit()
+        try:
+            count = pynvml.nvmlDeviceGetCount()
+            if not count: return None
+            gpus = []
+            for i in range(count):
+                h = pynvml.nvmlDeviceGetHandleByIndex(i)
+                name = pynvml.nvmlDeviceGetName(h)
+                if isinstance(name, bytes): name = name.decode()
+                cc = pynvml.nvmlDeviceGetCudaComputeCapability(h)
+                mem = pynvml.nvmlDeviceGetMemoryInfo(h)
+                gpus.append(GPUInfo(i, name, cc, compute_capability_to_architecture(*cc),
+                                   mem.total // (1024 * 1024), mem.free // (1024 * 1024)))
+            return gpus
+        finally:
+            pynvml.nvmlShutdown()
+    except Exception: return None
+
+
+def _detect_torch() -> list[GPUInfo] | None:
+    try:
+        import torch
+        if not torch.cuda.is_available(): return None
+        gpus = []
+        for i in range(torch.cuda.device_count()):
+            p = torch.cuda.get_device_properties(i)
+            gpus.append(GPUInfo(i, p.name, (p.major, p.minor),
+                               compute_capability_to_architecture(p.major, p.minor),
+                               p.total_memory // (1024 * 1024)))
+        return gpus or None
+    except Exception: return None
+
+
+def _detect_smi() -> list[GPUInfo] | None:
+    try:
+        r = subprocess.run(
+            ["nvidia-smi", "--query-gpu=index,name,uuid,pci.bus_id,compute_cap,memory.total,memory.free,driver_version",
+             "--format=csv,noheader,nounits"], capture_output=True, text=True, timeout=10)
+        if r.returncode != 0: return None
+        gpus = []
+        for line in r.stdout.strip().split("\n"):
+            if not line.strip(): continue
+            p = [x.strip() for x in line.split(",")]
+            if len(p) < 5: continue
+            cc = _parse_cc(p[4])
+            gpus.append(GPUInfo(
+                int(p[0]) if p[0].isdigit() else len(gpus), p[1], cc, compute_capability_to_architecture(*cc),
+                int(p[5]) if len(p) > 5 and p[5].isdigit() else 0,
+                int(p[6]) if len(p) > 6 and p[6].isdigit() else 0,
+                p[2] if len(p) > 2 else "", p[3] if len(p) > 3 else "", p[7] if len(p) > 7 else ""))
+        return gpus or None
+    except Exception: return None
+
+
+def _detect_sysfs() -> list[GPUInfo] | None:
+    try:
+        pci_path = Path("/sys/bus/pci/devices")
+        if not pci_path.exists(): return None
+        gpus = []
+        for d in sorted(pci_path.iterdir()):
+            vendor = (d / "vendor").read_text().strip().lower() if (d / "vendor").exists() else ""
+            if "10de" not in vendor: continue
+            cls = (d / "class").read_text().strip() if (d / "class").exists() else ""
+            if not (cls.startswith("0x0300") or cls.startswith("0x0302")): continue
+            gpus.append(GPUInfo(len(gpus), "NVIDIA GPU", (0, 0), "Unknown", pci_bus_id=d.name))
+        return gpus or None
+    except Exception: return None
+
+
+def _get_driver_version() -> str:
+    try:
+        import pynvml
+        pynvml.nvmlInit()
+        v = pynvml.nvmlSystemGetDriverVersion()
+        pynvml.nvmlShutdown()
+        return v.decode() if isinstance(v, bytes) else v
+    except Exception: pass
+    try:
+        r = subprocess.run(["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
+                          capture_output=True, text=True, timeout=5)
+        if r.returncode == 0: return r.stdout.strip().split("\n")[0]
+    except Exception: pass
+    return ""
+
+
+def _get_cuda_version() -> str:
+    try:
+        import torch
+        if torch.cuda.is_available() and torch.version.cuda: return torch.version.cuda
+    except Exception: pass
+    return ""

comfy_env/detection/platform.py

@@ -0,0 +1,70 @@
+"""Platform detection - OS, architecture, platform tags."""
+
+from __future__ import annotations
+
+import platform as platform_module
+import sys
+from dataclasses import dataclass
+
+
+@dataclass
+class PlatformInfo:
+    os_name: str       # linux, windows, darwin
+    arch: str          # x86_64, aarch64, arm64
+    platform_tag: str  # linux_x86_64, win_amd64, macosx_11_0_arm64
+
+
+def detect_platform() -> PlatformInfo:
+    return PlatformInfo(
+        os_name=_get_os_name(),
+        arch=platform_module.machine().lower(),
+        platform_tag=get_platform_tag(),
+    )
+
+
+def _get_os_name() -> str:
+    if sys.platform.startswith('linux'): return 'linux'
+    if sys.platform == 'win32': return 'windows'
+    if sys.platform == 'darwin': return 'darwin'
+    return sys.platform
+
+
+_PLATFORM_TAGS = {
+    ('linux', 'x86_64'): 'linux_x86_64', ('linux', 'amd64'): 'linux_x86_64',
+    ('linux', 'aarch64'): 'linux_aarch64',
+    ('win32', 'amd64'): 'win_amd64', ('win32', 'x86_64'): 'win_amd64',
+    ('darwin', 'arm64'): 'macosx_11_0_arm64',
+    ('darwin', 'x86_64'): 'macosx_10_9_x86_64',
+}
+
+
+def get_platform_tag() -> str:
+    key = (sys.platform if sys.platform != 'linux' else 'linux', platform_module.machine().lower())
+    return _PLATFORM_TAGS.get(key, f'{sys.platform}_{platform_module.machine().lower()}')
+
+
+_PIXI_PLATFORMS = {
+    ('linux', 'x86_64'): 'linux-64', ('linux', 'amd64'): 'linux-64',
+    ('linux', 'aarch64'): 'linux-aarch64',
+    ('windows', 'amd64'): 'win-64', ('windows', 'x86_64'): 'win-64',
+    ('darwin', 'arm64'): 'osx-arm64',
+    ('darwin', 'x86_64'): 'osx-64',
+}
+
+
+def get_pixi_platform() -> str:
+    key = (_get_os_name(), platform_module.machine().lower())
+    return _PIXI_PLATFORMS.get(key, f'{key[0]}-{key[1]}')
+
+
+def get_library_extension() -> str:
+    return {'.dll': 'windows', '.dylib': 'darwin'}.get(_get_os_name(), '.so')
+
+
+def get_executable_suffix() -> str:
+    return '.exe' if _get_os_name() == 'windows' else ''
+
+
+def is_linux() -> bool: return _get_os_name() == 'linux'
+def is_windows() -> bool: return _get_os_name() == 'windows'
+def is_macos() -> bool: return _get_os_name() == 'darwin'