Upload bucket_alone.py

bucket_alone.py

@@ -0,0 +1,105 @@
+# 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)"
+}