Streaming decode: uses cache for correct context between frames. Each call processes a single latent frame and produces audio samples.
(&self, latents: &Tensor, cache: &mut StreamingConvCache)
| 426 | /// Streaming decode: uses cache for correct context between frames. |
| 427 | /// Each call processes a single latent frame and produces audio samples. |
| 428 | pub fn decode_streaming(&self, latents: &Tensor, cache: &mut StreamingConvCache) -> Result<Tensor> { |
| 429 | let x = if latents.dim(1)? == 64 { |
| 430 | latents.clone() |
| 431 | } else { |
| 432 | latents.transpose(1, 2)? |
| 433 | }; |
| 434 | |
| 435 | cache.reset_counter(); |
| 436 | let mut h = x; |
| 437 | let profile = log::log_enabled!(log::Level::Trace); |
| 438 | let mut stage_times: Vec<(f64, f64, usize)> = Vec::new(); // (upsample_ms, blocks_ms, seq_len) |
| 439 | |
| 440 | for (i, (upsample, stage)) in self.upsample_layers.iter().zip(self.stages.iter()).enumerate() { |
| 441 | let t_stage = std::time::Instant::now(); |
| 442 | if i == 0 { |
| 443 | // Conv1d k=7: use cached context (6 samples) instead of zero-pad |
| 444 | let (slot, is_first) = cache.take_slot(); |
| 445 | let ctx = if is_first { |
| 446 | Tensor::zeros((h.dim(0)?, h.dim(1)?, 6), h.dtype(), h.device())? |
| 447 | } else { |
| 448 | cache.get(slot).unwrap().clone() |
| 449 | }; |
| 450 | let padded = Tensor::cat(&[&ctx, &h], 2)?; |
| 451 | let plen = padded.dim(2)?; |
| 452 | cache.set(slot, padded.narrow(2, plen.saturating_sub(6), 6.min(plen))?); |
| 453 | h = upsample.forward(&padded, &*self.backend)?; |
| 454 | } else { |
| 455 | // ConvTranspose1d: cache input history for context |
| 456 | let kernel_size = self.ratios[i] * 2; |
| 457 | let ctx_size = kernel_size - 1; |
| 458 | let stride = self.ratios[i]; |
| 459 | let new_len = h.dim(2)?; |
| 460 | |
| 461 | let (slot, is_first) = cache.take_slot(); |
| 462 | |
| 463 | let full_input = if is_first { |
| 464 | h.clone() |
| 465 | } else { |
| 466 | let cached = cache.get(slot).unwrap(); |
| 467 | Tensor::cat(&[cached, &h], 2)? |
| 468 | }; |
| 469 | |
| 470 | let full_output = upsample.forward(&full_input, &*self.backend)?; |
| 471 | let full_output = Self::causal_trim(&full_output, stride)?; |
| 472 | |
| 473 | h = if is_first { |
| 474 | full_output |
| 475 | } else { |
| 476 | let out_len = full_output.dim(2)?; |
| 477 | let new_out = new_len * stride; |
| 478 | full_output.narrow(2, out_len - new_out, new_out)? |
| 479 | }; |
| 480 | |
| 481 | // Update cache: last ctx_size samples of full_input |
| 482 | let fi_len = full_input.dim(2)?; |
| 483 | if fi_len > ctx_size { |
| 484 | cache.set(slot, full_input.narrow(2, fi_len - ctx_size, ctx_size)?); |
| 485 | } else { |