| 88 | |
| 89 | # Define scale schedule with start and end points at 1 |
| 90 | def scale_schedule(scale, sigma, step_index): |
| 91 | num_steps_period = num_steps // n_periods |
| 92 | num_half_steps_period = num_steps_period // 2 |
| 93 | |
| 94 | step_index_period = step_index % num_steps_period |
| 95 | if step_index_period < num_half_steps_period: |
| 96 | if leading_constant: |
| 97 | return scale |
| 98 | else: |
| 99 | # Linearly increase in the first half of the steps |
| 100 | return 1 + (scale - 1) * (step_index_period / (num_half_steps_period - 1)) |
| 101 | else: |
| 102 | # Linearly decrease in the second half of the steps |
| 103 | return scale - (scale - 1) * ((step_index_period - num_half_steps_period) / (num_steps_period - num_half_steps_period - 1)) |
| 104 | |
| 105 | # half_steps = num_steps // 2 |
| 106 | # if step_index < half_steps: |
| 107 | # if leading_constant: |
| 108 | # return scale |
| 109 | # else: |
| 110 | # # Linearly increase in the first half of the steps |
| 111 | # return 1 + (scale - 1) * (step_index / (half_steps - 1)) |
| 112 | # else: |
| 113 | # # Linearly decrease in the second half of the steps |
| 114 | # return scale - (scale - 1) * ((step_index - half_steps) / (num_steps - half_steps - 1)) |
| 115 | |
| 116 | # Partial function with predefined scale |
| 117 | self.scale_schedule = partial(scale_schedule, scale) |