From badcdf8f177b8b421c4cb88f40ea640bfffb3192 Mon Sep 17 00:00:00 2001 From: Peyton Date: Mon, 7 Aug 2023 22:39:14 -0500 Subject: [PATCH] Forgot import and fixed proper tabbing for code generation. --- comfyui_to_python.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/comfyui_to_python.py b/comfyui_to_python.py index b2585c2..7e79674 100644 --- a/comfyui_to_python.py +++ b/comfyui_to_python.py @@ -215,13 +215,12 @@ def assemble_python_code(import_statements: set, loader_code: List[str], code: L Returns: final_code (str): Generated final code as a string """ - static_imports = ['import random'] + static_imports = ['import random', 'import torch'] imports_code = [f"from nodes import {', '.join([class_name for class_name in import_statements])}" ] - main_function_code = f"def main():\n\t" + 'with torch.inference_mode():\n\t' + '\n\t'.join(loader_code) + f'\n\n\tfor q in range({queue_size}):\n\t' + '\n\t'.join(code) + main_function_code = f"def main():\n\t" + 'with torch.inference_mode():\n\t\t' + '\n\t\t'.join(loader_code) + f'\n\n\t\tfor q in range({queue_size}):\n\t\t' + '\n\t\t'.join(code) final_code = '\n'.join(static_imports + ['import sys\nsys.path.append("../")'] + imports_code + ['', main_function_code, '', 'if __name__ == "__main__":', '\tmain()']) final_code = black.format_str(final_code, mode=black.Mode()) - return final_code