695 lines
26 KiB
Rust
695 lines
26 KiB
Rust
//! High-level Generator façade.
|
||
|
||
use crate::audio_io::TARGET_SAMPLE_RATE;
|
||
use crate::config::ModelConfig;
|
||
use crate::error::Result;
|
||
use crate::hub;
|
||
use crate::mimi::Mimi;
|
||
use crate::model::CsmModel;
|
||
use crate::post::PostProcess;
|
||
use crate::prompt::{Segment, build_prompt};
|
||
use crate::repetition::{RepetitionConfig, RepetitionGuard};
|
||
use crate::sampler::{CsmSampler, DEFAULT_TEMPERATURE, DEFAULT_TOPK, DEFAULT_TOPP};
|
||
use crate::text_norm::TextNormalize;
|
||
use crate::tokenizer::CsmTokenizer;
|
||
use crate::util;
|
||
use candle_core::{DType, Device, Tensor};
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct GenerateOptions {
|
||
pub max_audio_ms: u32,
|
||
pub temperature: f64,
|
||
pub top_k: usize,
|
||
/// Nucleus (top-p) filter applied after top-k. Set to `1.0` (or `0.0`) to disable.
|
||
pub top_p: f64,
|
||
pub seed: u64,
|
||
/// Loop-escape guard. `None` disables (not recommended).
|
||
pub repetition: Option<RepetitionConfig>,
|
||
/// Classifier-Free Guidance scale on the codebook-0 head. `None` or `Some(1.0)`
|
||
/// disables CFG. Koel-TTS recommends 1.5–3.0. Requires the generator to
|
||
/// have been loaded via [`Generator::load_csm_1b_with_cfg`] AND the
|
||
/// generation call to provide non-empty context (without context the
|
||
/// uncond branch == cond branch and CFG has no effect).
|
||
///
|
||
/// Equivalent to `cfg_schedule = Some(CfgSchedule::Constant(s))` when
|
||
/// `cfg_schedule` is None. The schedule field takes precedence when set.
|
||
pub cfg_scale: Option<f64>,
|
||
/// Optional per-frame CFG scale schedule (Selective CFG, arXiv 2509.19668).
|
||
/// Overrides `cfg_scale` when set. Use `Step { early, late, transition }`
|
||
/// to start at one CFG scale for the first `transition` frames and drop
|
||
/// to a lower one — preserves text adherence late in synthesis while
|
||
/// keeping speaker fidelity early.
|
||
pub cfg_schedule: Option<crate::cfg_schedule::CfgSchedule>,
|
||
/// Optional control-token hint prepended to the text prompt, e.g.
|
||
/// `Some("[whisper]".into())`. Concatenated as `<hint> <text>` and fed
|
||
/// through the same Llama BPE path as ordinary text — the hint is plain
|
||
/// text from the model's perspective. Only useful when the LoRA / fine-tune
|
||
/// has been trained with the same tags in its prompts (per
|
||
/// `personal_voice_training_guide.md` §4): the adapter learns
|
||
/// `[whisper] <text>` → whispered prosody. Out of the box (no fine-tune)
|
||
/// this is a no-op cosmetic prefix.
|
||
pub emotion_hint: Option<String>,
|
||
}
|
||
|
||
impl Default for GenerateOptions {
|
||
fn default() -> Self {
|
||
Self {
|
||
max_audio_ms: 10_000,
|
||
temperature: DEFAULT_TEMPERATURE,
|
||
top_k: DEFAULT_TOPK,
|
||
top_p: DEFAULT_TOPP,
|
||
seed: 42,
|
||
repetition: Some(RepetitionConfig::default()),
|
||
cfg_scale: None,
|
||
cfg_schedule: None,
|
||
emotion_hint: None,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Prepend an optional control-token hint to the (already-normalized) text.
|
||
/// Joined with a single space when present so the BPE tokenizer treats the tag
|
||
/// as a separate sub-sequence from the body text. Centralized here so every
|
||
/// generate variant — and the LoRA trainer in `training.rs` — applies the
|
||
/// same convention. Training and inference *must* use identical prefix
|
||
/// formatting or the adapter won't transfer.
|
||
pub(crate) fn apply_emotion_hint(text: String, hint: Option<&str>) -> String {
|
||
match hint {
|
||
Some(tag) if !tag.is_empty() => format!("{} {}", tag.trim(), text),
|
||
_ => text,
|
||
}
|
||
}
|
||
|
||
pub struct Generator {
|
||
pub model: CsmModel,
|
||
pub mimi: Mimi,
|
||
pub tokenizer: CsmTokenizer,
|
||
pub config: ModelConfig,
|
||
pub device: Device,
|
||
/// Applied to every text input before tokenization. Set to
|
||
/// [`TextNormalize::passthrough`] if you've pre-normalized upstream.
|
||
pub text_normalize: TextNormalize,
|
||
/// Optional watermarker applied inside [`Self::generate_to_wav`] AFTER
|
||
/// post-processing and BEFORE WAV write. Set via
|
||
/// [`Self::set_watermarker`]. `None` = no-op (Sesame's reference TTS
|
||
/// also ships unwatermarked by default; this is the integration hook).
|
||
pub watermarker: Option<Box<dyn crate::watermark::Watermarker>>,
|
||
}
|
||
|
||
impl Generator {
|
||
pub fn new(model: CsmModel, mimi: Mimi, tokenizer: CsmTokenizer, device: Device) -> Self {
|
||
let config = model.config.clone();
|
||
Self {
|
||
model,
|
||
mimi,
|
||
tokenizer,
|
||
config,
|
||
device,
|
||
text_normalize: TextNormalize::default(),
|
||
watermarker: None,
|
||
}
|
||
}
|
||
|
||
/// Install a watermarker that runs inside `generate_to_wav` after
|
||
/// post-processing. Take ownership of the watermarker so the generator
|
||
/// can be moved into worker threads (Watermarker is `Send + Sync`).
|
||
pub fn set_watermarker(&mut self, wm: Box<dyn crate::watermark::Watermarker>) {
|
||
self.watermarker = Some(wm);
|
||
}
|
||
|
||
/// Drop any installed watermarker.
|
||
pub fn clear_watermarker(&mut self) {
|
||
self.watermarker = None;
|
||
}
|
||
|
||
/// Like [`Self::load_csm_1b`] but loads the CSM weights from an
|
||
/// explicit path instead of the HuggingFace cache. Mimi + Llama
|
||
/// tokenizer are still resolved through the hub. Use this with a
|
||
/// LoRA-merged checkpoint produced by
|
||
/// `lora::merge_into_safetensors`.
|
||
pub fn load_csm_1b_from_path<P: AsRef<std::path::Path>>(
|
||
csm_weights: P,
|
||
device: &Device,
|
||
) -> Result<Self> {
|
||
let assets = hub::resolve_csm_1b()?;
|
||
let csm_weights = csm_weights.as_ref().to_path_buf();
|
||
let config = ModelConfig::csm_1b();
|
||
let dtype = match device {
|
||
Device::Cpu => DType::F32,
|
||
Device::Metal(_) => DType::F16,
|
||
_ => DType::BF16,
|
||
};
|
||
let model = CsmModel::load_from_safetensors(&csm_weights, config.clone(), dtype, device)?;
|
||
let mimi = Mimi::load(&assets.mimi_weights, device)?;
|
||
let tokenizer = CsmTokenizer::from_file(&assets.tokenizer_json)?;
|
||
Ok(Self::new(model, mimi, tokenizer, device.clone()))
|
||
}
|
||
|
||
/// Download (cached) CSM-1B + Mimi + Llama tokenizer from HuggingFace and
|
||
/// build a ready-to-generate `Generator`.
|
||
pub fn load_csm_1b(device: &Device) -> Result<Self> {
|
||
let assets = hub::resolve_csm_1b()?;
|
||
let config = ModelConfig::csm_1b();
|
||
|
||
// Optional: audit tensor keys once before loading to surface naming drift.
|
||
if std::env::var("CSM_AUDIT_KEYS").is_ok() {
|
||
let descs = crate::model::dump_safetensors_keys(&assets.csm_weights)?;
|
||
let missing = crate::model::audit_csm_keys(&descs);
|
||
if missing.is_empty() {
|
||
tracing::info!("safetensors key audit: ok ({} tensors)", descs.len());
|
||
} else {
|
||
tracing::warn!("safetensors missing keys: {missing:?}");
|
||
}
|
||
}
|
||
|
||
// dtype selection by backend:
|
||
// CPU → F32 (candle's CPU backend has no BF16 matmul kernel)
|
||
// Metal → F16 (BF16 is ~50% slower than F16 on M1/M2; M3+ added hw bf16
|
||
// but f16 still ties or wins. CSM softmax is well-behaved
|
||
// post-RMSNorm so f16 dynamic range is fine in practice.)
|
||
// CUDA → BF16 (modern NVIDIA tensor cores prefer BF16)
|
||
let dtype = match device {
|
||
Device::Cpu => DType::F32,
|
||
Device::Metal(_) => DType::F16,
|
||
_ => DType::BF16,
|
||
};
|
||
let model =
|
||
CsmModel::load_from_safetensors(&assets.csm_weights, config.clone(), dtype, device)?;
|
||
let mimi = Mimi::load(&assets.mimi_weights, device)?;
|
||
let tokenizer = CsmTokenizer::from_file(&assets.tokenizer_json)?;
|
||
|
||
Ok(Self::new(model, mimi, tokenizer, device.clone()))
|
||
}
|
||
|
||
/// Like [`Self::load_csm_1b`] but allocates a second backbone for the
|
||
/// unconditional CFG branch. Adds ~70 MB of KV cache; weight tensors are
|
||
/// shared via mmap.
|
||
pub fn load_csm_1b_with_cfg(device: &Device, enable_cfg: bool) -> Result<Self> {
|
||
let assets = hub::resolve_csm_1b()?;
|
||
let config = ModelConfig::csm_1b();
|
||
let dtype = match device {
|
||
Device::Cpu => DType::F32,
|
||
Device::Metal(_) => DType::F16,
|
||
_ => DType::BF16,
|
||
};
|
||
let model = CsmModel::load_from_safetensors_with_cfg(
|
||
&assets.csm_weights,
|
||
config.clone(),
|
||
dtype,
|
||
device,
|
||
enable_cfg,
|
||
)?;
|
||
let mimi = Mimi::load(&assets.mimi_weights, device)?;
|
||
let tokenizer = CsmTokenizer::from_file(&assets.tokenizer_json)?;
|
||
Ok(Self::new(model, mimi, tokenizer, device.clone()))
|
||
}
|
||
|
||
/// Load a quantized CSM-1B from a GGUF file (the artifact emitted by
|
||
/// `examples/quantize`). Defaults to candle's raw QMatMul kernel path —
|
||
/// Q8 weights stay quantized at runtime, the Metal/CPU kernel performs
|
||
/// fused dequant+matmul. This is the actual quantization perf win:
|
||
/// smaller weight memory at runtime AND correct, well-pronounced output.
|
||
///
|
||
/// Verified 2026-04-25 via per-layer hidden-state bisection: hidden
|
||
/// states match between raw-QTensor and F16-dequant paths to <0.02% on
|
||
/// every backbone layer. Earlier "gibberish" was traced to the v1 GGUF
|
||
/// having an incomplete quant policy that left MLP weights at F16 while
|
||
/// attention was Q8 — the mixed-dtype model interacted with the kernel
|
||
/// paths. The kernel itself is correct.
|
||
///
|
||
/// To force F16-dequantization-on-load (occasionally useful for
|
||
/// numerical-precision debugging or to test against non-Q kernels), set
|
||
/// `CSM_DEQUANT_F16=1` in the environment before invoking. That path
|
||
/// gives F16 runtime weights — same memory profile as the F16 safetensors
|
||
/// path, no inference speedup.
|
||
pub fn load_csm_1b_quantized<P: AsRef<std::path::Path>>(
|
||
gguf_path: P,
|
||
device: &Device,
|
||
enable_cfg: bool,
|
||
) -> Result<Self> {
|
||
// Opt-in fallback: F16 dequantization at load time.
|
||
if std::env::var("CSM_DEQUANT_F16").is_ok() {
|
||
// SAFETY: candle reads this env var via a thread-local on first
|
||
// access. Setting before QMatMul construction is the supported
|
||
// pattern.
|
||
unsafe { std::env::set_var("CANDLE_DEQUANTIZE_ALL_F16", "1") };
|
||
}
|
||
|
||
let mimi_weights = hub::resolve_mimi()?;
|
||
let tokenizer_json = hub::resolve_llama_tokenizer()?;
|
||
let config = ModelConfig::csm_1b();
|
||
// F32 runtime is the safest choice for the quantized path even with
|
||
// DEQUANTIZE_ALL_F16. Reason: tensors stored as F16 in the GGUF
|
||
// (kept-native heads/embeds/MLPs etc.) get auto-dequantized to F16
|
||
// tensors by candle's QMatMul construction, which would clash with
|
||
// F16 runtime when also mixing TensorF16-wrapped Q8 weights. Routing
|
||
// everything through F32 activations + F16 dequantize_f16 paths
|
||
// (auto-cast inside TensorF16 forward) avoids any dtype mismatch.
|
||
let runtime_dtype = DType::F32;
|
||
let model =
|
||
CsmModel::load_from_gguf(gguf_path, config.clone(), runtime_dtype, device, enable_cfg)?;
|
||
let mimi = Mimi::load(&mimi_weights, device)?;
|
||
let tokenizer = CsmTokenizer::from_file(&tokenizer_json)?;
|
||
Ok(Self::new(model, mimi, tokenizer, device.clone()))
|
||
}
|
||
|
||
pub fn reset(&mut self) {
|
||
self.model.clear_kv_cache();
|
||
}
|
||
|
||
/// Install activation-steering vectors on the conditional Llama backbone.
|
||
/// Only the FP backend supports steering today; on a quantized backend
|
||
/// this returns an error rather than silently no-op'ing. See
|
||
/// `crate::steering`. Pass `None` to remove.
|
||
pub fn set_steering(
|
||
&mut self,
|
||
steering: Option<crate::steering::LayerSteering>,
|
||
) -> crate::Result<()> {
|
||
match &mut self.model.inner {
|
||
crate::model::ModelBackend::Fp(m) => {
|
||
m.set_backbone_steering(steering);
|
||
Ok(())
|
||
}
|
||
crate::model::ModelBackend::Quantized(_) => Err(crate::error::CsmError::Config(
|
||
"activation steering is only supported on the FP backbone \
|
||
(--quantized-gguf disables it)"
|
||
.into(),
|
||
)),
|
||
}
|
||
}
|
||
|
||
/// Number of backbone layers — useful when constructing
|
||
/// `LayerSteering::empty(n)` to match the loaded model. Returns `None`
|
||
/// for the quantized backend (steering not supported there).
|
||
pub fn backbone_num_layers(&self) -> Option<usize> {
|
||
match &self.model.inner {
|
||
crate::model::ModelBackend::Fp(m) => Some(m.backbone_num_layers()),
|
||
crate::model::ModelBackend::Quantized(_) => None,
|
||
}
|
||
}
|
||
|
||
/// Install steering on the depth decoder (acoustic codebooks). Vectors
|
||
/// must be `(decoder_embed_dim,)` — for CSM-1B that's 1024, NOT 2048.
|
||
/// Architectural hypothesis: backbone carries semantic content (what
|
||
/// the model says), decoder carries acoustic detail (how it sounds);
|
||
/// steering the decoder should shift prosody/timbre without disturbing
|
||
/// word-level fidelity.
|
||
pub fn set_decoder_steering(
|
||
&mut self,
|
||
steering: Option<crate::steering::LayerSteering>,
|
||
) -> crate::Result<()> {
|
||
match &mut self.model.inner {
|
||
crate::model::ModelBackend::Fp(m) => {
|
||
m.set_decoder_steering(steering);
|
||
Ok(())
|
||
}
|
||
crate::model::ModelBackend::Quantized(_) => Err(crate::error::CsmError::Config(
|
||
"decoder steering only supported on the FP backbone".into(),
|
||
)),
|
||
}
|
||
}
|
||
|
||
/// Number of decoder layers (CSM-1B = 4). `None` on quantized.
|
||
pub fn decoder_num_layers(&self) -> Option<usize> {
|
||
match &self.model.inner {
|
||
crate::model::ModelBackend::Fp(m) => Some(m.decoder_num_layers()),
|
||
crate::model::ModelBackend::Quantized(_) => None,
|
||
}
|
||
}
|
||
|
||
/// Full generation loop: prompt → backbone/decoder per-frame → Mimi decode → f32 PCM.
|
||
pub fn generate(
|
||
&mut self,
|
||
text: &str,
|
||
speaker: u32,
|
||
context: &[Segment],
|
||
opts: GenerateOptions,
|
||
) -> Result<Vec<f32>> {
|
||
self.reset();
|
||
|
||
let normalized = self.text_normalize.apply(text)?;
|
||
let prompted = apply_emotion_hint(normalized, opts.emotion_hint.as_deref());
|
||
let current = Segment::new_text(speaker, prompted);
|
||
let prompt = build_prompt(
|
||
context,
|
||
¤t,
|
||
&self.model,
|
||
&mut self.mimi,
|
||
&self.tokenizer,
|
||
)?;
|
||
|
||
let cb = self.config.audio_num_codebooks;
|
||
let mut sampler = CsmSampler::new(opts.seed, opts.temperature, opts.top_k, opts.top_p);
|
||
let inner_lp = sampler.inner_mut();
|
||
|
||
// CFG path: dual backbone if all preconditions are met. The
|
||
// schedule field takes precedence over the legacy `cfg_scale`.
|
||
let cfg_schedule = opts.cfg_schedule.or_else(|| {
|
||
opts.cfg_scale
|
||
.map(crate::cfg_schedule::CfgSchedule::Constant)
|
||
});
|
||
let cfg_active = cfg_schedule.map(|s| s.is_active()).unwrap_or(false)
|
||
&& !context.is_empty()
|
||
&& self.model.inner.cfg_enabled();
|
||
let cfg_schedule = if cfg_active {
|
||
cfg_schedule.unwrap()
|
||
} else {
|
||
crate::cfg_schedule::CfgSchedule::Constant(1.0)
|
||
};
|
||
|
||
let uncond_prompt_opt = if cfg_active {
|
||
Some(build_prompt(
|
||
&[],
|
||
¤t,
|
||
&self.model,
|
||
&mut self.mimi,
|
||
&self.tokenizer,
|
||
)?)
|
||
} else {
|
||
None
|
||
};
|
||
|
||
let mut pos: usize = 0;
|
||
let mut uncond_pos: usize = 0;
|
||
let mut all_frames: Vec<Vec<u32>> = Vec::new();
|
||
let max_frames =
|
||
((opts.max_audio_ms as f32) / self.config.frame_duration_ms()).ceil() as usize;
|
||
|
||
let mut input_tokens = prompt.tokens;
|
||
let mut input_mask = prompt.mask;
|
||
let mut uncond_tokens = uncond_prompt_opt.as_ref().map(|p| p.tokens.clone());
|
||
let mut uncond_mask = uncond_prompt_opt.as_ref().map(|p| p.mask.clone());
|
||
|
||
let mut rep_guard = opts.repetition.map(RepetitionGuard::new);
|
||
|
||
if cfg_active {
|
||
tracing::info!("CFG active (schedule={cfg_schedule:?})");
|
||
}
|
||
|
||
for frame_idx in 0..max_frames {
|
||
let sampled = if cfg_active {
|
||
let ut = uncond_tokens.as_ref().unwrap();
|
||
let um = uncond_mask.as_ref().unwrap();
|
||
let frame_scale = cfg_schedule.scale_at(frame_idx);
|
||
let r = self.model.inner.generate_frame_cfg(
|
||
&input_tokens,
|
||
&input_mask,
|
||
pos,
|
||
ut,
|
||
um,
|
||
uncond_pos,
|
||
frame_scale,
|
||
inner_lp,
|
||
)?;
|
||
uncond_pos += ut.dim(1)?;
|
||
r
|
||
} else {
|
||
self.model
|
||
.inner
|
||
.generate_frame(&input_tokens, &input_mask, pos, inner_lp)?
|
||
};
|
||
pos += input_tokens.dim(1)?;
|
||
|
||
let is_eot = frame_idx >= 1 && sampled.iter().all(|v| *v == 0);
|
||
if is_eot {
|
||
tracing::info!("EOT detected at frame {frame_idx}");
|
||
break;
|
||
}
|
||
|
||
if let Some(g) = rep_guard.as_mut()
|
||
&& g.observe(&sampled)
|
||
{
|
||
tracing::warn!(
|
||
"loop-escape: repetition guard tripped at frame {frame_idx}; ending generation"
|
||
);
|
||
break;
|
||
}
|
||
|
||
all_frames.push(sampled.clone());
|
||
|
||
let (t, m) = self.model.inner.audio_tokens_and_mask(sampled)?;
|
||
// Both branches feed the same sampled frame (sampling produces a single
|
||
// discrete decision; the uncond cache must track the chosen path too).
|
||
if cfg_active {
|
||
uncond_tokens = Some(t.clone());
|
||
uncond_mask = Some(m.clone());
|
||
}
|
||
input_tokens = t;
|
||
input_mask = m;
|
||
}
|
||
|
||
if all_frames.is_empty() {
|
||
return Ok(Vec::new());
|
||
}
|
||
|
||
// Pack frames into (1, cb, T) i64 and decode through Mimi.
|
||
let t = all_frames.len();
|
||
let mut flat: Vec<u32> = Vec::with_capacity(t * cb);
|
||
// Transpose: all_frames is Vec<frame: Vec<cb>>; we want row-major (cb, t).
|
||
for c in 0..cb {
|
||
for frame in &all_frames {
|
||
flat.push(frame[c]);
|
||
}
|
||
}
|
||
let codes = Tensor::from_vec(flat, (1, cb, t), &self.device)?.to_dtype(DType::U32)?;
|
||
let pcm = self.mimi.decode(&codes)?;
|
||
tracing::info!(
|
||
"generated {} samples (~{:.2}s at {} Hz)",
|
||
pcm.len(),
|
||
pcm.len() as f32 / TARGET_SAMPLE_RATE as f32,
|
||
TARGET_SAMPLE_RATE
|
||
);
|
||
Ok(pcm)
|
||
}
|
||
|
||
/// Streaming version of [`Self::generate`]: invokes `on_chunk` with new PCM
|
||
/// samples every `chunk_frames` frames (default 4 = ~320 ms) so a downstream
|
||
/// player can start audio output before generation completes.
|
||
///
|
||
/// Uses Mimi's `decode_step` (StreamTensor-based incremental decode), so
|
||
/// total decode cost is O(n) instead of O(n²). First-audio latency drops
|
||
/// from ~T_total to `chunk_frames × 80 ms + per_frame_compute × chunk_frames`.
|
||
pub fn generate_streaming<F>(
|
||
&mut self,
|
||
text: &str,
|
||
speaker: u32,
|
||
context: &[Segment],
|
||
opts: GenerateOptions,
|
||
chunk_frames: usize,
|
||
mut on_chunk: F,
|
||
) -> Result<Vec<f32>>
|
||
where
|
||
F: FnMut(&[f32]) -> Result<()>,
|
||
{
|
||
self.reset();
|
||
// Mimi's streaming state must be reset at the start of every stream;
|
||
// leftover state from a prior call corrupts the first chunk.
|
||
self.mimi.reset_state();
|
||
|
||
let chunk_frames = chunk_frames.max(1);
|
||
|
||
let normalized = self.text_normalize.apply(text)?;
|
||
let prompted = apply_emotion_hint(normalized, opts.emotion_hint.as_deref());
|
||
let current = Segment::new_text(speaker, prompted);
|
||
let prompt = build_prompt(
|
||
context,
|
||
¤t,
|
||
&self.model,
|
||
&mut self.mimi,
|
||
&self.tokenizer,
|
||
)?;
|
||
|
||
let cb = self.config.audio_num_codebooks;
|
||
let mut sampler = CsmSampler::new(opts.seed, opts.temperature, opts.top_k, opts.top_p);
|
||
let inner_lp = sampler.inner_mut();
|
||
|
||
let mut pos: usize = 0;
|
||
let max_frames =
|
||
((opts.max_audio_ms as f32) / self.config.frame_duration_ms()).ceil() as usize;
|
||
let mut input_tokens = prompt.tokens;
|
||
let mut input_mask = prompt.mask;
|
||
|
||
// Pending frames not yet sent to Mimi's decode_step.
|
||
let mut pending: Vec<Vec<u32>> = Vec::with_capacity(chunk_frames);
|
||
let mut full_pcm: Vec<f32> = Vec::new();
|
||
let mut rep_guard = opts.repetition.map(RepetitionGuard::new);
|
||
|
||
for frame_idx in 0..max_frames {
|
||
let sampled =
|
||
self.model
|
||
.inner
|
||
.generate_frame(&input_tokens, &input_mask, pos, inner_lp)?;
|
||
pos += input_tokens.dim(1)?;
|
||
|
||
let is_eot = frame_idx >= 1 && sampled.iter().all(|v| *v == 0);
|
||
if is_eot {
|
||
tracing::info!("EOT detected at frame {frame_idx}");
|
||
break;
|
||
}
|
||
if let Some(g) = rep_guard.as_mut()
|
||
&& g.observe(&sampled)
|
||
{
|
||
tracing::warn!(
|
||
"loop-escape: repetition guard tripped at frame {frame_idx}; ending"
|
||
);
|
||
break;
|
||
}
|
||
|
||
pending.push(sampled.clone());
|
||
|
||
if pending.len() >= chunk_frames {
|
||
stream_decode_pending(
|
||
&mut pending,
|
||
cb,
|
||
&mut self.mimi,
|
||
&self.device,
|
||
&mut full_pcm,
|
||
&mut on_chunk,
|
||
)?;
|
||
}
|
||
|
||
let (t, m) = self.model.inner.audio_tokens_and_mask(sampled)?;
|
||
input_tokens = t;
|
||
input_mask = m;
|
||
}
|
||
|
||
// Flush any frames left in the buffer.
|
||
if !pending.is_empty() {
|
||
stream_decode_pending(
|
||
&mut pending,
|
||
cb,
|
||
&mut self.mimi,
|
||
&self.device,
|
||
&mut full_pcm,
|
||
&mut on_chunk,
|
||
)?;
|
||
}
|
||
// Drain any internal Mimi buffering by feeding a `None` step.
|
||
if let Some(tail) = self.mimi.decode_step(None)?
|
||
&& !tail.is_empty()
|
||
{
|
||
on_chunk(&tail)?;
|
||
full_pcm.extend_from_slice(&tail);
|
||
}
|
||
|
||
tracing::info!(
|
||
"streaming generation finished: {} samples (~{:.2}s)",
|
||
full_pcm.len(),
|
||
full_pcm.len() as f32 / self.config.sample_rate as f32
|
||
);
|
||
Ok(full_pcm)
|
||
}
|
||
|
||
/// Generate using a `SpeakerProfile` as context. The profile is automatically
|
||
/// fit to budget; the caller's provided context segments come *after* the profile.
|
||
pub fn generate_with_profile(
|
||
&mut self,
|
||
profile: &crate::speaker::SpeakerProfile,
|
||
text: &str,
|
||
extra_context: &[Segment],
|
||
opts: GenerateOptions,
|
||
) -> Result<Vec<f32>> {
|
||
let mut profile = profile.clone();
|
||
profile.fit_within_budget(crate::speaker::DEFAULT_PROFILE_BUDGET_TOKENS);
|
||
let mut ctx = profile.segments().to_vec();
|
||
ctx.extend_from_slice(extra_context);
|
||
self.generate(text, profile.id, &ctx, opts)
|
||
}
|
||
|
||
/// Convenience: `generate` + apply default post-processing + watermark
|
||
/// (if installed) + write WAV. Pass [`PostProcess::disabled`] to skip
|
||
/// post-processing. Pass [`Self::clear_watermarker`] (or never install
|
||
/// one) to skip watermarking.
|
||
///
|
||
/// Order: model → post-process (HPF + declick + LUFS) → watermark.
|
||
/// Watermarking comes last so the loudness target the user sees on disk
|
||
/// is the loudness target the user requested (the watermark residual
|
||
/// is at most a few dB and well below LUFS measurement floor).
|
||
pub fn generate_to_wav(
|
||
&mut self,
|
||
text: &str,
|
||
speaker: u32,
|
||
context: &[Segment],
|
||
opts: GenerateOptions,
|
||
post: &PostProcess,
|
||
out_path: &std::path::Path,
|
||
) -> Result<()> {
|
||
let mut pcm = self.generate(text, speaker, context, opts)?;
|
||
post.apply(&mut pcm, self.config.sample_rate)?;
|
||
if let Some(wm) = self.watermarker.as_ref() {
|
||
pcm = wm.embed(&pcm)?;
|
||
}
|
||
crate::audio_io::write_wav_24k_mono(out_path, &pcm)?;
|
||
Ok(())
|
||
}
|
||
|
||
/// Choose the best available device based on enabled features.
|
||
pub fn default_device() -> Result<Device> {
|
||
util::pick_device()
|
||
}
|
||
}
|
||
|
||
/// Take all `pending` frames, build a (1, cb, k) tensor, push to Mimi's
|
||
/// streaming decoder, emit any audio it produces, and clear `pending`.
|
||
fn stream_decode_pending<F>(
|
||
pending: &mut Vec<Vec<u32>>,
|
||
num_codebooks: usize,
|
||
mimi: &mut Mimi,
|
||
device: &Device,
|
||
full_pcm: &mut Vec<f32>,
|
||
on_chunk: &mut F,
|
||
) -> Result<()>
|
||
where
|
||
F: FnMut(&[f32]) -> Result<()>,
|
||
{
|
||
if pending.is_empty() {
|
||
return Ok(());
|
||
}
|
||
let k = pending.len();
|
||
let mut flat: Vec<u32> = Vec::with_capacity(k * num_codebooks);
|
||
for c in 0..num_codebooks {
|
||
for frame in pending.iter() {
|
||
flat.push(frame[c]);
|
||
}
|
||
}
|
||
let codes = Tensor::from_vec(flat, (1, num_codebooks, k), device)?.to_dtype(DType::U32)?;
|
||
pending.clear();
|
||
if let Some(samples) = mimi.decode_step(Some(&codes))?
|
||
&& !samples.is_empty()
|
||
{
|
||
on_chunk(&samples)?;
|
||
full_pcm.extend_from_slice(&samples);
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn apply_emotion_hint_no_op_when_none() {
|
||
let out = apply_emotion_hint("Hello there".into(), None);
|
||
assert_eq!(out, "Hello there");
|
||
}
|
||
|
||
#[test]
|
||
fn apply_emotion_hint_no_op_when_empty() {
|
||
let out = apply_emotion_hint("Hello there".into(), Some(""));
|
||
assert_eq!(out, "Hello there");
|
||
}
|
||
|
||
#[test]
|
||
fn apply_emotion_hint_prepends_with_space() {
|
||
let out = apply_emotion_hint("Hello there".into(), Some("[whisper]"));
|
||
assert_eq!(out, "[whisper] Hello there");
|
||
}
|
||
|
||
#[test]
|
||
fn apply_emotion_hint_trims_outer_whitespace_on_tag() {
|
||
// Stray whitespace on the tag itself should not double up the separator.
|
||
let out = apply_emotion_hint("Hi".into(), Some(" [excited] "));
|
||
assert_eq!(out, "[excited] Hi");
|
||
}
|
||
}
|