110 lines
2.8 KiB
Python
110 lines
2.8 KiB
Python
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()
|