| import torch | |
| from einops import rearrange | |
| from vector_quantize_pytorch.residual_fsq import GroupedResidualFSQ | |
| def quantize( | |
| quantizer: GroupedResidualFSQ, | |
| audio_latents: torch.Tensor, # (batch_size, audio_len, audio_dim=1024) | |
| ) -> tuple[torch.Tensor, torch.Tensor]: | |
| # feat shape (batch_size, audio_len, audio_dim) | |
| # ind shape (GFSQ.G, batch_size, audio_len, GFSQ.R) | |
| # num_vq=GFSQ.G*GFSQ.R | |
| feat, ind = quantizer(audio_latents) | |
| audio_quantized_latents = feat # (batch_size, audio_len, audio_dim) | |
| audio_input_ids = rearrange( # (batch_size, audio_len, num_vq) | |
| ind, | |
| "g b t r ->b t (g r)", | |
| ) | |
| return audio_quantized_latents, audio_input_ids | |