Fix upscale model loader export bootstrap

2026-03-28 13:08:34 -05:00
parent 32d42bf472
commit 77b98ad5e1
4 changed files with 272 additions and 11 deletions
@@ -0,0 +1,109 @@
+import json
+import unittest
+from io import StringIO
+
+from comfyui_to_python import ComfyUItoPython
+
+
+class UpscaleModelLoader:
+    CATEGORY = "loaders"
+    FUNCTION = "load_model"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model_name": ("STRING",),
+            }
+        }
+
+    def load_model(self, model_name):
+        return (model_name,)
+
+
+class LoadImage:
+    CATEGORY = "image"
+    FUNCTION = "load_image"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("STRING",),
+            }
+        }
+
+    def load_image(self, image):
+        return (image,)
+
+
+class ImageUpscaleWithModel:
+    CATEGORY = "image/upscaling"
+    FUNCTION = "upscale"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "upscale_model": ("UPSCALE_MODEL",),
+                "image": ("IMAGE",),
+            }
+        }
+
+    def upscale(self, upscale_model, image):
+        return (image,)
+
+
+class UpscaleModelLoaderExportTest(unittest.TestCase):
+    def test_upscale_workflow_uses_direct_upscale_model_loader_init(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {
+                    "image": "example.png",
+                },
+            },
+            "2": {
+                "class_type": "UpscaleModelLoader",
+                "inputs": {
+                    "model_name": "RealESRGAN_x4plus.safetensors",
+                },
+            },
+            "3": {
+                "class_type": "ImageUpscaleWithModel",
+                "inputs": {
+                    "upscale_model": ["2", 0],
+                    "image": ["1", 0],
+                },
+            },
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={
+                "LoadImage": LoadImage,
+                "UpscaleModelLoader": UpscaleModelLoader,
+                "ImageUpscaleWithModel": ImageUpscaleWithModel,
+            },
+        )
+
+        generated = output.getvalue()
+
+        self.assertIn("from nodes import (", generated)
+        self.assertIn("LoadImage,", generated)
+        self.assertIn("UpscaleModelLoader,", generated)
+        self.assertIn("ImageUpscaleWithModel,", generated)
+        self.assertIn("NODE_CLASS_MAPPINGS,", generated)
+        self.assertIn("loadimage = LoadImage()", generated)
+        self.assertIn("upscalemodelloader = UpscaleModelLoader()", generated)
+        self.assertIn(
+            "imageupscalewithmodel_3 = imageupscalewithmodel.upscale(",
+            generated,
+        )
+        self.assertNotIn('NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()', generated)
+
+
+if __name__ == "__main__":
+    unittest.main()