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

Fix lr quantizer decode

Browse Source
This commit is contained in:
James Betker 2021-11-24 00:01:26 -07:00
parent 3f6ecfe0db
commit 7a3c4a4fc6
1 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
codes/models/gpt_voice/lucidrains_dvae.py
View File

@ -195,7 +195,10 @@ class DiscreteVAE(nn.Module):
img_seq
):
self.log_codes(img_seq)
image_embeds = self.codebook.embed_code(img_seq)
if hasattr(self.codebook, 'embed_code'):
image_embeds = self.codebook.embed_code(img_seq)
else:
image_embeds = F.embedding(img_seq, self.codebook.embed.transpose(1,0))
b, n, d = image_embeds.shape
kwargs = {}
@ -273,5 +276,6 @@ if __name__ == '__main__':
use_lr_quantizer=True)
#v.load_state_dict(torch.load('../experiments/clips_dvae_8192_rev2.pth'))
#v.eval()
o=v(torch.randn(1,80,256))
print(o[-1].shape)
r,l,o=v(torch.randn(1,80,256))
v.decode(torch.randint(0,8192,(1,256)))
print(o.shape, l.shape)