Sample the v-parameterized vector field at time t
(self, fm_x, fm_t, **kwargs)
| 80 | self.register_buffer('rectified_sqrt_alphas_cumprod_full', self.sqrt_one_minus_alphas_cumprod_full / (self.sqrt_alphas_cumprod_full + self.sqrt_one_minus_alphas_cumprod_full)) |
| 81 | |
| 82 | def sample_vt(self, fm_x, fm_t, **kwargs): |
| 83 | """ |
| 84 | Sample the v-parameterized vector field at time t |
| 85 | """ |
| 86 | dm_t = self.convert_fm_t_to_dm_t(fm_t) |
| 87 | # print(fm_t, dm_t) |
| 88 | dm_x = self.convert_fm_xt_to_dm_xt(fm_x, fm_t) |
| 89 | # vt = self.net(dm_x, dm_t, **kwargs) |
| 90 | vt = forward_with_cfg(dm_x, dm_t, self.net, **kwargs) |
| 91 | |
| 92 | # TODO: ugly fix for nan values!!! |
| 93 | if torch.isnan(vt).any(): |
| 94 | vt[torch.isnan(vt)] = 0 |
| 95 | |
| 96 | # vt = self.forward(x=dm_x, t=dm_t, **kwargs) |
| 97 | if self.diffusion_parameterization == 'v': |
| 98 | vector_field = self.get_vector_field_from_v(vt, dm_x, dm_t) |
| 99 | elif self.diffusion_parameterization == 'eps': |
| 100 | vector_field = self.get_vector_field_from_eps(vt, dm_x, dm_t) |
| 101 | return vector_field |
| 102 | |
| 103 | def convert_fm_t_to_dm_t(self, t): |
| 104 | """ |
no test coverage detected