From 1fa2db89e15dd74d984f0ed1112263765df54b9e Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sat, 28 Mar 2026 13:42:44 -0500 Subject: [PATCH] Preserve reimportable ComfyUI workflow metadata --- __init__.py | 7 +++- comfyui_to_python.py | 26 +++++++++++--- js/save-as-script.js | 6 +++- tests/test_upscale_model_loader_export.py | 44 +++++++++++++++++++++-- 4 files changed, 73 insertions(+), 10 deletions(-) diff --git a/__init__.py b/__init__.py index 69ce450..10cdf1c 100644 --- a/__init__.py +++ b/__init__.py @@ -45,9 +45,14 @@ async def save_as_script(request): data = await request.json() name = data["name"] workflow = data["workflow"] + frontend_workflow = data.get("frontend_workflow") sio = StringIO() - ComfyUItoPython(workflow=workflow, output_file=sio) + ComfyUItoPython( + workflow=workflow, + frontend_workflow=frontend_workflow, + output_file=sio, + ) sio.seek(0) data = sio.read() diff --git a/comfyui_to_python.py b/comfyui_to_python.py index 201ef8c..df08184 100644 --- a/comfyui_to_python.py +++ b/comfyui_to_python.py @@ -198,6 +198,7 @@ class CodeGenerator: self, load_order: List, workflow_data: Dict, + metadata_workflow_data: Dict | None = None, queue_size: int = 10, ) -> str: """Generate the execution code based on the load order. @@ -314,6 +315,7 @@ class CodeGenerator: special_functions_code, code, workflow_data, + metadata_workflow_data or workflow_data, queue_size, custom_nodes, ) @@ -384,6 +386,7 @@ class CodeGenerator: speical_functions_code: List[str], code: List[str], workflow_data: Dict, + metadata_workflow_data: Dict, queue_size: int, custom_nodes=False, ) -> str: @@ -393,7 +396,8 @@ class CodeGenerator: import_statements (Dict[str, set]): Import statements grouped by module. speical_functions_code (List[str]): A list of special functions code strings. code (List[str]): A list of code strings. - workflow_data (Dict): The original workflow data for metadata preservation. + workflow_data (Dict): The API workflow data used for runtime prompt execution. + metadata_workflow_data (Dict): The workflow metadata to embed into saved outputs. queue_size (int): Number of photos that will be generated by the script. custom_nodes (bool): Whether to include custom nodes in the code. @@ -427,9 +431,9 @@ class CodeGenerator: "prompt = json.loads(json.dumps(workflow))", ( "extra_pnginfo = " - '{"workflow": {"source": "workflow_api", "payload": json.loads(' - + json.dumps(json.dumps(workflow_data)) - + ")}}" + '{"workflow": json.loads(' + + json.dumps(json.dumps(metadata_workflow_data)) + + ")}" ), ] ) @@ -572,6 +576,7 @@ class ComfyUItoPython: def __init__( self, workflow: str = "", + frontend_workflow: str | Dict | None = None, input_file: str = "", output_file: str | TextIO = "", queue_size: int = 1, @@ -597,6 +602,7 @@ class ComfyUItoPython: raise ValueError("Needs output_file") self.workflow = workflow + self.frontend_workflow = frontend_workflow self.input_file = input_file self.output_file = output_file self.queue_size = queue_size @@ -622,6 +628,13 @@ class ComfyUItoPython: else: data = json.loads(self.workflow) + frontend_workflow_data = None + if self.frontend_workflow: + if isinstance(self.frontend_workflow, str): + frontend_workflow_data = json.loads(self.frontend_workflow) + else: + frontend_workflow_data = self.frontend_workflow + # Step 2: Initialize extra/custom nodes when requested or when the workflow references # a node class that is not currently loaded in the runtime. missing_node_types = { @@ -642,7 +655,10 @@ class ComfyUItoPython: self.node_class_mappings, self.base_node_class_mappings ) generated_code = code_generator.generate_workflow( - load_order, data, queue_size=self.queue_size + load_order, + data, + frontend_workflow_data, + queue_size=self.queue_size, ) # Step 5: Write the generated code to a file diff --git a/js/save-as-script.js b/js/save-as-script.js index 91a1a84..d09a0a3 100644 --- a/js/save-as-script.js +++ b/js/save-as-script.js @@ -29,7 +29,11 @@ const extension = { } 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 + const json = JSON.stringify({ + name: filename + ".json", + workflow: JSON.stringify(p.output, null, 2), + frontend_workflow: JSON.stringify(app.graph.serialize(), 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"}); diff --git a/tests/test_upscale_model_loader_export.py b/tests/test_upscale_model_loader_export.py index 754c69a..656a3e8 100644 --- a/tests/test_upscale_model_loader_export.py +++ b/tests/test_upscale_model_loader_export.py @@ -91,19 +91,57 @@ class UpscaleModelLoaderExportTest(unittest.TestCase): generated = output.getvalue() - self.assertIn("from nodes import (", generated) + self.assertIn("from nodes import NODE_CLASS_MAPPINGS", generated) + self.assertIn("from tests.test_upscale_model_loader_export import (", generated) self.assertIn("LoadImage,", generated) self.assertIn("UpscaleModelLoader,", generated) self.assertIn("ImageUpscaleWithModel,", generated) - self.assertIn("NODE_CLASS_MAPPINGS,", generated) self.assertIn("loadimage = LoadImage()", generated) self.assertIn("upscalemodelloader = UpscaleModelLoader()", generated) self.assertIn( - "imageupscalewithmodel_3 = imageupscalewithmodel.upscale(", + "imageupscalewithmodel__3 = imageupscalewithmodel.upscale(", generated, ) self.assertNotIn('NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()', generated) + def test_frontend_workflow_metadata_is_embedded_for_reimport(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": { + "image": "example.png", + }, + } + } + frontend_workflow = { + "version": 0.4, + "last_node_id": 1, + "last_link_id": 0, + "nodes": [], + "links": [], + "groups": [], + "config": {}, + "extra": {}, + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + frontend_workflow=json.dumps(frontend_workflow), + output_file=output, + node_class_mappings={ + "LoadImage": LoadImage, + }, + ) + + generated = output.getvalue() + + self.assertIn('extra_pnginfo = {', generated) + self.assertIn('"workflow": json.loads(', generated) + self.assertIn('"version": 0.4', generated) + self.assertIn('"nodes": []', generated) + self.assertNotIn('"source": "workflow_api"', generated) + if __name__ == "__main__": unittest.main()