diff --git a/LICENSE b/LICENSE index a8907c1..6cecb61 100644 --- a/LICENSE +++ b/LICENSE @@ -1,5 +1,6 @@ The MIT License (MIT) +Copyright (c) 2024 Anthony Maranto Copyright (c) 2023 Peyton DeNiro Permission is hereby granted, free of charge, to any person obtaining a copy diff --git a/README.md b/README.md index aa3adc3..3e7ddce 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,39 @@ -## ComfyUI-to-Python-Extension +## ComfyUI-to-Python-Extension (SaveAsScript version) + +This is a fork of the [ComfyUI-To-Python-Extension](https://github.com/pydn/ComfyUI-to-Python-Extension) by Peyton DeNiro. It provides a more robust command-line interface and the ability to export your current workflow as a script directly from the ComfyUI web interface. + +### Usage (Web) + +Upon installation, there will be a button labeled "Save as Script" on the interface, pictured below: + +![A button that exports the current script on the ComfyUI](images/save_as_script.png) + +Click that button and enter your preferred name for the downloaded script (no ".py" is appeneded automatically), and your file will be downloaded. + +Note that if you use the script this way, you do **not** need to enable developer mode. You *will* if you use the CLI (to get the API JSON). + +### Usage (CLI) + +Navigating to the extension directory, you can also use the basic CLI included in the script to save your script without changing the Python file. + +```bash +usage: comfyui_to_python.py [-h] [--queue-size QUEUE_SIZE] [--yes] workflow output + +Converts a ComfyUI-style workflow.json file to a Python file. Must have been exported with API calls + +positional arguments: + workflow The workflow.json file to convert + output The output file (defaults to [input file].py) + +options: + -h, --help show this help message and exit + --queue-size QUEUE_SIZE, -q QUEUE_SIZE + The queue size per run + --yes, --overwrite, -y + Overwrite the output file if it exists +``` + +## Original Description The `ComfyUI-to-Python-Extension` is a powerful tool that translates [ComfyUI](https://github.com/comfyanonymous/ComfyUI) workflows into executable Python code. Designed to bridge the gap between ComfyUI's visual interface and Python's programming environment, this script facilitates the seamless transition from design to code execution. Whether you're a data scientist, a software developer, or an AI enthusiast, this tool streamlines the process of implementing ComfyUI workflows in Python. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..ca899b4 --- /dev/null +++ b/__init__.py @@ -0,0 +1,42 @@ +import sys +import os + +from io import StringIO +import json + +import traceback + +from aiohttp import web + +ext_dir = os.path.dirname(__file__) +sys.path.append(ext_dir) + +# Prevent reimporting of custom nodes +os.environ["RUNNING_IN_COMFYUI"] = "TRUE" + +from comfyui_to_python import ComfyUItoPython + +sys.path.append(os.path.dirname(os.path.dirname(ext_dir))) + +import server + +WEB_DIRECTORY = "js" +NODE_CLASS_MAPPINGS = {} + +@server.PromptServer.instance.routes.post("/saveasscript") +async def save_as_script(request): + try: + data = await request.json() + name = data['name'] + workflow = data['workflow'] + + sio = StringIO() + ComfyUItoPython(workflow=workflow, output_file=sio) + + sio.seek(0) + data = sio.read() + + return web.Response(text=data, status=200) + except Exception as e: + traceback.print_exc() + return web.Response(text=str(e), status=500) diff --git a/comfyui_to_python.py b/comfyui_to_python.py index 33e6e53..47ae90e 100644 --- a/comfyui_to_python.py +++ b/comfyui_to_python.py @@ -6,25 +6,26 @@ import os import random import sys import re -from typing import Dict, List, Any, Callable, Tuple +from typing import Dict, List, Any, Callable, Tuple, TextIO import black -from utils import import_custom_nodes, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths, get_value_at_index +from comfyui_to_python_utils import import_custom_nodes, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths, get_value_at_index -sys.path.append('../') +add_comfyui_directory_to_sys_path() from nodes import NODE_CLASS_MAPPINGS +import nodes class FileHandler: """Handles reading and writing files. - This class provides methods to read JSON data from an input file and write code to an output file. + This class provides methods to read JSON data from an input file and write code to an output file (either file-like objects or string paths). """ @staticmethod - def read_json_file(file_path: str) -> dict: + def read_json_file(file_path: str | TextIO) -> dict: """ Reads a JSON file and returns its contents as a dictionary. @@ -39,33 +40,13 @@ class FileHandler: ValueError: If the file is not a valid JSON. """ - try: - with open(file_path, 'r') as file: - data = json.load(file) - return data - - except FileNotFoundError: - # Get the directory from the file_path - directory = os.path.dirname(file_path) - - # If the directory is an empty string (which means file is in the current directory), - # get the current working directory - if not directory: - directory = os.getcwd() - - # Find all JSON files in the directory - json_files = glob.glob(f"{directory}/*.json") - - # Format the list of JSON files as a string - json_files_str = "\n".join(json_files) - - raise FileNotFoundError(f"\n\nFile not found: {file_path}. JSON files in the directory:\n{json_files_str}") - - except json.JSONDecodeError: - raise ValueError(f"Invalid JSON format in file: {file_path}") + if hasattr(file_path, "read"): return json.load(file_path) + with open(file_path, 'r') as file: + data = json.load(file) + return data @staticmethod - def write_code_to_file(file_path: str, code: str) -> None: + def write_code_to_file(file_path: str | TextIO, code: str) -> None: """Write the specified code to a Python file. Args: @@ -75,16 +56,19 @@ class FileHandler: Returns: None """ - # Extract directory from the filename - directory = os.path.dirname(file_path) + if isinstance(file_path, str): + # Extract directory from the filename + directory = os.path.dirname(file_path) - # If the directory does not exist, create it - if directory and not os.path.exists(directory): - os.makedirs(directory) + # If the directory does not exist, create it + if directory and not os.path.exists(directory): + os.makedirs(directory) - # Save the code to a .py file - with open(file_path, 'w') as file: - file.write(code) + # Save the code to a .py file + with open(file_path, 'w') as file: + file.write(code) + else: + file_path.write(code) class LoadOrderDeterminer: @@ -181,6 +165,13 @@ class CodeGenerator: """ self.node_class_mappings = node_class_mappings self.base_node_class_mappings = base_node_class_mappings + + def can_be_imported(self, import_name: str): + if import_name in self.base_node_class_mappings.keys(): + if getattr(nodes, import_name, None) is not None: + return True + + return False def generate_workflow(self, load_order: List, filename: str = 'generated_code_workflow.py', queue_size: int = 10) -> str: """Generate the execution code based on the load order. @@ -212,10 +203,10 @@ class CodeGenerator: # No need to use preview image nodes since we are executing the script in a terminal if class_type == 'PreviewImage': continue - + class_type, import_statement, class_code = self.get_class_info(class_type) initialized_objects[class_type] = self.clean_variable_name(class_type) - if class_type in self.base_node_class_mappings.keys(): + if self.can_be_imported(class_type): import_statements.add(import_statement) if class_type not in self.base_node_class_mappings.keys(): custom_nodes = True @@ -337,7 +328,7 @@ class CodeGenerator: """ import_statement = class_type variable_name = self.clean_variable_name(class_type) - if class_type in self.base_node_class_mappings.keys(): + if self.can_be_imported(class_type): class_code = f'{variable_name} = {class_type.strip()}()' else: class_code = f'{variable_name} = NODE_CLASS_MAPPINGS["{class_type}"]()' @@ -408,19 +399,27 @@ class ComfyUItoPython: base_node_class_mappings (Dict): Base mappings of node classes. """ - def __init__(self, input_file: str, output_file: str, queue_size: int = 10, node_class_mappings: Dict = NODE_CLASS_MAPPINGS): + def __init__(self, workflow: str = "", input_file: str = "", output_file: (str | TextIO) = "", queue_size: int = 10, node_class_mappings: Dict = NODE_CLASS_MAPPINGS): """Initialize the ComfyUItoPython class with the given parameters. Args: + workflow (str): The input_file (str): Path to the input JSON file. output_file (str): Path to the output Python file. queue_size (int): The number of times a workflow will be executed by the script. Defaults to 10. node_class_mappings (Dict): Mappings of node classes. Defaults to NODE_CLASS_MAPPINGS. """ + if input_file and workflow: + raise ValueError("Can't provide both input_file and workflow") + elif not input_file and not workflow: + raise ValueError("Needs input_file or workflow") + + self.workflow = workflow self.input_file = input_file self.output_file = output_file self.queue_size = queue_size self.node_class_mappings = node_class_mappings + self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings) self.execute() @@ -430,11 +429,15 @@ class ComfyUItoPython: Returns: None """ - # Step 1: Import all custom nodes - import_custom_nodes() + #if self.should_init_custom_nodes: + # Step 1: Import all custom nodes (already done for us) + #import_custom_nodes() # Step 2: Read JSON data from the input file - data = FileHandler.read_json_file(self.input_file) + if self.input_file: + data = FileHandler.read_json_file(self.input_file) + else: + data = json.loads(self.workflow) # Step 3: Determine the load order load_order_determiner = LoadOrderDeterminer(data, self.node_class_mappings) @@ -451,10 +454,24 @@ class ComfyUItoPython: if __name__ == '__main__': - # Update class parameters here - input_file = 'workflow_api.json' - output_file = 'workflow_api.py' - queue_size = 10 + import argparse + + ap = argparse.ArgumentParser(description="Converts a ComfyUI-style workflow.json file to a Python file. Must have been exported with API calls") + + ap.add_argument("workflow", help="The workflow.json file to convert") + ap.add_argument("output", default=None, help="The output file (defaults to [input file].py)") + ap.add_argument("--queue-size", "-q", default=1, type=int, help="The queue size per run") + ap.add_argument("--yes", "--overwrite", "-y", action="store_true", help="Overwrite the output file if it exists") + + args = ap.parse_args() + + + output_file = args.output_file if args.output_file else args.input_file + ".py" + if os.path.isfile(output_file): + if not args.yes: + if input("Are you sure you want to overwrite " + output_file + "?\nY/n").strip().lower() not in ("y", "yes"): + print("Exiting.") + sys.exit(1) # Convert ComfyUI workflow to Python - ComfyUItoPython(input_file=input_file, output_file=output_file, queue_size=queue_size) + ComfyUItoPython(input_file=args.input_file, output_file=output_file, queue_size=args.queue_size) diff --git a/utils.py b/comfyui_to_python_utils.py similarity index 100% rename from utils.py rename to comfyui_to_python_utils.py diff --git a/images/save_as_script.png b/images/save_as_script.png new file mode 100644 index 0000000..3b8dc75 Binary files /dev/null and b/images/save_as_script.png differ diff --git a/install.py b/install.py new file mode 100644 index 0000000..8d6ba79 --- /dev/null +++ b/install.py @@ -0,0 +1,20 @@ +import os +import sys + +from subprocess import Popen, check_output, PIPE + +requirements = open(os.path.join(os.path.dirname(__file__), "requirements.txt")).read().split("\n") + +installed_packages = check_output( + [sys.executable, "-m", "pip", "list"], + universal_newlines=True +).split("\n") + +installed_packages = set([package.split(" ")[0].lower() for package in installed_packages if package.strip()]) + +for requirement in requirements: + if requirement.lower() not in installed_packages: + print(f"Installing requirements...") + Popen([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"], stdout=PIPE, stderr=PIPE, cwd=os.path.dirname(__file__)).communicate() + print(f"Installed.") + break diff --git a/js/save-as-script.js b/js/save-as-script.js new file mode 100644 index 0000000..60432a6 --- /dev/null +++ b/js/save-as-script.js @@ -0,0 +1,52 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; +import { $el } from "../../scripts/ui.js"; + +app.registerExtension({ + name: "Comfy.SaveAsScript", + init() { + $el("style", { + parent: document.head, + }); + }, + async setup() { + const menu = document.querySelector(".comfy-menu"); + const separator = document.createElement("hr"); + + separator.style.margin = "20px 0"; + separator.style.width = "100%"; + menu.append(separator); + + const saveButton = document.createElement("button"); + saveButton.textContent = "Save as Script"; + saveButton.onclick = () => { + const filename = prompt("Save script as:"); + if(filename === undefined || filename === null || filename === "") { + return + } + + app.graphToPrompt().then(async (p) => { + const json = JSON.stringify({name: filename, workflow: JSON.stringify(p.output, null, 2)}, null, 2); // convert the data to a JSON string + var response = await api.fetchApi(`/saveasscript`, { method: "POST", body: json }); + if(response.status == 200) { + const blob = new Blob([await response.text()], {type: "application/json"}); + const url = URL.createObjectURL(blob); + const a = $el("a", { + href: url, + download: filename, + style: {display: "none"}, + parent: document.body, + }); + a.click(); + setTimeout(function () { + a.remove(); + window.URL.revokeObjectURL(url); + }, 0); + } + }); + } + menu.append(saveButton); + + console.log("SaveAsScript loaded"); + } +});