Load encoder with explicit vae_dim (for semantic tokenizer which has different dim).
(
vb: VarBuilder,
cfg: &super::config::AcousticTokenizerConfig,
vae_dim: usize,
backend: Arc<dyn ComputeBackend>,
)
| 190 | |
| 191 | /// Load encoder with explicit vae_dim (for semantic tokenizer which has different dim). |
| 192 | pub fn load_with_vae_dim( |
| 193 | vb: VarBuilder, |
| 194 | cfg: &super::config::AcousticTokenizerConfig, |
| 195 | vae_dim: usize, |
| 196 | backend: Arc<dyn ComputeBackend>, |
| 197 | ) -> Result<Self> { |
| 198 | let eps = cfg.layernorm_eps; |
| 199 | let n_filters = cfg.encoder_n_filters; |
| 200 | // Encoder ratios are REVERSED from config (config has decoder order) |
| 201 | let ratios: Vec<usize> = cfg.encoder_ratios.iter().rev().copied().collect(); |
| 202 | |
| 203 | let num_stages = ratios.len() + 1; // 7 stages for 6 ratios |
| 204 | |
| 205 | // Channel progression: n_filters, n_filters*2, n_filters*4, ..., n_filters*2^6 |
| 206 | // For n_filters=32: [32, 64, 128, 256, 512, 1024, 2048] |
| 207 | let mut channels = Vec::with_capacity(num_stages); |
| 208 | for i in 0..num_stages { |
| 209 | channels.push(n_filters * (1 << i)); |
| 210 | } |
| 211 | |
| 212 | // Downsample layers |
| 213 | let mut downsample_convs = Vec::with_capacity(num_stages); |
| 214 | let mut downsample_paddings = Vec::with_capacity(num_stages); |
| 215 | let mut downsample_strides = Vec::with_capacity(num_stages); |
| 216 | |
| 217 | // Stem: Conv1d(1 → n_filters, kernel=7, stride=1) |
| 218 | // Causal padding = (kernel-1)*dilation - (stride-1) = 6 |
| 219 | let stem_vb = vb.pp("downsample_layers").pp("0").pp("0").pp("conv").pp("conv"); |
| 220 | let stem_weight = stem_vb.get((channels[0], 1, 7), "weight")?; |
| 221 | let stem_bias = stem_vb.get(channels[0], "bias").ok(); |
| 222 | downsample_convs.push(RawConv1d { |
| 223 | weight: stem_weight, bias: stem_bias, |
| 224 | padding: 0, stride: 1, dilation: 1, groups: 1, |
| 225 | }); |
| 226 | downsample_paddings.push(6); // (7-1) - (1-1) = 6 |
| 227 | downsample_strides.push(1); |
| 228 | |
| 229 | // Downsampling Conv1d layers with stride |
| 230 | for (i, &ratio) in ratios.iter().enumerate() { |
| 231 | let in_ch = channels[i]; |
| 232 | let out_ch = channels[i + 1]; |
| 233 | let kernel = ratio * 2; |
| 234 | // Causal padding = (kernel-1) - (stride-1) = kernel - stride = 2*ratio - ratio = ratio |
| 235 | let conv_vb = vb.pp("downsample_layers") |
| 236 | .pp(i + 1) |
| 237 | .pp("0") |
| 238 | .pp("conv") |
| 239 | .pp("conv"); |
| 240 | let conv_weight = conv_vb.get((out_ch, in_ch, kernel), "weight")?; |
| 241 | let conv_bias = conv_vb.get(out_ch, "bias").ok(); |
| 242 | downsample_convs.push(RawConv1d { |
| 243 | weight: conv_weight, bias: conv_bias, |
| 244 | padding: 0, stride: ratio, dilation: 1, groups: 1, |
| 245 | }); |
| 246 | downsample_paddings.push(kernel - ratio); // (2*ratio - 1) - (ratio - 1) = ratio |
| 247 | downsample_strides.push(ratio); |
| 248 | } |
| 249 |