Question about Mimi VAE codec release + encode/decode usage

[orech]

Hi team,

Thank you for releasing the model, it’s been really fun to experiment with.

I’m currently playing with it for speech generation and noticed it uses a Mimi codec with a VAE bottleneck (rather than the RVQ version). Do you have plans to release the Mimi VAE codec (standalone codec checkpoint) as well?

In the meantime, I tried extracting the Mimi module from Pocket-TTS to test a simple audio reconstruction path (encode - decode), but I’m clearly missing something: the reconstructed audio sounds heavily degraded / essentially unusable. Since I haven’t worked with VAE-style codecs before, I suspect I’m calling the interface incorrectly.

Here’s the minimal snippet I’m trying:

import torch
from pocket_tts.modules.stateful_module import init_states

audio_input = audio.unsqueeze(0).unsqueeze(0).float()  # [B, C, T]
mimi = tts_model.mimi

with torch.no_grad():
    mimi_state = init_states(mimi, batch_size=1, sequence_length=1000)

    latents = mimi.encode_to_latent(audio_input)
    reconstructed = mimi.decode_from_latent(latents, mimi_state)

reconstructed_audio = reconstructed.squeeze(0).squeeze(0)

Could you please point me to the correct way to do reconstruction with the Mimi VAE module (or share a reference snippet)?

[vvolhejn]

Hi! I'll get back to you regarding releasing the Mimi VAE codec separately, but my first response would be that this is useful to a limited amount of people so it's not likely to be an immediate priority.

Regarding your code snippet, I'll check later, but what I guess is missing is running the "dummy quantizer" (see here). It doesn't actually quantize - as you noted - but it does apply a linear projection, which is probably what you're missing. I'm not 100% sure if we release those weights as well, it might be the same issue as #30 (comment).

I'll check and we can provide that weight matrix later. It's just one projection so I'm sure it's also possible to recover it somehow, but it'd be a bit cumbersome.

[vvolhejn]

@orech
https://drive.google.com/file/d/1WP7RiUh5Niiz0na12vCSRuM1E7YGZYfQ/view?usp=drive_link Here is the input projection weight, hope this helps!

[orech]

Thanks a lot for sharing the projection weights! This is really valuable 👍 I’m learning continuous latent prediction for TTS right now, so it helps a ton

[hujingbin1]

@vvolhejn Hello, this is a very meaningful work! I have some questions about the VAE part. I would like to ask how to handle the resampling for audio reconstruction in this part. It seems that you only provided one project for the weights. How should the mean and variance prediction be handled?

[manukyutai]

Hello! Manu here who trained the VAE. We did train a variance projection but we don't use at inference time for simplicity - we could provide it if you want.

[hujingbin1]

Hello! Manu here who trained the VAE. We did train a variance projection but we don't use at inference time for simplicity - we could provide it if you want.

That's great! Could you please provide the corresponding projection weights? This is important for my study of continuous representation speech synthesis systems. Thank you so much!

[gabrieldemarmiesse]

Here is the link to the original weights as they appear in the training codebase: https://drive.google.com/drive/folders/1KIHnBjgWI_Q0zsYG1Dc_OMO_3Ioh9D0x?usp=sharing

I recommend analyzing them with a tool such as this https://github.com/EricLBuehler/safetensors_explorer

You should be able to find all the weights you're looking for there :)

[manukyutai]

The code to load those weights should look like this:

import torch
from torch import nn


class Quantizer(nn.Module):
    """Quantizer with input and logvar projs."""

    def __init__(self, dimension: int, output_dimension: int):
        super().__init__()
        self.dimension = dimension
        self.output_dimension = output_dimension
        self.output_proj = torch.nn.Conv1d(self.dimension, self.output_dimension, 1, bias=False)
        self.input_proj = torch.nn.Conv1d(self.output_dimension, self.dimension, 1, bias=False)
        self.logvar_proj = torch.nn.Conv1d(self.output_dimension, self.dimension, 1, bias=False)

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        return self.input_proj(x) + torch.randn_like(x) * torch.exp(.5 * self.logvar_proj(x))
    
    def decode(self, x: torch.Tensor) -> torch.Tensor:
        return self.output_proj(x)

[hujingbin1]

以下是训练代码库中原始权重文件的链接：https://drive.google.com/drive/folders/1KIHnBjgWI_Q0zsYG1Dc_OMO_3Ioh9D0x ?usp=sharing

我建议使用类似这样的工具来分析它们：https://github.com/EricLBuehler/safetensors_explorer

你应该能在那里找到你想要的所有重量数据 :)

Thank you very much!

[hujingbin1]

The code to load those weights should look like this:

import torch
from torch import nn

class Quantizer(nn.Module):
"""Quantizer with input and logvar projs."""

def __init__(self, dimension: int, output_dimension: int):
    super().__init__()
    self.dimension = dimension
    self.output_dimension = output_dimension
    self.output_proj = torch.nn.Conv1d(self.dimension, self.output_dimension, 1, bias=False)
    self.input_proj = torch.nn.Conv1d(self.output_dimension, self.dimension, 1, bias=False)
    self.logvar_proj = torch.nn.Conv1d(self.output_dimension, self.dimension, 1, bias=False)

def encode(self, x: torch.Tensor) -> torch.Tensor:
    return self.input_proj(x) + torch.randn_like(x) * torch.exp(.5 * self.logvar_proj(x))

def decode(self, x: torch.Tensor) -> torch.Tensor:
    return self.output_proj(x)

Thanks!

[cfrancesco]

Hi, thanks so much for providing the weights and the code!

The code to load those weights should look like this:

import torch
from torch import nn

class Quantizer(nn.Module):
"""Quantizer with input and logvar projs."""

def __init__(self, dimension: int, output_dimension: int):
    super().__init__()
    self.dimension = dimension
    self.output_dimension = output_dimension
    self.output_proj = torch.nn.Conv1d(self.dimension, self.output_dimension, 1, bias=False)
    self.input_proj = torch.nn.Conv1d(self.output_dimension, self.dimension, 1, bias=False)
    self.logvar_proj = torch.nn.Conv1d(self.output_dimension, self.dimension, 1, bias=False)

def encode(self, x: torch.Tensor) -> torch.Tensor:
    return self.input_proj(x) + torch.randn_like(x) * torch.exp(.5 * self.logvar_proj(x))

def decode(self, x: torch.Tensor) -> torch.Tensor:
    return self.output_proj(x)

Looking at this code, it seems like the encode function's shapes shouldn't work (randn_like has different shape than logvar_proj). Would this be the closer to the correct version?

     def encode(self, x: torch.Tensor) -> torch.Tensor:
         in_proj = self.input_proj(x)
         return in_proj + torch.randn_like(in_proj) * torch.exp(.5 * self.logvar_proj(x))

[manukyutai]

Hmm, I hope I didn't make a mistake, both self.input_proj(x) and self.logvar_proj(x) should have dimension self.dimension (=32)

[cfrancesco]

sorry, maybe this is clearer:

     def encode(self, x: torch.Tensor) -> torch.Tensor:
         logvar_proj = self.logvar_proj(x)
         return self.input_proj(x) + torch.randn_like(logvar_proj) * torch.exp(.5 * logvar_proj)  # instead of torch.randn_like(x)

the randn_like(x) shape mismatch with self.logvar_proj(x) was the issue

torch.randn_like(x) * torch.exp(.5 * self.logvar_proj(x))

[manukyutai]

Ah true! My bad! Your code is indeed correct