diff --git a/__init__.py b/__init__.py index 227d315..0888cb3 100644 --- a/__init__.py +++ b/__init__.py @@ -5,6 +5,9 @@ from .src.comfymath.float import NODE_CLASS_MAPPINGS as float_NCM from .src.comfymath.number import NODE_CLASS_MAPPINGS as number_NCM from .src.comfymath.vec import NODE_CLASS_MAPPINGS as vec_NCM from .src.comfymath.control import NODE_CLASS_MAPPINGS as control_NCM +from .src.comfymath.graphics import NODE_CLASS_MAPPINGS as graphics_NCM + + NODE_CLASS_MAPPINGS = { @@ -15,6 +18,7 @@ NODE_CLASS_MAPPINGS = { **number_NCM, **vec_NCM, **control_NCM, + **graphics_NCM, } diff --git a/src/comfymath/graphics.py b/src/comfymath/graphics.py new file mode 100644 index 0000000..6a8d4c1 --- /dev/null +++ b/src/comfymath/graphics.py @@ -0,0 +1,77 @@ +from typing import Any, Mapping + + +SDXL_SUPPORTED_RESOLUTIONS = [ + (1024, 1024, 1.0), + (1152, 896, 1.2857142857142858), + (896, 1152, 0.7777777777777778), + (1216, 832, 1.4615384615384615), + (832, 1216, 0.6842105263157895), + (1344, 768, 1.75), + (768, 1344, 0.5714285714285714), + (1536, 640, 2.4), + (640, 1536, 0.4166666666666667), +] + + +class SDXLResolution: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "resolution": ( + [f"{res[0]}x{res[1]}" for res in SDXL_SUPPORTED_RESOLUTIONS], + ) + } + } + + RETURN_TYPES = ("INT", "INT") + RETURN_NAMES = ("width", "height") + FUNCTION = "op" + CATEGORY = "math/graphics" + + def op(self, resolution: str) -> tuple[int, int]: + width, height = resolution.split("x") + return (int(width), int(height)) + + +class NearestSDXLResolution: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return {"required": {"image": ("IMAGE",)}} + + RETURN_TYPES = ("INT", "INT") + RETURN_NAMES = ("width", "height") + FUNCTION = "op" + CATEGORY = "math/graphics" + + def op(self, image) -> tuple[int, int]: + image_width = image.size()[2] + image_height = image.size()[1] + print(f"Input image resolution: {image_width}x{image_height}") + image_ratio = image_width / image_height + differences = [ + (abs(image_ratio - resolution[2]), resolution) + for resolution in SDXL_SUPPORTED_RESOLUTIONS + ] + smallest = None + for difference in differences: + if smallest is None: + smallest = difference + else: + if difference[0] < smallest[0]: + smallest = difference + if smallest is not None: + width = smallest[1][0] + height = smallest[1][1] + else: + width = 1024 + height = 1024 + print(f"Selected SDXL resolution: {width}x{height}") + return (width, height) + + +NODE_CLASS_MAPPINGS = { + "CM_SDXLResolution": SDXLResolution, + "CM_NearestSDXLResolution": NearestSDXLResolution, +}