comfy-env 0.1.2__tar.gz → 0.1.4__tar.gz

{comfy_env-0.1.2 → comfy_env-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comfy-env
-Version: 0.1.2
+Version: 0.1.4
 Summary: Environment management for ComfyUI custom nodes - CUDA wheel resolution and process isolation
 Project-URL: Homepage, https://github.com/PozzettiAndrea/comfy-env
 Project-URL: Repository, https://github.com/PozzettiAndrea/comfy-env

{comfy_env-0.1.2 → comfy_env-0.1.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "comfy-env"
-version = "0.1.2"
+version = "0.1.4"
 description = "Environment management for ComfyUI custom nodes - CUDA wheel resolution and process isolation"
 readme = "README.md"
 license = {text = "MIT"}

{comfy_env-0.1.2 → comfy_env-0.1.4}/src/comfy_env/isolation/wrap.py

@@ -67,6 +67,7 @@ def _get_worker(
     working_dir: Path,
     sys_path: list[str],
     lib_path: Optional[str] = None,
+    env_vars: Optional[dict] = None,
 ):
     """Get or create a persistent worker for the isolated environment."""
     cache_key = str(env_dir)
@@ -106,11 +107,14 @@ def _get_worker(
 
             print(f"[comfy-env] Starting isolated worker (MPWorker)")
             print(f"[comfy-env]   Env: {env_dir}")
+            if env_vars:
+                print(f"[comfy-env]   env_vars: {', '.join(f'{k}={v}' for k, v in env_vars.items())}")
 
             worker = MPWorker(
                 name=working_dir.name,
                 sys_path=sys_path,
                 lib_path=lib_path,
+                env_vars=env_vars,
             )
 
         _workers[cache_key] = worker
@@ -187,6 +191,7 @@ def _wrap_node_class(
     working_dir: Path,
     sys_path: list[str],
     lib_path: Optional[str] = None,
+    env_vars: Optional[dict] = None,
 ) -> type:
     """
     Wrap a node class so its FUNCTION method runs in the isolated environment.
@@ -232,7 +237,7 @@ def _wrap_node_class(
             print(f"[comfy-env] PROXY CALLED: {cls.__name__}.{func_name}", flush=True)
             print(f"[comfy-env]   kwargs keys: {list(kwargs.keys())}", flush=True)
 
-        worker = _get_worker(env_dir, working_dir, sys_path, lib_path)
+        worker = _get_worker(env_dir, working_dir, sys_path, lib_path, env_vars)
         if _DEBUG:
             print(f"[comfy-env]   worker alive: {worker.is_alive()}", flush=True)
 
@@ -322,6 +327,20 @@ def wrap_isolated_nodes(
         print(f"[comfy-env] Warning: No comfy-env.toml in {nodes_dir}")
         return node_class_mappings
 
+    # Read env_vars from comfy-env.toml
+    env_vars = {}
+    try:
+        if sys.version_info >= (3, 11):
+            import tomllib
+        else:
+            import tomli as tomllib
+        with open(config_file, "rb") as f:
+            config = tomllib.load(f)
+        env_vars_data = config.get("env_vars", {})
+        env_vars = {str(k): str(v) for k, v in env_vars_data.items()}
+    except Exception:
+        pass  # Ignore errors reading config
+
     # Find environment directory and paths
     env_dir = _find_env_dir(nodes_dir)
     site_packages, lib_dir = _find_env_paths(nodes_dir)
@@ -342,10 +361,12 @@ def wrap_isolated_nodes(
     print(f"[comfy-env]   site-packages: {site_packages}")
     if lib_path:
         print(f"[comfy-env]   lib: {lib_path}")
+    if env_vars:
+        print(f"[comfy-env]   env_vars: {', '.join(f'{k}={v}' for k, v in env_vars.items())}")
 
     # Wrap all node classes
     for node_name, node_cls in node_class_mappings.items():
         if hasattr(node_cls, "FUNCTION"):
-            _wrap_node_class(node_cls, env_dir, nodes_dir, sys_path, lib_path)
+            _wrap_node_class(node_cls, env_dir, nodes_dir, sys_path, lib_path, env_vars)
 
     return node_class_mappings

{comfy_env-0.1.2 → comfy_env-0.1.4}/src/comfy_env/pixi/core.py

@@ -9,10 +9,12 @@ See: https://pixi.sh/
 
 import copy
 import platform
+import re
 import shutil
 import stat
 import subprocess
 import sys
+import urllib.request
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional
 
@@ -28,15 +30,9 @@ PIXI_URLS = {
     ("Windows", "AMD64"): "https://github.com/prefix-dev/pixi/releases/latest/download/pixi-x86_64-pc-windows-msvc.exe",
 }
 
-# CUDA wheels index
+# CUDA wheels index (includes flash-attn, PyG packages, and custom wheels)
 CUDA_WHEELS_INDEX = "https://pozzettiandrea.github.io/cuda-wheels/"
 
-# Flash attention pre-built wheels (mjun0812)
-FLASH_ATTN_INDEX = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/expanded_assets/v0.7.16"
-
-# PyTorch Geometric wheels index
-PYG_WHEELS_INDEX = "https://data.pyg.org/whl/"
-
 # CUDA version -> PyTorch version mapping
 CUDA_TORCH_MAP = {
     "12.8": "2.8",
@@ -44,19 +40,87 @@ CUDA_TORCH_MAP = {
     "12.1": "2.4",
 }
 
+def find_matching_wheel(package: str, torch_version: str, cuda_version: str) -> Optional[str]:
+    """
+    Query cuda-wheels index to find a wheel matching the CUDA/torch version.
+    Returns the full version spec (e.g., "flash-attn===2.8.3+cu128torch2.8") or None.
+    """
+    cuda_short = cuda_version.replace(".", "")[:3]  # "12.8" -> "128"
+    torch_short = torch_version.replace(".", "")[:2]  # "2.8" -> "28"
+
+    # Try different directory name variants
+    pkg_variants = [package, package.replace("-", "_"), package.replace("_", "-")]
+
+    for pkg_dir in pkg_variants:
+        url = f"{CUDA_WHEELS_INDEX}{pkg_dir}/"
+        try:
+            with urllib.request.urlopen(url, timeout=10) as resp:
+                html = resp.read().decode("utf-8")
+        except Exception:
+            continue
+
+        # Parse wheel filenames from href attributes
+        # Pattern: package_name-version+localversion-cpXX-cpXX-platform.whl
+        wheel_pattern = re.compile(
+            r'href="[^"]*?([^"/]+\.whl)"',
+            re.IGNORECASE
+        )
+
+        # Local version patterns to match:
+        # flash-attn style: +cu128torch2.8
+        # PyG style: +pt28cu128
+        local_patterns = [
+            f"+cu{cuda_short}torch{torch_version}",  # flash-attn style
+            f"+pt{torch_short}cu{cuda_short}",       # PyG style
+        ]
+
+        best_match = None
+        best_version = None
+
+        for match in wheel_pattern.finditer(html):
+            wheel_name = match.group(1)
+            # URL decode
+            wheel_name = wheel_name.replace("%2B", "+")
+
+            # Check if wheel matches our CUDA/torch version
+            for local_pattern in local_patterns:
+                if local_pattern in wheel_name:
+                    # Extract version from wheel name
+                    # Format: name-version+local-cpXX-cpXX-platform.whl
+                    parts = wheel_name.split("-")
+                    if len(parts) >= 2:
+                        version_part = parts[1]  # e.g., "2.8.3+cu128torch2.8"
+                        if best_version is None or version_part > best_version:
+                            best_version = version_part
+                            best_match = f"{package}==={version_part}"
+                    break
+
+        if best_match:
+            return best_match
+
+    return None
+
+
+def get_package_spec(package: str, torch_version: str, cuda_version: str) -> str:
+    """
+    Get package spec with local version for CUDA/torch compatibility.
+    Queries the index to find matching wheels dynamically.
+    """
+    spec = find_matching_wheel(package, torch_version, cuda_version)
+    return spec if spec else package
+
 
 def get_all_find_links(package: str, torch_version: str, cuda_version: str) -> list:
     """Get all find-links URLs for a CUDA package."""
-    cuda_short = cuda_version.replace(".", "")[:3]
-
-    return [
-        # cuda-wheels index (package-specific page)
-        f"{CUDA_WHEELS_INDEX}{package}/",
-        # flash-attn GitHub releases
-        FLASH_ATTN_INDEX,
-        # PyG index (torch version specific)
-        f"{PYG_WHEELS_INDEX}torch-{torch_version}.0+cu{cuda_short}.html",
-    ]
+    # Try both underscore and hyphen variants since directory naming is inconsistent
+    pkg_underscore = package.replace("-", "_")
+    pkg_hyphen = package.replace("_", "-")
+    urls = [f"{CUDA_WHEELS_INDEX}{package}/"]
+    if pkg_underscore != package:
+        urls.append(f"{CUDA_WHEELS_INDEX}{pkg_underscore}/")
+    if pkg_hyphen != package:
+        urls.append(f"{CUDA_WHEELS_INDEX}{pkg_hyphen}/")
+    return urls
 
 
 def get_current_platform() -> str:
@@ -340,6 +404,7 @@ def pixi_install(
         for package in cfg.cuda_packages:
             pip_cmd = [
                 str(python_path), "-m", "pip", "install",
+                "--no-index",  # Only use --find-links, don't query PyPI
                 "--no-deps",
                 "--no-cache-dir",
             ]
@@ -348,8 +413,10 @@ def pixi_install(
             for url in get_all_find_links(package, torch_version, cuda_version):
                 pip_cmd.extend(["--find-links", url])
 
-            log(f"  Installing {package}")
-            pip_cmd.append(package)
+            # Get package spec with local version for CUDA/torch compatibility
+            pkg_spec = get_package_spec(package, torch_version, cuda_version)
+            log(f"  Installing {pkg_spec}")
+            pip_cmd.append(pkg_spec)
             result = subprocess.run(pip_cmd, capture_output=True, text=True)
             if result.returncode != 0:
                 log(f"CUDA package install failed for {package}:\n{result.stderr}")

{comfy_env-0.1.2 → comfy_env-0.1.4}/src/comfy_env/workers/mp.py

@@ -40,7 +40,7 @@ _SHUTDOWN = object()
 _CALL_METHOD = "call_method"
 
 
-def _worker_loop(queue_in, queue_out, sys_path_additions=None, lib_path=None):
+def _worker_loop(queue_in, queue_out, sys_path_additions=None, lib_path=None, env_vars=None):
     """
     Worker process main loop.
 
@@ -56,11 +56,16 @@ def _worker_loop(queue_in, queue_out, sys_path_additions=None, lib_path=None):
         queue_out: Output queue for sending results
         sys_path_additions: Paths to add to sys.path
         lib_path: Path to add to LD_LIBRARY_PATH (for conda libraries)
+        env_vars: Environment variables to set (from comfy-env.toml)
     """
     import os
     import sys
     from pathlib import Path
 
+    # Apply env_vars FIRST (before any library imports that might check them)
+    if env_vars:
+        os.environ.update(env_vars)
+
     # Set worker mode env var
     os.environ["COMFYUI_ISOLATION_WORKER"] = "1"
 
@@ -424,7 +429,7 @@ class MPWorker(Worker):
     interpreter without inherited state from the parent.
     """
 
-    def __init__(self, name: Optional[str] = None, sys_path: Optional[list] = None, lib_path: Optional[str] = None):
+    def __init__(self, name: Optional[str] = None, sys_path: Optional[list] = None, lib_path: Optional[str] = None, env_vars: Optional[dict] = None):
         """
         Initialize the worker.
 
@@ -432,10 +437,12 @@ class MPWorker(Worker):
             name: Optional name for logging/debugging.
             sys_path: Optional list of paths to add to sys.path in worker process.
             lib_path: Optional path to add to LD_LIBRARY_PATH (for conda libraries).
+            env_vars: Optional environment variables to set in worker process.
         """
         self.name = name or "MPWorker"
         self._sys_path = sys_path or []
         self._lib_path = lib_path
+        self._env_vars = env_vars or {}
         self._process = None
         self._queue_in = None
         self._queue_out = None
@@ -504,7 +511,7 @@ class MPWorker(Worker):
             self._queue_out = ctx.Queue()
             self._process = ctx.Process(
                 target=_worker_loop,
-                args=(self._queue_in, self._queue_out, self._sys_path, self._lib_path),
+                args=(self._queue_in, self._queue_out, self._sys_path, self._lib_path, self._env_vars),
                 daemon=True,
             )
             self._process.start()