comfy-env 0.1.9__py3-none-any.whl → 0.1.10__py3-none-any.whl

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/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/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/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/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

{comfy_env-0.1.9.dist-info → comfy_env-0.1.10.dist-info}/METADATA

@@ -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.dist-info → comfy_env-0.1.10.dist-info}/RECORD

@@ -3,14 +3,14 @@ comfy_env/cli.py,sha256=ty4HYlzollCUCS0o6Sha6eczPAsW_gHRVgvck3IfA2w,12723
 comfy_env/errors.py,sha256=q-C3vyrPa_kk_Ao8l17mIGfJiG2IR0hCFV0GFcNLmcI,9924
 comfy_env/install.py,sha256=N7eBj8wB2DrGepVYk-Hks2mSf6UuGzj34pfVLNYJgQ4,10357
 comfy_env/nodes.py,sha256=nR-LctovXYTqk-seHovq9W_Ll0hlLeGMw-SQsa6IkBM,5235
-comfy_env/prestartup.py,sha256=NML_nzg4h37YdEg7gHxA7ke_uXWZsldz-X3HldHEUsE,5338
+comfy_env/prestartup.py,sha256=aKuW07R0CmoeIXukENJZnGMPzCv7kU7azvHgIOt2usk,5813
 comfy_env/config/__init__.py,sha256=4Guylkb-FV8QxhFwschzpzbr2eu8y-KNgNT3_JOm9jc,403
 comfy_env/config/parser.py,sha256=dA1lX5ExBEfCqUJwe4V5i_jn2NJ69bMq3c3ji3lMSV8,4295
 comfy_env/config/types.py,sha256=Sb8HO34xsSZu5YAc2K4M7Gb3QNevJlngf12hHiwuU0w,2140
 comfy_env/isolation/__init__.py,sha256=vw9a4mpJ2CFjy-PLe_A3zQ6umBQklgqWNxwn9beNw3g,175
-comfy_env/isolation/wrap.py,sha256=fYsf3tWpzNPoZT7mvWHbD-LmgBMhMCIfzlfcOFgwhlU,12641
+comfy_env/isolation/wrap.py,sha256=9pxlrO2qUP-4Aov9LUXhgqkDw8V4FT-M-YPeDRXUjuE,13664
 comfy_env/pixi/__init__.py,sha256=BUrq7AQf3WDm0cHWh72B2xZbURNnDu2dCuELWiQCUiM,997
-comfy_env/pixi/core.py,sha256=zPEtuF8D2gnIaljmGEGT8Ttw0lTy_cV-hIIn1fvlbBU,18788
+comfy_env/pixi/core.py,sha256=jy1CAaryPlldCm9bGq6i2UnPckYqyqzQELSDhQsB8ho,19958
 comfy_env/pixi/cuda_detection.py,sha256=sqB3LjvGNdV4eFqiARQGfyecBM3ZiUmeh6nG0YCRYQw,9751
 comfy_env/pixi/resolver.py,sha256=U_A8rBDxCj4gUlJt2YJQniP4cCKqxJEiVFgXOoH7vM8,6339
 comfy_env/pixi/platform/__init__.py,sha256=Nb5MPZIEeanSMEWwqU4p4bnEKTJn1tWcwobnhq9x9IY,614
@@ -22,11 +22,11 @@ comfy_env/templates/comfy-env-instructions.txt,sha256=ve1RAthW7ouumU9h6DM7mIRX1M
 comfy_env/templates/comfy-env.toml,sha256=ROIqi4BlPL1MEdL1VgebfTHpdwPNYGHwWeigI9Kw-1I,4831
 comfy_env/workers/__init__.py,sha256=TMVG55d2XLP1mJ3x1d16H0SBDJZtk2kMC5P4HLk9TrA,1073
 comfy_env/workers/base.py,sha256=4ZYTaQ4J0kBHCoO_OfZnsowm4rJCoqinZUaOtgkOPbw,2307
-comfy_env/workers/mp.py,sha256=8Xqcr4R4KlHtP1P_UDAtSTPHA_5i-ySTlxxLpqyOk-w,29757
-comfy_env/workers/subprocess.py,sha256=a9gpaTppbL7ZRgv4grK__QioGUiN0Ci9RIZGTVTsQ78,57993
+comfy_env/workers/mp.py,sha256=vsxDGWepmSNgfqBhZPW7h8yOiKEyQcDFYP09masLTV4,32337
+comfy_env/workers/subprocess.py,sha256=neggSmCqJxfiIOCUQwn_TrDONFJoT5veBZuMjDYf_LA,58537
 comfy_env/workers/tensor_utils.py,sha256=TCuOAjJymrSbkgfyvcKtQ_KbVWTqSwP9VH_bCaFLLq8,6409
-comfy_env-0.1.9.dist-info/METADATA,sha256=qplbcTtKJj8y2m8rGUpIMcZyZymKdqA7wQNQf4t5HsI,6970
-comfy_env-0.1.9.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-comfy_env-0.1.9.dist-info/entry_points.txt,sha256=J4fXeqgxU_YenuW_Zxn_pEL7J-3R0--b6MS5t0QmAr0,49
-comfy_env-0.1.9.dist-info/licenses/LICENSE,sha256=E68QZMMpW4P2YKstTZ3QU54HRQO8ecew09XZ4_Vn870,1093
-comfy_env-0.1.9.dist-info/RECORD,,
+comfy_env-0.1.10.dist-info/METADATA,sha256=tfO_Jhe-jz6XY_AHhfnHroJAVuAzDAe4evf9FYY3nw0,6971
+comfy_env-0.1.10.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+comfy_env-0.1.10.dist-info/entry_points.txt,sha256=J4fXeqgxU_YenuW_Zxn_pEL7J-3R0--b6MS5t0QmAr0,49
+comfy_env-0.1.10.dist-info/licenses/LICENSE,sha256=E68QZMMpW4P2YKstTZ3QU54HRQO8ecew09XZ4_Vn870,1093
+comfy_env-0.1.10.dist-info/RECORD,,