fix code_emb

This commit is contained in:
James Betker 2022-06-19 17:54:08 -06:00
parent 368dca18b1
commit 8c8efbe131

View File

@ -210,8 +210,8 @@ class TransformerDiffusion(nn.Module):
def timestep_independent(self, prior, expected_seq_len): def timestep_independent(self, prior, expected_seq_len):
if self.new_code_expansion: if self.new_code_expansion:
code_emb = F.interpolate(prior.permute(0,2,1), size=expected_seq_len, mode='linear').permute(0,2,1) prior = F.interpolate(prior.permute(0,2,1), size=expected_seq_len, mode='linear').permute(0,2,1)
code_emb = self.ar_input(code_emb) if self.ar_prior else self.input_converter(code_emb) code_emb = self.ar_input(prior) if self.ar_prior else self.input_converter(prior)
code_emb = self.ar_prior_intg(code_emb) if self.ar_prior else self.code_converter(code_emb) code_emb = self.ar_prior_intg(code_emb) if self.ar_prior else self.code_converter(code_emb)
# Mask out the conditioning branch for whole batch elements, implementing something similar to classifier-free guidance. # Mask out the conditioning branch for whole batch elements, implementing something similar to classifier-free guidance.
@ -732,14 +732,14 @@ def test_cheater_model():
# For music: # For music:
model = TransformerDiffusionWithCheaterLatent(in_channels=256, out_channels=512, model = TransformerDiffusionWithCheaterLatent(in_channels=256, out_channels=512,
model_channels=1024, contraction_dim=512, model_channels=1536, contraction_dim=768,
prenet_channels=1024, num_heads=8, prenet_channels=1024, num_heads=12,
input_vec_dim=256, num_layers=12, prenet_layers=6, input_vec_dim=256, num_layers=20, prenet_layers=6,
dropout=.1, new_code_expansion=True, dropout=.1, new_code_expansion=True,
) )
diff_weights = torch.load('extracted_diff.pth') #diff_weights = torch.load('extracted_diff.pth')
model.diff.load_state_dict(diff_weights, strict=False) #model.diff.load_state_dict(diff_weights, strict=False)
cheater_ar_weights = torch.load('X:\\dlas\\experiments\\train_music_gpt_cheater\\models\\19500_generator_ema.pth') cheater_ar_weights = torch.load('X:\\dlas\\experiments\\train_music_gpt_cheater\\models\\60000_generator_ema.pth')
cheater_ar = GptMusicLower(dim=1024, encoder_out_dim=256, layers=16, fp16=False, num_target_vectors=8192, num_vaes=4, cheater_ar = GptMusicLower(dim=1024, encoder_out_dim=256, layers=16, fp16=False, num_target_vectors=8192, num_vaes=4,
vqargs= {'positional_dims': 1, 'channels': 64, vqargs= {'positional_dims': 1, 'channels': 64,
'hidden_dim': 512, 'num_resnet_blocks': 3, 'codebook_dim': 512, 'num_tokens': 8192, 'hidden_dim': 512, 'num_resnet_blocks': 3, 'codebook_dim': 512, 'num_tokens': 8192,