104 lines
3.4 KiB
Python
104 lines
3.4 KiB
Python
"""ComfyUI node implementation for ImageToMultipleOf."""
|
|
|
|
from typing import Dict, Any, Tuple
|
|
|
|
from torch import Tensor
|
|
|
|
from ...base import ComfyAssetsBaseNode
|
|
from .logic import process_image_to_multiple_of
|
|
|
|
|
|
class ImageToMultipleOfNode(ComfyAssetsBaseNode):
|
|
"""
|
|
Adjusts image dimensions to be multiples of a specified value.
|
|
|
|
Useful for models that require specific dimension constraints.
|
|
Supports both center cropping and rescaling methods.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"multiple_of": (
|
|
"INT",
|
|
{
|
|
"default": 64,
|
|
"min": 1,
|
|
"max": 256,
|
|
"step": 16,
|
|
"display": "number",
|
|
},
|
|
),
|
|
"method": (["center crop", "rescale"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "🫶 ComfyAssets/🖼️ Resolution"
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "process"
|
|
|
|
def process(self, image: Tensor, multiple_of: int, method: str) -> Tuple[Tensor]:
|
|
"""
|
|
Process image to ensure dimensions are multiples of specified value.
|
|
|
|
Args:
|
|
image: Input image tensor
|
|
multiple_of: Value that dimensions should be multiple of
|
|
method: Processing method - "center crop" or "rescale"
|
|
|
|
Returns:
|
|
Tuple containing processed image tensor
|
|
"""
|
|
try:
|
|
self.validate_inputs(image=image, multiple_of=multiple_of, method=method)
|
|
|
|
# Process the image
|
|
processed_image = process_image_to_multiple_of(image, multiple_of, method)
|
|
|
|
_, new_height, new_width, _ = processed_image.shape
|
|
self.log_info(
|
|
f"Processed image from {image.shape[1]}x{image.shape[2]} "
|
|
f"to {new_height}x{new_width} (multiple of {multiple_of}) "
|
|
f"using {method}"
|
|
)
|
|
|
|
return (processed_image,)
|
|
|
|
except Exception as e:
|
|
self.handle_error(f"Failed to process image: {str(e)}", e)
|
|
|
|
def validate_inputs(self, **kwargs) -> None:
|
|
"""Validate inputs for ImageToMultipleOf node."""
|
|
image = kwargs.get("image")
|
|
multiple_of = kwargs.get("multiple_of")
|
|
method = kwargs.get("method")
|
|
|
|
if image is None:
|
|
raise ValueError("Image input is required")
|
|
|
|
if not isinstance(image, Tensor) or len(image.shape) != 4:
|
|
raise ValueError(
|
|
f"Expected image tensor with shape (batch, height, width, channels), "
|
|
f"got shape {image.shape if isinstance(image, Tensor) else 'non-tensor'}"
|
|
)
|
|
|
|
if multiple_of <= 0:
|
|
raise ValueError(f"multiple_of must be positive, got {multiple_of}")
|
|
|
|
if method not in ["center crop", "rescale"]:
|
|
raise ValueError(f"Invalid method: {method}")
|
|
|
|
# Check if resulting dimensions would be too small
|
|
_, height, width, _ = image.shape
|
|
new_height = height - (height % multiple_of)
|
|
new_width = width - (width % multiple_of)
|
|
|
|
if new_height <= 0 or new_width <= 0:
|
|
raise ValueError(
|
|
f"Image dimensions ({height}x{width}) are too small "
|
|
f"to be adjusted to multiple of {multiple_of}"
|
|
)
|