| 229 | |
| 230 | |
| 231 | def space_timesteps(num_timesteps, section_counts): |
| 232 | if isinstance(section_counts, str): |
| 233 | if section_counts.startswith("ddim"): |
| 234 | desired_count = int(section_counts[len("ddim") :]) |
| 235 | for i in range(1, num_timesteps): |
| 236 | if len(range(0, num_timesteps, i)) == desired_count: |
| 237 | return set(range(0, num_timesteps, i)) |
| 238 | raise ValueError( |
| 239 | f"cannot create exactly {num_timesteps} steps with an integer stride" |
| 240 | ) |
| 241 | section_counts = [int(x) for x in section_counts.split(",")] # [250,] |
| 242 | size_per = num_timesteps // len(section_counts) |
| 243 | extra = num_timesteps % len(section_counts) |
| 244 | start_idx = 0 |
| 245 | all_steps = [] |
| 246 | for i, section_count in enumerate(section_counts): |
| 247 | size = size_per + (1 if i < extra else 0) |
| 248 | if size < section_count: |
| 249 | raise ValueError( |
| 250 | f"cannot divide section of {size} steps into {section_count}" |
| 251 | ) |
| 252 | if section_count <= 1: |
| 253 | frac_stride = 1 |
| 254 | else: |
| 255 | frac_stride = (size - 1) / (section_count - 1) |
| 256 | cur_idx = 0.0 |
| 257 | taken_steps = [] |
| 258 | for _ in range(section_count): |
| 259 | taken_steps.append(start_idx + round(cur_idx)) |
| 260 | cur_idx += frac_stride |
| 261 | all_steps += taken_steps |
| 262 | start_idx += size |
| 263 | return set(all_steps) |
| 264 | |
| 265 | |
| 266 | def chunk(it, size): |