From e7ab67733bbcec96bfcd82fd514b9aed7e0a9235 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Fri, 3 Apr 2026 21:54:51 +0000 Subject: [PATCH] Add ComfyUI Desktop worker support --- README.md | 8 +- tests/test_worker_process_runtime.py | 134 +++++++++++++++++++++++++++ workers/process/launch_builder.py | 71 +++++++++++++- workers/process/lifecycle.py | 1 + workers/process/root_discovery.py | 29 +++++- 5 files changed, 232 insertions(+), 11 deletions(-) create mode 100644 tests/test_worker_process_runtime.py diff --git a/README.md b/README.md index 465bf19..ebe5559 100644 --- a/README.md +++ b/README.md @@ -194,10 +194,10 @@ No, it does not combine VRAM of multiple GPUs. No, it does not speed up the generation of a single image or video. Instead, it enables the generation of more images or videos simultaneously. However, it can speed up the upscaling of a single image when using the Ultimate SD Upscale Distributed feature. -
-Does it work with the ComfyUI desktop app? -Currently, it is not compatible with the ComfyUI desktop app. -
+
+Does it work with the ComfyUI desktop app? +Yes. Local managed workers now inherit the active ComfyUI runtime paths, so ComfyUI Desktop installs can launch workers with the same base, user, input, output, and frontend directories as the master. +
Can I combine my RTX 5090 with a GTX 980 to get faster results? diff --git a/tests/test_worker_process_runtime.py b/tests/test_worker_process_runtime.py new file mode 100644 index 0000000..dec0659 --- /dev/null +++ b/tests/test_worker_process_runtime.py @@ -0,0 +1,134 @@ +import importlib.util +import sys +import types +import unittest +from argparse import Namespace +from pathlib import Path +from unittest.mock import patch + + +def _load_process_module(module_filename: str): + module_path = Path(__file__).resolve().parents[1] / "workers" / "process" / module_filename + package_name = "dist_proc_testpkg" + module_name = module_filename[:-3] + + for mod_name in list(sys.modules): + if mod_name == package_name or mod_name.startswith(f"{package_name}."): + del sys.modules[mod_name] + + root_pkg = types.ModuleType(package_name) + root_pkg.__path__ = [] + sys.modules[package_name] = root_pkg + + workers_pkg = types.ModuleType(f"{package_name}.workers") + workers_pkg.__path__ = [] + sys.modules[f"{package_name}.workers"] = workers_pkg + + process_pkg = types.ModuleType(f"{package_name}.workers.process") + process_pkg.__path__ = [] + sys.modules[f"{package_name}.workers.process"] = process_pkg + + utils_pkg = types.ModuleType(f"{package_name}.utils") + utils_pkg.__path__ = [] + sys.modules[f"{package_name}.utils"] = utils_pkg + + logging_module = types.ModuleType(f"{package_name}.utils.logging") + logging_module.debug_log = lambda *_args, **_kwargs: None + logging_module.log = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.utils.logging"] = logging_module + + process_module = types.ModuleType(f"{package_name}.utils.process") + process_module.get_python_executable = lambda: "/usr/bin/test-python" + sys.modules[f"{package_name}.utils.process"] = process_module + + spec = importlib.util.spec_from_file_location( + f"{package_name}.workers.process.{module_name}", + module_path, + ) + module = importlib.util.module_from_spec(spec) + assert spec is not None and spec.loader is not None + spec.loader.exec_module(module) + return module + + +root_discovery_module = _load_process_module("root_discovery.py") +launch_builder_module = _load_process_module("launch_builder.py") + + +class ComfyRootDiscoveryTests(unittest.TestCase): + def test_prefers_loaded_comfyui_module_path(self): + discovery = root_discovery_module.ComfyRootDiscovery() + server_module = types.SimpleNamespace(__file__="/opt/ComfyUI/server.py") + + def fake_exists(path): + return path == "/opt/ComfyUI/main.py" + + with patch.dict(sys.modules, {"server": server_module}, clear=False), \ + patch.object(root_discovery_module.os.path, "exists", side_effect=fake_exists), \ + patch.dict(root_discovery_module.os.environ, {}, clear=True): + self.assertEqual(discovery.find_comfy_root(), "/opt/ComfyUI") + + +class LaunchCommandBuilderTests(unittest.TestCase): + def test_inherits_runtime_layout_args_for_desktop(self): + builder = launch_builder_module.LaunchCommandBuilder() + runtime_args = Namespace( + listen="127.0.0.1", + base_directory="C:/Users/test/ComfyUI", + temp_directory=None, + input_directory="C:/Users/test/ComfyUI/input", + output_directory="C:/Users/test/ComfyUI/output", + user_directory="C:/Users/test/ComfyUI/user", + front_end_root="C:/Program Files/ComfyUI/web_custom_versions/desktop_app", + extra_model_paths_config=[["C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml"]], + enable_manager=True, + disable_manager_ui=False, + enable_manager_legacy_ui=False, + windows_standalone_build=True, + log_stdout=True, + verbose="INFO", + enable_cors_header="*", + ) + comfy_module = types.ModuleType("comfy") + comfy_cli_args = types.ModuleType("comfy.cli_args") + comfy_cli_args.args = runtime_args + + worker_config = { + "port": 9001, + "extra_args": "--preview-method auto", + } + + def fake_exists(path): + return path == "/desktop/ComfyUI/main.py" + + with patch.dict( + sys.modules, + {"comfy": comfy_module, "comfy.cli_args": comfy_cli_args}, + clear=False, + ), patch.object(launch_builder_module.os.path, "exists", side_effect=fake_exists): + cmd = builder.build_launch_command(worker_config, "/desktop/ComfyUI") + + self.assertEqual(cmd[:2], ["/usr/bin/test-python", "/desktop/ComfyUI/main.py"]) + self.assertIn("--listen", cmd) + self.assertIn("127.0.0.1", cmd) + self.assertIn("--base-directory", cmd) + self.assertIn("C:/Users/test/ComfyUI", cmd) + self.assertIn("--input-directory", cmd) + self.assertIn("--output-directory", cmd) + self.assertIn("--user-directory", cmd) + self.assertIn("--front-end-root", cmd) + self.assertIn("--extra-model-paths-config", cmd) + self.assertIn("C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml", cmd) + self.assertIn("--enable-manager", cmd) + self.assertIn("--windows-standalone-build", cmd) + self.assertIn("--log-stdout", cmd) + self.assertIn("--disable-auto-launch", cmd) + self.assertIn("--enable-cors-header", cmd) + self.assertIn("*", cmd) + self.assertIn("--port", cmd) + self.assertIn("9001", cmd) + self.assertNotIn("--auto-launch", cmd) + + +if __name__ == "__main__": + unittest.main() diff --git a/workers/process/launch_builder.py b/workers/process/launch_builder.py index 9c03e4a..c4d510d 100644 --- a/workers/process/launch_builder.py +++ b/workers/process/launch_builder.py @@ -10,6 +10,62 @@ from ...utils.process import get_python_executable class LaunchCommandBuilder: """Build command-lines for launching worker ComfyUI processes.""" + def _extend_arg(self, cmd, flag, value): + if value in (None, "", [], ()): + return + cmd.extend([flag, str(value)]) + + def _extend_grouped_args(self, cmd, flag, values): + for group in values or []: + flattened = [str(item) for item in group if item] + if flattened: + cmd.append(flag) + cmd.extend(flattened) + + def _get_runtime_args(self): + try: + from comfy.cli_args import args + return args + except Exception as exc: + debug_log(f"Could not read current ComfyUI CLI args for worker launch: {exc}") + return None + + def _build_runtime_launch_args(self): + args = self._get_runtime_args() + if args is None: + return [] + + inherited = [] + self._extend_arg(inherited, "--listen", getattr(args, "listen", None)) + self._extend_arg(inherited, "--base-directory", getattr(args, "base_directory", None)) + self._extend_arg(inherited, "--temp-directory", getattr(args, "temp_directory", None)) + self._extend_arg(inherited, "--input-directory", getattr(args, "input_directory", None)) + self._extend_arg(inherited, "--output-directory", getattr(args, "output_directory", None)) + self._extend_arg(inherited, "--user-directory", getattr(args, "user_directory", None)) + self._extend_arg(inherited, "--front-end-root", getattr(args, "front_end_root", None)) + self._extend_grouped_args( + inherited, + "--extra-model-paths-config", + getattr(args, "extra_model_paths_config", None), + ) + + if getattr(args, "enable_manager", False): + inherited.append("--enable-manager") + if getattr(args, "disable_manager_ui", False): + inherited.append("--disable-manager-ui") + if getattr(args, "enable_manager_legacy_ui", False): + inherited.append("--enable-manager-legacy-ui") + if getattr(args, "windows_standalone_build", False): + inherited.append("--windows-standalone-build") + if getattr(args, "log_stdout", False): + inherited.append("--log-stdout") + + verbose = getattr(args, "verbose", None) + if verbose and verbose != "INFO": + inherited.extend(["--verbose", str(verbose)]) + + return inherited + def _find_windows_terminal(self): """Find Windows Terminal executable.""" possible_paths = [ @@ -39,10 +95,19 @@ class LaunchCommandBuilder: cmd = [ get_python_executable(), main_py, - "--port", - str(worker_config["port"]), - "--enable-cors-header", ] + cmd.extend(self._build_runtime_launch_args()) + cmd.extend(["--port", str(worker_config["port"])]) + + current_args = self._get_runtime_args() + current_cors = getattr(current_args, "enable_cors_header", None) if current_args else None + cmd.append("--enable-cors-header") + if current_cors is not None: + cmd.append(str(current_cors)) + + if "--disable-auto-launch" not in cmd: + cmd.append("--disable-auto-launch") + debug_log(f"Using main.py: {main_py}") else: error_msg = f"Could not find main.py in {comfy_root}\n" diff --git a/workers/process/lifecycle.py b/workers/process/lifecycle.py index b5a04e7..62a3653 100644 --- a/workers/process/lifecycle.py +++ b/workers/process/lifecycle.py @@ -33,6 +33,7 @@ class ProcessLifecycle: env["CUDA_VISIBLE_DEVICES"] = str(worker_config.get("cuda_device", 0)) env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" env["COMFYUI_MASTER_PID"] = str(os.getpid()) + env["COMFYUI_IS_WORKER"] = "1" cmd = self._manager.build_launch_command(worker_config, comfy_root) cwd = comfy_root diff --git a/workers/process/root_discovery.py b/workers/process/root_discovery.py index f8a3201..a17c5bd 100644 --- a/workers/process/root_discovery.py +++ b/workers/process/root_discovery.py @@ -1,4 +1,5 @@ import os +import sys from ...utils.logging import debug_log, log @@ -6,6 +7,21 @@ from ...utils.logging import debug_log, log class ComfyRootDiscovery: """Resolve the ComfyUI root directory across local and container layouts.""" + def _find_root_from_loaded_modules(self): + """Use already-imported ComfyUI modules to locate the runtime root.""" + for module_name in ("server", "folder_paths", "main"): + module = sys.modules.get(module_name) + module_file = getattr(module, "__file__", None) + if not module_file: + continue + + candidate = os.path.dirname(os.path.abspath(module_file)) + if os.path.exists(os.path.join(candidate, "main.py")): + debug_log(f"Found ComfyUI root via loaded module {module_name}: {candidate}") + return candidate + + return None + def find_comfy_root(self): # Start from current file location. current_dir = os.path.dirname(os.path.abspath(__file__)) @@ -17,12 +33,17 @@ class ComfyRootDiscovery: debug_log(f"Found ComfyUI root via COMFYUI_ROOT environment variable: {env_root}") return env_root - # Method 2: Try going up from custom_nodes directory. + # Method 2: Inspect the already-loaded ComfyUI runtime modules. + runtime_root = self._find_root_from_loaded_modules() + if runtime_root: + return runtime_root + + # Method 3: Try going up from custom_nodes directory. if os.path.exists(os.path.join(potential_root, "main.py")): debug_log(f"Found ComfyUI root via directory traversal: {potential_root}") return potential_root - # Method 3: Look for common Docker paths. + # Method 4: Look for common Docker paths. docker_paths = [ "/basedir", "/ComfyUI", @@ -37,7 +58,7 @@ class ComfyRootDiscovery: debug_log(f"Found ComfyUI root in Docker path: {path}") return path - # Method 4: Search upwards for main.py. + # Method 5: Search upwards for main.py. search_dir = current_dir for _ in range(5): if os.path.exists(os.path.join(search_dir, "main.py")): @@ -48,7 +69,7 @@ class ComfyRootDiscovery: break search_dir = parent - # Method 5: Try to import and use folder_paths. + # Method 6: Try to import and use folder_paths. try: import folder_paths