# Save this as ResizeToClosestBucket_alone.py in your ComfyUI custom_nodes directory

import torch
import logging
from comfy.utils import common_upscale

# Define aspect ratio constants
ASPECT_RATIO_512 = {
    "1:1": [512, 512],
    "4:3": [512, 384],
    "3:4": [384, 512],
    "16:9": [512, 288],
    "9:16": [288, 512],
    "21:9": [512, 219],
    "9:21": [219, 512]
}

def get_closest_ratio(height, width, ratios):
    """Find the closest aspect ratio bucket for given dimensions"""
    input_ratio = height / width
    closest_ratio = None
    closest_diff = float('inf')
    closest_size = None
    
    for ratio_name, size in ratios.items():
        target_ratio = size[0] / size[1]
        diff = abs(input_ratio - target_ratio)
        if diff < closest_diff:
            closest_diff = diff
            closest_ratio = ratio_name
            closest_size = size
            
    return closest_size, closest_ratio

class ResizeToClosestBucket_alone:
    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE", ),
                "base_resolution": ("INT", {
                    "min": 64,
                    "max": 1280,
                    "step": 64,
                    "default": 512,
                    "tooltip": "Base resolution, closest training data bucket resolution is chosen based on the selection."
                }),
                "upscale_method": (cls.upscale_methods, {
                    "default": "lanczos",
                    "tooltip": "Upscale method to use"
                }),
                "crop": (["disabled", "center"],),
            }
        }

    RETURN_TYPES = ("IMAGE", "INT", "INT")
    RETURN_NAMES = ("images", "width", "height")
    FUNCTION = "resize"
    CATEGORY = "CogVideoWrapper"

    def resize(self, images, base_resolution, upscale_method, crop):
        """
        Resize images to the closest aspect ratio bucket based on base_resolution
        
        Args:
            images: Input tensor of shape (B, H, W, C)
            base_resolution: Base resolution to scale from
            upscale_method: Method to use for upscaling
            crop: Whether to center crop or not
            
        Returns:
            tuple: (resized_images, width, height)
        """
        # Get input dimensions
        B, H, W, C = images.shape
        
        # Calculate aspect ratio buckets based on base_resolution
        aspect_ratio_sample_size = {
            key: [x / 512 * base_resolution for x in ASPECT_RATIO_512[key]]
            for key in ASPECT_RATIO_512.keys()
        }

        # Find closest matching size
        closest_size, closest_ratio = get_closest_ratio(H, W, ratios=aspect_ratio_sample_size)
        height, width = [int(x / 16) * 16 for x in closest_size]  # Ensure dimensions are divisible by 16
        
        logging.info(f"Closest bucket size: {width}x{height}")

        # Perform resize operation
        resized_images = images.clone().movedim(-1, 1)  # Move channels to correct position
        resized_images = common_upscale(resized_images, width, height, upscale_method, crop)
        resized_images = resized_images.movedim(1, -1)  # Move channels back
        
        return (resized_images, width, height)

# Node class mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
    "ResizeToClosestBucket_alone": ResizeToClosestBucket_alone
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ResizeToClosestBucket_alone": "Resize to Closest Bucket (Standalone)"
}