Add multi gpus feature

comfy_cli/cmdline.py

@@ -450,10 +450,16 @@
         raise typer.Exit(code=1)
 
 
-@app.command(help="Stop background ComfyUI")
+@app.command(help="Stop background ComfyUI: ?[--name str]")
 @tracking.track_command()
-def stop():
-    if constants.CONFIG_KEY_BACKGROUND not in ConfigManager().config["DEFAULT"]:
+def stop(
+    name: Annotated[
+        Optional[str], typer.Option(help="If running multiple servers, a name should be specified for each of them")
+    ] = None,
+):
+    if name is not None:
+        ConfigManager().load(name=name)
+    if not ConfigManager().config.has_option(name or constants.CONFIG_DEFAULT_KEY, constants.CONFIG_KEY_BACKGROUND):
         rprint("[bold red]No ComfyUI is running in the background.[/bold red]\n")
         raise typer.Exit(code=1)
 
@@ -468,16 +474,19 @@
     else:
         rprint(f"[bold yellow]Background ComfyUI is stopped.[/bold yellow] ({bg_info[0]}:{bg_info[1]})")
 
-    ConfigManager().remove_background()
+    ConfigManager().remove_background(name=name)
 
 
-@app.command(help="Launch ComfyUI: ?[--background] ?[-- <extra args ...>]")
+@app.command(help="Launch ComfyUI: ?[--background] ?[--name str] ?[-- <extra args ...>]")
 @tracking.track_command()
 def launch(
     background: Annotated[bool, typer.Option(help="Launch ComfyUI in background")] = False,
     extra: List[str] = typer.Argument(None),
+    name: Annotated[
+        Optional[str], typer.Option(help="If running multiple servers, a name should be specified for each of them")
+    ] = None,
 ):
-    launch_command(background, extra)
+    launch_command(background, extra, name=name)
 
 
 @app.command("set-default", help="Set default ComfyUI path")

comfy_cli/command/launch.py

@@ -10,6 +10,7 @@
 import typer
 from rich import print
 from rich.console import Console
+from rich.markup import escape
 from rich.panel import Panel
 
 from comfy_cli import constants, utils
@@ -56,12 +57,12 @@
         def redirector_stderr():
             while True:
                 if process is not None:
-                    print(process.stderr.readline(), end="")
+                    print(escape(process.stderr.readline()), end="")
 
         def redirector_stdout():
             while True:
                 if process is not None:
-                    print(process.stdout.readline(), end="")
+                    print(escape(process.stdout.readline()), end="")
 
         threading.Thread(target=redirector_stderr).start()
         threading.Thread(target=redirector_stdout).start()
@@ -107,6 +108,7 @@
 def launch(
     background: bool = False,
     extra: list[str] | None = None,
+    name: str | None = None,
 ):
     check_for_updates()
     resolved_workspace = workspace_manager.workspace_path
@@ -119,7 +121,7 @@
         raise typer.Exit(code=1)
 
     if (extra is None or len(extra) == 0) and workspace_manager.workspace_type == WorkspaceType.DEFAULT:
-        launch_extras = workspace_manager.config_manager.config["DEFAULT"].get(
+        launch_extras = workspace_manager.config_manager.config[name or constants.CONFIG_DEFAULT_KEY].get(
             constants.CONFIG_KEY_DEFAULT_LAUNCH_EXTRAS, ""
         )
 
@@ -133,12 +135,14 @@
 
     os.chdir(resolved_workspace)
     if background:
-        background_launch(extra)
+        background_launch(extra, name=name)
     else:
         launch_comfyui(extra)
 
 
-def background_launch(extra):
+def background_launch(extra, name: str | None = None):
+    if name is not None:
+        ConfigManager().load(name=name)
     config_background = ConfigManager().background
     if config_background is not None and utils.is_running(config_background[2]):
         console.print(
@@ -174,7 +178,7 @@
     ] + extra
 
     loop = asyncio.get_event_loop()
-    log = loop.run_until_complete(launch_and_monitor(cmd, listen, port))
+    log = loop.run_until_complete(launch_and_monitor(cmd, listen, port, name=name))
 
     if log is not None:
         console.print(
@@ -190,7 +194,7 @@
     os._exit(1)
 
 
-async def launch_and_monitor(cmd, listen, port):
+async def launch_and_monitor(cmd, listen, port, name: str | None = None):
     """
     Monitor the process during the background launch.
 
@@ -238,7 +242,11 @@
                 print(
                     f"[bold yellow]ComfyUI is successfully launched in the background.[/bold yellow]\nTo see the GUI go to: http://{listen}:{port}"
                 )
-                ConfigManager().config["DEFAULT"][constants.CONFIG_KEY_BACKGROUND] = f"{(listen, port, process.pid)}"
+                if name is not None and name not in ConfigManager().config:
+                    ConfigManager().config.add_section(name)
+                ConfigManager().config[name or constants.CONFIG_DEFAULT_KEY][constants.CONFIG_KEY_BACKGROUND] = (
+                    f"{(listen, port, process.pid)}"
+                )
                 ConfigManager().write_config()
 
                 # NOTE: os.exit(0) doesn't work.

comfy_cli/command/run.py

@@ -71,6 +71,7 @@
         execution.queue()
         if wait:
             execution.watch_execution()
+            execution.close()
             end = time.time()
             progress.stop()
             progress = None
@@ -165,6 +166,10 @@
                 if not self.on_message(message):
                     break
 
+    def close(self):
+        if self.ws:
+            self.ws.close()
+
     def update_overall_progress(self):
         self.progress.update(self.overall_task, completed=self.total_nodes - len(self.remaining_nodes))
 
@@ -273,4 +278,5 @@
 
     def on_error(self, data):
         pprint(f"[bold red]Error running workflow\n{json.dumps(data, indent=2)}[/bold red]")
+        self.close()
         raise typer.Exit(code=1)

comfy_cli/config_manager.py

@@ -1,8 +1,11 @@
+import ast
 import configparser
 import os
 from importlib.metadata import version
 from typing import Optional, Tuple
 
+from filelock import FileLock
+
 from comfy_cli import constants, logging
 from comfy_cli.utils import get_os, is_running, singleton
 
@@ -23,56 +26,64 @@
 
     def write_config(self):
         config_file_path = os.path.join(self.get_config_path(), "config.ini")
+        config_file_path_lock = os.path.join(self.get_config_path(), "config.ini.lock")
         dir_path = os.path.dirname(config_file_path)
-        if not os.path.exists(dir_path):
-            os.mkdir(dir_path)
+        lock = FileLock(config_file_path_lock, timeout=10)
+        with lock:
+            if not os.path.exists(dir_path):
+                os.mkdir(dir_path)
 
-        with open(config_file_path, "w") as configfile:
-            self.config.write(configfile)
+            with open(config_file_path, "w") as configfile:
+                self.config.write(configfile)
 
-    def set(self, key, value):
+    def set(self, key, value, name: Optional[str] = None):
         """
         Set a key-value pair in the config file.
         """
-        self.config["DEFAULT"][key] = value
+        self.config[name or constants.CONFIG_DEFAULT_KEY][key] = value
         self.write_config()  # Write changes to file immediately
 
-    def get(self, key):
+    def get(self, key, name: Optional[str] = None):
         """
         Get a value from the config file. Returns None if the key does not exist.
         """
-        return self.config["DEFAULT"].get(key, None)  # Returns None if the key does not exist
+        return self.config[name or constants.CONFIG_DEFAULT_KEY].get(
+            key, None
+        )  # Returns None if the key does not exist
 
-    def load(self):
+    def load(self, name: Optional[str] = None):
         config_file_path = self.get_config_file_path()
         if os.path.exists(config_file_path):
             self.config = configparser.ConfigParser()
-            self.config.read(config_file_path)
+            config_file_path_lock = os.path.join(self.get_config_path(), "config.ini.lock")
+            lock = FileLock(config_file_path_lock, timeout=10)
+            with lock:
+                self.config.read(config_file_path)
 
         # TODO: We need a policy for clearing the tmp directory.
         tmp_path = os.path.join(self.get_config_path(), "tmp")
         if not os.path.exists(tmp_path):
             os.makedirs(tmp_path)
 
-        if "background" in self.config["DEFAULT"]:
-            bg_info = self.config["DEFAULT"]["background"].strip("()").split(",")
-            bg_info = [item.strip().strip("'") for item in bg_info]
-            self.background = bg_info[0], int(bg_info[1]), int(bg_info[2])
+        section = name or constants.CONFIG_DEFAULT_KEY
+        if self.config.has_option(section, constants.CONFIG_KEY_BACKGROUND):
+            self.background = ast.literal_eval(self.config[section][constants.CONFIG_KEY_BACKGROUND])
 
             if not is_running(self.background[2]):
-                self.remove_background()
+                self.remove_background(name=section)
 
-    def fill_print_env(self, table):
+    def fill_print_env(self, table, name: Optional[str] = None):
+        section = name or constants.CONFIG_DEFAULT_KEY
         table.add_row("Config Path", self.get_config_file_path())
 
         launch_extras = ""
-        if self.config.has_option("DEFAULT", "default_workspace"):
+        if self.config.has_option(section, "default_workspace"):
             table.add_row(
                 "Default ComfyUI workspace",
-                self.config["DEFAULT"][constants.CONFIG_KEY_DEFAULT_WORKSPACE],
+                self.config[section][constants.CONFIG_KEY_DEFAULT_WORKSPACE],
             )
 
-            launch_extras = self.config["DEFAULT"].get(constants.CONFIG_KEY_DEFAULT_LAUNCH_EXTRAS, "")
+            launch_extras = self.config[section].get(constants.CONFIG_KEY_DEFAULT_LAUNCH_EXTRAS, "")
         else:
             table.add_row("Default ComfyUI workspace", "No default ComfyUI workspace")
 
@@ -81,21 +92,21 @@
 
         table.add_row("Default ComfyUI launch extra options", launch_extras)
 
-        if self.config.has_option("DEFAULT", constants.CONFIG_KEY_RECENT_WORKSPACE):
+        if self.config.has_option(section, constants.CONFIG_KEY_RECENT_WORKSPACE):
             table.add_row(
                 "Recent ComfyUI workspace",
-                self.config["DEFAULT"][constants.CONFIG_KEY_RECENT_WORKSPACE],
+                self.config[section][constants.CONFIG_KEY_RECENT_WORKSPACE],
             )
         else:
             table.add_row("Recent ComfyUI workspace", "No recent run")
 
-        if self.config.has_option("DEFAULT", "enable_tracking"):
+        if self.config.has_option(section, "enable_tracking"):
             table.add_row(
                 "Tracking Analytics",
-                ("Enabled" if self.config["DEFAULT"]["enable_tracking"] == "True" else "Disabled"),
+                ("Enabled" if self.config[section]["enable_tracking"] == "True" else "Disabled"),
             )
 
-        if self.config.has_option("DEFAULT", constants.CONFIG_KEY_BACKGROUND):
+        if self.config.has_option(section, constants.CONFIG_KEY_BACKGROUND):
             bg_info = self.background
             if bg_info:
                 table.add_row(
@@ -105,11 +116,25 @@
         else:
             table.add_row("Background ComfyUI", "[bold red]No[/bold red]")
 
-    def remove_background(self):
-        del self.config["DEFAULT"]["background"]
+    def remove_background(self, name: Optional[str] = None):
+        section = name or constants.CONFIG_DEFAULT_KEY
+        del self.config[section][constants.CONFIG_KEY_BACKGROUND]
+        if name is not None:
+            self.cleanup_session(name)
         self.write_config()
         self.background = None
 
+    def cleanup_session(self, section: str):
+        if section not in self.config or section == constants.CONFIG_DEFAULT_KEY:
+            return
+        overridden_options = [
+            opt
+            for opt in self.config.options(section)
+            if section in self.config and opt in self.config[section] and opt not in self.config.defaults()
+        ]
+        if not overridden_options:
+            self.config.remove_section(section)
+
     def get_cli_version(self):
         # Note: this approach should work for users installing the CLI via
         # PyPi and Homebrew (e.g., pip install comfy-cli)

comfy_cli/constants.py

@@ -28,6 +28,7 @@ class PROC(str, Enum):
     OS.MACOS: os.path.join(os.path.expanduser("~"), "Library", "Application Support", "comfy-cli"),
     OS.LINUX: os.path.join(os.path.expanduser("~"), ".config", "comfy-cli"),
 }
+CONFIG_DEFAULT_KEY = "DEFAULT"
 
 CONTEXT_KEY_WORKSPACE = "workspace"
 CONTEXT_KEY_RECENT = "recent"

pyproject.toml

@@ -44,6 +44,7 @@ dependencies = [
   "ruff",
   "semver~=3.0.2",
   "cookiecutter",
+  "filelock",
 ]
 
 [project.optional-dependencies]