comfy-env 0.1.8__tar.gz → 0.1.10__tar.gz

{comfy_env-0.1.8 → comfy_env-0.1.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comfy-env
-Version: 0.1.8
+Version: 0.1.10
 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
@@ -17,6 +17,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.10
+Requires-Dist: pip>=21.0
 Requires-Dist: tomli-w>=1.0.0
 Requires-Dist: tomli>=2.0.0; python_version < '3.11'
 Requires-Dist: uv>=0.4.0

{comfy_env-0.1.8 → comfy_env-0.1.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "comfy-env"
-version = "0.1.8"
+version = "0.1.10"
 description = "Environment management for ComfyUI custom nodes - CUDA wheel resolution and process isolation"
 readme = "README.md"
 license = {text = "MIT"}
@@ -22,6 +22,7 @@ dependencies = [
     "tomli>=2.0.0; python_version < '3.11'",  # TOML parsing (built-in tomllib for 3.11+)
     "tomli-w>=1.0.0",  # TOML writing (no stdlib equivalent)
     "uv>=0.4.0",  # Fast Python package installer and venv creator
+    "pip>=21.0",  # Fallback package installer
 ]
 
 [project.optional-dependencies]

{comfy_env-0.1.8 → comfy_env-0.1.10}/src/comfy_env/isolation/wrap.py

@@ -34,6 +34,12 @@ from typing import Any, Dict, Optional
 # Debug logging (set COMFY_ENV_DEBUG=1 to enable)
 _DEBUG = os.environ.get("COMFY_ENV_DEBUG", "").lower() in ("1", "true", "yes")
 
+
+def get_env_name(dir_name: str) -> str:
+    """Convert directory name to env name: ComfyUI-UniRig → _env_unirig"""
+    name = dir_name.lower().replace("-", "_").lstrip("comfyui_")
+    return f"_env_{name}"
+
 # Global worker cache (one per isolated environment)
 _workers: Dict[str, Any] = {}
 _workers_lock = threading.Lock()
@@ -144,13 +150,27 @@ def _find_env_paths(node_dir: Path) -> tuple[Optional[Path], Optional[Path]]:
     """
     import glob
 
-    # Check pixi environment first
+    # Check _env_<name> directory first (new pattern)
+    env_name = get_env_name(node_dir.name)
+    env_dir = node_dir / env_name
+    if env_dir.exists():
+        if sys.platform == "win32":
+            site_packages = env_dir / "Lib" / "site-packages"
+            lib_dir = env_dir / "Library" / "bin"
+        else:
+            pattern = str(env_dir / "lib" / "python*" / "site-packages")
+            matches = glob.glob(pattern)
+            site_packages = Path(matches[0]) if matches else None
+            lib_dir = env_dir / "lib"
+        if site_packages and site_packages.exists():
+            return site_packages, lib_dir if lib_dir.exists() else None
+
+    # Fallback: Check old .pixi/envs/default (for backward compat)
     pixi_env = node_dir / ".pixi" / "envs" / "default"
     if pixi_env.exists():
-        # Find site-packages (pythonX.Y varies)
         if sys.platform == "win32":
             site_packages = pixi_env / "Lib" / "site-packages"
-            lib_dir = pixi_env / "Library" / "bin"  # Windows DLLs
+            lib_dir = pixi_env / "Library" / "bin"
         else:
             pattern = str(pixi_env / "lib" / "python*" / "site-packages")
             matches = glob.glob(pattern)
@@ -176,6 +196,12 @@ def _find_env_paths(node_dir: Path) -> tuple[Optional[Path], Optional[Path]]:
 
 def _find_env_dir(node_dir: Path) -> Optional[Path]:
     """Find the environment directory (for cache key)."""
+    # Check _env_<name> first
+    env_name = get_env_name(node_dir.name)
+    env_dir = node_dir / env_name
+    if env_dir.exists():
+        return env_dir
+    # Fallback to old paths
     pixi_env = node_dir / ".pixi" / "envs" / "default"
     if pixi_env.exists():
         return pixi_env
@@ -319,6 +345,13 @@ def wrap_isolated_nodes(
     if os.environ.get("COMFYUI_ISOLATION_WORKER") == "1":
         return node_class_mappings
 
+    # Get ComfyUI base path from folder_paths (canonical source)
+    try:
+        import folder_paths
+        comfyui_base = folder_paths.base_path
+    except ImportError:
+        comfyui_base = None
+
     nodes_dir = Path(nodes_dir).resolve()
 
     # Check for comfy-env.toml
@@ -341,6 +374,10 @@ def wrap_isolated_nodes(
     except Exception:
         pass  # Ignore errors reading config
 
+    # Set COMFYUI_BASE for worker to find ComfyUI modules
+    if comfyui_base:
+        env_vars["COMFYUI_BASE"] = str(comfyui_base)
+
     # Find environment directory and paths
     env_dir = _find_env_dir(nodes_dir)
     site_packages, lib_dir = _find_env_paths(nodes_dir)

{comfy_env-0.1.8 → comfy_env-0.1.10}/src/comfy_env/nodes.py

@@ -71,6 +71,31 @@ def clone_node(
     return node_path
 
 
+def install_requirements(
+    node_dir: Path,
+    log: Callable[[str], None],
+) -> None:
+    """
+    Install requirements.txt in a node directory if it exists.
+
+    Args:
+        node_dir: Path to the node directory
+        log: Logging callback
+    """
+    requirements_file = node_dir / "requirements.txt"
+
+    if requirements_file.exists():
+        log(f"  Installing requirements for {node_dir.name}...")
+        result = subprocess.run(
+            ["uv", "pip", "install", "-r", str(requirements_file), "--python", sys.executable],
+            cwd=node_dir,
+            capture_output=True,
+            text=True,
+        )
+        if result.returncode != 0:
+            log(f"  Warning: requirements.txt install failed for {node_dir.name}: {result.stderr.strip()[:200]}")
+
+
 def run_install_script(
     node_dir: Path,
     log: Callable[[str], None],
@@ -132,6 +157,9 @@ def install_node_deps(
             # Clone the repository
             clone_node(req.repo, req.name, custom_nodes_dir, log)
 
+            # Install requirements.txt if present
+            install_requirements(node_path, log)
+
             # Run install.py if present
             run_install_script(node_path, log)
 

{comfy_env-0.1.8 → comfy_env-0.1.10}/src/comfy_env/pixi/core.py

@@ -40,10 +40,87 @@ CUDA_TORCH_MAP = {
     "12.1": "2.4",
 }
 
+def find_wheel_url(
+    package: str,
+    torch_version: str,
+    cuda_version: str,
+    python_version: str,
+) -> Optional[str]:
+    """
+    Query cuda-wheels index and return the direct URL for the matching wheel.
+
+    This bypasses pip's version validation by providing a direct URL,
+    which is necessary for wheels where the filename has a local version
+    but the internal METADATA doesn't (e.g., flash-attn from mjun0812).
+
+    Args:
+        package: Package name (e.g., "flash-attn")
+        torch_version: PyTorch version (e.g., "2.8")
+        cuda_version: CUDA version (e.g., "12.8")
+        python_version: Python version (e.g., "3.10")
+
+    Returns:
+        Direct URL to the wheel file, or None if no match found.
+    """
+    cuda_short = cuda_version.replace(".", "")[:3]  # "12.8" -> "128"
+    torch_short = torch_version.replace(".", "")[:2]  # "2.8" -> "28"
+    py_tag = f"cp{python_version.replace('.', '')}"  # "3.10" -> "cp310"
+
+    # Platform tag for current system
+    if sys.platform == "linux":
+        platform_tag = "linux_x86_64"
+    elif sys.platform == "win32":
+        platform_tag = "win_amd64"
+    else:
+        platform_tag = None  # macOS doesn't typically have CUDA wheels
+
+    # Local version patterns to match:
+    # cuda-wheels style: +cu128torch28
+    # PyG style: +pt28cu128
+    local_patterns = [
+        f"+cu{cuda_short}torch{torch_short}",  # cuda-wheels style
+        f"+pt{torch_short}cu{cuda_short}",     # PyG style
+    ]
+
+    pkg_variants = [package, package.replace("-", "_"), package.replace("_", "-")]
+
+    for pkg_dir in pkg_variants:
+        index_url = f"{CUDA_WHEELS_INDEX}{pkg_dir}/"
+        try:
+            with urllib.request.urlopen(index_url, timeout=10) as resp:
+                html = resp.read().decode("utf-8")
+        except Exception:
+            continue
+
+        # Parse href and display name from HTML: <a href="URL">DISPLAY_NAME</a>
+        link_pattern = re.compile(r'href="([^"]+\.whl)"[^>]*>([^<]+)</a>', re.IGNORECASE)
+
+        for match in link_pattern.finditer(html):
+            wheel_url = match.group(1)
+            display_name = match.group(2)
+
+            # Match on display name (has normalized torch28 format)
+            matches_cuda_torch = any(p in display_name for p in local_patterns)
+            matches_python = py_tag in display_name
+            matches_platform = platform_tag is None or platform_tag in display_name
+
+            if matches_cuda_torch and matches_python and matches_platform:
+                # Return absolute URL
+                if wheel_url.startswith("http"):
+                    return wheel_url
+                # Relative URL - construct absolute
+                return f"{CUDA_WHEELS_INDEX}{pkg_dir}/{wheel_url}"
+
+    return None
+
+
 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.
+
+    Note: This is used as a fallback for packages with correct wheel metadata.
+    For packages with mismatched metadata (like flash-attn), use find_wheel_url() instead.
     """
     cuda_short = cuda_version.replace(".", "")[:3]  # "12.8" -> "128"
     torch_short = torch_version.replace(".", "")[:2]  # "2.8" -> "28"
@@ -67,11 +144,11 @@ def find_matching_wheel(package: str, torch_version: str, cuda_version: str) ->
         )
 
         # Local version patterns to match:
-        # flash-attn style: +cu128torch2.8
+        # cuda-wheels style: +cu128torch28
         # PyG style: +pt28cu128
         local_patterns = [
-            f"+cu{cuda_short}torch{torch_version}",  # flash-attn style
-            f"+pt{torch_short}cu{cuda_short}",       # PyG style
+            f"+cu{cuda_short}torch{torch_short}",  # cuda-wheels style
+            f"+pt{torch_short}cu{cuda_short}",     # PyG style
         ]
 
         best_match = None
@@ -204,6 +281,12 @@ def ensure_pixi(
     return pixi_path
 
 
+def get_env_name(dir_name: str) -> str:
+    """Convert directory name to env name: ComfyUI-UniRig → _env_unirig"""
+    name = dir_name.lower().replace("-", "_").lstrip("comfyui_")
+    return f"_env_{name}"
+
+
 def clean_pixi_artifacts(node_dir: Path, log: Callable[[str], None] = print) -> None:
     """Remove previous pixi installation artifacts."""
     for path in [node_dir / "pixi.toml", node_dir / "pixi.lock"]:
@@ -212,11 +295,21 @@ def clean_pixi_artifacts(node_dir: Path, log: Callable[[str], None] = print) ->
     pixi_dir = node_dir / ".pixi"
     if pixi_dir.exists():
         shutil.rmtree(pixi_dir)
+    # Also clean old _env_* directories
+    env_name = get_env_name(node_dir.name)
+    env_dir = node_dir / env_name
+    if env_dir.exists():
+        shutil.rmtree(env_dir)
 
 
 def get_pixi_python(node_dir: Path) -> Optional[Path]:
     """Get path to Python in the pixi environment."""
-    env_dir = node_dir / ".pixi" / "envs" / "default"
+    # Check new _env_<name> location first
+    env_name = get_env_name(node_dir.name)
+    env_dir = node_dir / env_name
+    if not env_dir.exists():
+        # Fallback to old .pixi path
+        env_dir = node_dir / ".pixi" / "envs" / "default"
     if sys.platform == "win32":
         python_path = env_dir / "python.exe"
     else:
@@ -394,29 +487,35 @@ def pixi_install(
         log(f"pixi install failed:\n{result.stderr}")
         raise RuntimeError(f"pixi install failed: {result.stderr}")
 
-    # Install CUDA packages with --find-links (searches all known sources)
+    # Install CUDA packages via direct URL or find-links fallback
     if cfg.cuda_packages and cuda_version:
         log(f"Installing CUDA packages: {cfg.cuda_packages}")
         python_path = get_pixi_python(node_dir)
         if not python_path:
             raise RuntimeError("Could not find Python in pixi environment")
 
+        # Get Python version from the pixi environment
+        py_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+
         for package in cfg.cuda_packages:
+            # Find direct wheel URL (bypasses metadata validation)
+            wheel_url = find_wheel_url(package, torch_version, cuda_version, py_version)
+
+            if not wheel_url:
+                raise RuntimeError(
+                    f"No wheel found for {package} with CUDA {cuda_version}, "
+                    f"torch {torch_version}, Python {py_version}. "
+                    f"Check cuda-wheels index."
+                )
+
+            log(f"  Installing {package} from {wheel_url}")
             pip_cmd = [
                 str(python_path), "-m", "pip", "install",
-                "--no-index",  # Only use --find-links, don't query PyPI
                 "--no-deps",
                 "--no-cache-dir",
+                wheel_url,
             ]
 
-            # Add all find-links sources
-            for url in get_all_find_links(package, torch_version, cuda_version):
-                pip_cmd.extend(["--find-links", url])
-
-            # 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}")
@@ -424,5 +523,20 @@ def pixi_install(
 
         log("CUDA packages installed")
 
+    # Move environment from .pixi/envs/default to _env_<name>
+    old_env = node_dir / ".pixi" / "envs" / "default"
+    env_name = get_env_name(node_dir.name)
+    new_env = node_dir / env_name
+
+    if old_env.exists():
+        if new_env.exists():
+            shutil.rmtree(new_env)  # Clean old env
+        shutil.move(str(old_env), str(new_env))
+        # Clean up .pixi directory (keep pixi.toml and pixi.lock)
+        pixi_dir = node_dir / ".pixi"
+        if pixi_dir.exists():
+            shutil.rmtree(pixi_dir)
+        log(f"Moved environment to {new_env}")
+
     log("Installation complete!")
     return True

{comfy_env-0.1.8 → comfy_env-0.1.10}/src/comfy_env/prestartup.py

@@ -11,6 +11,12 @@ from pathlib import Path
 from typing import Optional, Dict
 
 
+def get_env_name(dir_name: str) -> str:
+    """Convert directory name to env name: ComfyUI-UniRig → _env_unirig"""
+    name = dir_name.lower().replace("-", "_").lstrip("comfyui_")
+    return f"_env_{name}"
+
+
 def _load_env_vars(config_path: str) -> Dict[str, str]:
     """
     Load [env_vars] section from comfy-env.toml.
@@ -121,10 +127,15 @@ def setup_env(node_dir: Optional[str] = None) -> None:
     for key, value in env_vars.items():
         os.environ[key] = value
 
-    pixi_env = os.path.join(node_dir, ".pixi", "envs", "default")
+    # Check _env_<name> first, then fallback to old .pixi path
+    env_name = get_env_name(os.path.basename(node_dir))
+    pixi_env = os.path.join(node_dir, env_name)
 
     if not os.path.exists(pixi_env):
-        return  # No pixi environment
+        # Fallback to old .pixi path
+        pixi_env = os.path.join(node_dir, ".pixi", "envs", "default")
+        if not os.path.exists(pixi_env):
+            return  # No environment found
 
     if sys.platform == "win32":
         # Windows: add to PATH for DLL loading

{comfy_env-0.1.8 → comfy_env-0.1.10}/src/comfy_env/workers/mp.py

@@ -40,6 +40,53 @@ _SHUTDOWN = object()
 _CALL_METHOD = "call_method"
 
 
+# ---------------------------------------------------------------------------
+# Tensor file transfer - avoids CUDA IPC issues with cudaMallocAsync
+# ---------------------------------------------------------------------------
+
+def _save_tensors_to_files(obj, file_registry=None):
+    """Recursively save torch tensors to temp files for IPC."""
+    if file_registry is None:
+        file_registry = []
+
+    try:
+        import torch
+        if isinstance(obj, torch.Tensor):
+            import tempfile
+            f = tempfile.NamedTemporaryFile(suffix='.pt', delete=False)
+            torch.save(obj.cpu(), f.name)  # Always save as CPU tensor
+            f.close()
+            file_registry.append(f.name)
+            return {"__tensor_file__": f.name, "dtype": str(obj.dtype), "device": str(obj.device)}
+    except ImportError:
+        pass
+
+    if isinstance(obj, dict):
+        return {k: _save_tensors_to_files(v, file_registry) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [_save_tensors_to_files(v, file_registry) for v in obj]
+    elif isinstance(obj, tuple):
+        return tuple(_save_tensors_to_files(v, file_registry) for v in obj)
+    return obj
+
+
+def _load_tensors_from_files(obj):
+    """Recursively load torch tensors from temp files."""
+    if isinstance(obj, dict):
+        if "__tensor_file__" in obj:
+            import os
+            import torch
+            tensor = torch.load(obj["__tensor_file__"], weights_only=True)
+            os.unlink(obj["__tensor_file__"])  # Cleanup temp file
+            return tensor
+        return {k: _load_tensors_from_files(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [_load_tensors_from_files(v) for v in obj]
+    elif isinstance(obj, tuple):
+        return tuple(_load_tensors_from_files(v) for v in obj)
+    return obj
+
+
 def _dump_worker_env(worker_name: str = "unknown", print_to_terminal: bool = False):
     """Dump worker environment to .comfy-env/logs/ (always) and optionally print."""
     import json
@@ -205,14 +252,20 @@ def _worker_loop(queue_in, queue_out, sys_path_additions=None, lib_path=None, en
                 # Handle method call protocol
                 if isinstance(item, tuple) and len(item) == 6 and item[0] == _CALL_METHOD:
                     _, module_name, class_name, method_name, self_state, kwargs = item
+                    # Load tensors from files (saved by host to avoid cudaMallocAsync IPC issues)
+                    kwargs = _load_tensors_from_files(kwargs)
                     result = _execute_method_call(
                         module_name, class_name, method_name, self_state, kwargs
                     )
+                    # Save tensors to files to avoid CUDA IPC issues with cudaMallocAsync
+                    result = _save_tensors_to_files(result)
                     queue_out.put(("ok", result))
                 else:
                     # Direct function call (legacy)
                     func, args, kwargs = item
                     result = func(*args, **kwargs)
+                    # Save tensors to files to avoid CUDA IPC issues with cudaMallocAsync
+                    result = _save_tensors_to_files(result)
                     queue_out.put(("ok", result))
 
             except Exception as e:
@@ -646,6 +699,9 @@ class MPWorker(Worker):
         """
         self._ensure_started()
 
+        # Save tensors to files to avoid CUDA IPC issues with cudaMallocAsync
+        kwargs = _save_tensors_to_files(kwargs)
+
         # Send method call request using protocol
         self._queue_in.put((
             _CALL_METHOD,
@@ -672,6 +728,8 @@ class MPWorker(Worker):
 
         # Handle response
         if status == "ok":
+            # Load tensors from temp files
+            result = _load_tensors_from_files(result)
             return result
         elif status == "error":
             msg, tb = result

{comfy_env-0.1.8 → comfy_env-0.1.10}/src/comfy_env/workers/subprocess.py

@@ -221,10 +221,12 @@ def _to_shm(obj, registry, visited=None):
         visited[obj_id] = result
         return result
 
-    # torch.Tensor → convert to numpy → shared memory
+    # torch.Tensor → convert to numpy → shared memory (with marker to restore type)
     if t == 'Tensor':
         arr = obj.detach().cpu().numpy()
-        return _to_shm(arr, registry, visited)
+        result = _to_shm(arr, registry, visited)
+        result["__was_tensor__"] = True
+        return result
 
     # trimesh.Trimesh → vertices + faces arrays → shared memory
     if t == 'Trimesh':
@@ -279,13 +281,17 @@ def _from_shm(obj, unlink=True):
             return [_from_shm(v, unlink) for v in obj]
         return obj
 
-    # numpy array
+    # numpy array (or tensor that was converted to numpy)
     if "__shm_np__" in obj:
         block = shm.SharedMemory(name=obj["__shm_np__"])
         arr = np.ndarray(tuple(obj["shape"]), dtype=np.dtype(obj["dtype"]), buffer=block.buf).copy()
         block.close()
         if unlink:
             block.unlink()
+        # Convert back to tensor if it was originally a tensor
+        if obj.get("__was_tensor__"):
+            import torch
+            return torch.from_numpy(arr)
         return arr
 
     # trimesh
@@ -544,7 +550,9 @@ def _to_shm(obj, registry, visited=None):
 
     if t == 'Tensor':
         arr = obj.detach().cpu().numpy()
-        return _to_shm(arr, registry, visited)
+        result = _to_shm(arr, registry, visited)
+        result["__was_tensor__"] = True
+        return result
 
     if t == 'Trimesh':
         verts = np.ascontiguousarray(obj.vertices, dtype=np.float64)
@@ -587,6 +595,10 @@ def _from_shm(obj):
         block = shm.SharedMemory(name=obj["__shm_np__"])
         arr = np.ndarray(tuple(obj["shape"]), dtype=np.dtype(obj["dtype"]), buffer=block.buf).copy()
         block.close()
+        # Convert back to tensor if it was originally a tensor
+        if obj.get("__was_tensor__"):
+            import torch
+            return torch.from_numpy(arr)
         return arr
     if "__shm_trimesh__" in obj:
         import trimesh
@@ -1044,15 +1056,21 @@ class SubprocessWorker(Worker):
 
     def _find_comfyui_base(self) -> Optional[Path]:
         """Find ComfyUI base directory."""
-        # Check common child directories (for test environments)
-        # Also check parent's children (isolated venv is sibling to .comfy-test-env)
+        # Use folder_paths.base_path (canonical source) if available
+        try:
+            import folder_paths
+            return Path(folder_paths.base_path)
+        except ImportError:
+            pass
+
+        # Fallback: Check common child directories (for test environments)
         for base in [self.working_dir, self.working_dir.parent]:
             for child in [".comfy-test-env/ComfyUI", "ComfyUI"]:
                 candidate = base / child
                 if (candidate / "main.py").exists() and (candidate / "comfy").exists():
                     return candidate
 
-        # Walk up from working_dir (standard ComfyUI custom_nodes layout)
+        # Fallback: Walk up from working_dir (standard ComfyUI custom_nodes layout)
         current = self.working_dir.resolve()
         for _ in range(10):
             if (current / "main.py").exists() and (current / "comfy").exists():