From e6d17777461c2e89de5abf5c61c084244af7365a Mon Sep 17 00:00:00 2001
From: AbyssYuan0 <jiaxiangchang0@gmail.com>
Date: Tue, 24 Oct 2023 12:24:08 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E5=9B=BE=E7=89=87=E6=8C=89?=
 =?UTF-8?q?=E6=9D=A1=E4=BB=B6=E7=BC=A9=E6=94=BE=E2=80=9C=E6=9C=80=E7=9F=AD?=
 =?UTF-8?q?=E8=BE=B9=E2=80=9D?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 __init__.py | 77 ++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 76 insertions(+), 1 deletion(-)

diff --git a/__init__.py b/__init__.py
index d4b6617..0d29923 100644
--- a/__init__.py
+++ b/__init__.py
@@ -3,7 +3,12 @@ import math
 from PIL import Image
 import numpy as np
 import torch
+import comfy.utils
 
+def getImageSize(IMAGE) -> tuple[int, int]:
+    samples = IMAGE.movedim(-1, 1)
+    size = samples.shape[3], samples.shape[2]
+    return size
 
 def tensorToImg(imageTensor):
     imaget = imageTensor[0]
@@ -241,7 +246,76 @@ class ImageNormalization:
 
         return (nw, nh, top, left, bottom, right,)
 
+class ImageScaleToSide:
+    upscale_methods = ["nearest-exact", "bilinear", "area"]
+    crop_methods = ["disabled", "center"]
 
+    def __init__(self) -> None:
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "side_length": ("INT", {
+                    "default": 1,
+                    "min": 1,
+                    "max": 4096,
+                    "step": 1,
+                    "display": "number"
+                }),
+                "side": (["Longest", "Shortest", "Width", "Height"],),
+                "upscale_method": (cls.upscale_methods,),
+                "crop": (cls.crop_methods,)}}
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "imageUpscaleToSide"
+
+    CATEGORY = "badger"
+
+    def imageUpscaleToSide(self, image, upscale_method, side_length: int, side: str, crop):
+        samples = image.movedim(-1, 1)
+
+        size = getImageSize(image)
+
+        width_B = int(size[0])
+        height_B = int(size[1])
+
+        width = width_B
+        height = height_B
+
+        def determineSide(_side: str) -> tuple[int, int]:
+            width, height = 0, 0
+            if _side == "Width":
+                heigh_ratio = height_B / width_B
+                width = side_length
+                height = heigh_ratio * width
+            elif _side == "Height":
+                width_ratio = width_B / height_B
+                height = side_length
+                width = width_ratio * height
+            return width, height
+
+        if side == "Longest":
+            if width > height:
+                width, height = determineSide("Width")
+            else:
+                width, height = determineSide("Height")
+        elif side == "Shortest":
+            if width < height:
+                width, height = determineSide("Width")
+            else:
+                width, height = determineSide("Height")
+        else:
+            width, height = determineSide(side)
+
+        width = math.ceil(width)
+        height = math.ceil(height)
+
+        cls = comfy.utils.common_upscale(samples, width, height, upscale_method, crop)
+        cls = cls.movedim(1, -1)
+        return (cls,)
 
 
 
@@ -251,7 +325,8 @@ NODE_CLASS_MAPPINGS = {
     "FloatToInt-badger": FloatToInt,
     "IntToString-badger": IntToString,
     "FloatToString-badger": FloatToString,
-    "ImageNormalization-badger": ImageNormalization
+    "ImageNormalization-badger": ImageNormalization,
+    "ImageScaleToSide-badger": ImageScaleToSide
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {