The resolution of the generated video will be...

diffusers_helper/bucket_tools.py

@@ -1,30 +1,103 @@
-bucket_options = {
-    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),
-    ],
-}
-
-
-def find_nearest_bucket(h, w, resolution=640):
-    min_metric = float('inf')
-    best_bucket = None
-    for (bucket_h, bucket_w) in bucket_options[resolution]:
-        metric = abs(h * bucket_w - w * bucket_h)
-        if metric <= min_metric:
-            min_metric = metric
-            best_bucket = (bucket_h, bucket_w)
-    return best_bucket
-
+bucket_options = {
+    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),
+    ],
+    672: [
+        (480, 864),
+        (512, 832),
+        (544, 768),
+        (576, 704),
+        (608, 672),
+        (640, 640),
+        (672, 608),
+        (704, 576),
+        (768, 544),
+        (832, 512),
+        (864, 480),
+    ],
+    704: [
+        (480, 960),
+        (512, 864),
+        (544, 832),
+        (576, 768),
+        (608, 704),
+        (640, 672),
+        (672, 640),
+        (704, 608),
+        (768, 576),
+        (832, 544),
+        (864, 512),
+        (960, 480),
+    ],
+    768: [
+        (512, 960),
+        (544, 864),
+        (576, 832),
+        (608, 768),
+        (640, 704),
+        (672, 672),
+        (704, 640),
+        (768, 608),
+        (832, 576),
+        (864, 544),
+        (960, 512),
+    ],
+    832: [
+        (544, 960),
+        (576, 864),
+        (608, 832),
+        (640, 768),
+        (672, 704),
+        (704, 672),
+        (768, 640),
+        (832, 608),
+        (864, 576),
+        (960, 544),
+    ],
+    864: [
+        (576, 960),
+        (608, 864),
+        (640, 832),
+        (672, 768),
+        (704, 704),
+        (768, 672),
+        (832, 640),
+        (864, 608),
+        (960, 576),
+    ],
+    960: [
+        (608, 960),
+        (640, 864),
+        (672, 832),
+        (704, 768),
+        (768, 704),
+        (832, 672),
+        (864, 640),
+        (960, 608),
+    ],
+}
+
+
+def find_nearest_bucket(h, w, resolution=640):
+    min_metric = float('inf')
+    best_bucket = None
+    for (bucket_h, bucket_w) in bucket_options[resolution]:
+        metric = abs(h * bucket_w - w * bucket_h)
+        if metric <= min_metric:
+            min_metric = metric
+            best_bucket = (bucket_h, bucket_w)
+    print("The resolution of the generated video will be " + str(best_bucket))
+    return best_bucket