import torch import torch.nn as nn import torch.nn.functional as F import numpy as np import math from .layers import get_timestep_embedding def extract(v, t, x_shape): out = torch.gather(v, index=t, dim=0).float() return out.view([t.shape[0]] + [1] * (len(x_shape) - 1)) class GaussianDiffusionTrainer(nn.Module): def __init__(self, model, beta_1, beta_T, T): super().__init__() self.model = model self.T = T betas = torch.linspace(beta_1, beta_T, T, dtype=torch.float64).double() alphas = 1. - betas self.register_buffer('betas', betas) alphas_bar = torch.cumprod(alphas, dim=0) self.register_buffer( 'sqrt_alphas_bar', torch.sqrt(alphas_bar)) self.register_buffer( 'sqrt_one_minus_alphas_bar', torch.sqrt(1. - alphas_bar)) self.register_buffer( 'sqrt_recip_alphas_bar', torch.sqrt(1. / alphas_bar)) self.register_buffer( 'sqrt_recipm1_alphas_bar', torch.sqrt(1. / alphas_bar - 1)) def make_x_t(self, x_0_con, t, noise): x_t_con = ( extract(self.sqrt_alphas_bar, t, x_0_con.shape) * x_0_con + extract(self.sqrt_one_minus_alphas_bar, t, x_0_con.shape) * noise) return x_t_con def predict_xstart_from_eps(self, x_t, t, eps): assert x_t.shape == eps.shape return ( extract(self.sqrt_recip_alphas_bar, t, x_t.shape) * x_t - extract(self.sqrt_recipm1_alphas_bar, t, x_t.shape) * eps ) class GaussianDiffusionSampler(nn.Module): def __init__(self, model, beta_1, beta_T, T, mean_type='eps', var_type='fixedlarge'): assert mean_type in ['xprev' 'xstart', 'epsilon'] assert var_type in ['fixedlarge', 'fixedsmall'] super().__init__() self.model = model self.T = T self.mean_type = mean_type self.var_type = var_type betas = torch.linspace(beta_1, beta_T, T, dtype=torch.float64).double() alphas = 1. - betas self.register_buffer( 'betas', betas) alphas_bar = torch.cumprod(alphas, dim=0) alphas_bar_prev = F.pad(alphas_bar, [1, 0], value=1)[:T] self.register_buffer( 'sqrt_alphas_bar', torch.sqrt(alphas_bar)) self.register_buffer( 'sqrt_one_minus_alphas_bar', torch.sqrt(1. - alphas_bar)) self.register_buffer( 'sqrt_recip_alphas_bar', torch.sqrt(1. / alphas_bar)) self.register_buffer( 'sqrt_recipm1_alphas_bar', torch.sqrt(1. / alphas_bar - 1)) self.register_buffer( 'posterior_var', self.betas * (1. - alphas_bar_prev) / (1. - alphas_bar)) # Avoid log(0) at the first diffusion step. self.register_buffer( 'posterior_log_var_clipped', torch.log( torch.cat([self.posterior_var[1:2], self.posterior_var[1:]]))) self.register_buffer( 'posterior_mean_coef1', torch.sqrt(alphas_bar_prev) * self.betas / (1. - alphas_bar)) self.register_buffer( 'posterior_mean_coef2', torch.sqrt(alphas) * (1. - alphas_bar_prev) / (1. - alphas_bar)) def q_mean_variance(self, x_0, x_t, t): """Mean and log variance of q(x_{t-1} | x_t, x_0).""" assert x_0.shape == x_t.shape posterior_mean = ( extract(self.posterior_mean_coef1, t, x_t.shape) * x_0 + extract(self.posterior_mean_coef2, t, x_t.shape) * x_t ) posterior_log_var_clipped = extract( self.posterior_log_var_clipped, t, x_t.shape) return posterior_mean, posterior_log_var_clipped def predict_xstart_from_eps(self, x_t, t, eps): assert x_t.shape == eps.shape return ( extract(self.sqrt_recip_alphas_bar, t, x_t.shape) * x_t - extract(self.sqrt_recipm1_alphas_bar, t, x_t.shape) * eps ) def p_mean_variance(self, x_t, t, cond, trans): model_log_var = { 'fixedlarge': torch.log(torch.cat([self.posterior_var[1:2], self.betas[1:]])), 'fixedsmall': self.posterior_log_var_clipped, }[self.var_type] model_log_var = extract(model_log_var, t, x_t.shape) if self.mean_type == 'epsilon': eps = self.model(x_t, t, cond) x_0 = self.predict_xstart_from_eps(x_t, t, eps=eps) model_mean, _ = self.q_mean_variance(x_0, x_t, t) else: raise NotImplementedError(self.mean_type) return model_mean, model_log_var class DiffusionProteinModel(nn.Module): def __init__(self, base_model, FLAGS): super().__init__() self.base_model = base_model self._freeze_base_model_layers() self.embed_dim = FLAGS.nf self.time_embed = nn.Sequential( nn.Linear(self.embed_dim, self.embed_dim * 4), nn.SiLU(), nn.Linear(self.embed_dim * 4, 1280) ) self.cond_proj = nn.Linear(FLAGS.cond_size, 1280) self.output_proj = nn.Sequential( nn.Linear(1280, FLAGS.output_size * 2), nn.SiLU(), nn.Linear(FLAGS.output_size * 2, FLAGS.output_size) ) def _freeze_base_model_layers(self): for param in self.base_model.parameters(): param.requires_grad = False for i in range(23, 33): for param in self.base_model.esm_model.encoder.layer[i].parameters(): param.requires_grad = True for component in [self.base_model.lm_head, self.base_model.attn_head, self.base_model.attn_ln, self.base_model.attn_skip, self.base_model.linear_layers, self.base_model.reduction_layers, self.base_model.clf_ln]: for param in component.parameters(): param.requires_grad = True for param in self.base_model.esm_model.embeddings.parameters(): param.requires_grad = False def decode_embeddings(self, embeddings): mlm_logits = self.base_model.lm_head(embeddings) return mlm_logits def get_embeddings(self, x, detach=False): _, _, x_combined, _ = self.base_model( input_ids=x['input_ids'], attention_masks=x['attention_mask'] ) if detach: return x_combined.detach() return x_combined def forward(self, x, time_cond, cond): if isinstance(x, dict): _, _, x_combined, _ = self.base_model( input_ids=x['input_ids'], attention_masks=x['attention_mask'] ) x = x_combined seq_length = x.size(1) temb = get_timestep_embedding(time_cond, self.embed_dim) temb = self.time_embed(temb) temb = temb.unsqueeze(1).expand(-1, seq_length, -1) if len(cond.shape) == 3: cond = self.cond_proj(cond) else: cond = self.cond_proj(cond) cond = cond.unsqueeze(1).expand(-1, seq_length, -1) x = x + temb + cond return x def get_token_embeddings(self): return self.base_model.esm_model.embeddings.word_embeddings.weight