Open-Sora/opensora/datasets/aspect.py

import math


# Ours


def get_h_w(a, ts):
    h = (ts * a) ** 0.5
    h = math.ceil(h) if math.ceil(h) % 2 == 0 else math.floor(h)
    w = h / a
    w = math.ceil(w) if math.ceil(w) % 2 == 0 else math.floor(w)
    return h, w


def get_aspect_ratios_dict(
    ts=360 * 640,
    ars=[0.39, 0.42, 0.48, 0.50, 0.52, 0.56, 0.66, 0.75, 1.00, 1.33, 1.52, 1.78, 1.92, 2.00, 2.10],
):
    est = {f"{a:.2f}": get_h_w(a, ts) for a in ars}
    return est


# S = 2073600
ASPECT_RATIO_1080P = {
    "0.39": (900, 2308),
    "0.42": (934, 2224),
    "0.48": (998, 2080),
    "0.50": (1018, 2036),
    "0.52": (1038, 1996),
    "0.56": (1080, 1920),  # base
    "0.66": (1170, 1772),
    "0.75": (1248, 1664),
    "1.00": (1440, 1440),
    "1.33": (1660, 1248),
    "1.52": (1776, 1168),
    "1.78": (1922, 1080),
    "1.92": (1996, 1040),
    "2.00": (2036, 1018),
    "2.10": (2086, 994),
}

# S = 921600
ASPECT_RATIO_720P = {
    "0.39": (600, 1538),
    "0.42": (622, 1480),
    "0.48": (666, 1388),
    "0.50": (678, 1356),
    "0.52": (692, 1330),
    "0.56": (720, 1280),  # base
    "0.66": (780, 1182),
    "0.75": (832, 1110),
    "1.00": (960, 960),
    "1.33": (1108, 834),
    "1.52": (1184, 778),
    "1.78": (1280, 720),
    "1.92": (1330, 692),
    "2.00": (1358, 679),
    "2.10": (1392, 662),
}

# S = 409920
ASPECT_RATIO_480P = {
    "0.39": (400, 1026),
    "0.42": (414, 986),
    "0.48": (444, 925),
    "0.50": (452, 904),
    "0.52": (462, 888),
    "0.56": (480, 854),  # base
    "0.66": (520, 788),
    "0.75": (554, 738),
    "1.00": (640, 640),
    "1.33": (738, 554),
    "1.52": (790, 520),
    "1.78": (854, 480),
    "1.92": (888, 462),
    "2.00": (906, 453),
    "2.10": (928, 442),
}

# S = 230400
ASPECT_RATIO_360P = {
    "0.39": (300, 770),
    "0.42": (312, 742),
    "0.48": (332, 692),
    "0.50": (340, 680),
    "0.52": (346, 666),
    "0.56": (360, 640),  # base
    "0.66": (390, 590),
    "0.75": (416, 554),
    "1.00": (480, 480),
    "1.33": (554, 416),
    "1.52": (592, 390),
    "1.78": (640, 360),
    "1.92": (666, 346),
    "2.00": (678, 339),
    "2.10": (696, 332),
}

# S = 102240
ASPECT_RATIO_240P = {
    "0.39": (200, 512),
    "0.42": (208, 496),
    "0.48": (222, 462),
    "0.50": (226, 452),
    "0.52": (230, 442),
    "0.56": (240, 426),
    "0.66": (260, 394),
    "0.75": (276, 368),
    "1.00": (320, 320),
    "1.33": (368, 276),
    "1.52": (394, 260),
    "1.78": (426, 240),
    "1.92": (444, 232),
    "2.00": (452, 226),
    "2.10": (464, 220),
}


# from PixArt
# S = 1048576
ASPECT_RATIO_1024 = {
    "0.25": [512.0, 2048.0],
    "0.26": [512.0, 1984.0],
    "0.27": [512.0, 1920.0],
    "0.28": [512.0, 1856.0],
    "0.32": [576.0, 1792.0],
    "0.33": [576.0, 1728.0],
    "0.35": [576.0, 1664.0],
    "0.4": [640.0, 1600.0],
    "0.42": [640.0, 1536.0],
    "0.48": [704.0, 1472.0],
    "0.5": [704.0, 1408.0],
    "0.52": [704.0, 1344.0],
    "0.57": [768.0, 1344.0],
    "0.6": [768.0, 1280.0],
    "0.68": [832.0, 1216.0],
    "0.72": [832.0, 1152.0],
    "0.78": [896.0, 1152.0],
    "0.82": [896.0, 1088.0],
    "0.88": [960.0, 1088.0],
    "0.94": [960.0, 1024.0],
    "1.0": [1024.0, 1024.0],
    "1.07": [1024.0, 960.0],
    "1.13": [1088.0, 960.0],
    "1.21": [1088.0, 896.0],
    "1.29": [1152.0, 896.0],
    "1.38": [1152.0, 832.0],
    "1.46": [1216.0, 832.0],
    "1.67": [1280.0, 768.0],
    "1.75": [1344.0, 768.0],
    "2.0": [1408.0, 704.0],
    "2.09": [1472.0, 704.0],
    "2.4": [1536.0, 640.0],
    "2.5": [1600.0, 640.0],
    "2.89": [1664.0, 576.0],
    "3.0": [1728.0, 576.0],
    "3.11": [1792.0, 576.0],
    "3.62": [1856.0, 512.0],
    "3.75": [1920.0, 512.0],
    "3.88": [1984.0, 512.0],
    "4.0": [2048.0, 512.0],
}

# S = 262144
ASPECT_RATIO_512 = {
    "0.25": [256.0, 1024.0],
    "0.26": [256.0, 992.0],
    "0.27": [256.0, 960.0],
    "0.28": [256.0, 928.0],
    "0.32": [288.0, 896.0],
    "0.33": [288.0, 864.0],
    "0.35": [288.0, 832.0],
    "0.4": [320.0, 800.0],
    "0.42": [320.0, 768.0],
    "0.48": [352.0, 736.0],
    "0.5": [352.0, 704.0],
    "0.52": [352.0, 672.0],
    "0.57": [384.0, 672.0],
    "0.6": [384.0, 640.0],
    "0.68": [416.0, 608.0],
    "0.72": [416.0, 576.0],
    "0.78": [448.0, 576.0],
    "0.82": [448.0, 544.0],
    "0.88": [480.0, 544.0],
    "0.94": [480.0, 512.0],
    "1.0": [512.0, 512.0],
    "1.07": [512.0, 480.0],
    "1.13": [544.0, 480.0],
    "1.21": [544.0, 448.0],
    "1.29": [576.0, 448.0],
    "1.38": [576.0, 416.0],
    "1.46": [608.0, 416.0],
    "1.67": [640.0, 384.0],
    "1.75": [672.0, 384.0],
    "2.0": [704.0, 352.0],
    "2.09": [736.0, 352.0],
    "2.4": [768.0, 320.0],
    "2.5": [800.0, 320.0],
    "2.89": [832.0, 288.0],
    "3.0": [864.0, 288.0],
    "3.11": [896.0, 288.0],
    "3.62": [928.0, 256.0],
    "3.75": [960.0, 256.0],
    "3.88": [992.0, 256.0],
    "4.0": [1024.0, 256.0],
}

# S = 65536
ASPECT_RATIO_256 = {
    "0.25": [128.0, 512.0],
    "0.26": [128.0, 496.0],
    "0.27": [128.0, 480.0],
    "0.28": [128.0, 464.0],
    "0.32": [144.0, 448.0],
    "0.33": [144.0, 432.0],
    "0.35": [144.0, 416.0],
    "0.4": [160.0, 400.0],
    "0.42": [160.0, 384.0],
    "0.48": [176.0, 368.0],
    "0.5": [176.0, 352.0],
    "0.52": [176.0, 336.0],
    "0.57": [192.0, 336.0],
    "0.6": [192.0, 320.0],
    "0.68": [208.0, 304.0],
    "0.72": [208.0, 288.0],
    "0.78": [224.0, 288.0],
    "0.82": [224.0, 272.0],
    "0.88": [240.0, 272.0],
    "0.94": [240.0, 256.0],
    "1.0": [256.0, 256.0],
    "1.07": [256.0, 240.0],
    "1.13": [272.0, 240.0],
    "1.21": [272.0, 224.0],
    "1.29": [288.0, 224.0],
    "1.38": [288.0, 208.0],
    "1.46": [304.0, 208.0],
    "1.67": [320.0, 192.0],
    "1.75": [336.0, 192.0],
    "2.0": [352.0, 176.0],
    "2.09": [368.0, 176.0],
    "2.4": [384.0, 160.0],
    "2.5": [400.0, 160.0],
    "2.89": [416.0, 144.0],
    "3.0": [432.0, 144.0],
    "3.11": [448.0, 144.0],
    "3.62": [464.0, 128.0],
    "3.75": [480.0, 128.0],
    "3.88": [496.0, 128.0],
    "4.0": [512.0, 128.0],
}


def get_closest_ratio(height: float, width: float, ratios: dict):
    aspect_ratio = height / width
    closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - aspect_ratio))
    return ratios[closest_ratio], float(closest_ratio)


ASPECT_RATIOS = {
    "256": (65536, ASPECT_RATIO_256),
    "240p": (102240, ASPECT_RATIO_240P),
    "360p": (230400, ASPECT_RATIO_360P),
    "512": (262144, ASPECT_RATIO_512),
    "480p": (409920, ASPECT_RATIO_480P),
    "720p": (921600, ASPECT_RATIO_720P),
    "1024": (1048576, ASPECT_RATIO_1024),
    "1080p": (2073600, ASPECT_RATIO_1080P),
}