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