comfy-env 0.0.64__py3-none-any.whl → 0.0.66__py3-none-any.whl

comfy_env/stub_imports.py

@@ -1,270 +0,0 @@
-"""
-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. Scan pixi environment's site-packages for installed packages
-2. Look up import names from top_level.txt in .dist-info directories
-3. Inject stub modules directly into sys.modules for missing packages
-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 List, Set
-
-
-def _log(msg: str) -> None:
-    """Log with immediate flush to stderr (visible on Windows subprocess)."""
-    print(msg, file=sys.stderr, flush=True)
-
-
-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
-
-
-def _get_import_names_from_pixi(node_dir: Path) -> Set[str]:
-    """
-    Get import names from pixi environment using top_level.txt metadata.
-
-    This properly maps package names to import names (e.g., libigl -> igl,
-    PyYAML -> yaml) by reading the canonical top_level.txt files.
-
-    Returns:
-        Set of import names that should be stubbed.
-    """
-    import_names = set()
-
-    pixi_base = node_dir / ".pixi" / "envs" / "default"
-
-    # Find site-packages (different paths on Windows vs Linux)
-    site_packages = None
-    win_site = pixi_base / "Lib" / "site-packages"
-    if win_site.exists():
-        site_packages = win_site
-    else:
-        pixi_lib = pixi_base / "lib"
-        if pixi_lib.exists():
-            python_dirs = list(pixi_lib.glob("python3.*"))
-            if python_dirs:
-                site_packages = python_dirs[0] / "site-packages"
-
-    if site_packages is None or not site_packages.exists():
-        return import_names
-
-    _log(f"[comfy-env] Scanning: {site_packages}")
-
-    # PRIMARY: Read top_level.txt from all .dist-info directories
-    for dist_info in site_packages.glob("*.dist-info"):
-        top_level_file = dist_info / "top_level.txt"
-        if top_level_file.exists():
-            try:
-                for line in top_level_file.read_text(encoding="utf-8").splitlines():
-                    name = line.strip()
-                    if name and not name.startswith('#'):
-                        # Extract just the top-level name
-                        top_name = name.replace('\\', '/').split('/')[0]
-                        if top_name:
-                            import_names.add(top_name)
-            except Exception:
-                pass
-
-    # FALLBACK: Scan for packages/modules not covered by dist-info
-    for item in site_packages.iterdir():
-        name = item.name
-
-        if name.startswith('_') or name.startswith('.'):
-            continue
-        if name.endswith('.dist-info') or name.endswith('.egg-info'):
-            continue
-        if name in {'bin', 'share', 'include', 'etc'}:
-            continue
-
-        # Package directory (has __init__.py)
-        if item.is_dir() and (item / "__init__.py").exists():
-            import_names.add(name)
-            continue
-
-        # Namespace package (directory without __init__.py but has submodules)
-        if item.is_dir():
-            has_py = any(item.glob("*.py"))
-            has_subpkg = any((item / d / "__init__.py").exists() for d in item.iterdir() if d.is_dir())
-            if has_py or has_subpkg:
-                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 name.endswith('.so') or name.endswith('.pyd'):
-            module_name = name.split('.')[0]
-            import_names.add(module_name)
-
-    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 (DLL load, etc.) - stub these too
-            missing.add(name)
-
-    return missing
-
-
-# Track what we stubbed for cleanup
-_stubbed_modules: Set[str] = set()
-
-
-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 _stubbed_modules
-
-    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:
-        _log("[comfy-env] No pixi environment found")
-        return []
-
-    # Filter to only those missing in host
-    missing = _filter_to_missing(pixi_imports)
-
-    if not missing:
-        _log("[comfy-env] All packages available in host")
-        return []
-
-    # Direct injection into sys.modules - simple and reliable
-    for name in missing:
-        if name not in sys.modules:
-            sys.modules[name] = _StubModule(name)
-            _stubbed_modules.add(name)
-
-    stubbed = sorted(_stubbed_modules)
-    if len(stubbed) <= 10:
-        _log(f"[comfy-env] Injected {len(stubbed)} stubs: {', '.join(stubbed)}")
-    else:
-        _log(f"[comfy-env] Injected {len(stubbed)} stubs: {', '.join(stubbed[:10])}... +{len(stubbed)-10} more")
-
-    return stubbed
-
-
-def cleanup_stubs():
-    """Remove injected stub modules from sys.modules."""
-    global _stubbed_modules
-
-    for name in list(_stubbed_modules):
-        if name in sys.modules and isinstance(sys.modules[name], _StubModule):
-            del sys.modules[name]
-
-    _stubbed_modules.clear()

comfy_env/stubs/__init__.py

@@ -1 +0,0 @@
-# ComfyUI stubs for isolated workers

comfy_env/stubs/comfy/__init__.py

@@ -1,6 +0,0 @@
-"""
-Comfy stubs package for isolated worker processes.
-
-Provides minimal implementations of commonly used comfy modules
-without requiring the full ComfyUI installation.
-"""

comfy_env/stubs/comfy/model_management.py

@@ -1,58 +0,0 @@
-"""
-Stub for comfy.model_management in isolated worker processes.
-
-Provides device detection and memory management functions without
-requiring the full ComfyUI installation.
-"""
-
-import torch
-
-
-def get_torch_device():
-    """Return the best available torch device."""
-    if torch.cuda.is_available():
-        return torch.device("cuda")
-    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-        return torch.device("mps")
-    return torch.device("cpu")
-
-
-def get_free_memory(device=None, torch_free_too=False):
-    """Return free VRAM in bytes."""
-    if device is None:
-        device = get_torch_device()
-    if device.type == "cuda":
-        free, total = torch.cuda.mem_get_info(device)
-        return free
-    return 0
-
-
-def get_total_memory(device=None, torch_total_too=False):
-    """Return total VRAM in bytes."""
-    if device is None:
-        device = get_torch_device()
-    if device.type == "cuda":
-        free, total = torch.cuda.mem_get_info(device)
-        return total
-    return 0
-
-
-def soft_empty_cache(force=False):
-    """Clear CUDA cache."""
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
-
-def unload_all_models():
-    """No-op in isolated worker - models managed by the node itself."""
-    pass
-
-
-def interrupt_current_processing(value=True):
-    """No-op in isolated worker."""
-    pass
-
-
-def processing_interrupted():
-    """Always returns False in isolated worker."""
-    return False

comfy_env/stubs/comfy/utils.py

@@ -1,29 +0,0 @@
-"""
-Stub for comfy.utils in isolated worker processes.
-
-Provides utility classes like ProgressBar without requiring
-the full ComfyUI installation.
-"""
-
-
-class ProgressBar:
-    """
-    No-op progress bar for isolated workers.
-
-    In isolated subprocess, we can't update the main ComfyUI progress bar,
-    so this just tracks progress internally. Nodes can still use the same API.
-    """
-
-    def __init__(self, total):
-        self.total = total
-        self.current = 0
-
-    def update(self, value):
-        """Increment progress by value."""
-        self.current += value
-
-    def update_absolute(self, value, total=None, preview=None):
-        """Set progress to absolute value."""
-        self.current = value
-        if total is not None:
-            self.total = total

comfy_env/stubs/folder_paths.py

@@ -1,71 +0,0 @@
-"""
-Minimal folder_paths stub for isolated worker processes.
-
-Provides the same interface as ComfyUI's folder_paths module
-without importing any ComfyUI dependencies.
-"""
-
-import os
-from pathlib import Path
-
-_comfyui_base = None
-
-def _find_comfyui_base():
-    """Find ComfyUI base from COMFYUI_BASE env var, child dirs, or by walking up."""
-    global _comfyui_base
-    if _comfyui_base:
-        return _comfyui_base
-
-    # Check env var first
-    if os.environ.get("COMFYUI_BASE"):
-        _comfyui_base = Path(os.environ["COMFYUI_BASE"])
-        return _comfyui_base
-
-    # Check common child directories (for test environments)
-    # Also check parent's children (isolated venv is sibling to .comfy-test-env)
-    cwd = Path.cwd().resolve()
-    for base in [cwd, cwd.parent]:
-        for child in [".comfy-test-env/ComfyUI", "ComfyUI"]:
-            candidate = base / child
-            if (candidate / "main.py").exists() and (candidate / "comfy").exists():
-                _comfyui_base = candidate
-                return _comfyui_base
-
-    # Walk up from cwd looking for ComfyUI
-    current = cwd
-    for _ in range(10):
-        if (current / "main.py").exists() and (current / "comfy").exists():
-            _comfyui_base = current
-            return _comfyui_base
-        current = current.parent
-
-    return None
-
-# Models directory
-@property
-def models_dir():
-    base = _find_comfyui_base()
-    return str(base / "models") if base else None
-
-# Make models_dir work as both attribute and property
-class _ModuleProxy:
-    @property
-    def models_dir(self):
-        base = _find_comfyui_base()
-        return str(base / "models") if base else None
-
-    def get_output_directory(self):
-        base = _find_comfyui_base()
-        return str(base / "output") if base else None
-
-    def get_input_directory(self):
-        base = _find_comfyui_base()
-        return str(base / "input") if base else None
-
-    def get_temp_directory(self):
-        base = _find_comfyui_base()
-        return str(base / "temp") if base else None
-
-# Replace module with proxy instance
-import sys
-sys.modules[__name__] = _ModuleProxy()

comfy_env/workers/pool.py

@@ -1,241 +0,0 @@
-"""
-WorkerPool - Global registry and management of named workers.
-
-Provides a simple API for getting workers by name:
-
-    from comfy_env.workers import get_worker
-
-    worker = get_worker("sam3d")
-    result = worker.call_module("my_module", "my_func", image=tensor)
-
-Workers are registered at startup and reused across calls:
-
-    from comfy_env.workers import register_worker, TorchMPWorker
-
-    register_worker("default", TorchMPWorker())
-    register_worker("sam3d", PersistentVenvWorker(
-        python="/path/to/venv/bin/python",
-        working_dir="/path/to/nodes",
-    ))
-"""
-
-import atexit
-import threading
-from typing import Dict, Optional, Union
-from pathlib import Path
-
-from .base import Worker
-
-
-class WorkerPool:
-    """
-    Singleton pool of named workers.
-
-    Manages worker lifecycle, provides access by name, handles cleanup.
-    """
-
-    _instance: Optional["WorkerPool"] = None
-    _lock = threading.Lock()
-
-    def __new__(cls):
-        if cls._instance is None:
-            with cls._lock:
-                if cls._instance is None:
-                    cls._instance = super().__new__(cls)
-                    cls._instance._initialized = False
-        return cls._instance
-
-    def __init__(self):
-        if self._initialized:
-            return
-        self._initialized = True
-        self._workers: Dict[str, Worker] = {}
-        self._factories: Dict[str, callable] = {}
-        self._worker_lock = threading.Lock()
-
-    def register(
-        self,
-        name: str,
-        worker: Optional[Worker] = None,
-        factory: Optional[callable] = None,
-    ) -> None:
-        """
-        Register a worker or worker factory.
-
-        Args:
-            name: Name to register under.
-            worker: Pre-created worker instance.
-            factory: Callable that creates worker on first use (lazy).
-
-        Only one of worker or factory should be provided.
-        """
-        if worker is not None and factory is not None:
-            raise ValueError("Provide either worker or factory, not both")
-        if worker is None and factory is None:
-            raise ValueError("Must provide worker or factory")
-
-        with self._worker_lock:
-            # Shutdown existing worker if replacing
-            if name in self._workers:
-                try:
-                    self._workers[name].shutdown()
-                except:
-                    pass
-
-            if worker is not None:
-                self._workers[name] = worker
-                self._factories.pop(name, None)
-            else:
-                self._factories[name] = factory
-                self._workers.pop(name, None)
-
-    def get(self, name: str) -> Worker:
-        """
-        Get a worker by name.
-
-        Args:
-            name: Registered worker name.
-
-        Returns:
-            The worker instance.
-
-        Raises:
-            KeyError: If no worker registered with that name.
-        """
-        with self._worker_lock:
-            # Check for existing worker
-            if name in self._workers:
-                worker = self._workers[name]
-                if worker.is_alive():
-                    return worker
-                # Worker died, try to recreate from factory
-                if name not in self._factories:
-                    raise RuntimeError(f"Worker '{name}' died and no factory to recreate")
-
-            # Create from factory
-            if name in self._factories:
-                worker = self._factories[name]()
-                self._workers[name] = worker
-                return worker
-
-            raise KeyError(f"No worker registered with name: {name}")
-
-    def shutdown(self, name: Optional[str] = None) -> None:
-        """
-        Shutdown workers.
-
-        Args:
-            name: If provided, shutdown only this worker.
-                  If None, shutdown all workers.
-        """
-        with self._worker_lock:
-            if name is not None:
-                if name in self._workers:
-                    try:
-                        self._workers[name].shutdown()
-                    except:
-                        pass
-                    del self._workers[name]
-            else:
-                for worker in self._workers.values():
-                    try:
-                        worker.shutdown()
-                    except:
-                        pass
-                self._workers.clear()
-
-    def list_workers(self) -> Dict[str, str]:
-        """
-        List all registered workers.
-
-        Returns:
-            Dict of name -> status string.
-        """
-        with self._worker_lock:
-            result = {}
-            for name, worker in self._workers.items():
-                status = "alive" if worker.is_alive() else "dead"
-                result[name] = f"{type(worker).__name__} ({status})"
-            for name in self._factories:
-                if name not in result:
-                    result[name] = f"factory (not started)"
-            return result
-
-
-# Global pool instance
-_pool = WorkerPool()
-
-
-def get_worker(name: str) -> Worker:
-    """
-    Get a worker by name from the global pool.
-
-    Args:
-        name: Registered worker name.
-
-    Returns:
-        Worker instance.
-
-    Example:
-        worker = get_worker("sam3d")
-        result = worker.call_module("my_module", "my_func", image=tensor)
-    """
-    return _pool.get(name)
-
-
-def register_worker(
-    name: str,
-    worker: Optional[Worker] = None,
-    factory: Optional[callable] = None,
-) -> None:
-    """
-    Register a worker in the global pool.
-
-    Args:
-        name: Name to register under.
-        worker: Pre-created worker instance.
-        factory: Callable that creates worker on demand.
-
-    Example:
-        # Register pre-created worker
-        register_worker("default", TorchMPWorker())
-
-        # Register factory for lazy creation
-        register_worker("sam3d", factory=lambda: PersistentVenvWorker(
-            python="/path/to/venv/bin/python",
-        ))
-    """
-    _pool.register(name, worker=worker, factory=factory)
-
-
-def shutdown_workers(name: Optional[str] = None) -> None:
-    """
-    Shutdown workers in the global pool.
-
-    Args:
-        name: If provided, shutdown only this worker.
-              If None, shutdown all workers.
-    """
-    _pool.shutdown(name)
-
-
-def list_workers() -> Dict[str, str]:
-    """
-    List all registered workers.
-
-    Returns:
-        Dict of name -> status description.
-    """
-    return _pool.list_workers()
-
-
-# Register default worker (TorchMPWorker) on import
-def _register_default():
-    from .torch_mp import TorchMPWorker
-    register_worker("default", factory=lambda: TorchMPWorker(name="default"))
-
-
-_register_default()
-
-# Cleanup on exit
-atexit.register(lambda: shutdown_workers())