Create a list of timesteps to use from an original diffusion process, given the number of timesteps we want to take from equally-sized portions of the original process. For example, if there's 300 timesteps and the section counts are [10,15,20] then the first 100 timesteps are s
(num_timesteps, section_counts)
| 10 | |
| 11 | |
| 12 | def space_timesteps(num_timesteps, section_counts): |
| 13 | """ |
| 14 | Create a list of timesteps to use from an original diffusion process, |
| 15 | given the number of timesteps we want to take from equally-sized portions |
| 16 | of the original process. |
| 17 | For example, if there's 300 timesteps and the section counts are [10,15,20] |
| 18 | then the first 100 timesteps are strided to be 10 timesteps, the second 100 |
| 19 | are strided to be 15 timesteps, and the final 100 are strided to be 20. |
| 20 | If the stride is a string starting with "ddim", then the fixed striding |
| 21 | from the DDIM paper is used, and only one section is allowed. |
| 22 | :param num_timesteps: the number of diffusion steps in the original |
| 23 | process to divide up. |
| 24 | :param section_counts: either a list of numbers, or a string containing |
| 25 | comma-separated numbers, indicating the step count |
| 26 | per section. As a special case, use "ddimN" where N |
| 27 | is a number of steps to use the striding from the |
| 28 | DDIM paper. |
| 29 | :return: a set of diffusion steps from the original process to use. |
| 30 | """ |
| 31 | if isinstance(section_counts, str): |
| 32 | if section_counts.startswith("ddim"): |
| 33 | desired_count = int(section_counts[len("ddim") :]) |
| 34 | for i in range(1, num_timesteps): |
| 35 | if len(range(0, num_timesteps, i)) == desired_count: |
| 36 | return set(range(0, num_timesteps, i)) |
| 37 | raise ValueError( |
| 38 | f"cannot create exactly {num_timesteps} steps with an integer stride" |
| 39 | ) |
| 40 | section_counts = [int(x) for x in section_counts.split(",")] |
| 41 | size_per = num_timesteps // len(section_counts) |
| 42 | extra = num_timesteps % len(section_counts) |
| 43 | start_idx = 0 |
| 44 | all_steps = [] |
| 45 | for i, section_count in enumerate(section_counts): |
| 46 | size = size_per + (1 if i < extra else 0) |
| 47 | if size < section_count: |
| 48 | raise ValueError( |
| 49 | f"cannot divide section of {size} steps into {section_count}" |
| 50 | ) |
| 51 | if section_count <= 1: |
| 52 | frac_stride = 1 |
| 53 | else: |
| 54 | frac_stride = (size - 1) / (section_count - 1) |
| 55 | cur_idx = 0.0 |
| 56 | taken_steps = [] |
| 57 | for _ in range(section_count): |
| 58 | taken_steps.append(start_idx + round(cur_idx)) |
| 59 | cur_idx += frac_stride |
| 60 | all_steps += taken_steps |
| 61 | start_idx += size |
| 62 | return set(all_steps) |
| 63 | |
| 64 | |
| 65 | class SpacedDiffusion(GaussianDiffusion): |