diff --git a/comfy_cli/command/custom_nodes/command.py b/comfy_cli/command/custom_nodes/command.py index bac6c16f..c99cb196 100644 --- a/comfy_cli/command/custom_nodes/command.py +++ b/comfy_cli/command/custom_nodes/command.py @@ -15,6 +15,7 @@ extract_node_configuration, upload_file_to_signed_url, zip_files, + initialize_project_config, ) app = typer.Typer() @@ -434,3 +435,17 @@ def publish( # Upload the zip file to the signed URL typer.echo("Uploading zip file...") upload_file_to_signed_url(signed_url, zip_filename) + + +@app.command("init", help="Init scaffolding for custom node") +@tracking.track_command("node") +def scaffold(): + if os.path.exists("comfynode.toml"): + typer.echo("Warning: 'comfynode.toml' already exists. Will not overwrite.") + raise typer.Exit(code=1) + + typer.echo("Initializing metadata...") + initialize_project_config() + typer.echo( + "comfynode.toml created successfully. Defaults were filled in. Please check before publishing." + ) diff --git a/comfy_cli/registry/__init__.py b/comfy_cli/registry/__init__.py index 2ba64a35..26cc3027 100644 --- a/comfy_cli/registry/__init__.py +++ b/comfy_cli/registry/__init__.py @@ -1,7 +1,6 @@ from .api import publish_node_version, upload_file_to_signed_url -# Import specific functions from the config_parser module -from .config_parser import extract_node_configuration +from .config_parser import extract_node_configuration, initialize_project_config from .types import PyProjectConfig, PublishNodeVersionResponse, NodeVersion from .zip import zip_files @@ -13,4 +12,5 @@ "NodeVersion", "zip_files", "upload_file_to_signed_url", + "initialize_project_config", ] diff --git a/comfy_cli/registry/config_parser.py b/comfy_cli/registry/config_parser.py index f9f8c3f8..7c682308 100644 --- a/comfy_cli/registry/config_parser.py +++ b/comfy_cli/registry/config_parser.py @@ -1,4 +1,6 @@ import os + +import tomlkit.exceptions from comfy_cli.registry.types import ( PyProjectConfig, ProjectConfig, @@ -6,6 +8,103 @@ Model, ComfyConfig, ) +import tomlkit +import subprocess + + +def create_comfynode_config(): + # Create the initial structure of the TOML document + document = tomlkit.document() + + project = tomlkit.table() + project["name"] = "" + project["description"] = "" + project["version"] = "1.0.0" + project["dependencies"] = tomlkit.aot() + project["license"] = "LICENSE" + + urls = tomlkit.table() + urls["Repository"] = "" + + project.add("urls", urls) + document.add("project", project) + + tool = tomlkit.table() + comfy = tomlkit.table() + comfy["PublisherId"] = "" + comfy["DisplayName"] = "" + comfy["Icon"] = "" + + # Add the default model + models = tomlkit.array() + model = tomlkit.inline_table() + model["location"] = "/checkpoints/model.safetensor" + model["model_url"] = "https://example.com/model.zip" + models.append(model) + comfy["Models"] = models + + tool.add("comfy", comfy) + document.add("tool", tool) + + # Write the TOML document to a file + try: + with open("comfynode.toml", "w") as toml_file: + toml_file.write(tomlkit.dumps(document)) + print("comfynode.toml has been created successfully in the current directory.") + except IOError as e: + raise Exception(f"Failed to write 'comfynode.toml': {str(e)}") + + +def initialize_project_config(): + create_comfynode_config() + + with open("comfynode.toml", "r") as file: + document = tomlkit.parse(file.read()) + + # Get the current git remote URL + try: + git_remote_url = ( + subprocess.check_output(["git", "remote", "get-url", "origin"]) + .decode() + .strip() + ) + except subprocess.CalledProcessError: + raise Exception( + "Could not retrieve Git remote URL. Are you in a Git repository?" + ) + + repo_name = git_remote_url.split("/")[-1].replace(".git", "") + + project = document.get("project", tomlkit.table()) + urls = project.get("urls", tomlkit.table()) + urls["Repository"] = git_remote_url + project["urls"] = urls + project["name"] = repo_name.lower() + project["description"] = "" + project["version"] = "1.0.0" + project["license"] = "LICENSE" + + tool = document.get("tool", tomlkit.table()) + comfy = tool.get("comfy", tomlkit.table()) + comfy["DisplayName"] = repo_name + tool["comfy"] = comfy + document["tool"] = tool + + # Handle dependencies + if os.path.exists("requirements.txt"): + with open("requirements.txt", "r") as req_file: + dependencies = [line.strip() for line in req_file if line.strip()] + project["dependencies"] = dependencies + else: + print("Warning: 'requirements.txt' not found. No dependencies will be added.") + + # Write the updated config to a new file in the current directory + try: + with open("comfynode.toml", "w") as toml_file: + toml_file.write(tomlkit.dumps(document)) + print("comfynode.toml has been created successfully in the current directory.") + except IOError as e: + raise IOError(f"Failed to write 'comfynode.toml': {str(e)}") def extract_node_configuration(