| 8 | |
| 9 | |
| 10 | class AbstractLowScaleModel(nn.Module): |
| 11 | # for concatenating a downsampled image to the latent representation |
| 12 | def __init__(self, noise_schedule_config=None): |
| 13 | super(AbstractLowScaleModel, self).__init__() |
| 14 | if noise_schedule_config is not None: |
| 15 | self.register_schedule(**noise_schedule_config) |
| 16 | |
| 17 | def register_schedule(self, beta_schedule="linear", timesteps=1000, |
| 18 | linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): |
| 19 | betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end, |
| 20 | cosine_s=cosine_s) |
| 21 | alphas = 1. - betas |
| 22 | alphas_cumprod = np.cumprod(alphas, axis=0) |
| 23 | alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1]) |
| 24 | |
| 25 | timesteps, = betas.shape |
| 26 | self.num_timesteps = int(timesteps) |
| 27 | self.linear_start = linear_start |
| 28 | self.linear_end = linear_end |
| 29 | assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep' |
| 30 | |
| 31 | to_torch = partial(torch.tensor, dtype=torch.float32) |
| 32 | |
| 33 | self.register_buffer('betas', to_torch(betas)) |
| 34 | self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) |
| 35 | self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev)) |
| 36 | |
| 37 | # calculations for diffusion q(x_t | x_{t-1}) and others |
| 38 | self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod))) |
| 39 | self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod))) |
| 40 | self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod))) |
| 41 | self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod))) |
| 42 | self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1))) |
| 43 | |
| 44 | def q_sample(self, x_start, t, noise=None): |
| 45 | noise = default(noise, lambda: torch.randn_like(x_start)) |
| 46 | return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + |
| 47 | extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise) |
| 48 | |
| 49 | def forward(self, x): |
| 50 | return x, None |
| 51 | |
| 52 | def decode(self, x): |
| 53 | return x |
| 54 | |
| 55 | |
| 56 | class SimpleImageConcat(AbstractLowScaleModel): |
nothing calls this directly
no outgoing calls
no test coverage detected