osobhandClaude Opus 4.7
f642132f4a
Mirror of --steering-layers for the decoder: comma-separated layer indices to actually steer (decoder has 4 layers; useful subsets are [3], [2,3], [1,2]). Sweep at seed 7, [email protected]: [0,1,2,3] cos 0.86 WER 0.93 "I'm not that tall ×3" ← repetition [3] cos 0.58 WER 0.86 "I'll be off and offense..." [2,3] cos 0.80 WER 1.43 "My daughter, Penny Ryan, and I have…" [0,1] cos 0.82 WER 1.57 "And I'll check on them..." [1,2] cos 0.81 WER 1.43 "I'm going to call him an X-Man..." The repetition is specific to all-layers-at-once steering. Any 2-layer subset eliminates it while preserving most of the cosine boost. Same pattern as the backbone's [8,10,12] finding: partial perturbation lets the unsteered layers act as a stabilizing prior. [2,3] (decoder last 2) is the new recommended recipe — best cosine of the no-repetition subsets and the longest fluent transcript. Documented in docs/perf_history.md. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
383 lines
14 KiB
Rust
383 lines
14 KiB
Rust
//! CSM generation CLI.
|
||
//!
|
||
//! Downloads CSM-1B + Mimi + Llama tokenizer from HuggingFace (cached), then
|
||
//! generates speech for the given `--text`.
|
||
//!
|
||
//! Usage:
|
||
//! ```
|
||
//! cargo run -p rtx-csm --release --example generate -- \
|
||
//! --text "Hello from Rust." --speaker 0 --out /tmp/hello.wav
|
||
//! ```
|
||
//!
|
||
//! Context conditioning (optional, from a prior utterance you have WAV for).
|
||
//! `--context-wav` and `--context-text` are repeatable and zipped in order:
|
||
//! ```
|
||
//! cargo run -p rtx-csm --release --example generate -- \
|
||
//! --text "Nice to talk to you." \
|
||
//! --context-wav prior_a.wav --context-text "First clip transcript." \
|
||
//! --context-wav prior_b.wav --context-text "Second clip transcript." \
|
||
//! --context-speaker 1 \
|
||
//! --speaker 0 --out /tmp/reply.wav
|
||
//! ```
|
||
//! Tip: rank manifest clips for context suitability with `scripts/pick_context.sh`.
|
||
|
||
use anyhow::Result;
|
||
use clap::Parser;
|
||
use rtx_csm::{GenerateOptions, Generator, PostProcess, Segment, audio_io};
|
||
|
||
#[derive(Debug, Parser)]
|
||
#[command(name = "csm-generate", about = "Generate speech with rtx-csm")]
|
||
struct Cli {
|
||
/// Text to synthesize.
|
||
#[arg(long)]
|
||
text: String,
|
||
|
||
/// Speaker id (0 or 1).
|
||
#[arg(long, default_value_t = 0)]
|
||
speaker: u32,
|
||
|
||
/// Output WAV file (24 kHz mono 16-bit).
|
||
#[arg(long)]
|
||
out: std::path::PathBuf,
|
||
|
||
/// Max audio length in milliseconds.
|
||
#[arg(long, default_value_t = 10_000)]
|
||
max_audio_ms: u32,
|
||
|
||
/// Sampling temperature.
|
||
#[arg(long, default_value_t = 0.9)]
|
||
temperature: f64,
|
||
|
||
/// Top-K sampling.
|
||
#[arg(long, default_value_t = 50)]
|
||
top_k: usize,
|
||
|
||
/// Top-p (nucleus) cutoff applied after top-k. 1.0 disables.
|
||
#[arg(long, default_value_t = 0.9)]
|
||
top_p: f64,
|
||
|
||
/// RNG seed.
|
||
#[arg(long, default_value_t = 42)]
|
||
seed: u64,
|
||
|
||
/// Optional emotion control tag prepended to the text prompt, e.g.
|
||
/// `--emotion-hint "[whisper]"`. Only meaningful when the model has been
|
||
/// fine-tuned with matching tags in its training prompts (see
|
||
/// `docs/personal_voice_training_guide.md`); on the un-adapted base it is
|
||
/// just an extra cosmetic prefix the BPE tokenizer encodes as ordinary
|
||
/// tokens.
|
||
#[arg(long)]
|
||
emotion_hint: Option<String>,
|
||
|
||
/// Prior utterance audio (WAV) for context conditioning. Repeatable;
|
||
/// each `--context-wav` must be paired with a matching `--context-text`
|
||
/// in the same order.
|
||
#[arg(long)]
|
||
context_wav: Vec<std::path::PathBuf>,
|
||
|
||
/// Transcription of the context WAV. Repeatable; must have the same
|
||
/// count as `--context-wav`.
|
||
#[arg(long)]
|
||
context_text: Vec<String>,
|
||
|
||
/// Speaker id assigned to all context utterances. CSM only distinguishes
|
||
/// 0/1; pick something different from `--speaker` if you want the model
|
||
/// to treat the cloning reference as a separate voice.
|
||
#[arg(long, default_value_t = 1)]
|
||
context_speaker: u32,
|
||
|
||
/// Force CPU device even if cuda/metal features are enabled.
|
||
#[arg(long)]
|
||
cpu: bool,
|
||
|
||
/// Load a quantized GGUF model instead of the default safetensors path.
|
||
/// Path should point at the file produced by `examples/quantize`.
|
||
#[arg(long)]
|
||
quantized_gguf: Option<std::path::PathBuf>,
|
||
|
||
/// Apply a LoRA adapter (safetensors) trained via `examples/lora_train`.
|
||
/// Adapter is injected into the FP backbone before generation.
|
||
#[arg(long)]
|
||
lora: Option<std::path::PathBuf>,
|
||
|
||
/// LoRA rank override. When omitted, auto-detected from the adapter
|
||
/// file's embedded metadata (Phase 12.5); falls back to 8 for older
|
||
/// adapters without metadata.
|
||
#[arg(long)]
|
||
lora_rank: Option<usize>,
|
||
|
||
/// LoRA alpha override. When omitted, auto-detected from the adapter's
|
||
/// embedded metadata; falls back to 16.0 for older adapters.
|
||
#[arg(long)]
|
||
lora_alpha: Option<f32>,
|
||
|
||
/// Force extended LoRA coverage (q+k+v+output_proj + MLP). When omitted,
|
||
/// auto-detected from the adapter's embedded metadata. Setting this on a
|
||
/// classic q+v adapter just allocates extra unused B=0 slots.
|
||
#[arg(long, default_value_t = false)]
|
||
extended_lora: bool,
|
||
|
||
/// Activation-steering safetensors (one (embed_dim,) tensor per layer
|
||
/// keyed `layer_<i>_steering`). Vectors are added to each backbone
|
||
/// layer's residual stream output during generation. See
|
||
/// `crate::steering` and Sprint 2 of the Phase 8 roadmap. Quantized
|
||
/// GGUF backbone is not supported.
|
||
#[arg(long)]
|
||
steering_vec: Option<std::path::PathBuf>,
|
||
|
||
/// Global multiplier on every steering vector. EmoSteer-TTS uses 2.0
|
||
/// for emotion conversion, 2.5 for erasure. 0.0 disables steering
|
||
/// without dropping the file.
|
||
#[arg(long, default_value_t = 1.0)]
|
||
steering_scale: f32,
|
||
|
||
/// Comma-separated layer indices to actually steer (e.g. "4,8,12").
|
||
/// Defaults to all layers in the loaded safetensors. EmoSteer-TTS
|
||
/// found that perturbing a small spaced subset of middle-to-deep
|
||
/// layers preserves fluency better than perturbing every layer.
|
||
#[arg(long, value_delimiter = ',')]
|
||
steering_layers: Option<Vec<usize>>,
|
||
|
||
/// Steering vectors for the depth decoder (acoustic codebooks).
|
||
/// Decoder is Llama100M for CSM-1B → 4 layers × 1024 embed_dim, so
|
||
/// these are SMALLER than backbone vectors. Architectural hypothesis:
|
||
/// decoder steering shifts prosody/timbre without disturbing word
|
||
/// content the way backbone steering does.
|
||
#[arg(long)]
|
||
decoder_steering_vec: Option<std::path::PathBuf>,
|
||
|
||
/// Scale on the decoder steering vector. Default 1.0.
|
||
#[arg(long, default_value_t = 1.0)]
|
||
decoder_steering_scale: f32,
|
||
|
||
/// Comma-separated layer indices to actually steer on the decoder
|
||
/// (e.g. "2,3"). The decoder has 4 layers (CSM-1B); useful subsets
|
||
/// to try: `3` (deepest), `2,3` (last two), `1,2` (middle two).
|
||
/// Defaults to all 4 layers.
|
||
#[arg(long, value_delimiter = ',')]
|
||
decoder_steering_layers: Option<Vec<usize>>,
|
||
|
||
/// Per-frame CFG scale schedule (Selective CFG, arXiv 2509.19668).
|
||
/// Forms:
|
||
/// - `const:<scale>` (constant scale, equivalent to --cfg-scale)
|
||
/// - `step:<early>:<late>:<T>` (early until frame T, then late)
|
||
/// - `linear:<start>:<end>:<R>` (interpolate over R frames, hold)
|
||
///
|
||
/// Requires non-empty context AND CFG-enabled load (--enable-cfg).
|
||
/// Recommended: step:3.0:1.5:12 — full CFG for the first ~1s, then
|
||
/// drop to preserve text adherence.
|
||
#[arg(long)]
|
||
cfg_schedule: Option<String>,
|
||
|
||
/// Constant CFG scale (legacy single-value form). Ignored if
|
||
/// `--cfg-schedule` is also passed.
|
||
#[arg(long)]
|
||
cfg_scale: Option<f64>,
|
||
|
||
/// Load with the CFG-capable dual backbone. Required for any cfg_*
|
||
/// flag to take effect. Adds ~1.5x memory; only worth it when you're
|
||
/// actually using CFG.
|
||
#[arg(long, default_value_t = false)]
|
||
enable_cfg: bool,
|
||
|
||
/// Disable audio post-processing (HPF + declick + LUFS normalize).
|
||
#[arg(long)]
|
||
raw: bool,
|
||
|
||
/// LUFS target for loudness normalization. Ignored if --raw.
|
||
#[arg(long, default_value_t = -16.0)]
|
||
lufs: f32,
|
||
|
||
/// Path to converted AudioSeal generator safetensors (run
|
||
/// `audioseal_convert` first). When set together with
|
||
/// `--watermark-detector`, the watermarker is wired into
|
||
/// `generate_to_wav` so output is automatically watermarked.
|
||
#[arg(long)]
|
||
watermark_generator: Option<std::path::PathBuf>,
|
||
|
||
/// Path to converted AudioSeal detector safetensors. Required for the
|
||
/// watermarker even if you only want to embed (the detector is part of
|
||
/// AudioSealWatermarker construction; future builds may make it
|
||
/// optional).
|
||
#[arg(long)]
|
||
watermark_detector: Option<std::path::PathBuf>,
|
||
|
||
/// 16-bit watermark message (decimal or 0xHEX).
|
||
#[arg(long, default_value = "0")]
|
||
watermark_message: String,
|
||
}
|
||
|
||
fn main() -> Result<()> {
|
||
tracing_subscriber::fmt().init();
|
||
|
||
let cli = Cli::parse();
|
||
let device = if cli.cpu {
|
||
candle_core::Device::Cpu
|
||
} else {
|
||
Generator::default_device()?
|
||
};
|
||
tracing::info!("device: {device:?}");
|
||
|
||
let mut generator = if let Some(gguf) = cli.quantized_gguf.as_ref() {
|
||
tracing::info!("loading quantized CSM from {}", gguf.display());
|
||
Generator::load_csm_1b_quantized(gguf, &device, false)?
|
||
} else if cli.enable_cfg {
|
||
tracing::info!("loading CSM-1B with CFG-enabled dual backbone");
|
||
Generator::load_csm_1b_with_cfg(&device, true)?
|
||
} else {
|
||
Generator::load_csm_1b(&device)?
|
||
};
|
||
|
||
if let Some(steering_path) = cli.steering_vec.as_ref() {
|
||
let n = generator
|
||
.backbone_num_layers()
|
||
.ok_or_else(|| anyhow::anyhow!("steering not supported on the quantized backbone"))?;
|
||
let mut steering =
|
||
rtx_csm::steering::LayerSteering::load_safetensors(steering_path, n, &device)?;
|
||
steering.set_scale(cli.steering_scale);
|
||
if let Some(layers) = cli.steering_layers.as_ref() {
|
||
steering.restrict_to_layers(layers);
|
||
}
|
||
tracing::info!(
|
||
"loaded backbone steering from {} (scale={}, active_layers={:?})",
|
||
steering_path.display(),
|
||
cli.steering_scale,
|
||
steering.active_layers(),
|
||
);
|
||
generator.set_steering(Some(steering))?;
|
||
}
|
||
|
||
if let Some(dec_path) = cli.decoder_steering_vec.as_ref() {
|
||
let n = generator
|
||
.decoder_num_layers()
|
||
.ok_or_else(|| anyhow::anyhow!("decoder steering needs the FP backbone"))?;
|
||
let mut steering =
|
||
rtx_csm::steering::LayerSteering::load_safetensors(dec_path, n, &device)?;
|
||
steering.set_scale(cli.decoder_steering_scale);
|
||
if let Some(layers) = cli.decoder_steering_layers.as_ref() {
|
||
steering.restrict_to_layers(layers);
|
||
}
|
||
tracing::info!(
|
||
"loaded decoder steering from {} (scale={}, active_layers={:?})",
|
||
dec_path.display(),
|
||
cli.decoder_steering_scale,
|
||
steering.active_layers(),
|
||
);
|
||
generator.set_decoder_steering(Some(steering))?;
|
||
}
|
||
|
||
if let Some(lora_path) = cli.lora.as_ref() {
|
||
// `--extended-lora` is a force flag (no way to clap-distinguish "off"
|
||
// from "unset" on a bare bool). When the user passes it we override
|
||
// to true; otherwise we let the file metadata or the default decide.
|
||
let extended_override = if cli.extended_lora { Some(true) } else { None };
|
||
rtx_csm::training::apply_lora_adapter(
|
||
&mut generator,
|
||
lora_path,
|
||
cli.lora_rank,
|
||
cli.lora_alpha,
|
||
extended_override,
|
||
&device,
|
||
)?;
|
||
}
|
||
tracing::info!(
|
||
"model loaded (sr={} Hz, frame_rate={} Hz, codebooks={})",
|
||
generator.config.sample_rate,
|
||
generator.config.frame_rate_hz,
|
||
generator.config.audio_num_codebooks,
|
||
);
|
||
|
||
if cli.context_wav.len() != cli.context_text.len() {
|
||
anyhow::bail!(
|
||
"--context-wav and --context-text must be repeated the same number of times \
|
||
(got {} wav(s), {} text(s))",
|
||
cli.context_wav.len(),
|
||
cli.context_text.len(),
|
||
);
|
||
}
|
||
let mut context: Vec<Segment> = Vec::with_capacity(cli.context_wav.len());
|
||
for (wav, txt) in cli.context_wav.iter().zip(cli.context_text.iter()) {
|
||
let audio = audio_io::load_mono_24k(wav)?;
|
||
tracing::info!(
|
||
"loaded context #{}: {} samples ({:.2} s) from {}",
|
||
context.len(),
|
||
audio.len(),
|
||
audio.len() as f32 / generator.config.sample_rate as f32,
|
||
wav.display()
|
||
);
|
||
context.push(Segment::new(cli.context_speaker, txt, audio));
|
||
}
|
||
|
||
// Optional watermarker wiring (AudioSeal + 24k↔16k resample adapter).
|
||
if let (Some(gen_path), Some(det_path)) = (
|
||
cli.watermark_generator.as_ref(),
|
||
cli.watermark_detector.as_ref(),
|
||
) {
|
||
let msg_str = cli.watermark_message.trim();
|
||
let message: u16 = if let Some(rest) = msg_str
|
||
.strip_prefix("0x")
|
||
.or_else(|| msg_str.strip_prefix("0X"))
|
||
{
|
||
u16::from_str_radix(rest, 16)?
|
||
} else {
|
||
msg_str.parse::<u16>()?
|
||
};
|
||
let gen_vb = unsafe {
|
||
candle_nn::VarBuilder::from_mmaped_safetensors(
|
||
&[gen_path],
|
||
candle_core::DType::F32,
|
||
&device,
|
||
)
|
||
}?;
|
||
let det_vb = unsafe {
|
||
candle_nn::VarBuilder::from_mmaped_safetensors(
|
||
&[det_path],
|
||
candle_core::DType::F32,
|
||
&device,
|
||
)
|
||
}?;
|
||
let inner = rtx_csm::AudioSealWatermarker::from_var_builders(
|
||
gen_vb,
|
||
det_vb,
|
||
device.clone(),
|
||
message,
|
||
)?;
|
||
// CSM produces 24 kHz; AudioSeal native is 16 kHz.
|
||
let wm = rtx_csm::ResampledWatermarker::new(inner, generator.config.sample_rate, 16_000);
|
||
generator.set_watermarker(Box::new(wm));
|
||
tracing::info!(
|
||
"watermarker installed (message=0x{:04X}, model 16 kHz, output {} Hz)",
|
||
message,
|
||
generator.config.sample_rate
|
||
);
|
||
}
|
||
|
||
let cfg_schedule = match (cli.cfg_schedule.as_ref(), cli.cfg_scale) {
|
||
(Some(s), _) => Some(rtx_csm::cfg_schedule::CfgSchedule::parse(s)?),
|
||
(None, Some(_)) => None, // legacy cfg_scale path handled by GenerateOptions::cfg_scale
|
||
(None, None) => None,
|
||
};
|
||
let opts = GenerateOptions {
|
||
max_audio_ms: cli.max_audio_ms,
|
||
temperature: cli.temperature,
|
||
top_k: cli.top_k,
|
||
top_p: cli.top_p,
|
||
seed: cli.seed,
|
||
emotion_hint: cli.emotion_hint.clone(),
|
||
cfg_scale: cli.cfg_scale,
|
||
cfg_schedule,
|
||
..GenerateOptions::default()
|
||
};
|
||
let post = if cli.raw {
|
||
PostProcess::disabled()
|
||
} else {
|
||
PostProcess {
|
||
lufs_target: Some(cli.lufs),
|
||
..PostProcess::default()
|
||
}
|
||
};
|
||
generator.generate_to_wav(&cli.text, cli.speaker, &context, opts, &post, &cli.out)?;
|
||
println!("wrote {}", cli.out.display());
|
||
Ok(())
|
||
}
|