Add ComfyUI Desktop worker support

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