A single function for NovelAI aspect ratio bucketing. No need for classes or objects.

gen_buckets.py

def gen_buckets(
    min_dim: int = 256,
    base_res: tuple[int, int] = (512, 512),
    max_size: tuple[int, int] = (768, 512),
    dim_limit: int = 1024,
    divisible=64,
):
    """
    Adapted from:
    https://github.com/NovelAI/novelai-aspect-ratio-bucketing/blob/main/bucketmanager.py
    """

    f = 8
    max_tokens = (max_size[0] / f) * (max_size[1] / f)
    resolutions = []
    aspects = []
    w = min_dim
    while (w / f) * (min_dim / f) <= max_tokens and w <= dim_limit:
        h = min_dim
        got_base = False
        while (w / f) * ((h + divisible) / f) <= max_tokens and (
            h + divisible
        ) <= dim_limit:
            if w == base_res[0] and h == base_res[1]:
                got_base = True
            h += divisible
        if (w != base_res[0] or h != base_res[1]) and got_base:
            resolutions.append(base_res)
            aspects.append(1)
        resolutions.append((w, h))
        aspects.append(float(w) / float(h))
        w += divisible

    h = min_dim
    while (h / f) * (min_dim / f) <= max_tokens and h <= dim_limit:
        w = min_dim
        got_base = False
        while (h / f) * ((w + divisible) / f) <= max_tokens and (
            w + divisible
        ) <= dim_limit:
            if w == base_res[0] and h == base_res[1]:
                got_base = True
            w += divisible
        resolutions.append((w, h))
        aspects.append(float(w) / float(h))
        h += divisible
    res_map = {}
    for i, res in enumerate(resolutions):
        res_map[res] = aspects[i]

    resolutions = sorted(res_map.keys(), key=lambda x: x[0] * 4096 - x[1])
    aspects = list(map(lambda x: res_map[x], resolutions))

    return resolutions, aspects