meng2003/music2dance
0
1from math import log2, sqrt2import torch3from torch import nn, einsum4import torch.nn.functional as F5 6from models.transformer import BasicTransformerModel, EncDecTransformerModel, EncDecXTransformer7 8from axial_positional_embedding import AxialPositionalEmbedding9from einops import rearrange10 11# from dalle_pytorch import distributed_utils12# from dalle_pytorch.vae import OpenAIDiscreteVAE13# from dalle_pytorch.vae import VQGanVAE102414# from dalle_pytorch.transformer import Transformer15 16# helpers17 18def exists(val):19 return val is not None20 21def default(val, d):22 return val if exists(val) else d23 24def always(val):25 def inner(*args, **kwargs):26 return val27 return inner28 29def is_empty(t):30 return t.nelement() == 031 32def masked_mean(t, mask, dim = 1):33 t = t.masked_fill(~mask[:, :, None], 0.)34 return t.sum(dim = 1) / mask.sum(dim = 1)[..., None]35 36def eval_decorator(fn):37 def inner(model, *args, **kwargs):38 was_training = model.training39 model.eval()40 out = fn(model, *args, **kwargs)41 model.train(was_training)42 return out43 return inner44 45# sampling helpers46 47def top_k(logits, thres = 0.5):48 num_logits = logits.shape[-1]49 k = max(int((1 - thres) * num_logits), 1)50 val, ind = torch.topk(logits, k)51 probs = torch.full_like(logits, float('-inf'))52 probs.scatter_(1, ind, val)53 return probs54 55# discrete vae class56 57class ResBlock(nn.Module):58 def __init__(self, chan):59 super().__init__()60 self.net = nn.Sequential(61 nn.Conv2d(chan, chan, 3, padding = 1),62 nn.ReLU(),63 nn.Conv2d(chan, chan, 3, padding = 1),64 nn.ReLU(),65 nn.Conv2d(chan, chan, 1)66 )67 68 def forward(self, x):69 return self.net(x) + x70 71class ConditionalDiscreteVAEVision(nn.Module):72 def __init__(73 self,74 image_shape = (256,256),75 num_tokens = 512,76 codebook_dim = 512,77 num_layers = 3,78 num_resnet_blocks = 0,79 hidden_dim = 64,80 conditioning_dim = 64,81 channels = 3,82 smooth_l1_loss = False,83 temperature = 0.9,84 straight_through = False,85 kl_div_loss_weight = 0.,86 normalization = ((0.5,) * 3, (0.5,) * 3)87 ):88 super().__init__()89 assert log2(image_shape[0]).is_integer(), 'image size must be a power of 2'90 assert log2(image_shape[1]).is_integer(), 'image size must be a power of 2'91 assert num_layers >= 1, 'number of layers must be greater than or equal to 1'92 has_resblocks = num_resnet_blocks > 093 94 self.image_shape = image_shape95 self.num_tokens = num_tokens96 self.num_layers = num_layers97 self.temperature = temperature98 self.straight_through = straight_through99 self.codebook = nn.Embedding(num_tokens, codebook_dim)100 101 hdim = hidden_dim102 103 enc_chans = [hidden_dim] * num_layers104 dec_chans = list(reversed(enc_chans))105 106 enc_chans = [channels, *enc_chans]107 108 if not has_resblocks:109 dec_init_chan = codebook_dim110 else:111 dec_init_chan = dec_chans[0]112 dec_chans = [dec_init_chan, *dec_chans]113 114 enc_chans_io, dec_chans_io = map(lambda t: list(zip(t[:-1], t[1:])), (enc_chans, dec_chans))115 116 enc_layers = []117 dec_layers = []118 119 for (enc_in, enc_out), (dec_in, dec_out) in zip(enc_chans_io, dec_chans_io):120 enc_layers.append(nn.Sequential(nn.Conv2d(enc_in, enc_out, 4, stride = 2, padding = 1), nn.ReLU()))121 dec_layers.append(nn.Sequential(nn.ConvTranspose2d(dec_in, dec_out, 4, stride = 2, padding = 1), nn.ReLU()))122 123 for _ in range(num_resnet_blocks):124 dec_layers.insert(0, ResBlock(dec_chans[1]))125 enc_layers.append(ResBlock(enc_chans[-1]))126 127 if num_resnet_blocks > 0:128 dec_layers.insert(0, nn.Conv2d(codebook_dim, dec_chans[1], 1))129 130 enc_layers.append(nn.Conv2d(enc_chans[-1], num_tokens, 1))131 dec_layers.append(nn.Conv2d(dec_chans[-1], channels, 1))132 133 self.encoder = nn.Sequential(*enc_layers)134 self.decoder = nn.Sequential(*dec_layers)135 136 self.loss_fn = F.smooth_l1_loss if smooth_l1_loss else F.mse_loss137 self.kl_div_loss_weight = kl_div_loss_weight138 139 # take care of normalization within class140 self.normalization = normalization141 142 # self._register_external_parameters()143 144 # def _register_external_parameters(self):145 # """Register external parameters for DeepSpeed partitioning."""146 # if (147 # not distributed_utils.is_distributed148 # or not distributed_utils.using_backend(149 # distributed_utils.DeepSpeedBackend)150 # ):151 # return152 #153 # deepspeed = distributed_utils.backend.backend_module154 # deepspeed.zero.register_external_parameters(self, self.codebook.weight)155 156 def norm(self, images):157 if not exists(self.normalization):158 return images159 160 means, stds = map(lambda t: torch.as_tensor(t).to(images), self.normalization)161 means, stds = map(lambda t: rearrange(t, 'c -> () c () ()'), (means, stds))162 images = images.clone()163 images.sub_(means).div_(stds)164 return images165 166 @torch.no_grad()167 @eval_decorator168 def get_codebook_indices(self, images):169 logits = self(images, return_logits = True)170 codebook_indices = logits.argmax(dim = 1).flatten(1)171 return codebook_indices172 173 def decode(174 self,175 img_seq176 ):177 image_embeds = self.codebook(img_seq)178 b, n, d = image_embeds.shape179 h = w = int(sqrt(n))180 181 image_embeds = rearrange(image_embeds, 'b (h w) d -> b d h w', h = h, w = w)182 images = self.decoder(image_embeds)183 return images184 185 def forward(186 self,187 img,188 return_loss = False,189 return_recons = False,190 return_logits = False,191 temp = None192 ):193 device, num_tokens, image_shape, kl_div_loss_weight = img.device, self.num_tokens, self.image_shape, self.kl_div_loss_weight194 assert img.shape[-1] == image_shape[1] and img.shape[-2] == image_shape[0], f'input must have the correct image size {image_shape[0]}x{image_shape[1]}'195 196 img = self.norm(img)197 198 logits = self.encoder(img)199 200 if return_logits:201 return logits # return logits for getting hard image indices for DALL-E training202 203 temp = default(temp, self.temperature)204 soft_one_hot = F.gumbel_softmax(logits, tau = temp, dim = 1, hard = self.straight_through)205 sampled = einsum('b n h w, n d -> b d h w', soft_one_hot, self.codebook.weight)206 out = self.decoder(sampled)207 208 if not return_loss:209 return out210 211 # reconstruction loss212 213 recon_loss = self.loss_fn(img, out)214 215 # kl divergence216 217 logits = rearrange(logits, 'b n h w -> b (h w) n')218 log_qy = F.log_softmax(logits, dim = -1)219 log_uniform = torch.log(torch.tensor([1. / num_tokens], device = device))220 kl_div = F.kl_div(log_uniform, log_qy, None, None, 'batchmean', log_target = True)221 222 loss = recon_loss + (kl_div * kl_div_loss_weight)223 224 if not return_recons:225 return loss226 227 return loss, out228 229class ConditionalDiscreteVAE(nn.Module):230 def __init__(231 self,232 input_shape = (256,256),233 num_tokens = 512,234 codebook_dim = 512,235 num_layers = 3,236 num_resnet_blocks = 0,237 hidden_dim = 64,238 cond_dim = 0,239 channels = 3,240 smooth_l1_loss = False,241 temperature = 0.9,242 straight_through = False,243 kl_div_loss_weight = 0.,244 normalization = None,245 prior_nhead = 8,246 prior_dhid = 512,247 prior_nlayers = 8,248 prior_dropout = 0,249 prior_use_pos_emb = True,250 prior_use_x_transformers = False,251 opt = None,252 cond_vae = False253 ):254 super().__init__()255 assert num_layers >= 1, 'number of layers must be greater than or equal to 1'256 has_resblocks = num_resnet_blocks > 0257 258 self.input_shape = input_shape259 self.num_tokens = num_tokens260 self.num_layers = num_layers261 self.temperature = temperature262 self.straight_through = straight_through263 self.codebook = nn.Embedding(num_tokens, codebook_dim)264 self.cond_dim = cond_dim265 self.cond_vae = cond_vae266 267 hdim = hidden_dim268 269 enc_chans = [hidden_dim] * num_layers270 dec_chans = list(reversed(enc_chans))271 272 if cond_vae:273 enc_chans = [channels + cond_dim, *enc_chans]274 else:275 enc_chans = [channels, *enc_chans]276 277 if not has_resblocks:278 if cond_vae:279 dec_init_chan = codebook_dim + cond_dim280 else:281 dec_init_chan = codebook_dim282 else:283 dec_init_chan = dec_chans[0]284 dec_chans = [dec_init_chan, *dec_chans]285 286 enc_chans_io, dec_chans_io = map(lambda t: list(zip(t[:-1], t[1:])), (enc_chans, dec_chans))287 288 enc_layers = []289 dec_layers = []290 291 292 if input_shape[0] == 1:293 kernel_size1 = 1294 padding_size1 = 0295 codebook_layer_shape1 = 1296 elif input_shape[0] in [2,3,4]:297 kernel_size1 = 3298 padding_size1 = 1299 codebook_layer_shape1 = input_shape[0]300 else:301 #kernel_size1 = 4302 kernel_size1 = 3303 padding_size1 = 1304 #codebook_layer_shape1 = input_shape[0] - num_layers305 codebook_layer_shape1 = input_shape[0]306 307 if input_shape[1] == 1:308 kernel_size2 = 1309 padding_size2 = 0310 codebook_layer_shape2 = 1311 elif input_shape[1] in [2,3,4]:312 kernel_size2 = 3313 padding_size2 = 1314 codebook_layer_shape2 = input_shape[1]315 else:316 #kernel_size2 = 4317 kernel_size2 = 3318 padding_size2 = 1319 #codebook_layer_shape2 = input_shape[1] - num_layers320 codebook_layer_shape2 = input_shape[1]321 322 self.codebook_layer_shape = (codebook_layer_shape1,codebook_layer_shape2)323 kernel_shape = (kernel_size1, kernel_size2)324 padding_shape = (padding_size1, padding_size2)325 for (enc_in, enc_out), (dec_in, dec_out) in zip(enc_chans_io, dec_chans_io):326 enc_layers.append(nn.Sequential(nn.Conv2d(enc_in, enc_out, kernel_shape, stride = 1, padding = padding_shape), nn.ReLU()))327 dec_layers.append(nn.Sequential(nn.ConvTranspose2d(dec_in, dec_out, kernel_shape, stride = 1, padding = padding_shape), nn.ReLU()))328 329 for _ in range(num_resnet_blocks):330 dec_layers.insert(0, ResBlock(dec_chans[1]))331 enc_layers.append(ResBlock(enc_chans[-1]))332 333 if num_resnet_blocks > 0:334 if cond_vae:335 dec_layers.insert(0, nn.Conv2d(codebook_dim + cond_dim, dec_chans[1], 1))336 else:337 dec_layers.insert(0, nn.Conv2d(codebook_dim, dec_chans[1], 1))338 339 enc_layers.append(nn.Conv2d(enc_chans[-1], num_tokens, 1))340 dec_layers.append(nn.Conv2d(dec_chans[-1], channels, 1))341 342 self.cond_upsampler = torch.nn.Upsample(size=input_shape) #upsampler to feed the conditioning to the input of the encoder343 self.encoder = nn.Sequential(*enc_layers)344 self.decoder = nn.Sequential(*dec_layers)345 346 self.loss_fn = F.smooth_l1_loss if smooth_l1_loss else F.mse_loss347 self.kl_div_loss_weight = kl_div_loss_weight348 349 # take care of normalization within class350 self.normalization = normalization351 352 latent_size = codebook_layer_shape1*codebook_layer_shape2353 self.latent_size = latent_size354 if cond_dim > 0:355 self.prior_transformer = ContDiscTransformer(cond_dim, num_tokens, codebook_dim, prior_nhead, prior_dhid, prior_nlayers, prior_dropout,356 use_pos_emb=prior_use_pos_emb,357 src_length=latent_size,358 tgt_length=latent_size,359 use_x_transformers=prior_use_x_transformers,360 opt=opt)361 362 # self._register_external_parameters()363 364 # def _register_external_parameters(self):365 # """Register external parameters for DeepSpeed partitioning."""366 # if (367 # not distributed_utils.is_distributed368 # or not distributed_utils.using_backend(369 # distributed_utils.DeepSpeedBackend)370 # ):371 # return372 #373 # deepspeed = distributed_utils.backend.backend_module374 # deepspeed.zero.register_external_parameters(self, self.codebook.weight)375 376 def norm(self, images):377 if not exists(self.normalization):378 return images379 380 means, stds = map(lambda t: torch.as_tensor(t).to(images), self.normalization)381 means, stds = map(lambda t: rearrange(t, 'c -> () c () ()'), (means, stds))382 images = images.clone()383 images.sub_(means).div_(stds)384 return images385 386 @torch.no_grad()387 @eval_decorator388 def get_codebook_indices(self, inputs, cond=None):389 logits = self(inputs, cond, return_logits = True)390 codebook_indices = logits.argmax(dim = 1).flatten(1)391 return codebook_indices392 393 def decode(394 self,395 img_seq,396 cond = None397 ):398 image_embeds = self.codebook(img_seq)399 b, n, d = image_embeds.shape400 h = w = int(sqrt(n))401 402 image_embeds = rearrange(image_embeds, 'b (h w) d -> b d h w', h = h, w = w)403 if cond is not None:404 image_embeds_cond = torch.cat([image_embeds, cond], dim = 1)405 images = self.decoder(image_embeds_cond)406 else:407 images = self.decoder(image_embeds)408 409 return images410 411 def prior_logp(412 self,413 inputs,414 cond = None,415 return_accuracy = False,416 detach_cond = False417 ):418 # import pdb;pdb.set_trace()419 #if cond is None: raise NotImplementedError("Haven't implemented non-conditional DVAEs")420 if len(inputs.shape) == 3:421 inputs = inputs.reshape(inputs.shape[0], inputs.shape[1],*self.input_shape)422 if len(cond.shape) == 3:423 cond = cond.reshape(cond.shape[0], cond.shape[1],*self.codebook_layer_shape)424 with torch.no_grad():425 if self.cond_vae:426 labels = self.get_codebook_indices(inputs, cond)427 else:428 labels = self.get_codebook_indices(inputs)429 if detach_cond:430 cond = cond.detach()431 logits = self.prior_transformer(cond.squeeze(-1).permute(2,0,1), labels.permute(1,0)).permute(1,2,0)432 loss = F.cross_entropy(logits, labels)433 if not return_accuracy:434 return loss435 # import pdb;pdb.set_trace()436 predicted = logits.argmax(dim = 1).flatten(1)437 accuracy = (predicted == labels).sum()/predicted.nelement()438 return loss, accuracy439 440 def generate(self, cond, temp=1.0, filter_thres = 0.5):441 #if cond is None: raise NotImplementedError("Haven't implemented non-conditional DVAEs")442 if len(cond.shape) == 3:443 cond = cond.reshape(cond.shape[0], cond.shape[1],*self.codebook_layer_shape)444 dummy = torch.zeros(1,1).long().to(cond.device)445 tokens = []446 for i in range(self.latent_size):447 # print(i)448 logits = self.prior_transformer(cond.squeeze(-1).permute(2,0,1), torch.cat(tokens+[dummy], 0)).permute(1,2,0)[:,-1,:]449 filtered_logits = top_k(logits, thres = filter_thres)450 probs = F.softmax(filtered_logits / temp, dim = -1)451 sampled = torch.multinomial(probs, 1)452 tokens.append(sampled)453 print(tokens)454 embs = self.codebook(torch.cat(tokens, 0))455 # import pdb;pdb.set_trace()456 if self.cond_vae:457 sampled_cond = torch.cat([embs.permute(2,0,1).unsqueeze(0),cond], dim=1)458 else:459 sampled_cond = embs.permute(2,0,1).unsqueeze(0)460 out = self.decoder(sampled_cond)461 return out462 463 def forward(464 self,465 inp,466 cond = None,467 return_loss = False,468 return_recons = False,469 return_logits = False,470 temp = None471 ):472 if len(inp.shape) == 3:473 inp = inp.reshape(inp.shape[0], inp.shape[1],*self.input_shape)474 device, num_tokens, input_shape, kl_div_loss_weight = inp.device, self.num_tokens, self.input_shape, self.kl_div_loss_weight475 assert inp.shape[-1] == input_shape[1] and inp.shape[-2] == input_shape[0], f'input must have the correct image size {input_shape[0]}x{input_shape[1]}. Instead got {inp.shape[0]}x{inp.shape[1]}'476 477 inp = self.norm(inp)478 if cond is not None:479 if len(cond.shape) == 3:480 cond = cond.reshape(cond.shape[0], cond.shape[1],*self.codebook_layer_shape)481 cond_upsampled = self.cond_upsampler(cond)482 inp_cond = torch.cat([inp,cond_upsampled], dim=1)483 inp_cond = self.norm(inp_cond)484 else:485 inp_cond = self.norm(inp)486 487 logits = self.encoder(inp_cond)488 # codebook_indices = logits.argmax(dim = 1).flatten(1)489 # print(codebook_indices.shape)490 # print(codebook_indices)491 # print(list(self.encoder.parameters())[1].data)492 # for p in self.prior_transformer.parameters():493 # print(p.norm())494 495 if return_logits:496 return logits # return logits for getting hard image indices for DALL-E training497 498 temp = default(temp, self.temperature)499 soft_one_hot = F.gumbel_softmax(logits, tau = temp, dim = 1, hard = self.straight_through)500 sampled = einsum('b n h w, n d -> b d h w', soft_one_hot, self.codebook.weight)501 if cond is not None:502 sampled_cond = torch.cat([sampled,cond], dim=1)503 out = self.decoder(sampled_cond)504 else:505 out = self.decoder(sampled)506 507 if not return_loss:508 return out509 510 # reconstruction loss511 512 # import pdb;pdb.set_trace()513 recon_loss = self.loss_fn(inp, out)514 515 # kl divergence516 517 logits = rearrange(logits, 'b n h w -> b (h w) n')518 log_qy = F.log_softmax(logits, dim = -1)519 log_uniform = torch.log(torch.tensor([1. / num_tokens], device = device))520 kl_div = F.kl_div(log_uniform, log_qy, None, None, 'batchmean', log_target = True)521 522 loss = recon_loss + (kl_div * kl_div_loss_weight)523 524 if not return_recons:525 return loss526 527 return loss, out528 529class ContDiscTransformer(nn.Module):530 531 def __init__(self, src_d, tgt_num_tokens, tgt_emb_dim, nhead, dhid, nlayers, dropout=0.5,use_pos_emb=False,src_length=0,tgt_length=0,use_x_transformers=False,opt=None):532 super(ContDiscTransformer, self).__init__()533 self.transformer = EncDecTransformerModel(tgt_num_tokens, src_d, tgt_emb_dim, nhead, dhid, nlayers, dropout=dropout,use_pos_emb=use_pos_emb,src_length=src_length,tgt_length=tgt_length,use_x_transformers=use_x_transformers,opt=opt)534 #self.transformer = EncDecTransformerModel(tgt_num_tokens, src_d, tgt_emb_dim, nhead, dhid, nlayers, dropout=dropout,use_pos_emb=False,src_length=src_length,tgt_length=tgt_length,use_x_transformers=use_x_transformers,opt=opt)535 # self.transformer = EncDecXTransformer(dim=dhid, dec_dim_out=tgt_num_tokens, enc_dim_in=src_d, enc_dim_out=tgt_emb_dim, dec_din_in=tgt_emb_dim, enc_heads=nhead, dec_heads=nhead, enc_depth=nlayers, dec_depth=nlayers, enc_dropout=dropout, dec_dropout=dropout, enc_max_seq_len=1024, dec_max_seq_len=1024)536 self.embedding = nn.Embedding(tgt_num_tokens, tgt_emb_dim)537 self.first_input = nn.Parameter((torch.randn(1,1,tgt_emb_dim)))538 539 def forward(self, src, tgt):540 tgt = tgt[:-1]541 embs = self.embedding(tgt)542 embs = torch.cat([torch.tile(self.first_input, (1,embs.shape[1],1)), embs], 0)543 output = self.transformer(src,embs)544 return output545