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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comfy-env
-Version: 0.1.9
+Version: 0.1.11
 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.11}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "comfy-env"
-version = "0.1.9"
+version = "0.1.11"
 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.11}/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
@@ -185,6 +211,34 @@ def _find_env_dir(node_dir: Path) -> Optional[Path]:
     return None
 
 
+def _find_custom_node_root(nodes_dir: Path) -> Optional[Path]:
+    """
+    Find the custom node root (direct child of custom_nodes/).
+
+    Uses folder_paths to find custom_nodes directories, then finds
+    which one is an ancestor of nodes_dir.
+
+    Example: /path/custom_nodes/ComfyUI-UniRig/nodes/nodes_gpu
+             -> returns /path/custom_nodes/ComfyUI-UniRig
+    """
+    try:
+        import folder_paths
+        custom_nodes_dirs = folder_paths.get_folder_paths("custom_nodes")
+    except (ImportError, KeyError):
+        return None
+
+    for cn_dir in custom_nodes_dirs:
+        cn_path = Path(cn_dir)
+        try:
+            rel = nodes_dir.relative_to(cn_path)
+            if rel.parts:
+                return cn_path / rel.parts[0]
+        except ValueError:
+            continue
+
+    return None
+
+
 def _wrap_node_class(
     cls: type,
     env_dir: Path,
@@ -362,7 +416,9 @@ def wrap_isolated_nodes(
         print(f"[comfy-env] Run 'comfy-env install' in {nodes_dir}")
         return node_class_mappings
 
-    # Build sys.path for the worker - site-packages first, then node dir
+    # Build sys.path - site-packages first, then nodes_dir
+    # Note: isolated modules should use absolute imports (their dir is in sys.path)
+    # Relative imports would require importing parent package which may have host-only deps
     sys_path = [str(site_packages), str(nodes_dir)]
 
     # lib_dir for LD_LIBRARY_PATH (conda libraries)

{comfy_env-0.1.9 → comfy_env-0.1.11}/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:
@@ -441,11 +457,14 @@ def pixi_install(
     # Build pypi-dependencies section (CUDA packages excluded - installed separately)
     pypi_deps = pixi_data.get("pypi-dependencies", {})
 
-    # Add torch if we have CUDA packages
+    # Enforce torch version if we have CUDA packages (must match cuda_packages wheels)
     if cfg.has_cuda and torch_version:
         torch_major = torch_version.split(".")[0]
         torch_minor = int(torch_version.split(".")[1])
-        pypi_deps.setdefault("torch", f">={torch_version},<{torch_major}.{torch_minor + 1}")
+        required_torch = f">={torch_version},<{torch_major}.{torch_minor + 1}"
+        if "torch" in pypi_deps and pypi_deps["torch"] != required_torch:
+            log(f"Overriding torch={pypi_deps['torch']} with {required_torch} (required for cuda_packages)")
+        pypi_deps["torch"] = required_torch
 
     # NOTE: CUDA packages are NOT added here - they're installed with --no-deps after pixi
 
@@ -478,8 +497,16 @@ def pixi_install(
         if not python_path:
             raise RuntimeError("Could not find Python in pixi environment")
 
-        # Get Python version from the pixi environment
-        py_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+        # Get Python version from the pixi environment (not host Python)
+        result = subprocess.run(
+            [str(python_path), "-c", "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')"],
+            capture_output=True, text=True
+        )
+        if result.returncode == 0:
+            py_version = result.stdout.strip()
+        else:
+            py_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+            log(f"Warning: Could not detect pixi Python version, using host: {py_version}")
 
         for package in cfg.cuda_packages:
             # Find direct wheel URL (bypasses metadata validation)
@@ -507,5 +534,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.11}/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.11}/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.11}/src/comfy_env/workers/subprocess.py

@@ -221,28 +221,26 @@ 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
+    # trimesh.Trimesh → pickle → shared memory (preserves visual, metadata, normals)
     if t == 'Trimesh':
-        verts = np.ascontiguousarray(obj.vertices, dtype=np.float64)
-        faces = np.ascontiguousarray(obj.faces, dtype=np.int64)
-
-        v_block = shm.SharedMemory(create=True, size=verts.nbytes)
-        np.ndarray(verts.shape, verts.dtype, buffer=v_block.buf)[:] = verts
-        registry.append(v_block)
+        import pickle
+        mesh_bytes = pickle.dumps(obj)
 
-        f_block = shm.SharedMemory(create=True, size=faces.nbytes)
-        np.ndarray(faces.shape, faces.dtype, buffer=f_block.buf)[:] = faces
-        registry.append(f_block)
+        block = shm.SharedMemory(create=True, size=len(mesh_bytes))
+        block.buf[:len(mesh_bytes)] = mesh_bytes
+        registry.append(block)
 
         result = {
             "__shm_trimesh__": True,
-            "v_name": v_block.name, "v_shape": list(verts.shape),
-            "f_name": f_block.name, "f_shape": list(faces.shape),
+            "name": block.name,
+            "size": len(mesh_bytes),
         }
         visited[obj_id] = result
         return result
@@ -279,31 +277,28 @@ 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
+    # trimesh (pickled to preserve visual, metadata, normals)
     if "__shm_trimesh__" in obj:
-        import trimesh
-        v_block = shm.SharedMemory(name=obj["v_name"])
-        verts = np.ndarray(tuple(obj["v_shape"]), dtype=np.float64, buffer=v_block.buf).copy()
-        v_block.close()
-        if unlink:
-            v_block.unlink()
-
-        f_block = shm.SharedMemory(name=obj["f_name"])
-        faces = np.ndarray(tuple(obj["f_shape"]), dtype=np.int64, buffer=f_block.buf).copy()
-        f_block.close()
+        import pickle
+        block = shm.SharedMemory(name=obj["name"])
+        mesh_bytes = bytes(block.buf[:obj["size"]])
+        block.close()
         if unlink:
-            f_block.unlink()
-
-        return trimesh.Trimesh(vertices=verts, faces=faces, process=False)
+            block.unlink()
+        return pickle.loads(mesh_bytes)
 
     # regular dict - recurse
     return {k: _from_shm(v, unlink) for k, v in obj.items()}
@@ -420,42 +415,29 @@ faulthandler.enable(file=sys.stderr, all_threads=True)
 # Debug logging (set COMFY_ENV_DEBUG=1 to enable)
 _DEBUG = os.environ.get("COMFY_ENV_DEBUG", "").lower() in ("1", "true", "yes")
 
-# Pre-import bpy FIRST to avoid DLL conflicts with numpy/torch/MKL
-# bpy's DLLs must be loaded before other packages load conflicting versions
-try:
-    import bpy
-    if _DEBUG:
-        print("[worker] Pre-imported bpy successfully", file=sys.stderr, flush=True)
-except ImportError as e:
-    # bpy not available in this environment - that's fine
-    pass
-except Exception as e:
-    if _DEBUG:
-        print(f"[worker] bpy pre-import warning: {e}", file=sys.stderr, flush=True)
-
 # Watchdog: dump all thread stacks every 60 seconds to catch hangs
 import threading
 import tempfile as _tempfile
 _watchdog_log = os.path.join(_tempfile.gettempdir(), "comfy_worker_watchdog.log")
 def _watchdog():
     import time
-    import io
     tick = 0
     while True:
         time.sleep(60)
         tick += 1
-        # Capture stack dump to string
-        buf = io.StringIO()
-        faulthandler.dump_traceback(file=buf, all_threads=True)
-        dump = buf.getvalue()
-
-        # Write to file
+        # Dump to temp file first (faulthandler needs real file descriptor)
+        tmp_path = _watchdog_log + ".tmp"
+        with open(tmp_path, "w", encoding="utf-8") as tmp:
+            faulthandler.dump_traceback(file=tmp, all_threads=True)
+        with open(tmp_path, "r", encoding="utf-8") as tmp:
+            dump = tmp.read()
+
+        # Write to persistent log
         with open(_watchdog_log, "a", encoding="utf-8") as f:
             f.write(f"\\n=== WATCHDOG TICK {tick} ({time.strftime('%H:%M:%S')}) ===\\n")
             f.write(dump)
             f.write("=== END ===\\n")
             f.flush()
-            os.fsync(f.fileno())
 
         # Also print
         print(f"\\n=== WATCHDOG TICK {tick} ===", flush=True)
@@ -544,24 +526,23 @@ 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
 
+    # trimesh.Trimesh → pickle → shared memory (preserves visual, metadata, normals)
     if t == 'Trimesh':
-        verts = np.ascontiguousarray(obj.vertices, dtype=np.float64)
-        faces = np.ascontiguousarray(obj.faces, dtype=np.int64)
-
-        v_block = shm.SharedMemory(create=True, size=verts.nbytes)
-        np.ndarray(verts.shape, verts.dtype, buffer=v_block.buf)[:] = verts
-        registry.append(v_block)
+        import pickle
+        mesh_bytes = pickle.dumps(obj)
 
-        f_block = shm.SharedMemory(create=True, size=faces.nbytes)
-        np.ndarray(faces.shape, faces.dtype, buffer=f_block.buf)[:] = faces
-        registry.append(f_block)
+        block = shm.SharedMemory(create=True, size=len(mesh_bytes))
+        block.buf[:len(mesh_bytes)] = mesh_bytes
+        registry.append(block)
 
         result = {
             "__shm_trimesh__": True,
-            "v_name": v_block.name, "v_shape": list(verts.shape),
-            "f_name": f_block.name, "f_shape": list(faces.shape),
+            "name": block.name,
+            "size": len(mesh_bytes),
         }
         visited[obj_id] = result
         return result
@@ -587,18 +568,18 @@ 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
+    # trimesh (pickled to preserve visual, metadata, normals)
     if "__shm_trimesh__" in obj:
-        import trimesh
-        v_block = shm.SharedMemory(name=obj["v_name"])
-        verts = np.ndarray(tuple(obj["v_shape"]), dtype=np.float64, buffer=v_block.buf).copy()
-        v_block.close()
-
-        f_block = shm.SharedMemory(name=obj["f_name"])
-        faces = np.ndarray(tuple(obj["f_shape"]), dtype=np.int64, buffer=f_block.buf).copy()
-        f_block.close()
-
-        return trimesh.Trimesh(vertices=verts, faces=faces, process=False)
+        import pickle
+        block = shm.SharedMemory(name=obj["name"])
+        mesh_bytes = bytes(block.buf[:obj["size"]])
+        block.close()
+        return pickle.loads(mesh_bytes)
     return {k: _from_shm(v) for k, v in obj.items()}
 
 def _cleanup_shm(registry):