Upload bucket_final.py

bucket_final.py

@@ -0,0 +1,124 @@
+# Save this as ResizeToClosetBucket.py in your ComfyUI custom_nodes directory
+
+import torch
+import logging
+from comfy.utils import common_upscale
+
+# Define aspect ratio constants for bucket mode
+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 ResizeToClosetBucket:
+    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: In bucket mode, chooses closest training bucket; In proportional mode, targets longest side."
+                }),
+                "mode": (["bucket", "proportional"], {
+                    "default": "bucket",
+                    "tooltip": "Bucket: Use predefined aspect ratios; Proportional: Maintain original proportions"
+                }),
+                "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 = "ImageProcessing"
+
+    def resize(self, images, base_resolution, mode, upscale_method, crop):
+        """
+        Resize images either to closest bucket or maintaining proportions
+        
+        Args:
+            images: Input tensor of shape (B, H, W, C)
+            base_resolution: Target resolution
+            mode: "bucket" or "proportional"
+            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
+        aspect_ratio = H / W
+
+        if mode == "bucket":
+            # Bucket mode: Use predefined aspect ratios
+            aspect_ratio_sample_size = {
+                key: [x / 512 * base_resolution for x in ASPECT_RATIO_512[key]]
+                for key in ASPECT_RATIO_512.keys()
+            }
+            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]
+            logging.info(f"Bucket mode - Closest bucket size: {width}x{height}")
+        
+        else:  # mode == "proportional"
+            # Proportional mode: Maintain original aspect ratio
+            if H > W:
+                target_height = base_resolution
+                target_width = int(target_height / aspect_ratio)
+            else:
+                target_width = base_resolution
+                target_height = int(target_width * aspect_ratio)
+            
+            height = int(round(target_height / 16) * 16)
+            width = int(round(target_width / 16) * 16)
+            height = max(64, height)
+            width = max(64, width)
+            logging.info(f"Proportional mode - Resized to: {width}x{height} (aspect ratio: {aspect_ratio:.2f})")
+
+        # 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 = {
+    "ResizeToClosetBucket": ResizeToClosetBucket
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ResizeToClosetBucket": "Resize to Closet Bucket or Proportional"
+}