Skip to content

Commit 77965f8

Browse files
author
Peyton
committed
Updated to script generation to use torch inference_mode and reverted back to using VAEDecode instead of VAEDecodeTiled
1 parent 654e6a2 commit 77965f8

File tree

1 file changed

+1
-5
lines changed

1 file changed

+1
-5
lines changed

comfyui_to_python.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -197,10 +197,6 @@ def get_class_info(class_type: str) -> (str, str, str):
197197
import_statement (str): Import statement string
198198
class_code (str): Class initialization code
199199
"""
200-
# If the class is 'VAEDecode', adjust the class name
201-
if class_type == 'VAEDecode':
202-
class_type = 'VAEDecodeTiled'
203-
204200
import_statement = class_type
205201
class_code = f'{class_type.lower()} = {class_type}()'
206202

@@ -221,7 +217,7 @@ def assemble_python_code(import_statements: set, loader_code: List[str], code: L
221217
"""
222218
static_imports = ['import random']
223219
imports_code = [f"from nodes import {', '.join([class_name for class_name in import_statements])}" ]
224-
main_function_code = f"def main():\n\t" + '\n\t'.join(loader_code) + f'\n\n\tfor q in range({queue_size}):\n\t' + '\n\t'.join(code)
220+
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)
225221
final_code = '\n'.join(static_imports + ['import sys\nsys.path.append("../")'] + imports_code + ['', main_function_code, '', 'if __name__ == "__main__":', '\tmain()'])
226222
final_code = black.format_str(final_code, mode=black.Mode())
227223

0 commit comments

Comments
 (0)