Skip to content

Commit

Permalink
feat: comfynode.toml init command.
Browse files Browse the repository at this point in the history
  • Loading branch information
robinjhuang committed May 3, 2024
1 parent 3bc9f64 commit 18ebbeb
Show file tree
Hide file tree
Showing 3 changed files with 116 additions and 2 deletions.
15 changes: 15 additions & 0 deletions comfy_cli/command/custom_nodes/command.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
extract_node_configuration,
upload_file_to_signed_url,
zip_files,
initialize_project_config,
)

app = typer.Typer()
Expand Down Expand Up @@ -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."
)
4 changes: 2 additions & 2 deletions comfy_cli/registry/__init__.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -13,4 +12,5 @@
"NodeVersion",
"zip_files",
"upload_file_to_signed_url",
"initialize_project_config",
]
99 changes: 99 additions & 0 deletions comfy_cli/registry/config_parser.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,110 @@
import os

import tomlkit.exceptions
from comfy_cli.registry.types import (
PyProjectConfig,
ProjectConfig,
URLs,
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(
Expand Down

0 comments on commit 18ebbeb

Please sign in to comment.