mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Allow bi-directional clipping

Browse Source
This commit is contained in:
James Betker 2022-01-08 22:18:04 -07:00
parent 894d245062
commit 15d9517e26
1 changed files with 16 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
codes/models/gpt_voice/text_voice_clip.py
View File

@ -40,8 +40,9 @@ class VoiceCLIP(nn.Module):
speech_enc_depth=6, speech_enc_depth=6,
speech_heads=8, speech_heads=8,
speech_seq_len=250, speech_seq_len=250,
text_mask_percentage: 0, text_mask_percentage=0,
wav_token_compression = 1024, voice_mask_percentage=0,
wav_token_compression=1024,
): ):
super().__init__() super().__init__()
self.text_emb = nn.Embedding(num_text_tokens, dim_text) self.text_emb = nn.Embedding(num_text_tokens, dim_text)
@ -58,6 +59,7 @@ class VoiceCLIP(nn.Module):
self.temperature = nn.Parameter(torch.tensor(1.)) self.temperature = nn.Parameter(torch.tensor(1.))
self.text_mask_percentage = text_mask_percentage self.text_mask_percentage = text_mask_percentage
self.voice_mask_percentage = voice_mask_percentage
self.wav_token_compression = wav_token_compression self.wav_token_compression = wav_token_compression
def forward( def forward(
@ -76,7 +78,12 @@ class VoiceCLIP(nn.Module):
speech_tokens = speech_tokens[:, :max_mel_len] speech_tokens = speech_tokens[:, :max_mel_len]
b, device = text.shape[0], text.device b, device = text.shape[0], text.device
text_mask = torch.rand_like(text.float()) > self.text_mask_percentage if self.training:
text_mask = torch.rand_like(text.float()) > self.text_mask_percentage
voice_mask = torch.rand_like(speech_tokens.float()) > self.voice_mask_percentage
else:
text_mask = torch.ones_like(text.float()).bool()
voice_mask = torch.ones_like(speech_tokens.float()).bool()
text_emb = self.text_emb(text) text_emb = self.text_emb(text)
text_emb += self.text_pos_emb(torch.arange(text.shape[1], device=device)) text_emb += self.text_pos_emb(torch.arange(text.shape[1], device=device))
@ -85,14 +92,10 @@ class VoiceCLIP(nn.Module):
speech_emb += self.speech_pos_emb(torch.arange(speech_emb.shape[1], device=device)) speech_emb += self.speech_pos_emb(torch.arange(speech_emb.shape[1], device=device))
enc_text = self.text_transformer(text_emb, mask=text_mask) enc_text = self.text_transformer(text_emb, mask=text_mask)
enc_speech = self.speech_transformer(speech_emb) enc_speech = self.speech_transformer(speech_emb, mask=voice_mask)
if self.text_mask_percentage > 0: text_latents = masked_mean(enc_text, text_mask, dim=1)
text_latents = masked_mean(enc_text, text_mask, dim=1) speech_latents = masked_mean(enc_speech, voice_mask, dim=1)
else:
text_latents = enc_text.mean(dim=1)
speech_latents = enc_speech.mean(dim=1)
text_latents = self.to_text_latent(text_latents) text_latents = self.to_text_latent(text_latents)
speech_latents = self.to_speech_latent(speech_latents) speech_latents = self.to_speech_latent(speech_latents)
@ -117,7 +120,9 @@ def register_voice_clip(opt_net, opt):
if __name__ == '__main__': if __name__ == '__main__':
clip = VoiceCLIP(text_mask_percentage=.2) clip = VoiceCLIP(text_mask_percentage=.2, voice_mask_percentage=.2)
clip(torch.randint(0,256,(2,120)), clip(torch.randint(0,256,(2,120)),
torch.tensor([50,100]),
torch.randint(0,8192,(2,250)), torch.randint(0,8192,(2,250)),
torch.tensor([101,102]),
return_loss=True) return_loss=True)