comfy-env 0.0.48__py3-none-any.whl → 0.0.50__py3-none-any.whl

comfy_env/stub_imports.py

@@ -0,0 +1,310 @@
+"""
+Import stub system for isolated node packs.
+
+This module provides automatic import stubbing for packages that exist only
+in the isolated pixi environment, not in the host ComfyUI Python.
+
+How it works:
+1. Read package names from comfy-env.toml
+2. Look up their import names from top_level.txt in the pixi environment
+3. Register import hooks that provide stub modules for those imports
+4. Stubs allow class definitions to parse without the real packages
+5. Real packages are used when FUNCTION runs in the isolated worker
+
+Usage:
+    # In node pack's __init__.py, BEFORE importing nodes:
+    from comfy_env import setup_isolated_imports
+    setup_isolated_imports(__file__)
+
+    from .nodes import NODE_CLASS_MAPPINGS  # Now works!
+"""
+
+import sys
+import types
+from pathlib import Path
+from typing import Dict, List, Optional, Set
+
+
+class _StubModule(types.ModuleType):
+    """
+    A stub module that accepts any attribute access or call.
+    """
+
+    def __init__(self, name: str):
+        super().__init__(name)
+        self.__path__ = []  # Make it a package
+        self.__file__ = f"<stub:{name}>"
+        self._stub_name = name
+
+    def __getattr__(self, name: str):
+        if name.startswith('_'):
+            raise AttributeError(name)
+        return _StubObject(f"{self._stub_name}.{name}")
+
+    def __repr__(self):
+        return f"<StubModule '{self._stub_name}'>"
+
+
+class _StubObject:
+    """
+    A stub object that accepts any operation.
+    """
+
+    def __init__(self, name: str = "stub"):
+        self._stub_name = name
+
+    def __getattr__(self, name: str):
+        if name.startswith('_'):
+            raise AttributeError(name)
+        return _StubObject(f"{self._stub_name}.{name}")
+
+    def __call__(self, *args, **kwargs):
+        return _StubObject(f"{self._stub_name}()")
+
+    def __iter__(self):
+        return iter([])
+
+    def __len__(self):
+        return 0
+
+    def __bool__(self):
+        return False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+    def __repr__(self):
+        return f"<Stub '{self._stub_name}'>"
+
+    def __add__(self, other): return self
+    def __radd__(self, other): return self
+    def __sub__(self, other): return self
+    def __rsub__(self, other): return self
+    def __mul__(self, other): return self
+    def __rmul__(self, other): return self
+    def __truediv__(self, other): return self
+    def __rtruediv__(self, other): return self
+    def __eq__(self, other): return False
+    def __ne__(self, other): return True
+    def __lt__(self, other): return False
+    def __le__(self, other): return False
+    def __gt__(self, other): return False
+    def __ge__(self, other): return False
+    def __hash__(self): return hash(self._stub_name)
+    def __getitem__(self, key): return _StubObject(f"{self._stub_name}[{key}]")
+    def __setitem__(self, key, value): pass
+    def __contains__(self, item): return False
+
+
+class _StubFinder:
+    """Import hook finder that provides stub modules for specified packages."""
+
+    def __init__(self, stub_packages: Set[str]):
+        self.stub_packages = stub_packages
+
+    def find_module(self, fullname: str, path=None):
+        top_level = fullname.split('.')[0]
+        if top_level in self.stub_packages:
+            return _StubLoader(self.stub_packages)
+        return None
+
+
+class _StubLoader:
+    """Import hook loader that creates stub modules."""
+
+    def __init__(self, stub_packages: Set[str]):
+        self.stub_packages = stub_packages
+
+    def load_module(self, fullname: str):
+        if fullname in sys.modules:
+            return sys.modules[fullname]
+
+        module = _StubModule(fullname)
+        module.__loader__ = self
+
+        if '.' in fullname:
+            parent = fullname.rsplit('.', 1)[0]
+            module.__package__ = parent
+            if parent not in sys.modules:
+                self.load_module(parent)
+        else:
+            module.__package__ = fullname
+
+        sys.modules[fullname] = module
+        return module
+
+
+def _normalize_package_name(name: str) -> str:
+    """Normalize package name for comparison (PEP 503)."""
+    return name.lower().replace('-', '_').replace('.', '_')
+
+
+def _get_import_names_from_pixi(node_dir: Path) -> Set[str]:
+    """
+    Get import names by scanning the pixi environment's site-packages.
+
+    Finds all importable packages by looking for:
+    1. Directories with __init__.py (packages)
+    2. .py files (single-file modules)
+    3. .so/.pyd files (extension modules)
+
+    Returns:
+        Set of import names that should be stubbed.
+    """
+    import_names = set()
+
+    # Find the pixi site-packages
+    pixi_lib = node_dir / ".pixi" / "envs" / "default" / "lib"
+
+    if not pixi_lib.exists():
+        return import_names
+
+    # Find the python version directory (e.g., python3.11)
+    python_dirs = list(pixi_lib.glob("python3.*"))
+    if not python_dirs:
+        return import_names
+
+    site_packages = python_dirs[0] / "site-packages"
+    if not site_packages.exists():
+        return import_names
+
+    # Scan for importable modules
+    for item in site_packages.iterdir():
+        name = item.name
+
+        # Skip private/internal items
+        if name.startswith('_') or name.startswith('.'):
+            continue
+
+        # Skip dist-info and egg-info directories
+        if name.endswith('.dist-info') or name.endswith('.egg-info'):
+            continue
+
+        # Skip common non-module items
+        if name in {'bin', 'share', 'include', 'etc'}:
+            continue
+
+        # Package directory (has __init__.py)
+        if item.is_dir():
+            if (item / "__init__.py").exists():
+                import_names.add(name)
+            continue
+
+        # Single-file module (.py)
+        if name.endswith('.py'):
+            import_names.add(name[:-3])
+            continue
+
+        # Extension module (.so on Linux, .pyd on Windows)
+        if '.cpython-' in name and (name.endswith('.so') or name.endswith('.pyd')):
+            # Extract module name: foo.cpython-311-x86_64-linux-gnu.so -> foo
+            module_name = name.split('.')[0]
+            import_names.add(module_name)
+            continue
+
+    return import_names
+
+
+def _filter_to_missing(import_names: Set[str]) -> Set[str]:
+    """Filter to only imports not available in host Python."""
+    missing = set()
+
+    for name in import_names:
+        # Skip if already in sys.modules
+        if name in sys.modules:
+            continue
+
+        # Try to import
+        try:
+            __import__(name)
+        except ImportError:
+            missing.add(name)
+        except Exception:
+            # Other errors - don't stub, let real error surface
+            pass
+
+    return missing
+
+
+# Track whether we've already set up stubs
+_stub_finder: Optional[_StubFinder] = None
+
+
+def setup_isolated_imports(init_file: str) -> List[str]:
+    """
+    Set up import stubs for packages in the pixi environment but not in host Python.
+
+    Call this BEFORE importing your nodes module.
+
+    Args:
+        init_file: The __file__ of the calling module (usually __file__ from __init__.py)
+
+    Returns:
+        List of import names that were stubbed.
+
+    Example:
+        from comfy_env import setup_isolated_imports
+        setup_isolated_imports(__file__)
+
+        from .nodes import NODE_CLASS_MAPPINGS  # Now works!
+    """
+    global _stub_finder
+
+    node_dir = Path(init_file).resolve().parent
+
+    # Get all import names from pixi environment
+    pixi_imports = _get_import_names_from_pixi(node_dir)
+
+    if not pixi_imports:
+        print("[comfy-env] No pixi environment found, skipping import stubbing")
+        return []
+
+    # Filter to only those missing in host
+    missing = _filter_to_missing(pixi_imports)
+
+    if not missing:
+        print("[comfy-env] All pixi packages available in host, no stubbing needed")
+        return []
+
+    # Remove old finder if exists
+    if _stub_finder is not None:
+        try:
+            sys.meta_path.remove(_stub_finder)
+        except ValueError:
+            pass
+
+    # Register new finder
+    _stub_finder = _StubFinder(missing)
+    sys.meta_path.insert(0, _stub_finder)
+
+    stubbed = sorted(missing)
+    if len(stubbed) <= 10:
+        print(f"[comfy-env] Stubbed {len(stubbed)} imports: {', '.join(stubbed)}")
+    else:
+        print(f"[comfy-env] Stubbed {len(stubbed)} imports: {', '.join(stubbed[:10])}... and {len(stubbed)-10} more")
+
+    return stubbed
+
+
+def cleanup_stubs():
+    """Remove the stub import hooks."""
+    global _stub_finder
+
+    if _stub_finder is not None:
+        try:
+            sys.meta_path.remove(_stub_finder)
+        except ValueError:
+            pass
+
+        # Remove stubbed modules from sys.modules
+        to_remove = [
+            name for name in sys.modules
+            if isinstance(sys.modules[name], _StubModule)
+        ]
+        for name in to_remove:
+            del sys.modules[name]
+
+        _stub_finder = None

comfy_env/templates/comfy-env-instructions.txt

@@ -0,0 +1,103 @@
+comfy-env Setup Instructions
+============================
+
+comfy-env handles CUDA wheel installation and process isolation for ComfyUI
+custom nodes.
+
+QUICK START
+-----------
+1. Install comfy-env:
+   pip install comfy-env
+
+2. Initialize config in your node directory:
+   comfy-env init
+
+3. Edit comfy-env.toml to add your dependencies
+
+4. Add to your __init__.py (at the top, before other imports):
+   from comfy_env import install
+   install()
+
+5. Test locally:
+   comfy-env install --dry-run  # Preview what will be installed
+   comfy-env install            # Actually install
+   comfy-env doctor             # Verify all packages work
+
+
+COMMON USE CASES
+----------------
+
+Case 1: Just need CUDA packages (nvdiffrast, pytorch3d, etc.)
+  - Add packages to [cuda] section
+  - Call install() in your __init__.py
+
+Case 2: Need process isolation (conflicting dependencies, conda packages)
+  - Define an isolated environment with `isolated = true`
+  - Use enable_isolation(NODE_CLASS_MAPPINGS) in your __init__.py
+  - See PROCESS ISOLATION section below
+
+Case 3: Need system packages (apt)
+  - Add to [system] linux = ["package1", "package2"]
+
+
+CLI COMMANDS
+------------
+  comfy-env init           Create comfy-env.toml template
+  comfy-env install        Install dependencies from config
+  comfy-env install --dry-run  Preview without installing
+  comfy-env info           Show detected environment (Python, CUDA, PyTorch)
+  comfy-env resolve PKG    Show resolved wheel URL for a package
+  comfy-env doctor         Verify installation
+  comfy-env list-packages  Show all packages in built-in registry
+
+
+PROCESS ISOLATION
+-----------------
+For nodes that need isolated dependencies:
+
+RECOMMENDED: Pack-wide isolation (all nodes in same isolated env)
+
+  from comfy_env import setup_isolated_imports, enable_isolation
+
+  # Setup import stubs BEFORE importing nodes
+  setup_isolated_imports(__file__)
+
+  from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+  # Enable isolation for all nodes
+  enable_isolation(NODE_CLASS_MAPPINGS)
+
+Requires `isolated = true` in comfy-env.toml:
+
+  [mypack]
+  python = "3.11"
+  isolated = true
+
+  [mypack.packages]
+  requirements = ["trimesh", "scipy"]
+
+ALTERNATIVE: Per-node isolation (for multiple isolated envs)
+
+  from comfy_env import isolated
+
+  @isolated(env="myenv")
+  class MyNode:
+      FUNCTION = "process"
+      def process(self, image):
+          import conflicting_lib
+          return (result,)
+
+How it works:
+  - Runs FUNCTION methods in a separate Python process
+  - Tensors/numpy arrays passed by value (efficient)
+  - Complex objects (meshes, etc.) passed by reference
+
+
+TROUBLESHOOTING
+---------------
+- "Package X not found in registry": Add custom wheel URL to [wheel_sources]
+- "CUDA not detected": Ensure PyTorch with CUDA is installed in ComfyUI
+- "Worker failed to connect": Check the isolated env was set up correctly
+- Import errors: Run `comfy-env doctor` to verify packages
+
+For more help: https://github.com/PozzettiAndrea/comfy-env

comfy_env/templates/comfy-env.toml

@@ -0,0 +1,186 @@
+# =============================================================================
+# comfy-env.toml - Environment configuration for ComfyUI custom nodes
+# Documentation: https://github.com/PozzettiAndrea/comfy-env
+# =============================================================================
+#
+# This file configures dependencies for your ComfyUI custom node.
+# comfy-env handles two main challenges:
+#   1. Installing CUDA packages that require compilation (nvdiffrast, pytorch3d, etc.)
+#   2. Running nodes in isolated environments with their own dependencies
+#
+# Quick start:
+#   pip install comfy-env
+#   comfy-env init          # Creates this template
+#   comfy-env install       # Installs dependencies
+#   comfy-env doctor        # Verifies installation
+
+
+# =============================================================================
+# SYSTEM PACKAGES (optional)
+# =============================================================================
+# System-level packages installed via apt (Linux) before Python packages.
+# Useful for packages that need system libraries to compile or run.
+
+[system]
+linux = []
+# Examples:
+# linux = ["libgl1", "libopengl0"]     # For OpenGL rendering (nvdiffrast)
+# linux = ["python3-dev"]               # For packages that compile C extensions
+# linux = ["ffmpeg"]                    # For video processing nodes
+
+
+# =============================================================================
+# CUDA PACKAGES
+# =============================================================================
+# CUDA packages from the built-in registry. These are pre-compiled wheels
+# that install without needing a compiler. Run `comfy-env list-packages`
+# to see all available packages.
+
+[cuda]
+# Use exact package names from `comfy-env list-packages` (hyphens vs underscores matter!)
+# nvdiffrast = "0.4.0"       # NVIDIA differentiable rasterizer
+# pytorch3d = "0.7.9"        # PyTorch3D - 3D deep learning
+# gsplat = "1.5.3"           # Gaussian splatting rasterization
+# torch-scatter = "2.1.2"    # PyG scatter operations
+# torch-cluster = "1.6.3"    # PyG clustering algorithms
+# flash-attn = "2.7.4"       # Flash Attention (Linux x86_64 only)
+# spconv = "2.3.8"           # Sparse convolution library
+# sageattention = "2.2.0"    # SageAttention - faster than FlashAttention
+# dpvo_cuda = "0.0.0"        # DPVO CUDA extensions
+
+
+# =============================================================================
+# REGULAR PACKAGES
+# =============================================================================
+# Standard pip packages (no special CUDA handling needed).
+
+[packages]
+requirements = []
+# Examples:
+# requirements = [
+#     "transformers>=4.56",
+#     "pillow",
+#     "opencv-python-headless",
+#     "trimesh",
+# ]
+
+
+# =============================================================================
+# CUSTOM WHEEL SOURCES (optional)
+# =============================================================================
+# Override built-in wheel URLs or add packages not in the registry.
+# Template variables: {version}, {cuda_short}, {torch_mm}, {py_tag}, {platform}
+
+[wheel_sources]
+# my-custom-package = "https://my-server.com/my-package-{version}+cu{cuda_short}-{py_tag}-{platform}.whl"
+
+
+# =============================================================================
+# NODE DEPENDENCIES (optional)
+# =============================================================================
+# Other ComfyUI custom nodes this node depends on.
+
+[node_reqs]
+# VideoHelperSuite = "Kosinkadink/ComfyUI-VideoHelperSuite"
+# ComfyUI-Impact-Pack = "ltdrdata/ComfyUI-Impact-Pack"
+
+
+# =============================================================================
+# EXTERNAL TOOLS (optional)
+# =============================================================================
+# External applications required by the node.
+
+[tools]
+# blender = "4.2"
+
+
+# #############################################################################
+#
+#                       PROCESS ISOLATION (ADVANCED)
+#
+# #############################################################################
+#
+# For nodes that need completely isolated dependencies (different Python
+# version, conda packages, conflicting native libraries), define an isolated
+# environment with `isolated = true`.
+#
+# RECOMMENDED: Pack-wide isolation (all nodes in one environment)
+# ----------------------------------------------------------------
+# This is the simplest approach - all your nodes run in the same isolated env.
+#
+# Step 1: Define environment in comfy-env.toml (this file)
+# Step 2: In __init__.py:
+#
+#     from comfy_env import setup_isolated_imports, enable_isolation
+#
+#     # Setup import stubs BEFORE importing nodes
+#     setup_isolated_imports(__file__)
+#
+#     from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+#
+#     # Enable isolation for all nodes
+#     enable_isolation(NODE_CLASS_MAPPINGS)
+#
+# =============================================================================
+
+
+# -----------------------------------------------------------------------------
+# Example: Full pack isolation with conda packages (RECOMMENDED)
+# -----------------------------------------------------------------------------
+# Uses pixi to create an isolated environment with conda + pip packages.
+
+# [mypack]
+# python = "3.11"
+# isolated = true           # Required for enable_isolation()
+#
+# [mypack.conda]
+# channels = ["conda-forge"]
+# packages = ["cgal", "openmesh"]
+#
+# [mypack.packages]
+# requirements = ["trimesh[easy]>=4.0", "numpy", "scipy"]
+
+
+# -----------------------------------------------------------------------------
+# Example: Multiple isolated environments (per-node control)
+# -----------------------------------------------------------------------------
+# Use @isolated(env="envname") decorator when different nodes need different envs.
+#
+#     from comfy_env import isolated
+#
+#     @isolated(env="env-preprocessing")
+#     class PreprocessNode: ...
+#
+#     @isolated(env="env-inference")
+#     class InferenceNode: ...
+
+# [env-preprocessing]
+# python = "3.11"
+#
+# [env-preprocessing.packages]
+# requirements = ["opencv-python-headless", "pillow"]
+
+# [env-inference]
+# python = "3.10"
+#
+# [env-inference.cuda]
+# torch-scatter = "2.1.2"
+
+
+# -----------------------------------------------------------------------------
+# Example: Platform-specific packages
+# -----------------------------------------------------------------------------
+# Different packages for Windows vs Linux.
+
+# [crossplatform]
+# python = "3.11"
+# isolated = true
+#
+# [crossplatform.packages]
+# requirements = ["numpy", "pillow"]
+#
+# [crossplatform.packages.windows]
+# requirements = ["pywin32"]
+#
+# [crossplatform.packages.linux]
+# requirements = ["python-xlib"]

comfy_env/wheel_sources.yml

@@ -107,7 +107,7 @@ packages:
     default_version: "2.2.0"
     description: SageAttention - 2-5x faster than FlashAttention with quantized kernels
 
-  dpvo_cuda:
+  dpvo-cuda:
     wheel_template: "https://github.com/PozzettiAndrea/cuda-wheels/releases/download/dpvo_cuda-latest/dpvo_cuda-{version}%2Bcu{cuda_short}torch{torch_mm}-{py_tag}-{py_tag}-{platform}.whl"
     default_version: "0.0.0"
     description: DPVO CUDA extensions (cuda_corr, cuda_ba, lietorch_backends) - torch 2.4 only

comfy_env/workers/torch_mp.py

@@ -50,8 +50,12 @@ def _worker_loop(queue_in, queue_out, sys_path_additions=None):
     - _SHUTDOWN: Shutdown the worker
 
     Runs until receiving _SHUTDOWN sentinel.
+
+    Args:
+        queue_in: Input queue for receiving work items
+        queue_out: Output queue for sending results
+        sys_path_additions: Paths to add to sys.path
     """
-    import importlib
     import os
     import sys
     from pathlib import Path