A diffusion process which can skip steps in a base diffusion process. :param use_timesteps: a collection (sequence or set) of timesteps from the original diffusion process to retain. :param kwargs: the kwargs to create the base diffusion process.
| 294 | |
| 295 | |
| 296 | class SpacedDiffusion(GaussianDiffusion): |
| 297 | """ |
| 298 | A diffusion process which can skip steps in a base diffusion process. |
| 299 | :param use_timesteps: a collection (sequence or set) of timesteps from the |
| 300 | original diffusion process to retain. |
| 301 | :param kwargs: the kwargs to create the base diffusion process. |
| 302 | """ |
| 303 | |
| 304 | def __init__(self, use_timesteps, **kwargs): |
| 305 | self.use_timesteps = set(use_timesteps) |
| 306 | self.timestep_map = [] |
| 307 | self.original_num_steps = len(kwargs["betas"]) |
| 308 | |
| 309 | base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa |
| 310 | last_alpha_cumprod = 1.0 |
| 311 | new_betas = [] |
| 312 | for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): |
| 313 | if i in self.use_timesteps: |
| 314 | new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) |
| 315 | last_alpha_cumprod = alpha_cumprod |
| 316 | self.timestep_map.append(i) |
| 317 | kwargs["betas"] = np.array(new_betas) |
| 318 | super().__init__(**kwargs) |
| 319 | |
| 320 | def p_mean_variance( |
| 321 | self, model, *args, **kwargs |
| 322 | ): # pylint: disable=signature-differs |
| 323 | return super().p_mean_variance(self._wrap_model(model), *args, **kwargs) |
| 324 | |
| 325 | def training_losses( |
| 326 | self, model, *args, **kwargs |
| 327 | ): # pylint: disable=signature-differs |
| 328 | return super().training_losses(self._wrap_model(model), *args, **kwargs) |
| 329 | |
| 330 | def condition_mean(self, cond_fn, *args, **kwargs): |
| 331 | return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs) |
| 332 | |
| 333 | def condition_score(self, cond_fn, *args, **kwargs): |
| 334 | return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs) |
| 335 | |
| 336 | def _wrap_model(self, model): |
| 337 | if isinstance(model, _WrappedModel): |
| 338 | return model |
| 339 | return _WrappedModel( |
| 340 | model, self.timestep_map, self.rescale_timesteps, self.original_num_steps |
| 341 | ) |
| 342 | |
| 343 | def _scale_timesteps(self, t): |
| 344 | # Scaling is done by the wrapped model. |
| 345 | return t |
| 346 | |
| 347 | |
| 348 | class _WrappedModel: |
nothing calls this directly
no outgoing calls
no test coverage detected