comfy-env 0.0.8__py3-none-any.whl

comfy_env/registry.py

@@ -0,0 +1,252 @@
+"""Built-in registry of CUDA packages and their wheel sources.
+
+This module provides a mapping of well-known CUDA packages to their
+installation sources, eliminating the need for users to specify
+wheel_sources in their comfyui_env.toml.
+
+Install method types:
+- "index": Use pip --extra-index-url (PEP 503 simple repository)
+- "github_index": GitHub Pages index (--find-links)
+- "find_links": Use pip --find-links (for PyG, etc.)
+- "pypi_variant": Package name varies by CUDA version (e.g., spconv-cu124)
+- "github_release": Direct wheel URL from GitHub releases with fallback sources
+"""
+
+from typing import Dict, Any, Optional
+
+
+def get_cuda_short2(cuda_version: str) -> str:
+    """Convert CUDA version to 2-3 digit format for spconv.
+
+    spconv uses "cu124" not "cu1240" for CUDA 12.4.
+
+    Args:
+        cuda_version: CUDA version string (e.g., "12.4", "12.8")
+
+    Returns:
+        Short format string (e.g., "124", "128")
+
+    Examples:
+        >>> get_cuda_short2("12.4")
+        '124'
+        >>> get_cuda_short2("12.8")
+        '128'
+        >>> get_cuda_short2("11.8")
+        '118'
+    """
+    parts = cuda_version.split(".")
+    major = parts[0]
+    minor = parts[1] if len(parts) > 1 else "0"
+    return f"{major}{minor}"
+
+
+# =============================================================================
+# Package Registry
+# =============================================================================
+# Maps package names to their installation configuration.
+#
+# Template variables available:
+#   {cuda_version}  - Full CUDA version (e.g., "12.8")
+#   {cuda_short}    - CUDA without dot (e.g., "128")
+#   {cuda_short2}   - CUDA short for spconv (e.g., "124" not "1240")
+#   {torch_version} - Full PyTorch version (e.g., "2.8.0")
+#   {torch_short}   - PyTorch without dots (e.g., "280")
+#   {torch_mm}      - PyTorch major.minor (e.g., "28")
+#   {py_version}    - Python version (e.g., "3.10")
+#   {py_short}      - Python without dot (e.g., "310")
+#   {py_minor}      - Python minor version only (e.g., "10")
+#   {platform}      - Platform tag (e.g., "linux_x86_64")
+# =============================================================================
+
+PACKAGE_REGISTRY: Dict[str, Dict[str, Any]] = {
+    # =========================================================================
+    # PyTorch Geometric (PyG) packages - official index
+    # https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html
+    # Uses --find-links (not --extra-index-url) for proper wheel discovery
+    # =========================================================================
+    "torch-scatter": {
+        "method": "find_links",
+        "index_url": "https://data.pyg.org/whl/torch-{torch_version}+cu{cuda_short}.html",
+        "description": "Scatter operations for PyTorch",
+    },
+    "torch-cluster": {
+        "method": "find_links",
+        "index_url": "https://data.pyg.org/whl/torch-{torch_version}+cu{cuda_short}.html",
+        "description": "Clustering algorithms for PyTorch",
+    },
+    "torch-sparse": {
+        "method": "find_links",
+        "index_url": "https://data.pyg.org/whl/torch-{torch_version}+cu{cuda_short}.html",
+        "description": "Sparse tensor operations for PyTorch",
+    },
+    "torch-spline-conv": {
+        "method": "find_links",
+        "index_url": "https://data.pyg.org/whl/torch-{torch_version}+cu{cuda_short}.html",
+        "description": "Spline convolutions for PyTorch",
+    },
+
+    # =========================================================================
+    # pytorch3d - Facebook's official wheels
+    # https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md
+    # =========================================================================
+    "pytorch3d": {
+        "method": "index",
+        "index_url": "https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py3{py_minor}_cu{cuda_short}_pyt{torch_short}/download.html",
+        "description": "PyTorch3D - 3D deep learning library",
+    },
+
+    # =========================================================================
+    # PozzettiAndrea wheel repos (GitHub Pages indexes)
+    # =========================================================================
+    # nvdiffrast - wheels are now at cu{cuda}-torch{torch_short} releases
+    "nvdiffrast": {
+        "method": "github_index",
+        "index_url": "https://pozzettiandrea.github.io/nvdiffrast-full-wheels/cu{cuda_short}-torch{torch_short}/",
+        "description": "NVIDIA differentiable rasterizer",
+    },
+    # cumesh, o_voxel, flex_gemm, nvdiffrec_render use torch_short (3 digits: 280)
+    "cumesh": {
+        "method": "github_index",
+        "index_url": "https://pozzettiandrea.github.io/cumesh-wheels/cu{cuda_short}-torch{torch_short}/",
+        "description": "CUDA-accelerated mesh utilities",
+    },
+    "o_voxel": {
+        "method": "github_index",
+        "index_url": "https://pozzettiandrea.github.io/ovoxel-wheels/cu{cuda_short}-torch{torch_short}/",
+        "description": "O-Voxel CUDA extension for TRELLIS",
+    },
+    "flex_gemm": {
+        "method": "github_index",
+        "index_url": "https://pozzettiandrea.github.io/flexgemm-wheels/cu{cuda_short}-torch{torch_short}/",
+        "description": "Flexible GEMM operations",
+    },
+    "nvdiffrec_render": {
+        "method": "github_release",
+        "sources": [
+            {
+                "name": "PozzettiAndrea",
+                "url_template": "https://github.com/PozzettiAndrea/nvdiffrec_render-wheels/releases/download/cu{cuda_short}-torch{torch_short}/nvdiffrec_render-{version}%2Bcu{cuda_short}torch{torch_mm}-{py_tag}-{py_tag}-linux_x86_64.whl",
+                "platforms": ["linux_x86_64"],
+            },
+            {
+                "name": "PozzettiAndrea-windows",
+                "url_template": "https://github.com/PozzettiAndrea/nvdiffrec_render-wheels/releases/download/cu{cuda_short}-torch{torch_short}/nvdiffrec_render-{version}%2Bcu{cuda_short}torch{torch_mm}-{py_tag}-{py_tag}-win_amd64.whl",
+                "platforms": ["win_amd64", "windows_amd64"],
+            },
+        ],
+        "description": "NVDiffRec rendering utilities",
+    },
+
+    # =========================================================================
+    # spconv - PyPI with CUDA-versioned package names
+    # Package names: spconv-cu118, spconv-cu120, spconv-cu121, spconv-cu124, spconv-cu126
+    # Note: Max available is cu126 as of Jan 2026, use explicit spconv-cu126 in config
+    # =========================================================================
+    "spconv": {
+        "method": "pypi_variant",
+        "package_template": "spconv-cu{cuda_short2}",
+        "description": "Sparse convolution library (use spconv-cu126 for CUDA 12.6+)",
+    },
+
+    # =========================================================================
+    # sageattention - Fast quantized attention (2-5x faster than FlashAttention)
+    # Linux: Prebuilt wheels from Kijai/PrecompiledWheels (v2.2.0, cp312)
+    # Windows: Prebuilt wheels from woct0rdho (v2.2.0, cp39-abi3)
+    # =========================================================================
+    "sageattention": {
+        "method": "github_release",
+        "sources": [
+            # Linux: Kijai's precompiled wheels on HuggingFace (Python 3.12)
+            {
+                "name": "kijai-hf",
+                "url_template": "https://huggingface.co/Kijai/PrecompiledWheels/resolve/main/sageattention-{version}-cp312-cp312-linux_x86_64.whl",
+                "platforms": ["linux_x86_64"],
+            },
+            # Windows: woct0rdho prebuilt wheels (ABI3: Python >= 3.9)
+            # Format: sageattention-2.2.0+cu128torch2.8.0.post3-cp39-abi3-win_amd64.whl
+            {
+                "name": "woct0rdho",
+                "url_template": "https://github.com/woct0rdho/SageAttention/releases/download/v2.2.0-windows.post3/sageattention-2.2.0%2Bcu{cuda_short}torch{torch_version}.post3-cp39-abi3-win_amd64.whl",
+                "platforms": ["win_amd64"],
+            },
+        ],
+        "description": "SageAttention - 2-5x faster than FlashAttention with quantized kernels",
+    },
+
+    # =========================================================================
+    # triton - Required for sageattention on Linux (usually bundled with PyTorch)
+    # =========================================================================
+    "triton": {
+        "method": "pypi",
+        "description": "Triton compiler for custom CUDA kernels (required by sageattention)",
+    },
+
+    # =========================================================================
+    # flash-attn - Multi-source prebuilt wheels
+    # Required for UniRig and other transformer-based models
+    # Sources: Dao-AILab (official), mjun0812 (Linux), bdashore3 (Windows)
+    # =========================================================================
+    "flash-attn": {
+        "method": "github_release",
+        "sources": [
+            # Linux: Dao-AILab official wheels (CUDA 12.x, PyTorch 2.4-2.8)
+            # Format: flash_attn-2.8.3+cu12torch2.8cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
+            {
+                "name": "Dao-AILab",
+                "url_template": "https://github.com/Dao-AILab/flash-attention/releases/download/v{version}/flash_attn-{version}%2Bcu{cuda_major}torch{torch_dotted_mm}cxx11abiTRUE-{py_tag}-{py_tag}-linux_x86_64.whl",
+                "platforms": ["linux_x86_64"],
+            },
+            # Linux: mjun0812 prebuilt wheels (CUDA 12.4-13.0, PyTorch 2.5-2.9)
+            # Format: flash_attn-2.8.3+cu128torch2.8-cp310-cp310-linux_x86_64.whl
+            # Note: Release v0.7.2 contains multiple flash_attn versions
+            {
+                "name": "mjun0812",
+                "url_template": "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.2/flash_attn-{version}%2Bcu{cuda_short}torch{torch_dotted_mm}-{py_tag}-{py_tag}-linux_x86_64.whl",
+                "platforms": ["linux_x86_64"],
+            },
+            # Windows: bdashore3 prebuilt wheels (CUDA 12.4/12.8, PyTorch 2.6-2.8)
+            {
+                "name": "bdashore3",
+                "url_template": "https://github.com/bdashore3/flash-attention/releases/download/v{version}/flash_attn-{version}%2Bcu{cuda_short}torch{torch_version}cxx11abiFALSE-{py_tag}-{py_tag}-win_amd64.whl",
+                "platforms": ["win_amd64"],
+            },
+        ],
+        "description": "Flash Attention for fast transformer inference",
+    },
+}
+
+
+def get_package_info(package: str) -> Optional[Dict[str, Any]]:
+    """Get registry info for a package.
+
+    Args:
+        package: Package name (case-insensitive)
+
+    Returns:
+        Registry entry dict or None if not found
+    """
+    return PACKAGE_REGISTRY.get(package.lower())
+
+
+def list_packages() -> Dict[str, str]:
+    """List all registered packages with their descriptions.
+
+    Returns:
+        Dict mapping package name to description
+    """
+    return {
+        name: info.get("description", "No description")
+        for name, info in PACKAGE_REGISTRY.items()
+    }
+
+
+def is_registered(package: str) -> bool:
+    """Check if a package is in the registry.
+
+    Args:
+        package: Package name (case-insensitive)
+
+    Returns:
+        True if package is registered
+    """
+    return package.lower() in PACKAGE_REGISTRY

comfy_env/resolver.py

@@ -0,0 +1,399 @@
+"""
+Wheel URL resolver for CUDA-compiled packages.
+
+This module provides deterministic wheel URL construction based on the runtime
+environment (CUDA version, PyTorch version, Python version, platform).
+
+Unlike pip's constraint solver, this module constructs exact URLs from templates
+and validates that they exist. If a wheel doesn't exist, it fails fast with
+a clear error message.
+
+Example:
+    from comfy_env.resolver import WheelResolver, RuntimeEnv
+
+    env = RuntimeEnv.detect()
+    resolver = WheelResolver()
+
+    url = resolver.resolve("nvdiffrast", version="0.4.0", env=env)
+    # Returns: https://github.com/.../nvdiffrast-0.4.0+cu128torch28-cp310-...-linux_x86_64.whl
+"""
+
+import platform
+import re
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+from urllib.parse import urlparse
+
+from .env.detection import detect_cuda_version, detect_gpu_info
+
+
+@dataclass
+class RuntimeEnv:
+    """
+    Detected runtime environment for wheel resolution.
+
+    Contains all variables needed for wheel URL template expansion.
+    """
+    # OS/Platform
+    os_name: str  # linux, windows, darwin
+    platform_tag: str  # linux_x86_64, win_amd64, macosx_...
+
+    # Python
+    python_version: str  # 3.10, 3.11, 3.12
+    python_short: str  # 310, 311, 312
+
+    # CUDA
+    cuda_version: Optional[str]  # 12.8, 12.4, None
+    cuda_short: Optional[str]  # 128, 124, None
+
+    # PyTorch (detected or configured)
+    torch_version: Optional[str]  # 2.8.0, 2.5.1
+    torch_short: Optional[str]  # 280, 251
+    torch_mm: Optional[str]  # 28, 25 (major.minor without dot)
+
+    # GPU info
+    gpu_name: Optional[str] = None
+    gpu_compute: Optional[str] = None  # sm_89, sm_100
+
+    @classmethod
+    def detect(cls, torch_version: Optional[str] = None) -> "RuntimeEnv":
+        """
+        Detect runtime environment from current system.
+
+        Args:
+            torch_version: Optional PyTorch version override. If not provided,
+                          attempts to detect from installed torch.
+
+        Returns:
+            RuntimeEnv with detected values.
+        """
+        # OS detection
+        os_name = sys.platform
+        if os_name.startswith('linux'):
+            os_name = 'linux'
+        elif os_name == 'win32':
+            os_name = 'windows'
+        elif os_name == 'darwin':
+            os_name = 'darwin'
+
+        # Platform tag
+        platform_tag = _get_platform_tag()
+
+        # Python version
+        py_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+        py_short = f"{sys.version_info.major}{sys.version_info.minor}"
+
+        # CUDA version
+        cuda_version = detect_cuda_version()
+        cuda_short = cuda_version.replace(".", "") if cuda_version else None
+
+        # PyTorch version
+        if torch_version is None:
+            torch_version = _detect_torch_version()
+
+        torch_short = None
+        torch_mm = None
+        if torch_version:
+            torch_short = torch_version.replace(".", "")
+            parts = torch_version.split(".")[:2]
+            torch_mm = "".join(parts)
+
+        # GPU info
+        gpu_name = None
+        gpu_compute = None
+        try:
+            gpu_info = detect_gpu_info()
+            if gpu_info:
+                gpu_name = gpu_info.get("name")
+                gpu_compute = gpu_info.get("compute_capability")
+        except Exception:
+            pass
+
+        return cls(
+            os_name=os_name,
+            platform_tag=platform_tag,
+            python_version=py_version,
+            python_short=py_short,
+            cuda_version=cuda_version,
+            cuda_short=cuda_short,
+            torch_version=torch_version,
+            torch_short=torch_short,
+            torch_mm=torch_mm,
+            gpu_name=gpu_name,
+            gpu_compute=gpu_compute,
+        )
+
+    def as_dict(self) -> Dict[str, str]:
+        """Convert to dict for template substitution."""
+        # Extract py_minor from python_version (e.g., "3.10" -> "10")
+        py_minor = self.python_version.split(".")[-1] if self.python_version else ""
+
+        result = {
+            "os": self.os_name,
+            "platform": self.platform_tag,
+            "python_version": self.python_version,
+            "py_version": self.python_version,
+            "py_short": self.python_short,
+            "py_minor": py_minor,
+        }
+
+        if self.cuda_version:
+            result["cuda_version"] = self.cuda_version
+            result["cuda_short"] = self.cuda_short
+
+        if self.torch_version:
+            result["torch_version"] = self.torch_version
+            result["torch_short"] = self.torch_short
+            result["torch_mm"] = self.torch_mm
+            # torch_dotted_mm: "2.8" format (major.minor with dot) for flash-attn URLs
+            parts = self.torch_version.split(".")[:2]
+            result["torch_dotted_mm"] = ".".join(parts)
+
+        return result
+
+    def __str__(self) -> str:
+        parts = [
+            f"Python {self.python_version}",
+            f"CUDA {self.cuda_version}" if self.cuda_version else "CPU",
+        ]
+        if self.torch_version:
+            parts.append(f"PyTorch {self.torch_version}")
+        if self.gpu_name:
+            parts.append(f"GPU: {self.gpu_name}")
+        return ", ".join(parts)
+
+
+def _get_platform_tag() -> str:
+    """Get wheel platform tag for current system."""
+    machine = platform.machine().lower()
+
+    if sys.platform.startswith('linux'):
+        # Use manylinux tag
+        if machine in ('x86_64', 'amd64'):
+            return 'linux_x86_64'
+        elif machine == 'aarch64':
+            return 'linux_aarch64'
+        return f'linux_{machine}'
+
+    elif sys.platform == 'win32':
+        if machine in ('amd64', 'x86_64'):
+            return 'win_amd64'
+        return 'win32'
+
+    elif sys.platform == 'darwin':
+        # macOS - use generic tag
+        if machine == 'arm64':
+            return 'macosx_11_0_arm64'
+        return 'macosx_10_9_x86_64'
+
+    return f'{sys.platform}_{machine}'
+
+
+def _detect_torch_version() -> Optional[str]:
+    """Detect installed PyTorch version."""
+    try:
+        import torch
+        version = torch.__version__
+        # Strip CUDA suffix (e.g., "2.8.0+cu128" -> "2.8.0")
+        if '+' in version:
+            version = version.split('+')[0]
+        return version
+    except ImportError:
+        return None
+
+
+@dataclass
+class WheelSource:
+    """Configuration for a wheel source (GitHub releases, custom index, etc.)."""
+    name: str
+    url_template: str
+    packages: List[str] = field(default_factory=list)  # Empty = all packages
+
+    def supports(self, package: str) -> bool:
+        """Check if this source provides the given package."""
+        if not self.packages:
+            return True  # Empty list = supports all
+        return package.lower() in [p.lower() for p in self.packages]
+
+
+# Default wheel sources for common CUDA packages
+DEFAULT_WHEEL_SOURCES = [
+    WheelSource(
+        name="nvdiffrast-wheels",
+        url_template="https://github.com/PozzettiAndrea/nvdiffrast-full-wheels/releases/download/v{version}/nvdiffrast-{version}%2Bcu{cuda_short}torch{torch_mm}-cp{py_short}-cp{py_short}-{platform}.whl",
+        packages=["nvdiffrast"],
+    ),
+    WheelSource(
+        name="cumesh-wheels",
+        url_template="https://github.com/PozzettiAndrea/cumesh-wheels/releases/download/v{version}/{package}-{version}%2Bcu{cuda_short}torch{torch_mm}-cp{py_short}-cp{py_short}-{platform}.whl",
+        packages=["pytorch3d", "torch-cluster", "torch-scatter", "torch-sparse"],
+    ),
+]
+
+
+class WheelResolver:
+    """
+    Resolves CUDA wheel URLs from package name and runtime environment.
+
+    Resolution strategy:
+    1. Check explicit overrides in config
+    2. Try configured wheel sources in order
+    3. Fail with actionable error message
+
+    This is NOT a constraint solver. It constructs deterministic URLs
+    based on exact version matches.
+    """
+
+    def __init__(
+        self,
+        sources: Optional[List[WheelSource]] = None,
+        overrides: Optional[Dict[str, str]] = None,
+    ):
+        """
+        Initialize resolver.
+
+        Args:
+            sources: List of WheelSource configurations.
+            overrides: Package-specific URL overrides (package -> template).
+        """
+        self.sources = sources or DEFAULT_WHEEL_SOURCES
+        self.overrides = overrides or {}
+
+    def resolve(
+        self,
+        package: str,
+        version: str,
+        env: RuntimeEnv,
+        verify: bool = False,
+    ) -> str:
+        """
+        Resolve wheel URL for a package.
+
+        Args:
+            package: Package name (e.g., "nvdiffrast").
+            version: Package version (e.g., "0.4.0").
+            env: Runtime environment for template expansion.
+
+        Returns:
+            Fully resolved wheel URL.
+
+        Raises:
+            WheelNotFoundError: If no wheel URL could be constructed or verified.
+        """
+        from .errors import WheelNotFoundError
+
+        # Prepare template variables
+        variables = env.as_dict()
+        variables["package"] = package
+        variables["version"] = version
+
+        # 1. Check explicit override
+        if package.lower() in self.overrides:
+            url = self._substitute(self.overrides[package.lower()], variables)
+            if verify and not self._url_exists(url):
+                raise WheelNotFoundError(
+                    package=package,
+                    version=version,
+                    env=env,
+                    tried_urls=[url],
+                    reason="Override URL returned 404",
+                )
+            return url
+
+        # 2. Try wheel sources
+        tried_urls = []
+        for source in self.sources:
+            if not source.supports(package):
+                continue
+
+            url = self._substitute(source.url_template, variables)
+            tried_urls.append(url)
+
+            if verify:
+                if self._url_exists(url):
+                    return url
+            else:
+                return url
+
+        # 3. Fail with helpful error
+        raise WheelNotFoundError(
+            package=package,
+            version=version,
+            env=env,
+            tried_urls=tried_urls,
+            reason="No wheel source found for package",
+        )
+
+    def resolve_all(
+        self,
+        packages: Dict[str, str],
+        env: RuntimeEnv,
+        verify: bool = False,
+    ) -> Dict[str, str]:
+        """
+        Resolve URLs for multiple packages.
+
+        Args:
+            packages: Dict of package -> version.
+            env: Runtime environment.
+            verify: Whether to verify URLs exist.
+
+        Returns:
+            Dict of package -> resolved URL.
+
+        Raises:
+            WheelNotFoundError: If any package cannot be resolved.
+        """
+        results = {}
+        for package, version in packages.items():
+            results[package] = self.resolve(package, version, env, verify=verify)
+        return results
+
+    def _substitute(self, template: str, variables: Dict[str, str]) -> str:
+        """
+        Substitute variables into URL template.
+
+        Handles both {var} and {var_name} style placeholders.
+        Missing variables are left as-is (caller should validate).
+        """
+        result = template
+        for key, value in variables.items():
+            result = result.replace(f"{{{key}}}", str(value))
+        return result
+
+    def _url_exists(self, url: str, timeout: float = 10.0) -> bool:
+        """
+        Check if a URL exists using HTTP HEAD request.
+
+        Returns True if URL returns 200 OK.
+        """
+        try:
+            import urllib.request
+            request = urllib.request.Request(url, method='HEAD')
+            with urllib.request.urlopen(request, timeout=timeout) as response:
+                return response.status == 200
+        except Exception:
+            return False
+
+
+def parse_wheel_requirement(req: str) -> Tuple[str, Optional[str]]:
+    """
+    Parse a wheel requirement string.
+
+    Examples:
+        "nvdiffrast==0.4.0" -> ("nvdiffrast", "0.4.0")
+        "pytorch3d>=0.7.8" -> ("pytorch3d", "0.7.8")
+        "torch-cluster" -> ("torch-cluster", None)
+
+    Returns:
+        Tuple of (package_name, version_or_None).
+    """
+    # Handle version specifiers
+    for op in ['==', '>=', '<=', '~=', '!=', '>', '<']:
+        if op in req:
+            parts = req.split(op, 1)
+            return (parts[0].strip(), parts[1].strip())
+
+    return (req.strip(), None)