| 22 | |
| 23 | |
| 24 | class GaussianDiffusion(nn.Module): |
| 25 | def __init__(self, model, T, schedule): |
| 26 | super().__init__() |
| 27 | self.visual = False |
| 28 | if self.visual: |
| 29 | self.num = 0 |
| 30 | self.model = model |
| 31 | self.T = T |
| 32 | self.schedule = schedule |
| 33 | betas = self.schedule.get_betas() |
| 34 | self.register_buffer('betas', betas.float()) |
| 35 | alphas = 1. - self.betas |
| 36 | alphas_bar = torch.cumprod(alphas, dim=0) |
| 37 | alphas_bar_prev = F.pad(alphas_bar, [1, 0], value=1)[:T] |
| 38 | gammas = alphas_bar |
| 39 | |
| 40 | self.register_buffer('coeff1', torch.sqrt(1. / alphas)) |
| 41 | self.register_buffer('coeff2', self.coeff1 * (1. - alphas) / torch.sqrt(1. - alphas_bar)) |
| 42 | self.register_buffer('posterior_var', self.betas * (1. - alphas_bar_prev) / (1. - alphas_bar)) |
| 43 | |
| 44 | # calculation for q(y_t|y_{t-1}) |
| 45 | self.register_buffer('gammas', gammas) |
| 46 | self.register_buffer('sqrt_one_minus_gammas', np.sqrt(1 - gammas)) |
| 47 | self.register_buffer('sqrt_gammas', np.sqrt(gammas)) |
| 48 | |
| 49 | def predict_xt_prev_mean_from_eps(self, x_t, t, eps): |
| 50 | assert x_t.shape == eps.shape |
| 51 | return extract(self.coeff1, t, x_t.shape) * x_t - extract(self.coeff2, t, x_t.shape) * eps |
| 52 | |
| 53 | def p_mean_variance(self, x_t, cond_, t): |
| 54 | # below: only log_variance is used in the KL computations |
| 55 | var = torch.cat([self.posterior_var[1:2], self.betas[1:]]) |
| 56 | #var = self.betas |
| 57 | var = extract(var, t, x_t.shape) |
| 58 | eps = self.model(torch.cat((x_t, cond_), dim=1), t) |
| 59 | # nonEps = self.model(x_t, t, torch.zeros_like(labels).to(labels.device)) |
| 60 | # eps = (1. + self.w) * eps - self.w * nonEps |
| 61 | xt_prev_mean = self.predict_xt_prev_mean_from_eps(x_t, t, eps=eps) |
| 62 | return xt_prev_mean, var |
| 63 | |
| 64 | def noisy_image(self, t, y): |
| 65 | """ Compute y_noisy according to (6) p15 of [2]""" |
| 66 | noise = torch.randn_like(y) |
| 67 | y_noisy = extract_(self.sqrt_gammas, t, y.shape) * y + extract_(self.sqrt_one_minus_gammas, t, noise.shape) * noise |
| 68 | return y_noisy, noise |
| 69 | |
| 70 | def forward(self, x_T, cond, pre_ori='False'): |
| 71 | """ |
| 72 | Algorithm 2. |
| 73 | """ |
| 74 | x_t = x_T |
| 75 | cond_ = cond |
| 76 | for time_step in reversed(range(self.T)): |
| 77 | print("time_step: ", time_step) |
| 78 | t = x_t.new_ones([x_T.shape[0], ], dtype=torch.long) * time_step |
| 79 | if pre_ori == 'False': |
| 80 | mean, var = self.p_mean_variance(x_t=x_t, t=t, cond_=cond_) |
| 81 | if time_step > 0: |
no outgoing calls
no test coverage detected