Method __init__

tokenizer.py:321–361

(self, c: Config)

        mode: Literal['encoder', 'decoder'] = 'encoder'
        
    def __init__(self, c: Config):
        super().__init__()
        assert c.mode in ('encoder', 'decoder'), f"Mode ({c.mode}) is not supported!"
        
        self.mode = c.mode

        assert len(c.channel_ratios) == len(c.strides)
        channel_ratios = (1,) + c.channel_ratios
        strides = c.strides
        self.middle_channels = c.encode_channels * channel_ratios[-1]
        if c.mode == 'decoder':
            channel_ratios = tuple(reversed(channel_ratios))
            strides = tuple(reversed(strides))

        self.multiplier = c.decode_channel_multiplier if c.mode == 'decoder' else 1
        res_blocks = [ResNetBlock(
            c.encode_channels * channel_ratios[s_idx] * self.multiplier, 
            c.encode_channels * channel_ratios[s_idx+1] * self.multiplier, 
            stride,
            kernel_size=c.kernel_size,
            bias=c.bias,
            mode=c.mode,
        ) for s_idx, stride in enumerate(strides)]

        data_conv = CausalConv1d(
            in_channels=c.input_channels if c.mode == 'encoder' else c.encode_channels * self.multiplier, 
            out_channels=c.encode_channels if c.mode == 'encoder' else c.output_channels, 
            kernel_size=c.kernel_size, 
            stride=1, 
            bias=False,
        )

        if c.mode == 'encoder':
            self.res_stack = nn.Sequential(data_conv, *res_blocks)
        elif c.mode == 'decoder':
            self.res_stack = nn.Sequential(*res_blocks, data_conv)

        if c.latent_dim is not None:
            self.latent_proj = Conv1d1x1(self.middle_channels, c.latent_dim, bias=c.bias) if c.mode == 'encoder' else Conv1d1x1(c.latent_dim, self.middle_channels, bias=c.bias)
        if self.multiplier != 1:
            self.multiplier_proj = Conv1d1x1(self.middle_channels, self.middle_channels * self.multiplier, bias=c.bias)

    def forward(self, x, return_feats=False):
        if self.c.latent_dim is not None and self.mode == 'decoder':