Add runtime E2E validation harness

comfyui_to_python.py

@@ -196,6 +196,7 @@ class CodeGenerator:
     def generate_workflow(
         self,
         load_order: List,
+        workflow_data: Dict,
         queue_size: int = 10,
     ) -> str:
         """Generate the execution code based on the load order.
@@ -209,7 +210,7 @@ class CodeGenerator:
         """
         # Create the necessary data structures to hold imports and generated code
         import_statements, executed_variables, special_functions_code, code = (
-            set(["NODE_CLASS_MAPPINGS"]),
+            {"nodes": {"NODE_CLASS_MAPPINGS"}},
             {},
             [],
             [],
@@ -245,8 +246,9 @@ class CodeGenerator:
                 )
                 initialized_objects[class_type] = self.clean_variable_name(class_type)
                 if class_type in self.base_node_class_mappings.keys():
-                    import_statements.add(import_statement)
-                if class_type not in self.base_node_class_mappings.keys():
+                    module_name, import_name = import_statement
+                    import_statements.setdefault(module_name, set()).add(import_name)
+                if 'NODE_CLASS_MAPPINGS["' in class_code:
                     custom_nodes = True
                 special_functions_code.append(class_code)
 
@@ -268,6 +270,11 @@ class CodeGenerator:
                 and "unique_id" in input_types["hidden"].keys()
             ):
                 inputs["unique_id"] = random.randint(1, 2**64)
+            if "hidden" in input_types.keys():
+                if "prompt" in input_types["hidden"].keys():
+                    inputs["prompt"] = {"variable_name": "prompt"}
+                if "extra_pnginfo" in input_types["hidden"].keys():
+                    inputs["extra_pnginfo"] = {"variable_name": "extra_pnginfo"}
             elif class_def_params is not None:
                 if "unique_id" in class_def_params:
                     inputs["unique_id"] = random.randint(1, 2**64)
@@ -302,6 +309,7 @@ class CodeGenerator:
             import_statements,
             special_functions_code,
             code,
+            workflow_data,
             queue_size,
             custom_nodes,
         )
@@ -361,18 +369,20 @@ class CodeGenerator:
 
     def assemble_python_code(
         self,
-        import_statements: set,
+        import_statements: Dict[str, set],
         speical_functions_code: List[str],
         code: List[str],
+        workflow_data: Dict,
         queue_size: int,
         custom_nodes=False,
     ) -> str:
         """Generates the final code string.
 
         Args:
-            import_statements (set): A set of unique import statements.
+            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.
             queue_size (int): Number of photos that will be generated by the script.
             custom_nodes (bool): Whether to include custom nodes in the code.
 
@@ -392,6 +402,7 @@ class CodeGenerator:
         # Define static import statements required for the script
         static_imports = (
             [
+                "import json",
                 "import os",
                 "import random",
                 "import sys",
@@ -399,7 +410,17 @@ class CodeGenerator:
                 "import torch",
             ]
             + func_strings
-            + ["\n\nadd_comfyui_directory_to_sys_path()\nadd_extra_model_paths()\n"]
+            + [
+                "\n\nadd_comfyui_directory_to_sys_path()\nadd_extra_model_paths()\n",
+                f"workflow = json.loads({json.dumps(json.dumps(workflow_data))})",
+                "prompt = json.loads(json.dumps(workflow))",
+                (
+                    "extra_pnginfo = "
+                    '{"workflow": {"source": "workflow_api", "payload": json.loads('
+                    + json.dumps(json.dumps(workflow_data))
+                    + ")}}"
+                ),
+            ]
         )
         # Check if custom nodes should be included
         if custom_nodes:
@@ -408,9 +429,10 @@ class CodeGenerator:
         else:
             custom_nodes = ""
         # Create import statements for node classes
-        imports_code = [
-            f"from nodes import {', '.join([class_name for class_name in import_statements])}"
-        ]
+        imports_code = []
+        for module_name in sorted(import_statements.keys()):
+            class_names = ", ".join(sorted(import_statements[module_name]))
+            imports_code.append(f"from {module_name} import {class_names}")
         # Assemble the main function code, including custom nodes if applicable
         main_function_code = (
             "def main():\n\t"
@@ -430,20 +452,31 @@ class CodeGenerator:
 
         return final_code
 
-    def get_class_info(self, class_type: str) -> Tuple[str, str, str]:
+    def get_class_info(self, class_type: str) -> Tuple[str, Tuple[str, str], str]:
         """Generates and returns necessary information about class type.
 
         Args:
             class_type (str): Class type.
 
         Returns:
-            Tuple[str, str, str]: Updated class type, import statement string, class initialization code.
+            Tuple[str, Tuple[str, str], str]: Updated class type, import statement and initialization code.
         """
+        class_obj = self.base_node_class_mappings.get(class_type)
+        module_name = "nodes"
+        if class_obj is not None:
+            module_name = class_obj.__module__
         variable_name = self.clean_variable_name(class_type)
-        import_statement = class_type
-        if class_type in self.base_node_class_mappings.keys():
+        is_importable_module = bool(
+            module_name
+            and "/" not in module_name
+            and "\\" not in module_name
+            and all(part.isidentifier() for part in module_name.split("."))
+        )
+        if class_type in self.base_node_class_mappings.keys() and is_importable_module:
+            import_statement = (module_name, class_type)
             class_code = f"{variable_name} = {class_type.strip()}()"
         else:
+            import_statement = ("nodes", "NODE_CLASS_MAPPINGS")
             class_code = f'{variable_name} = NODE_CLASS_MAPPINGS["{class_type}"]()'
 
         return class_type, import_statement, class_code
@@ -596,7 +629,7 @@ class ComfyUItoPython:
             self.node_class_mappings, self.base_node_class_mappings
         )
         generated_code = code_generator.generate_workflow(
-            load_order, queue_size=self.queue_size
+            load_order, data, queue_size=self.queue_size
         )
 
         # Step 5: Write the generated code to a file