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 timestep
(num_timesteps, section_counts)
| 57 | return normalized_feat * style_std.expand(size) + style_mean.expand(size) |
| 58 | |
| 59 | def space_timesteps(num_timesteps, section_counts): |
| 60 | """ |
| 61 | Create a list of timesteps to use from an original diffusion process, |
| 62 | given the number of timesteps we want to take from equally-sized portions |
| 63 | of the original process. |
| 64 | |
| 65 | For example, if there's 300 timesteps and the section counts are [10,15,20] |
| 66 | then the first 100 timesteps are strided to be 10 timesteps, the second 100 |
| 67 | are strided to be 15 timesteps, and the final 100 are strided to be 20. |
| 68 | |
| 69 | If the stride is a string starting with "ddim", then the fixed striding |
| 70 | from the DDIM paper is used, and only one section is allowed. |
| 71 | |
| 72 | :param num_timesteps: the number of diffusion steps in the original |
| 73 | process to divide up. |
| 74 | :param section_counts: either a list of numbers, or a string containing |
| 75 | comma-separated numbers, indicating the step count |
| 76 | per section. As a special case, use "ddimN" where N |
| 77 | is a number of steps to use the striding from the |
| 78 | DDIM paper. |
| 79 | :return: a set of diffusion steps from the original process to use. |
| 80 | """ |
| 81 | if isinstance(section_counts, str): |
| 82 | if section_counts.startswith("ddim"): |
| 83 | desired_count = int(section_counts[len("ddim"):]) |
| 84 | for i in range(1, num_timesteps): |
| 85 | if len(range(0, num_timesteps, i)) == desired_count: |
| 86 | return set(range(0, num_timesteps, i)) |
| 87 | raise ValueError( |
| 88 | f"cannot create exactly {num_timesteps} steps with an integer stride" |
| 89 | ) |
| 90 | section_counts = [int(x) for x in section_counts.split(",")] #[250,] |
| 91 | size_per = num_timesteps // len(section_counts) |
| 92 | extra = num_timesteps % len(section_counts) |
| 93 | start_idx = 0 |
| 94 | all_steps = [] |
| 95 | for i, section_count in enumerate(section_counts): |
| 96 | size = size_per + (1 if i < extra else 0) |
| 97 | if size < section_count: |
| 98 | raise ValueError( |
| 99 | f"cannot divide section of {size} steps into {section_count}" |
| 100 | ) |
| 101 | if section_count <= 1: |
| 102 | frac_stride = 1 |
| 103 | else: |
| 104 | frac_stride = (size - 1) / (section_count - 1) |
| 105 | cur_idx = 0.0 |
| 106 | taken_steps = [] |
| 107 | for _ in range(section_count): |
| 108 | taken_steps.append(start_idx + round(cur_idx)) |
| 109 | cur_idx += frac_stride |
| 110 | all_steps += taken_steps |
| 111 | start_idx += size |
| 112 | return set(all_steps) |
| 113 | |
| 114 | def chunk(it, size): |
| 115 | it = iter(it) |