comfy-env 0.1.12__py3-none-any.whl → 0.1.13__py3-none-any.whl

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}"
 

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}"
 

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}"
 

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

{comfy_env-0.1.12.dist-info → comfy_env-0.1.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comfy-env
-Version: 0.1.12
+Version: 0.1.13
 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.dist-info → comfy_env-0.1.13.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=tUbsUdjnJCUUoxM7NpsdUuawuIz1exfOmWdsLGILXiY,5391
-comfy_env/prestartup.py,sha256=aKuW07R0CmoeIXukENJZnGMPzCv7kU7azvHgIOt2usk,5813
+comfy_env/prestartup.py,sha256=aNiTgkDeTdez7bsX0Ew3AoPoUKiTM3fmjjhCGWzZbLY,5812
 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=cTRzbdHJncmg4yUAyH5n9TPBu88grv6D1XezZwXKpok,14638
+comfy_env/isolation/wrap.py,sha256=A6kv4p4xBxx9-1yCNDKuKJi6esnT3hF2Y5wqLQZB_ig,14637
 comfy_env/pixi/__init__.py,sha256=BUrq7AQf3WDm0cHWh72B2xZbURNnDu2dCuELWiQCUiM,997
-comfy_env/pixi/core.py,sha256=fqozZZEoub8dAQaE6v6gyqw1mMiAYiCcdzH0gjCANBo,20637
+comfy_env/pixi/core.py,sha256=ZpjAmwh7EE1jqBK1Q41GJEEXIsAtyqdNsr_ygoPQdUY,20636
 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=vsxDGWepmSNgfqBhZPW7h8yOiKEyQcDFYP09masLTV4,32337
-comfy_env/workers/subprocess.py,sha256=bNwHpgz_EIIwZVRlgsx_ZEowYRrWbp-8uUIboacc-7M,57136
+comfy_env/workers/mp.py,sha256=R0XWsiHv8gswxa_-iNHU14o_9Og0RFG0QnY9DRZzn2c,34060
+comfy_env/workers/subprocess.py,sha256=B0PsHLuywPIJSRpuA4E8Dg6tGzuro6wqIr8czqMqtPE,57128
 comfy_env/workers/tensor_utils.py,sha256=TCuOAjJymrSbkgfyvcKtQ_KbVWTqSwP9VH_bCaFLLq8,6409
-comfy_env-0.1.12.dist-info/METADATA,sha256=J41pwCOTdj_2zckDDKjse-HfGWErDaDX3Hire4lKwgM,6971
-comfy_env-0.1.12.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-comfy_env-0.1.12.dist-info/entry_points.txt,sha256=J4fXeqgxU_YenuW_Zxn_pEL7J-3R0--b6MS5t0QmAr0,49
-comfy_env-0.1.12.dist-info/licenses/LICENSE,sha256=E68QZMMpW4P2YKstTZ3QU54HRQO8ecew09XZ4_Vn870,1093
-comfy_env-0.1.12.dist-info/RECORD,,
+comfy_env-0.1.13.dist-info/METADATA,sha256=xBivq9U8DeIzi-fwdAHZ1GjSyG-p724kEg42FqMdSlc,6971
+comfy_env-0.1.13.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+comfy_env-0.1.13.dist-info/entry_points.txt,sha256=J4fXeqgxU_YenuW_Zxn_pEL7J-3R0--b6MS5t0QmAr0,49
+comfy_env-0.1.13.dist-info/licenses/LICENSE,sha256=E68QZMMpW4P2YKstTZ3QU54HRQO8ecew09XZ4_Vn870,1093
+comfy_env-0.1.13.dist-info/RECORD,,