From 7c2ed17dd21c312da32fcef0882af4eb3bc69279 Mon Sep 17 00:00:00 2001 From: Yoland Y <4950057+yoland68@users.noreply.github.com> Date: Fri, 3 May 2024 17:55:39 -0700 Subject: [PATCH 1/3] Feature: save yaml file --- comfy_cli/cmdline.py | 8 ++ comfy_cli/constants.py | 6 +- comfy_cli/workspace_manager.py | 183 ++++++++++++++++++++++++--------- 3 files changed, 146 insertions(+), 51 deletions(-) diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py index 68888fa..8d60b3a 100644 --- a/comfy_cli/cmdline.py +++ b/comfy_cli/cmdline.py @@ -505,6 +505,14 @@ def which(): ) raise typer.Exit(code=1) + if workspace_manager.workspace_type == WorkspaceType.DEFAULT: + print("Default Comfy path is used.") + elif workspace_manager.workspace_type == WorkspaceType.CURRENT_DIR: + print("Current directory is used.") + elif workspace_manager.workspace_type == WorkspaceType.RECENT: + print("Most recently used ComfyUP is used") + elif workspace_manager.workspace_type == WorkspaceType.SPECIFIED: + print("Specified path is used.") print(f"Target ComfyUI path: {comfy_path}") diff --git a/comfy_cli/constants.py b/comfy_cli/constants.py index 2c77561..3a711bf 100644 --- a/comfy_cli/constants.py +++ b/comfy_cli/constants.py @@ -38,9 +38,7 @@ class OS(Enum): CONFIG_KEY_INSTALL_EVENT_TRIGGERED = "install_event_triggered" CONFIG_KEY_BACKGROUND = "background" -DEFAULT_TRACKING_VALUE = True - -COMFY_LOCK_YAML_FILE = "comfy.lock.yaml" +COMFY_LOCK_YAML_FILE = "comfy-lock.yaml" # TODO: figure out a better way to check if this is a comfy repo COMFY_ORIGIN_URL_CHOICES = [ @@ -65,4 +63,6 @@ class GPU_OPTION(Enum): # https://github.com/comfyanonymous/ComfyUI/blob/a88b0ebc2d2f933c94e42aa689c42e836eedaf3c/folder_paths.py#L5 SUPPORTED_PT_EXTENSIONS = (".ckpt", ".pt", ".bin", ".pth", ".safetensors") +IGNORE_CUSTOM_NODE_FOLDERS = ("__pycache__",) + COMFY_REGISTRY_URL_ROOT = "https://api-frontend-dev-qod3oz2v2q-uc.a.run.app" diff --git a/comfy_cli/workspace_manager.py b/comfy_cli/workspace_manager.py index 4350e6c..8434fbc 100644 --- a/comfy_cli/workspace_manager.py +++ b/comfy_cli/workspace_manager.py @@ -1,7 +1,8 @@ import concurrent.futures import os +import pdb import sys -from dataclasses import dataclass, field +from dataclasses import asdict, dataclass, field, is_dataclass from datetime import datetime from enum import Enum from pathlib import Path @@ -34,6 +35,7 @@ class Model: @dataclass class Basics: + remote: str name: Optional[str] = None updated_at: datetime = None @@ -41,7 +43,10 @@ class Basics: @dataclass class CustomNode: # Todo: Add custom node fields for comfy-lock.yaml - pass + path: str + enabled: bool + is_git: Optional[bool] = False + remote: Optional[str] = None @dataclass @@ -68,52 +73,72 @@ def check_comfy_repo(path): return False, None -# Generate and update this following method using chatGPT +def _serialize_dataclass(data): + """Serialize dataclasses, lists, and datetimes, filtering out None values.""" + if is_dataclass(data): + # Convert dataclass to dict, omitting keys with None values + return { + k: _serialize_dataclass(v) for k, v in asdict(data).items() if v is not None + } + elif isinstance(data, list): + # Serialize each item in the list + return [_serialize_dataclass(item) for item in data] + # TODO: keep in mind for any other data type to create condition to serialize it properly + elif isinstance(data, datetime): + return data.isoformat() + return data + + def load_yaml(file_path: str) -> ComfyLockYAMLStruct: with open(file_path, "r", encoding="utf-8") as file: data = yaml.safe_load(file) + basics = Basics( + remote=data.get("basics", {}).get("remote", ""), name=data.get("basics", {}).get("name"), - updated_at=( - datetime.fromisoformat(data.get("basics", {}).get("updated_at")) - if data.get("basics", {}).get("updated_at") - else None - ), + updated_at=datetime.fromisoformat(data.get("basics", {}).get("updated_at")) + if data.get("basics", {}).get("updated_at") + else None, ) + models = [ Model( - name=m.get("model"), + name=m.get("name"), url=m.get("url"), - paths=[ModelPath(path=p.get("path")) for p in m.get("paths", [])], + paths=[ModelPath(path=p["path"]) for p in m.get("paths", [])], hash=m.get("hash"), type=m.get("type"), ) for m in data.get("models", []) ] - custom_nodes = [] + + custom_nodes = [ + CustomNode( + path=cn.get("path"), + enabled=cn.get("enabled", False), + is_git=cn.get("is_git"), + remote=cn.get("remote"), + ) + for cn in data.get("custom_nodes", []) + ] + + return ComfyLockYAMLStruct( + basics=basics, models=models, custom_nodes=custom_nodes + ) -# Generate and update this following method using chatGPT def save_yaml(file_path: str, metadata: ComfyLockYAMLStruct): - data = { - "basics": { - "name": metadata.basics.name, - "updated_at": metadata.basics.updated_at.isoformat(), - }, - "models": [ - { - "model": m.name, - "url": m.url, - "paths": [{"path": p.path} for p in m.paths], - "hash": m.hash, - "type": m.type, - } - for m in metadata.models - ], - "custom_nodes": [], - } + # Serialize the dataclass to a dictionary, handling nested dataclasses and lists + data = _serialize_dataclass(metadata) + + # Write the dictionary to a YAML file with open(file_path, "w", encoding="utf-8") as file: - yaml.safe_dump(data, file, default_flow_style=False, allow_unicode=True) + file.write( + "# Beta Feature: This file is generated with comfy-cli to track ComfyUI state\n" + ) + yaml.safe_dump( + data, file, default_flow_style=False, allow_unicode=True, sort_keys=True + ) # Function to check if the file is config.json @@ -122,6 +147,14 @@ def check_file_is_model(path): return str(path) +def check_folder_is_git(path) -> Tuple[bool, Optional[git.Repo]]: + try: + repo = git.Repo(path) + return True, repo + except git.exc.InvalidGitRepositoryError: + return False, None + + class WorkspaceType(Enum): CURRENT_DIR = "current_dir" DEFAULT = "default" @@ -136,12 +169,13 @@ def __init__( self, ): self.config_manager = ConfigManager() - self.metadata = ComfyLockYAMLStruct(basics=Basics(), models=[]) + self.metadata = ComfyLockYAMLStruct(basics=Basics(remote=""), models=[]) self.specified_workspace = None self.use_here = None self.use_recent = None self.workspace_path = None self.workspace_type = None + self.workspace_repo = None self.skip_prompting = None def setup_workspace_manager( @@ -154,7 +188,13 @@ def setup_workspace_manager( self.specified_workspace = specified_workspace self.use_here = use_here self.use_recent = use_recent - self.workspace_path, self.workspace_type = self.get_workspace_path() + ( + self.workspace_path, + self.workspace_type, + self.workspace_repo, + ) = self.get_workspace_path() + if self.workspace_type != WorkspaceType.NOT_FOUND: + self.metadata.basics.remote = self.workspace_repo.remotes[0].url self.skip_prompting = skip_prompting def set_recent_workspace(self, path: str): @@ -187,7 +227,7 @@ def get_specified_workspace(self): return os.path.abspath(os.path.expanduser(self.specified_workspace)) - def get_workspace_path(self) -> Tuple[str, WorkspaceType]: + def get_workspace_path(self) -> Tuple[str, WorkspaceType, Optional[git.Repo]]: """ Retrieves the workspace path and type based on the following precedence: 1. Specified Workspace (--workspace) @@ -204,8 +244,9 @@ def get_workspace_path(self) -> Tuple[str, WorkspaceType]: # Check for explicitly specified workspace first specified_workspace = self.get_specified_workspace() if specified_workspace: - if check_comfy_repo(specified_workspace): - return specified_workspace, WorkspaceType.SPECIFIED + found_comfy_repo, repo = check_comfy_repo(specified_workspace) + if found_comfy_repo: + return specified_workspace, WorkspaceType.SPECIFIED, repo print( "[bold red]warn: The specified workspace is not ComfyUI directory.[/bold red]" @@ -218,8 +259,9 @@ def get_workspace_path(self) -> Tuple[str, WorkspaceType]: constants.CONFIG_KEY_RECENT_WORKSPACE ) if recent_workspace: - if check_comfy_repo(recent_workspace): - return recent_workspace, WorkspaceType.RECENT + found_comfy_repo, repo = check_comfy_repo(recent_workspace) + if found_comfy_repo: + return recent_workspace, WorkspaceType.RECENT, repo else: print( "[bold red]warn: No recent workspace has been set.[/bold red]" @@ -236,7 +278,7 @@ def get_workspace_path(self) -> Tuple[str, WorkspaceType]: current_directory = os.getcwd() found_comfy_repo, comfy_repo = check_comfy_repo(current_directory) if found_comfy_repo: - return comfy_repo.working_dir, WorkspaceType.CURRENT_DIR + return comfy_repo.working_dir, WorkspaceType.CURRENT_DIR, comfy_repo print( "[bold red]warn: you are not current in a ComfyUI directory.[/bold red]" @@ -251,30 +293,34 @@ def get_workspace_path(self) -> Tuple[str, WorkspaceType]: ) # If it's in a sub dir of the ComfyUI repo, get the repo working dir if found_comfy_repo: - return comfy_repo.working_dir, WorkspaceType.CURRENT_DIR + return comfy_repo.working_dir, WorkspaceType.CURRENT_DIR, comfy_repo # Check for user-set default workspace default_workspace = self.config_manager.get( constants.CONFIG_KEY_DEFAULT_WORKSPACE ) - if default_workspace and check_comfy_repo(default_workspace): - return default_workspace, WorkspaceType.DEFAULT + if default_workspace: + found_comfy_repo, repo = check_comfy_repo(default_workspace) + return default_workspace, WorkspaceType.DEFAULT, repo # Fallback to the most recent workspace if it exists if self.use_recent is None: recent_workspace = self.config_manager.get( constants.CONFIG_KEY_RECENT_WORKSPACE ) - if recent_workspace and check_comfy_repo(recent_workspace): - return recent_workspace, WorkspaceType.RECENT + if recent_workspace: + found_comfy_repo, repo = check_comfy_repo(recent_workspace) + if found_comfy_repo: + return recent_workspace, WorkspaceType.RECENT, repo # Check for comfy-cli default workspace default_workspace = utils.get_not_user_set_default_workspace() - if check_comfy_repo(default_workspace): - return default_workspace, WorkspaceType.DEFAULT + found_comfy_repo, repo = check_comfy_repo(default_workspace) + if found_comfy_repo: + return default_workspace, WorkspaceType.DEFAULT, repo - return None, WorkspaceType.NOT_FOUND + return None, WorkspaceType.NOT_FOUND, None def get_comfyui_manager_path(self): if self.workspace_path is None: @@ -299,11 +345,52 @@ def is_comfyui_manager_installed(self): def scan_dir(self): logging.info(f"Scanning directory: {self.workspace_path}") model_files = [] + custom_node_folders = [] + counter = 0 for root, _dirs, files in os.walk(self.workspace_path): for file in files: + counter += 1 if file.endswith(constants.SUPPORTED_PT_EXTENSIONS): - model_files.append(os.path.join(root, file)) - return model_files + model_files.append( + Model( + name=os.path.basename(file), + paths=[ModelPath(path=file)], + ) + ) + for custom_node_path in os.listdir( + os.path.join(self.workspace_path, "custom_nodes") + ): + if not os.path.isdir(custom_node_path): + continue + if custom_node_path in constants.IGNORE_CUSTOM_NODE_FOLDERS: + continue + is_git_custom_node, repo = check_folder_is_git(custom_node_path) + if is_git_custom_node: + custom_node_folders.append( + CustomNode( + path=custom_node_path, + enabled=not custom_node_path.endswith(".disabled"), + is_git=True, + remote=repo.remotes[0].url if repo.remotes else None, + ) + ) + else: + custom_node_folders.append( + CustomNode( + path=custom_node_path, + enabled=not custom_node_path.endswith(".disabled"), + is_git=False, + remote=None, + ) + ) + + self.metadata.custom_nodes = custom_node_folders + self.metadata.models = model_files + self.metadata.basics.updated_at = datetime.now() + save_yaml( + os.path.join(self.workspace_path, constants.COMFY_LOCK_YAML_FILE), + self.metadata, + ) def scan_dir_concur(self): base_path = Path(".") From 814999a0df0907c8cbc32d2a55d01eaec01390c4 Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" Date: Mon, 6 May 2024 12:12:29 +0900 Subject: [PATCH 2/3] reformat --- comfy_cli/workspace_manager.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/comfy_cli/workspace_manager.py b/comfy_cli/workspace_manager.py index 8434fbc..70cf5ff 100644 --- a/comfy_cli/workspace_manager.py +++ b/comfy_cli/workspace_manager.py @@ -96,9 +96,11 @@ def load_yaml(file_path: str) -> ComfyLockYAMLStruct: basics = Basics( remote=data.get("basics", {}).get("remote", ""), name=data.get("basics", {}).get("name"), - updated_at=datetime.fromisoformat(data.get("basics", {}).get("updated_at")) - if data.get("basics", {}).get("updated_at") - else None, + updated_at=( + datetime.fromisoformat(data.get("basics", {}).get("updated_at")) + if data.get("basics", {}).get("updated_at") + else None + ), ) models = [ From a6681d09549c7d88d8b90283f8b2bda29dab7589 Mon Sep 17 00:00:00 2001 From: Yoland Y <4950057+yoland68@users.noreply.github.com> Date: Tue, 14 May 2024 16:21:43 -0700 Subject: [PATCH 3/3] Add CLI version info --- comfy_cli/cmdline.py | 21 ++++++++++++++++++--- comfy_cli/env_checker.py | 2 ++ 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py index 8d60b3a..0449b02 100644 --- a/comfy_cli/cmdline.py +++ b/comfy_cli/cmdline.py @@ -19,7 +19,7 @@ from comfy_cli.config_manager import ConfigManager from comfy_cli.constants import GPU_OPTION from comfy_cli.env_checker import EnvChecker, check_comfy_server_running -from comfy_cli.update import check_for_updates +from comfy_cli.update import check_for_updates, get_version_from_pyproject from comfy_cli.workspace_manager import ( WorkspaceManager, WorkspaceType, @@ -94,9 +94,10 @@ def entry( skip_prompt: Annotated[ Optional[bool], typer.Option( + "--skip-prompt", show_default=False, is_flag=True, - help="Do not prompt user for input, use default options", + help="Do not prompt user for input, use default options and command line args", ), ] = None, enable_telemetry: Annotated[ @@ -108,12 +109,26 @@ def entry( help="Enable tracking", ), ] = True, + version: Annotated[ + Optional[bool], + typer.Option( + "-v", + "--version", + show_default=False, + is_flag=True, + help="Display version", + ), + ] = False, ): workspace_manager.setup_workspace_manager(workspace, here, recent, skip_prompt) tracking.prompt_tracking_consent(skip_prompt, default_value=enable_telemetry) - if ctx.invoked_subcommand is None: + if version: + version = get_version_from_pyproject() + print(f"Comfy CLI version: {version}") + + if ctx.invoked_subcommand is None and not version: print( "[bold yellow]Welcome to Comfy CLI![/bold yellow]: https://github.com/Comfy-Org/comfy-cli" ) diff --git a/comfy_cli/env_checker.py b/comfy_cli/env_checker.py index 2a426ce..9814a75 100644 --- a/comfy_cli/env_checker.py +++ b/comfy_cli/env_checker.py @@ -8,6 +8,7 @@ from rich.table import Table import requests +from comfy_cli import update from comfy_cli.utils import singleton from comfy_cli.config_manager import ConfigManager @@ -109,6 +110,7 @@ def fill_print_table(self): self.virtualenv_path if self.virtualenv_path else "Not Used", ) table.add_row("Conda Env", self.conda_env if self.conda_env else "Not Used") + table.add_row("Comfy CLI Version", update.get_version_from_pyproject()) ConfigManager().fill_print_env(table)