Merge pull request kijai#3 from Crimsonfart/patch-2

Adds the ability to support "any" resolution.

Author: kijai

diffusers_helper/bucket_tools.py

@@ -1,52 +1,36 @@
 bucket_options = {
-    240: [
-        (224, 336),
-        (256, 304),
-        (288, 272),
-        (320, 240),
-        (352, 208),
-        (384, 176),
-    ],
-    320: [
-        (256, 576),
-        (288, 544),
-        (320, 512),
-        (352, 480),
-        (384, 448),
-        (416, 416),
-        (448, 384),
-        (480, 352),
-        (512, 320),
-        (544, 288),
-        (576, 256),
-    ],
-    640: [
-        (416, 960),
-        (448, 864),
-        (480, 832),
-        (512, 768),
-        (544, 704),
-        (576, 672),
-        (608, 640),
-        (640, 608),
-        (672, 576),
-        (704, 544),
-        (768, 512),
-        (832, 480),
-        (864, 448),
-        (960, 416),
-    ],
+    (416, 960),
+    (448, 864),
+    (480, 832),
+    (512, 768),
+    (544, 704),
+    (576, 672),
+    (608, 640),
+    (640, 608),
+    (672, 576),
+    (704, 544),
+    (768, 512),
+    (832, 480),
+    (864, 448),
+    (960, 416),
 }
 
 
 def find_nearest_bucket(h, w, resolution=640):
-    min_metric = float("inf")
+    min_metric = float('inf')
     best_bucket = None
-    for bucket_h, bucket_w in bucket_options[resolution]:
+    for (bucket_h, bucket_w) in bucket_options:
         metric = abs(h * bucket_w - w * bucket_h)
         if metric <= min_metric:
             min_metric = metric
             best_bucket = (bucket_h, bucket_w)
-    print(f"Best bucket: {best_bucket}")
+
+    if resolution != 640:
+        scale_factor = resolution / 640.0
+        scaled_height = round(best_bucket[0] * scale_factor / 16) * 16
+        scaled_width = round(best_bucket[1] * scale_factor / 16) * 16
+        best_bucket = (scaled_height, scaled_width)
+        print(f'Resolution: {best_bucket[1]} x {best_bucket[0]}')
+
     return best_bucket