diffusers_helper/bucket_tools.py

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