Class WaveCodec

tokenizer.py:471–538

@si_module
class WaveCodec(nn.Module):
    class Config:
        resnet_config: ResNetStack.Config = None
        sample_rate: int = 16_000
        use_weight_norm: bool = False

        compressor_config: dataclass = None

        norm_stddev: float = 1.0

    def __init__(self, c: Config):
        super().__init__()
        self.norm_stddev = c.norm_stddev
        self.encoder = c.resnet_config(mode='encoder')
        self.sample_rate = c.sample_rate

        self.total_stride = 1
        for stride in c.resnet_config.strides:
            self.total_stride *= stride
        self.tokens_per_second = self.sample_rate / self.total_stride
        
        self.compressor = c.compressor_config(dim=self.encoder.middle_channels)

        self.decoder = c.resnet_config(mode='decoder')        

        if c.use_weight_norm:
            self.encoder.apply_weight_norm()
            self.decoder.apply_weight_norm()
            self.encoder.reset_parameters()
            self.decoder.reset_parameters()

    def encode(self, data):
        return self.encoder(data/self.norm_stddev)

    def decode(self, latent):
        return self.decoder(latent.transpose(1, 2))*self.norm_stddev

    @T.no_grad()
    def latent_from_data(self, data, get_parameters=False):
        x = self.encode(data)
        l_in = x.transpose(1, 2)
        l, latent = self.compressor(l_in)
        return latent['z'] if not get_parameters else {
            'mu': latent['mu'],
            'logvar': latent['logvar'],
            'z': latent['z'],
        }
    
    @T.no_grad()
    def data_from_latent(self, latent):
        l = self.compressor.repr_from_latent(latent)
        x = self.decode(l)
        return x
    
    def process(self, x):
        return self.latent_from_data(x)

    def unprocess(self, latent):