Preserve reimportable ComfyUI workflow metadata

__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()

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

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

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()