From 42d0526661409c77ae0120bb2eea6944fe119119 Mon Sep 17 00:00:00 2001
From: Peyton DeNiro <25550995+pydn@users.noreply.github.com>
Date: Sun, 29 Mar 2026 15:15:52 -0500
Subject: [PATCH] Fix generated workflow torch import order

Move generated torch imports until after add_extra_model_paths runs so exported scripts do not trigger ComfyUI's early-torch warning during startup. Update export assertions and checked-in runtime generated fixtures to match the new main() ordering.
---
 comfyui_to_python/generator/render.py         |  4 +--
 tests/runtime/generated/text-to-image.py      |  4 +--
 .../runtime/generated/upscale-model-loader.py | 23 ++++++++---------
 tests/test_upscale_model_loader_export.py     | 25 ++++++++++++++++---
 4 files changed, 36 insertions(+), 20 deletions(-)

diff --git a/comfyui_to_python/generator/render.py b/comfyui_to_python/generator/render.py
index e3c6238..1c531bc 100644
--- a/comfyui_to_python/generator/render.py
+++ b/comfyui_to_python/generator/render.py
@@ -76,8 +76,6 @@ class WorkflowRenderer:
             "# Workflow execution",
             "def main():",
             "    bootstrap_comfyui_runtime()",
-            "    import torch",
-            "",
             "    add_extra_model_paths()",
         ]
         if custom_nodes_call:
@@ -87,6 +85,8 @@ class WorkflowRenderer:
             execution_section.extend(f"    {line}" for line in imports_code)
         execution_section.extend(
             [
+                "",
+                "    import torch",
                 "",
                 "    with torch.inference_mode():",
             ]
diff --git a/tests/runtime/generated/text-to-image.py b/tests/runtime/generated/text-to-image.py
index 5f9b0e0..e57de44 100644
--- a/tests/runtime/generated/text-to-image.py
+++ b/tests/runtime/generated/text-to-image.py
@@ -165,8 +165,6 @@ extra_pnginfo = build_extra_pnginfo()
 # Workflow execution
 def main():
     bootstrap_comfyui_runtime()
-    import torch
-
     add_extra_model_paths()
 
     # Node imports
@@ -180,6 +178,8 @@ def main():
         VAEDecode,
     )
 
+    import torch
+
     with torch.inference_mode():
         checkpointloadersimple = CheckpointLoaderSimple()
         checkpointloadersimple_1 = checkpointloadersimple.load_checkpoint(
diff --git a/tests/runtime/generated/upscale-model-loader.py b/tests/runtime/generated/upscale-model-loader.py
index b7b0a09..ed56d7a 100644
--- a/tests/runtime/generated/upscale-model-loader.py
+++ b/tests/runtime/generated/upscale-model-loader.py
@@ -105,14 +105,6 @@ def bootstrap_comfyui_runtime() -> None:
         os.environ["OCL_SET_SVM_SIZE"] = "262144"
 
 
-# Runtime support
-
-bootstrap_comfyui_runtime()
-import torch
-
-add_extra_model_paths()
-
-
 def import_custom_nodes() -> None:
     """Initialize ComfyUI custom nodes in the exporter runtime."""
     comfyui_path = get_comfyui_path()
@@ -166,14 +158,19 @@ workflow = build_workflow()
 prompt = json.loads(json.dumps(workflow))
 extra_pnginfo = build_extra_pnginfo()
 
+
 # Workflow execution
-# Node imports
-from nodes import LoadImage, NODE_CLASS_MAPPINGS, SaveImage
-
-
 def main():
+    bootstrap_comfyui_runtime()
+    add_extra_model_paths()
+    import_custom_nodes()
+
+    # Node imports
+    from nodes import LoadImage, NODE_CLASS_MAPPINGS, SaveImage
+
+    import torch
+
     with torch.inference_mode():
-        import_custom_nodes()
         loadimage = LoadImage()
         loadimage_1 = loadimage.load_image(image="e2e_upscale_input.png")
         upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()
diff --git a/tests/test_upscale_model_loader_export.py b/tests/test_upscale_model_loader_export.py
index d5d2643..804de12 100644
--- a/tests/test_upscale_model_loader_export.py
+++ b/tests/test_upscale_model_loader_export.py
@@ -149,9 +149,28 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
         self.assertIn("comfy.options.enable_args_parsing()", generated)
         self.assertIn("import cuda_malloc", generated)
         self.assertNotIn("\nbootstrap_comfyui_runtime()\n", generated)
-        self.assertIn("def main():\n    bootstrap_comfyui_runtime()\n    import torch", generated)
-        self.assertLess(generated.index("def bootstrap_comfyui_runtime()"), generated.index("def main():"))
-        self.assertLess(generated.index("bootstrap_comfyui_runtime()"), generated.index("import torch"))
+        self.assertIn(
+            "def main():\n    bootstrap_comfyui_runtime()\n    add_extra_model_paths()",
+            generated,
+        )
+        self.assertLess(
+            generated.index("def bootstrap_comfyui_runtime()"), generated.index("def main():")
+        )
+        main_section = generated[generated.index("def main():") :]
+        self.assertIn(
+            "def main():\n"
+            "    bootstrap_comfyui_runtime()\n"
+            "    add_extra_model_paths()",
+            main_section,
+        )
+        self.assertLess(
+            main_section.index("bootstrap_comfyui_runtime()"),
+            main_section.index("add_extra_model_paths()"),
+        )
+        self.assertLess(
+            main_section.index("add_extra_model_paths()"),
+            main_section.index("import torch"),
+        )
         self.assertLess(generated.index("import cuda_malloc"), generated.index("import torch"))
 
     def test_generated_module_import_does_not_parse_cli_args(self):