Reworked the bucketing system to precisly match model specific divisibility. The old SDXL bucket sizes needed to go.
This commit is contained in:
Jaret Burkett
parent
68e9b38220
commit
c9264bdd0b
|
|
@ -1,4 +1,4 @@
|
|||
from typing import Type, List, Union, TypedDict
|
||||
from typing import TypedDict
|
||||
|
||||
|
||||
class BucketResolution(TypedDict):
|
||||
|
|
@ -6,124 +6,23 @@ class BucketResolution(TypedDict):
|
|||
height: int
|
||||
|
||||
|
||||
# resolutions SDXL was trained on with a 1024x1024 base resolution
|
||||
resolutions_1024: List[BucketResolution] = [
|
||||
# SDXL Base resolution
|
||||
{"width": 1024, "height": 1024},
|
||||
# SDXL Resolutions, widescreen
|
||||
{"width": 2048, "height": 512},
|
||||
{"width": 1984, "height": 512},
|
||||
{"width": 1920, "height": 512},
|
||||
{"width": 1856, "height": 512},
|
||||
{"width": 1792, "height": 576},
|
||||
{"width": 1728, "height": 576},
|
||||
{"width": 1664, "height": 576},
|
||||
{"width": 1600, "height": 640},
|
||||
{"width": 1536, "height": 640},
|
||||
{"width": 1472, "height": 704},
|
||||
{"width": 1408, "height": 704},
|
||||
{"width": 1344, "height": 704},
|
||||
{"width": 1344, "height": 768},
|
||||
{"width": 1280, "height": 768},
|
||||
{"width": 1216, "height": 832},
|
||||
{"width": 1152, "height": 832},
|
||||
{"width": 1152, "height": 896},
|
||||
{"width": 1088, "height": 896},
|
||||
{"width": 1088, "height": 960},
|
||||
{"width": 1024, "height": 960},
|
||||
# SDXL Resolutions, portrait
|
||||
{"width": 960, "height": 1024},
|
||||
{"width": 960, "height": 1088},
|
||||
{"width": 896, "height": 1088},
|
||||
{"width": 896, "height": 1152}, # 2:3
|
||||
{"width": 832, "height": 1152},
|
||||
{"width": 832, "height": 1216},
|
||||
{"width": 768, "height": 1280},
|
||||
{"width": 768, "height": 1344},
|
||||
{"width": 704, "height": 1408},
|
||||
{"width": 704, "height": 1472},
|
||||
{"width": 640, "height": 1536},
|
||||
{"width": 640, "height": 1600},
|
||||
{"width": 576, "height": 1664},
|
||||
{"width": 576, "height": 1728},
|
||||
{"width": 576, "height": 1792},
|
||||
{"width": 512, "height": 1856},
|
||||
{"width": 512, "height": 1920},
|
||||
{"width": 512, "height": 1984},
|
||||
{"width": 512, "height": 2048},
|
||||
# extra wides
|
||||
{"width": 8192, "height": 128},
|
||||
{"width": 128, "height": 8192},
|
||||
]
|
||||
|
||||
def get_bucket_sizes(resolution: int = 512, divisibility: int = 8) -> List[BucketResolution]:
|
||||
# determine scaler form 1024 to resolution
|
||||
scaler = resolution / 1024
|
||||
|
||||
bucket_size_list = []
|
||||
for bucket in resolutions_1024:
|
||||
# must be divisible by 8
|
||||
width = int(bucket["width"] * scaler)
|
||||
height = int(bucket["height"] * scaler)
|
||||
if width % divisibility != 0:
|
||||
width = width - (width % divisibility)
|
||||
if height % divisibility != 0:
|
||||
height = height - (height % divisibility)
|
||||
bucket_size_list.append({"width": width, "height": height})
|
||||
|
||||
return bucket_size_list
|
||||
|
||||
|
||||
def get_resolution(width, height):
|
||||
num_pixels = width * height
|
||||
# determine same number of pixels for square image
|
||||
square_resolution = int(num_pixels ** 0.5)
|
||||
square_resolution = int(num_pixels**0.5)
|
||||
return square_resolution
|
||||
|
||||
|
||||
def get_bucket_for_image_size(
|
||||
width: int,
|
||||
height: int,
|
||||
bucket_size_list: List[BucketResolution] = None,
|
||||
resolution: Union[int, None] = None,
|
||||
divisibility: int = 8
|
||||
width: int, height: int, resolution: int = 512, divisibility: int = 8
|
||||
) -> BucketResolution:
|
||||
total_pixels = width * height
|
||||
max_pixels = resolution * resolution
|
||||
|
||||
if bucket_size_list is None and resolution is None:
|
||||
# get resolution from width and height
|
||||
resolution = get_resolution(width, height)
|
||||
if bucket_size_list is None:
|
||||
# if real resolution is smaller, use that instead
|
||||
real_resolution = get_resolution(width, height)
|
||||
resolution = min(resolution, real_resolution)
|
||||
bucket_size_list = get_bucket_sizes(resolution=resolution, divisibility=divisibility)
|
||||
target_pixels = min(total_pixels, max_pixels)
|
||||
|
||||
# Check for exact match first
|
||||
for bucket in bucket_size_list:
|
||||
if bucket["width"] == width and bucket["height"] == height:
|
||||
return bucket
|
||||
scaler = target_pixels / total_pixels
|
||||
new_width = int((width * scaler) // divisibility * divisibility)
|
||||
new_height = int((height * scaler) // divisibility * divisibility)
|
||||
|
||||
# If exact match not found, find the closest bucket
|
||||
closest_bucket = None
|
||||
min_removed_pixels = float("inf")
|
||||
|
||||
for bucket in bucket_size_list:
|
||||
scale_w = bucket["width"] / width
|
||||
scale_h = bucket["height"] / height
|
||||
|
||||
# To minimize pixels, we use the larger scale factor to minimize the amount that has to be cropped.
|
||||
scale = max(scale_w, scale_h)
|
||||
|
||||
new_width = int(width * scale)
|
||||
new_height = int(height * scale)
|
||||
|
||||
removed_pixels = (new_width - bucket["width"]) * new_height + (new_height - bucket["height"]) * new_width
|
||||
|
||||
if removed_pixels < min_removed_pixels:
|
||||
min_removed_pixels = removed_pixels
|
||||
closest_bucket = bucket
|
||||
|
||||
if closest_bucket is None:
|
||||
raise ValueError("No suitable bucket found")
|
||||
|
||||
return closest_bucket
|
||||
return {"width": new_width, "height": new_height}
|
||||
|
|
|
|||
Loading…
Reference in New Issue