A diffusion process which can skip steps in a base diffusion process. :param use_timesteps: (unordered) timesteps from the original diffusion process to retain. :param kwargs: the kwargs to create the base diffusion process.
| 1002 | |
| 1003 | |
| 1004 | class SpacedDiffusion(GaussianDiffusion): |
| 1005 | """ |
| 1006 | A diffusion process which can skip steps in a base diffusion process. |
| 1007 | :param use_timesteps: (unordered) timesteps from the original diffusion |
| 1008 | process to retain. |
| 1009 | :param kwargs: the kwargs to create the base diffusion process. |
| 1010 | """ |
| 1011 | |
| 1012 | def __init__(self, use_timesteps: Iterable[int], **kwargs): |
| 1013 | self.use_timesteps = set(use_timesteps) |
| 1014 | self.timestep_map = [] |
| 1015 | self.original_num_steps = len(kwargs["betas"]) |
| 1016 | |
| 1017 | base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa |
| 1018 | last_alpha_cumprod = 1.0 |
| 1019 | new_betas = [] |
| 1020 | for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): |
| 1021 | if i in self.use_timesteps: |
| 1022 | new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) |
| 1023 | last_alpha_cumprod = alpha_cumprod |
| 1024 | self.timestep_map.append(i) |
| 1025 | kwargs["betas"] = np.array(new_betas) |
| 1026 | super().__init__(**kwargs) |
| 1027 | |
| 1028 | def p_mean_variance(self, model, *args, **kwargs): |
| 1029 | return super().p_mean_variance(self._wrap_model(model), *args, **kwargs) |
| 1030 | |
| 1031 | def training_losses(self, model, *args, **kwargs): |
| 1032 | return super().training_losses(self._wrap_model(model), *args, **kwargs) |
| 1033 | |
| 1034 | def condition_mean(self, cond_fn, *args, **kwargs): |
| 1035 | return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs) |
| 1036 | |
| 1037 | def condition_score(self, cond_fn, *args, **kwargs): |
| 1038 | return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs) |
| 1039 | |
| 1040 | def _wrap_model(self, model): |
| 1041 | if isinstance(model, _WrappedModel): |
| 1042 | return model |
| 1043 | return _WrappedModel(model, self.timestep_map, self.original_num_steps) |
| 1044 | |
| 1045 | |
| 1046 | class _WrappedModel: |
no outgoing calls
no test coverage detected