mirror of
https://github.com/ostris/ai-toolkit.git
synced 2026-07-11 17:52:04 +00:00
49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
import math
|
|
from typing import TypedDict
|
|
|
|
|
|
class BucketResolution(TypedDict):
|
|
width: int
|
|
height: int
|
|
|
|
|
|
def get_resolution(width, height):
|
|
num_pixels = width * height
|
|
# determine same number of pixels for square image
|
|
square_resolution = int(num_pixels**0.5)
|
|
return square_resolution
|
|
|
|
|
|
def get_bucket_for_image_size(
|
|
width: int, height: int, resolution: int = 512, divisibility: int = 8
|
|
) -> BucketResolution:
|
|
total_pixels = width * height
|
|
max_pixels = resolution * resolution
|
|
|
|
target_pixels = min(total_pixels, max_pixels)
|
|
|
|
scaler = (target_pixels / total_pixels) ** 0.5
|
|
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}
|