Preserve reimportable ComfyUI workflow metadata

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