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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comfy-env
-Version: 0.1.9
+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

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

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

{comfy_env-0.1.9 → 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

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

@@ -281,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"]:
@@ -289,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:
@@ -507,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.9 → 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.9 → 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.9 → 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