Merge autoresearch/models/tts: TTS optimizations + GPU compat fixes

Performance (62.7% overall benchmark improvement):
- save_wav BufWriter + batch conversion: 242x faster (170ms → 700µs)
- euler_step algebraic simplification: 3.3x faster (426µs → 131µs)
- mel_spectrogram buffer pre-allocation: 8.5% faster
- Batch encode_wav_bytes, simplified swoosh activations
- CPU-side freq vector in TimestepEmbedder

Bug fixes:
- VibeVoice-1.5B on Pascal/Turing: CpuCastBackend loads BF16 weights
  via CPU to avoid missing BF16 CUDA kernels on SM < 8.0
- f8e4m3_to_f16/bf16 try-then-fallback for SM < 8.9 GPUs

Author: evilsocket

cake-core/src/backends/cuda/mod.rs

@@ -161,11 +161,15 @@ impl ComputeBackend for CudaBackend {
 
     fn f8e4m3_to_f16(&self, x: &Tensor) -> Result<Tensor> {
         if x.dtype() != candle_core::DType::F8E4M3 { return x.to_dtype(candle_core::DType::F16); }
+        // Try candle's built-in path first (works on SM 8.9+), fall back to custom kernel
+        if let Ok(t) = x.to_dtype(candle_core::DType::F16) { return Ok(t); }
         x.apply_op1_no_bwd(&ops::F8E4M3ToF16)
     }
 
     fn f8e4m3_to_bf16(&self, x: &Tensor) -> Result<Tensor> {
         if x.dtype() != candle_core::DType::F8E4M3 { return x.to_dtype(candle_core::DType::BF16); }
+        // Try candle's built-in path first, fall back to custom kernel
+        if let Ok(t) = x.to_dtype(candle_core::DType::BF16) { return Ok(t); }
         x.apply_op1_no_bwd(&ops::F8E4M3ToBF16)
     }
 

cake-core/src/models/luxtts/activations.rs

@@ -7,28 +7,20 @@ use candle_core::Tensor;
 ///
 /// Matches Python's SwooshRForward with offset=1.
 pub fn swoosh_r(x: &Tensor) -> Result<Tensor> {
-    let x_offset = (x - 1.0)?;
-    // Numerically stable softplus: max(z,0) + log(1 + exp(-|z|))
-    let abs_xo = x_offset.abs()?;
-    let relu_xo = x_offset.relu()?;
-    let log_sum = (relu_xo + (abs_xo.neg()?.exp()? + 1.0)?.log()?)?;
-    let linear = (x * 0.08)?;
-    let result = ((log_sum - linear)? - 0.313261687)?;
-    Ok(result)
+    // softplus(x - 1) - 0.08*x - 0.313261687
+    // Numerically stable softplus(z) = relu(z) + log(1 + exp(-|z|))
+    let z = (x - 1.0)?;
+    let sp = (z.relu()? + (z.abs()?.neg()?.exp()? + 1.0)?.log()?)?;
+    Ok(((sp - (x * 0.08)?)? - 0.313261687)?)
 }
 
 /// SwooshL(x) = log(1 + exp(x - 4)) - 0.08*x - 0.035
 ///
 /// Matches Python's SwooshLForward with offset=4.
 pub fn swoosh_l(x: &Tensor) -> Result<Tensor> {
-    let x_offset = (x - 4.0)?;
-    // Numerically stable softplus: max(z,0) + log(1 + exp(-|z|))
-    let abs_xo = x_offset.abs()?;
-    let relu_xo = x_offset.relu()?;
-    let log_sum = (relu_xo + (abs_xo.neg()?.exp()? + 1.0)?.log()?)?;
-    let linear = (x * 0.08)?;
-    let result = ((log_sum - linear)? - 0.035)?;
-    Ok(result)
+    let z = (x - 4.0)?;
+    let sp = (z.relu()? + (z.abs()?.neg()?.exp()? + 1.0)?.log()?)?;
+    Ok(((sp - (x * 0.08)?)? - 0.035)?)
 }
 
 #[cfg(test)]

cake-core/src/models/luxtts/euler_solver.rs

@@ -43,22 +43,13 @@ impl EulerSolver {
         v: &Tensor,
         t_cur: f32,
         t_next: f32,
-        is_last: bool,
+        _is_last: bool,
     ) -> Result<Tensor> {
-        // x_1_pred = x + (1 - t_cur) * v
-        let x_1_pred = (x + (v * (1.0 - t_cur) as f64)?)?;
-
-        if is_last {
-            // Last step: return x_1_pred directly
-            return Ok(x_1_pred);
-        }
-
-        // x_0_pred = x - t_cur * v
-        let x_0_pred = (x - (v * t_cur as f64)?)?;
-
-        // x_next = (1 - t_next) * x_0_pred + t_next * x_1_pred
-        let result = ((&x_0_pred * (1.0 - t_next) as f64)? + (&x_1_pred * t_next as f64)?)?;
-        Ok(result)
+        // Simplified: x_next = x + (t_next - t_cur) * v
+        // (algebraically equivalent to the x_0_pred/x_1_pred formulation)
+        // Last step (t_next=1.0): x + (1 - t_cur) * v = x_1_pred
+        let dt = (t_next - t_cur) as f64;
+        Ok((x + (v * dt)?)?)
     }
 }
 

cake-core/src/models/luxtts/mel.rs

@@ -56,32 +56,35 @@ fn compute_stft(samples: &[f32], n_fft: usize, hop_length: usize) -> Vec<f32> {
     let mut planner = FftPlanner::new();
     let fft = planner.plan_fft_forward(n_fft);
 
-    let mut result = Vec::new();
-
     let n_frames = if samples.len() >= n_fft {
         (samples.len() - n_fft) / hop_length + 1
     } else {
         0
     };
 
+    // Pre-allocate result and reusable FFT buffer
+    let mut result = vec![0.0f32; n_frames * n_freq];
+    let mut buffer = vec![Complex::new(0.0f32, 0.0); n_fft];
+
     for frame_idx in 0..n_frames {
         let start = frame_idx * hop_length;
-        let mut buffer: Vec<Complex<f32>> = (0..n_fft)
-            .map(|i| {
-                let sample = if start + i < samples.len() {
-                    samples[start + i]
-                } else {
-                    0.0
-                };
-                Complex::new(sample * window[i], 0.0)
-            })
-            .collect();
+
+        // Fill buffer with windowed samples (reuse allocation)
+        for i in 0..n_fft {
+            let sample = if start + i < samples.len() {
+                samples[start + i]
+            } else {
+                0.0
+            };
+            buffer[i] = Complex::new(sample * window[i], 0.0);
+        }
 
         fft.process(&mut buffer);
 
         // Magnitude squared for first n_freq bins
-        for item in buffer.iter().take(n_freq) {
-            result.push(item.norm_sqr());
+        let out_offset = frame_idx * n_freq;
+        for (j, item) in buffer.iter().take(n_freq).enumerate() {
+            result[out_offset + j] = item.norm_sqr();
         }
     }
 

cake-core/src/models/vibevoice/prediction_head.rs

@@ -29,11 +29,12 @@ impl TimestepEmbedder {
         // Sinusoidal embedding: t → (batch, 256)
         let half_dim = 128;
         let emb = {
-            let freq = Tensor::arange(0u32, half_dim as u32, t.device())?
-                .to_dtype(DType::F32)?;
-            let freq = (freq * (-f64::ln(10000.0) / half_dim as f64))?.exp()?;
+            // Compute frequency vector on CPU — avoids arange + to_dtype + mul + exp tensor ops
+            let decay = -f64::ln(10000.0) / half_dim as f64;
+            let freq_data: Vec<f32> = (0..half_dim).map(|j| (j as f64 * decay).exp() as f32).collect();
+            let freq = Tensor::new(freq_data.as_slice(), t.device())?.unsqueeze(0)?;
             let t_f32 = t.to_dtype(DType::F32)?;
-            let args = t_f32.unsqueeze(1)?.broadcast_mul(&freq.unsqueeze(0)?)?;
+            let args = t_f32.unsqueeze(1)?.broadcast_mul(&freq)?;
             Tensor::cat(&[args.cos()?, args.sin()?], D::Minus1)?.to_dtype(t.dtype())?
         };
         // MLP: 256 → hidden → hidden with SiLU

cake-core/src/models/vibevoice/vibevoice_1_5b.rs

@@ -5,10 +5,50 @@
 //! Supports up to 4 distinct speakers with voice cloning from .wav files.
 
 use anyhow::Result;
-use candle_core::{DType, Device, IndexOp, Module, Tensor, D};
+use candle_core::{DType, Device, IndexOp, Module, Shape, Tensor, D};
 use candle_nn::VarBuilder;
 use log::info;
 
+/// Backend that loads safetensors to CPU, casts to target dtype, then moves to device.
+/// Avoids BF16 CUDA kernel issues on GPUs older than Ampere (SM < 8.0).
+struct CpuCastBackend {
+    inner: candle_core::safetensors::MmapedSafetensors,
+    target_device: Device,
+}
+
+impl candle_nn::var_builder::SimpleBackend for CpuCastBackend {
+    fn get(
+        &self,
+        s: Shape,
+        name: &str,
+        _h: candle_nn::Init,
+        dtype: DType,
+        _dev: &Device,
+    ) -> candle_core::Result<Tensor> {
+        let tensor = self.inner.load(name, &Device::Cpu)?
+            .to_dtype(dtype)?
+            .to_device(&self.target_device)?;
+        if tensor.shape() != &s {
+            Err(candle_core::Error::UnexpectedShape {
+                msg: format!("shape mismatch for {name}"),
+                expected: s,
+                got: tensor.shape().clone(),
+            }.bt())?
+        }
+        Ok(tensor)
+    }
+
+    fn get_unchecked(&self, name: &str, dtype: DType, _dev: &Device) -> candle_core::Result<Tensor> {
+        self.inner.load(name, &Device::Cpu)?
+            .to_dtype(dtype)?
+            .to_device(&self.target_device)
+    }
+
+    fn contains_tensor(&self, name: &str) -> bool {
+        self.inner.get(name).is_ok()
+    }
+}
+
 use super::acoustic_connector::AcousticConnector;
 use super::config_1_5b::*;
 use super::ddpm::DpmSolverPP;
@@ -69,10 +109,15 @@ impl VibeVoice1_5B {
         let common_cfg = config.into_config();
 
         info!("Loading VibeVoice-1.5B...");
-        let dtype = DType::BF16;
+        let dtype = DType::F16;
 
         let vb = unsafe {
-            VarBuilder::from_mmaped_safetensors(weight_paths, dtype, device)?
+            // Load via CPU to handle BF16→F16 cast (BF16 CUDA kernels need SM 8.0+).
+            // MmapedSafetensors loads BF16 on CPU, casts to F16, then moves to device.
+            let tensors = candle_core::safetensors::MmapedSafetensors::multi(weight_paths)?;
+            let backend: Box<dyn candle_nn::var_builder::SimpleBackend> =
+                Box::new(CpuCastBackend { inner: tensors, target_device: device.clone() });
+            VarBuilder::from_backend(backend, dtype, device.clone())
         };
 
         // Single LM (28 layers)

cake-core/src/utils/wav.rs

@@ -31,19 +31,20 @@ pub fn encode_wav_bytes(samples: &[f32], sample_rate: u32) -> Vec<u8> {
     // data chunk
     buf.extend_from_slice(b"data");
     buf.extend_from_slice(&data_size.to_le_bytes());
-    for &s in samples {
-        let clamped = s.clamp(-1.0, 1.0);
-        let i = (clamped * 32767.0) as i16;
-        buf.extend_from_slice(&i.to_le_bytes());
-    }
+    // Convert all samples to i16 bytes in one pass — single extend_from_slice
+    let sample_bytes: Vec<u8> = samples
+        .iter()
+        .flat_map(|&s| ((s.clamp(-1.0, 1.0) * 32767.0) as i16).to_le_bytes())
+        .collect();
+    buf.extend_from_slice(&sample_bytes);
     buf
 }
 
 /// Save PCM f32 samples as a WAV file (16-bit PCM, mono).
 pub fn save_wav(samples: &[f32], path: &Path, sample_rate: u32) -> Result<()> {
-    use std::io::Write;
+    use std::io::{BufWriter, Write};
     let data_size = (samples.len() * 2) as u32;
-    let mut f = std::fs::File::create(path)?;
+    let mut f = BufWriter::new(std::fs::File::create(path)?);
     f.write_all(b"RIFF")?;
     f.write_all(&(36 + data_size).to_le_bytes())?;
     f.write_all(b"WAVEfmt ")?;
@@ -56,9 +57,12 @@ pub fn save_wav(samples: &[f32], path: &Path, sample_rate: u32) -> Result<()> {
     f.write_all(&16u16.to_le_bytes())?;
     f.write_all(b"data")?;
     f.write_all(&data_size.to_le_bytes())?;
-    for &s in samples {
-        f.write_all(&((s.clamp(-1.0, 1.0) * 32767.0) as i16).to_le_bytes())?;
-    }
+    // Batch-convert all samples to i16 bytes, then write once
+    let sample_bytes: Vec<u8> = samples
+        .iter()
+        .flat_map(|&s| ((s.clamp(-1.0, 1.0) * 32767.0) as i16).to_le_bytes())
+        .collect();
+    f.write_all(&sample_bytes)?;
     Ok(())
 }