fred-dev/comfy_ui_ali

1import torch2 3def calc_mantissa(abs_x, exponent, normal_mask, MANTISSA_BITS, EXPONENT_BIAS, generator=None):4    mantissa_scaled = torch.where(5        normal_mask,6        (abs_x / (2.0 ** (exponent - EXPONENT_BIAS)) - 1.0) * (2**MANTISSA_BITS),7        (abs_x / (2.0 ** (-EXPONENT_BIAS + 1 - MANTISSA_BITS)))8    )9 10    mantissa_scaled += torch.rand(mantissa_scaled.size(), dtype=mantissa_scaled.dtype, layout=mantissa_scaled.layout, device=mantissa_scaled.device, generator=generator)11    return mantissa_scaled.floor() / (2**MANTISSA_BITS)12 13#Not 100% sure about this14def manual_stochastic_round_to_float8(x, dtype, generator=None):15    if dtype == torch.float8_e4m3fn:16        EXPONENT_BITS, MANTISSA_BITS, EXPONENT_BIAS = 4, 3, 717    elif dtype == torch.float8_e5m2:18        EXPONENT_BITS, MANTISSA_BITS, EXPONENT_BIAS = 5, 2, 1519    else:20        raise ValueError("Unsupported dtype")21 22    x = x.half()23    sign = torch.sign(x)24    abs_x = x.abs()25    sign = torch.where(abs_x == 0, 0, sign)26 27    # Combine exponent calculation and clamping28    exponent = torch.clamp(29        torch.floor(torch.log2(abs_x)) + EXPONENT_BIAS,30        0, 2**EXPONENT_BITS - 131    )32 33    # Combine mantissa calculation and rounding34    normal_mask = ~(exponent == 0)35 36    abs_x[:] = calc_mantissa(abs_x, exponent, normal_mask, MANTISSA_BITS, EXPONENT_BIAS, generator=generator)37 38    sign *= torch.where(39        normal_mask,40        (2.0 ** (exponent - EXPONENT_BIAS)) * (1.0 + abs_x),41        (2.0 ** (-EXPONENT_BIAS + 1)) * abs_x42    )43 44    inf = torch.finfo(dtype)45    torch.clamp(sign, min=inf.min, max=inf.max, out=sign)46    return sign47 48 49 50def stochastic_rounding(value, dtype, seed=0):51    if dtype == torch.float32:52        return value.to(dtype=torch.float32)53    if dtype == torch.float16:54        return value.to(dtype=torch.float16)55    if dtype == torch.bfloat16:56        return value.to(dtype=torch.bfloat16)57    if dtype == torch.float8_e4m3fn or dtype == torch.float8_e5m2:58        generator = torch.Generator(device=value.device)59        generator.manual_seed(seed)60        output = torch.empty_like(value, dtype=dtype)61        num_slices = max(1, (value.numel() / (4096 * 4096)))62        slice_size = max(1, round(value.shape[0] / num_slices))63        for i in range(0, value.shape[0], slice_size):64            output[i:i+slice_size].copy_(manual_stochastic_round_to_float8(value[i:i+slice_size], dtype, generator=generator))65        return output66 67    return value.to(dtype=dtype)68