Narrow runtime compatibility fix

Author: async-metric

README.md

@@ -96,7 +96,7 @@ Flags:
 - `--output_file`: output Python file, default `workflow_api.py`
 - `--queue_size`: default execution count in the generated script, default `10`
 
-![Dev Mode Options](images/dev_mode_options.png)
+![Dev Mode Options](images/dev_mode_options.PNG)
 
 ## Generated Scripts
 

comfyui_to_python.py

@@ -209,7 +209,7 @@ def generate_workflow(
         """
         # Create the necessary data structures to hold imports and generated code
         import_statements, executed_variables, special_functions_code, code = (
-            {"nodes": {"NODE_CLASS_MAPPINGS"}},
+            set(["NODE_CLASS_MAPPINGS"]),
             {},
             [],
             [],
@@ -245,9 +245,8 @@ def generate_workflow(
                 )
                 initialized_objects[class_type] = self.clean_variable_name(class_type)
                 if class_type 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:
+                    import_statements.add(import_statement)
+                if class_type not in self.base_node_class_mappings.keys():
                     custom_nodes = True
                 special_functions_code.append(class_code)
 
@@ -362,7 +361,7 @@ def format_arg(self, key: str, value: any) -> str:
 
     def assemble_python_code(
         self,
-        import_statements: Dict[str, set],
+        import_statements: set,
         speical_functions_code: List[str],
         code: List[str],
         queue_size: int,
@@ -371,7 +370,7 @@ def assemble_python_code(
         """Generates the final code string.
 
         Args:
-            import_statements (Dict[str, set]): Import statements grouped by module.
+            import_statements (set): A set of unique import statements.
             speical_functions_code (List[str]): A list of special functions code strings.
             code (List[str]): A list of code strings.
             queue_size (int): Number of photos that will be generated by the script.
@@ -409,10 +408,9 @@ def assemble_python_code(
         else:
             custom_nodes = ""
         # Create import statements for node classes
-        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}")
+        imports_code = [
+            f"from nodes import {', '.join([class_name for class_name in import_statements])}"
+        ]
         # Assemble the main function code, including custom nodes if applicable
         main_function_code = (
             "def main():\n\t"
@@ -432,31 +430,20 @@ def assemble_python_code(
 
         return final_code
 
-    def get_class_info(self, class_type: str) -> Tuple[str, Tuple[str, str], str]:
+    def get_class_info(self, class_type: str) -> Tuple[str, str, str]:
         """Generates and returns necessary information about class type.
 
         Args:
             class_type (str): Class type.
 
         Returns:
-            Tuple[str, Tuple[str, str], str]: Updated class type, import statement and initialization code.
+            Tuple[str, str, str]: Updated class type, import statement string, class 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)
-        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)
+        import_statement = class_type
+        if class_type in self.base_node_class_mappings.keys():
             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