import torch import torch.nn as nn import torch.nn.functional as F import numpy as np import numpy.random as npr import copy from functools import partial from contextlib import contextmanager from lib.model_zoo.common.get_model import get_model, register from lib.log_service import print_log version = '0' symbol = 'sd' from .diffusion_utils import \ count_params, extract_into_tensor, make_beta_schedule from .distributions import normal_kl, DiagonalGaussianDistribution from .ema import LitEma def highlight_print(info): print_log('') print_log(''.join(['#']*(len(info)+4))) print_log('# '+info+' #') print_log(''.join(['#']*(len(info)+4))) print_log('') class DDPM(nn.Module): def __init__(self, unet_config, timesteps=1000, use_ema=True, beta_schedule="linear", beta_linear_start=1e-4, beta_linear_end=2e-2, loss_type="l2", clip_denoised=True, cosine_s=8e-3, given_betas=None, l_simple_weight=1., original_elbo_weight=0., v_posterior=0., # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta parameterization="eps", use_positional_encodings=False, learn_logvar=False, logvar_init=0, ): super().__init__() assert parameterization in ["eps", "x0"], \ 'currently only supporting "eps" and "x0"' self.parameterization = parameterization highlight_print("Running in {} mode".format(self.parameterization)) self.cond_stage_model = None self.clip_denoised = clip_denoised self.use_positional_encodings = use_positional_encodings from collections import OrderedDict self.model = nn.Sequential(OrderedDict([('diffusion_model', get_model()(unet_config))])) # TODO: Remove this ugly trick to match SD with deprecated version, after no bug with the module. self.use_ema = use_ema if self.use_ema: self.model_ema = LitEma(self.model) print_log(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.") self.v_posterior = v_posterior self.l_simple_weight = l_simple_weight self.original_elbo_weight = original_elbo_weight self.register_schedule( given_betas=given_betas, beta_schedule=beta_schedule, timesteps=timesteps, linear_start=beta_linear_start, linear_end=beta_linear_end, cosine_s=cosine_s) self.loss_type = loss_type self.learn_logvar = learn_logvar self.logvar = torch.full( fill_value=logvar_init, size=(self.num_timesteps,)) if self.learn_logvar: self.logvar = nn.Parameter(self.logvar, requires_grad=True) def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): if given_betas is not None: betas = given_betas else: betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s) alphas = 1. - betas alphas_cumprod = np.cumprod(alphas, axis=0) alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1]) timesteps, = betas.shape self.num_timesteps = int(timesteps) self.linear_start = linear_start self.linear_end = linear_end assert alphas_cumprod.shape[0] == self.num_timesteps, \ 'alphas have to be defined for each timestep' to_torch = partial(torch.tensor, dtype=torch.float32) self.register_buffer('betas', to_torch(betas)) self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev)) # calculations for diffusion q(x_t | x_{t-1}) and others self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod))) self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod))) self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod))) self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod))) self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1))) # calculations for posterior q(x_{t-1} | x_t, x_0) posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / ( 1. - alphas_cumprod) + self.v_posterior * betas # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t) self.register_buffer('posterior_variance', to_torch(posterior_variance)) # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20)))) self.register_buffer('posterior_mean_coef1', to_torch( betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))) self.register_buffer('posterior_mean_coef2', to_torch( (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))) if self.parameterization == "eps": lvlb_weights = self.betas ** 2 / ( 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod)) elif self.parameterization == "x0": lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod)) else: raise NotImplementedError("mu not supported") # TODO how to choose this term lvlb_weights[0] = lvlb_weights[1] self.register_buffer('lvlb_weights', lvlb_weights, persistent=False) assert not torch.isnan(self.lvlb_weights).all() @contextmanager def ema_scope(self, context=None): if self.use_ema: self.model_ema.store(self.model.parameters()) self.model_ema.copy_to(self.model) if context is not None: print_log(f"{context}: Switched to EMA weights") try: yield None finally: if self.use_ema: self.model_ema.restore(self.model.parameters()) if context is not None: print_log(f"{context}: Restored training weights") def q_mean_variance(self, x_start, t): """ Get the distribution q(x_t | x_0). :param x_start: the [N x C x ...] tensor of noiseless inputs. :param t: the number of diffusion steps (minus 1). Here, 0 means one step. :return: A tuple (mean, variance, log_variance), all of x_start's shape. """ mean = (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start) variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) log_variance = extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape) return mean, variance, log_variance def predict_start_from_noise(self, x_t, t, noise): value1 = extract_into_tensor( self.sqrt_recip_alphas_cumprod, t, x_t.shape) value2 = extract_into_tensor( self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) return value1*x_t -value2*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, x, t, clip_denoised: bool): model_out = self.model(x, t) if self.parameterization == "eps": x_recon = self.predict_start_from_noise(x, t=t, noise=model_out) elif self.parameterization == "x0": x_recon = model_out if 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 @torch.no_grad() def p_sample(self, x, t, clip_denoised=True, repeat_noise=False): b, *_, device = *x.shape, x.device model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised) noise = noise_like(x.shape, device, repeat_noise) # 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 @torch.no_grad() def p_sample_loop(self, shape, return_intermediates=False): device = self.betas.device b = shape[0] img = torch.randn(shape, device=device) intermediates = [img] for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps): img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long), clip_denoised=self.clip_denoised) if i % self.log_every_t == 0 or i == self.num_timesteps - 1: intermediates.append(img) if return_intermediates: return img, intermediates return img @torch.no_grad() def sample(self, batch_size=16, return_intermediates=False): image_size = self.image_size channels = self.channels return self.p_sample_loop((batch_size, channels, image_size, image_size), return_intermediates=return_intermediates) def q_sample(self, x_start, t, noise=None): noise = torch.randn_like(x_start) if noise is None else noise 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) def get_loss(self, pred, target, mean=True): if self.loss_type == 'l1': loss = (target - pred).abs() if mean: loss = loss.mean() elif self.loss_type == 'l2': if mean: loss = torch.nn.functional.mse_loss(target, pred) else: loss = torch.nn.functional.mse_loss(target, pred, reduction='none') else: raise NotImplementedError("unknown loss type '{loss_type}'") return loss def p_losses(self, x_start, t, noise=None): noise = default(noise, lambda: torch.randn_like(x_start)) x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) model_out = self.model(x_noisy, t) loss_dict = {} if self.parameterization == "eps": target = noise elif self.parameterization == "x0": target = x_start else: raise NotImplementedError(f"Paramterization {self.parameterization} not yet supported") loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3]) log_prefix = 'train' if self.training else 'val' loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()}) loss_simple = loss.mean() * self.l_simple_weight loss_vlb = (self.lvlb_weights[t] * loss).mean() loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb}) loss = loss_simple + self.original_elbo_weight * loss_vlb loss_dict.update({f'{log_prefix}/loss': loss}) return loss, loss_dict def forward(self, x, *args, **kwargs): # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size # assert h == img_size and w == img_size, f'height and width of image must be {img_size}' t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long() return self.p_losses(x, t, *args, **kwargs) def on_train_batch_end(self, *args, **kwargs): if self.use_ema: self.model_ema(self.model) @register('sd_t2i', version) class SD_T2I(DDPM): def __init__(self, first_stage_config, cond_stage_config, num_timesteps_cond=None, cond_stage_trainable=False, scale_factor=1.0, scale_by_std=False, *args, **kwargs): self.num_timesteps_cond = num_timesteps_cond \ if num_timesteps_cond is not None else 1 self.scale_by_std = scale_by_std assert self.num_timesteps_cond <= kwargs['timesteps'] super().__init__(*args, **kwargs) self.first_stage_model = get_model()(first_stage_config) self.cond_stage_model = get_model()(cond_stage_config) self.concat_mode = 'crossattn' self.cond_stage_trainable = cond_stage_trainable if not scale_by_std: self.scale_factor = scale_factor else: self.register_buffer('scale_factor', torch.tensor(scale_factor)) self.device = 'cpu' def to(self, device): self.device = device super().to(device) @torch.no_grad() def on_train_batch_start(self, x): # only for very first batch if self.scale_by_std: assert self.scale_factor == 1., \ 'rather not use custom rescaling and std-rescaling simultaneously' # set rescale weight to 1./std of encodings encoder_posterior = self.encode_first_stage(x) z = self.get_first_stage_encoding(encoder_posterior).detach() del self.scale_factor self.register_buffer('scale_factor', 1. / z.flatten().std()) highlight_print("setting self.scale_factor to {}".format(self.scale_factor)) def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): super().register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s) self.shorten_cond_schedule = self.num_timesteps_cond > 1 if self.shorten_cond_schedule: self.make_cond_schedule() def make_cond_schedule(self, ): self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long) ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long() self.cond_ids[:self.num_timesteps_cond] = ids @torch.no_grad() def encode_image(self, im): encoder_posterior = self.first_stage_model.encode(im) z = self.get_first_stage_encoding(encoder_posterior).detach() return z def get_first_stage_encoding(self, encoder_posterior): if isinstance(encoder_posterior, DiagonalGaussianDistribution): z = encoder_posterior.sample() elif isinstance(encoder_posterior, torch.Tensor): z = encoder_posterior else: raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented") return self.scale_factor * z @torch.no_grad() def decode_image(self, z, predict_cids=False, force_not_quantize=False): z = 1. / self.scale_factor * z return self.first_stage_model.decode(z) @torch.no_grad() def encode_text(self, text): return self.get_learned_conditioning(text) def get_learned_conditioning(self, c): if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode): c = self.cond_stage_model.encode(c) if isinstance(c, DiagonalGaussianDistribution): c = c.mode() else: c = self.cond_stage_model(c) return c def forward(self, x, c, noise=None): t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=x.device).long() if self.cond_stage_trainable: c = self.get_learned_conditioning(c) return self.p_losses(x, c, t, noise) def apply_model(self, x_noisy, t, cond): return self.model.diffusion_model(x_noisy, t, cond) def p_losses(self, x_start, cond, t, noise=None): noise = torch.randn_like(x_start) if noise is None else noise x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) model_output = self.apply_model(x_noisy, t, cond) loss_dict = {} prefix = 'train' if self.training else 'val' if self.parameterization == "x0": target = x_start elif self.parameterization == "eps": target = noise else: raise NotImplementedError() loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3]) loss_dict['loss_simple'] = loss_simple.mean() logvar_t = self.logvar[t].to(self.device) loss = loss_simple / torch.exp(logvar_t) + logvar_t if self.learn_logvar: loss_dict['loss_gamma'] = loss.mean() loss_dict['logvar' ] = self.logvar.data.mean() loss = self.l_simple_weight * loss.mean() loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3)) loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean() loss_dict['loss_vlb'] = loss_vlb loss += (self.original_elbo_weight * loss_vlb) loss_dict.update({'Loss': loss}) return loss, loss_dict def _predict_eps_from_xstart(self, x_t, t, pred_xstart): return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart) / \ extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) def _prior_bpd(self, x_start): """ Get the prior KL term for the variational lower-bound, measured in bits-per-dim. This term can't be optimized, as it only depends on the encoder. :param x_start: the [N x C x ...] tensor of inputs. :return: a batch of [N] KL values (in bits), one per batch element. """ batch_size = x_start.shape[0] t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t) kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0) return mean_flat(kl_prior) / np.log(2.0) def p_mean_variance(self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False, return_x0=False, score_corrector=None, corrector_kwargs=None): t_in = t model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids) if score_corrector is not None: assert self.parameterization == "eps" model_out = score_corrector.modify_score(self, model_out, x, t, c, **corrector_kwargs) if return_codebook_ids: model_out, logits = model_out if self.parameterization == "eps": x_recon = self.predict_start_from_noise(x, t=t, noise=model_out) elif self.parameterization == "x0": x_recon = model_out else: raise NotImplementedError() if clip_denoised: x_recon.clamp_(-1., 1.) if quantize_denoised: x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon) model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t) if return_codebook_ids: return model_mean, posterior_variance, posterior_log_variance, logits elif return_x0: return model_mean, posterior_variance, posterior_log_variance, x_recon else: return model_mean, posterior_variance, posterior_log_variance @torch.no_grad() def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False, return_codebook_ids=False, quantize_denoised=False, return_x0=False, temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None): b, *_, device = *x.shape, x.device outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised, return_codebook_ids=return_codebook_ids, quantize_denoised=quantize_denoised, return_x0=return_x0, score_corrector=score_corrector, corrector_kwargs=corrector_kwargs) if return_codebook_ids: raise DeprecationWarning("Support dropped.") model_mean, _, model_log_variance, logits = outputs elif return_x0: model_mean, _, model_log_variance, x0 = outputs else: model_mean, _, model_log_variance = outputs noise = noise_like(x.shape, device, repeat_noise) * temperature if noise_dropout > 0.: noise = torch.nn.functional.dropout(noise, p=noise_dropout) # no noise when t == 0 nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1))) if return_codebook_ids: return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, logits.argmax(dim=1) if return_x0: return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0 else: return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise @torch.no_grad() def progressive_denoising(self, cond, shape, verbose=True, callback=None, quantize_denoised=False, img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None, log_every_t=None): if not log_every_t: log_every_t = self.log_every_t timesteps = self.num_timesteps if batch_size is not None: b = batch_size if batch_size is not None else shape[0] shape = [batch_size] + list(shape) else: b = batch_size = shape[0] if x_T is None: img = torch.randn(shape, device=self.device) else: img = x_T intermediates = [] if cond is not None: if isinstance(cond, dict): cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else list(map(lambda x: x[:batch_size], cond[key])) for key in cond} else: cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size] if start_T is not None: timesteps = min(timesteps, start_T) iterator = tqdm(reversed(range(0, timesteps)), desc='Progressive Generation', total=timesteps) if verbose else reversed( range(0, timesteps)) if type(temperature) == float: temperature = [temperature] * timesteps for i in iterator: ts = torch.full((b,), i, device=self.device, dtype=torch.long) if self.shorten_cond_schedule: assert self.model.conditioning_key != 'hybrid' tc = self.cond_ids[ts].to(cond.device) cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) img, x0_partial = self.p_sample(img, cond, ts, clip_denoised=self.clip_denoised, quantize_denoised=quantize_denoised, return_x0=True, temperature=temperature[i], noise_dropout=noise_dropout, score_corrector=score_corrector, corrector_kwargs=corrector_kwargs) if mask is not None: assert x0 is not None img_orig = self.q_sample(x0, ts) img = img_orig * mask + (1. - mask) * img if i % log_every_t == 0 or i == timesteps - 1: intermediates.append(x0_partial) if callback: callback(i) if img_callback: img_callback(img, i) return img, intermediates @torch.no_grad() def p_sample_loop(self, cond, shape, return_intermediates=False, x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False, mask=None, x0=None, img_callback=None, start_T=None, log_every_t=None): if not log_every_t: log_every_t = self.log_every_t device = self.betas.device b = shape[0] if x_T is None: img = torch.randn(shape, device=device) else: img = x_T intermediates = [img] if timesteps is None: timesteps = self.num_timesteps if start_T is not None: timesteps = min(timesteps, start_T) iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed( range(0, timesteps)) if mask is not None: assert x0 is not None assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match for i in iterator: ts = torch.full((b,), i, device=device, dtype=torch.long) if self.shorten_cond_schedule: assert self.model.conditioning_key != 'hybrid' tc = self.cond_ids[ts].to(cond.device) cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) img = self.p_sample(img, cond, ts, clip_denoised=self.clip_denoised, quantize_denoised=quantize_denoised) if mask is not None: img_orig = self.q_sample(x0, ts) img = img_orig * mask + (1. - mask) * img if i % log_every_t == 0 or i == timesteps - 1: intermediates.append(img) if callback: callback(i) if img_callback: img_callback(img, i) if return_intermediates: return img, intermediates return img @torch.no_grad() def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None, verbose=True, timesteps=None, quantize_denoised=False, mask=None, x0=None, shape=None,**kwargs): if shape is None: shape = (batch_size, self.channels, self.image_size, self.image_size) if cond is not None: if isinstance(cond, dict): cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else list(map(lambda x: x[:batch_size], cond[key])) for key in cond} else: cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size] return self.p_sample_loop(cond, shape, return_intermediates=return_intermediates, x_T=x_T, verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised, mask=mask, x0=x0) @register('sd_t2i_split_trans_pg', version) class SD_T2I_SplitTransPG(SD_T2I): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.parameter_group = { # 'first_stage_model' : self.first_stage_model, # 'cond_stage_model' : self.cond_stage_model, 'transformers' : [v for n, v in self.model.named_parameters() if n.find('transformer_blocks')!=-1], 'other' :[v for n, v in self.model.named_parameters() if n.find('transformer_blocks')==-1], } @register('sd_dual_crossattn', version) class SD_Dual_CrossAttn(SD_T2I): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def is_part_of_trans(name): if name.find('.1.norm')!=-1: return True if name.find('.1.proj_in')!=-1: return True if name.find('.1.transformer_blocks')!=-1: return True if name.find('.1.proj_out')!=-1: return True return False self.parameter_group = { 'transformers' : [v for n, v in self.model.named_parameters() if is_part_of_trans(n)], 'other' :[v for n, v in self.model.named_parameters() if not is_part_of_trans(n)], } def apply_model(self, x_noisy, t, cond, cond_type): if cond_type in ['prompt', 'text']: which_attn = 0 elif cond_type in ['vision', 'visual', 'image']: which_attn = 1 elif isinstance(cond_type, float): assert 0 < cond_type < 1, \ 'A special cond_type that will doing a random mix between two input condition, '\ 'rand() < cond_type is text, else visual' which_attn = cond_type else: assert False return self.model.diffusion_model(x_noisy, t, cond, which_attn=which_attn) def p_losses(self, x_start, cond, t, noise=None, cond_type=None): noise = torch.randn_like(x_start) if noise is None else noise x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) model_output = self.apply_model(x_noisy, t, cond, cond_type=cond_type) loss_dict = {} prefix = 'train' if self.training else 'val' if self.parameterization == "x0": target = x_start elif self.parameterization == "eps": target = noise else: raise NotImplementedError() loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3]) loss_dict['loss_simple'] = loss_simple.mean() logvar_t = self.logvar[t].to(self.device) loss = loss_simple / torch.exp(logvar_t) + logvar_t if self.learn_logvar: loss_dict['loss_gamma'] = loss.mean() loss_dict['logvar' ] = self.logvar.data.mean() loss = self.l_simple_weight * loss.mean() loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3)) loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean() loss_dict['loss_vlb'] = loss_vlb loss += (self.original_elbo_weight * loss_vlb) loss_dict.update({'Loss': loss}) return loss, loss_dict @torch.no_grad() def clip_encode_text(self, text): clip_encode_type = self.cond_stage_model.encode_type self.cond_stage_model.encode_type = 'encode_text' embedding = self.get_learned_conditioning(text) self.cond_stage_model.encode_type = clip_encode_type return embedding @torch.no_grad() def clip_encode_vision(self, vision, encode_type='encode_vision'): clip_encode_type = self.cond_stage_model.encode_type self.cond_stage_model.encode_type = encode_type if isinstance(vision, torch.Tensor): vision = ((vision+1)/2).to('cpu').numpy() vision = np.transpose(vision, (0, 2, 3, 1)) vision = [vi for vi in vision] embedding = self.get_learned_conditioning(vision) self.cond_stage_model.encode_type = clip_encode_type return embedding def get_learned_conditioning(self, c): if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode): c = self.cond_stage_model.encode(c) if isinstance(c, DiagonalGaussianDistribution): c = c.mode() else: c = self.cond_stage_model(c) return c def forward(self, x, c, noise=None, cond_type=None): t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=x.device).long() if self.cond_stage_trainable: c = self.get_learned_conditioning(c) return self.p_losses(x, c, t, noise, cond_type=cond_type) @register('sd_variation', version) class SD_Variation(SD_T2I): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def is_part_of_trans(name): if name.find('.1.norm')!=-1: return True if name.find('.1.proj_in')!=-1: return True if name.find('.1.transformer_blocks')!=-1: return True if name.find('.1.proj_out')!=-1: return True return False self.parameter_group = { 'transformers' : [v for n, v in self.model.named_parameters() if is_part_of_trans(n)], 'other' :[v for n, v in self.model.named_parameters() if not is_part_of_trans(n)], } self.encode_image = None self.encode_text = None self._predict_eps_from_xstart = None self._prior_bpd = None self.p_mean_variance = None self.p_sample = None self.progressive_denoising = None self.p_sample_loop = None self.sample = None @torch.no_grad() def encode_input(self, im): encoder_posterior = self.first_stage_model.encode(im) if isinstance(encoder_posterior, DiagonalGaussianDistribution): z = encoder_posterior.sample() elif isinstance(encoder_posterior, torch.Tensor): z = encoder_posterior else: raise NotImplementedError("Encoder_posterior of type '{}' not yet implemented".format(type(encoder_posterior))) return z * self.scale_factor @torch.no_grad() def decode_latent(self, z): z = 1. / self.scale_factor * z return self.first_stage_model.decode(z) @torch.no_grad() def clip_encode_vision(self, vision): if isinstance(vision, list): if not isinstance(vision[0], torch.Tensor): import torchvision.transforms as tvtrans vision = [tvtrans.ToTensor()(i) for i in vision] vh = torch.stack(vision) elif isinstance(vision, torch.Tensor): vh = vision.unsqueeze(0) if (vision.shape==3) else vision assert len(vh.shape) == 4 else: raise ValueError vh = vh.to(self.device) return self.encode_conditioning(vh) # legacy def get_learned_conditioning(self, c): return self.encode_conditioning(c) def encode_conditioning(self, c): return self.cond_stage_model.encode(c) def forward(self, x, c, noise=None): t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=x.device).long() if self.cond_stage_trainable: c = self.encode_conditioning(c) return self.p_losses(x, c, t, noise) @register('sd_all_in_one', version) class SD_ALL_IN_ONE(DDPM): def __init__(self, autokl_cfg, optimus_cfg, clip_cfg, scale_factor=1.0, scale_by_std=False, *args, **kwargs): self.scale_by_std = scale_by_std super().__init__(*args, **kwargs) self.autokl = get_model()(autokl_cfg) self.optimus = get_model()(optimus_cfg) self.clip = get_model()(clip_cfg) self.concat_mode = 'crossattn' if not scale_by_std: self.scale_factor = scale_factor else: self.register_buffer('scale_factor', torch.tensor(scale_factor)) self.device = 'cpu' self.parameter_group = self.create_parameter_group() debug = 1 def create_parameter_group(self): def is_part_of_unet_image(name): if name.find('.unet_image.')!=-1: return True return False def is_part_of_unet_text(name): if name.find('.unet_text.')!=-1: return True return False def is_part_of_trans(name): if name.find('.1.norm')!=-1: return True if name.find('.1.proj_in')!=-1: return True if name.find('.1.transformer_blocks')!=-1: return True if name.find('.1.proj_out')!=-1: return True return False parameter_group = { 'image_trans' : [], 'image_rest' : [], 'text_trans' : [], 'text_rest' : [], 'rest' : [],} for pname, para in self.model.named_parameters(): if is_part_of_unet_image(pname): if is_part_of_trans(pname): parameter_group['image_trans'].append(para) else: parameter_group['image_rest'].append(para) elif is_part_of_unet_text(pname): if is_part_of_trans(pname): parameter_group['text_trans'].append(para) else: parameter_group['text_rest'].append(para) else: parameter_group['rest'].append(para) return parameter_group def to(self, device): self.device = device super().to(device) @torch.no_grad() def on_train_batch_start(self, x): # only for very first batch if self.scale_by_std: assert self.scale_factor == 1., \ 'rather not use custom rescaling and std-rescaling simultaneously' # set rescale weight to 1./std of encodings encoder_posterior = self.encode_first_stage(x) z = self.get_first_stage_encoding(encoder_posterior).detach() del self.scale_factor self.register_buffer('scale_factor', 1. / z.flatten().std()) highlight_print("setting self.scale_factor to {}".format(self.scale_factor)) @torch.no_grad() def autokl_encode(self, image): encoder_posterior = self.autokl.encode(image) z = encoder_posterior.sample() return self.scale_factor * z @torch.no_grad() def autokl_decode(self, z): z = 1. / self.scale_factor * z return self.autokl.decode(z) def mask_tokens(inputs, tokenizer, args): labels = inputs.clone() # We sample a few tokens in each sequence for masked-LM training (with probability args.mlm_probability defaults to 0.15 in Bert/RoBERTa) masked_indices = torch.bernoulli(torch.full(labels.shape, args.mlm_probability)).to(torch.uint8) labels[masked_indices==1] = -1 # We only compute loss on masked tokens # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK]) indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).to(torch.uint8) & masked_indices inputs[indices_replaced] = tokenizer.convert_tokens_to_ids(tokenizer.mask_token) # 10% of the time, we replace masked input tokens with random word indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).to(torch.uint8) & masked_indices & ~indices_replaced indices_random = indices_random random_words = torch.randint(len(tokenizer), labels.shape, dtype=torch.long) inputs[indices_random] = random_words[indices_random] # The rest of the time (10% of the time) we keep the masked input tokens unchanged return inputs, labels @torch.no_grad() def optimus_encode(self, text): tokenizer = self.optimus.tokenizer_encoder token = [tokenizer.tokenize(sentence.lower()) for sentence in text] token_id = [] for tokeni in token: token_sentence = [tokenizer._convert_token_to_id(i) for i in tokeni] token_sentence = tokenizer.add_special_tokens_single_sentence(token_sentence) token_id.append(torch.LongTensor(token_sentence)) token_id = torch._C._nn.pad_sequence(token_id, batch_first=True, padding_value=0.0) token_id = token_id.to(self.device) z = self.optimus.encoder(token_id, attention_mask=(token_id > 0).float())[1] z_mu, z_logvar = self.optimus.encoder.linear(z).chunk(2, -1) # z_sampled = self.optimus.reparameterize(z_mu, z_logvar, 1) return z_mu.squeeze(1) @torch.no_grad() def optimus_decode(self, z, temperature=1.0): bos_token = self.optimus.tokenizer_decoder.encode('') eos_token = self.optimus.tokenizer_decoder.encode('') context_tokens = torch.LongTensor(bos_token).to(z.device) from .optimus import sample_single_sequence_conditional sentenses = [] for zi in z: out = sample_single_sequence_conditional( model=self.optimus.decoder, context=context_tokens, past=zi, temperature=temperature, top_k=0, top_p=1.0, max_length=30, eos_token = eos_token[0],) text = self.optimus.tokenizer_decoder.decode(out.tolist(), clean_up_tokenization_spaces=True) text = text.split()[1:-1] text = ' '.join(text) sentenses.append(text) return sentenses @torch.no_grad() def clip_encode_text(self, text, encode_type='encode_text'): swap_type = self.clip.encode_type self.clip.encode_type = encode_type embedding = self.clip.encode(text) self.clip.encode_type = swap_type return embedding @torch.no_grad() def clip_encode_vision(self, vision, encode_type='encode_vision'): swap_type = self.clip.encode_type self.clip.encode_type = encode_type if isinstance(vision, torch.Tensor): vision = ((vision+1)/2).to('cpu').numpy() vision = np.transpose(vision, (0, 2, 3, 1)) vision = [vi for vi in vision] embedding = self.clip.encode(vision) self.clip.encode_type = swap_type return embedding def forward(self, x, c, noise=None, xtype='image', ctype='prompt'): t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=x.device).long() return self.p_losses(x, c, t, noise, xtype, ctype) def apply_model(self, x_noisy, t, cond, xtype='image', ctype='prompt'): return self.model.diffusion_model(x_noisy, t, cond, xtype, ctype) def get_image_loss(self, pred, target, mean=True): if self.loss_type == 'l1': loss = (target - pred).abs() if mean: loss = loss.mean() elif self.loss_type == 'l2': if mean: loss = torch.nn.functional.mse_loss(target, pred) else: loss = torch.nn.functional.mse_loss(target, pred, reduction='none') else: raise NotImplementedError("unknown loss type '{loss_type}'") return loss def get_text_loss(self, pred, target): if self.loss_type == 'l1': loss = (target - pred).abs() elif self.loss_type == 'l2': loss = torch.nn.functional.mse_loss(target, pred, reduction='none') return loss def p_losses(self, x_start, cond, t, noise=None, xtype='image', ctype='prompt'): noise = torch.randn_like(x_start) if noise is None else noise x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) model_output = self.apply_model(x_noisy, t, cond, xtype, ctype) loss_dict = {} if self.parameterization == "x0": target = x_start elif self.parameterization == "eps": target = noise else: raise NotImplementedError() if xtype == 'image': loss_simple = self.get_image_loss(model_output, target, mean=False).mean([1, 2, 3]) elif xtype == 'text': loss_simple = self.get_text_loss(model_output, target).mean([1]) logvar_t = self.logvar[t].to(self.device) if logvar_t.sum().item() != 0: assert False, "Default SD training has logvar fixed at 0" if self.learn_logvar: assert False, "Default SD training don't learn logvar" if self.l_simple_weight != 1: assert False, "Default SD training always set l_simple_weight==1" loss = loss_simple.mean() loss_dict['loss_simple'] = loss_simple.mean().item() loss_dict['Loss'] = loss.item() return loss, loss_dict def apply_model_ex(self, x_noisy, t, c_in, c_ex, xtype='image', c_in_type='image', c_ex_type='text', mixed_ratio=0.5): return self.model.diffusion_model.forward_ex(x_noisy, t, c_in, c_ex, xtype, c_in_type, c_ex_type, mixed_ratio) def apply_model_dc(self, x_noisy, t, first_c, second_c, xtype='image', first_ctype='vision', second_ctype='prompt', mixed_ratio=0.5): return self.model.diffusion_model.forward_dc(x_noisy, t, first_c, second_c, xtype, first_ctype, second_ctype, mixed_ratio)