38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
import comfy.utils
|
|
import os
|
|
|
|
# Scales an image based on width and aspect ratio. Ensures it's divisible by 16
|
|
class AutoScaleImage:
|
|
ESSENTIALS_CATEGORY = "Image Tools"
|
|
upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "image": ("IMAGE",), "upscale_method": (s.upscale_methods,),
|
|
"width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 16}),}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "upscale"
|
|
|
|
CATEGORY = "image/upscaling"
|
|
|
|
def upscale(self, image, upscale_method, width):
|
|
samples = image.movedim(-1,1)
|
|
# Calculate the original height and width
|
|
original_height, original_width = samples.shape[2], samples.shape[3]
|
|
# Calculate the aspect ratio
|
|
aspect_ratio = original_height / original_width
|
|
# Scale height based on aspect ratio and round to nearest factor of 16
|
|
height = round(aspect_ratio * width / 16) * 16
|
|
s = comfy.utils.common_upscale(samples, width, height, upscale_method, "disabled")
|
|
s = s.movedim(1,-1)
|
|
return (s,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AutoScaleImage": AutoScaleImage
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AutoScaleImage": "Auto Scale Image"
|
|
}
|