Add ComfyUI Desktop worker support

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.
 </details>
 
-<details>
-<summary>Does it work with the ComfyUI desktop app?</summary>
-Currently, it is not compatible with the ComfyUI desktop app.
-</details>
+<details>
+<summary>Does it work with the ComfyUI desktop app?</summary>
+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.
+</details>
 
 <details>
 <summary>Can I combine my RTX 5090 with a GTX 980 to get faster results?</summary>

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

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"

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

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