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.
| 1238 | |
| 1239 | |
| 1240 | class SpacedDiffusion(GaussianDiffusion): |
| 1241 | """ |
| 1242 | A diffusion process which can skip steps in a base diffusion process. |
| 1243 | |
| 1244 | :param use_timesteps: a collection (sequence or set) of timesteps from the |
| 1245 | original diffusion process to retain. |
| 1246 | :param kwargs: the kwargs to create the base diffusion process. |
| 1247 | """ |
| 1248 | |
| 1249 | def __init__(self, use_timesteps, **kwargs): |
| 1250 | self.use_timesteps = set(use_timesteps) |
| 1251 | self.timestep_map = [] |
| 1252 | self.original_num_steps = len(kwargs["betas"]) |
| 1253 | |
| 1254 | base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa |
| 1255 | last_alpha_cumprod = 1.0 |
| 1256 | new_betas = [] |
| 1257 | for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): |
| 1258 | if i in self.use_timesteps: |
| 1259 | new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) |
| 1260 | last_alpha_cumprod = alpha_cumprod |
| 1261 | self.timestep_map.append(i) |
| 1262 | kwargs["betas"] = np.array(new_betas) |
| 1263 | super().__init__(**kwargs) |
| 1264 | |
| 1265 | def p_mean_variance(self, model, *args, **kwargs): |
| 1266 | return super().p_mean_variance(self._wrap_model(model), *args, |
| 1267 | **kwargs) |
| 1268 | |
| 1269 | def condition_mean(self, cond_fn, *args, **kwargs): |
| 1270 | return super().condition_mean(self._wrap_model(cond_fn), *args, |
| 1271 | **kwargs) |
| 1272 | |
| 1273 | def condition_score(self, cond_fn, *args, **kwargs): |
| 1274 | return super().condition_score(self._wrap_model(cond_fn), *args, |
| 1275 | **kwargs) |
| 1276 | |
| 1277 | def _wrap_model(self, model): |
| 1278 | if isinstance(model, _WrappedModel): |
| 1279 | return model |
| 1280 | return _WrappedModel(model, self.timestep_map, self.original_num_steps) |
| 1281 | |
| 1282 | |
| 1283 | class _WrappedModel: |