Method encode_streaming

cake-core/src/models/vibevoice/vae_encoder.rs:362–414

Streaming encode: uses cache for correct context between frames.

(
        &self,
        audio: &Tensor,
        cache: &mut super::vae_decoder::StreamingConvCache,
    )

    /// Streaming encode: uses cache for correct context between frames.
    pub fn encode_streaming(
        &self,
        audio: &Tensor,
        cache: &mut super::vae_decoder::StreamingConvCache,
    ) -> Result<Tensor> {
        let x = if audio.rank() == 2 {
            audio.unsqueeze(1)?
        } else {
            audio.clone()
        };

        cache.reset_counter();
        let mut h = x;

        for (i, (conv, stage)) in self
            .downsample_convs
            .iter()
            .zip(self.stages.iter())
            .enumerate()
        {
            // Streaming: use cached context instead of zero-padding
            let ctx_size = self.downsample_paddings[i];
            let (slot, is_first) = cache.take_slot();
            let context = if is_first {
                Tensor::zeros((h.dim(0)?, h.dim(1)?, ctx_size), h.dtype(), h.device())?
            } else {
                cache.get(slot).unwrap().clone()
            };
            let padded = Tensor::cat(&[&context, &h], 2)?;

            // Update cache: last ctx_size samples of padded
            let plen = padded.dim(2)?;
            let start = plen.saturating_sub(ctx_size);
            cache.set(slot, padded.narrow(2, start, plen - start)?);

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

        // Head conv: streaming context
        let (slot, is_first) = cache.take_slot();
        let ctx = if is_first {
            Tensor::zeros((h.dim(0)?, h.dim(1)?, 6), h.dtype(), h.device())?
        } else {
            cache.get(slot).unwrap().clone()
        };
        let padded = Tensor::cat(&[&ctx, &h], 2)?;
        let plen = padded.dim(2)?;
        cache.set(slot, padded.narrow(2, plen.saturating_sub(6), 6.min(plen))?);
        h = self.backend.conv1d(&padded, &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)?;

        h.transpose(1, 2)
    }
}

#[cfg(test)]