comfy-env 0.1.12__tar.gz → 0.1.14__tar.gz

{comfy_env-0.1.12 → comfy_env-0.1.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comfy-env
-Version: 0.1.12
+Version: 0.1.14
 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.12 → comfy_env-0.1.14}/pyproject.toml

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

{comfy_env-0.1.12 → comfy_env-0.1.14}/src/comfy_env/__init__.py

@@ -56,6 +56,15 @@ from .install import install, verify_installation
 # Prestartup helpers
 from .prestartup import setup_env
 
+# Cache management
+from .cache import (
+    get_cache_dir,
+    cleanup_orphaned_envs,
+    resolve_env_path,
+    CACHE_DIR,
+    MARKER_FILE,
+)
+
 # Errors
 from .errors import (
     EnvManagerError,
@@ -105,4 +114,20 @@ __all__ = [
     "DependencyError",
     "CUDANotFoundError",
     "InstallError",
+    # Cache
+    "get_cache_dir",
+    "cleanup_orphaned_envs",
+    "resolve_env_path",
+    "CACHE_DIR",
+    "MARKER_FILE",
 ]
+
+# Run orphan cleanup once on module load (silently)
+def _run_startup_cleanup():
+    """Clean orphaned envs on startup. Runs silently, never fails startup."""
+    try:
+        cleanup_orphaned_envs(log=lambda x: None)  # Silent
+    except Exception:
+        pass  # Never fail startup due to cleanup
+
+_run_startup_cleanup()

comfy_env-0.1.14/src/comfy_env/cache.py

@@ -0,0 +1,331 @@
+"""
+Central environment cache management for comfy-env.
+
+Stores environments in ~/.comfy-env/envs/<nodename>_<subfolder>_<hash>/
+to avoid Windows MAX_PATH (260 char) issues.
+
+Marker files in node folders link to central envs and enable orphan cleanup.
+"""
+
+import hashlib
+import os
+import shutil
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import Optional, Tuple, Callable
+
+# Import version
+try:
+    from . import __version__
+except ImportError:
+    __version__ = "0.0.0-dev"
+
+# Lazy import tomli/tomllib
+def _get_tomli():
+    if sys.version_info >= (3, 11):
+        import tomllib
+        return tomllib
+    else:
+        try:
+            import tomli
+            return tomli
+        except ImportError:
+            return None
+
+def _get_tomli_w():
+    try:
+        import tomli_w
+        return tomli_w
+    except ImportError:
+        return None
+
+
+# Constants
+CACHE_DIR = Path.home() / ".comfy-env" / "envs"
+MARKER_FILE = ".comfy-env-marker.toml"
+METADATA_FILE = ".comfy-env-metadata.toml"
+
+
+def get_cache_dir() -> Path:
+    """Get central cache directory, create if needed."""
+    CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    return CACHE_DIR
+
+
+def compute_config_hash(config_path: Path) -> str:
+    """Compute hash of comfy-env.toml content (first 8 chars of SHA256)."""
+    content = config_path.read_bytes()
+    return hashlib.sha256(content).hexdigest()[:8]
+
+
+def sanitize_name(name: str) -> str:
+    """Sanitize a name for use in filesystem paths."""
+    # Lowercase and replace problematic chars
+    name = name.lower()
+    for prefix in ("comfyui-", "comfyui_"):
+        if name.startswith(prefix):
+            name = name[len(prefix):]
+    return name.replace("-", "_").replace(" ", "_")
+
+
+def get_env_name(node_dir: Path, config_path: Path) -> str:
+    """
+    Generate env name: <nodename>_<subfolder>_<hash>
+
+    Args:
+        node_dir: The custom node directory (e.g., custom_nodes/ComfyUI-UniRig)
+        config_path: Path to comfy-env.toml
+
+    Examples:
+        - ComfyUI-UniRig/nodes/comfy-env.toml -> unirig_nodes_a1b2c3d4
+        - ComfyUI-Pack/comfy-env.toml -> pack__f5e6d7c8 (double underscore = no subfolder)
+    """
+    # Get node name
+    node_name = sanitize_name(node_dir.name)
+
+    # Get subfolder (relative path from node_dir to config parent)
+    config_parent = config_path.parent
+    if config_parent == node_dir:
+        subfolder = ""
+    else:
+        try:
+            rel_path = config_parent.relative_to(node_dir)
+            subfolder = rel_path.as_posix().replace("/", "_")
+        except ValueError:
+            # config_path not under node_dir - use parent folder name
+            subfolder = sanitize_name(config_parent.name)
+
+    # Compute hash
+    config_hash = compute_config_hash(config_path)
+
+    return f"{node_name}_{subfolder}_{config_hash}"
+
+
+def get_central_env_path(node_dir: Path, config_path: Path) -> Path:
+    """Get path to central environment for this config."""
+    env_name = get_env_name(node_dir, config_path)
+    return get_cache_dir() / env_name
+
+
+def write_marker(config_path: Path, env_path: Path) -> None:
+    """
+    Write marker file linking node to central env.
+
+    Args:
+        config_path: Path to comfy-env.toml
+        env_path: Path to central environment
+    """
+    tomli_w = _get_tomli_w()
+    if not tomli_w:
+        # Fallback to manual TOML writing
+        marker_path = config_path.parent / MARKER_FILE
+        content = f'''[env]
+name = "{env_path.name}"
+path = "{env_path}"
+config_hash = "{compute_config_hash(config_path)}"
+created = "{datetime.now().isoformat()}"
+comfy_env_version = "{__version__}"
+'''
+        marker_path.write_text(content)
+        return
+
+    marker_path = config_path.parent / MARKER_FILE
+    marker_data = {
+        "env": {
+            "name": env_path.name,
+            "path": str(env_path),
+            "config_hash": compute_config_hash(config_path),
+            "created": datetime.now().isoformat(),
+            "comfy_env_version": __version__,
+        }
+    }
+    marker_path.write_text(tomli_w.dumps(marker_data))
+
+
+def write_env_metadata(env_path: Path, marker_path: Path) -> None:
+    """
+    Write metadata file in env for orphan detection.
+
+    Args:
+        env_path: Path to central environment
+        marker_path: Path to marker file in node folder
+    """
+    tomli_w = _get_tomli_w()
+    metadata_path = env_path / METADATA_FILE
+
+    if not tomli_w:
+        # Fallback to manual TOML writing
+        content = f'''marker_path = "{marker_path}"
+created = "{datetime.now().isoformat()}"
+'''
+        metadata_path.write_text(content)
+        return
+
+    metadata = {
+        "marker_path": str(marker_path),
+        "created": datetime.now().isoformat(),
+    }
+    metadata_path.write_text(tomli_w.dumps(metadata))
+
+
+def read_marker(marker_path: Path) -> Optional[dict]:
+    """
+    Read marker file, return None if invalid/missing.
+
+    Args:
+        marker_path: Path to .comfy-env-marker.toml
+
+    Returns:
+        Parsed marker data or None
+    """
+    if not marker_path.exists():
+        return None
+
+    tomli = _get_tomli()
+    if not tomli:
+        return None
+
+    try:
+        with open(marker_path, "rb") as f:
+            return tomli.load(f)
+    except Exception:
+        return None
+
+
+def read_env_metadata(env_path: Path) -> Optional[dict]:
+    """
+    Read metadata file from env, return None if invalid/missing.
+
+    Args:
+        env_path: Path to central environment
+
+    Returns:
+        Parsed metadata or None
+    """
+    metadata_path = env_path / METADATA_FILE
+    if not metadata_path.exists():
+        return None
+
+    tomli = _get_tomli()
+    if not tomli:
+        return None
+
+    try:
+        with open(metadata_path, "rb") as f:
+            return tomli.load(f)
+    except Exception:
+        return None
+
+
+def resolve_env_path(node_dir: Path) -> Tuple[Optional[Path], Optional[Path], Optional[Path]]:
+    """
+    Resolve environment path with fallback chain.
+
+    Args:
+        node_dir: Directory containing comfy-env.toml or marker file
+
+    Returns:
+        (env_path, site_packages, lib_dir) or (None, None, None)
+
+    Fallback order:
+        1. Marker file -> central cache
+        2. _env_<name> (current location)
+        3. .pixi/envs/default (old pixi)
+        4. .venv (venv support)
+    """
+    # 1. Check marker file -> central cache
+    marker_path = node_dir / MARKER_FILE
+    marker = read_marker(marker_path)
+    if marker and "env" in marker:
+        env_path = Path(marker["env"]["path"])
+        if env_path.exists():
+            return _get_env_paths(env_path)
+
+    # 2. Check _env_<name>
+    node_name = sanitize_name(node_dir.name)
+    env_name = f"_env_{node_name}"
+    local_env = node_dir / env_name
+    if local_env.exists():
+        return _get_env_paths(local_env)
+
+    # 3. Check .pixi/envs/default
+    pixi_env = node_dir / ".pixi" / "envs" / "default"
+    if pixi_env.exists():
+        return _get_env_paths(pixi_env)
+
+    # 4. Check .venv
+    venv_dir = node_dir / ".venv"
+    if venv_dir.exists():
+        return _get_env_paths(venv_dir)
+
+    return None, None, None
+
+
+def _get_env_paths(env_path: Path) -> Tuple[Path, Optional[Path], Optional[Path]]:
+    """
+    Get site-packages and lib paths from an environment.
+
+    Args:
+        env_path: Path to environment root
+
+    Returns:
+        (env_path, site_packages, lib_dir)
+    """
+    import glob
+
+    if sys.platform == "win32":
+        site_packages = env_path / "Lib" / "site-packages"
+        lib_dir = env_path / "Library" / "bin"
+    else:
+        # Linux/Mac: lib/python*/site-packages
+        matches = glob.glob(str(env_path / "lib" / "python*" / "site-packages"))
+        site_packages = Path(matches[0]) if matches else None
+        lib_dir = env_path / "lib"
+
+    return env_path, site_packages, lib_dir
+
+
+def cleanup_orphaned_envs(log: Callable[[str], None] = print) -> int:
+    """
+    Scan central cache and remove orphaned environments.
+
+    An env is orphaned if its marker file no longer exists
+    (meaning the node was deleted).
+
+    Args:
+        log: Logging function
+
+    Returns:
+        Number of envs cleaned up
+    """
+    cache_dir = get_cache_dir()
+    if not cache_dir.exists():
+        return 0
+
+    cleaned = 0
+    for env_dir in cache_dir.iterdir():
+        if not env_dir.is_dir():
+            continue
+
+        # Skip if no metadata (might be manually created or old format)
+        metadata = read_env_metadata(env_dir)
+        if not metadata:
+            continue
+
+        # Check if marker file still exists
+        marker_path_str = metadata.get("marker_path", "")
+        if not marker_path_str:
+            continue
+
+        marker_path = Path(marker_path_str)
+        if not marker_path.exists():
+            # Marker gone = node was deleted = orphan
+            log(f"[comfy-env] Cleaning orphaned env: {env_dir.name}")
+            try:
+                shutil.rmtree(env_dir)
+                cleaned += 1
+            except Exception as e:
+                log(f"[comfy-env] Failed to cleanup {env_dir.name}: {e}")
+
+    return cleaned

{comfy_env-0.1.12 → comfy_env-0.1.14}/src/comfy_env/isolation/wrap.py

@@ -36,7 +36,7 @@ _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"""
+    """Convert directory name to env name: ComfyUI-UniRig -> _env_unirig"""
     name = dir_name.lower().replace("-", "_").lstrip("comfyui_")
     return f"_env_{name}"
 
@@ -145,15 +145,19 @@ def _find_env_paths(node_dir: Path) -> tuple[Optional[Path], Optional[Path]]:
     """
     Find site-packages and lib directories for the isolated environment.
 
+    Fallback order:
+        1. Marker file -> central cache
+        2. _env_<name> (local)
+        3. .pixi/envs/default (old pixi)
+        4. .venv
+
     Returns:
         (site_packages, lib_dir) - lib_dir is for LD_LIBRARY_PATH
     """
     import glob
 
-    # 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():
+    def _get_paths_from_env(env_dir: Path) -> tuple[Optional[Path], Optional[Path]]:
+        """Extract site-packages and lib_dir from an env directory."""
         if sys.platform == "win32":
             site_packages = env_dir / "Lib" / "site-packages"
             lib_dir = env_dir / "Library" / "bin"
@@ -163,23 +167,45 @@ def _find_env_paths(node_dir: Path) -> tuple[Optional[Path], Optional[Path]]:
             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
+            return site_packages, lib_dir if lib_dir and lib_dir.exists() else None
+        return None, None
 
-    # Fallback: Check old .pixi/envs/default (for backward compat)
+    # 1. Check marker file -> central cache
+    marker_path = node_dir / ".comfy-env-marker.toml"
+    if marker_path.exists():
+        try:
+            if sys.version_info >= (3, 11):
+                import tomllib
+            else:
+                import tomli as tomllib
+            with open(marker_path, "rb") as f:
+                marker = tomllib.load(f)
+            env_path = marker.get("env", {}).get("path")
+            if env_path:
+                env_dir = Path(env_path)
+                if env_dir.exists():
+                    result = _get_paths_from_env(env_dir)
+                    if result[0]:
+                        return result
+        except Exception:
+            pass  # Fall through to other options
+
+    # 2. Check _env_<name> directory (local)
+    env_name = get_env_name(node_dir.name)
+    env_dir = node_dir / env_name
+    if env_dir.exists():
+        result = _get_paths_from_env(env_dir)
+        if result[0]:
+            return result
+
+    # 3. Fallback: Check old .pixi/envs/default (for backward compat)
     pixi_env = node_dir / ".pixi" / "envs" / "default"
     if pixi_env.exists():
-        if sys.platform == "win32":
-            site_packages = pixi_env / "Lib" / "site-packages"
-            lib_dir = pixi_env / "Library" / "bin"
-        else:
-            pattern = str(pixi_env / "lib" / "python*" / "site-packages")
-            matches = glob.glob(pattern)
-            site_packages = Path(matches[0]) if matches else None
-            lib_dir = pixi_env / "lib"
-        if site_packages and site_packages.exists():
-            return site_packages, lib_dir if lib_dir.exists() else None
+        result = _get_paths_from_env(pixi_env)
+        if result[0]:
+            return result
 
-    # Check .venv directory
+    # 4. Check .venv directory
     venv_dir = node_dir / ".venv"
     if venv_dir.exists():
         if sys.platform == "win32":
@@ -195,19 +221,49 @@ 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
+    """
+    Find the environment directory (for cache key).
+
+    Fallback order:
+        1. Marker file -> central cache
+        2. _env_<name> (local)
+        3. .pixi/envs/default (old pixi)
+        4. .venv
+    """
+    # 1. Check marker file -> central cache
+    marker_path = node_dir / ".comfy-env-marker.toml"
+    if marker_path.exists():
+        try:
+            if sys.version_info >= (3, 11):
+                import tomllib
+            else:
+                import tomli as tomllib
+            with open(marker_path, "rb") as f:
+                marker = tomllib.load(f)
+            env_path = marker.get("env", {}).get("path")
+            if env_path:
+                env_dir = Path(env_path)
+                if env_dir.exists():
+                    return env_dir
+        except Exception:
+            pass
+
+    # 2. 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
+
+    # 3. Fallback to old .pixi path
     pixi_env = node_dir / ".pixi" / "envs" / "default"
     if pixi_env.exists():
         return pixi_env
+
+    # 4. Check .venv
     venv_dir = node_dir / ".venv"
     if venv_dir.exists():
         return venv_dir
+
     return None
 
 

{comfy_env-0.1.12 → comfy_env-0.1.14}/src/comfy_env/pixi/core.py

@@ -282,7 +282,7 @@ def ensure_pixi(
 
 
 def get_env_name(dir_name: str) -> str:
-    """Convert directory name to env name: ComfyUI-UniRig → _env_unirig"""
+    """Convert directory name to env name: ComfyUI-UniRig -> _env_unirig"""
     name = dir_name.lower().replace("-", "_").lstrip("comfyui_")
     return f"_env_{name}"
 
@@ -534,20 +534,55 @@ def pixi_install(
 
         log("CUDA packages installed")
 
-    # Move environment from .pixi/envs/default to _env_<name>
+    # Move environment from .pixi/envs/default to central cache
+    from ..cache import (
+        get_central_env_path, write_marker, write_env_metadata,
+        MARKER_FILE, get_cache_dir
+    )
+
     old_env = node_dir / ".pixi" / "envs" / "default"
-    env_name = get_env_name(node_dir.name)
-    new_env = node_dir / env_name
+    config_path = node_dir / "comfy-env.toml"
+
+    # Determine the main node directory (for naming)
+    # If node_dir is custom_nodes/NodeName/subdir, main_node_dir is custom_nodes/NodeName
+    # If node_dir is custom_nodes/NodeName, main_node_dir is custom_nodes/NodeName
+    if node_dir.parent.name == "custom_nodes":
+        main_node_dir = node_dir
+    else:
+        # Walk up to find custom_nodes parent
+        main_node_dir = node_dir
+        for parent in node_dir.parents:
+            if parent.parent.name == "custom_nodes":
+                main_node_dir = parent
+                break
+
+    # Get central env path
+    central_env = get_central_env_path(main_node_dir, config_path)
 
     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)
+        # Ensure cache directory exists
+        get_cache_dir()
+
+        # Remove old central env if exists
+        if central_env.exists():
+            shutil.rmtree(central_env)
+
+        # Move to central cache
+        shutil.move(str(old_env), str(central_env))
+
+        # Write marker file in node directory
+        write_marker(config_path, central_env)
+
+        # Write metadata in env for orphan detection
+        marker_path = config_path.parent / MARKER_FILE
+        write_env_metadata(central_env, marker_path)
+
+        # Clean up .pixi directory
         pixi_dir = node_dir / ".pixi"
         if pixi_dir.exists():
             shutil.rmtree(pixi_dir)
-        log(f"Moved environment to {new_env}")
+
+        log(f"Environment created at: {central_env}")
 
     log("Installation complete!")
     return True

{comfy_env-0.1.12 → comfy_env-0.1.14}/src/comfy_env/prestartup.py

@@ -12,7 +12,7 @@ from typing import Optional, Dict
 
 
 def get_env_name(dir_name: str) -> str:
-    """Convert directory name to env name: ComfyUI-UniRig → _env_unirig"""
+    """Convert directory name to env name: ComfyUI-UniRig -> _env_unirig"""
     name = dir_name.lower().replace("-", "_").lstrip("comfyui_")
     return f"_env_{name}"
 
@@ -127,15 +127,43 @@ def setup_env(node_dir: Optional[str] = None) -> None:
     for key, value in env_vars.items():
         os.environ[key] = value
 
-    # 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):
-        # 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
+    # Resolve environment path with fallback chain:
+    # 1. Marker file -> central cache
+    # 2. _env_<name> (current local)
+    # 3. .pixi/envs/default (old pixi)
+    pixi_env = None
+
+    # 1. Check marker file -> central cache
+    marker_path = os.path.join(node_dir, ".comfy-env-marker.toml")
+    if os.path.exists(marker_path):
+        try:
+            if sys.version_info >= (3, 11):
+                import tomllib
+            else:
+                import tomli as tomllib
+            with open(marker_path, "rb") as f:
+                marker = tomllib.load(f)
+            env_path = marker.get("env", {}).get("path")
+            if env_path and os.path.exists(env_path):
+                pixi_env = env_path
+        except Exception:
+            pass  # Fall through to other options
+
+    # 2. Check _env_<name> (local)
+    if not pixi_env:
+        env_name = get_env_name(os.path.basename(node_dir))
+        local_env = os.path.join(node_dir, env_name)
+        if os.path.exists(local_env):
+            pixi_env = local_env
+
+    # 3. Fallback to old .pixi path
+    if not pixi_env:
+        old_pixi = os.path.join(node_dir, ".pixi", "envs", "default")
+        if os.path.exists(old_pixi):
+            pixi_env = old_pixi
+
+    if not pixi_env:
+        return  # No environment found
 
     if sys.platform == "win32":
         # Windows: add to PATH for DLL loading

{comfy_env-0.1.12 → comfy_env-0.1.14}/src/comfy_env/workers/mp.py

@@ -29,6 +29,7 @@ from queue import Empty as QueueEmpty
 from typing import Any, Callable, Optional
 
 from .base import Worker, WorkerError
+from .tensor_utils import prepare_for_ipc_recursive, keep_tensors_recursive
 
 logger = logging.getLogger("comfy_env")
 
@@ -40,8 +41,20 @@ _SHUTDOWN = object()
 _CALL_METHOD = "call_method"
 
 
+def _can_use_cuda_ipc():
+    """
+    Check if CUDA IPC is available.
+
+    CUDA IPC works with native allocator but breaks with cudaMallocAsync.
+    If no backend is specified, CUDA IPC should work (PyTorch default is native).
+    """
+    import os
+    conf = os.environ.get('PYTORCH_CUDA_ALLOC_CONF', '')
+    return 'cudaMallocAsync' not in conf
+
+
 # ---------------------------------------------------------------------------
-# Tensor file transfer - avoids CUDA IPC issues with cudaMallocAsync
+# Tensor file transfer - fallback for cudaMallocAsync (CUDA IPC doesn't work)
 # ---------------------------------------------------------------------------
 
 def _save_tensors_to_files(obj, file_registry=None):
@@ -252,20 +265,31 @@ 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)
+                    # Load tensors from files if using file-based transfer
+                    if not _can_use_cuda_ipc():
+                        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)
+                    # Handle result based on allocator
+                    if _can_use_cuda_ipc():
+                        keep_tensors_recursive(result)
+                    else:
+                        result = _save_tensors_to_files(result)
                     queue_out.put(("ok", result))
                 else:
                     # Direct function call (legacy)
                     func, args, kwargs = item
+                    # Load tensors from files if using file-based transfer
+                    if not _can_use_cuda_ipc():
+                        args = tuple(_load_tensors_from_files(a) for a in args)
+                        kwargs = _load_tensors_from_files(kwargs)
                     result = func(*args, **kwargs)
-                    # Save tensors to files to avoid CUDA IPC issues with cudaMallocAsync
-                    result = _save_tensors_to_files(result)
+                    # Handle result based on allocator
+                    if _can_use_cuda_ipc():
+                        keep_tensors_recursive(result)
+                    else:
+                        result = _save_tensors_to_files(result)
                     queue_out.put(("ok", result))
 
             except Exception as e:
@@ -661,6 +685,16 @@ class MPWorker(Worker):
         """
         self._ensure_started()
 
+        # Handle tensors based on allocator
+        if _can_use_cuda_ipc():
+            # CUDA IPC - zero copy (works with native allocator)
+            kwargs = {k: prepare_for_ipc_recursive(v) for k, v in kwargs.items()}
+            args = tuple(prepare_for_ipc_recursive(a) for a in args)
+        else:
+            # File-based transfer (fallback for cudaMallocAsync)
+            kwargs = _save_tensors_to_files(kwargs)
+            args = tuple(_save_tensors_to_files(a) for a in args)
+
         # Send work item
         self._queue_in.put((func, args, kwargs))
 
@@ -699,8 +733,13 @@ class MPWorker(Worker):
         """
         self._ensure_started()
 
-        # Save tensors to files to avoid CUDA IPC issues with cudaMallocAsync
-        kwargs = _save_tensors_to_files(kwargs)
+        # Handle tensors based on allocator
+        if _can_use_cuda_ipc():
+            # CUDA IPC - zero copy (works with native allocator)
+            kwargs = prepare_for_ipc_recursive(kwargs)
+        else:
+            # File-based transfer (fallback for cudaMallocAsync)
+            kwargs = _save_tensors_to_files(kwargs)
 
         # Send method call request using protocol
         self._queue_in.put((
@@ -728,8 +767,9 @@ class MPWorker(Worker):
 
         # Handle response
         if status == "ok":
-            # Load tensors from temp files
-            result = _load_tensors_from_files(result)
+            # Load tensors from temp files if using file-based transfer
+            if not _can_use_cuda_ipc():
+                result = _load_tensors_from_files(result)
             return result
         elif status == "error":
             msg, tb = result

{comfy_env-0.1.12 → comfy_env-0.1.14}/src/comfy_env/workers/subprocess.py

@@ -211,7 +211,7 @@ def _to_shm(obj, registry, visited=None):
 
     t = type(obj).__name__
 
-    # numpy array → direct shared memory
+    # numpy array -> direct shared memory
     if t == 'ndarray':
         arr = np.ascontiguousarray(obj)
         block = shm.SharedMemory(create=True, size=arr.nbytes)
@@ -221,14 +221,14 @@ def _to_shm(obj, registry, visited=None):
         visited[obj_id] = result
         return result
 
-    # torch.Tensor → convert to numpy → shared memory (with marker to restore type)
+    # torch.Tensor -> convert to numpy -> shared memory (with marker to restore type)
     if t == 'Tensor':
         arr = obj.detach().cpu().numpy()
         result = _to_shm(arr, registry, visited)
         result["__was_tensor__"] = True
         return result
 
-    # trimesh.Trimesh → pickle → shared memory (preserves visual, metadata, normals)
+    # trimesh.Trimesh -> pickle -> shared memory (preserves visual, metadata, normals)
     if t == 'Trimesh':
         import pickle
         mesh_bytes = pickle.dumps(obj)
@@ -245,7 +245,7 @@ def _to_shm(obj, registry, visited=None):
         visited[obj_id] = result
         return result
 
-    # Path → string
+    # Path -> string
     from pathlib import PurePath
     if isinstance(obj, PurePath):
         return str(obj)
@@ -530,7 +530,7 @@ def _to_shm(obj, registry, visited=None):
         result["__was_tensor__"] = True
         return result
 
-    # trimesh.Trimesh → pickle → shared memory (preserves visual, metadata, normals)
+    # trimesh.Trimesh -> pickle -> shared memory (preserves visual, metadata, normals)
     if t == 'Trimesh':
         import pickle
         mesh_bytes = pickle.dumps(obj)