nodes/pydn-ComfyUI-to-Python-Extension
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #74 from pydn/atmaranto-main

Browse Source 
Add UI button to save script directly form ComfyUI web app
This commit is contained in:
Peyton DeNiro
2024-09-14 14:01:52 -05:00
committed by GitHub
parent d0c9b00336 aa5416634c
commit af9a0dfe75
7 changed files with 193 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+58
View File
@@ -0,0 +1,58 @@
import sys
import os
from io import StringIO
import traceback
from aiohttp import web
ext_dir = os.path.dirname(__file__)
sys.path.append(ext_dir)
try:
import black
except ImportError:
print("Unable to import requirements for ComfyUI-SaveAsScript.")
print("Installing...")
import importlib
spec = importlib.util.spec_from_file_location(
"impact_install", os.path.join(os.path.dirname(__file__), "install.py")
)
impact_install = importlib.util.module_from_spec(spec)
spec.loader.exec_module(impact_install)
print("Successfully installed. Hopefully, at least.")
# 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)
comfyui_to_python.py
+59 -54
View File
@@ -6,7 +6,7 @@ import os
import random
import sys
import re
from typing import Dict, List, Callable, Tuple
from typing import Dict, List, Any, Callable, Tuple, TextIO
from argparse import ArgumentParser
import black
@@ -20,7 +20,7 @@ from comfyui_to_python_utils import (
get_value_at_index,
)
sys.path.append("../")
add_comfyui_directory_to_sys_path()
from nodes import NODE_CLASS_MAPPINGS
@@ -36,7 +36,7 @@ class FileHandler:
"""
@staticmethod
def read_json_file(file_path: str) -> dict:
def read_json_file(file_path: str | TextIO, encoding: str = "utf-8") -> dict:
"""
Reads a JSON file and returns its contents as a dictionary.
@@ -51,35 +51,14 @@ 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", encoding="utf-8") 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:
@@ -89,16 +68,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", encoding="utf-8") as file:
file.write(code)
else:
file_path.write(code)
class LoadOrderDeterminer:
@@ -203,15 +185,12 @@ class CodeGenerator:
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.
Args:
load_order (List): A list of tuples representing the load order.
filename (str): The name of the Python file to which the code should be saved.
Defaults to 'generated_code_workflow.py'.
queue_size (int): The number of photos that will be created by the script.
Returns:
@@ -515,23 +494,37 @@ class ComfyUItoPython:
def __init__(
self,
input_file: str,
output_file: str,
queue_size: int = 10,
workflow: str = "",
input_file: str = "",
output_file: str | TextIO = "",
queue_size: int = 1,
node_class_mappings: Dict = NODE_CLASS_MAPPINGS,
needs_init_custom_nodes: bool = False,
):
"""Initialize the ComfyUItoPython class with the given parameters.
"""Initialize the ComfyUItoPython class with the given parameters. Exactly one of workflow or input_file must be specified.
Args:
workflow (str): The workflow's JSON.
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.
output_file (str | TextIO): Path to the output file or a file-like object.
queue_size (int): The number of times a workflow will be executed by the script. Defaults to 1.
node_class_mappings (Dict): Mappings of node classes. Defaults to NODE_CLASS_MAPPINGS.
needs_init_custom_nodes (bool): Whether to initialize custom nodes. Defaults to False.
"""
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")
if not output_file:
raise ValueError("Needs output_file")
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.needs_init_custom_nodes = needs_init_custom_nodes
self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings)
self.execute()
@@ -541,11 +534,18 @@ class ComfyUItoPython:
Returns:
None
"""
# Step 1: Import all custom nodes
import_custom_nodes()
# Step 1: Import all custom nodes if we need to
if self.needs_init_custom_nodes:
import_custom_nodes()
else:
# If they're already imported, we don't know which nodes are custom nodes, so we need to import all of them
self.base_node_class_mappings = {}
# 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)
@@ -556,7 +556,7 @@ class ComfyUItoPython:
self.node_class_mappings, self.base_node_class_mappings
)
generated_code = code_generator.generate_workflow(
load_order, filename=self.output_file, queue_size=self.queue_size
load_order, queue_size=self.queue_size
)
# Step 5: Write the generated code to a file
@@ -582,7 +582,12 @@ def run(
Returns:
None
"""
ComfyUItoPython(input_file, output_file, queue_size)
ComfyUItoPython(
input_file=input_file,
output_file=output_file,
queue_size=queue_size,
needs_init_custom_nodes=True,
)
def main() -> None:
@@ -612,7 +617,7 @@ def main() -> None:
default=DEFAULT_QUEUE_SIZE,
)
pargs = parser.parse_args()
ComfyUItoPython(**vars(pargs))
run(**vars(pargs))
print("Done.")
comfyui_to_python_utils.py
-2
View File
@@ -2,8 +2,6 @@ import os
from typing import Sequence, Mapping, Any, Union
import sys
sys.path.append("../")
def import_custom_nodes() -> None:
"""Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS
images/save_as_script.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 58 KiB

install.py
+20
View File
@@ -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
js/save-as-script.js
+56
View File
@@ -0,0 +1,56 @@
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 = () => {
var filename = prompt("Save script as:");
if(filename === undefined || filename === null || filename === "") {
return
}
app.graphToPrompt().then(async (p) => {
const json = JSON.stringify({name: filename + ".json", 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: "text/python;charset=utf-8"});
const url = URL.createObjectURL(blob);
if(!filename.endsWith(".py")) {
filename += ".py";
}
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");
}
});
requirements.txt
-12
View File
@@ -1,13 +1 @@
torch
torchdiffeq
torchsde
einops
transformers>=4.25.1
safetensors>=0.3.0
aiohttp
accelerate
pyyaml
Pillow
scipy
tqdm
black