Add prompt data and handle kwargs

Author: atmaranto

comfyui_to_python.py

@@ -166,7 +166,7 @@ class CodeGenerator:
         base_node_class_mappings (Dict): Base mappings of node classes.
     """
 
-    def __init__(self, node_class_mappings: Dict, base_node_class_mappings: Dict):
+    def __init__(self, node_class_mappings: Dict, base_node_class_mappings: Dict, prompt: Dict):
         """Initialize the CodeGenerator with given node class mappings.
 
         Args:
@@ -175,6 +175,7 @@ def __init__(self, node_class_mappings: Dict, base_node_class_mappings: Dict):
         """
         self.node_class_mappings = node_class_mappings
         self.base_node_class_mappings = base_node_class_mappings
+        self.prompt = prompt
 
     def can_be_imported(self, import_name: str):
         if import_name in self.base_node_class_mappings.keys():
@@ -195,6 +196,7 @@ def generate_workflow(self, load_order: List, queue_size: int = 1) -> str:
         Returns:
             str: Generated execution code as a string.
         """
+        include_prompt_data = False
         # Create the necessary data structures to hold imports and generated code
         import_statements, executed_variables, arg_inputs, special_functions_code, code = set(['NODE_CLASS_MAPPINGS']), {}, [], [], []
         # This dictionary will store the names of the objects that we have already initialized
@@ -206,8 +208,9 @@ def generate_workflow(self, load_order: List, queue_size: int = 1) -> str:
             # Generate class definition and inputs from the data
             inputs, class_type = data['inputs'], data['class_type']
 
+            input_types = self.node_class_mappings[class_type].INPUT_TYPES()
             missing = []
-            for i, input in enumerate(self.node_class_mappings[class_type].INPUT_TYPES().get("required", {}).keys()):
+            for i, input in enumerate(input_types.get("required", {}).keys()):
                 if input not in inputs:
                     input_var = f"{input}{len(arg_inputs)+1}"
                     arg_inputs.append((input_var, f"Argument {i}, input `{input}` for node \\\"{data['_meta'].get('title', class_type)}\\\" id {idx}"))
@@ -233,14 +236,18 @@ def generate_workflow(self, load_order: List, queue_size: int = 1) -> str:
 
             # Get all possible parameters for class_def
             class_def_params = self.get_function_parameters(getattr(class_def, class_def.FUNCTION))
+            no_params = class_def_params is None
 
             # Remove any keyword arguments from **inputs if they are not in class_def_params
-            inputs = {key: value for key, value in inputs.items() if key in class_def_params}
+            inputs = {key: value for key, value in inputs.items() if no_params or key in class_def_params}
             for input, input_var, arg in missing:
                 inputs[input] = {"variable_name": f"parse_arg(args." + input_var + ")"}
             # Deal with hidden variables
-            if 'unique_id' in class_def_params:
+            if no_params or 'unique_id' in class_def_params:
                 inputs['unique_id'] = random.randint(1, 2**64)
+            if no_params or 'prompt' in class_def_params:
+                inputs["prompt"] = {"variable_name": "PROMPT_DATA"}
+                include_prompt_data = True
 
             # Create executed variable and generate code
             executed_variables[idx] = f'{self.clean_variable_name(class_type)}_{idx}'
@@ -261,7 +268,7 @@ def generate_workflow(self, load_order: List, queue_size: int = 1) -> str:
                     code.append(self.create_function_call_code(initialized_objects[class_type], class_def.FUNCTION, executed_variables[idx], is_special_function, **inputs))
 
         # Generate final code by combining imports and code, and wrap them in a main function
-        final_code = self.assemble_python_code(import_statements, special_functions_code, arg_inputs, code, queue_size, custom_nodes)
+        final_code = self.assemble_python_code(import_statements, special_functions_code, arg_inputs, code, queue_size, custom_nodes, include_prompt_data)
 
         return final_code
 
@@ -304,7 +311,7 @@ def format_arg(self, key: str, value: any) -> str:
             return f'{key}={value["variable_name"]}'
         return f'{key}={value}'
 
-    def assemble_python_code(self, import_statements: set, special_functions_code: List[str], arg_inputs: List[Tuple[str, str]], code: List[str], queue_size: int, custom_nodes=False) -> str:
+    def assemble_python_code(self, import_statements: set, special_functions_code: List[str], arg_inputs: List[Tuple[str, str]], code: List[str], queue_size: int, custom_nodes=False, include_prompt_data=True) -> str:
         """Generates the final code string.
 
         Args:
@@ -349,6 +356,8 @@ def assemble_python_code(self, import_statements: set, special_functions_code: L
         # Define static import statements required for the script
         static_imports = ['import os', 'import random', 'import sys', 'import json', 'import argparse', 'import contextlib', 'from typing import Sequence, Mapping, Any, Union', 
                           'import torch'] + func_strings + argparse_code
+        if include_prompt_data:
+            static_imports.append(f'PROMPT_DATA = json.loads({repr(json.dumps(self.prompt))})')
         # Check if custom nodes should be included
         if custom_nodes:
             static_imports.append(f'\n{inspect.getsource(import_custom_nodes)}\n')
@@ -459,7 +468,8 @@ def get_function_parameters(self, func: Callable) -> List:
         signature = inspect.signature(func)
         parameters = {name: param.default if param.default != param.empty else None
                     for name, param in signature.parameters.items()}
-        return list(parameters.keys())  
+        catch_all = any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values())
+        return list(parameters.keys()) if not catch_all else None
 
     def update_inputs(self, inputs: Dict, executed_variables: Dict) -> Dict:
         """Update inputs based on the executed variables.
@@ -542,7 +552,7 @@ def execute(self):
         load_order = load_order_determiner.determine_load_order()
 
         # Step 4: Generate the workflow code
-        code_generator = CodeGenerator(self.node_class_mappings, self.base_node_class_mappings)
+        code_generator = CodeGenerator(self.node_class_mappings, self.base_node_class_mappings, data)
         generated_code = code_generator.generate_workflow(load_order, queue_size=self.queue_size)
 
         # Step 5: Write the generated code to a file