From f265b5cc6f70eea840beca6ee2644fbdf7883947 Mon Sep 17 00:00:00 2001 From: dbataille Date: Fri, 25 Apr 2025 21:30:33 +0000 Subject: [PATCH 1/5] Add multi gpus feature --- comfy_cli/cmdline.py | 17 ++++++--- comfy_cli/command/launch.py | 22 +++++++---- comfy_cli/config_manager.py | 74 ++++++++++++++++++++++++------------- comfy_cli/constants.py | 1 + pyproject.toml | 1 + 5 files changed, 75 insertions(+), 40 deletions(-) diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py index 247ba95a..8beb7538 100644 --- a/comfy_cli/cmdline.py +++ b/comfy_cli/cmdline.py @@ -450,10 +450,14 @@ def validate_comfyui(_env_checker): 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[str | None, 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 +472,17 @@ def stop(): 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] ?[-- ]") +@app.command(help="Launch ComfyUI: ?[--background] ?[--name str] ?[-- ]") @tracking.track_command() def launch( background: Annotated[bool, typer.Option(help="Launch ComfyUI in background")] = False, extra: List[str] = typer.Argument(None), + name: Annotated[str | None, 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") diff --git a/comfy_cli/command/launch.py b/comfy_cli/command/launch.py index 52885e66..e0958489 100644 --- a/comfy_cli/command/launch.py +++ b/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 launch_comfyui(extra): 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 redirector_stdout(): 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 @@ def launch( 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 @@ def launch( 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 @@ def background_launch(extra): ] + 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 @@ def background_launch(extra): 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,9 @@ def msg_hook(stream): 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. diff --git a/comfy_cli/config_manager.py b/comfy_cli/config_manager.py index 616bf090..cdd29ca7 100644 --- a/comfy_cli/config_manager.py +++ b/comfy_cli/config_manager.py @@ -1,4 +1,6 @@ +import ast import configparser +from filelock import FileLock import os from importlib.metadata import version from typing import Optional, Tuple @@ -23,56 +25,62 @@ def get_config_file_path(self): 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: str | None = 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: str | None = 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: str | None = 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: str | None = 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 +89,21 @@ def fill_print_env(self, table): 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 +113,25 @@ def fill_print_env(self, table): 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: str | None = 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) diff --git a/comfy_cli/constants.py b/comfy_cli/constants.py index 37b1a24e..443f5616 100644 --- a/comfy_cli/constants.py +++ b/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" diff --git a/pyproject.toml b/pyproject.toml index 328c087c..a1d87d6e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ dependencies = [ "ruff", "semver~=3.0.2", "cookiecutter", + "filelock", ] [project.optional-dependencies] From 4e0ddd79dad90d6f531f21b3a0018ec69719a078 Mon Sep 17 00:00:00 2001 From: dbataille Date: Fri, 25 Apr 2025 21:54:13 +0000 Subject: [PATCH 2/5] Format and fix typing because old Python --- comfy_cli/cmdline.py | 8 ++++++-- comfy_cli/command/launch.py | 4 +++- comfy_cli/config_manager.py | 20 +++++++++++--------- 3 files changed, 20 insertions(+), 12 deletions(-) diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py index 8beb7538..a2a4dc67 100644 --- a/comfy_cli/cmdline.py +++ b/comfy_cli/cmdline.py @@ -453,7 +453,9 @@ def validate_comfyui(_env_checker): @app.command(help="Stop background ComfyUI: ?[--name str]") @tracking.track_command() def stop( - name: Annotated[str | None, typer.Option(help="If running multiple servers, a name should be specified for each of them")] = None + name: Annotated[ + str | None, 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) @@ -480,7 +482,9 @@ def stop( def launch( background: Annotated[bool, typer.Option(help="Launch ComfyUI in background")] = False, extra: List[str] = typer.Argument(None), - name: Annotated[str | None, typer.Option(help="If running multiple servers, a name should be specified for each of them")] = None + name: Annotated[ + str | None, typer.Option(help="If running multiple servers, a name should be specified for each of them") + ] = None, ): launch_command(background, extra, name=name) diff --git a/comfy_cli/command/launch.py b/comfy_cli/command/launch.py index e0958489..f7a42c10 100644 --- a/comfy_cli/command/launch.py +++ b/comfy_cli/command/launch.py @@ -244,7 +244,9 @@ def msg_hook(stream): ) 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().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. diff --git a/comfy_cli/config_manager.py b/comfy_cli/config_manager.py index cdd29ca7..6f2597cf 100644 --- a/comfy_cli/config_manager.py +++ b/comfy_cli/config_manager.py @@ -35,20 +35,22 @@ def write_config(self): with open(config_file_path, "w") as configfile: self.config.write(configfile) - def set(self, key, value, name: str | None = None): + def set(self, key, value, name: Optional[str] = None): """ Set a key-value pair in the config file. """ self.config[name or constants.CONFIG_DEFAULT_KEY][key] = value self.write_config() # Write changes to file immediately - def get(self, key, name: str | None = None): + 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[name or constants.CONFIG_DEFAULT_KEY].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, name: str | None = None): + 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() @@ -69,7 +71,7 @@ def load(self, name: str | None = None): if not is_running(self.background[2]): self.remove_background(name=section) - def fill_print_env(self, table, name: str | None = None): + 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()) @@ -113,7 +115,7 @@ def fill_print_env(self, table, name: str | None = None): else: table.add_row("Background ComfyUI", "[bold red]No[/bold red]") - def remove_background(self, name: str | None = None): + 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: @@ -125,9 +127,9 @@ 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() + 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) From 1faf38f9659a916dae38b1fd359cc5fb27344b2a Mon Sep 17 00:00:00 2001 From: dbataille Date: Fri, 25 Apr 2025 22:00:17 +0000 Subject: [PATCH 3/5] Format again --- comfy_cli/cmdline.py | 2 +- comfy_cli/config_manager.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py index a2a4dc67..27a9065d 100644 --- a/comfy_cli/cmdline.py +++ b/comfy_cli/cmdline.py @@ -454,7 +454,7 @@ def validate_comfyui(_env_checker): @tracking.track_command() def stop( name: Annotated[ - str | None, typer.Option(help="If running multiple servers, a name should be specified for each of them") + Optional[str], typer.Option(help="If running multiple servers, a name should be specified for each of them") ] = None, ): if name is not None: diff --git a/comfy_cli/config_manager.py b/comfy_cli/config_manager.py index 6f2597cf..3ad7224d 100644 --- a/comfy_cli/config_manager.py +++ b/comfy_cli/config_manager.py @@ -1,10 +1,11 @@ import ast import configparser -from filelock import FileLock 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 From d622c5eb4f613f5493c2f228a546f4f5b13b66e2 Mon Sep 17 00:00:00 2001 From: dbataille Date: Fri, 25 Apr 2025 22:03:19 +0000 Subject: [PATCH 4/5] Format again --- comfy_cli/cmdline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py index 27a9065d..1a2861b6 100644 --- a/comfy_cli/cmdline.py +++ b/comfy_cli/cmdline.py @@ -483,7 +483,7 @@ def launch( background: Annotated[bool, typer.Option(help="Launch ComfyUI in background")] = False, extra: List[str] = typer.Argument(None), name: Annotated[ - str | None, typer.Option(help="If running multiple servers, a name should be specified for each of them") + Optional[str], typer.Option(help="If running multiple servers, a name should be specified for each of them") ] = None, ): launch_command(background, extra, name=name) From 4a679a7417aca3f91508152be99214cd33f9047b Mon Sep 17 00:00:00 2001 From: dbataille Date: Tue, 29 Apr 2025 01:34:42 +0000 Subject: [PATCH 5/5] Close connection when the workflow is finished to be executed --- comfy_cli/command/run.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/comfy_cli/command/run.py b/comfy_cli/command/run.py index 78eb1e47..e8696a7e 100644 --- a/comfy_cli/command/run.py +++ b/comfy_cli/command/run.py @@ -71,6 +71,7 @@ def execute(workflow: str, host, port, wait=True, verbose=False, local_paths=Fal execution.queue() if wait: execution.watch_execution() + execution.close() end = time.time() progress.stop() progress = None @@ -165,6 +166,10 @@ def watch_execution(self): 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_executed(self, data): 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)