Spaces:
Runtime error
Runtime error
File size: 3,152 Bytes
f9a674e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import torch
import torch.nn as nn
import numpy as np
from tqdm import tqdm
from ldm.util import default
from ldm.modules.diffusionmodules.util import extract_into_tensor
from .ddpm import DDPM
class LatentDiffusion(DDPM):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# hardcoded
self.clip_denoised = False
def q_sample(self, x_start, t, noise=None):
noise = default(noise, lambda: torch.randn_like(x_start))
return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
"Does not support DDPM sampling anymore. Only do DDIM or PLMS"
# = = = = = = = = = = = = Below is for sampling = = = = = = = = = = = = #
# def predict_start_from_noise(self, x_t, t, noise):
# return ( extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
# extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise )
# def q_posterior(self, x_start, x_t, t):
# posterior_mean = (
# extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start +
# extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
# )
# posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape)
# posterior_log_variance_clipped = extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape)
# return posterior_mean, posterior_variance, posterior_log_variance_clipped
# def p_mean_variance(self, model, x, c, t):
# model_out = model(x, t, c)
# x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
# if self.clip_denoised:
# x_recon.clamp_(-1., 1.)
# model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
# return model_mean, posterior_variance, posterior_log_variance, x_recon
# @torch.no_grad()
# def p_sample(self, model, x, c, t):
# b, *_, device = *x.shape, x.device
# model_mean, _, model_log_variance, x0 = self.p_mean_variance(model, x=x, c=c, t=t, )
# noise = torch.randn_like(x)
# # no noise when t == 0
# nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
# return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0
# @torch.no_grad()
# def p_sample_loop(self, model, shape, c):
# device = self.betas.device
# b = shape[0]
# img = torch.randn(shape, device=device)
# iterator = tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps)
# for i in iterator:
# ts = torch.full((b,), i, device=device, dtype=torch.long)
# img, x0 = self.p_sample(model, img, c, ts)
# return img
# @torch.no_grad()
# def sample(self, model, shape, c, uc=None, guidance_scale=None):
# return self.p_sample_loop(model, shape, c)
|