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