diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py index da347437..d929a782 100644 --- a/comfy_cli/cmdline.py +++ b/comfy_cli/cmdline.py @@ -252,7 +252,13 @@ def install( @app.command(help="Update ComfyUI Environment [all|comfy]") @tracking.track_command() -def update(target: str = typer.Argument("comfy", help="[all|comfy]")): +def update( + target: str = typer.Argument( + "comfy", + help="[all|comfy]", + autocompletion=utils.create_choice_completer(["all", "comfy"]), + ) +): if target not in ["all", "comfy"]: typer.echo( f"Invalid target: {target}. Allowed targets are 'all', 'comfy'.", @@ -277,6 +283,8 @@ def update(target: str = typer.Argument("comfy", help="[all|comfy]")): check=True, ) + custom_nodes.command.update_node_id_cache() + # @app.command(help="Run workflow file") # @tracking.track_command() diff --git a/comfy_cli/command/custom_nodes/command.py b/comfy_cli/command/custom_nodes/command.py index f08c3464..545069c8 100644 --- a/comfy_cli/command/custom_nodes/command.py +++ b/comfy_cli/command/custom_nodes/command.py @@ -9,7 +9,7 @@ from rich import print from typing_extensions import List, Annotated -from comfy_cli import ui, logging, tracking +from comfy_cli import ui, logging, tracking, utils from comfy_cli.config_manager import ConfigManager from comfy_cli.file_utils import ( download_file, @@ -168,8 +168,9 @@ def clear(path: str): execute_cm_cli(["clear", path]) -def show_completer(incomplete: str) -> list[str]: - valid_choices = [ +# completers +show_completer = utils.create_choice_completer( + [ "installed", "enabled", "not-installed", @@ -178,17 +179,33 @@ def show_completer(incomplete: str) -> list[str]: "snapshot", "snapshot-list", ] - return [choice for choice in valid_choices if choice.startswith(incomplete)] +) + + +mode_completer = utils.create_choice_completer(["remote", "local", "cache"]) + + +channel_completer = utils.create_choice_completer( + ["default", "recent", "dev", "forked", "tutorial", "legacy"] +) -def mode_completer(incomplete: str) -> list[str]: - modes = ["remote", "local", "cache"] - return [mode for mode in modes if mode.startswith(incomplete)] +def node_completer(incomplete: str) -> list[str]: + try: + config_manager = ConfigManager() + tmp_path = os.path.join( + config_manager.get_config_path(), "tmp", "node-cache.list" + ) + with open(tmp_path, "r", encoding="UTF-8", errors="ignore") as cache_file: + return [ + node_id + for node_id in cache_file.readlines() + if node_id.startswith(incomplete) + ] -def channel_completer(incomplete: str) -> list[str]: - modes = ["default", "recent", "dev", "forked", "tutorial", "legacy"] - return [mode for mode in modes if mode.startswith(incomplete)] + except Exception: + return [] @app.command(help="Show node list") @@ -285,7 +302,9 @@ def simple_show( @app.command(help="Install custom nodes") @tracking.track_command("node") def install( - args: List[str] = typer.Argument(..., help="install custom nodes"), + args: List[str] = typer.Argument( + ..., help="install custom nodes", autocompletion=node_completer + ), channel: Annotated[ str, typer.Option( @@ -318,7 +337,9 @@ def install( @app.command(help="Reinstall custom nodes") @tracking.track_command("node") def reinstall( - args: List[str] = typer.Argument(..., help="reinstall custom nodes"), + args: List[str] = typer.Argument( + ..., help="reinstall custom nodes", autocompletion=node_completer + ), channel: Annotated[ str, typer.Option( @@ -351,7 +372,9 @@ def reinstall( @app.command(help="Uninstall custom nodes") @tracking.track_command("node") def uninstall( - args: List[str] = typer.Argument(..., help="uninstall custom nodes"), + args: List[str] = typer.Argument( + ..., help="uninstall custom nodes", autocompletion=node_completer + ), channel: Annotated[ str, typer.Option( @@ -381,13 +404,36 @@ def uninstall( execute_cm_cli(["uninstall"] + args, channel, mode) -# `update, disable, enable, fix` allows `all` param +def update_node_id_cache(): + try: + config_manager = ConfigManager() + workspace_path = workspace_manager.workspace_path + + cm_cli_path = os.path.join( + workspace_path, "custom_nodes", "ComfyUI-Manager", "cm-cli.py" + ) + + tmp_path = os.path.join(config_manager.get_config_path(), "tmp") + if not os.path.exists(tmp_path): + os.makedirs(tmp_path) + + cache_path = os.path.join(tmp_path, "node-cache.list") + cmd = [sys.executable, cm_cli_path, "export-custom-node-ids", cache_path] + + new_env = os.environ.copy() + new_env["COMFYUI_PATH"] = workspace_path + subprocess.check_output(cmd, env=new_env) + except Exception: + pass +# `update, disable, enable, fix` allows `all` param @app.command(help="Update custom nodes or ComfyUI") @tracking.track_command("node") def update( - args: List[str] = typer.Argument(..., help="update custom nodes"), + args: List[str] = typer.Argument( + ..., help="update custom nodes", autocompletion=node_completer + ), channel: Annotated[ str, typer.Option( @@ -412,11 +458,15 @@ def update( execute_cm_cli(["update"] + args, channel, mode) + update_node_id_cache() + @app.command(help="Disable custom nodes") @tracking.track_command("node") def disable( - args: List[str] = typer.Argument(..., help="disable custom nodes"), + args: List[str] = typer.Argument( + ..., help="disable custom nodes", autocompletion=node_completer + ), channel: Annotated[ str, typer.Option( @@ -445,7 +495,9 @@ def disable( @app.command(help="Enable custom nodes") @tracking.track_command("node") def enable( - args: List[str] = typer.Argument(..., help="enable custom nodes"), + args: List[str] = typer.Argument( + ..., help="enable custom nodes", autocompletion=node_completer + ), channel: Annotated[ str, typer.Option( @@ -475,7 +527,9 @@ def enable( @tracking.track_command("node") def fix( args: List[str] = typer.Argument( - ..., help="fix dependencies for specified custom nodes" + ..., + help="fix dependencies for specified custom nodes", + autocompletion=node_completer, ), channel: Annotated[ str, diff --git a/comfy_cli/command/install.py b/comfy_cli/command/install.py index af5d6c48..79567a5f 100644 --- a/comfy_cli/command/install.py +++ b/comfy_cli/command/install.py @@ -8,6 +8,7 @@ from comfy_cli import constants from comfy_cli.constants import GPU_OPTION from comfy_cli.workspace_manager import WorkspaceManager, check_comfy_repo +from comfy_cli.command.custom_nodes.command import update_node_id_cache def get_os_details(): @@ -173,6 +174,8 @@ def execute( subprocess.run(["git", "clone", manager_url, manager_repo_dir]) install_manager_dependencies(repo_dir) + update_node_id_cache() + os.chdir(repo_dir) print("") diff --git a/comfy_cli/utils.py b/comfy_cli/utils.py index 49054aab..23973d0a 100644 --- a/comfy_cli/utils.py +++ b/comfy_cli/utils.py @@ -57,3 +57,10 @@ def is_running(pid): return True except psutil.NoSuchProcess: return False + + +def create_choice_completer(opts): + def f(incomplete: str) -> list[str]: + return [opt for opt in opts if opt.startswith(incomplete)] + + return f