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
(num_timesteps, section_counts)
| 1177 | |
| 1178 | |
| 1179 | def space_timesteps(num_timesteps, section_counts): |
| 1180 | """ |
| 1181 | Create a list of timesteps to use from an original diffusion process, |
| 1182 | given the number of timesteps we want to take from equally-sized portions |
| 1183 | of the original process. |
| 1184 | |
| 1185 | For example, if there's 300 timesteps and the section counts are [10,15,20] |
| 1186 | then the first 100 timesteps are strided to be 10 timesteps, the second 100 |
| 1187 | are strided to be 15 timesteps, and the final 100 are strided to be 20. |
| 1188 | |
| 1189 | :param num_timesteps: the number of diffusion steps in the original |
| 1190 | process to divide up. |
| 1191 | :param section_counts: either a list of numbers, or a string containing |
| 1192 | comma-separated numbers, indicating the step count |
| 1193 | per section. As a special case, use "ddimN" where N |
| 1194 | is a number of steps to use the striding from the |
| 1195 | DDIM paper. |
| 1196 | :return: a set of diffusion steps from the original process to use. |
| 1197 | """ |
| 1198 | if isinstance(section_counts, str): |
| 1199 | if section_counts.startswith("ddim"): |
| 1200 | desired_count = int(section_counts[len("ddim"):]) |
| 1201 | for i in range(1, num_timesteps): |
| 1202 | if len(range(0, num_timesteps, i)) == desired_count: |
| 1203 | return set(range(0, num_timesteps, i)) |
| 1204 | raise ValueError( |
| 1205 | f"cannot create exactly {num_timesteps} steps with an integer stride" |
| 1206 | ) |
| 1207 | elif section_counts == "fast27": |
| 1208 | # steps = space_timesteps(num_timesteps, "10,10,3,2,2") |
| 1209 | # steps = space_timesteps(num_timesteps, "30,30,16,12,12") |
| 1210 | steps = space_timesteps(num_timesteps, "15,15,8,6,6") |
| 1211 | |
| 1212 | # Help reduce DDIM artifacts from noisiest timesteps. |
| 1213 | steps.remove(num_timesteps - 1) |
| 1214 | steps.add(num_timesteps - 3) |
| 1215 | return steps |
| 1216 | section_counts = [int(x) for x in section_counts.split(",")] |
| 1217 | size_per = num_timesteps // len(section_counts) |
| 1218 | extra = num_timesteps % len(section_counts) |
| 1219 | start_idx = 0 |
| 1220 | all_steps = [] |
| 1221 | for i, section_count in enumerate(section_counts): |
| 1222 | size = size_per + (1 if i < extra else 0) |
| 1223 | if size < section_count: |
| 1224 | raise ValueError( |
| 1225 | f"cannot divide section of {size} steps into {section_count}") |
| 1226 | if section_count <= 1: |
| 1227 | frac_stride = 1 |
| 1228 | else: |
| 1229 | frac_stride = (size - 1) / (section_count - 1) |
| 1230 | cur_idx = 0.0 |
| 1231 | taken_steps = [] |
| 1232 | for _ in range(section_count): |
| 1233 | taken_steps.append(start_idx + round(cur_idx)) |
| 1234 | cur_idx += frac_stride |
| 1235 | all_steps += taken_steps |
| 1236 | start_idx += size |