RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Adjust bucket sizes to achieve maximum pixels without going over.

This commit is contained in:
Jaret Burkett 2026-05-28 09:36:17 -06:00
parent ef57c1077c
commit 3e86d81fc6
1 changed files with 22 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
toolkit/buckets.py
View File

@ -1,3 +1,4 @@
import math
from typing import TypedDict
@ -22,7 +23,26 @@ def get_bucket_for_image_size(
target_pixels = min(total_pixels, max_pixels)
scaler = target_pixels / total_pixels
new_width = int(round((width * scaler) / divisibility) * divisibility)
new_height = int(round((height * scaler) / divisibility) * divisibility)
w_raw = (width * scaler) / divisibility
h_raw = (height * scaler) / divisibility
candidates = [
(math.floor(w_raw) * divisibility, math.floor(h_raw) * divisibility),
(math.floor(w_raw) * divisibility, math.ceil(h_raw) * divisibility),
(math.ceil(w_raw) * divisibility, math.floor(h_raw) * divisibility),
(math.ceil(w_raw) * divisibility, math.ceil(h_raw) * divisibility),
]
capped = [(w, h) for w, h in candidates if w > 0 and h > 0 and w * h <= max_pixels]
if not capped:
capped = [
(
max(divisibility, math.floor(w_raw) * divisibility),
max(divisibility, math.floor(h_raw) * divisibility),
)
]
new_width, new_height = min(
capped, key=lambda wh: abs(wh[0] * wh[1] - target_pixels)
)
return {"width": new_width, "height": new_height}