Method encode

cake-core/src/models/vibevoice/vae_encoder.rs:304–359

Encode audio waveform to latent representation. Input: (batch, 1, samples) or (batch, samples) Output: (batch, frames, vae_dim)

(&self, audio: &Tensor)

    /// Input: (batch, 1, samples) or (batch, samples)
    /// Output: (batch, frames, vae_dim)
    pub fn encode(&self, audio: &Tensor) -> Result<Tensor> {
        // Ensure input is (batch, channels=1, samples)
        let x = if audio.rank() == 2 {
            audio.unsqueeze(1)?
        } else {
            audio.clone()
        };

        let mut h = x;

        // Interleave: downsample → stage blocks
        for (i, (conv, stage)) in self
            .downsample_convs
            .iter()
            .zip(self.stages.iter())
            .enumerate()
        {
            // Causal left-pad then conv
            h = Self::causal_pad(&h, self.downsample_paddings[i])?;

            // For strided convolutions, add extra padding for alignment
            if self.downsample_strides[i] > 1 {
                let kernel = conv.weight.dim(2)?;
                let stride = self.downsample_strides[i];
                let length = h.dim(2)?;
                let n_frames = (length - kernel) / stride + 1;
                let ideal_length = n_frames * stride + kernel;
                if ideal_length < length {
                    // This shouldn't happen with proper causal padding
                } else if ideal_length > length {
                    let extra = ideal_length - length;
                    h = Tensor::cat(
                        &[
                            &h,
                            &Tensor::zeros(
                                (h.dim(0)?, h.dim(1)?, extra),
                                h.dtype(),
                                h.device(),
                            )?,
                        ],
                        2,
                    )?;
                }
            }

            h = self.backend.conv1d(&h, &conv.weight, conv.bias.as_ref(), conv.padding, conv.stride, conv.dilation, conv.groups)?;
            h = stage.forward(&h)?;
        }

        // Head conv: causal left-pad 6
        h = Self::causal_pad(&h, 6)?;
        h = self.backend.conv1d(&h, &self.head_conv.weight, self.head_conv.bias.as_ref(), self.head_conv.padding, self.head_conv.stride, self.head_conv.dilation, self.head_conv.groups)?;

        // Return as (batch, frames, vae_dim)
        h.transpose(1, 2)
    }

    /// Streaming encode: uses cache for correct context between frames.