File size: 660 Bytes
1cc592d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
import torch, torch.nn as nn
from vector_quantize_pytorch import ResidualVQ
class RVQWrapper(nn.Module):
def __init__(self, dim, num_quantizers, codebook_size):
super().__init__()
self.ln_in = nn.LayerNorm(dim)
self.proj_in = nn.Linear(dim, dim)
self.rvq = ResidualVQ(dim=dim, num_quantizers=num_quantizers, codebook_size=codebook_size)
self.ln_out = nn.LayerNorm(dim)
self.proj_out = nn.Linear(dim, dim)
def forward(self, x):
x = self.proj_in(self.ln_in(x))
q, indices, commit_loss = self.rvq(x)
y = self.proj_out(self.ln_out(q))
return y, indices, commit_loss
|