comfy-env 0.1.2__tar.gz → 0.1.3__tar.gz

{comfy_env-0.1.2 → comfy_env-0.1.3}/PKG-INFO

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

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

{comfy_env-0.1.2 → comfy_env-0.1.3}/src/comfy_env/isolation/wrap.py

@@ -67,6 +67,7 @@ def _get_worker(
     working_dir: Path,
     sys_path: list[str],
     lib_path: Optional[str] = None,
+    env_vars: Optional[dict] = None,
 ):
     """Get or create a persistent worker for the isolated environment."""
     cache_key = str(env_dir)
@@ -106,11 +107,14 @@ def _get_worker(
 
             print(f"[comfy-env] Starting isolated worker (MPWorker)")
             print(f"[comfy-env]   Env: {env_dir}")
+            if env_vars:
+                print(f"[comfy-env]   env_vars: {', '.join(f'{k}={v}' for k, v in env_vars.items())}")
 
             worker = MPWorker(
                 name=working_dir.name,
                 sys_path=sys_path,
                 lib_path=lib_path,
+                env_vars=env_vars,
             )
 
         _workers[cache_key] = worker
@@ -187,6 +191,7 @@ def _wrap_node_class(
     working_dir: Path,
     sys_path: list[str],
     lib_path: Optional[str] = None,
+    env_vars: Optional[dict] = None,
 ) -> type:
     """
     Wrap a node class so its FUNCTION method runs in the isolated environment.
@@ -232,7 +237,7 @@ def _wrap_node_class(
             print(f"[comfy-env] PROXY CALLED: {cls.__name__}.{func_name}", flush=True)
             print(f"[comfy-env]   kwargs keys: {list(kwargs.keys())}", flush=True)
 
-        worker = _get_worker(env_dir, working_dir, sys_path, lib_path)
+        worker = _get_worker(env_dir, working_dir, sys_path, lib_path, env_vars)
         if _DEBUG:
             print(f"[comfy-env]   worker alive: {worker.is_alive()}", flush=True)
 
@@ -322,6 +327,20 @@ def wrap_isolated_nodes(
         print(f"[comfy-env] Warning: No comfy-env.toml in {nodes_dir}")
         return node_class_mappings
 
+    # Read env_vars from comfy-env.toml
+    env_vars = {}
+    try:
+        if sys.version_info >= (3, 11):
+            import tomllib
+        else:
+            import tomli as tomllib
+        with open(config_file, "rb") as f:
+            config = tomllib.load(f)
+        env_vars_data = config.get("env_vars", {})
+        env_vars = {str(k): str(v) for k, v in env_vars_data.items()}
+    except Exception:
+        pass  # Ignore errors reading config
+
     # Find environment directory and paths
     env_dir = _find_env_dir(nodes_dir)
     site_packages, lib_dir = _find_env_paths(nodes_dir)
@@ -342,10 +361,12 @@ def wrap_isolated_nodes(
     print(f"[comfy-env]   site-packages: {site_packages}")
     if lib_path:
         print(f"[comfy-env]   lib: {lib_path}")
+    if env_vars:
+        print(f"[comfy-env]   env_vars: {', '.join(f'{k}={v}' for k, v in env_vars.items())}")
 
     # Wrap all node classes
     for node_name, node_cls in node_class_mappings.items():
         if hasattr(node_cls, "FUNCTION"):
-            _wrap_node_class(node_cls, env_dir, nodes_dir, sys_path, lib_path)
+            _wrap_node_class(node_cls, env_dir, nodes_dir, sys_path, lib_path, env_vars)
 
     return node_class_mappings

{comfy_env-0.1.2 → comfy_env-0.1.3}/src/comfy_env/workers/mp.py

@@ -40,7 +40,7 @@ _SHUTDOWN = object()
 _CALL_METHOD = "call_method"
 
 
-def _worker_loop(queue_in, queue_out, sys_path_additions=None, lib_path=None):
+def _worker_loop(queue_in, queue_out, sys_path_additions=None, lib_path=None, env_vars=None):
     """
     Worker process main loop.
 
@@ -56,11 +56,16 @@ def _worker_loop(queue_in, queue_out, sys_path_additions=None, lib_path=None):
         queue_out: Output queue for sending results
         sys_path_additions: Paths to add to sys.path
         lib_path: Path to add to LD_LIBRARY_PATH (for conda libraries)
+        env_vars: Environment variables to set (from comfy-env.toml)
     """
     import os
     import sys
     from pathlib import Path
 
+    # Apply env_vars FIRST (before any library imports that might check them)
+    if env_vars:
+        os.environ.update(env_vars)
+
     # Set worker mode env var
     os.environ["COMFYUI_ISOLATION_WORKER"] = "1"
 
@@ -424,7 +429,7 @@ class MPWorker(Worker):
     interpreter without inherited state from the parent.
     """
 
-    def __init__(self, name: Optional[str] = None, sys_path: Optional[list] = None, lib_path: Optional[str] = None):
+    def __init__(self, name: Optional[str] = None, sys_path: Optional[list] = None, lib_path: Optional[str] = None, env_vars: Optional[dict] = None):
         """
         Initialize the worker.
 
@@ -432,10 +437,12 @@ class MPWorker(Worker):
             name: Optional name for logging/debugging.
             sys_path: Optional list of paths to add to sys.path in worker process.
             lib_path: Optional path to add to LD_LIBRARY_PATH (for conda libraries).
+            env_vars: Optional environment variables to set in worker process.
         """
         self.name = name or "MPWorker"
         self._sys_path = sys_path or []
         self._lib_path = lib_path
+        self._env_vars = env_vars or {}
         self._process = None
         self._queue_in = None
         self._queue_out = None
@@ -504,7 +511,7 @@ class MPWorker(Worker):
             self._queue_out = ctx.Queue()
             self._process = ctx.Process(
                 target=_worker_loop,
-                args=(self._queue_in, self._queue_out, self._sys_path, self._lib_path),
+                args=(self._queue_in, self._queue_out, self._sys_path, self._lib_path, self._env_vars),
                 daemon=True,
             )
             self._process.start()