Compute what the p_mean_variance output would have been, should the model's score function be conditioned by cond_fn. See condition_mean() for details on cond_fn. Unlike condition_mean(), this instead uses the conditioning strategy from Song et al (2020).
(self, cond_fn, p_mean_var, x, t, model_kwargs=None)
| 385 | return new_mean |
| 386 | |
| 387 | def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None): |
| 388 | """ |
| 389 | Compute what the p_mean_variance output would have been, should the |
| 390 | model's score function be conditioned by cond_fn. |
| 391 | |
| 392 | See condition_mean() for details on cond_fn. |
| 393 | |
| 394 | Unlike condition_mean(), this instead uses the conditioning strategy |
| 395 | from Song et al (2020). |
| 396 | """ |
| 397 | alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) |
| 398 | |
| 399 | eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"]) |
| 400 | eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **model_kwargs) |
| 401 | |
| 402 | out = p_mean_var.copy() |
| 403 | out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps) |
| 404 | out["mean"], _, _ = self.q_posterior_mean_variance(x_start=out["pred_xstart"], x_t=x, t=t) |
| 405 | return out |
| 406 | |
| 407 | def p_sample( |
| 408 | self, |
no test coverage detected