146 lines
5.7 KiB
Python
146 lines
5.7 KiB
Python
import numpy as np
|
|
from typing import Tuple
|
|
|
|
class ConstrainResolution:
|
|
"""
|
|
A ComfyUI node that analyzes images and suggests optimal dimensions
|
|
while preserving aspect ratio.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""Define input types with tooltips"""
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", {
|
|
"tooltip": "Input image to analyze for resolution constraints"
|
|
}),
|
|
"min_res": ("INT", {
|
|
"default": 704,
|
|
"min": 1,
|
|
"tooltip": "Minimum resolution in pixels for both width and height"
|
|
}),
|
|
"max_res": ("INT", {
|
|
"default": 1280,
|
|
"min": 1,
|
|
"tooltip": "Maximum resolution in pixels for both width and height"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "INT", "FLOAT", "FLOAT")
|
|
RETURN_NAMES = ("image_passthrough", "constrained_width", "constrained_height", "constrained_aspect_ratio", "original_aspect_ratio")
|
|
|
|
# Add tooltips for outputs
|
|
RETURN_TYPES_TOOLTIPS = {
|
|
"image_passthrough": "The unchanged input image (passed through)",
|
|
"constrained_width": "Suggested width that fits within min/max constraints while preserving aspect ratio",
|
|
"constrained_height": "Suggested height that fits within min/max constraints while preserving aspect ratio",
|
|
"constrained_aspect_ratio": "The aspect ratio of the constrained dimensions",
|
|
"original_aspect_ratio": "The original aspect ratio of the input image"
|
|
}
|
|
|
|
FUNCTION = "process"
|
|
CATEGORY = "image/resolution"
|
|
|
|
@staticmethod
|
|
def calculate_aspect_ratio(width: int, height: int) -> float:
|
|
return round(width / height, 4)
|
|
|
|
@staticmethod
|
|
def calculate_optimal_dimensions(width: int, height: int, min_res: int, max_res: int) -> Tuple[int, int]:
|
|
aspect_ratio = width / height
|
|
|
|
# Calculate maximum allowed pixels based on max_res x min_res
|
|
max_pixels = max_res * min_res
|
|
|
|
if abs(aspect_ratio - 1.0) < 0.01: # Square image (with small tolerance)
|
|
# Calculate optimal square dimension based on max pixels
|
|
optimal_dim = int(np.sqrt(max_pixels))
|
|
# Round up to nearest even number
|
|
optimal_dim = optimal_dim + (optimal_dim % 2)
|
|
# Ensure it's within min/max bounds
|
|
optimal_dim = min(max(optimal_dim, min_res), max_res)
|
|
return optimal_dim, optimal_dim
|
|
|
|
if aspect_ratio > 1: # Landscape
|
|
# First try to set width to max_res
|
|
width = max_res
|
|
height = int(width / aspect_ratio)
|
|
|
|
# If total pixels exceed max_pixels, scale down proportionally
|
|
if width * height > max_pixels:
|
|
width = int(np.sqrt(max_pixels * aspect_ratio))
|
|
# Round up to nearest even number
|
|
width = width + (width % 2)
|
|
height = int(width / aspect_ratio)
|
|
# Round up to nearest even number
|
|
height = height + (height % 2)
|
|
|
|
# Check minimum bounds
|
|
if height < min_res:
|
|
height = min_res
|
|
width = int(height * aspect_ratio)
|
|
width = width + (width % 2)
|
|
else: # Portrait
|
|
# First try to set height to max_res
|
|
height = max_res
|
|
width = int(height * aspect_ratio)
|
|
|
|
# If total pixels exceed max_pixels, scale down proportionally
|
|
if width * height > max_pixels:
|
|
height = int(np.sqrt(max_pixels / aspect_ratio))
|
|
# Round up to nearest even number
|
|
height = height + (height % 2)
|
|
width = int(height * aspect_ratio)
|
|
# Round up to nearest even number
|
|
width = width + (width % 2)
|
|
|
|
# Check minimum bounds
|
|
if width < min_res:
|
|
width = min_res
|
|
height = int(width / aspect_ratio)
|
|
height = height + (height % 2)
|
|
|
|
# Final bounds check
|
|
width = min(max(width, min_res), max_res)
|
|
height = min(max(height, min_res), max_res)
|
|
|
|
# Final pixel count check (just in case rounding pushed us over)
|
|
while width * height > max_pixels:
|
|
if width > height:
|
|
width = width - 2 # Reduce by 2 to keep even numbers
|
|
else:
|
|
height = height - 2
|
|
|
|
return width, height
|
|
|
|
def process(self, image, min_res, max_res):
|
|
# Get image dimensions
|
|
if len(image.shape) == 4:
|
|
height = image.shape[1]
|
|
width = image.shape[2]
|
|
else:
|
|
height = image.shape[0]
|
|
width = image.shape[1]
|
|
|
|
# Calculate original aspect ratio
|
|
original_aspect_ratio = self.calculate_aspect_ratio(width, height)
|
|
|
|
# Calculate optimal dimensions
|
|
suggested_width, suggested_height = self.calculate_optimal_dimensions(width, height, min_res, max_res)
|
|
|
|
# Calculate final aspect ratio
|
|
aspect_ratio = self.calculate_aspect_ratio(suggested_width, suggested_height)
|
|
|
|
# Return image first, followed by other outputs
|
|
return (image, suggested_width, suggested_height, aspect_ratio, original_aspect_ratio)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ConstrainResolution": ConstrainResolution
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ConstrainResolution": "Constrain Resolution"
|
|
}
|