Function run_vibevoice_0_5b

cake-cli/src/main.rs:365–436

(
    ctx: Context,
    model_path: &std::path::Path,
    config_path: &std::path::Path,
)

#[cfg(feature = "master")]
#[cfg(feature = "vibevoice")]
async fn run_vibevoice_0_5b(
    ctx: Context,
    model_path: &std::path::Path,
    config_path: &std::path::Path,
) -> Result<()> {
    use cake_core::models::vibevoice;

    let weights_path = model_path.join("model.safetensors");
    println!("[VibeVoice-0.5B] Loading from {}", model_path.display());

    let model = vibevoice::VibeVoiceTTS::load(
        config_path,
        &weights_path,
        &ctx.device,
        Some(ctx.args.tts_diffusion_steps),
        &ctx.topology,
        ctx.args.cluster_key.as_deref(),
    )
    .await?;

    let prompt = &ctx.args.prompt;
    println!("[VibeVoice-0.5B] Generating speech for: \"{}\"", prompt);

    let tokenizer = {
        let local = model_path.join("tokenizer.json");
        if local.exists() {
            tokenizers::Tokenizer::from_file(&local)
                .map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?
        } else {
            println!("[VibeVoice-0.5B] Downloading Qwen2.5 tokenizer...");
            let qwen_path = utils::hf::ensure_model_downloaded("Qwen/Qwen2.5-0.5B")?;
            tokenizers::Tokenizer::from_file(qwen_path.join("tokenizer.json"))
                .map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?
        }
    };

    let text_with_newline = format!("{}\n", prompt.trim());
    let encoding = tokenizer
        .encode(text_with_newline.as_str(), false)
        .map_err(|e| anyhow::anyhow!("tokenize: {e}"))?;
    let token_ids = encoding.get_ids();
    println!("[VibeVoice-0.5B] Tokenized: {} tokens", token_ids.len());

    let voice_path = ctx
        .args
        .voice_prompt
        .as_ref()
        .ok_or_else(|| anyhow::anyhow!("--voice-prompt required for TTS"))?;
    let voice_prompt = vibevoice::VoicePrompt::load_f32(
        std::path::Path::new(voice_path),
        &ctx.device,
    )?;

    let mut model = model;
    let samples = model.generate(
        token_ids,
        &voice_prompt,
        ctx.args.max_audio_frames,