Cosine Alpha Bar Schedule

Problem

Evaluate the improved cosine noise schedule at a given timestep.

alpha_bar(t) = cos((t / T + s) / (1 + s) * pi / 2)^2 divided by cos(s / (1 + s) * pi / 2)^2, where T is timesteps and s the small offset that keeps the schedule smooth near t = 0. Clamp t to at most timesteps, return 1.0 when t <= 0, and otherwise round to 6 decimal places.

Test cases

case 1

in: [0,1000,0.008]

expected: 1

case 2

in: [500,1000,0.008]

expected: 0.493844

case 3

in: [999,1000,0.008]

expected: 0.000002

case 4

in: [2000,1000,0.008]

expected: 0

Starter codePython

def cosine_alpha_bar(t, timesteps, s):
    # Your code here
    pass

4 test cases · run your code against them in the editor