Merge pull request #74 from pydn/atmaranto-main

Add UI button to save script directly form ComfyUI web app

Author: pydn

__init__.py

@@ -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

@@ -6,7 +6,7 @@
 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 @@
     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 @@ def read_json_file(file_path: str) -> dict:
             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 @@ def write_code_to_file(file_path: str, code: str) -> None:
         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 @@ def __init__(self, node_class_mappings: Dict, base_node_class_mappings: Dict):
     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 @@ def execute(self):
         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 @@ def execute(self):
             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,8 +2,6 @@
 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

@@ -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

install.py

@@ -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");
+	}
+});

js/save-as-script.js

@@ -1,13 +1 @@
-torch
-torchdiffeq
-torchsde
-einops
-transformers>=4.25.1
-safetensors>=0.3.0
-aiohttp
-accelerate
-pyyaml
-Pillow
-scipy
-tqdm
 black