110 lines
3.8 KiB
Python
110 lines
3.8 KiB
Python
import math
|
|
from PIL import Image as PILImage
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
class SmartResize:
|
|
"""
|
|
Resize images so the longest side equals max_pixels.
|
|
Supports lanczos interpolation and multiple_of snapping.
|
|
Works with single images and batches.
|
|
"""
|
|
|
|
INTERP_MAP = {
|
|
"lanczos": PILImage.LANCZOS,
|
|
"bicubic": PILImage.BICUBIC,
|
|
"bilinear": PILImage.BILINEAR,
|
|
"nearest": PILImage.NEAREST,
|
|
}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"max_pixels": ("INT", {
|
|
"default": 1536,
|
|
"min": 64,
|
|
"max": 8192,
|
|
"step": 8,
|
|
"tooltip": "Longest side will be resized to this value (aspect ratio preserved)"
|
|
}),
|
|
"multiple_of": ("INT", {
|
|
"default": 8,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1,
|
|
"tooltip": "Round dimensions to a multiple of this value (8 = SD compatible)"
|
|
}),
|
|
"interpolation": (["lanczos", "bicubic", "bilinear", "nearest"], {
|
|
"default": "lanczos"
|
|
}),
|
|
"upscale_if_smaller": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Also upscale images that are already smaller than max_pixels"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "INT")
|
|
RETURN_NAMES = ("image", "width", "height")
|
|
FUNCTION = "run"
|
|
CATEGORY = "image/transform"
|
|
|
|
def _resize_one(self, img_tensor, max_pixels, multiple_of, interp, upscale_if_smaller):
|
|
"""Resize a single [H, W, C] tensor."""
|
|
h, w = img_tensor.shape[0], img_tensor.shape[1]
|
|
longest = max(w, h)
|
|
|
|
# Skip if already small enough and upscale disabled
|
|
if longest <= max_pixels and not upscale_if_smaller:
|
|
# Still snap to multiple_of
|
|
new_w = max(multiple_of, round(w / multiple_of) * multiple_of) if multiple_of > 1 else w
|
|
new_h = max(multiple_of, round(h / multiple_of) * multiple_of) if multiple_of > 1 else h
|
|
if new_w == w and new_h == h:
|
|
return img_tensor, w, h
|
|
else:
|
|
scale = max_pixels / longest
|
|
new_w = int(w * scale)
|
|
new_h = int(h * scale)
|
|
|
|
# Snap to multiple_of
|
|
if multiple_of > 1:
|
|
new_w = max(multiple_of, round(new_w / multiple_of) * multiple_of)
|
|
new_h = max(multiple_of, round(new_h / multiple_of) * multiple_of)
|
|
|
|
# Convert tensor [H, W, C] float32 → PIL → resize → tensor
|
|
arr = (img_tensor.numpy() * 255).astype(np.uint8)
|
|
pil = PILImage.fromarray(arr, mode="RGB")
|
|
pil_resized = pil.resize((new_w, new_h), interp)
|
|
arr_resized = np.array(pil_resized).astype(np.float32) / 255.0
|
|
return torch.from_numpy(arr_resized), new_w, new_h
|
|
|
|
def run(self, image, max_pixels, multiple_of, interpolation, upscale_if_smaller):
|
|
# image shape: [B, H, W, C]
|
|
interp = self.INTERP_MAP[interpolation]
|
|
results = []
|
|
last_w, last_h = 0, 0
|
|
|
|
for i in range(image.shape[0]):
|
|
resized, w, h = self._resize_one(
|
|
image[i], max_pixels, multiple_of, interp, upscale_if_smaller
|
|
)
|
|
results.append(resized)
|
|
last_w, last_h = w, h
|
|
|
|
# Stack back to [B, H, W, C] — all frames must be same size
|
|
# (they will be if all input frames are same size, which is the ComfyUI norm)
|
|
output = torch.stack(results, dim=0)
|
|
return (output, last_w, last_h)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SmartResize": SmartResize,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SmartResize": "Smart Resize 📐",
|
|
}
|