Fabrice-TIERCELIN commited on
Commit
23c4bff
·
verified ·
1 Parent(s): dafe49c
Files changed (1) hide show
  1. diffusers_helper/bucket_tools.py +28 -3
diffusers_helper/bucket_tools.py CHANGED
@@ -1,4 +1,29 @@
1
- def find_nearest_bucket(h, w, resolution=640):
2
- rate = ((resolution**2) / (h * w))**(1/2)
3
- return (int(h * rate), int(w * rate))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4
 
 
 
 
 
 
 
 
 
 
 
1
+ bucket_options = {
2
+ 640: [
3
+ (416, 960),
4
+ (448, 864),
5
+ (480, 832),
6
+ (512, 768),
7
+ (544, 704),
8
+ (576, 672),
9
+ (608, 640),
10
+ (640, 608),
11
+ (672, 576),
12
+ (704, 544),
13
+ (768, 512),
14
+ (832, 480),
15
+ (864, 448),
16
+ (960, 416),
17
+ ],
18
+ }
19
+
20
 
21
+ def find_nearest_bucket(h, w, resolution=640):
22
+ min_metric = float('inf')
23
+ best_bucket = None
24
+ for (bucket_h, bucket_w) in bucket_options[resolution]:
25
+ metric = abs(h * bucket_w - w * bucket_h)
26
+ if metric <= min_metric:
27
+ min_metric = metric
28
+ best_bucket = (bucket_h, bucket_w)
29
+ return best_bucket