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

comfy_env/config/parser.py

@@ -1,135 +1,42 @@
-"""Load configuration from comfy-env.toml.
-
-comfy-env.toml is a superset of pixi.toml. Custom sections we handle:
-- python = "3.11" - Python version for isolated envs
-- [cuda] packages = [...] - CUDA packages (triggers find-links + PyTorch detection)
-- [node_reqs] - Other ComfyUI nodes to clone
-
-Everything else passes through to pixi.toml directly.
-
-Example config:
-
-    python = "3.11"
-
-    [cuda]
-    packages = ["cumesh"]
-
-    [dependencies]
-    mesalib = "*"
-    cgal = "*"
-
-    [pypi-dependencies]
-    numpy = ">=1.21.0,<2"
-    trimesh = { version = ">=4.0.0", extras = ["easy"] }
-
-    [target.linux-64.pypi-dependencies]
-    embreex = "*"
-
-    [node_reqs]
-    SomeNode = "owner/repo"
-"""
+"""Configuration parsing for comfy-env."""
 
 import copy
-import sys
 from pathlib import Path
-from typing import Optional, Dict, Any, List
-
-# Use built-in tomllib (Python 3.11+) or tomli fallback
-if sys.version_info >= (3, 11):
-    import tomllib
-else:
-    try:
-        import tomli as tomllib
-    except ImportError:
-        tomllib = None  # type: ignore
+from typing import Any, Dict, List, Optional
 
-from .types import ComfyEnvConfig, NodeReq
+import tomli
 
+from .types import ComfyEnvConfig, NodeDependency
 
 CONFIG_FILE_NAME = "comfy-env.toml"
 
-# Sections we handle specially (not passed through to pixi.toml)
-CUSTOM_SECTIONS = {"python", "cuda", "node_reqs", "apt", "env_vars"}
-
 
 def load_config(path: Path) -> ComfyEnvConfig:
-    """
-    Load configuration from a TOML file.
-
-    Args:
-        path: Path to comfy-env.toml
-
-    Returns:
-        ComfyEnvConfig instance
-
-    Raises:
-        FileNotFoundError: If config file doesn't exist
-        ImportError: If tomli not installed (Python < 3.11)
-    """
-    if tomllib is None:
-        raise ImportError(
-            "TOML parsing requires tomli for Python < 3.11. "
-            "Install with: pip install tomli"
-        )
-
+    """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 = tomllib.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.
-
-    Args:
-        node_dir: Directory to search
-
-    Returns:
-        ComfyEnvConfig if found, None otherwise
-    """
-    if tomllib is None:
-        return None
-
+    """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 None
+    return load_config(config_path) if config_path.exists() else 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,
@@ -137,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

@@ -1,70 +1,37 @@
-"""Configuration types for comfy-env."""
-
-from dataclasses import dataclass, field
-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"
-
-
-@dataclass
-class ComfyEnvConfig:
-    """
-    Configuration from comfy-env.toml.
-
-    comfy-env.toml is a superset of pixi.toml. Custom sections we handle:
-    - python = "3.11" - Python version for isolated envs
-    - [cuda] packages = [...] - CUDA packages (triggers find-links + PyTorch detection)
-    - [node_reqs] - Other ComfyUI nodes to clone
-
-    Everything else passes through to pixi.toml directly:
-    - [dependencies] - conda packages
-    - [pypi-dependencies] - pip packages
-    - [target.linux-64.pypi-dependencies] - platform-specific deps
-    - Any other pixi.toml syntax
-
-    Example config:
-        python = "3.11"
-
-        [cuda]
-        packages = ["cumesh"]
-
-        [dependencies]
-        mesalib = "*"
-        cgal = "*"
-
-        [pypi-dependencies]
-        numpy = ">=1.21.0,<2"
-        trimesh = { version = ">=4.0.0", extras = ["easy"] }
-
-        [target.linux-64.pypi-dependencies]
-        embreex = "*"
-
-        [node_reqs]
-        SomeNode = "owner/repo"
-    """
-    # python = "3.11" - Python version (for isolated envs)
-    python: Optional[str] = None
-
-    # [cuda] - CUDA packages (installed via find-links index)
-    cuda_packages: List[str] = field(default_factory=list)
-
-    # [apt] - System packages to install via apt (Linux only)
-    apt_packages: List[str] = field(default_factory=list)
-
-    # [env_vars] - Environment variables to set early (in prestartup)
-    env_vars: Dict[str, str] = field(default_factory=dict)
-
-    # [node_reqs] - other ComfyUI nodes to clone
-    node_reqs: List[NodeReq] = field(default_factory=list)
-
-    # Everything else from comfy-env.toml passes through to pixi.toml
-    pixi_passthrough: Dict[str, Any] = field(default_factory=dict)
-
-    @property
-    def has_cuda(self) -> bool:
-        return bool(self.cuda_packages)
+"""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 ""