From 89ed3d68c134387b767b07c6315fe72e72fccdcb Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Sun, 5 Jul 2026 00:32:41 +0800 Subject: [PATCH 01/34] cosyvoice: add preliminary streaming TTS support --- include/cosyvoice-interface.h | 41 +++- include/cosyvoice-lowlevel.h | 50 +++++ include/cosyvoice.h | 56 +++++ src/cosyvoice-graph.cpp | 68 ++++-- src/cosyvoice-model.cpp | 18 ++ src/cosyvoice-model.h | 5 + src/cosyvoice-modules.h | 21 +- src/cosyvoice-tts.cpp | 376 +++++++++++++++++++++++----------- src/cosyvoice.cpp | 80 +++++++- 9 files changed, 555 insertions(+), 160 deletions(-) diff --git a/include/cosyvoice-interface.h b/include/cosyvoice-interface.h index 7f06756..c640fd1 100644 --- a/include/cosyvoice-interface.h +++ b/include/cosyvoice-interface.h @@ -50,6 +50,8 @@ struct cosyvoice_model_context // KV Cache virtual uint32_t llm_get_kv_cache_len() = 0; ///< Get the current KV-cache sequence length. virtual bool llm_set_kv_cache_len(uint32_t len) = 0; ///< Trim the KV-cache sequence length. + virtual void llm_offload_kv_cache() = 0; ///< Offload the KV cache to CPU memory. + virtual void llm_load_kv_cache() = 0; ///< Load the KV cache from CPU memory back to the backend. // Token Sampling and Acceptance virtual int llm_sample_token() = 0; ///< Sample the next token from the current logits. @@ -60,14 +62,41 @@ struct cosyvoice_model_context virtual const int* llm_get_accepted_tokens() = 0; ///< Get the accepted-token buffer. // Inference - virtual bool llm_job(const int* text, uint32_t text_len, cosyvoice_prompt_t prompt) = 0; ///< Run low-level LLM inference for a prompt and tokenized text. + virtual bool llm_job( + const int* text, + uint32_t text_len, + cosyvoice_prompt_t prompt + ) = 0; ///< Run low-level LLM inference for a prompt and tokenized text. + virtual bool token2wav( - const int* token_ids, - uint32_t n_tokens, - float speed, - cosyvoice_prompt_t prompt, + const int* token_ids, + uint32_t n_tokens, + float speed, + cosyvoice_prompt_t prompt, cosyvoice_generated_speech_ptr result ) = 0; ///< Convert speech tokens into waveform samples. + + // Extended Inference + virtual bool llm_job_ext( + const int* text, + uint32_t text_len, + cosyvoice_prompt_t prompt, + uint32_t max_new_tokens, + bool* final + ) = 0; ///< Run low-level LLM inference for a prompt and tokenized text with additional options. + + virtual bool token2wav_ext( + const int* token_ids, + uint32_t n_tokens, + float speed, + cosyvoice_prompt_t prompt, + uint32_t* speech_offset_ptr, + bool streaming, + bool finalize, + cosyvoice_generated_speech_ptr result + ) = 0; ///< Convert speech tokens into waveform samples with additional options. + + virtual uint32_t get_chunk_tokens() = 0; ///< Get the number of tokens processed in each chunk during streaming inference. // Status virtual ggml_status get_last_status() = 0; ///< Get the status of the most recent backend operation. @@ -76,7 +105,7 @@ struct cosyvoice_model_context virtual void set_prompt( cosyvoice_prompt_t prompt, cosyvoice_inference_mode_t mode, - const int* instruction, + const int* instruction, uint32_t instruction_length ) = 0; diff --git a/include/cosyvoice-lowlevel.h b/include/cosyvoice-lowlevel.h index 3d0795e..5731607 100644 --- a/include/cosyvoice-lowlevel.h +++ b/include/cosyvoice-lowlevel.h @@ -217,6 +217,16 @@ COSYVOICE_API uint32_t cosyvoice_llm_get_kv_cache_len(cosyvoice_context_t ctx); */ COSYVOICE_API bool cosyvoice_llm_set_kv_cache_len(cosyvoice_context_t ctx, uint32_t len); +/** +* @brief Offload the KV cache to CPU memory. +*/ +COSYVOICE_API void cosyvoice_llm_offload_kv_cache(cosyvoice_context_t ctx); + +/** +* @brief Load the KV cache from CPU memory back to the backend. +*/ +COSYVOICE_API void cosyvoice_llm_load_kv_cache(cosyvoice_context_t ctx); + // ---------------------------------------------------------------------------- // Sampling and Token Management // ---------------------------------------------------------------------------- @@ -266,6 +276,20 @@ COSYVOICE_API bool cosyvoice_llm_job( cosyvoice_prompt_t prompt ); +/** +* @brief Run the LLM with the given input tokens and prompt, with additional options. + * @param max_new_tokens Maximum number of new tokens to generate. If 0, no new tokens are generated. + * @param final Output parameter indicating whether the generation is complete (true) or more tokens can be generated (false). + */ +COSYVOICE_API bool cosyvoice_llm_job_ext( + cosyvoice_context_t ctx, + const int* text, + uint32_t text_len, + cosyvoice_prompt_t prompt, + uint32_t max_new_tokens, + bool* final +); + /** * @brief Convert generated speech tokens to waveform data. */ @@ -278,6 +302,21 @@ COSYVOICE_API bool cosyvoice_token2wav( cosyvoice_generated_speech_ptr generated_speech ); +/** +* @brief Convert generated speech tokens to waveform data with additional options. +*/ +COSYVOICE_API bool cosyvoice_token2wav_ext( + cosyvoice_context_t ctx, + const int* token_ids, + uint32_t n_tokens, + float speed, + cosyvoice_prompt_t prompt, + uint32_t* speech_offset_ptr, + bool streaming, + bool finalize, + cosyvoice_generated_speech_ptr result +); + /** * @brief Run the full TTS pipeline from input tokens to waveform output. */ @@ -290,6 +329,17 @@ COSYVOICE_API bool cosyvoice_tts( cosyvoice_generated_speech_ptr result ); + +COSYVOICE_API bool cosyvoice_tts_stream( + cosyvoice_context_t ctx, + const int* text, + uint32_t text_len, + float speed, + cosyvoice_prompt_t prompt, + cosyvoice_tts_audio_callback_t callback, + void* user_data +); + // ---------------------------------------------------------------------------- // Tokenizer Operations // ---------------------------------------------------------------------------- diff --git a/include/cosyvoice.h b/include/cosyvoice.h index 684e8ed..d40ece2 100644 --- a/include/cosyvoice.h +++ b/include/cosyvoice.h @@ -608,6 +608,62 @@ COSYVOICE_API bool cosyvoice_tts_cross_lingual( cosyvoice_generated_speech_ptr result ); +/** + * @brief Callback invoked by the streaming TTS API for each generated audio chunk. + * @param audio PCM samples in 32-bit floating point format, 1 channel. + * @param n_samples Number of samples in this chunk. + * @param user_data Opaque context passed when registering the callback. + * @return True to continue streaming, false to abort the synthesis. + */ +typedef bool (*cosyvoice_tts_audio_callback_t)(const float* audio, uint32_t n_samples, void* user_data); + +/** + * @brief Generate speech with streaming output, delivering audio chunks via a callback. + * @details The function synthesizes speech incrementally and invokes @p callback for + * each chunk as it becomes available. The callback receives the PCM data + * sequentially; the previous chunk's data is no longer valid after the + * callback returns. + * @param ctx The TTS context. + * @param text Input text to synthesize. + * @param instruction Optional instruction (used in instruct mode). Pass NULL for zero-shot or cross-lingual. + * @param speed Speed multiplier (1.0 = normal). + * @param callback Callback receiving each audio chunk. + * @param user_data Opaque context passed to @p callback. + * @return True on success, false on failure. + */ +COSYVOICE_API bool cosyvoice_tts_zero_shot_stream( + cosyvoice_tts_context_t ctx, + const char* text, + float speed, + cosyvoice_tts_audio_callback_t callback, + void* user_data +); + +/** + * @brief Generate speech with streaming output in instruct mode. + * @see cosyvoice_tts_zero_shot_stream + */ +COSYVOICE_API bool cosyvoice_tts_instruct_stream( + cosyvoice_tts_context_t ctx, + const char* text, + const char* instruction, + float speed, + cosyvoice_tts_audio_callback_t callback, + void* user_data +); + +/** + * @brief Generate speech with streaming output in cross-lingual mode. + * @see cosyvoice_tts_zero_shot_stream + */ +COSYVOICE_API bool cosyvoice_tts_cross_lingual_stream( + cosyvoice_tts_context_t ctx, + const char* text, + float speed, + cosyvoice_tts_audio_callback_t callback, + void* user_data +); + // ---------------------------------------------------------------------------- // Audio Output Utilities // ---------------------------------------------------------------------------- diff --git a/src/cosyvoice-graph.cpp b/src/cosyvoice-graph.cpp index c0fc50d..30b7f3e 100644 --- a/src/cosyvoice-graph.cpp +++ b/src/cosyvoice-graph.cpp @@ -323,7 +323,7 @@ std::array AdaLayerNormZero::build_cgraph(ggml_context* ctx, gg return { x, gate_msa, shift_mlp, scale_mlp, gate_mlp }; } -ggml_tensor* Attention::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* position_ids, int64_t cut_len) const +ggml_tensor* Attention::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* position_ids, int64_t cut_len, ggml_tensor* attn_mask) const { const auto full_seq_len = position_ids->ne[0]; const auto seq_len = full_seq_len - (cut_len > 0 ? cut_len : 0); @@ -366,12 +366,11 @@ ggml_tensor* Attention::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_ten ggml_tensor* attn_output; if (fattn) - attn_output = ggml_flash_attn_ext(ctx, query, key, value, nullptr, 1.f / sqrtf(static_cast(head_dim)), 0.f, 0.f); + attn_output = ggml_flash_attn_ext(ctx, query, key, value, attn_mask, 1.f / sqrtf(static_cast(head_dim)), 0.f, 0.f); else { auto attn_scores = ggml_mul_mat(ctx, key, query); - attn_scores = ggml_scale(ctx, attn_scores, 1.f / sqrtf(static_cast(head_dim))); - auto attn_weights = ggml_soft_max(ctx, attn_scores); + auto attn_weights = ggml_soft_max_ext_inplace(ctx, attn_scores, attn_mask, 1.f / std::sqrt(static_cast(key->ne[0])), 0.f); value = ggml_permute(ctx, value, 1, 0, 2, 3); value = ggml_cont(ctx, value); attn_output = ggml_mul_mat(ctx, value, attn_weights); @@ -396,11 +395,11 @@ ggml_tensor* FeedForward::build_cgraph(ggml_context* ctx, ggml_tensor* x) const return x; } -ggml_tensor* DiTBlock::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* time_emb, ggml_tensor* position_ids, int64_t cut_len) const +ggml_tensor* DiTBlock::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* time_emb, ggml_tensor* position_ids, int64_t cut_len, ggml_tensor* attn_mask) const { auto [norm, gate_msa, shift_mlp, scale_mlp, gate_mlp] = attn_norm.build_cgraph(ctx, x, time_emb); - auto attn_output = attn.build_cgraph(ctx, norm, position_ids, cut_len); + auto attn_output = attn.build_cgraph(ctx, norm, position_ids, cut_len, attn_mask); gate_msa = unsqueeze(ctx, gate_msa, 1); attn_output = ggml_mul(ctx, attn_output, gate_msa); if (cut_len > 0) @@ -431,7 +430,7 @@ ggml_tensor* AdaLayerNorm_Final::build_cgraph(ggml_context* ctx, ggml_tensor* x, return x; } -ggml_tensor* DiT::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* mu, ggml_tensor* t, ggml_tensor* spks, ggml_tensor* cond, int64_t cut_len, ggml_tensor*& ref_position_ids, ggml_backend_op_capabilities capabilities) const +ggml_tensor* DiT::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* mu, ggml_tensor* t, ggml_tensor* spks, ggml_tensor* cond, int64_t cut_len, ggml_tensor*& ref_position_ids, ggml_backend_op_capabilities capabilities, bool streaming, ggml_tensor*& ref_attn_mask) const { x = ggml_permute(ctx, x, 1, 0, 2, 3); @@ -443,10 +442,19 @@ ggml_tensor* DiT::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* m ref_position_ids = position_ids; position_ids = ggml_dup(ctx, position_ids); + ggml_tensor* attn_mask = nullptr; + if (streaming) + { + const auto seq_len = x->ne[1]; + attn_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, seq_len, seq_len); + ggml_set_input(attn_mask); + ref_attn_mask = attn_mask; + } + for (const auto& block : std::span(transformer_blocks.cbegin(), transformer_blocks.size() - 1)) - x = block.build_cgraph(ctx, x, t, position_ids, 0); + x = block.build_cgraph(ctx, x, t, position_ids, 0, attn_mask); // Apply `cut_len` only on the final block. - x = transformer_blocks.back().build_cgraph(ctx, x, t, position_ids, cut_len); + x = transformer_blocks.back().build_cgraph(ctx, x, t, position_ids, cut_len, attn_mask); x = norm_out.build_cgraph(ctx, x, t); @@ -485,7 +493,7 @@ std::array CausalConditionalCFM::get_t_and_dt(ggml_context* ctx, int s return { t, dt }; } -ggml_tensor* CausalConditionalCFM::build_cgraph_one_step(ggml_context* ctx, const DiTContext& ditctx, int step, ggml_backend_op_capabilities capabilities, int64_t cut_len, ggml_tensor*& t_tensor, ggml_tensor*& position_ids) const +ggml_tensor* CausalConditionalCFM::build_cgraph_one_step(ggml_context* ctx, const DiTContext& ditctx, int step, ggml_backend_op_capabilities capabilities, int64_t cut_len, ggml_tensor*& t_tensor, ggml_tensor*& position_ids, bool streaming, ggml_tensor*& ref_attn_mask) const { auto x = ditctx.x; auto [t, dt] = get_t_and_dt(ctx, step); @@ -503,7 +511,9 @@ ggml_tensor* CausalConditionalCFM::build_cgraph_one_step(ggml_context* ctx, cons ditctx.cond_in, step == t_span.size() - 1 ? cut_len : 0, position_ids, - capabilities), + capabilities, + streaming, + ref_attn_mask), 2); dphi_dt = ggml_permute(ctx, dphi_dt, 1, 0, 2, 3); @@ -526,11 +536,14 @@ ggml_tensor* CausalConditionalCFM::build_cgraph_one_step(ggml_context* ctx, cons return x; } -ggml_tensor* PreLookaheadLayer::build_cgraph(ggml_context* ctx, ggml_tensor* inputs) const +ggml_tensor* PreLookaheadLayer::build_cgraph(ggml_context* ctx, ggml_tensor* inputs, bool streaming) const { auto outputs = ggml_permute(ctx, inputs, 1, 0, 2, 3); outputs = ggml_cont(ctx, outputs); - outputs = ggml_pad(ctx, outputs, pre_lookahead_len, 0, 0, 0); + if (streaming) + inputs = ggml_view_2d(ctx, inputs, inputs->ne[0], inputs->ne[1] - pre_lookahead_len, inputs->nb[1], 0); + else + outputs = ggml_pad(ctx, outputs, pre_lookahead_len, 0, 0, 0); outputs = conv1.build_cgraph(ctx, outputs, 1, 0, 1, 1, {}); outputs = ggml_leaky_relu(ctx, outputs, 0.01f, false); @@ -543,7 +556,7 @@ ggml_tensor* PreLookaheadLayer::build_cgraph(ggml_context* ctx, ggml_tensor* inp return outputs; } -CausalMaskedDiffWithDiT::EncodeResult CausalMaskedDiffWithDiT::build_cgraph_encode(ggml_context* ctx, ggml_tensor* token, ggml_tensor* prompt_token, ggml_tensor* prompt_feat, ggml_tensor* embedding, ggml_backend_op_capabilities capabilities) const +CausalMaskedDiffWithDiT::EncodeResult CausalMaskedDiffWithDiT::build_cgraph_encode(ggml_context* ctx, ggml_tensor* token, ggml_tensor* prompt_token, ggml_tensor* prompt_feat, ggml_tensor* embedding, ggml_backend_op_capabilities capabilities, bool streaming) const { embedding = ggml_l2_norm(ctx, embedding, 1e-6f); embedding = spk_embed_affine_layer.build_cgraph(ctx, embedding); @@ -551,7 +564,7 @@ CausalMaskedDiffWithDiT::EncodeResult CausalMaskedDiffWithDiT::build_cgraph_enco token = concat_tensors(ctx, std::array{ prompt_token, token }, 0, capabilities); token = ggml_get_rows(ctx, input_embedding, token); - auto h = pre_lookahead_layer.build_cgraph(ctx, token); + ggml_tensor* h = pre_lookahead_layer.build_cgraph(ctx, token, streaming); h = unsqueeze(ctx, h, 1); h = ggml_repeat_4d(ctx, h, h->ne[0], token_mel_ratio, h->ne[2], h->ne[3]); h = ggml_reshape_3d(ctx, h, h->ne[0], h->ne[1] * h->ne[2], h->ne[3]); @@ -575,14 +588,19 @@ int CausalConv1d::causal_padding() const return causal_padding; } -ggml_tensor* CausalConv1d::build_cgraph(ggml_context* ctx, ggml_tensor* x) const +ggml_tensor* CausalConv1d::build_cgraph(ggml_context* ctx, ggml_tensor* x, bool finalize) const { + GGML_ASSERT(finalize || causal_type == right); + if (!ggml_is_contiguous(x)) x = ggml_cont(ctx, x); if (causal_type == left) x = ggml_pad_ext(ctx, x, causal_padding(), 0, 0, 0, 0, 0, 0, 0); else if (causal_type == right) - x = ggml_pad_ext(ctx, x, 0, causal_padding(), 0, 0, 0, 0, 0, 0); + { + if (finalize) + x = ggml_pad_ext(ctx, x, 0, causal_padding(), 0, 0, 0, 0, 0, 0); + } else GGML_ABORT("CausalConv1d: invalid causal type %d", static_cast(causal_type)); @@ -599,9 +617,9 @@ CausalConvRNNF0Predictor::CausalConvRNNF0Predictor() condnet_8.causal_type = CausalConv1d::left; } -ggml_tensor* CausalConvRNNF0Predictor::build_cgraph(ggml_context* ctx, ggml_tensor* x) const +ggml_tensor* CausalConvRNNF0Predictor::build_cgraph(ggml_context* ctx, ggml_tensor* x, bool finalize) const { - x = condnet_0.build_cgraph(ctx, x); + x = condnet_0.build_cgraph(ctx, x, finalize); x = ggml_elu(ctx, x); x = condnet_2.build_cgraph(ctx, x); x = ggml_elu(ctx, x); @@ -707,9 +725,9 @@ ggml_tensor* CausalConv1dUpsample::build_cgraph(ggml_context* ctx, ggml_tensor* return Conv1d::build_cgraph(ctx, x, 1, 0, 1, 1, {}); } -std::array CausalHiFTGenerator::build_cgraph(ggml_context* ctx, ggml_tensor* speech_feat) const +std::array CausalHiFTGenerator::build_cgraph(ggml_context* ctx, ggml_tensor* speech_feat, bool finalize) const { - auto f0 = f0_predictor.build_cgraph(ctx, speech_feat); + auto f0 = f0_predictor.build_cgraph(ctx, speech_feat, finalize); auto s_input = ggml_permute(ctx, f0, 1, 0, 2, 3); s_input = ggml_interpolate(ctx, s_input, s_input->ne[0], s_input->ne[1] * scale_factor, s_input->ne[2], 1, GGML_SCALE_MODE_NEAREST); auto [s, noise] = m_source.build_cgraph(ctx, s_input, nb_harmonics, sampling_rate, scale_factor, nsf_alpha, nsf_voiced_threshold, nsf_sigma); @@ -718,11 +736,15 @@ std::array CausalHiFTGenerator::build_cgraph(ggml_context* ctx, // decode ggml_tensor* x; { - x = conv_pre.build_cgraph(ctx, speech_feat); + if (!finalize) + speech_feat = ggml_view_2d(ctx, speech_feat, speech_feat->ne[0] - f0_predictor.condnet_0.causal_padding(), speech_feat->ne[1], speech_feat->nb[1], 0); + x = conv_pre.build_cgraph(ctx, speech_feat, finalize); auto s_stft = ggml_stft(ctx, s, window, hop_len, true, fctx.get()); s_stft = ggml_cont(ctx, s_stft); s_stft = ggml_reshape_2d(ctx, s_stft, s_stft->ne[0], s_stft->ne[1] * 2); + if (!finalize) + s_stft = ggml_view_2d(ctx, s_stft, s_stft->ne[0] - (scale_factor / hop_len * conv_pre_look_right), s_stft->ne[1], s_stft->nb[1], 0); const auto num_upsamples = ups.size(); const auto num_kernels = resblocks.size() / num_upsamples; @@ -761,6 +783,8 @@ std::array CausalHiFTGenerator::build_cgraph(ggml_context* ctx, ggml_sin(ctx, phase)); x = ggml_istft(ctx, real, imag, window, hop_len, true, ictx.get()); + if (!finalize) + x = ggml_view_2d(ctx, x, x->ne[0] - scale_factor, x->ne[1], x->nb[1], 0); x = ggml_clamp(ctx, x, -audio_limit, audio_limit); ggml_set_cpu(x); return { x, noise }; diff --git a/src/cosyvoice-model.cpp b/src/cosyvoice-model.cpp index b36f218..d60633d 100644 --- a/src/cosyvoice-model.cpp +++ b/src/cosyvoice-model.cpp @@ -568,3 +568,21 @@ bool cosyvoice_model::llm_set_kv_cache_len(uint32_t len) } return false; } + +void cosyvoice_model::llm_offload_kv_cache() +{ + auto sched = worker->sched.get(); + auto& kv_cache = worker->kv_cache; + ggml_backend_sched_reset(sched); + kv_cache.offload_cache(worker->backend.get(), sched, kv_cache.cur_len); + ggml_backend_sched_synchronize(sched); +} + +void cosyvoice_model::llm_load_kv_cache() +{ + auto sched = worker->sched.get(); + auto& kv_cache = worker->kv_cache; + ggml_backend_sched_reset(sched); + kv_cache.offload_cache(worker->backend.get(), sched, kv_cache.cur_len); + ggml_backend_sched_synchronize(sched); +} diff --git a/src/cosyvoice-model.h b/src/cosyvoice-model.h index 28b6025..7e2ad42 100644 --- a/src/cosyvoice-model.h +++ b/src/cosyvoice-model.h @@ -90,6 +90,8 @@ struct cosyvoice_model : virtual cosyvoice_model_context, virtual cosyvoice_obje uint32_t llm_get_kv_cache_len(); bool llm_set_kv_cache_len(uint32_t len); + void llm_offload_kv_cache(); + void llm_load_kv_cache(); int llm_sample_token(); void llm_accept_token(int token); @@ -180,7 +182,10 @@ struct cosyvoice_model_3 : cosyvoice_model void set_hift_rand_ini(const float* data); bool llm_job(const int* text, uint32_t text_len, cosyvoice_prompt_t prompt); + bool llm_job_ext(const int* text, uint32_t text_len, cosyvoice_prompt_t prompt, uint32_t max_new_tokens, bool* final); bool token2wav(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, cosyvoice_generated_speech_ptr result); + bool token2wav_ext(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* speech_offset, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result); + uint32_t get_chunk_tokens(); void empty_buffer_cache(); void get_memory_usage(cosyvoice_memory_usage_t* usage); diff --git a/src/cosyvoice-modules.h b/src/cosyvoice-modules.h index 99da835..286713b 100644 --- a/src/cosyvoice-modules.h +++ b/src/cosyvoice-modules.h @@ -130,7 +130,7 @@ struct Attention : Module void OnLoad(const gguf_loader& loader, const std::string& prefix); - ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* position_ids, int64_t cut_len) const; + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* position_ids, int64_t cut_len, ggml_tensor* attn_mask = nullptr) const; }; struct FeedForward : Module @@ -153,7 +153,7 @@ struct DiTBlock : Module void OnLoad(const gguf_loader& loader, const std::string& prefix); - ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* time_emb, ggml_tensor* position_ids, int64_t cut_len) const; + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* time_emb, ggml_tensor* position_ids, int64_t cut_len, ggml_tensor* attn_mask = nullptr) const; }; struct AdaLayerNorm_Final : Module @@ -178,10 +178,11 @@ struct DiT : Module Linear proj_out; int mel_dim; + constexpr static int static_chunk_size = 50; void OnLoad(const gguf_loader& loader, const std::string& prefix); - ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* mu, ggml_tensor* t, ggml_tensor* spks, ggml_tensor* cond, int64_t cut_len, ggml_tensor*& position_ids, ggml_backend_op_capabilities capabilities) const; + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* mu, ggml_tensor* t, ggml_tensor* spks, ggml_tensor* cond, int64_t cut_len, ggml_tensor*& position_ids, ggml_backend_op_capabilities capabilities, bool streaming, ggml_tensor*& ref_attn_mask) const; }; struct CausalConditionalCFM : Module @@ -204,7 +205,7 @@ struct CausalConditionalCFM : Module DiTContext prepare_context(ggml_context* ctx, ggml_tensor* mu, ggml_tensor* spks, ggml_tensor* cond) const; std::array get_t_and_dt(ggml_context* ctx, int step) const; - ggml_tensor* build_cgraph_one_step(ggml_context* ctx, const DiTContext& ditctx, int step, ggml_backend_op_capabilities capabilities, int64_t cut_len, ggml_tensor*& t_tensor, ggml_tensor*& position_ids) const; + ggml_tensor* build_cgraph_one_step(ggml_context* ctx, const DiTContext& ditctx, int step, ggml_backend_op_capabilities capabilities, int64_t cut_len, ggml_tensor*& t_tensor, ggml_tensor*& position_ids, bool streaming, ggml_tensor*& ref_attn_mask) const; }; struct PreLookaheadLayer : Module @@ -216,7 +217,7 @@ struct PreLookaheadLayer : Module void OnLoad(const gguf_loader& loader, const std::string& prefix); - ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* inputs) const; + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* inputs, bool streaming) const; }; struct CausalMaskedDiffWithDiT : Module @@ -238,7 +239,7 @@ struct CausalMaskedDiffWithDiT : Module int64_t cut_len; }; - EncodeResult build_cgraph_encode(ggml_context* ctx, ggml_tensor* token, ggml_tensor* prompt_token, ggml_tensor* prompt_feat, ggml_tensor* embedding, ggml_backend_op_capabilities capabilities) const; + EncodeResult build_cgraph_encode(ggml_context* ctx, ggml_tensor* token, ggml_tensor* prompt_token, ggml_tensor* prompt_feat, ggml_tensor* embedding, ggml_backend_op_capabilities capabilities, bool streaming = false) const; }; struct CausalConv1dBase : Conv1d @@ -258,7 +259,8 @@ struct CausalConv1d : CausalConv1dBase int causal_padding() const; - ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x) const; + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x) const { return build_cgraph(ctx, x, true); } + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, bool finalize) const; }; struct CausalConvRNNF0Predictor : Module @@ -275,7 +277,7 @@ struct CausalConvRNNF0Predictor : Module Linear classifier; - ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x) const; + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, bool finalize) const; }; struct SineGen2 : Module @@ -339,10 +341,11 @@ struct CausalHiFTGenerator : Module std::vector resblocks; CausalConv1d conv_post; - std::array build_cgraph(ggml_context* ctx, ggml_tensor* speech_feat) const; + std::array build_cgraph(ggml_context* ctx, ggml_tensor* speech_feat, bool finalize) const; void set_rand_ini(const float* data) const; + constexpr static int conv_pre_look_right = 4; float lrelu_slope; int scale_factor; int nb_harmonics; diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index c11b141..717f094 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -5,7 +5,6 @@ #include "ggml-cpu-flag.h" #include -#include #include #include #include @@ -17,6 +16,12 @@ constexpr uint32_t COSYVOICE_TTS_FLAG_MASK = COSYVOICE_TTS_FLAG_TEXT_NORMALIZATI constexpr double COSYVOICE_TTS_FADE_IN_SECONDS = 0.02; bool cosyvoice_model_3::llm_job(const int* text, uint32_t text_len, cosyvoice_prompt_t prompt) +{ + llm_clear_accepted_tokens(); + return llm_job_ext(text, text_len, prompt, UINT32_MAX, nullptr); +} + +bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoice_prompt_t prompt, uint32_t max_new_tokens, bool* final) { const auto& params = shared->params; auto& sched = worker->sched; @@ -24,9 +29,7 @@ bool cosyvoice_model_3::llm_job(const int* text, uint32_t text_len, cosyvoice_pr auto& batch_buffer = worker->batch_buffer; auto& prompt_crc32 = worker->prompt_crc32; auto last_prompt_crc32 = prompt_crc32; - - if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) - ggml_backend_sched_synchronize(sched.get()); + bool stop_reached = false; try { @@ -37,103 +40,117 @@ bool cosyvoice_model_3::llm_job(const int* text, uint32_t text_len, cosyvoice_pr const auto token_emb = reinterpret_cast(llm.embed_tokens_weight->data); const char* cur; uint32_t offset = 0; - if (speech_type == llm.embed_tokens_weight->type) + + if (text != nullptr) { - auto prefill_embedding = [&](const char* data, int token_id) + if (speech_type == llm.embed_tokens_weight->type) { - if (offset == n_batch) + auto prefill_embedding = [&](const char* data, int token_id) { - if (!llm_prefill(speech_type, batch_buffer.get(), n_batch)) - throw std::runtime_error("Failed to prefill LLM KV cache.\n"); - offset = 0; - } + if (offset == n_batch) + { + if (!llm_prefill(speech_type, batch_buffer.get(), n_batch)) + throw std::runtime_error("Failed to prefill LLM KV cache.\n"); + offset = 0; + } - memcpy(batch_buffer.get() + offset++ * speech_row_size, data + token_id * speech_row_size, speech_row_size); - }; + memcpy(batch_buffer.get() + offset++ * speech_row_size, data + token_id * speech_row_size, speech_row_size); + }; - if (llm_get_kv_cache_len() == 0) - { - prefill_embedding(speech_emb, llm.sos_token_id); - prompt_crc32 = 0; - } - // The first token is assumed to be the SOS token already stored in the KV cache. - if (prompt_crc32 != prompt->prompt_crc32) - { - llm_set_kv_cache_len(1); - for (const auto& i : prompt->prompt_text) - prefill_embedding(token_emb, i); - prompt_crc32 = prompt->prompt_crc32; - } - else llm_set_kv_cache_len(1 + static_cast(prompt->prompt_text.size())); + if (llm_get_kv_cache_len() == 0) + { + prefill_embedding(speech_emb, llm.sos_token_id); + prompt_crc32 = 0; + } + // The first token is assumed to be the SOS token already stored in the KV cache. + if (prompt_crc32 != prompt->prompt_crc32) + { + llm_set_kv_cache_len(1); + for (const auto& i : prompt->prompt_text) + prefill_embedding(token_emb, i); + prompt_crc32 = prompt->prompt_crc32; + } + else llm_set_kv_cache_len(1 + static_cast(prompt->prompt_text.size())); - for (uint32_t i = 0; i != text_len; ++i) - prefill_embedding(token_emb, text[i]); + for (uint32_t i = 0; i != text_len; ++i) + prefill_embedding(token_emb, text[i]); - if (prompt->llm_prompt_speech_tokens.second != 0) - { - prefill_embedding(speech_emb, llm.task_token_id); + if (prompt->llm_prompt_speech_tokens.second != 0) + { + prefill_embedding(speech_emb, llm.task_token_id); - const auto end = prompt->llm_prompt_speech_tokens.second - 1; - for (uint32_t i = 0; i != end; ++i) - prefill_embedding(speech_emb, prompt->llm_prompt_speech_tokens.first[i]); - cur = speech_emb + prompt->llm_prompt_speech_tokens.first[end] * speech_row_size; + const auto end = prompt->llm_prompt_speech_tokens.second - 1; + for (uint32_t i = 0; i != end; ++i) + prefill_embedding(speech_emb, prompt->llm_prompt_speech_tokens.first[i]); + cur = speech_emb + prompt->llm_prompt_speech_tokens.first[end] * speech_row_size; + } + else cur = speech_emb + llm.task_token_id * speech_row_size; } - else cur = speech_emb + llm.task_token_id * speech_row_size; - } - else - { - const auto token_type = llm.embed_tokens_weight->type; - const auto token_row_size = static_cast(llm.embed_tokens_weight->nb[1]); - - auto prefill_embedding = [&](const char* data, int token_id, uint32_t row_size, ggml_type type) + else { - if (offset == n_batch) + const auto token_type = llm.embed_tokens_weight->type; + const auto token_row_size = static_cast(llm.embed_tokens_weight->nb[1]); + + auto prefill_embedding = [&](const char* data, int token_id, uint32_t row_size, ggml_type type) { - if (!llm_prefill(type, batch_buffer.get(), n_batch)) - throw std::runtime_error("Failed to prefill LLM KV cache.\n"); - offset = 0; + if (offset == n_batch) + { + if (!llm_prefill(type, batch_buffer.get(), n_batch)) + throw std::runtime_error("Failed to prefill LLM KV cache.\n"); + offset = 0; + } + + memcpy(batch_buffer.get() + offset++ * row_size, data + token_id * row_size, row_size); + }; + + if (llm_get_kv_cache_len() == 0) + llm_prefill(speech_type, speech_emb + llm.sos_token_id * speech_row_size, 1); + // The first token is assumed to be the SOS token already stored in the KV cache. + if (prompt_crc32 != prompt->prompt_crc32) + { + llm_set_kv_cache_len(1); + for (const auto& i : prompt->prompt_text) + prefill_embedding(token_emb, i, token_row_size, token_type); + prompt_crc32 = prompt->prompt_crc32; } + else llm_set_kv_cache_len(1 + static_cast(prompt->prompt_text.size())); - memcpy(batch_buffer.get() + offset++ * row_size, data + token_id * row_size, row_size); - }; + for (uint32_t i = 0; i != text_len; ++i) + prefill_embedding(token_emb, text[i], token_row_size, token_type); - if (llm_get_kv_cache_len() == 0) - llm_prefill(speech_type, speech_emb + llm.sos_token_id * speech_row_size, 1); - // The first token is assumed to be the SOS token already stored in the KV cache. - if (prompt_crc32 != prompt->prompt_crc32) - { - llm_set_kv_cache_len(1); - for (const auto& i : prompt->prompt_text) - prefill_embedding(token_emb, i, token_row_size, token_type); - prompt_crc32 = prompt->prompt_crc32; - } - else llm_set_kv_cache_len(1 + static_cast(prompt->prompt_text.size())); - - for (uint32_t i = 0; i != text_len; ++i) - prefill_embedding(token_emb, text[i], token_row_size, token_type); - - if (offset != 0 && !llm_prefill(token_type, batch_buffer.get(), offset)) - throw std::runtime_error("Failed to prefill LLM KV cache.\n"); - offset = 0; + if (offset != 0 && !llm_prefill(token_type, batch_buffer.get(), offset)) + throw std::runtime_error("Failed to prefill LLM KV cache.\n"); + offset = 0; - if (prompt->llm_prompt_speech_tokens.second != 0) - { - prefill_embedding(speech_emb, llm.task_token_id, speech_row_size, speech_type); + if (prompt->llm_prompt_speech_tokens.second != 0) + { + prefill_embedding(speech_emb, llm.task_token_id, speech_row_size, speech_type); - const auto end = prompt->llm_prompt_speech_tokens.second - 1; - for (uint32_t i = 0; i != end; ++i) - prefill_embedding(speech_emb, prompt->llm_prompt_speech_tokens.first[i], speech_row_size, speech_type); - cur = speech_emb + prompt->llm_prompt_speech_tokens.first[end] * speech_row_size; + const auto end = prompt->llm_prompt_speech_tokens.second - 1; + for (uint32_t i = 0; i != end; ++i) + prefill_embedding(speech_emb, prompt->llm_prompt_speech_tokens.first[i], speech_row_size, speech_type); + cur = speech_emb + prompt->llm_prompt_speech_tokens.first[end] * speech_row_size; + } + else cur = speech_emb + llm.task_token_id * speech_row_size; } - else cur = speech_emb + llm.task_token_id * speech_row_size; - } - if (offset != 0 && !llm_prefill(speech_type, batch_buffer.get(), offset)) - throw std::runtime_error("Failed to prefill LLM KV cache.\n"); + if (offset != 0 && !llm_prefill(speech_type, batch_buffer.get(), offset)) + throw std::runtime_error("Failed to prefill LLM KV cache.\n"); + } + else + { + // Continue from existing KV cache — get last accepted token's embedding + const auto n_acc = llm_get_n_accepted_tokens(); + if (n_acc == 0) + throw std::runtime_error("llm_job_ext: cannot continue generation with empty token history.\n"); + const auto last_token_id = llm_get_accepted_tokens()[n_acc - 1]; + cur = speech_emb + last_token_id * speech_row_size; + } + // First call: min/max from text input const auto min_len = static_cast(text_len * worker->config.min_token_text_ratio); const auto max_len = static_cast(text_len * worker->config.max_token_text_ratio); - llm_clear_accepted_tokens(); + const auto limit = std::min(max_new_tokens, max_len); // FSQ silent/breath token filtering: allow up to 5 consecutive silent tokens, // then drop any further consecutive ones to avoid generating long silence. @@ -142,7 +159,7 @@ bool cosyvoice_model_3::llm_job(const int* text, uint32_t text_len, cosyvoice_pr constexpr uint32_t max_silent_token_num = 5; const auto& silent_tokens = cv3_shared->silent_tokens; - for (uint32_t n = 0; n != max_len; ++n) + for (uint32_t n = 0; n != limit; ++n) { if (!llm_decode(speech_type, cur)) throw std::runtime_error("Failed to decode LLM output.\n"); @@ -152,7 +169,10 @@ bool cosyvoice_model_3::llm_job(const int* text, uint32_t text_len, cosyvoice_pr if (token_id == -1) throw std::runtime_error("Failed to sample token from LLM output. This might be wrong with the model or caused by an issue with the sampling parameters.\n"); if (n > min_len && llm_is_stop_token(token_id)) + { + stop_reached = true; break; + } if (silent_tokens.count(token_id)) { @@ -176,18 +196,25 @@ bool cosyvoice_model_3::llm_job(const int* text, uint32_t text_len, cosyvoice_pr } catch (const std::exception& e) { + worker->llm_input = nullptr; cosyvoice_call_ggml_log_callback(GGML_LOG_LEVEL_ERROR, e.what()); if (params.builtin_sampler_rng_policy == COSYVOICE_BUILTIN_SAMPLER_RNG_POLICY_RESET_PER_SESSION) reset_builtin_sampler_rng(); return false; } - if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED && last_prompt_crc32 != prompt_crc32) + worker->llm_input = nullptr; + if (final) + *final = stop_reached; + + if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED && text != nullptr && max_new_tokens == UINT32_MAX && last_prompt_crc32 != prompt_crc32) { - ggml_backend_sched_reset(sched.get()); - worker->kv_cache.offload_cache(worker->backend.get(), sched.get(), 1 + static_cast(prompt->prompt_text.size())); + llm_set_kv_cache_len(1 + static_cast(prompt->prompt_text.size())); + llm_offload_kv_cache(); } - if (params.builtin_sampler_rng_policy == COSYVOICE_BUILTIN_SAMPLER_RNG_POLICY_RESET_PER_SESSION) + + // Reset RNG on stop (matches non-streaming behavior) + if (stop_reached && params.builtin_sampler_rng_policy == COSYVOICE_BUILTIN_SAMPLER_RNG_POLICY_RESET_PER_SESSION) reset_builtin_sampler_rng(); return true; @@ -280,6 +307,12 @@ static void set_graph_backends(ggml_cgraph* gf, ggml_backend_sched_t sched, ggml bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, cosyvoice_generated_speech_ptr result) { + return token2wav_ext(token_ids, n_tokens, speed, prompt, nullptr, false, true, result); +} + +bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* speech_offset_ptr, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result) +{ + const uint32_t speech_offset = speech_offset_ptr ? *speech_offset_ptr : 0; const auto& params = shared->params; auto& sched = worker->sched; auto& flow = cv3_shared->flow; @@ -295,8 +328,6 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float ggml_reset(ctx0.get()); ggml_reset(ctx1.get()); - if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) - ggml_backend_sched_synchronize(sched.get()); ggml_backend_sched_reset(sched.get()); ggml_tensor* token = ggml_new_tensor_1d(ctx0.get(), GGML_TYPE_I32, n_tokens); @@ -304,8 +335,9 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float ggml_tensor* prompt_feat = ggml_new_tensor_2d(ctx0.get(), GGML_TYPE_F32, prompt->prompt_speech_feat.shape[1], prompt->prompt_speech_feat.shape[0]); ggml_tensor* embedding = ggml_new_tensor_2d(ctx0.get(), GGML_TYPE_F32, prompt->flow_embedding.shape[1], prompt->flow_embedding.shape[0]); + // Phase 1: Flow encode ggml_cgraph* gf = new_cgraph(ctx0.get()); - auto [mu, spks, conds, cut_len] = flow.build_cgraph_encode(ctx0.get(), token, prompt_token, prompt_feat, embedding, op_caps); + auto [mu, spks, conds, cut_len] = flow.build_cgraph_encode(ctx0.get(), token, prompt_token, prompt_feat, embedding, op_caps, streaming); auto ditctx = flow.decoder.prepare_context(ctx1.get(), mu, spks, conds); do { @@ -333,13 +365,16 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float float* noise_buffer = shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_BEFORE_FLOW, noise_len, nullptr, shared->noise_callback_ctx); ggml_backend_tensor_set_async(backend.get(), ditctx.x, noise_buffer, 0, ggml_nbytes(ditctx.x)); + // Phase 2: Flow decode steps ggml_tensor* t_leaf; ggml_tensor* position_ids; - auto feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, 1, op_caps, cut_len, t_leaf, position_ids); + ggml_tensor* attn_mask = nullptr; + auto feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, 1, op_caps, cut_len, t_leaf, position_ids, streaming, attn_mask); ggml_build_forward_expand(gf, feat); set_graph_backends(gf, sched.get(), backend.get(), cpu_backend.get(), op_caps); ggml_backend_sched_synchronize(sched.get()); ggml_backend_sched_alloc_graph(sched.get(), gf); + if (!op_caps.fill) ggml_set_zero(t_leaf); for (int64_t i = 0; i < position_ids->ne[1]; ++i) { @@ -347,6 +382,17 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float for (int32_t j = 0; j < position_ids->ne[0]; ++j) cur_row[j] = j; } + if (attn_mask) + { + const auto cs = flow.decoder.estimator.static_chunk_size; + for (int64_t i = 0; i < attn_mask->ne[0]; ++i) + { + int64_t block_end = ((i / cs) + 1) * cs; + auto row = reinterpret_cast(attn_mask->data) + i * attn_mask->ne[1]; + for (int64_t j = 0; j < attn_mask->ne[1]; ++j) + row[j] = j < block_end ? 0 : 0xFC00; + } + } ggml_backend_tensor_set_async(backend.get(), token, token_ids, 0, token->nb[1]); ggml_backend_tensor_set_async(backend.get(), prompt_token, prompt->flow_prompt_speech_tokens.first.get(), 0, prompt_token->nb[1]); @@ -356,7 +402,7 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_AFTER_FLOW, noise_len, noise_buffer, shared->noise_callback_ctx); if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_SHARED) - llm_set_kv_cache_len(0); // Reset the visible KV length when sharing the buffer. + llm_set_kv_cache_len(0); if (worker->status != GGML_STATUS_SUCCESS) return false; @@ -373,7 +419,8 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float ggml_reset(ctx0.get()); ggml_backend_sched_reset(sched.get()); gf = new_cgraph(ctx0.get()); - feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, 2, op_caps, cut_len, t_leaf, position_ids); + attn_mask = nullptr; + feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, 2, op_caps, cut_len, t_leaf, position_ids, streaming, attn_mask); ggml_build_forward_expand(gf, feat); set_graph_backends(gf, sched.get(), backend.get(), cpu_backend.get(), op_caps); @@ -386,6 +433,17 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float for (int32_t j = 0; j < position_ids->ne[0]; ++j) cur_row[j] = j; } + if (attn_mask) + { + const auto cs = flow.decoder.estimator.static_chunk_size; + for (int64_t i = 0; i < attn_mask->ne[0]; ++i) + { + int64_t block_end = ((i / cs) + 1) * cs; + auto row = reinterpret_cast(attn_mask->data) + i * attn_mask->ne[1]; + for (int64_t j = 0; j < attn_mask->ne[1]; ++j) + row[j] = j < block_end ? 0 : 0xFC00; + } + } worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); if (worker->status != GGML_STATUS_SUCCESS) return false; @@ -409,7 +467,8 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float ggml_reset(ctx0.get()); ggml_backend_sched_reset(sched.get()); gf = new_cgraph(ctx0.get()); - feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, static_cast(flow.decoder.t_span.size() - 1), op_caps, cut_len, t_leaf, position_ids); + attn_mask = nullptr; + feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, static_cast(flow.decoder.t_span.size() - 1), op_caps, cut_len, t_leaf, position_ids, streaming, attn_mask); ggml_build_forward_expand(gf, feat); set_graph_backends(gf, sched.get(), backend.get(), cpu_backend.get(), op_caps); @@ -422,12 +481,25 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float for (int32_t j = 0; j < position_ids->ne[0]; ++j) cur_row[j] = j; } + if (attn_mask) + { + const auto cs = flow.decoder.estimator.static_chunk_size; + for (int64_t i = 0; i < attn_mask->ne[0]; ++i) + { + int64_t block_end = ((i / cs) + 1) * cs; + auto row = reinterpret_cast(attn_mask->data) + i * attn_mask->ne[1]; + for (int64_t j = 0; j < attn_mask->ne[1]; ++j) + row[j] = j < block_end ? 0 : 0xFC00; + } + } worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); if (worker->status != GGML_STATUS_SUCCESS) return false; + // Phase 3: Copy flow output to speech_feat (persistent buffer) ggml_reset(ctx1.get()); ggml_tensor* speech_feat = ggml_new_tensor(ctx1.get(), feat->type, GGML_MAX_DIMS, feat->ne); - ggml_backend_tensor_alloc(token2wav_buffer.get(), speech_feat, ggml_backend_buffer_get_base(token2wav_buffer.get())); + ggml_backend_tensor_alloc(token2wav_buffer.get(), speech_feat, + ggml_backend_buffer_get_base(token2wav_buffer.get())); ggml_backend_tensor_copy_async(backend.get(), backend.get(), feat, speech_feat); ggml_reset(ctx0.get()); @@ -439,7 +511,7 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float ne0, speech_feat->ne[1], speech_feat->ne[2], speech_feat->ne[3], GGML_SCALE_MODE_BILINEAR); - auto [generated_speech, noise] = hift.build_cgraph(ctx0.get(), speech_feat); + auto [generated_speech, noise] = hift.build_cgraph(ctx0.get(), speech_feat, finalize); ggml_build_forward_expand(gf, generated_speech); set_graph_backends(gf, sched.get(), backend.get(), cpu_backend.get(), op_caps); @@ -453,10 +525,15 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float noise_buffer = shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_BEFORE_HIFT, noise_len, nullptr, shared->noise_callback_ctx); ggml_backend_tensor_set_async(backend.get(), noise, noise_buffer, 0, noise->nb[2]); - result->data = reinterpret_cast(generated_speech->data); - result->length = static_cast(generated_speech->ne[0]); + const auto total_speech_len = static_cast(generated_speech->ne[0]); + result->data = reinterpret_cast(generated_speech->data) + speech_offset; + result->length = speech_offset < total_speech_len ? total_speech_len - speech_offset : 0; worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_AFTER_HIFT, noise_len, noise_buffer, shared->noise_callback_ctx); + + if (worker->status == GGML_STATUS_SUCCESS && speech_offset_ptr) + *speech_offset_ptr += result->length; + if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) { ggml_backend_sched_reset(sched.get()); @@ -465,6 +542,11 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float return worker->status == GGML_STATUS_SUCCESS; } +uint32_t cosyvoice_model_3::get_chunk_tokens() +{ + return 25 + cv3_shared->flow.pre_lookahead_layer.pre_lookahead_len; +} + struct cosyvoice_tts_context : cosyvoice_tokenization_result_impl, cosyvoice_prompt, std::string { cosyvoice_tts_context(cosyvoice_context_t ctx, cosyvoice_prompt_t prompt) @@ -486,7 +568,7 @@ struct cosyvoice_tts_context : cosyvoice_tokenization_result_impl, cosyvoice_pro else prefix_len = 0; } - bool tts_job(const char* text, const char* instruction, float speed, cosyvoice_inference_mode mode, cosyvoice_generated_speech_ptr result) + bool tts_job(const char* text, const char* instruction, float speed, cosyvoice_inference_mode mode, cosyvoice_generated_speech_ptr result, cosyvoice_tts_audio_callback_t callback, void* user_data) { instruction_cache.resize(prefix_len); if (instruction) @@ -540,7 +622,8 @@ struct cosyvoice_tts_context : cosyvoice_tokenization_result_impl, cosyvoice_pro if (params.n_max_seq <= o + 1u || !(r > 0.0f)) { ctx->tokenize(effective_text, this, true); - return cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, result); + return result ? cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, result) + : cosyvoice_tts_stream_with_postprocess(get_tokens(), get_n_tokens(), speed, callback, user_data); } const auto max_text_tokens = static_cast( static_cast(params.n_max_seq - o - 1u) / (1.0f + r)); @@ -548,7 +631,8 @@ struct cosyvoice_tts_context : cosyvoice_tokenization_result_impl, cosyvoice_pro if (max_text_tokens == 0) { ctx->tokenize(effective_text, this, true); - return cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, result); + return result ? cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, result) + : cosyvoice_tts_stream_with_postprocess(get_tokens(), get_n_tokens(), speed, callback, user_data); } // Fast-split path: tokenize each fragment once and merge token-ID vectors, @@ -573,26 +657,37 @@ struct cosyvoice_tts_context : cosyvoice_tokenization_result_impl, cosyvoice_pro { const auto& t = chunk_token_list.empty() ? fragment_tokens[0] : chunk_token_list[0]; tokens = t; - return cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, result); + return result ? cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, result) + : cosyvoice_tts_stream_with_postprocess(get_tokens(), get_n_tokens(), speed, callback, user_data); + } + + if (result) + { + combined_pcm.clear(); + for (auto& chunk_tokens : chunk_token_list) + { + tokens = std::move(chunk_tokens); + cosyvoice_generated_speech part = {}; + if (!cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, &part) + || !part.data || part.length == 0) + { + result->data = nullptr; + result->length = 0; + return false; + } + combined_pcm.insert(combined_pcm.end(), part.data, part.data + part.length); + } + result->data = combined_pcm.data(); + result->length = static_cast(combined_pcm.size()); + return true; } - combined_pcm.clear(); for (auto& chunk_tokens : chunk_token_list) { tokens = std::move(chunk_tokens); - cosyvoice_generated_speech part = {}; - if (!cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, &part) - || !part.data || part.length == 0) - { - result->data = nullptr; - result->length = 0; + if (!cosyvoice_tts_stream_with_postprocess(get_tokens(), get_n_tokens(), speed, callback, user_data)) return false; - } - combined_pcm.insert(combined_pcm.end(), part.data, part.data + part.length); } - result->data = combined_pcm.data(); - result->length = static_cast(combined_pcm.size()); - return true; } // Slow-split path: tokenize each fragment only for counting, @@ -635,6 +730,42 @@ struct cosyvoice_tts_context : cosyvoice_tokenization_result_impl, cosyvoice_pro return true; } + bool cosyvoice_tts_stream_with_postprocess(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_tts_audio_callback_t callback, void* user_data) + { + bool ok; + if (flags & COSYVOICE_TTS_FLAG_FADE_IN) + ok = cosyvoice_tts_stream(ctx, token_ids, n_tokens, speed, this, callback, user_data); + else + { + struct callback_wrapper_context + { + cosyvoice_tts_audio_callback_t callback; + void* user_data; + uint32_t fade_samples; + uint32_t processed_samples = 0; + } cb_ctx{ callback, user_data, fade_samples }; + + auto callback_wrapper = [](const float* data, uint32_t length, void* user_data) -> bool + { + auto ctx = static_cast(user_data); + if (ctx->processed_samples < ctx->fade_samples) + { + const uint32_t fade_len = std::min(ctx->fade_samples - ctx->processed_samples, length); + for (uint32_t i = 0; i < fade_len; ++i) + { + const float ramp = static_cast(ctx->processed_samples + i + 1) / static_cast(ctx->fade_samples + 1); + const_cast(data)[i] *= ramp; + } + ctx->processed_samples += fade_len; + } + return ctx->callback(data, length, ctx->user_data); + }; + ok = cosyvoice_tts_stream(ctx, token_ids, n_tokens, speed, this, callback_wrapper, &cb_ctx); + } + + return ok; + } + bool cosyvoice_tts_with_postprocess(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_generated_speech_ptr result) { if (!cosyvoice_tts(ctx, token_ids, n_tokens, speed, this, result)) @@ -735,15 +866,30 @@ uint32_t cosyvoice_tts_context_set_flags(cosyvoice_tts_context_t ctx, uint32_t f bool cosyvoice_tts_zero_shot(cosyvoice_tts_context_t ctx, const char* text, float speed, cosyvoice_generated_speech_ptr result) { - return ctx->tts_job(text, nullptr, speed, COSYVOICE_INFERENCE_MODE_ZERO_SHOT, result); + return ctx->tts_job(text, nullptr, speed, COSYVOICE_INFERENCE_MODE_ZERO_SHOT, result, nullptr, nullptr); } bool cosyvoice_tts_instruct(cosyvoice_tts_context_t ctx, const char* text, const char* instruction, float speed, cosyvoice_generated_speech_ptr result) { - return ctx->tts_job(text, instruction, speed, COSYVOICE_INFERENCE_MODE_INSTRUCT, result); + return ctx->tts_job(text, instruction, speed, COSYVOICE_INFERENCE_MODE_INSTRUCT, result, nullptr, nullptr); } bool cosyvoice_tts_cross_lingual(cosyvoice_tts_context_t ctx, const char* text, float speed, cosyvoice_generated_speech_ptr result) { - return ctx->tts_job(text, nullptr, speed, COSYVOICE_INFERENCE_MODE_CROSS_LINGUAL, result); + return ctx->tts_job(text, nullptr, speed, COSYVOICE_INFERENCE_MODE_CROSS_LINGUAL, result, nullptr, nullptr); +} + +bool cosyvoice_tts_zero_shot_stream(cosyvoice_tts_context_t ctx, const char* text, float speed, cosyvoice_tts_audio_callback_t callback, void* user_data) +{ + return ctx->tts_job(text, nullptr, speed, COSYVOICE_INFERENCE_MODE_ZERO_SHOT, nullptr, callback, user_data); +} + +bool cosyvoice_tts_instruct_stream(cosyvoice_tts_context_t ctx, const char* text, const char* instruction, float speed, cosyvoice_tts_audio_callback_t callback, void* user_data) +{ + return ctx->tts_job(text, instruction, speed, COSYVOICE_INFERENCE_MODE_INSTRUCT, nullptr, callback, user_data); +} + +bool cosyvoice_tts_cross_lingual_stream(cosyvoice_tts_context_t ctx, const char* text, float speed, cosyvoice_tts_audio_callback_t callback, void* user_data) +{ + return ctx->tts_job(text, nullptr, speed, COSYVOICE_INFERENCE_MODE_CROSS_LINGUAL, nullptr, callback, user_data); } diff --git a/src/cosyvoice.cpp b/src/cosyvoice.cpp index a190672..b41c996 100644 --- a/src/cosyvoice.cpp +++ b/src/cosyvoice.cpp @@ -247,6 +247,16 @@ bool cosyvoice_llm_set_kv_cache_len(cosyvoice_context_t ctx, uint32_t len) return ctx->llm_set_kv_cache_len(len); } +void cosyvoice_llm_offload_kv_cache(cosyvoice_context_t ctx) +{ + ctx->llm_offload_kv_cache(); +} + +void cosyvoice_llm_load_kv_cache(cosyvoice_context_t ctx) +{ + ctx->llm_load_kv_cache(); +} + int cosyvoice_llm_sample_token(cosyvoice_context_t ctx) { return ctx->llm_sample_token(); @@ -282,12 +292,22 @@ bool cosyvoice_llm_job(cosyvoice_context_t ctx, const int* text, uint32_t text_l return ctx->llm_job(text, text_len, prompt); } +bool cosyvoice_llm_job_ext(cosyvoice_context_t ctx, const int* text, uint32_t text_len, cosyvoice_prompt_t prompt, uint32_t max_new_tokens, bool* final) +{ + return ctx->llm_job_ext(text, text_len, prompt, max_new_tokens, final); +} + bool cosyvoice_token2wav(cosyvoice_context_t ctx, const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, cosyvoice_generated_speech_ptr generated_speech) { return ctx->token2wav(token_ids, n_tokens, speed, prompt, generated_speech); } -bool cosyvoice_tts(cosyvoice_context_t ctx, const int* text, uint32_t text_len, float speed, cosyvoice_prompt_t prompt, cosyvoice_generated_speech_ptr result) +bool cosyvoice_token2wav_ext(cosyvoice_context_t ctx, const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* speech_offset_ptr, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result) +{ + return ctx->token2wav_ext(token_ids, n_tokens, speed, prompt, speech_offset_ptr, streaming, finalize, result); +} + +static bool check_length(cosyvoice_prompt_t prompt, uint32_t text_len, uint32_t n_max_seq) { // Defensive guard: bail out before llm_prefill / llm_decode would refuse the request because // the prefill layout exceeds n_max_seq. Without this, llm_job throws and emits an error log @@ -297,22 +317,31 @@ bool cosyvoice_tts(cosyvoice_context_t ctx, const int* text, uint32_t text_len, // length. llm_job internally truncates the cache back to either SOS or SOS+prompt_text // before re-prefilling, so a large residual cache from a previous call (e.g. when callers // synthesize chunked text via tts_job) is not a real budget consumer. - cosyvoice_context_params_t params; - ctx->get_context_params(¶ms); const uint32_t prompt_text_len = static_cast(prompt->prompt_text.size()); const uint32_t prompt_speech_len = prompt->llm_prompt_speech_tokens.second; // Worst-case prefill: SOS(1) + prompt_text + text + (task_token(1) + prompt_speech_tokens) when present. uint64_t prefill_len = 1ull + prompt_text_len + text_len; if (prompt_speech_len != 0) prefill_len += 1ull + prompt_speech_len; - if (prefill_len + 1 > params.n_max_seq) - { - result->data = nullptr; - result->length = 0; + if (prefill_len + 1 > n_max_seq) return false; + return true; +} + +bool cosyvoice_tts(cosyvoice_context_t ctx, const int* text, uint32_t text_len, float speed, cosyvoice_prompt_t prompt, cosyvoice_generated_speech_ptr result) +{ + { + cosyvoice_context_params_t params; + ctx->get_context_params(¶ms); + if (!check_length(prompt, text_len, params.n_max_seq)) + { + result->data = nullptr; + result->length = 0; + return false; + } } - if(ctx->llm_job(text, text_len, prompt) + if (ctx->llm_job(text, text_len, prompt) && ctx->token2wav(ctx->llm_get_accepted_tokens(), ctx->llm_get_n_accepted_tokens(), speed, prompt, result)) return true; result->data = nullptr; @@ -320,6 +349,41 @@ bool cosyvoice_tts(cosyvoice_context_t ctx, const int* text, uint32_t text_len, return false; } +bool cosyvoice_tts_stream(cosyvoice_context_t ctx, const int* text, uint32_t text_len, float speed, cosyvoice_prompt_t prompt, cosyvoice_tts_audio_callback_t callback, void* user_data) +{ + cosyvoice_context_params_t params; + ctx->get_context_params(¶ms); + if (!check_length(prompt, text_len, params.n_max_seq)) + return false; + + const auto chunk_tokens = ctx->get_chunk_tokens(); + + ctx->llm_clear_accepted_tokens(); + uint32_t speech_offset = 0; + bool final = false; + do + { + if (!ctx->llm_job_ext(text, text_len, prompt, chunk_tokens, &final)) + return false; + text = nullptr; + const uint32_t n_tokens = ctx->llm_get_n_accepted_tokens(); + const int* tokens = ctx->llm_get_accepted_tokens(); + + if (!final && params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) + ctx->llm_offload_kv_cache(); + + cosyvoice_generated_speech result = {}; + if (!ctx->token2wav_ext(tokens, n_tokens, speed, prompt, &speech_offset, true, final, &result)) + return false; + + if (result.data && result.length > 0 + && !callback(result.data, result.length, user_data)) + return false; + } while (!final); + + return true; +} + ggml_status cosyvoice_get_last_status(cosyvoice_context_t ctx) { return ctx->get_last_status(); From 177ac670b48a96d7379917b0c8d6562afb284729 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Sun, 5 Jul 2026 00:36:03 +0800 Subject: [PATCH 02/34] cli: add streaming audio playback and interactive /stream toggle --- tools/cli/cli_audio_player.cpp | 103 +++++++++++++++++++-------------- tools/cli/cli_audio_player.h | 3 + tools/cli/cosyvoice-cli.cpp | 102 +++++++++++++++++++++++++++----- 3 files changed, 150 insertions(+), 58 deletions(-) diff --git a/tools/cli/cli_audio_player.cpp b/tools/cli/cli_audio_player.cpp index 165a522..2d5b9e0 100644 --- a/tools/cli/cli_audio_player.cpp +++ b/tools/cli/cli_audio_player.cpp @@ -15,59 +15,63 @@ #define MA_NO_MP3 #define MA_NO_FLAC #define MA_NO_WAV +#define NOMINMAX #include "miniaudio.h" #include -#include -#include +#include +#include + #include -#include -struct playback_state +struct playback_context { - const float* data = nullptr; - uint32_t length = 0; - std::atomic cursor{0}; -}; + const streaming_callback_t& procedure; + std::mutex mutex; + std::condition_variable cv; + bool done = false; -static void data_callback(ma_device* device, void* output, const void* input, ma_uint32 frame_count) -{ - auto* state = reinterpret_cast(device->pUserData); - float* out = reinterpret_cast(output); - if (!state || !out) - return; + void call_procedure(float* output, uint32_t frame_count) + { + if (!done && !procedure(output, frame_count) + || is_playback_interrupted()) + { + { + std::unique_lock lock(mutex); + done = true; + } + cv.notify_one(); + } + } - const uint32_t cursor = state->cursor.load(std::memory_order_relaxed); - const uint32_t remaining = cursor < state->length ? (state->length - cursor) : 0; - const uint32_t frames_to_copy = std::min(frame_count, remaining); - if (frames_to_copy > 0) + void wait_for_completion() { - std::memcpy(out, state->data + cursor, frames_to_copy * sizeof(float)); - state->cursor.store(cursor + frames_to_copy, std::memory_order_relaxed); + std::unique_lock lock(mutex); + cv.wait(lock, [this] { return done; }); } - if (frames_to_copy < frame_count) - std::fill(out + frames_to_copy, out + frame_count, 0.0f); +}; + +static void streaming_callback(ma_device* device, void* output, const void* input, ma_uint32 frame_count) +{ + reinterpret_cast(device->pUserData)->call_procedure(reinterpret_cast(output), frame_count); } -bool cli_audio_play_pcm_blocking(const float* data, uint32_t length, uint32_t sample_rate, std::string* error) +bool cli_audio_play_pcm_streaming(uint32_t sample_rate, std::string* error, streaming_callback_t callback) { - if (!data || length == 0 || sample_rate == 0) + if (sample_rate == 0) { if (error) - *error = "Audio playback input is empty."; + *error = "Invalid sample rate."; return false; } - playback_state state; - state.data = data; - state.length = length; - + playback_context context{ callback }; ma_device_config config = ma_device_config_init(ma_device_type_playback); config.playback.format = ma_format_f32; config.playback.channels = 1; config.sampleRate = sample_rate; - config.dataCallback = data_callback; - config.pUserData = &state; + config.dataCallback = streaming_callback; + config.pUserData = &context; ma_device device; if (ma_device_init(nullptr, &config, &device) != MA_SUCCESS) @@ -85,20 +89,31 @@ bool cli_audio_play_pcm_blocking(const float* data, uint32_t length, uint32_t sa return false; } - const auto sleep_chunk = std::chrono::milliseconds(20); - bool interrupted = false; - while (state.cursor.load(std::memory_order_relaxed) < state.length) + context.wait_for_completion(); + ma_device_stop(&device); + ma_device_uninit(&device); + return !is_playback_interrupted(); +} + +bool cli_audio_play_pcm_blocking(const float* data, uint32_t length, uint32_t sample_rate, std::string* error) +{ + if (!data || length == 0) { - interrupted = is_playback_interrupted(); - if (interrupted) - break; - std::this_thread::sleep_for(sleep_chunk); + if (error) + *error = "Audio playback input is empty."; + return false; } - if (!interrupted && !is_playback_interrupted()) - std::this_thread::sleep_for(std::chrono::milliseconds(50)); - - ma_device_stop(&device); - ma_device_uninit(&device); - return !interrupted; + uint32_t offset = 0; + return cli_audio_play_pcm_streaming(sample_rate, error, [data, length, &offset](float* output, uint32_t frame_count) + { + uint32_t remaining_frames = length - offset; + if (remaining_frames == 0) return false; + uint32_t frames_to_copy = std::min(frame_count, remaining_frames); + std::memcpy(output, data + offset, frames_to_copy * sizeof(float)); + if (frame_count > frames_to_copy) + std::memset(output + frames_to_copy, 0, (frame_count - frames_to_copy) * sizeof(float)); + offset += frames_to_copy; + return true; + }); } diff --git a/tools/cli/cli_audio_player.h b/tools/cli/cli_audio_player.h index e638bb1..46e8fd2 100644 --- a/tools/cli/cli_audio_player.h +++ b/tools/cli/cli_audio_player.h @@ -3,6 +3,7 @@ #include #include #include +#include inline std::atomic_bool g_playback_interrupted; @@ -16,4 +17,6 @@ inline bool is_playback_interrupted() return g_playback_interrupted.load(std::memory_order_relaxed); } +using streaming_callback_t = std::function; +bool cli_audio_play_pcm_streaming(uint32_t sample_rate, std::string* error, streaming_callback_t callback); bool cli_audio_play_pcm_blocking(const float* data, uint32_t length, uint32_t sample_rate, std::string* error); diff --git a/tools/cli/cosyvoice-cli.cpp b/tools/cli/cosyvoice-cli.cpp index 47befe9..38ae62b 100644 --- a/tools/cli/cosyvoice-cli.cpp +++ b/tools/cli/cosyvoice-cli.cpp @@ -1,4 +1,4 @@ -#ifdef _MSC_VER +#ifdef _MSC_VER #define _CRT_SECURE_NO_WARNINGS #endif @@ -18,11 +18,13 @@ #include #include -#include -#include #include #include #include +#include +#include + +#include #include #include #include @@ -263,6 +265,7 @@ static void print_interactive_commands() printf(" /seed [value] Show or set next seed.\n"); printf(" /seed-policy Show or set seed policy.\n"); printf(" /help Show command list.\n"); + printf(" /stream Toggle streaming playback.\n"); printf(" /exit Exit interactive mode. Ctrl+C also exits.\n"); } @@ -848,6 +851,7 @@ static void run_interactive_loop( uint32_t sample_rate) { audio_cache cache; + bool streaming = false; std::string line; if (seed_state && !seed_state->has_next_seed) { @@ -856,9 +860,9 @@ static void run_interactive_loop( } printf("\nInteractive mode. Type text to synthesize, /exit to quit, or press Ctrl+C to quit" #ifndef COSYVOICE_CLI_NO_PLAYBACK - " (during playback, Ctrl+C stops playback)" + " (during playback, Ctrl+C stops playback)" #endif - ".\n"); + ".\n"); print_interactive_commands(); interactive_ctrl_c_guard ctrl_c_guard; @@ -1010,6 +1014,11 @@ static void run_interactive_loop( printf("Saved audio %u to %s\n", id, path.c_str()); } #ifndef COSYVOICE_CLI_NO_PLAYBACK + else if (cmd == "/stream") + { + streaming = !streaming; + printf("Streaming playback %s.\n", streaming ? "enabled" : "disabled"); + } else if (cmd == "/play") { bool id_ok = false; @@ -1046,18 +1055,83 @@ static void run_interactive_loop( uint32_t used_seed = 0; update_tts_seed(ctx, seed_state, &used_seed); - std::string error; - auto pcm = generate_tts_audio(tts_ctx, options, trimmed, &error); - if (!pcm.data) - { - if (!error.empty()) - print_error_log("Error: %s\n", error.c_str()); - continue; - } cached_audio entry; entry.seed = used_seed; entry.text = trimmed; - entry.pcm.insert(entry.pcm.end(), pcm.data, pcm.data + pcm.length); + + if (streaming) + { + struct stream_play_state + { + cached_audio* entry; + uint32_t sample_rate; + uint32_t cursor = 0; + std::mutex mutex; + }; + stream_play_state play_state{ &entry, sample_rate }; + + auto stream_callback = [](const float* audio, uint32_t n_samples, void* user_data) + { + auto state = static_cast(user_data); + std::lock_guard lock(state->mutex); + state->entry->pcm.insert(state->entry->pcm.end(), audio, audio + n_samples); + return !is_playback_interrupted(); + }; + + std::atomic_bool inferencing_done = false; + std::string error_string; + std::thread playback_thread(cli_audio_play_pcm_streaming, sample_rate, &error_string, + streaming_callback_t([&](float* data, uint32_t n_samples) + { + std::lock_guard lock(play_state.mutex); + if (play_state.cursor >= play_state.entry->pcm.size()) + if (inferencing_done) + return false; + else + memset(data, 0, n_samples * sizeof(float)); + else + { + uint32_t available = static_cast(play_state.entry->pcm.size()) - play_state.cursor; + uint32_t to_copy = std::min(n_samples, available); + memcpy(data, play_state.entry->pcm.data() + play_state.cursor, to_copy * sizeof(float)); + if (to_copy < n_samples) + memset(data + to_copy, 0, (n_samples - to_copy) * sizeof(float)); + play_state.cursor += n_samples; + } + return true; + })); + + bool ok; + if (options.mode == "cross-lingual") + ok = cosyvoice_tts_cross_lingual_stream(tts_ctx, trimmed.c_str(), options.speed, stream_callback, &play_state); + else if (options.mode == "zero-shot") + ok = cosyvoice_tts_zero_shot_stream(tts_ctx, trimmed.c_str(), options.speed, stream_callback, &play_state); + else + ok = cosyvoice_tts_instruct_stream(tts_ctx, trimmed.c_str(), options.instruction.c_str(), options.speed, stream_callback, &play_state); + + inferencing_done = true; + playback_thread.join(); + ctrl_c_guard._register(); + + if (!ok) + { + print_error_log("Error: TTS streaming generation failed.\n"); + continue; + } + } + else + { + std::string error; + auto pcm = generate_tts_audio(tts_ctx, options, trimmed, &error); + if (!pcm.data) + { + if (!error.empty()) + print_error_log("Error: %s\n", error.c_str()); + continue; + } + entry.pcm.insert(entry.pcm.end(), pcm.data, pcm.data + pcm.length); + } + const auto id = cache.next_id++; cache.last_success_id = id; const auto duration = pcm_length_seconds(sample_rate, entry.pcm.size()); From 872199dc436c2edfd96ad2cb5f8292649b756eb1 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Sun, 5 Jul 2026 09:43:42 +0800 Subject: [PATCH 03/34] cosyvoice: fix streaming LLM token limit calculation for continuation calls --- src/cosyvoice-tts.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index 717f094..9fa8fd6 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -150,7 +150,7 @@ bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoic // First call: min/max from text input const auto min_len = static_cast(text_len * worker->config.min_token_text_ratio); const auto max_len = static_cast(text_len * worker->config.max_token_text_ratio); - const auto limit = std::min(max_new_tokens, max_len); + const auto limit = std::min(llm_get_n_accepted_tokens() + max_new_tokens, max_len); // FSQ silent/breath token filtering: allow up to 5 consecutive silent tokens, // then drop any further consecutive ones to avoid generating long silence. @@ -159,7 +159,7 @@ bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoic constexpr uint32_t max_silent_token_num = 5; const auto& silent_tokens = cv3_shared->silent_tokens; - for (uint32_t n = 0; n != limit; ++n) + for (uint32_t n = llm_get_n_accepted_tokens(); n != limit; ++n) { if (!llm_decode(speech_type, cur)) throw std::runtime_error("Failed to decode LLM output.\n"); From a054711e76fa510516fc2da18e92c02410cd67e6 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Wed, 8 Jul 2026 16:47:43 +0800 Subject: [PATCH 04/34] cosyvoice: add KV cache for DiT --- include/cosyvoice-interface.h | 4 +- include/cosyvoice-lowlevel.h | 23 +++- src/cosyvoice-graph.cpp | 96 +++++++------- ...lm-kv-cache.cpp => cosyvoice-kv-cache.cpp} | 89 +++++++++---- ...ce-llm-kv-cache.h => cosyvoice-kv-cache.h} | 17 +-- src/cosyvoice-llm.cpp | 10 +- src/cosyvoice-loader.cpp | 45 ++++--- src/cosyvoice-model.cpp | 54 +++++--- src/cosyvoice-model.h | 13 +- src/cosyvoice-modules.h | 14 +- src/cosyvoice-tts.cpp | 120 +++++++++--------- src/cosyvoice.cpp | 44 +++++-- 12 files changed, 321 insertions(+), 208 deletions(-) rename src/{cosyvoice-llm-kv-cache.cpp => cosyvoice-kv-cache.cpp} (72%) rename src/{cosyvoice-llm-kv-cache.h => cosyvoice-kv-cache.h} (78%) diff --git a/include/cosyvoice-interface.h b/include/cosyvoice-interface.h index c640fd1..5d6e248 100644 --- a/include/cosyvoice-interface.h +++ b/include/cosyvoice-interface.h @@ -90,13 +90,15 @@ struct cosyvoice_model_context uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, - uint32_t* speech_offset_ptr, + uint32_t* offset, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result ) = 0; ///< Convert speech tokens into waveform samples with additional options. virtual uint32_t get_chunk_tokens() = 0; ///< Get the number of tokens processed in each chunk during streaming inference. + virtual uint32_t get_flow_overlap_tokens() = 0; ///< Get the number of overlapping tokens used for streaming inference. + virtual uint32_t get_hift_overlap_tokens() = 0; ///< Get the number of extra tokens to overlap for HiFT lookahead in streaming inference. // Status virtual ggml_status get_last_status() = 0; ///< Get the status of the most recent backend operation. diff --git a/include/cosyvoice-lowlevel.h b/include/cosyvoice-lowlevel.h index 5731607..45f37a4 100644 --- a/include/cosyvoice-lowlevel.h +++ b/include/cosyvoice-lowlevel.h @@ -311,7 +311,7 @@ COSYVOICE_API bool cosyvoice_token2wav_ext( uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, - uint32_t* speech_offset_ptr, + uint32_t* offset, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result @@ -329,7 +329,11 @@ COSYVOICE_API bool cosyvoice_tts( cosyvoice_generated_speech_ptr result ); - +/** + * @brief Run the full TTS pipeline from input tokens to waveform output with streaming support. + * @param callback Callback function to receive generated audio chunks. + * @param user_data User-defined data passed to the callback. + */ COSYVOICE_API bool cosyvoice_tts_stream( cosyvoice_context_t ctx, const int* text, @@ -340,6 +344,21 @@ COSYVOICE_API bool cosyvoice_tts_stream( void* user_data ); +/** +* @brief Get the number of tokens processed in each chunk during streaming inference. +*/ +COSYVOICE_API uint32_t cosyvoice_get_chunk_tokens(cosyvoice_context_t ctx); + +/** + * @brief Get the number of tokens that overlap between consecutive chunks during streaming inference. + */ +COSYVOICE_API uint32_t cosyvoice_get_flow_overlap_tokens(cosyvoice_context_t ctx); + +/** +* @brief Get the number of extra tokens to overlap for HiFT lookahead in streaming inference. +*/ +COSYVOICE_API uint32_t cosyvoice_get_hift_overlap_tokens(cosyvoice_context_t ctx); + // ---------------------------------------------------------------------------- // Tokenizer Operations // ---------------------------------------------------------------------------- diff --git a/src/cosyvoice-graph.cpp b/src/cosyvoice-graph.cpp index 30b7f3e..adb43ac 100644 --- a/src/cosyvoice-graph.cpp +++ b/src/cosyvoice-graph.cpp @@ -323,7 +323,7 @@ std::array AdaLayerNormZero::build_cgraph(ggml_context* ctx, gg return { x, gate_msa, shift_mlp, scale_mlp, gate_mlp }; } -ggml_tensor* Attention::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* position_ids, int64_t cut_len, ggml_tensor* attn_mask) const +ggml_tensor* Attention::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* position_ids, int64_t cut_len, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, int layer_idx) const { const auto full_seq_len = position_ids->ne[0]; const auto seq_len = full_seq_len - (cut_len > 0 ? cut_len : 0); @@ -332,6 +332,7 @@ ggml_tensor* Attention::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_ten auto key = to_k.build_cgraph(ctx, x); auto value = to_v.build_cgraph(ctx, x); + auto original_position_ids = position_ids; auto full_position_ids = position_ids; if (cut_len > 0) { @@ -361,21 +362,29 @@ ggml_tensor* Attention::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_ten key = ggml_permute(ctx, key, 0, 2, 1, 3); value = ggml_permute(ctx, value, 0, 2, 1, 3); - query = ggml_cont(ctx, query); - key = ggml_cont(ctx, key); - ggml_tensor* attn_output; - if (fattn) - attn_output = ggml_flash_attn_ext(ctx, query, key, value, attn_mask, 1.f / sqrtf(static_cast(head_dim)), 0.f, 0.f); + if (kv_cache) + { + kv_cache->update_cache(ctx, gf, key, value, original_position_ids, layer_idx); + attn_output = kv_cache->attention_forward(ctx, query, key, value, nullptr); + } else { - auto attn_scores = ggml_mul_mat(ctx, key, query); - auto attn_weights = ggml_soft_max_ext_inplace(ctx, attn_scores, attn_mask, 1.f / std::sqrt(static_cast(key->ne[0])), 0.f); - value = ggml_permute(ctx, value, 1, 0, 2, 3); - value = ggml_cont(ctx, value); - attn_output = ggml_mul_mat(ctx, value, attn_weights); - attn_output = ggml_permute(ctx, attn_output, 0, 2, 1, 3); - attn_output = ggml_cont(ctx, attn_output); + query = ggml_cont(ctx, query); + key = ggml_cont(ctx, key); + + if (fattn) + attn_output = ggml_flash_attn_ext(ctx, query, key, value, nullptr, 1.f / std::sqrt(static_cast(head_dim)), 0.f, 0.f); + else + { + value = ggml_permute(ctx, value, 0, 2, 1, 3); + value = ggml_cont(ctx, value); + auto attn_scores = ggml_mul_mat(ctx, key, query); + auto attn_weights = ggml_soft_max_ext_inplace(ctx, attn_scores, nullptr, 1.f / std::sqrt(static_cast(head_dim)), 0.f); + auto attn_output = ggml_mul_mat(ctx, value, attn_weights); + attn_output = ggml_permute(ctx, attn_output, 0, 2, 1, 3); + attn_output = ggml_cont(ctx, attn_output); + } } x = ggml_reshape_3d(ctx, attn_output, @@ -395,11 +404,11 @@ ggml_tensor* FeedForward::build_cgraph(ggml_context* ctx, ggml_tensor* x) const return x; } -ggml_tensor* DiTBlock::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* time_emb, ggml_tensor* position_ids, int64_t cut_len, ggml_tensor* attn_mask) const +ggml_tensor* DiTBlock::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* time_emb, ggml_tensor* position_ids, int64_t cut_len, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, int layer_idx) const { auto [norm, gate_msa, shift_mlp, scale_mlp, gate_mlp] = attn_norm.build_cgraph(ctx, x, time_emb); - auto attn_output = attn.build_cgraph(ctx, norm, position_ids, cut_len, attn_mask); + auto attn_output = attn.build_cgraph(ctx, norm, position_ids, cut_len, kv_cache, gf, layer_idx); gate_msa = unsqueeze(ctx, gate_msa, 1); attn_output = ggml_mul(ctx, attn_output, gate_msa); if (cut_len > 0) @@ -430,10 +439,8 @@ ggml_tensor* AdaLayerNorm_Final::build_cgraph(ggml_context* ctx, ggml_tensor* x, return x; } -ggml_tensor* DiT::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* mu, ggml_tensor* t, ggml_tensor* spks, ggml_tensor* cond, int64_t cut_len, ggml_tensor*& ref_position_ids, ggml_backend_op_capabilities capabilities, bool streaming, ggml_tensor*& ref_attn_mask) const +ggml_tensor* DiT::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* mu, ggml_tensor* t, ggml_tensor* spks, ggml_tensor* cond, int64_t cut_len, ggml_tensor*& ref_position_ids, ggml_backend_op_capabilities capabilities, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, int step) const { - x = ggml_permute(ctx, x, 1, 0, 2, 3); - t = time_embed.build_cgraph(ctx, t); x = input_embed.build_cgraph(ctx, x, cond, mu, spks, capabilities); @@ -442,19 +449,11 @@ ggml_tensor* DiT::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* m ref_position_ids = position_ids; position_ids = ggml_dup(ctx, position_ids); - ggml_tensor* attn_mask = nullptr; - if (streaming) - { - const auto seq_len = x->ne[1]; - attn_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, seq_len, seq_len); - ggml_set_input(attn_mask); - ref_attn_mask = attn_mask; - } - - for (const auto& block : std::span(transformer_blocks.cbegin(), transformer_blocks.size() - 1)) - x = block.build_cgraph(ctx, x, t, position_ids, 0, attn_mask); + int layer_idx_base = step * static_cast(transformer_blocks.size()); + for (int i = 0; i < transformer_blocks.size() - 1; ++i) + x = transformer_blocks[i].build_cgraph(ctx, x, t, position_ids, 0, kv_cache, gf, layer_idx_base + i); // Apply `cut_len` only on the final block. - x = transformer_blocks.back().build_cgraph(ctx, x, t, position_ids, cut_len, attn_mask); + x = transformer_blocks.back().build_cgraph(ctx, x, t, position_ids, cut_len, kv_cache, gf, layer_idx_base + static_cast(transformer_blocks.size()) - 1); x = norm_out.build_cgraph(ctx, x, t); @@ -468,7 +467,7 @@ CausalConditionalCFM::DiTContext CausalConditionalCFM::prepare_context(ggml_cont .x = ggml_new_tensor_2d( ctx, GGML_TYPE_F32, - mu->ne[1], 80), + 80, mu->ne[1]), .mu_in = ggml_pad(ctx, mu, 0, 0, 1, 0), @@ -493,7 +492,7 @@ std::array CausalConditionalCFM::get_t_and_dt(ggml_context* ctx, int s return { t, dt }; } -ggml_tensor* CausalConditionalCFM::build_cgraph_one_step(ggml_context* ctx, const DiTContext& ditctx, int step, ggml_backend_op_capabilities capabilities, int64_t cut_len, ggml_tensor*& t_tensor, ggml_tensor*& position_ids, bool streaming, ggml_tensor*& ref_attn_mask) const +ggml_tensor* CausalConditionalCFM::build_cgraph_one_step(ggml_context* ctx, const DiTContext& ditctx, int step, ggml_backend_op_capabilities capabilities, int64_t cut_len, ggml_tensor*& t_tensor, ggml_tensor*& position_ids, ggml_cgraph* gf, cosyvoice_kv_cache* kv_cache) const { auto x = ditctx.x; auto [t, dt] = get_t_and_dt(ctx, step); @@ -512,14 +511,12 @@ ggml_tensor* CausalConditionalCFM::build_cgraph_one_step(ggml_context* ctx, cons step == t_span.size() - 1 ? cut_len : 0, position_ids, capabilities, - streaming, - ref_attn_mask), + kv_cache, + gf, + step - 1), 2); - dphi_dt = ggml_permute(ctx, dphi_dt, 1, 0, 2, 3); - cfg_dphi_dt = ggml_permute(ctx, cfg_dphi_dt, 1, 0, 2, 3); - dphi_dt = ggml_cont(ctx, dphi_dt); - cfg_dphi_dt = ggml_cont(ctx, cfg_dphi_dt); + cfg_dphi_dt->nb[3] = dphi_dt->nb[3] = dphi_dt->nb[2]; cfg_dphi_dt = ggml_scale(ctx, cfg_dphi_dt, inference_cfg_rate); dphi_dt = ggml_scale(ctx, dphi_dt, 1.f + inference_cfg_rate); dphi_dt = ggml_sub(ctx, @@ -527,7 +524,7 @@ ggml_tensor* CausalConditionalCFM::build_cgraph_one_step(ggml_context* ctx, cons cfg_dphi_dt); if (step == t_span.size() - 1 && cut_len > 0) - x = ggml_view_3d(ctx, x, x->ne[0] - cut_len, x->ne[1], x->ne[2], x->nb[1], x->nb[2], x->nb[0] * cut_len); + x = ggml_view_3d(ctx, x, x->ne[0], x->ne[1] - cut_len, x->ne[2], x->nb[1], x->nb[2], x->nb[1] * cut_len); x = ggml_add(ctx, x, ggml_scale(ctx, dphi_dt, dt)); @@ -536,7 +533,7 @@ ggml_tensor* CausalConditionalCFM::build_cgraph_one_step(ggml_context* ctx, cons return x; } -ggml_tensor* PreLookaheadLayer::build_cgraph(ggml_context* ctx, ggml_tensor* inputs, bool streaming) const +ggml_tensor* PreLookaheadLayer::build_cgraph(ggml_context* ctx, ggml_tensor* inputs, bool streaming, uint32_t cut_len) const { auto outputs = ggml_permute(ctx, inputs, 1, 0, 2, 3); outputs = ggml_cont(ctx, outputs); @@ -552,11 +549,17 @@ ggml_tensor* PreLookaheadLayer::build_cgraph(ggml_context* ctx, ggml_tensor* inp outputs = ggml_permute(ctx, outputs, 1, 0, 2, 3); outputs = ggml_cont(ctx, outputs); + if (cut_len) + { + outputs = ggml_view_3d(ctx, outputs, outputs->ne[0], outputs->ne[1] - cut_len, outputs->ne[2], outputs->nb[1], outputs->nb[2], outputs->nb[1] * cut_len); + inputs = ggml_view_3d(ctx, inputs, inputs->ne[0], inputs->ne[1] - cut_len, inputs->ne[2], inputs->nb[1], inputs->nb[2], inputs->nb[1] * cut_len); + } + outputs = ggml_add(ctx, outputs, inputs); return outputs; } -CausalMaskedDiffWithDiT::EncodeResult CausalMaskedDiffWithDiT::build_cgraph_encode(ggml_context* ctx, ggml_tensor* token, ggml_tensor* prompt_token, ggml_tensor* prompt_feat, ggml_tensor* embedding, ggml_backend_op_capabilities capabilities, bool streaming) const +CausalMaskedDiffWithDiT::EncodeResult CausalMaskedDiffWithDiT::build_cgraph_encode(ggml_context* ctx, ggml_tensor* token, ggml_tensor* prompt_token, ggml_tensor* prompt_feat, ggml_tensor* embedding, ggml_backend_op_capabilities capabilities, uint32_t cut_len, bool streaming) const { embedding = ggml_l2_norm(ctx, embedding, 1e-6f); embedding = spk_embed_affine_layer.build_cgraph(ctx, embedding); @@ -564,20 +567,23 @@ CausalMaskedDiffWithDiT::EncodeResult CausalMaskedDiffWithDiT::build_cgraph_enco token = concat_tensors(ctx, std::array{ prompt_token, token }, 0, capabilities); token = ggml_get_rows(ctx, input_embedding, token); - ggml_tensor* h = pre_lookahead_layer.build_cgraph(ctx, token, streaming); + ggml_tensor* h = pre_lookahead_layer.build_cgraph(ctx, token, streaming, cut_len); h = unsqueeze(ctx, h, 1); h = ggml_repeat_4d(ctx, h, h->ne[0], token_mel_ratio, h->ne[2], h->ne[3]); h = ggml_reshape_3d(ctx, h, h->ne[0], h->ne[1] * h->ne[2], h->ne[3]); + cut_len *= token_mel_ratio; const auto mel_len1 = prompt_feat->ne[1]; - const auto mel_len2 = h->ne[1] - mel_len1; + const auto mel_len2 = h->ne[1] - mel_len1 + cut_len; auto conds = ggml_pad(ctx, prompt_feat, 0, static_cast(mel_len2), 0, 0); + if (cut_len) + conds = ggml_view_3d(ctx, conds, conds->ne[0], conds->ne[1] - cut_len, conds->ne[2], conds->nb[1], conds->nb[2], conds->nb[1] * cut_len); return EncodeResult{ .mu = h, .spks = embedding, .conds = conds, - .cut_len = mel_len1 + .cut_len = cut_len ? 0 : mel_len1 }; } @@ -744,7 +750,7 @@ std::array CausalHiFTGenerator::build_cgraph(ggml_context* ctx, s_stft = ggml_cont(ctx, s_stft); s_stft = ggml_reshape_2d(ctx, s_stft, s_stft->ne[0], s_stft->ne[1] * 2); if (!finalize) - s_stft = ggml_view_2d(ctx, s_stft, s_stft->ne[0] - (scale_factor / hop_len * conv_pre_look_right), s_stft->ne[1], s_stft->nb[1], 0); + s_stft = ggml_view_2d(ctx, s_stft, s_stft->ne[0] - (scale_factor / hop_len * conv_pre.causal_padding()), s_stft->ne[1], s_stft->nb[1], 0); const auto num_upsamples = ups.size(); const auto num_kernels = resblocks.size() / num_upsamples; diff --git a/src/cosyvoice-llm-kv-cache.cpp b/src/cosyvoice-kv-cache.cpp similarity index 72% rename from src/cosyvoice-llm-kv-cache.cpp rename to src/cosyvoice-kv-cache.cpp index 80a6366..660ad2a 100644 --- a/src/cosyvoice-llm-kv-cache.cpp +++ b/src/cosyvoice-kv-cache.cpp @@ -1,22 +1,22 @@ #include "cosyvoice-model.h" -#include "cosyvoice-llm-kv-cache.h" +#include "cosyvoice-kv-cache.h" #include #include #include -void cosyvoice_llm_kv_cache::build_kv_cache( +void cosyvoice_kv_cache::build_kv_cache( ggml_backend_t backend, - ggml_backend_buffer_ptr& shared_buffer, + ggml_backend_buffer_ptr& buffer, int layers, int k_head_dim, int v_head_dim, - int num_attention_heads, int num_key_value_heads, uint32_t max_seq, ggml_type k_type, ggml_type v_type, + int batch_size, bool fattn) { cur_len = 0; @@ -34,27 +34,27 @@ void cosyvoice_llm_kv_cache::build_kv_cache( }; ctx = ggml_init(params); - shared_buffer.reset(initialize_buffer(backend, k_head_dim, v_head_dim, num_attention_heads, num_key_value_heads, max_seq, k_type, v_type, fattn)); + buffer.reset(initialize_buffer(backend, k_head_dim, v_head_dim, num_attention_heads, num_key_value_heads, max_seq, k_type, v_type, batch_size, fattn)); } -ggml_backend_buffer* cosyvoice_llm_kv_cache::initialize_buffer(ggml_backend_t backend, int k_head_dim, int v_head_dim, int num_attention_heads, int num_key_value_heads, uint32_t max_seq, ggml_type k_type, ggml_type v_type, bool fattn) +ggml_backend_buffer* cosyvoice_kv_cache::initialize_buffer(ggml_backend_t backend, int k_head_dim, int v_head_dim, int num_attention_heads, int num_key_value_heads, uint32_t max_seq, ggml_type k_type, ggml_type v_type, int batch_size, bool fattn) { - int64_t k_ne[3] = { k_head_dim, max_seq, num_key_value_heads }; - int64_t v_ne[3] = { v_head_dim, max_seq, num_key_value_heads }; + int64_t k_ne[4] = { k_head_dim, max_seq, num_key_value_heads, batch_size }; + int64_t v_ne[4] = { v_head_dim, max_seq, num_key_value_heads, batch_size }; if (!fattn) std::swap(v_ne[0], v_ne[1]); ggml_reset(ctx); for (auto& [k, v, k_view, v_view] : std::span(kv_cache_layers, layers)) { - k = ggml_new_tensor(ctx, k_type, 3, k_ne); - v = ggml_new_tensor(ctx, v_type, 3, v_ne); + k = ggml_new_tensor(ctx, k_type, 4, k_ne); + v = ggml_new_tensor(ctx, v_type, 4, v_ne); k_view = nullptr; v_view = nullptr; } return ggml_backend_alloc_ctx_tensors(ctx, backend); } -uint32_t cosyvoice_llm_kv_cache::reset_buffer(ggml_backend_buffer* buffer) +uint32_t cosyvoice_kv_cache::reset_buffer(ggml_backend_buffer* buffer) { cur_len = 0; auto alignment = ggml_backend_buffer_get_alignment(buffer); @@ -106,7 +106,7 @@ struct offloaded_kv_cache }; #pragma pack(pop) -void cosyvoice_llm_kv_cache::offload_cache(ggml_backend_t backend, ggml_backend_sched* sched, uint32_t n_tokens) +void cosyvoice_kv_cache::offload_cache(ggml_backend_t backend, ggml_backend_sched* sched, uint32_t n_tokens) { char* buffer_base; size_t nbytes = kv_cache_layers[0].k->ne[0] * sizeof(float) * n_tokens * kv_cache_layers[0].v->ne[2]; @@ -176,9 +176,10 @@ void cosyvoice_llm_kv_cache::offload_cache(ggml_backend_t backend, ggml_backend_ ggml_backend_tensor_get_async(backend, offloaded_layer.v_view, offloaded_layer.v, 0, nbytes); buffer_base += nbytes * 2; } + ggml_backend_sched_synchronize(sched); } -void cosyvoice_llm_kv_cache::load_cache(ggml_backend_t backend, ggml_backend_sched* sched) +void cosyvoice_kv_cache::load_cache(ggml_backend_t backend, ggml_backend_sched* sched) { cur_len = offloaded_cache->len; ggml_reset(offloaded_cache->ctx); @@ -212,17 +213,17 @@ void cosyvoice_llm_kv_cache::load_cache(ggml_backend_t backend, ggml_backend_sch ggml_backend_tensor_set_async(backend, offloaded_layer.v_view, offloaded_layer.v, 0, offloaded_layer.v_view->nb[3]); } - ggml_backend_sched_graph_compute_async(sched, gf); + ggml_backend_sched_graph_compute(sched, gf); } -size_t cosyvoice_llm_kv_cache::get_offloaded_cache_size() const +size_t cosyvoice_kv_cache::get_offloaded_cache_size() const { if (offloaded_cache) return offloaded_cache->offloaded_kv_layers[0].k_view->nb[3] * 2 * layers; else return 0; } -void cosyvoice_llm_kv_cache::clear_offloaded_cache() +void cosyvoice_kv_cache::clear_offloaded_cache() { if (offloaded_cache) { @@ -233,14 +234,14 @@ void cosyvoice_llm_kv_cache::clear_offloaded_cache() } } -cosyvoice_llm_kv_cache::~cosyvoice_llm_kv_cache() +cosyvoice_kv_cache::~cosyvoice_kv_cache() { ggml_free(ctx); delete[] kv_cache_layers; clear_offloaded_cache(); } -void cosyvoice_llm_kv_cache::update_cache(ggml_context* ctx0, ggml_cgraph* gf, ggml_tensor*& k, ggml_tensor*& v, ggml_tensor* position_ids, int layer_idx) +void cosyvoice_kv_cache::update_cache(ggml_context* ctx0, ggml_cgraph* gf, ggml_tensor*& k, ggml_tensor*& v, ggml_tensor* position_ids, int layer_idx) { GGML_ASSERT(ggml_are_same_shape(k, v)); @@ -249,22 +250,22 @@ void cosyvoice_llm_kv_cache::update_cache(ggml_context* ctx0, ggml_cgraph* gf, g if (fattn) { layer.k_view = ggml_set_rows(ctx0, layer.k, k, position_ids); - layer.k_view = ggml_view_3d(ctx0, layer.k_view, k->ne[0], cur_len + position_ids->ne[0], k->ne[2], layer.k_view->nb[1], layer.k_view->nb[2], 0); + layer.k_view = ggml_view_4d(ctx0, layer.k_view, k->ne[0], cur_len + position_ids->ne[0], k->ne[2], k->ne[3], layer.k_view->nb[1], layer.k_view->nb[2], layer.k_view->nb[3], 0); layer.v_view = ggml_set_rows(ctx0, layer.v, v, position_ids); - layer.v_view = ggml_view_3d(ctx0, layer.v_view, v->ne[0], cur_len + position_ids->ne[0], v->ne[2], layer.v_view->nb[1], layer.v_view->nb[2], 0); + layer.v_view = ggml_view_4d(ctx0, layer.v_view, v->ne[0], cur_len + position_ids->ne[0], v->ne[2], v->ne[3], layer.v_view->nb[1], layer.v_view->nb[2], layer.v_view->nb[3], 0); } else { - auto k_view = ggml_view_3d(ctx0, layer.k, k->ne[0], k->ne[1], k->ne[2], layer.k->nb[1], layer.k->nb[2], layer.k->nb[1] * cur_len); + auto k_view = ggml_view_4d(ctx0, layer.k, k->ne[0], k->ne[1], k->ne[2], k->ne[3], layer.k->nb[1], layer.k->nb[2], layer.k->nb[3], layer.k->nb[1] * cur_len); k_view = ggml_cpy(ctx0, k, k_view); - layer.k_view = ggml_view_3d(ctx0, layer.k, k->ne[0], cur_len + k->ne[1], k->ne[2], layer.k->nb[1], layer.k->nb[2], 0); + layer.k_view = ggml_view_4d(ctx0, layer.k, k->ne[0], cur_len + k->ne[1], k->ne[2], k->ne[3], layer.k->nb[1], layer.k->nb[2], layer.k->nb[3], 0); layer.k_view->src[0] = k_view; v = ggml_permute(ctx0, v, 1, 0, 2, 3); - auto v_view = ggml_view_3d(ctx0, layer.v, v->ne[0], v->ne[1], v->ne[2], layer.v->nb[1], layer.v->nb[2], layer.v->nb[0] * cur_len); + auto v_view = ggml_view_4d(ctx0, layer.v, v->ne[0], v->ne[1], v->ne[2], v->ne[3], layer.v->nb[1], layer.v->nb[2], layer.v->nb[3], layer.v->nb[0] * cur_len); v_view = ggml_cpy(ctx0, v, v_view); - layer.v_view = ggml_view_3d(ctx0, layer.v, cur_len + v->ne[0], v->ne[1], v->ne[2], layer.v->nb[1], layer.v->nb[2], 0); + layer.v_view = ggml_view_4d(ctx0, layer.v, cur_len + v->ne[0], v->ne[1], v->ne[2], v->ne[3], layer.v->nb[1], layer.v->nb[2], layer.v->nb[3], 0); layer.v_view->src[0] = v_view; } @@ -275,7 +276,7 @@ void cosyvoice_llm_kv_cache::update_cache(ggml_context* ctx0, ggml_cgraph* gf, g ggml_build_forward_expand(gf, v); } -ggml_tensor* cosyvoice_llm_kv_cache::attention_forward(ggml_context* ctx0, ggml_cgraph* gf, ggml_tensor* query_states, ggml_tensor* key_states, ggml_tensor* value_states, ggml_tensor* attention_mask, int layer_idx) const +ggml_tensor* cosyvoice_kv_cache::attention_forward(ggml_context* ctx0, ggml_tensor* query_states, ggml_tensor* key_states, ggml_tensor* value_states, ggml_tensor* attention_mask) const { if (fattn) return ggml_flash_attn_ext(ctx0, query_states, key_states, value_states, attention_mask, 1.f / std::sqrt(static_cast(key_states->ne[0])), 0.f, 0.f); @@ -289,7 +290,7 @@ ggml_tensor* cosyvoice_llm_kv_cache::attention_forward(ggml_context* ctx0, ggml_ } } -void cosyvoice_llm_kv_cache::shift_kv_node_pos(uint32_t shift_pos) +void cosyvoice_kv_cache::shift_kv_node_pos(uint32_t shift_pos) { GGML_ASSERT(fattn); cur_len += shift_pos; @@ -301,7 +302,41 @@ void cosyvoice_llm_kv_cache::shift_kv_node_pos(uint32_t shift_pos) } } -bool cosyvoice_llm_kv_cache::can_reuse(bool prefill) const +bool cosyvoice_kv_cache::can_reuse(bool prefill) const { return fattn; } + +void cosyvoice_kv_cache::slide_kv_layers(int layer_idx, int stride) +{ + GGML_ASSERT(layer_idx + stride <= layers); + if (fattn) + { + const auto end = layer_idx + stride; + for (int cur = layer_idx; cur != end; ++cur) + { + auto& layer = kv_cache_layers[cur]; + auto& next_layer = kv_cache_layers[cur + stride]; + + auto k = next_layer.k; + auto k_view = layer.k_view; + next_layer.k_view = k_view; + k_view->data = k->data; + k_view->view_src = k; + k_view = k_view->src[0]; + k_view->data = k->data; + k_view->src[2] = k_view->view_src = k; + + auto v = next_layer.v; + auto v_view = layer.v_view; + next_layer.v_view = v_view; + v_view->data = v->data; + v_view->view_src = v; + v_view = v_view->src[0]; + v_view->data = v->data; + v_view->src[2] = v_view->view_src = v; + } + } + else + GGML_ABORT("slide_kv_layers is not supported for non-flash attention"); +} diff --git a/src/cosyvoice-llm-kv-cache.h b/src/cosyvoice-kv-cache.h similarity index 78% rename from src/cosyvoice-llm-kv-cache.h rename to src/cosyvoice-kv-cache.h index c3b1763..0e9581b 100644 --- a/src/cosyvoice-llm-kv-cache.h +++ b/src/cosyvoice-kv-cache.h @@ -5,11 +5,11 @@ #include // Non-FlashAttention mode currently rebuilds parts of the graph instead of reusing one. -class cosyvoice_llm_kv_cache +class cosyvoice_kv_cache { public: - cosyvoice_llm_kv_cache() = default; - ~cosyvoice_llm_kv_cache(); + cosyvoice_kv_cache() = default; + ~cosyvoice_kv_cache(); struct kv_cache_layer { @@ -22,26 +22,27 @@ class cosyvoice_llm_kv_cache void build_kv_cache( ggml_backend_t backend, - ggml_backend_buffer_ptr& shared_buffer, + ggml_backend_buffer_ptr& buffer, int layers, int k_head_dim, int v_head_dim, - int num_attention_heads, int num_key_value_heads, uint32_t max_seq, ggml_type k_type, ggml_type v_type, + int batch_size, bool fattn); void update_cache(ggml_context* ctx0, ggml_cgraph* gf, ggml_tensor*& k, ggml_tensor*& v, ggml_tensor* position_ids, int layer_idx); - ggml_tensor* attention_forward(ggml_context* ctx0, ggml_cgraph* gf, ggml_tensor* query_states, ggml_tensor* key_states, ggml_tensor* value_states, ggml_tensor* attention_mask, int layer_idx) const; + ggml_tensor* attention_forward(ggml_context* ctx0, ggml_tensor* query_states, ggml_tensor* key_states, ggml_tensor* value_states, ggml_tensor* attention_mask) const; void shift_kv_node_pos(uint32_t shift_pos); - bool can_reuse(bool prefill) const; - ggml_backend_buffer* initialize_buffer(ggml_backend_t backend, int k_head_dim, int v_head_dim, int num_attention_heads, int num_key_value_heads, uint32_t max_seq, ggml_type k_type, ggml_type v_type, bool fattn); + void slide_kv_layers(int layer_idx, int stride); + + ggml_backend_buffer* initialize_buffer(ggml_backend_t backend, int k_head_dim, int v_head_dim, int num_attention_heads, int num_key_value_heads, uint32_t max_seq, ggml_type k_type, ggml_type v_type, int batch_size, bool fattn); uint32_t reset_buffer(ggml_backend_buffer* buffer); void offload_cache(ggml_backend_t backend, ggml_backend_sched* sched, uint32_t n_tokens); diff --git a/src/cosyvoice-llm.cpp b/src/cosyvoice-llm.cpp index 5f4a4c7..1a8ba4b 100644 --- a/src/cosyvoice-llm.cpp +++ b/src/cosyvoice-llm.cpp @@ -1,6 +1,6 @@ #include "cosyvoice-internal.h" #include "cosyvoice-model.h" -#include "cosyvoice-llm-kv-cache.h" +#include "cosyvoice-kv-cache.h" #include #include @@ -27,7 +27,7 @@ static void build_causal_mask(ggml_fp16_t* mask, uint32_t n_batch, uint32_t seq_ } } -static ggml_tensor* build_qwen2_decoder_layer(const Qwen2DecoderLayer& layer, ggml_context* ctx0, ggml_cgraph* gf, ggml_tensor* hidden_states, ggml_tensor* position_ids, cosyvoice_llm_kv_cache& kv_cache, ggml_tensor* attention_mask, float rope_theta, float rms_norm_eps, int num_attention_heads, int num_key_value_heads, int layer_idx) +static ggml_tensor* build_qwen2_decoder_layer(const Qwen2DecoderLayer& layer, ggml_context* ctx0, ggml_cgraph* gf, ggml_tensor* hidden_states, ggml_tensor* position_ids, cosyvoice_kv_cache& kv_cache, ggml_tensor* attention_mask, float rope_theta, float rms_norm_eps, int num_attention_heads, int num_key_value_heads, int layer_idx) { auto residual = hidden_states; hidden_states = layer.input_layernorm.build_cgraph(ctx0, hidden_states, rms_norm_eps); @@ -62,7 +62,7 @@ static ggml_tensor* build_qwen2_decoder_layer(const Qwen2DecoderLayer& layer, gg query_states = ggml_cont(ctx0, query_states); kv_cache.update_cache(ctx0, gf, key_states, value_states, position_ids, static_cast(layer_idx)); - ggml_tensor* attn_output = kv_cache.attention_forward(ctx0, gf, query_states, key_states, value_states, attention_mask, layer_idx); + ggml_tensor* attn_output = kv_cache.attention_forward(ctx0, query_states, key_states, value_states, attention_mask); attn_output = ggml_reshape_2d( ctx0, attn_output, attn_output->ne[0] * attn_output->ne[1], @@ -99,7 +99,7 @@ bool cosyvoice_model_3::llm_prefill( uint32_t n_tokens ) { - auto kv_cache = &worker->kv_cache; + auto kv_cache = &worker->llm_kv_cache; auto total_len = n_tokens + kv_cache->cur_len; if (total_len > shared->params.n_max_seq - 1) return false; if (n_tokens > shared->params.n_batch) return false; @@ -203,7 +203,7 @@ bool cosyvoice_model_3::llm_prefill( bool cosyvoice_model_3::llm_decode(ggml_type type, const void* data) { - auto kv_cache = &worker->kv_cache; + auto kv_cache = &worker->llm_kv_cache; if (kv_cache->cur_len + 1 > shared->params.n_max_seq) return false; auto& position_ids = worker->position_ids; diff --git a/src/cosyvoice-loader.cpp b/src/cosyvoice-loader.cpp index 0ea3e85..33fb865 100644 --- a/src/cosyvoice-loader.cpp +++ b/src/cosyvoice-loader.cpp @@ -1,6 +1,6 @@ #include "cosyvoice-model.h" #include "cosyvoice-loader.h" -#include "cosyvoice-llm-kv-cache.h" +#include "cosyvoice-kv-cache.h" #include #include @@ -859,6 +859,16 @@ void cosyvoice_model_3::load(gguf_loader& loader) } while (false); } + if (shared->params.flow_use_flash_attn) + { + int heads = flow.decoder.estimator.transformer_blocks[0].attn.heads; + auto q = ggml_new_tensor_4d(worker->ctx0.get(), GGML_TYPE_F32, flow.decoder.estimator.transformer_blocks[0].attn.to_q.weight->ne[1] / heads, 1, heads, 2); + auto k = ggml_new_tensor_4d(worker->ctx0.get(), GGML_TYPE_F32, flow.decoder.estimator.transformer_blocks[0].attn.to_k.weight->ne[1] / heads, 1, heads, 2); + auto v = ggml_new_tensor_4d(worker->ctx0.get(), GGML_TYPE_F32, flow.decoder.estimator.transformer_blocks[0].attn.to_v.weight->ne[1] / heads, 1, heads, 2); + auto o = ggml_flash_attn_ext(worker->ctx0.get(), q, k, v, nullptr, 1.f / std::sqrt(static_cast(k->ne[0])), 0.f, 0.f); + shared->params.flow_use_flash_attn = ggml_backend_supports_op(backend, o); + } + for (auto& worker : std::span(workers, shared->params.n_workers)) { worker.nucleus_probs_capacity = static_cast(sampling.top_k * 2); @@ -867,29 +877,34 @@ void cosyvoice_model_3::load(gguf_loader& loader) worker.probs.reset(new float[llm.llm_decoder.weight->ne[1]]); worker.batch_buffer.reset(new char[shared->params.n_batch * std::max(llm.embed_tokens_weight->nb[1], llm.speech_embedding_weight->nb[1])]); - worker.kv_cache.build_kv_cache( + worker.llm_kv_cache.build_kv_cache( backend, - worker.kv_buffer, + worker.llm_kv_buffer, static_cast(llm.layers.size()), static_cast(llm.layers[0].self_attn.k_proj.weight->ne[1] / llm.num_key_value_heads), static_cast(llm.layers[0].self_attn.v_proj.weight->ne[1] / llm.num_key_value_heads), - llm.num_attention_heads, llm.num_key_value_heads, shared->params.n_max_seq, cv3_shared->k_type, cv3_shared->v_type, + 1, shared->params.llm_use_flash_attn ); - } - - if (shared->params.flow_use_flash_attn) - { - int heads = flow.decoder.estimator.transformer_blocks[0].attn.heads; - auto q = ggml_new_tensor_4d(worker->ctx0.get(), GGML_TYPE_F32, flow.decoder.estimator.transformer_blocks[0].attn.to_q.weight->ne[1] / heads, 1, heads, 2); - auto k = ggml_new_tensor_4d(worker->ctx0.get(), GGML_TYPE_F32, flow.decoder.estimator.transformer_blocks[0].attn.to_k.weight->ne[1] / heads, 1, heads, 2); - auto v = ggml_new_tensor_4d(worker->ctx0.get(), GGML_TYPE_F32, flow.decoder.estimator.transformer_blocks[0].attn.to_v.weight->ne[1] / heads, 1, heads, 2); - auto o = ggml_flash_attn_ext(worker->ctx0.get(), q, k, v, nullptr, 1.f / std::sqrt(static_cast(k->ne[0])), 0.f, 0.f); - shared->params.flow_use_flash_attn = ggml_backend_supports_op(backend, o); + + const auto dit_blocks = flow.decoder.estimator.transformer_blocks; + worker.dit_kv_cache.build_kv_cache( + backend, + worker.dit_kv_buffer, + static_cast(dit_blocks.size()* (flow.decoder.t_span.size() - 1)), + static_cast(dit_blocks[0].attn.to_k.weight->ne[1] / dit_blocks[0].attn.heads), + static_cast(dit_blocks[0].attn.to_v.weight->ne[1] / dit_blocks[0].attn.heads), + dit_blocks[0].attn.heads, + shared->params.n_max_seq, + GGML_TYPE_Q8_0, + GGML_TYPE_Q4_0, + 2, + shared->params.flow_use_flash_attn + ); } { @@ -910,7 +925,7 @@ void cosyvoice_model_3::load(gguf_loader& loader) { case COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED: case COSYVOICE_INFERENCE_BUFFER_POLICY_SHARED: - cv3_worker->token2wav_buffer.reset(worker->kv_buffer.get()); + cv3_worker->token2wav_buffer.reset(worker->llm_kv_buffer.get()); case COSYVOICE_INFERENCE_BUFFER_POLICY_DEDICATED: break; default: diff --git a/src/cosyvoice-model.cpp b/src/cosyvoice-model.cpp index d60633d..7f76c7b 100644 --- a/src/cosyvoice-model.cpp +++ b/src/cosyvoice-model.cpp @@ -66,7 +66,7 @@ cosyvoice_model_shared::cosyvoice_model_shared(const cosyvoice_context_params_v2 cosyvoice_worker_context::cosyvoice_worker_context(ggml_backend_t backend) : backend(backend), cpu_backend(backend), ctx0(ggml_init(ggml_init_params{ .mem_size = ggml_graph_overhead() * kCosyVoiceGraphSize, .no_alloc = true })), - gf(nullptr), llm_input(nullptr), llm_probs(nullptr), position_ids(nullptr), causal_mask(nullptr), kv_cache(), + gf(nullptr), llm_input(nullptr), llm_probs(nullptr), position_ids(nullptr), causal_mask(nullptr), llm_kv_cache(), status(GGML_STATUS_SUCCESS), prompt_crc32(0), sampler_seed(0), sampler(nullptr), sampler_ctx(nullptr), builtin_sampler_rng_policy(COSYVOICE_BUILTIN_SAMPLER_RNG_POLICY_RESET_PER_SESSION), nucleus_probs_capacity(0), nucleus_probs_len(0) {} cosyvoice_model::cosyvoice_model(ggml_backend_t backend, const cosyvoice_context_params_v2_cpp& params) @@ -222,7 +222,7 @@ cosyvoice_model::~cosyvoice_model() { case COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED: case COSYVOICE_INFERENCE_BUFFER_POLICY_SHARED: - worker->kv_buffer.release(); + worker->llm_kv_buffer.release(); } worker->~cosyvoice_worker_context(); @@ -251,7 +251,7 @@ void cosyvoice_model::empty_buffer_cache() shared->rand_noise_len = 0; } - worker->kv_cache.clear_offloaded_cache(); + worker->llm_kv_cache.clear_offloaded_cache(); worker->prompt_crc32 = 0; } @@ -262,8 +262,8 @@ void cosyvoice_model_3::empty_buffer_cache() if (cv3_worker->orig_max_seq_len != shared->params.n_max_seq) { shared->params.n_max_seq = cv3_worker->orig_max_seq_len; - worker->kv_buffer.reset( - worker->kv_cache.initialize_buffer( + worker->llm_kv_buffer.reset( + worker->llm_kv_cache.initialize_buffer( worker->backend.get(), static_cast(cv3_shared->llm.layers[0].self_attn.k_proj.weight->ne[1] / cv3_shared->llm.num_key_value_heads), static_cast(cv3_shared->llm.layers[0].self_attn.v_proj.weight->ne[1] / cv3_shared->llm.num_key_value_heads), @@ -272,6 +272,7 @@ void cosyvoice_model_3::empty_buffer_cache() shared->params.n_max_seq, cv3_shared->k_type, cv3_shared->v_type, + 1, shared->params.llm_use_flash_attn)); switch (shared->params.inference_buffer_policy) @@ -279,7 +280,7 @@ void cosyvoice_model_3::empty_buffer_cache() case COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED: case COSYVOICE_INFERENCE_BUFFER_POLICY_SHARED: cv3_worker->token2wav_buffer.release(); - cv3_worker->token2wav_buffer.reset(worker->kv_buffer.get()); + cv3_worker->token2wav_buffer.reset(worker->llm_kv_buffer.get()); break; case COSYVOICE_INFERENCE_BUFFER_POLICY_DEDICATED: cv3_worker->token2wav_buffer.reset(); @@ -295,9 +296,9 @@ void cosyvoice_model_3::get_memory_usage(cosyvoice_memory_usage_t* usage) usage->parameters = ggml_backend_buffer_get_size(shared->buffer.get()); usage->buffers = ggml_backend_sched_get_buffer_size(worker->sched.get(), worker->backend.get()); usage->cpu_buffers = ggml_backend_sched_get_buffer_size(worker->sched.get(), worker->cpu_backend.get()); - usage->offloaded_kv_cache = worker->kv_cache.get_offloaded_cache_size(); + usage->offloaded_kv_cache = worker->llm_kv_cache.get_offloaded_cache_size(); usage->random_noise = sizeof(float) * shared->rand_noise_len; - usage->kv_cache = worker->kv_buffer.get() ? ggml_backend_buffer_get_size(worker->kv_buffer.get()) : 0; + usage->kv_cache = worker->llm_kv_buffer.get() ? ggml_backend_buffer_get_size(worker->llm_kv_buffer.get()) : 0; usage->token2wav = cv3_worker->token2wav_buffer.get() ? ggml_backend_buffer_get_size(cv3_worker->token2wav_buffer.get()) : 0; } @@ -311,9 +312,9 @@ void cosyvoice_model_3::get_total_memory_usage(cosyvoice_memory_usage_t* usage) usage->buffers = ggml_backend_sched_get_buffer_size(worker->sched.get(), worker->backend.get()); usage->cpu_buffers = ggml_backend_sched_get_buffer_size(worker->sched.get(), worker->cpu_backend.get()); - usage->offloaded_kv_cache = worker->kv_cache.get_offloaded_cache_size(); + usage->offloaded_kv_cache = worker->llm_kv_cache.get_offloaded_cache_size(); usage->random_noise = sizeof(float) * shared->rand_noise_len; - usage->kv_cache = worker->kv_buffer.get() ? ggml_backend_buffer_get_size(worker->kv_buffer.get()) : 0; + usage->kv_cache = worker->llm_kv_buffer.get() ? ggml_backend_buffer_get_size(worker->llm_kv_buffer.get()) : 0; usage->token2wav = cv3_worker->token2wav_buffer.get() ? ggml_backend_buffer_get_size(cv3_worker->token2wav_buffer.get()) : 0; } } @@ -323,12 +324,29 @@ void cosyvoice_model_3::reset_shared_buffer(ggml_backend_buffer* new_buffer) cv3_worker->token2wav_buffer.reset(new_buffer); if (shared->params.inference_buffer_policy != COSYVOICE_INFERENCE_BUFFER_POLICY_DEDICATED) { - worker->kv_buffer.release(); - worker->kv_buffer.reset(new_buffer); - shared->params.n_max_seq = worker->kv_cache.reset_buffer(new_buffer); + worker->llm_kv_buffer.release(); + worker->llm_kv_buffer.reset(new_buffer); + shared->params.n_max_seq = worker->llm_kv_cache.reset_buffer(new_buffer); } } +uint32_t cosyvoice_model_3::get_chunk_tokens() +{ + return 25 + cv3_shared->flow.pre_lookahead_layer.pre_lookahead_len; +} + +uint32_t cosyvoice_model_3::get_flow_overlap_tokens() +{ + // TODO: calculate the actual overlap tokens + return 0; +} + +uint32_t cosyvoice_model_3::get_hift_overlap_tokens() +{ + // TODO: calculate the actual overlap tokens + return 0; +} + cosyvoice_3_worker_context::cosyvoice_3_worker_context() : ctx1(ggml_init(ggml_init_params{ .mem_size = ggml_tensor_overhead() * 4, .no_alloc = true })) {} @@ -555,12 +573,12 @@ uint32_t cosyvoice_model::get_n_workers() uint32_t cosyvoice_model::llm_get_kv_cache_len() { - return worker->kv_cache.cur_len; + return worker->llm_kv_cache.cur_len; } bool cosyvoice_model::llm_set_kv_cache_len(uint32_t len) { - auto& cur_len = worker->kv_cache.cur_len; + auto& cur_len = worker->llm_kv_cache.cur_len; if (len <= cur_len) { cur_len = len; @@ -572,17 +590,15 @@ bool cosyvoice_model::llm_set_kv_cache_len(uint32_t len) void cosyvoice_model::llm_offload_kv_cache() { auto sched = worker->sched.get(); - auto& kv_cache = worker->kv_cache; + auto& kv_cache = worker->llm_kv_cache; ggml_backend_sched_reset(sched); kv_cache.offload_cache(worker->backend.get(), sched, kv_cache.cur_len); - ggml_backend_sched_synchronize(sched); } void cosyvoice_model::llm_load_kv_cache() { auto sched = worker->sched.get(); - auto& kv_cache = worker->kv_cache; + auto& kv_cache = worker->llm_kv_cache; ggml_backend_sched_reset(sched); kv_cache.offload_cache(worker->backend.get(), sched, kv_cache.cur_len); - ggml_backend_sched_synchronize(sched); } diff --git a/src/cosyvoice-model.h b/src/cosyvoice-model.h index 7e2ad42..4cc4c87 100644 --- a/src/cosyvoice-model.h +++ b/src/cosyvoice-model.h @@ -3,7 +3,7 @@ #include "cosyvoice-modules.h" #include "cosyvoice-lowlevel.h" #include "cosyvoice-interface.h" -#include "cosyvoice-llm-kv-cache.h" +#include "cosyvoice-kv-cache.h" #include @@ -60,7 +60,8 @@ struct cosyvoice_worker_context std::unique_ptr full_position_ids; std::unique_ptr causal_mask_buffer; - cosyvoice_llm_kv_cache kv_cache; + cosyvoice_kv_cache llm_kv_cache; + cosyvoice_kv_cache dit_kv_cache; ggml_status status; uint32_t prompt_crc32; @@ -70,12 +71,14 @@ struct cosyvoice_worker_context void* sampler_ctx; cosyvoice_builtin_sampler_rng_policy_t builtin_sampler_rng_policy; + std::vector flow_cache; std::unique_ptr batch_buffer; std::unique_ptr nucleus_probs; uint32_t nucleus_probs_capacity; int nucleus_probs_len; std::unique_ptr probs; - ggml_backend_buffer_ptr kv_buffer; + ggml_backend_buffer_ptr llm_kv_buffer; + ggml_backend_buffer_ptr dit_kv_buffer; }; struct cosyvoice_model : virtual cosyvoice_model_context, virtual cosyvoice_object_ref_counter @@ -184,8 +187,10 @@ struct cosyvoice_model_3 : cosyvoice_model bool llm_job(const int* text, uint32_t text_len, cosyvoice_prompt_t prompt); bool llm_job_ext(const int* text, uint32_t text_len, cosyvoice_prompt_t prompt, uint32_t max_new_tokens, bool* final); bool token2wav(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, cosyvoice_generated_speech_ptr result); - bool token2wav_ext(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* speech_offset, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result); + bool token2wav_ext(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* offset, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result); uint32_t get_chunk_tokens(); + uint32_t get_flow_overlap_tokens(); + uint32_t get_hift_overlap_tokens(); void empty_buffer_cache(); void get_memory_usage(cosyvoice_memory_usage_t* usage); diff --git a/src/cosyvoice-modules.h b/src/cosyvoice-modules.h index 286713b..110d1be 100644 --- a/src/cosyvoice-modules.h +++ b/src/cosyvoice-modules.h @@ -1,6 +1,7 @@ #pragma once #include "cosyvoice-internal.h" +#include "cosyvoice-kv-cache.h" #include "ggml-fft.h" #include @@ -130,7 +131,7 @@ struct Attention : Module void OnLoad(const gguf_loader& loader, const std::string& prefix); - ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* position_ids, int64_t cut_len, ggml_tensor* attn_mask = nullptr) const; + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* position_ids, int64_t cut_len, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, int layer_idx) const; }; struct FeedForward : Module @@ -153,7 +154,7 @@ struct DiTBlock : Module void OnLoad(const gguf_loader& loader, const std::string& prefix); - ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* time_emb, ggml_tensor* position_ids, int64_t cut_len, ggml_tensor* attn_mask = nullptr) const; + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* time_emb, ggml_tensor* position_ids, int64_t cut_len, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, int layer_idx) const; }; struct AdaLayerNorm_Final : Module @@ -182,7 +183,7 @@ struct DiT : Module void OnLoad(const gguf_loader& loader, const std::string& prefix); - ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* mu, ggml_tensor* t, ggml_tensor* spks, ggml_tensor* cond, int64_t cut_len, ggml_tensor*& position_ids, ggml_backend_op_capabilities capabilities, bool streaming, ggml_tensor*& ref_attn_mask) const; + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* mu, ggml_tensor* t, ggml_tensor* spks, ggml_tensor* cond, int64_t cut_len, ggml_tensor*& position_ids, ggml_backend_op_capabilities capabilities, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, int step) const; }; struct CausalConditionalCFM : Module @@ -205,7 +206,7 @@ struct CausalConditionalCFM : Module DiTContext prepare_context(ggml_context* ctx, ggml_tensor* mu, ggml_tensor* spks, ggml_tensor* cond) const; std::array get_t_and_dt(ggml_context* ctx, int step) const; - ggml_tensor* build_cgraph_one_step(ggml_context* ctx, const DiTContext& ditctx, int step, ggml_backend_op_capabilities capabilities, int64_t cut_len, ggml_tensor*& t_tensor, ggml_tensor*& position_ids, bool streaming, ggml_tensor*& ref_attn_mask) const; + ggml_tensor* build_cgraph_one_step(ggml_context* ctx, const DiTContext& ditctx, int step, ggml_backend_op_capabilities capabilities, int64_t cut_len, ggml_tensor*& t_tensor, ggml_tensor*& position_ids, ggml_cgraph* gf, cosyvoice_kv_cache* kv_cache) const; }; struct PreLookaheadLayer : Module @@ -217,7 +218,7 @@ struct PreLookaheadLayer : Module void OnLoad(const gguf_loader& loader, const std::string& prefix); - ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* inputs, bool streaming) const; + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* inputs, bool streaming, uint32_t cut_len) const; }; struct CausalMaskedDiffWithDiT : Module @@ -239,7 +240,7 @@ struct CausalMaskedDiffWithDiT : Module int64_t cut_len; }; - EncodeResult build_cgraph_encode(ggml_context* ctx, ggml_tensor* token, ggml_tensor* prompt_token, ggml_tensor* prompt_feat, ggml_tensor* embedding, ggml_backend_op_capabilities capabilities, bool streaming = false) const; + EncodeResult build_cgraph_encode(ggml_context* ctx, ggml_tensor* token, ggml_tensor* prompt_token, ggml_tensor* prompt_feat, ggml_tensor* embedding, ggml_backend_op_capabilities capabilities, uint32_t cut_len = 0, bool streaming = false) const; }; struct CausalConv1dBase : Conv1d @@ -345,7 +346,6 @@ struct CausalHiFTGenerator : Module void set_rand_ini(const float* data) const; - constexpr static int conv_pre_look_right = 4; float lrelu_slope; int scale_factor; int nb_harmonics; diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index 9fa8fd6..143500b 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -1,6 +1,6 @@ #include "cosyvoice-internal.h" #include "cosyvoice-model.h" -#include "cosyvoice-llm-kv-cache.h" +#include "cosyvoice-kv-cache.h" #include "cosyvoice-text-chunk.h" #include "ggml-cpu-flag.h" @@ -184,6 +184,7 @@ bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoic // model state consistent, but skip llm_accept_token so // it does not appear in the output. cur = speech_emb + token_id * speech_row_size; + --n; continue; } } @@ -310,9 +311,8 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float return token2wav_ext(token_ids, n_tokens, speed, prompt, nullptr, false, true, result); } -bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* speech_offset_ptr, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result) +bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* offset, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result) { - const uint32_t speech_offset = speech_offset_ptr ? *speech_offset_ptr : 0; const auto& params = shared->params; auto& sched = worker->sched; auto& flow = cv3_shared->flow; @@ -334,10 +334,12 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f ggml_tensor* prompt_token = ggml_new_tensor_1d(ctx0.get(), GGML_TYPE_I32, prompt->flow_prompt_speech_tokens.second); ggml_tensor* prompt_feat = ggml_new_tensor_2d(ctx0.get(), GGML_TYPE_F32, prompt->prompt_speech_feat.shape[1], prompt->prompt_speech_feat.shape[0]); ggml_tensor* embedding = ggml_new_tensor_2d(ctx0.get(), GGML_TYPE_F32, prompt->flow_embedding.shape[1], prompt->flow_embedding.shape[0]); + auto kv_cache = streaming ? &worker->dit_kv_cache : nullptr; // Phase 1: Flow encode ggml_cgraph* gf = new_cgraph(ctx0.get()); - auto [mu, spks, conds, cut_len] = flow.build_cgraph_encode(ctx0.get(), token, prompt_token, prompt_feat, embedding, op_caps, streaming); + auto [mu, spks, conds, cut_len] = flow.build_cgraph_encode(ctx0.get(), token, prompt_token, prompt_feat, embedding, op_caps, offset ? *offset : 0, streaming); + if (offset) *offset += mu->ne[1] / flow.token_mel_ratio; auto ditctx = flow.decoder.prepare_context(ctx1.get(), mu, spks, conds); do { @@ -361,6 +363,7 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f } } while (false); + int position_offset = kv_cache ? static_cast(kv_cache->cur_len) : 0; uint32_t noise_len = static_cast(ggml_nelements(ditctx.x)); float* noise_buffer = shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_BEFORE_FLOW, noise_len, nullptr, shared->noise_callback_ctx); ggml_backend_tensor_set_async(backend.get(), ditctx.x, noise_buffer, 0, ggml_nbytes(ditctx.x)); @@ -368,8 +371,7 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f // Phase 2: Flow decode steps ggml_tensor* t_leaf; ggml_tensor* position_ids; - ggml_tensor* attn_mask = nullptr; - auto feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, 1, op_caps, cut_len, t_leaf, position_ids, streaming, attn_mask); + auto feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, 1, op_caps, cut_len, t_leaf, position_ids, gf, kv_cache); ggml_build_forward_expand(gf, feat); set_graph_backends(gf, sched.get(), backend.get(), cpu_backend.get(), op_caps); ggml_backend_sched_synchronize(sched.get()); @@ -380,18 +382,7 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f { auto cur_row = reinterpret_cast(position_ids->data) + i * position_ids->ne[0]; for (int32_t j = 0; j < position_ids->ne[0]; ++j) - cur_row[j] = j; - } - if (attn_mask) - { - const auto cs = flow.decoder.estimator.static_chunk_size; - for (int64_t i = 0; i < attn_mask->ne[0]; ++i) - { - int64_t block_end = ((i / cs) + 1) * cs; - auto row = reinterpret_cast(attn_mask->data) + i * attn_mask->ne[1]; - for (int64_t j = 0; j < attn_mask->ne[1]; ++j) - row[j] = j < block_end ? 0 : 0xFC00; - } + cur_row[j] = j + position_offset; } ggml_backend_tensor_set_async(backend.get(), token, token_ids, 0, token->nb[1]); @@ -419,8 +410,7 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f ggml_reset(ctx0.get()); ggml_backend_sched_reset(sched.get()); gf = new_cgraph(ctx0.get()); - attn_mask = nullptr; - feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, 2, op_caps, cut_len, t_leaf, position_ids, streaming, attn_mask); + feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, 2, op_caps, cut_len, t_leaf, position_ids, gf, kv_cache); ggml_build_forward_expand(gf, feat); set_graph_backends(gf, sched.get(), backend.get(), cpu_backend.get(), op_caps); @@ -431,28 +421,19 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f { auto cur_row = reinterpret_cast(position_ids->data) + i * position_ids->ne[0]; for (int32_t j = 0; j < position_ids->ne[0]; ++j) - cur_row[j] = j; - } - if (attn_mask) - { - const auto cs = flow.decoder.estimator.static_chunk_size; - for (int64_t i = 0; i < attn_mask->ne[0]; ++i) - { - int64_t block_end = ((i / cs) + 1) * cs; - auto row = reinterpret_cast(attn_mask->data) + i * attn_mask->ne[1]; - for (int64_t j = 0; j < attn_mask->ne[1]; ++j) - row[j] = j < block_end ? 0 : 0xFC00; - } + cur_row[j] = j + position_offset; } worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); if (worker->status != GGML_STATUS_SUCCESS) return false; auto scale_node = feat->src[1]; GGML_ASSERT(scale_node->op == GGML_OP_SCALE); - + int num_blocks = static_cast(flow.decoder.estimator.transformer_blocks.size()); for (int step = 3; step != flow.decoder.t_span.size() - 1; ++step) { ggml_backend_tensor_copy_async(backend.get(), backend.get(), feat, ditctx.x); + if (kv_cache) + kv_cache->slide_kv_layers((step - 2) * num_blocks, num_blocks); auto [t, dt] = flow.decoder.get_t_and_dt(ctx0.get(), step); reinterpret_cast(t_leaf->op_params)[op_caps.fill ? 0 : 1] = t; @@ -467,8 +448,7 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f ggml_reset(ctx0.get()); ggml_backend_sched_reset(sched.get()); gf = new_cgraph(ctx0.get()); - attn_mask = nullptr; - feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, static_cast(flow.decoder.t_span.size() - 1), op_caps, cut_len, t_leaf, position_ids, streaming, attn_mask); + feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, static_cast(flow.decoder.t_span.size() - 1), op_caps, cut_len, t_leaf, position_ids, gf, kv_cache); ggml_build_forward_expand(gf, feat); set_graph_backends(gf, sched.get(), backend.get(), cpu_backend.get(), op_caps); @@ -479,37 +459,59 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f { auto cur_row = reinterpret_cast(position_ids->data) + i * position_ids->ne[0]; for (int32_t j = 0; j < position_ids->ne[0]; ++j) - cur_row[j] = j; - } - if (attn_mask) - { - const auto cs = flow.decoder.estimator.static_chunk_size; - for (int64_t i = 0; i < attn_mask->ne[0]; ++i) - { - int64_t block_end = ((i / cs) + 1) * cs; - auto row = reinterpret_cast(attn_mask->data) + i * attn_mask->ne[1]; - for (int64_t j = 0; j < attn_mask->ne[1]; ++j) - row[j] = j < block_end ? 0 : 0xFC00; - } + cur_row[j] = j + position_offset; } worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); if (worker->status != GGML_STATUS_SUCCESS) return false; // Phase 3: Copy flow output to speech_feat (persistent buffer) ggml_reset(ctx1.get()); - ggml_tensor* speech_feat = ggml_new_tensor(ctx1.get(), feat->type, GGML_MAX_DIMS, feat->ne); - ggml_backend_tensor_alloc(token2wav_buffer.get(), speech_feat, - ggml_backend_buffer_get_base(token2wav_buffer.get())); - ggml_backend_tensor_copy_async(backend.get(), backend.get(), feat, speech_feat); + auto cache_length = static_cast(worker->flow_cache.size() / feat->ne[0]); + ggml_tensor* speech_feat = ggml_new_tensor_2d(ctx1.get(), feat->type, feat->ne[0], feat->ne[1] + cache_length); + if (cache_length != 0) + { + auto buft = ggml_backend_get_default_buffer_type(backend.get()); + auto size = ggml_backend_buft_get_alloc_size(buft, speech_feat); + if (size > ggml_backend_buffer_get_size(token2wav_buffer.get())) + reset_shared_buffer(ggml_backend_alloc_ctx_tensors_from_buft(ctx1.get(), buft)); + else + ggml_backend_tensor_alloc(token2wav_buffer.get(), speech_feat, + ggml_backend_buffer_get_base(token2wav_buffer.get())); + + ggml_backend_tensor_set_async(backend.get(), speech_feat, worker->flow_cache.data(), 0, speech_feat->nb[2]); + ggml_backend_tensor_copy_async(backend.get(), backend.get(), feat, + ggml_view_2d(ctx0.get(), speech_feat, speech_feat->ne[0], feat->ne[1], speech_feat->nb[1], speech_feat->nb[1] * cache_length)); + } + else + { + ggml_backend_tensor_alloc(token2wav_buffer.get(), speech_feat, + ggml_backend_buffer_get_base(token2wav_buffer.get())); + ggml_backend_tensor_copy_async(backend.get(), backend.get(), feat, speech_feat); + } + if (!finalize) + { + auto n_cached_elements = cache_length * feat->ne[0]; + worker->flow_cache.resize(n_cached_elements + static_cast(ggml_nelements(feat))); + ggml_backend_tensor_get_async(backend.get(), feat, worker->flow_cache.data() + n_cached_elements, 0, feat->nb[2]); + } + + if (kv_cache) + if (finalize) + kv_cache->cur_len = 0; + else + kv_cache->cur_len += static_cast(position_ids->ne[0]); ggml_reset(ctx0.get()); ggml_backend_sched_reset(sched.get()); gf = new_cgraph(ctx0.get()); + speech_feat = ggml_permute(ctx0.get(), speech_feat, 1, 0, 2, 3); if (auto ne0 = static_cast(speech_feat->ne[0] / speed); ne0 != speech_feat->ne[0]) speech_feat = ggml_interpolate(ctx0.get(), speech_feat, ne0, speech_feat->ne[1], speech_feat->ne[2], speech_feat->ne[3], GGML_SCALE_MODE_BILINEAR); + else + speech_feat = ggml_cont(ctx0.get(), speech_feat); auto [generated_speech, noise] = hift.build_cgraph(ctx0.get(), speech_feat, finalize); ggml_build_forward_expand(gf, generated_speech); @@ -525,28 +527,22 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f noise_buffer = shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_BEFORE_HIFT, noise_len, nullptr, shared->noise_callback_ctx); ggml_backend_tensor_set_async(backend.get(), noise, noise_buffer, 0, noise->nb[2]); - const auto total_speech_len = static_cast(generated_speech->ne[0]); - result->data = reinterpret_cast(generated_speech->data) + speech_offset; - result->length = speech_offset < total_speech_len ? total_speech_len - speech_offset : 0; + result->data = reinterpret_cast(generated_speech->data); + result->length = static_cast(generated_speech->ne[0]); worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_AFTER_HIFT, noise_len, noise_buffer, shared->noise_callback_ctx); - if (worker->status == GGML_STATUS_SUCCESS && speech_offset_ptr) - *speech_offset_ptr += result->length; + if (finalize) + worker->flow_cache.clear(); if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) { ggml_backend_sched_reset(sched.get()); - worker->kv_cache.load_cache(backend.get(), sched.get()); + worker->llm_kv_cache.load_cache(backend.get(), sched.get()); } return worker->status == GGML_STATUS_SUCCESS; } -uint32_t cosyvoice_model_3::get_chunk_tokens() -{ - return 25 + cv3_shared->flow.pre_lookahead_layer.pre_lookahead_len; -} - struct cosyvoice_tts_context : cosyvoice_tokenization_result_impl, cosyvoice_prompt, std::string { cosyvoice_tts_context(cosyvoice_context_t ctx, cosyvoice_prompt_t prompt) diff --git a/src/cosyvoice.cpp b/src/cosyvoice.cpp index b41c996..bc8922a 100644 --- a/src/cosyvoice.cpp +++ b/src/cosyvoice.cpp @@ -302,9 +302,9 @@ bool cosyvoice_token2wav(cosyvoice_context_t ctx, const int* token_ids, uint32_t return ctx->token2wav(token_ids, n_tokens, speed, prompt, generated_speech); } -bool cosyvoice_token2wav_ext(cosyvoice_context_t ctx, const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* speech_offset_ptr, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result) +bool cosyvoice_token2wav_ext(cosyvoice_context_t ctx, const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* offset, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result) { - return ctx->token2wav_ext(token_ids, n_tokens, speed, prompt, speech_offset_ptr, streaming, finalize, result); + return ctx->token2wav_ext(token_ids, n_tokens, speed, prompt, offset, streaming, finalize, result); } static bool check_length(cosyvoice_prompt_t prompt, uint32_t text_len, uint32_t n_max_seq) @@ -359,31 +359,49 @@ bool cosyvoice_tts_stream(cosyvoice_context_t ctx, const int* text, uint32_t tex const auto chunk_tokens = ctx->get_chunk_tokens(); ctx->llm_clear_accepted_tokens(); - uint32_t speech_offset = 0; + uint32_t offset = 0; + uint32_t n_tokens = 0; + uint32_t last_length = 0; bool final = false; do { - if (!ctx->llm_job_ext(text, text_len, prompt, chunk_tokens, &final)) - return false; - text = nullptr; - const uint32_t n_tokens = ctx->llm_get_n_accepted_tokens(); - const int* tokens = ctx->llm_get_accepted_tokens(); + if (!final) + if (!ctx->llm_job_ext(text, text_len, prompt, chunk_tokens + 1, &final)) + return false; - if (!final && params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) + n_tokens = std::min(n_tokens + chunk_tokens, ctx->llm_get_n_accepted_tokens()); + text = nullptr; + if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) ctx->llm_offload_kv_cache(); cosyvoice_generated_speech result = {}; - if (!ctx->token2wav_ext(tokens, n_tokens, speed, prompt, &speech_offset, true, final, &result)) + if (!ctx->token2wav_ext(ctx->llm_get_accepted_tokens(), n_tokens, speed, prompt, &offset, true, ctx->llm_get_n_accepted_tokens() == n_tokens, &result)) return false; - if (result.data && result.length > 0 - && !callback(result.data, result.length, user_data)) + if (result.data && result.length > last_length + && !callback(result.data + last_length, result.length - last_length, user_data)) return false; - } while (!final); + last_length = result.length; + } while (n_tokens != ctx->llm_get_n_accepted_tokens()); return true; } +uint32_t cosyvoice_get_chunk_tokens(cosyvoice_context_t ctx) +{ + return ctx->get_chunk_tokens(); +} + +uint32_t cosyvoice_get_flow_overlap_tokens(cosyvoice_context_t ctx) +{ + return ctx->get_flow_overlap_tokens(); +} + +uint32_t cosyvoice_get_hift_overlap_tokens(cosyvoice_context_t ctx) +{ + return ctx->get_hift_overlap_tokens(); +} + ggml_status cosyvoice_get_last_status(cosyvoice_context_t ctx) { return ctx->get_last_status(); From e507d1f469c17c29aa47728321f65d7d0be384b2 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Thu, 9 Jul 2026 14:30:38 +0800 Subject: [PATCH 05/34] cosyvoice: optimize KV cache offload/load and add multi-batch support --- src/cosyvoice-kv-cache.cpp | 158 +++++++++++++++++++++---------------- src/cosyvoice-kv-cache.h | 4 +- src/cosyvoice-model.cpp | 7 +- 3 files changed, 93 insertions(+), 76 deletions(-) diff --git a/src/cosyvoice-kv-cache.cpp b/src/cosyvoice-kv-cache.cpp index 660ad2a..7fd39e5 100644 --- a/src/cosyvoice-kv-cache.cpp +++ b/src/cosyvoice-kv-cache.cpp @@ -25,7 +25,7 @@ void cosyvoice_kv_cache::build_kv_cache( this->k_type = k_type; this->v_type = v_type; offloaded_cache = nullptr; - this->num_attention_heads = num_attention_heads; + this->num_heads = num_key_value_heads; kv_cache_layers = new kv_cache_layer[layers]; ggml_init_params params = { @@ -34,10 +34,10 @@ void cosyvoice_kv_cache::build_kv_cache( }; ctx = ggml_init(params); - buffer.reset(initialize_buffer(backend, k_head_dim, v_head_dim, num_attention_heads, num_key_value_heads, max_seq, k_type, v_type, batch_size, fattn)); + buffer.reset(initialize_buffer(backend, k_head_dim, v_head_dim, num_key_value_heads, max_seq, k_type, v_type, batch_size, fattn)); } -ggml_backend_buffer* cosyvoice_kv_cache::initialize_buffer(ggml_backend_t backend, int k_head_dim, int v_head_dim, int num_attention_heads, int num_key_value_heads, uint32_t max_seq, ggml_type k_type, ggml_type v_type, int batch_size, bool fattn) +ggml_backend_buffer* cosyvoice_kv_cache::initialize_buffer(ggml_backend_t backend, int k_head_dim, int v_head_dim, int num_key_value_heads, uint32_t max_seq, ggml_type k_type, ggml_type v_type, int batch_size, bool fattn) { int64_t k_ne[4] = { k_head_dim, max_seq, num_key_value_heads, batch_size }; int64_t v_ne[4] = { v_head_dim, max_seq, num_key_value_heads, batch_size }; @@ -95,13 +95,13 @@ uint32_t cosyvoice_kv_cache::reset_buffer(ggml_backend_buffer* buffer) struct offloaded_kv_cache { uint32_t len; + size_t buffer_size; ggml_context* ctx; struct offloaded_kv_layer { - ggml_tensor* k_view; - ggml_tensor* v_view; - void* k; - void* v; + ggml_tensor* v_tensor; + char* k; + char* v; } offloaded_kv_layers[]; }; #pragma pack(pop) @@ -109,7 +109,10 @@ struct offloaded_kv_cache void cosyvoice_kv_cache::offload_cache(ggml_backend_t backend, ggml_backend_sched* sched, uint32_t n_tokens) { char* buffer_base; - size_t nbytes = kv_cache_layers[0].k->ne[0] * sizeof(float) * n_tokens * kv_cache_layers[0].v->ne[2]; + const auto batch_size = kv_cache_layers[0].k->ne[3]; + const size_t k_head_nbytes = kv_cache_layers[0].k->nb[1] * n_tokens; + const size_t v_head_nbytes = fattn ? kv_cache_layers[0].v->nb[1] * n_tokens : ggml_row_size(kv_cache_layers[0].v->type, n_tokens) * kv_cache_layers[0].v->ne[1]; + const size_t kv_nbytes = (k_head_nbytes + v_head_nbytes) * num_heads * batch_size; if (!offloaded_cache) { ggml_init_params params = @@ -117,109 +120,120 @@ void cosyvoice_kv_cache::offload_cache(ggml_backend_t backend, ggml_backend_sche .mem_size = (ggml_tensor_overhead() + ggml_graph_overhead()) * layers * 4, .no_alloc = true }; - offloaded_cache = reinterpret_cast(malloc(sizeof(offloaded_kv_cache::offloaded_kv_layer) * layers + sizeof(uint32_t) + sizeof(ggml_context*))); + offloaded_cache = reinterpret_cast(malloc(sizeof(offloaded_kv_cache::offloaded_kv_layer) * layers + sizeof(uint32_t) + sizeof(ggml_context*) + sizeof(size_t))); offloaded_cache->ctx = ggml_init(params); + offloaded_cache->buffer_size = 0; goto alloc_buffer; } else { ggml_reset(offloaded_cache->ctx); - if (offloaded_cache->len < n_tokens) + if (offloaded_cache->buffer_size < kv_nbytes * layers) { free(offloaded_cache->offloaded_kv_layers[0].k); alloc_buffer: - buffer_base = reinterpret_cast(malloc(nbytes * 2 * layers)); + offloaded_cache->buffer_size = kv_nbytes * layers; + buffer_base = reinterpret_cast(malloc(offloaded_cache->buffer_size)); } else buffer_base = reinterpret_cast(offloaded_cache->offloaded_kv_layers[0].k); } offloaded_cache->len = n_tokens; - auto gf = ggml_new_graph_custom(offloaded_cache->ctx, layers * 4, false); - ggml_tensor k_placeholder_tensor{ .ne{ kv_cache_layers[0].k->ne[0], n_tokens, kv_cache_layers[0].k->ne[2], 1 } }; - ggml_tensor v_placeholder_tensor{ .ne{ fattn ? kv_cache_layers[0].v->ne[0] : n_tokens, fattn ? n_tokens : kv_cache_layers[0].v->ne[0], k_placeholder_tensor.ne[2], 1 } }; + const auto total_heads = num_heads * batch_size; for (int i = 0; i != layers; ++i) { auto& offloaded_layer = offloaded_cache->offloaded_kv_layers[i]; auto& layer = kv_cache_layers[i]; - ggml_tensor* k_view = ggml_view_3d(offloaded_cache->ctx, layer.k, layer.k->ne[0], n_tokens, layer.k->ne[2], layer.k->nb[1], layer.k->nb[2], 0); - ggml_tensor* v_view; - if (fattn) - v_view = ggml_view_3d(offloaded_cache->ctx, layer.v, layer.v->ne[0], n_tokens, layer.v->ne[2], layer.v->nb[1], layer.v->nb[2], 0); - else - v_view = ggml_view_3d(offloaded_cache->ctx, layer.v, n_tokens, layer.v->ne[1], layer.v->ne[2], layer.v->nb[1], layer.v->nb[2], 0); - - // The ggml_cast uses GGML_OP_CPY to implement the cast, which has a bug that the non-contiguous tensor causes the bad result. - // But most of the backends use cpy to implement the dup, so we have to use ggml_cast and then set the op to GGML_OP_DUP to avoid the bug. - offloaded_layer.k_view = ggml_cast(offloaded_cache->ctx, k_view, GGML_TYPE_F32); - offloaded_layer.v_view = ggml_cast(offloaded_cache->ctx, v_view, GGML_TYPE_F32); - offloaded_layer.k_view->op = GGML_OP_DUP; - offloaded_layer.v_view->op = GGML_OP_DUP; - // IMPORTANT: src[1] must be set to the placeholder tensor - // or it will cause the bug that the non-contiguous tensor causes the bad result. - offloaded_layer.k_view->src[1] = &k_placeholder_tensor; - offloaded_layer.v_view->src[1] = &v_placeholder_tensor; - - ggml_build_forward_expand(gf, offloaded_layer.k_view); - ggml_build_forward_expand(gf, offloaded_layer.v_view); + offloaded_layer.k = buffer_base + i * kv_nbytes; + offloaded_layer.v = offloaded_layer.k + k_head_nbytes * total_heads; + for (int h = 0; h != total_heads; ++h) + { + ggml_backend_tensor_get_async(backend, layer.k, offloaded_layer.k + h * k_head_nbytes, h * layer.k->nb[2], k_head_nbytes); + if (fattn) + ggml_backend_tensor_get_async(backend, layer.v, offloaded_layer.v + h * v_head_nbytes, h * layer.v->nb[2], v_head_nbytes); + } } - ggml_backend_sched_alloc_graph(sched, gf); - ggml_backend_sched_graph_compute_async(sched, gf); - - for (int i = 0; i != layers; ++i) + if (!fattn) { - auto& offloaded_layer = offloaded_cache->offloaded_kv_layers[i]; - offloaded_layer.k = buffer_base; - offloaded_layer.v = buffer_base + nbytes; - ggml_backend_tensor_get_async(backend, offloaded_layer.k_view, offloaded_layer.k, 0, nbytes); - ggml_backend_tensor_get_async(backend, offloaded_layer.v_view, offloaded_layer.v, 0, nbytes); - buffer_base += nbytes * 2; + ggml_reset(offloaded_cache->ctx); + auto gf = ggml_new_graph_custom(offloaded_cache->ctx, layers * 3, false); + for (int i = 0; i != layers; ++i) + { + auto& offloaded_layer = offloaded_cache->offloaded_kv_layers[i]; + auto& layer = kv_cache_layers[i]; + auto v = layer.v; + v = ggml_view_4d(offloaded_cache->ctx, v, n_tokens, v->ne[1], v->ne[2], v->ne[3], v->nb[1], v->nb[2], v->nb[3], 0); + v = ggml_cont(offloaded_cache->ctx, v); + ggml_backend_sched_set_tensor_backend(sched, v, backend); + ggml_build_forward_expand(gf, v); + offloaded_layer.v_tensor = v; + } + + ggml_backend_sched_alloc_graph(sched, gf); + ggml_backend_sched_graph_compute_async(sched, gf); + + auto v_nbytes = v_head_nbytes * total_heads; + for (auto& layer : std::span(offloaded_cache->offloaded_kv_layers, layers)) + ggml_backend_tensor_get_async(backend, layer.v_tensor, layer.v, 0, v_nbytes); } + + cur_len = 0; ggml_backend_sched_synchronize(sched); } void cosyvoice_kv_cache::load_cache(ggml_backend_t backend, ggml_backend_sched* sched) { + const auto batch_size = kv_cache_layers[0].k->ne[3]; + const auto total_heads = num_heads * batch_size; + const size_t k_head_nbytes = kv_cache_layers[0].k->nb[1] * offloaded_cache->len; + const size_t v_head_nbytes = fattn ? kv_cache_layers[0].v->nb[1] * offloaded_cache->len : ggml_row_size(kv_cache_layers[0].v->type, offloaded_cache->len) * kv_cache_layers[0].v->ne[1]; cur_len = offloaded_cache->len; - ggml_reset(offloaded_cache->ctx); - auto gf = ggml_new_graph_custom(offloaded_cache->ctx, layers * 4, false); - for (int i = 0; i != layers; ++i) + + if (!fattn) { - auto& offloaded_layer = offloaded_cache->offloaded_kv_layers[i]; - auto& layer = kv_cache_layers[i]; - ggml_tensor* k_view = ggml_view_3d(offloaded_cache->ctx, layer.k, layer.k->ne[0], offloaded_cache->len, layer.k->ne[2], layer.k->nb[1], layer.k->nb[2], 0); - ggml_tensor* v_view; - if (fattn) - v_view = ggml_view_3d(offloaded_cache->ctx, layer.v, layer.v->ne[0], offloaded_cache->len, layer.v->ne[2], layer.v->nb[1], layer.v->nb[2], 0); - else - v_view = ggml_view_3d(offloaded_cache->ctx, layer.v, offloaded_cache->len, layer.v->ne[1], layer.v->ne[2], layer.v->nb[1], layer.v->nb[2], 0); - - offloaded_layer.k_view = ggml_new_tensor(offloaded_cache->ctx, GGML_TYPE_F32, 3, k_view->ne); - offloaded_layer.v_view = ggml_new_tensor(offloaded_cache->ctx, GGML_TYPE_F32, 3, v_view->ne); - ggml_backend_sched_set_tensor_backend(sched, k_view, backend); - ggml_backend_sched_set_tensor_backend(sched, v_view, backend); - ggml_backend_sched_set_tensor_backend(sched, offloaded_layer.k_view, backend); - ggml_backend_sched_set_tensor_backend(sched, offloaded_layer.v_view, backend); - ggml_build_forward_expand(gf, ggml_cpy(offloaded_cache->ctx, offloaded_layer.k_view, k_view)); - ggml_build_forward_expand(gf, ggml_cpy(offloaded_cache->ctx, offloaded_layer.v_view, v_view)); + ggml_reset(offloaded_cache->ctx); + auto gf = ggml_new_graph_custom(offloaded_cache->ctx, layers * 4, false); + for (int i = 0; i != layers; ++i) + { + auto& offloaded_layer = offloaded_cache->offloaded_kv_layers[i]; + auto& layer = kv_cache_layers[i]; + ggml_tensor* v = ggml_new_tensor_4d(offloaded_cache->ctx, v_type, cur_len, layer.v->ne[1], num_heads, batch_size); + ggml_backend_sched_set_tensor_backend(sched, v, backend); + offloaded_layer.v_tensor = v; + ggml_tensor* v_view = ggml_view_4d(offloaded_cache->ctx, layer.v, cur_len, layer.v->ne[1], num_heads, batch_size, layer.v->nb[1], layer.v->nb[2], layer.v->nb[3], 0); + ggml_build_forward_expand(gf, ggml_cpy(offloaded_cache->ctx, v, v_view)); + } + + ggml_backend_sched_alloc_graph(sched, gf); + const auto v_nbytes = v_head_nbytes * total_heads; + for (auto& layer : std::span(offloaded_cache->offloaded_kv_layers, layers)) + ggml_backend_tensor_set_async(backend, layer.v_tensor, layer.v, 0, v_nbytes); + + ggml_backend_sched_graph_compute_async(sched, gf); } - ggml_backend_sched_alloc_graph(sched, gf); for (int i = 0; i != layers; ++i) { auto& offloaded_layer = offloaded_cache->offloaded_kv_layers[i]; - ggml_backend_tensor_set_async(backend, offloaded_layer.k_view, offloaded_layer.k, 0, offloaded_layer.k_view->nb[3]); - ggml_backend_tensor_set_async(backend, offloaded_layer.v_view, offloaded_layer.v, 0, offloaded_layer.v_view->nb[3]); + auto& layer = kv_cache_layers[i]; + for (int h = 0; h != total_heads; ++h) + { + ggml_backend_tensor_set_async(backend, layer.k, offloaded_layer.k + h * k_head_nbytes, h * layer.k->nb[2], k_head_nbytes); + if (fattn) + ggml_backend_tensor_set_async(backend, layer.v, offloaded_layer.v + h * v_head_nbytes, h * layer.v->nb[2], v_head_nbytes); + } } - ggml_backend_sched_graph_compute(sched, gf); + offloaded_cache->len = 0; + ggml_backend_sched_synchronize(sched); } size_t cosyvoice_kv_cache::get_offloaded_cache_size() const { if (offloaded_cache) - return offloaded_cache->offloaded_kv_layers[0].k_view->nb[3] * 2 * layers; + return offloaded_cache->buffer_size; else return 0; } @@ -304,7 +318,11 @@ void cosyvoice_kv_cache::shift_kv_node_pos(uint32_t shift_pos) bool cosyvoice_kv_cache::can_reuse(bool prefill) const { - return fattn; + if (!fattn) return false; + for (auto& layer : std::span(kv_cache_layers, layers)) + if (!layer.k_view || !layer.v_view) + return false; + return true; } void cosyvoice_kv_cache::slide_kv_layers(int layer_idx, int stride) diff --git a/src/cosyvoice-kv-cache.h b/src/cosyvoice-kv-cache.h index 0e9581b..5bf2671 100644 --- a/src/cosyvoice-kv-cache.h +++ b/src/cosyvoice-kv-cache.h @@ -42,7 +42,7 @@ class cosyvoice_kv_cache void slide_kv_layers(int layer_idx, int stride); - ggml_backend_buffer* initialize_buffer(ggml_backend_t backend, int k_head_dim, int v_head_dim, int num_attention_heads, int num_key_value_heads, uint32_t max_seq, ggml_type k_type, ggml_type v_type, int batch_size, bool fattn); + ggml_backend_buffer* initialize_buffer(ggml_backend_t backend, int k_head_dim, int v_head_dim, int num_key_value_heads, uint32_t max_seq, ggml_type k_type, ggml_type v_type, int batch_size, bool fattn); uint32_t reset_buffer(ggml_backend_buffer* buffer); void offload_cache(ggml_backend_t backend, ggml_backend_sched* sched, uint32_t n_tokens); @@ -53,7 +53,7 @@ class cosyvoice_kv_cache uint32_t cur_len; private: int layers; - int num_attention_heads; + int num_heads; bool fattn; ggml_type k_type; ggml_type v_type; diff --git a/src/cosyvoice-model.cpp b/src/cosyvoice-model.cpp index 7f76c7b..acab562 100644 --- a/src/cosyvoice-model.cpp +++ b/src/cosyvoice-model.cpp @@ -267,7 +267,6 @@ void cosyvoice_model_3::empty_buffer_cache() worker->backend.get(), static_cast(cv3_shared->llm.layers[0].self_attn.k_proj.weight->ne[1] / cv3_shared->llm.num_key_value_heads), static_cast(cv3_shared->llm.layers[0].self_attn.v_proj.weight->ne[1] / cv3_shared->llm.num_key_value_heads), - cv3_shared->llm.num_attention_heads, cv3_shared->llm.num_key_value_heads, shared->params.n_max_seq, cv3_shared->k_type, @@ -337,13 +336,13 @@ uint32_t cosyvoice_model_3::get_chunk_tokens() uint32_t cosyvoice_model_3::get_flow_overlap_tokens() { - // TODO: calculate the actual overlap tokens + // TODO: calculate the actual overlap tokens return 0; } uint32_t cosyvoice_model_3::get_hift_overlap_tokens() { - // TODO: calculate the actual overlap tokens + // TODO: calculate the actual overlap tokens return 0; } @@ -600,5 +599,5 @@ void cosyvoice_model::llm_load_kv_cache() auto sched = worker->sched.get(); auto& kv_cache = worker->llm_kv_cache; ggml_backend_sched_reset(sched); - kv_cache.offload_cache(worker->backend.get(), sched, kv_cache.cur_len); + kv_cache.load_cache(worker->backend.get(), sched); } From 1394562c367a09648efb93fc8bb739ea18e239a0 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Mon, 13 Jul 2026 22:22:57 +0800 Subject: [PATCH 06/34] cosyvoice: add multi-slot DiT KV cache with CPU offload scheduling --- include/cosyvoice-interface.h | 2 - include/cosyvoice-lowlevel.h | 22 +- include/cosyvoice.h | 131 ++++++++-- src/cosyvoice-graph.cpp | 41 ++-- src/cosyvoice-kv-cache.cpp | 206 ++++++++++------ src/cosyvoice-kv-cache.h | 24 +- src/cosyvoice-loader.cpp | 446 +++++++++++++++++++--------------- src/cosyvoice-model.cpp | 26 +- src/cosyvoice-model.h | 13 +- src/cosyvoice-modules.h | 11 +- src/cosyvoice-tts.cpp | 209 +++++++++++----- src/cosyvoice.cpp | 57 +++-- 12 files changed, 733 insertions(+), 455 deletions(-) diff --git a/include/cosyvoice-interface.h b/include/cosyvoice-interface.h index 5d6e248..ca570cd 100644 --- a/include/cosyvoice-interface.h +++ b/include/cosyvoice-interface.h @@ -97,8 +97,6 @@ struct cosyvoice_model_context ) = 0; ///< Convert speech tokens into waveform samples with additional options. virtual uint32_t get_chunk_tokens() = 0; ///< Get the number of tokens processed in each chunk during streaming inference. - virtual uint32_t get_flow_overlap_tokens() = 0; ///< Get the number of overlapping tokens used for streaming inference. - virtual uint32_t get_hift_overlap_tokens() = 0; ///< Get the number of extra tokens to overlap for HiFT lookahead in streaming inference. // Status virtual ggml_status get_last_status() = 0; ///< Get the status of the most recent backend operation. diff --git a/include/cosyvoice-lowlevel.h b/include/cosyvoice-lowlevel.h index 45f37a4..4860380 100644 --- a/include/cosyvoice-lowlevel.h +++ b/include/cosyvoice-lowlevel.h @@ -70,6 +70,7 @@ typedef float* (*cosyvoice_noise_callback_t)( #define COSYVOICE_CONTEXT_PARAMS_VERSION (0ul) #define COSYVOICE_CONTEXT_PARAMS_V2_VERSION (1ul) +#define COSYVOICE_CONTEXT_PARAMS_V3_VERSION (2ul) // ---------------------------------------------------------------------------- // Logging Utilities @@ -117,7 +118,7 @@ struct _is_same { static constexpr bool value = true; }; * @note This overload is only available in C++ and requires the full definition of all context parameter types. */ template -inline cosyvoice_context_t cosyvoice_load_from_file_ext( +constexpr cosyvoice_context_t cosyvoice_load_from_file_ext( const char* filename, const params_t* params, ggml_backend_t backend, @@ -130,10 +131,15 @@ inline cosyvoice_context_t cosyvoice_load_from_file_ext( return cosyvoice_load_from_file_ext(filename, params, backend, n_threads, COSYVOICE_CONTEXT_PARAMS_VERSION); else if constexpr (_is_same::value) return cosyvoice_load_from_file_ext(filename, ¶ms->base_params, backend, n_threads, COSYVOICE_CONTEXT_PARAMS_V2_VERSION); + else if constexpr (_is_same::value) + return cosyvoice_load_from_file_ext(filename, ¶ms->base_params.base_params, backend, n_threads, COSYVOICE_CONTEXT_PARAMS_V3_VERSION); else { - static_assert(_is_same::value, "Unsupported context parameter type"); - return cosyvoice_load_from_file_ext(filename, params, backend, n_threads, COSYVOICE_CONTEXT_PARAMS_V2_VERSION); + static_assert(_is_same::value || _is_same::value, "Unsupported context parameter type"); + constexpr auto version = _is_same::value + ? COSYVOICE_CONTEXT_PARAMS_V3_VERSION + : COSYVOICE_CONTEXT_PARAMS_V2_VERSION; + return cosyvoice_load_from_file_ext(filename, reinterpret_cast(params), backend, n_threads, version); } } @@ -349,16 +355,6 @@ COSYVOICE_API bool cosyvoice_tts_stream( */ COSYVOICE_API uint32_t cosyvoice_get_chunk_tokens(cosyvoice_context_t ctx); -/** - * @brief Get the number of tokens that overlap between consecutive chunks during streaming inference. - */ -COSYVOICE_API uint32_t cosyvoice_get_flow_overlap_tokens(cosyvoice_context_t ctx); - -/** -* @brief Get the number of extra tokens to overlap for HiFT lookahead in streaming inference. -*/ -COSYVOICE_API uint32_t cosyvoice_get_hift_overlap_tokens(cosyvoice_context_t ctx); - // ---------------------------------------------------------------------------- // Tokenizer Operations // ---------------------------------------------------------------------------- diff --git a/include/cosyvoice.h b/include/cosyvoice.h index d40ece2..e6eddb3 100644 --- a/include/cosyvoice.h +++ b/include/cosyvoice.h @@ -27,19 +27,28 @@ extern "C" { // ---------------------------------------------------------------------------- /** - * @brief Supported KV-cache storage formats for the LLM module. + * @brief Supported KV-cache storage formats. */ -typedef enum cosyvoice_llm_kv_cache_type +typedef enum cosyvoice_kv_cache_type { - COSYVOICE_LLM_KV_CACHE_TYPE_F32, ///< Store KV cache in 32-bit floating point format. - COSYVOICE_LLM_KV_CACHE_TYPE_F16, ///< Store KV cache in 16-bit floating point format. - COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0, ///< Store KV cache in GGML Q8_0 quantized format. - COSYVOICE_LLM_KV_CACHE_TYPE_Q5_1, ///< Store KV cache in GGML Q5_1 quantized format. - COSYVOICE_LLM_KV_CACHE_TYPE_Q5_0, ///< Store KV cache in GGML Q5_0 quantized format. - COSYVOICE_LLM_KV_CACHE_TYPE_Q4_1, ///< Store KV cache in GGML Q4_1 quantized format. - COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0, ///< Store KV cache in GGML Q4_0 quantized format. - COSYVOICE_LLM_KV_CACHE_TYPE_COUNT // Sentinel value. -} cosyvoice_llm_kv_cache_type_t; + COSYVOICE_KV_CACHE_TYPE_F32, ///< Store KV cache in 32-bit floating point format. + COSYVOICE_KV_CACHE_TYPE_F16, ///< Store KV cache in 16-bit floating point format. + COSYVOICE_KV_CACHE_TYPE_Q8_0, ///< Store KV cache in GGML Q8_0 quantized format. + COSYVOICE_KV_CACHE_TYPE_Q5_1, ///< Store KV cache in GGML Q5_1 quantized format. + COSYVOICE_KV_CACHE_TYPE_Q5_0, ///< Store KV cache in GGML Q5_0 quantized format. + COSYVOICE_KV_CACHE_TYPE_Q4_1, ///< Store KV cache in GGML Q4_1 quantized format. + COSYVOICE_KV_CACHE_TYPE_Q4_0, ///< Store KV cache in GGML Q4_0 quantized format. + COSYVOICE_KV_CACHE_TYPE_COUNT // Sentinel value. +} cosyvoice_kv_cache_type_t, cosyvoice_llm_kv_cache_type_t; + +#define COSYVOICE_LLM_KV_CACHE_TYPE_F32 COSYVOICE_KV_CACHE_TYPE_F32 +#define COSYVOICE_LLM_KV_CACHE_TYPE_F16 COSYVOICE_KV_CACHE_TYPE_F16 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0 COSYVOICE_KV_CACHE_TYPE_Q8_0 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q5_1 COSYVOICE_KV_CACHE_TYPE_Q5_1 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q5_0 COSYVOICE_KV_CACHE_TYPE_Q5_0 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q4_1 COSYVOICE_KV_CACHE_TYPE_Q4_1 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0 COSYVOICE_KV_CACHE_TYPE_Q4_0 +#define COSYVOICE_LLM_KV_CACHE_TYPE_COUNT COSYVOICE_KV_CACHE_TYPE_COUNT /** * @brief Buffer allocation strategies for inference workloads. @@ -181,14 +190,14 @@ typedef int (*cosyvoice_sampler_ext_t)( * bits 15–30 unused * bit 31 separate-K/V flag (must be 1) * - * When bit 31 is 0, the value is a plain `cosyvoice_llm_kv_cache_type_t` + * When bit 31 is 0, the value is a plain `cosyvoice_kv_cache_type_t` * that applies to both K and V (backward compatible). * * When the preferred K or V type is not supported by the backend, the * fallback type is tried first; if that also fails, auto-fallback applies. */ #define COSYVOICE_MAKE_SEPARATE_KV_CACHE(k_type, v_type, fallback_type) \ - ((cosyvoice_llm_kv_cache_type_t)((k_type) | ((v_type) << 5) | ((fallback_type) << 10) | (1U << 31))) + ((cosyvoice_kv_cache_type_t)((k_type) | ((v_type) << 5) | ((fallback_type) << 10) | (1U << 31))) /** * @brief Test whether llm_kv_cache_type is in separate-K/V mode. @@ -198,17 +207,17 @@ typedef int (*cosyvoice_sampler_ext_t)( /** * @brief Extract the K cache type from a packed value. */ -#define COSYVOICE_K_CACHE_TYPE(t) ((cosyvoice_llm_kv_cache_type_t)((t) & 0x1F)) +#define COSYVOICE_K_CACHE_TYPE(t) ((cosyvoice_kv_cache_type_t)((t) & 0x1F)) /** * @brief Extract the V cache type from a packed value. */ -#define COSYVOICE_V_CACHE_TYPE(t) ((cosyvoice_llm_kv_cache_type_t)(((t) >> 5) & 0x1F)) +#define COSYVOICE_V_CACHE_TYPE(t) ((cosyvoice_kv_cache_type_t)(((t) >> 5) & 0x1F)) /** * @brief Extract the fallback cache type from a packed value. */ -#define COSYVOICE_KV_CACHE_FALLBACK(t) ((cosyvoice_llm_kv_cache_type_t)(((t) >> 10) & 0x1F)) +#define COSYVOICE_KV_CACHE_FALLBACK(t) ((cosyvoice_kv_cache_type_t)(((t) >> 10) & 0x1F)) /** * @brief Parameters used when creating a model context. @@ -223,20 +232,20 @@ typedef struct cosyvoice_context_params struct { #if defined(__BYTE_ORDER__) && (__BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) || defined(_BYTE_ORDER) && (_BYTE_ORDER == _BIG_ENDIAN) || defined(__BIG_ENDIAN__) && !defined(__LITTLE_ENDIAN__) || defined(__ARMEB__) || defined(__MIPSEB__) || defined(__sparc__) - uint32_t llm_kv_cache_separate_buffers : 1; ///< If true, allocate separate buffers for the K and V caches in the LLM module. Ignored when `llm_kv_cache_type` is set to a unified format. - cosyvoice_llm_kv_cache_type_t : 16; ///< Unused bits. - cosyvoice_llm_kv_cache_type_t llm_kv_cache_fallback : 5; ///< Fallback type when the preferred K or V type is unsupported. - cosyvoice_llm_kv_cache_type_t llm_v_cache_type : 5; ///< The data type of the V cache in the LLM module. - cosyvoice_llm_kv_cache_type_t llm_k_cache_type : 5; ///< The data type of the K cache in the LLM module. + uint32_t llm_kv_cache_separate_buffers : 1; ///< If true, allocate separate buffers for the K and V caches in the LLM module. Ignored when `llm_kv_cache_type` is set to a unified format. + cosyvoice_kv_cache_type_t : 16; ///< Unused bits. + cosyvoice_kv_cache_type_t llm_kv_cache_fallback : 5; ///< Fallback type when the preferred K or V type is unsupported. + cosyvoice_kv_cache_type_t llm_v_cache_type : 5; ///< The data type of the V cache in the LLM module. + cosyvoice_kv_cache_type_t llm_k_cache_type : 5; ///< The data type of the K cache in the LLM module. #else - cosyvoice_llm_kv_cache_type_t llm_k_cache_type : 5; ///< The data type of the K cache in the LLM module. - cosyvoice_llm_kv_cache_type_t llm_v_cache_type : 5; ///< The data type of the V cache in the LLM module. - cosyvoice_llm_kv_cache_type_t llm_kv_cache_fallback : 5; ///< Fallback type when the preferred K or V type is unsupported. - cosyvoice_llm_kv_cache_type_t : 16; ///< Unused bits. - uint32_t llm_kv_cache_separate_buffers : 1; ///< If true, allocate separate buffers for the K and V caches in the LLM module. Ignored when `llm_kv_cache_type` is set to a unified format. + cosyvoice_kv_cache_type_t llm_k_cache_type : 5; ///< The data type of the K cache in the LLM module. + cosyvoice_kv_cache_type_t llm_v_cache_type : 5; ///< The data type of the V cache in the LLM module. + cosyvoice_kv_cache_type_t llm_kv_cache_fallback : 5; ///< Fallback type when the preferred K or V type is unsupported. + cosyvoice_kv_cache_type_t : 16; ///< Unused bits. + uint32_t llm_kv_cache_separate_buffers : 1; ///< If true, allocate separate buffers for the K and V caches in the LLM module. Ignored when `llm_kv_cache_type` is set to a unified format. #endif }; - cosyvoice_llm_kv_cache_type_t llm_kv_cache_type; ///< The data type of the KV cache in the LLM module. + cosyvoice_kv_cache_type_t llm_kv_cache_type; ///< The data type of the KV cache in the LLM module. }; bool llm_allow_kv_cache_fallback; ///< If true, fall back to a Flash Attention-compatible KV cache type when the requested one is unsupported. cosyvoice_inference_buffer_policy_t inference_buffer_policy; ///< Controls how inference buffers are allocated and reused, which affects performance and memory usage. @@ -261,11 +270,71 @@ typedef struct cosyvoice_context_params_v2 uint32_t n_workers; ///< Number of workers to use for inference. } cosyvoice_context_params_v2_t; +/** + * @brief Extended context parameters that add DiT KV cache configuration options. + */ +typedef struct cosyvoice_context_params_v3 +{ + cosyvoice_context_params_v2_t base_params; ///< V2 base parameters. + + union + { + struct + { +#if defined(__BYTE_ORDER__) && (__BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) || defined(_BYTE_ORDER) && (_BYTE_ORDER == _BIG_ENDIAN) || defined(__BIG_ENDIAN__) && !defined(__LITTLE_ENDIAN__) || defined(__ARMEB__) || defined(__MIPSEB__) || defined(__sparc__) + uint32_t dit_kv_cache_separate_buffers : 1; ///< If true, allocate separate buffers for the K and V caches in the DiT module. Ignored when `dit_kv_cache_type` is set to a unified format. + cosyvoice_kv_cache_type_t : 16; ///< Unused bits. + cosyvoice_kv_cache_type_t dit_kv_cache_fallback : 5; ///< Fallback type when the preferred K or V type is unsupported. + cosyvoice_kv_cache_type_t dit_v_cache_type : 5; ///< The data type of the V cache in the DiT module. + cosyvoice_kv_cache_type_t dit_k_cache_type : 5; ///< The data type of the K cache in the DiT module. +#else + cosyvoice_kv_cache_type_t dit_k_cache_type : 5; ///< The data type of the K cache in the DiT module. + cosyvoice_kv_cache_type_t dit_v_cache_type : 5; ///< The data type of the V cache in the DiT module. + cosyvoice_kv_cache_type_t dit_kv_cache_fallback : 5; ///< Fallback type when the preferred K or V type is unsupported. + cosyvoice_kv_cache_type_t : 16; ///< Unused bits. + uint32_t dit_kv_cache_separate_buffers : 1; ///< If true, allocate separate buffers for the K and V caches in the DiT module. Ignored when `dit_kv_cache_type` is set to a unified format. +#endif + }; + cosyvoice_kv_cache_type_t dit_kv_cache_type; ///< The data type of the KV cache in the DiT module. + }; + bool dit_allow_kv_cache_fallback; ///< If true, fall back to a Flash Attention-compatible KV cache type when the requested one is unsupported. + uint32_t dit_kv_fixed_slots; ///< Number of fixed (non-offloadable) DiT KV slots. + uint32_t dit_kv_offloadable_slots; ///< Number of offloadable DiT KV slots. + uint32_t dit_kv_cache_length; ///< Maximum sequence length for the DiT KV cache. 0 to use default (n_max_seq * 10). +} cosyvoice_context_params_v3_t; #ifdef __cplusplus struct cosyvoice_context_params_v2_cpp : cosyvoice_context_params_t { uint32_t n_workers; ///< Number of workers to use for inference. }; + +struct cosyvoice_context_params_v3_cpp : cosyvoice_context_params_v2_cpp +{ + union + { + struct + { +#if defined(__BYTE_ORDER__) && (__BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) || defined(_BYTE_ORDER) && (_BYTE_ORDER == _BIG_ENDIAN) || defined(__BIG_ENDIAN__) && !defined(__LITTLE_ENDIAN__) || defined(__ARMEB__) || defined(__MIPSEB__) || defined(__sparc__) + uint32_t dit_kv_cache_separate_buffers : 1; ///< If true, allocate separate buffers for the K and V caches in the DiT module. Ignored when `dit_kv_cache_type` is set to a unified format. + cosyvoice_kv_cache_type_t : 16; ///< Unused bits. + cosyvoice_kv_cache_type_t dit_kv_cache_fallback : 5; ///< Fallback type when the preferred K or V type is unsupported. + cosyvoice_kv_cache_type_t dit_v_cache_type : 5; ///< The data type of the V cache in the DiT module. + cosyvoice_kv_cache_type_t dit_k_cache_type : 5; ///< The data type of the K cache in the DiT module. +#else + cosyvoice_kv_cache_type_t dit_k_cache_type : 5; ///< The data type of the K cache in the DiT module. + cosyvoice_kv_cache_type_t dit_v_cache_type : 5; ///< The data type of the V cache in the DiT module. + cosyvoice_kv_cache_type_t dit_kv_cache_fallback : 5; ///< Fallback type when the preferred K or V type is unsupported. + cosyvoice_kv_cache_type_t : 16; ///< Unused bits. + uint32_t dit_kv_cache_separate_buffers : 1; ///< If true, allocate separate buffers for the K and V caches in the DiT module. Ignored when `dit_kv_cache_type` is set to a unified format. +#endif + }; + cosyvoice_kv_cache_type_t dit_kv_cache_type; ///< The data type of the KV cache in the DiT module. + }; + bool dit_allow_kv_cache_fallback; ///< If true, fall back to a Flash Attention-compatible KV cache type when the requested one is unsupported. + uint32_t dit_kv_fixed_slots; ///< Number of fixed (non-offloadable) DiT KV slots. + uint32_t dit_kv_offloadable_slots; ///< Number of offloadable DiT KV slots. + uint32_t dit_kv_cache_length; ///< Maximum sequence length for the DiT KV cache. 0 to use default (n_max_seq * 10). +}; #endif // ---------------------------------------------------------------------------- @@ -324,6 +393,14 @@ COSYVOICE_API cosyvoice_context_t cosyvoice_load_from_file_with_params_v2( const cosyvoice_context_params_v2_t* params ); +/** + * @brief Load a model context from a GGUF file using V3 context parameters with DiT KV cache configuration. + */ +COSYVOICE_API cosyvoice_context_t cosyvoice_load_from_file_with_params_v3( + const char* filename, + const cosyvoice_context_params_v3_t* params +); + /** * @brief Duplicate a loaded model context handle. * @note The duplicate shares the loaded model resources with the original context. It starts with the same active worker binding as the original context, and can then be rebound independently with `cosyvoice_set_worker_no()`. diff --git a/src/cosyvoice-graph.cpp b/src/cosyvoice-graph.cpp index adb43ac..df9236b 100644 --- a/src/cosyvoice-graph.cpp +++ b/src/cosyvoice-graph.cpp @@ -323,7 +323,7 @@ std::array AdaLayerNormZero::build_cgraph(ggml_context* ctx, gg return { x, gate_msa, shift_mlp, scale_mlp, gate_mlp }; } -ggml_tensor* Attention::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* position_ids, int64_t cut_len, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, int layer_idx) const +ggml_tensor* Attention::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* position_ids, int64_t cut_len, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, ggml_tensor* attn_mask, int layer_idx) const { const auto full_seq_len = position_ids->ne[0]; const auto seq_len = full_seq_len - (cut_len > 0 ? cut_len : 0); @@ -366,7 +366,7 @@ ggml_tensor* Attention::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_ten if (kv_cache) { kv_cache->update_cache(ctx, gf, key, value, original_position_ids, layer_idx); - attn_output = kv_cache->attention_forward(ctx, query, key, value, nullptr); + attn_output = kv_cache->attention_forward(ctx, query, key, value, attn_mask); } else { @@ -374,13 +374,13 @@ ggml_tensor* Attention::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_ten key = ggml_cont(ctx, key); if (fattn) - attn_output = ggml_flash_attn_ext(ctx, query, key, value, nullptr, 1.f / std::sqrt(static_cast(head_dim)), 0.f, 0.f); + attn_output = ggml_flash_attn_ext(ctx, query, key, value, attn_mask, 1.f / std::sqrt(static_cast(head_dim)), 0.f, 0.f); else { value = ggml_permute(ctx, value, 0, 2, 1, 3); value = ggml_cont(ctx, value); auto attn_scores = ggml_mul_mat(ctx, key, query); - auto attn_weights = ggml_soft_max_ext_inplace(ctx, attn_scores, nullptr, 1.f / std::sqrt(static_cast(head_dim)), 0.f); + auto attn_weights = ggml_soft_max_ext_inplace(ctx, attn_scores, attn_mask, 1.f / std::sqrt(static_cast(head_dim)), 0.f); auto attn_output = ggml_mul_mat(ctx, value, attn_weights); attn_output = ggml_permute(ctx, attn_output, 0, 2, 1, 3); attn_output = ggml_cont(ctx, attn_output); @@ -404,11 +404,11 @@ ggml_tensor* FeedForward::build_cgraph(ggml_context* ctx, ggml_tensor* x) const return x; } -ggml_tensor* DiTBlock::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* time_emb, ggml_tensor* position_ids, int64_t cut_len, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, int layer_idx) const +ggml_tensor* DiTBlock::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* time_emb, ggml_tensor* position_ids, int64_t cut_len, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, ggml_tensor* attn_mask, int layer_idx) const { auto [norm, gate_msa, shift_mlp, scale_mlp, gate_mlp] = attn_norm.build_cgraph(ctx, x, time_emb); - auto attn_output = attn.build_cgraph(ctx, norm, position_ids, cut_len, kv_cache, gf, layer_idx); + auto attn_output = attn.build_cgraph(ctx, norm, position_ids, cut_len, kv_cache, gf, attn_mask, layer_idx); gate_msa = unsqueeze(ctx, gate_msa, 1); attn_output = ggml_mul(ctx, attn_output, gate_msa); if (cut_len > 0) @@ -439,7 +439,7 @@ ggml_tensor* AdaLayerNorm_Final::build_cgraph(ggml_context* ctx, ggml_tensor* x, return x; } -ggml_tensor* DiT::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* mu, ggml_tensor* t, ggml_tensor* spks, ggml_tensor* cond, int64_t cut_len, ggml_tensor*& ref_position_ids, ggml_backend_op_capabilities capabilities, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, int step) const +ggml_tensor* DiT::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* mu, ggml_tensor* t, ggml_tensor* spks, ggml_tensor* cond, int64_t cut_len, ggml_tensor*& ref_position_ids, ggml_backend_op_capabilities capabilities, cosyvoice_kv_cache* kv_cache, ggml_tensor** attn_mask, ggml_cgraph* gf) const { t = time_embed.build_cgraph(ctx, t); x = input_embed.build_cgraph(ctx, x, cond, mu, spks, capabilities); @@ -449,11 +449,19 @@ ggml_tensor* DiT::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* m ref_position_ids = position_ids; position_ids = ggml_dup(ctx, position_ids); - int layer_idx_base = step * static_cast(transformer_blocks.size()); + ggml_tensor* mask = nullptr; + if (attn_mask) + { + const auto seq_len = x->ne[1]; + mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, seq_len, seq_len); + ggml_set_input(mask); + *attn_mask = mask; + } + for (int i = 0; i < transformer_blocks.size() - 1; ++i) - x = transformer_blocks[i].build_cgraph(ctx, x, t, position_ids, 0, kv_cache, gf, layer_idx_base + i); + x = transformer_blocks[i].build_cgraph(ctx, x, t, position_ids, 0, kv_cache, gf, mask, i); // Apply `cut_len` only on the final block. - x = transformer_blocks.back().build_cgraph(ctx, x, t, position_ids, cut_len, kv_cache, gf, layer_idx_base + static_cast(transformer_blocks.size()) - 1); + x = transformer_blocks.back().build_cgraph(ctx, x, t, position_ids, cut_len, kv_cache, gf, mask, static_cast(transformer_blocks.size()) - 1); x = norm_out.build_cgraph(ctx, x, t); @@ -492,7 +500,7 @@ std::array CausalConditionalCFM::get_t_and_dt(ggml_context* ctx, int s return { t, dt }; } -ggml_tensor* CausalConditionalCFM::build_cgraph_one_step(ggml_context* ctx, const DiTContext& ditctx, int step, ggml_backend_op_capabilities capabilities, int64_t cut_len, ggml_tensor*& t_tensor, ggml_tensor*& position_ids, ggml_cgraph* gf, cosyvoice_kv_cache* kv_cache) const +ggml_tensor* CausalConditionalCFM::build_cgraph_one_step(ggml_context* ctx, const DiTContext& ditctx, int step, ggml_backend_op_capabilities capabilities, int64_t cut_len, ggml_tensor*& t_tensor, ggml_tensor*& position_ids, ggml_cgraph* gf, cosyvoice_kv_cache* kv_cache, ggml_tensor** attn_mask) const { auto x = ditctx.x; auto [t, dt] = get_t_and_dt(ctx, step); @@ -508,12 +516,12 @@ ggml_tensor* CausalConditionalCFM::build_cgraph_one_step(ggml_context* ctx, cons ditctx.mu_in, t_in, ditctx.spks_in, ditctx.cond_in, - step == t_span.size() - 1 ? cut_len : 0, + cut_len, position_ids, capabilities, kv_cache, - gf, - step - 1), + attn_mask, + gf), 2); cfg_dphi_dt->nb[3] = dphi_dt->nb[3] = dphi_dt->nb[2]; @@ -523,7 +531,7 @@ ggml_tensor* CausalConditionalCFM::build_cgraph_one_step(ggml_context* ctx, cons dphi_dt, cfg_dphi_dt); - if (step == t_span.size() - 1 && cut_len > 0) + if (cut_len > 0) x = ggml_view_3d(ctx, x, x->ne[0], x->ne[1] - cut_len, x->ne[2], x->nb[1], x->nb[2], x->nb[1] * cut_len); x = ggml_add(ctx, x, @@ -567,12 +575,11 @@ CausalMaskedDiffWithDiT::EncodeResult CausalMaskedDiffWithDiT::build_cgraph_enco token = concat_tensors(ctx, std::array{ prompt_token, token }, 0, capabilities); token = ggml_get_rows(ctx, input_embedding, token); - ggml_tensor* h = pre_lookahead_layer.build_cgraph(ctx, token, streaming, cut_len); + ggml_tensor* h = pre_lookahead_layer.build_cgraph(ctx, token, streaming, cut_len / token_mel_ratio); h = unsqueeze(ctx, h, 1); h = ggml_repeat_4d(ctx, h, h->ne[0], token_mel_ratio, h->ne[2], h->ne[3]); h = ggml_reshape_3d(ctx, h, h->ne[0], h->ne[1] * h->ne[2], h->ne[3]); - cut_len *= token_mel_ratio; const auto mel_len1 = prompt_feat->ne[1]; const auto mel_len2 = h->ne[1] - mel_len1 + cut_len; auto conds = ggml_pad(ctx, prompt_feat, 0, static_cast(mel_len2), 0, 0); diff --git a/src/cosyvoice-kv-cache.cpp b/src/cosyvoice-kv-cache.cpp index 7fd39e5..1e66906 100644 --- a/src/cosyvoice-kv-cache.cpp +++ b/src/cosyvoice-kv-cache.cpp @@ -5,6 +5,42 @@ #include #include +#include + +struct kv_cache_layer +{ + ggml_tensor* k; + ggml_tensor* v; + + ggml_tensor* k_view; + ggml_tensor* v_view; +}; + +struct offloaded_kv_cache +{ + ~offloaded_kv_cache() = default; + + uint32_t len; + std::vector buffer; + ggml_context* ctx; + struct offloaded_kv_layer + { + ggml_tensor* v_tensor; + char* k; + char* v; + } offloaded_kv_layers[]; +}; + +static constexpr size_t get_offloaded_kv_cache_struct_size(int layers) +{ + return sizeof(offloaded_kv_cache) + sizeof(offloaded_kv_cache::offloaded_kv_layer) * layers; +} + +template +static constexpr T* advance_ptr(T* ptr, std::ptrdiff_t offset) +{ + return reinterpret_cast(reinterpret_cast(ptr) + offset); +} void cosyvoice_kv_cache::build_kv_cache( ggml_backend_t backend, @@ -17,6 +53,8 @@ void cosyvoice_kv_cache::build_kv_cache( ggml_type k_type, ggml_type v_type, int batch_size, + int n_slots, + int n_offloaded_kv_slots, bool fattn) { cur_len = 0; @@ -24,27 +62,53 @@ void cosyvoice_kv_cache::build_kv_cache( this->fattn = fattn; this->k_type = k_type; this->v_type = v_type; - offloaded_cache = nullptr; this->num_heads = num_key_value_heads; - kv_cache_layers = new kv_cache_layer[layers]; + this->n_slots = n_slots; + this->n_offloaded_kv_slots = n_offloaded_kv_slots; + cur_slot_idx = 0; + kv_cache_layers = new kv_cache_layer[layers * n_slots]; - ggml_init_params params = { - .mem_size = layers * 2 * ggml_tensor_overhead(), - .no_alloc = true - }; - ctx = ggml_init(params); + { + ggml_init_params params = { + .mem_size = layers * 2 * n_slots * ggml_tensor_overhead(), + .no_alloc = true + }; + ctx = ggml_init(params); + } - buffer.reset(initialize_buffer(backend, k_head_dim, v_head_dim, num_key_value_heads, max_seq, k_type, v_type, batch_size, fattn)); + buffer.reset(initialize_buffer(backend, k_head_dim, v_head_dim, max_seq, batch_size)); + + if (n_offloaded_kv_slots > 0) + { + ggml_init_params params = + { + .mem_size = (ggml_tensor_overhead() + ggml_graph_overhead()) * layers * 4 * n_offloaded_kv_slots, + .no_alloc = true + }; + const auto object_size = get_offloaded_kv_cache_struct_size(layers); + offloaded_cache = reinterpret_cast(malloc(object_size * n_offloaded_kv_slots)); + + for (int i = 0; i != n_offloaded_kv_slots; ++i) + { + auto slot = advance_ptr(offloaded_cache, object_size * i); + new(slot) offloaded_kv_cache(); + slot->len = 0; + slot->buffer.resize(0); + slot->ctx = ggml_init(params); + } + } + else + offloaded_cache = nullptr; } -ggml_backend_buffer* cosyvoice_kv_cache::initialize_buffer(ggml_backend_t backend, int k_head_dim, int v_head_dim, int num_key_value_heads, uint32_t max_seq, ggml_type k_type, ggml_type v_type, int batch_size, bool fattn) +ggml_backend_buffer* cosyvoice_kv_cache::initialize_buffer(ggml_backend_t backend, int k_head_dim, int v_head_dim, uint32_t max_seq, int batch_size) { - int64_t k_ne[4] = { k_head_dim, max_seq, num_key_value_heads, batch_size }; - int64_t v_ne[4] = { v_head_dim, max_seq, num_key_value_heads, batch_size }; + int64_t k_ne[4] = { k_head_dim, max_seq, num_heads, batch_size }; + int64_t v_ne[4] = { v_head_dim, max_seq, num_heads, batch_size }; if (!fattn) std::swap(v_ne[0], v_ne[1]); ggml_reset(ctx); - for (auto& [k, v, k_view, v_view] : std::span(kv_cache_layers, layers)) + for (auto& [k, v, k_view, v_view] : std::span(kv_cache_layers, layers * n_slots)) { k = ggml_new_tensor(ctx, k_type, 4, k_ne); v = ggml_new_tensor(ctx, v_type, 4, v_ne); @@ -64,7 +128,7 @@ uint32_t cosyvoice_kv_cache::reset_buffer(ggml_backend_buffer* buffer) auto current_k_size = get_aligned_size(ggml_backend_buffer_get_alloc_size(buffer, kv_cache_layers[0].k), alignment); auto current_v_size = get_aligned_size(ggml_backend_buffer_get_alloc_size(buffer, kv_cache_layers[0].v), alignment); - GGML_ASSERT(ggml_backend_buffer_get_size(buffer) >= static_cast(layers) * (current_k_size + current_v_size)); + GGML_ASSERT(ggml_backend_buffer_get_size(buffer) >= static_cast(layers * n_slots) * (current_k_size + current_v_size)); // Rebind the existing KV layout into the larger buffer without changing the // cached sequence capacity. @@ -75,7 +139,7 @@ uint32_t cosyvoice_kv_cache::reset_buffer(ggml_backend_buffer* buffer) ggml_reset(ctx); auto buffer_base = reinterpret_cast(ggml_backend_buffer_get_base(buffer)); - for (auto& [k, v, k_view, v_view] : std::span(kv_cache_layers, layers)) + for (auto& [k, v, k_view, v_view] : std::span(kv_cache_layers, layers * n_slots)) { k = ggml_new_tensor(ctx, k_type, GGML_MAX_DIMS, k_ne.data()); v = ggml_new_tensor(ctx, v_type, GGML_MAX_DIMS, v_ne.data()); @@ -90,53 +154,16 @@ uint32_t cosyvoice_kv_cache::reset_buffer(ggml_backend_buffer* buffer) return max_seq_len; } -#pragma pack(push) -#pragma pack(1) -struct offloaded_kv_cache -{ - uint32_t len; - size_t buffer_size; - ggml_context* ctx; - struct offloaded_kv_layer - { - ggml_tensor* v_tensor; - char* k; - char* v; - } offloaded_kv_layers[]; -}; -#pragma pack(pop) - void cosyvoice_kv_cache::offload_cache(ggml_backend_t backend, ggml_backend_sched* sched, uint32_t n_tokens) { - char* buffer_base; const auto batch_size = kv_cache_layers[0].k->ne[3]; const size_t k_head_nbytes = kv_cache_layers[0].k->nb[1] * n_tokens; const size_t v_head_nbytes = fattn ? kv_cache_layers[0].v->nb[1] * n_tokens : ggml_row_size(kv_cache_layers[0].v->type, n_tokens) * kv_cache_layers[0].v->ne[1]; const size_t kv_nbytes = (k_head_nbytes + v_head_nbytes) * num_heads * batch_size; - if (!offloaded_cache) - { - ggml_init_params params = - { - .mem_size = (ggml_tensor_overhead() + ggml_graph_overhead()) * layers * 4, - .no_alloc = true - }; - offloaded_cache = reinterpret_cast(malloc(sizeof(offloaded_kv_cache::offloaded_kv_layer) * layers + sizeof(uint32_t) + sizeof(ggml_context*) + sizeof(size_t))); - offloaded_cache->ctx = ggml_init(params); - offloaded_cache->buffer_size = 0; - goto alloc_buffer; - } - else - { - ggml_reset(offloaded_cache->ctx); - if (offloaded_cache->buffer_size < kv_nbytes * layers) - { - free(offloaded_cache->offloaded_kv_layers[0].k); - alloc_buffer: - offloaded_cache->buffer_size = kv_nbytes * layers; - buffer_base = reinterpret_cast(malloc(offloaded_cache->buffer_size)); - } - else buffer_base = reinterpret_cast(offloaded_cache->offloaded_kv_layers[0].k); - } + + ggml_reset(offloaded_cache->ctx); + offloaded_cache->buffer.resize(kv_nbytes * layers); + char* buffer_base = offloaded_cache->buffer.data(); offloaded_cache->len = n_tokens; const auto total_heads = num_heads * batch_size; @@ -198,7 +225,7 @@ void cosyvoice_kv_cache::load_cache(ggml_backend_t backend, ggml_backend_sched* for (int i = 0; i != layers; ++i) { auto& offloaded_layer = offloaded_cache->offloaded_kv_layers[i]; - auto& layer = kv_cache_layers[i]; + auto& layer = kv_cache_layers[i + cur_slot_idx * layers]; ggml_tensor* v = ggml_new_tensor_4d(offloaded_cache->ctx, v_type, cur_len, layer.v->ne[1], num_heads, batch_size); ggml_backend_sched_set_tensor_backend(sched, v, backend); offloaded_layer.v_tensor = v; @@ -217,7 +244,7 @@ void cosyvoice_kv_cache::load_cache(ggml_backend_t backend, ggml_backend_sched* for (int i = 0; i != layers; ++i) { auto& offloaded_layer = offloaded_cache->offloaded_kv_layers[i]; - auto& layer = kv_cache_layers[i]; + auto& layer = kv_cache_layers[i + cur_slot_idx * layers]; for (int h = 0; h != total_heads; ++h) { ggml_backend_tensor_set_async(backend, layer.k, offloaded_layer.k + h * k_head_nbytes, h * layer.k->nb[2], k_head_nbytes); @@ -233,7 +260,13 @@ void cosyvoice_kv_cache::load_cache(ggml_backend_t backend, ggml_backend_sched* size_t cosyvoice_kv_cache::get_offloaded_cache_size() const { if (offloaded_cache) - return offloaded_cache->buffer_size; + { + size_t nbytes = 0; + auto object_size = get_offloaded_kv_cache_struct_size(layers); + for (auto i = 0; i != layers; ++i) + nbytes += advance_ptr(offloaded_cache, object_size)->buffer.capacity(); + return nbytes; + } else return 0; } @@ -242,24 +275,51 @@ void cosyvoice_kv_cache::clear_offloaded_cache() if (offloaded_cache) { ggml_free(offloaded_cache->ctx); - free(offloaded_cache->offloaded_kv_layers[0].k); - free(offloaded_cache); - offloaded_cache = nullptr; + auto object_size = get_offloaded_kv_cache_struct_size(layers); + for (auto i = 0; i != layers; ++i) + { + auto cur = advance_ptr(offloaded_cache, object_size); + std::vector().swap(cur->buffer); + cur->len = 0; + } } } +void cosyvoice_kv_cache::offload_slot(ggml_backend_t backend, ggml_backend_sched* sched, int offloaded_slot_idx) +{ + auto offset = static_cast(get_offloaded_kv_cache_struct_size(layers) * offloaded_slot_idx); + offloaded_cache = advance_ptr(offloaded_cache, offset); + offload_cache(backend, sched, cur_len); + offloaded_cache = advance_ptr(offloaded_cache, -offset); +} + +void cosyvoice_kv_cache::load_slot(ggml_backend_t backend, ggml_backend_sched* sched, int offloaded_slot_idx) +{ + auto offset = static_cast(get_offloaded_kv_cache_struct_size(layers) * offloaded_slot_idx); + offloaded_cache = advance_ptr(offloaded_cache, offset); + load_cache(backend, sched); + offloaded_cache = advance_ptr(offloaded_cache, -offset); +} + cosyvoice_kv_cache::~cosyvoice_kv_cache() { ggml_free(ctx); delete[] kv_cache_layers; - clear_offloaded_cache(); + auto object_size = get_offloaded_kv_cache_struct_size(layers); + for (int i = 0; i != n_offloaded_kv_slots; ++i) + { + auto cur = advance_ptr(offloaded_cache, object_size * i); + ggml_free(cur->ctx); + cur->~offloaded_kv_cache(); + } + free(offloaded_cache); } void cosyvoice_kv_cache::update_cache(ggml_context* ctx0, ggml_cgraph* gf, ggml_tensor*& k, ggml_tensor*& v, ggml_tensor* position_ids, int layer_idx) { GGML_ASSERT(ggml_are_same_shape(k, v)); - auto& layer = kv_cache_layers[layer_idx]; + auto& layer = kv_cache_layers[cur_slot_idx * layers + layer_idx]; if (fattn) { @@ -309,7 +369,7 @@ void cosyvoice_kv_cache::shift_kv_node_pos(uint32_t shift_pos) GGML_ASSERT(fattn); cur_len += shift_pos; - for (auto& layer : std::span(kv_cache_layers, layers)) + for (auto& layer : std::span(kv_cache_layers + cur_slot_idx * layers, layers)) { layer.k_view->ne[1] += shift_pos; layer.v_view->ne[1] += shift_pos; @@ -319,22 +379,30 @@ void cosyvoice_kv_cache::shift_kv_node_pos(uint32_t shift_pos) bool cosyvoice_kv_cache::can_reuse(bool prefill) const { if (!fattn) return false; - for (auto& layer : std::span(kv_cache_layers, layers)) + for (auto& layer : std::span(kv_cache_layers + cur_slot_idx * layers, layers)) if (!layer.k_view || !layer.v_view) return false; return true; } -void cosyvoice_kv_cache::slide_kv_layers(int layer_idx, int stride) +bool cosyvoice_kv_cache::bind_slot(int slot_idx) +{ + if (slot_idx >= n_slots) return false; + cur_slot_idx = slot_idx; + return true; +} + +void cosyvoice_kv_cache::slide_kv_slot() { - GGML_ASSERT(layer_idx + stride <= layers); + GGML_ASSERT(cur_slot_idx < n_slots); if (fattn) { - const auto end = layer_idx + stride; + auto layer_idx = cur_slot_idx * layers; + const auto end = layer_idx + layers; for (int cur = layer_idx; cur != end; ++cur) { auto& layer = kv_cache_layers[cur]; - auto& next_layer = kv_cache_layers[cur + stride]; + auto& next_layer = kv_cache_layers[cur + layers]; auto k = next_layer.k; auto k_view = layer.k_view; diff --git a/src/cosyvoice-kv-cache.h b/src/cosyvoice-kv-cache.h index 5bf2671..667c318 100644 --- a/src/cosyvoice-kv-cache.h +++ b/src/cosyvoice-kv-cache.h @@ -11,15 +11,6 @@ class cosyvoice_kv_cache cosyvoice_kv_cache() = default; ~cosyvoice_kv_cache(); - struct kv_cache_layer - { - ggml_tensor* k; - ggml_tensor* v; - - ggml_tensor* k_view; - ggml_tensor* v_view; - }; - void build_kv_cache( ggml_backend_t backend, ggml_backend_buffer_ptr& buffer, @@ -31,6 +22,8 @@ class cosyvoice_kv_cache ggml_type k_type, ggml_type v_type, int batch_size, + int n_kv_slots, + int n_offloaded_kv_slots, bool fattn); void update_cache(ggml_context* ctx0, ggml_cgraph* gf, ggml_tensor*& k, ggml_tensor*& v, ggml_tensor* position_ids, int layer_idx); @@ -40,9 +33,10 @@ class cosyvoice_kv_cache void shift_kv_node_pos(uint32_t shift_pos); bool can_reuse(bool prefill) const; - void slide_kv_layers(int layer_idx, int stride); + bool bind_slot(int slot_idx); + void slide_kv_slot(); - ggml_backend_buffer* initialize_buffer(ggml_backend_t backend, int k_head_dim, int v_head_dim, int num_key_value_heads, uint32_t max_seq, ggml_type k_type, ggml_type v_type, int batch_size, bool fattn); + ggml_backend_buffer* initialize_buffer(ggml_backend_t backend, int k_head_dim, int v_head_dim, uint32_t max_seq, int batch_size); uint32_t reset_buffer(ggml_backend_buffer* buffer); void offload_cache(ggml_backend_t backend, ggml_backend_sched* sched, uint32_t n_tokens); @@ -50,14 +44,20 @@ class cosyvoice_kv_cache size_t get_offloaded_cache_size() const; void clear_offloaded_cache(); + void offload_slot(ggml_backend_t backend, ggml_backend_sched* sched, int offloaded_slot_idx); + void load_slot(ggml_backend_t backend, ggml_backend_sched* sched, int offloaded_slot_idx); + uint32_t cur_len; private: int layers; int num_heads; bool fattn; + int n_slots; + int n_offloaded_kv_slots; ggml_type k_type; ggml_type v_type; + int cur_slot_idx; ggml_context* ctx; - kv_cache_layer* kv_cache_layers; + struct kv_cache_layer* kv_cache_layers; struct offloaded_kv_cache* offloaded_cache; }; diff --git a/src/cosyvoice-loader.cpp b/src/cosyvoice-loader.cpp index 33fb865..0cd14f2 100644 --- a/src/cosyvoice-loader.cpp +++ b/src/cosyvoice-loader.cpp @@ -2,6 +2,7 @@ #include "cosyvoice-loader.h" #include "cosyvoice-kv-cache.h" +#include #include #include #include @@ -449,32 +450,32 @@ void CosyVoice3LM::OnLoad(const gguf_loader& loader, const std::string& prefix, } } -static ggml_type cosyvoice_llm_kv_cache_type_to_ggml(cosyvoice_llm_kv_cache_type_t t) +static ggml_type cosyvoice_kv_cache_type_to_ggml(cosyvoice_kv_cache_type_t t) { switch (t) { - case COSYVOICE_LLM_KV_CACHE_TYPE_F32: return GGML_TYPE_F32; - case COSYVOICE_LLM_KV_CACHE_TYPE_F16: return GGML_TYPE_F16; - case COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0: return GGML_TYPE_Q8_0; - case COSYVOICE_LLM_KV_CACHE_TYPE_Q5_1: return GGML_TYPE_Q5_1; - case COSYVOICE_LLM_KV_CACHE_TYPE_Q5_0: return GGML_TYPE_Q5_0; - case COSYVOICE_LLM_KV_CACHE_TYPE_Q4_1: return GGML_TYPE_Q4_1; - case COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0: return GGML_TYPE_Q4_0; + case COSYVOICE_KV_CACHE_TYPE_F32: return GGML_TYPE_F32; + case COSYVOICE_KV_CACHE_TYPE_F16: return GGML_TYPE_F16; + case COSYVOICE_KV_CACHE_TYPE_Q8_0: return GGML_TYPE_Q8_0; + case COSYVOICE_KV_CACHE_TYPE_Q5_1: return GGML_TYPE_Q5_1; + case COSYVOICE_KV_CACHE_TYPE_Q5_0: return GGML_TYPE_Q5_0; + case COSYVOICE_KV_CACHE_TYPE_Q4_1: return GGML_TYPE_Q4_1; + case COSYVOICE_KV_CACHE_TYPE_Q4_0: return GGML_TYPE_Q4_0; default: GGML_ABORT("unexpected kv cache type"); } } -static cosyvoice_llm_kv_cache_type_t cosyvoice_ggml_to_llm_kv_cache_type(ggml_type t) +static cosyvoice_kv_cache_type_t cosyvoice_ggml_to_kv_cache_type(ggml_type t) { switch (t) { - case GGML_TYPE_F32: return COSYVOICE_LLM_KV_CACHE_TYPE_F32; - case GGML_TYPE_F16: return COSYVOICE_LLM_KV_CACHE_TYPE_F16; - case GGML_TYPE_Q8_0: return COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0; - case GGML_TYPE_Q5_1: return COSYVOICE_LLM_KV_CACHE_TYPE_Q5_1; - case GGML_TYPE_Q5_0: return COSYVOICE_LLM_KV_CACHE_TYPE_Q5_0; - case GGML_TYPE_Q4_1: return COSYVOICE_LLM_KV_CACHE_TYPE_Q4_1; - case GGML_TYPE_Q4_0: return COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0; + case GGML_TYPE_F32: return COSYVOICE_KV_CACHE_TYPE_F32; + case GGML_TYPE_F16: return COSYVOICE_KV_CACHE_TYPE_F16; + case GGML_TYPE_Q8_0: return COSYVOICE_KV_CACHE_TYPE_Q8_0; + case GGML_TYPE_Q5_1: return COSYVOICE_KV_CACHE_TYPE_Q5_1; + case GGML_TYPE_Q5_0: return COSYVOICE_KV_CACHE_TYPE_Q5_0; + case GGML_TYPE_Q4_1: return COSYVOICE_KV_CACHE_TYPE_Q4_1; + case GGML_TYPE_Q4_0: return COSYVOICE_KV_CACHE_TYPE_Q4_0; default: GGML_ABORT("unexpected ggml type for kv cache"); } } @@ -493,12 +494,220 @@ static ggml_type cosyvoice_get_kv_fallback_type(ggml_type t) } } +union kv_cache_type_union +{ + struct + { +#if defined(__BYTE_ORDER__) && (__BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) || defined(_BYTE_ORDER) && (_BYTE_ORDER == _BIG_ENDIAN) || defined(__BIG_ENDIAN__) && !defined(__LITTLE_ENDIAN__) || defined(__ARMEB__) || defined(__MIPSEB__) || defined(__sparc__) + uint32_t kv_cache_separate_buffers : 1; + cosyvoice_kv_cache_type_t : 16; + cosyvoice_kv_cache_type_t kv_cache_fallback : 5; + cosyvoice_kv_cache_type_t v_cache_type : 5; + cosyvoice_kv_cache_type_t k_cache_type : 5; +#else + cosyvoice_kv_cache_type_t k_cache_type : 5; + cosyvoice_kv_cache_type_t v_cache_type : 5; + cosyvoice_kv_cache_type_t kv_cache_fallback : 5; + cosyvoice_kv_cache_type_t : 16; + uint32_t kv_cache_separate_buffers : 1; +#endif + }; + cosyvoice_kv_cache_type_t kv_cache_type; +}; + +static std::tuple cosyvoice_check_kv_cache_types( + ggml_context* ctx, ggml_backend_t backend, + bool& fattn, int num_attn_heads, int num_kv_heads, + const Linear& q_proj, const Linear& k_proj, const Linear& v_proj, + bool kv_fallback, kv_cache_type_union& kv_type_union) +{ + ggml_type k_type, v_type; + + if (fattn) + { + auto fattn_check = [&](ggml_type check_k, ggml_type check_v) -> bool + { + auto q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, q_proj.weight->ne[1] / num_attn_heads, 1, num_attn_heads); + auto k = ggml_new_tensor_3d(ctx, check_k, k_proj.weight->ne[1] / num_kv_heads, 1, num_kv_heads); + auto v = ggml_new_tensor_3d(ctx, check_v, v_proj.weight->ne[1] / num_kv_heads, 1, num_kv_heads); + auto o = ggml_flash_attn_ext(ctx, q, k, v, nullptr, 1.f / std::sqrt(static_cast(k->ne[0])), 0.f, 0.f); + return ggml_backend_supports_op(backend, o); + }; + + fattn = fattn_check(GGML_TYPE_F32, GGML_TYPE_F32); + if (fattn) + { + if (kv_fallback) + { + ggml_type cur_type; + + do + { + if (kv_type_union.kv_cache_separate_buffers) + { + if (k_type = cosyvoice_kv_cache_type_to_ggml(kv_type_union.k_cache_type), + v_type = cosyvoice_kv_cache_type_to_ggml(kv_type_union.v_cache_type); + fattn_check(k_type, v_type)) + { + kv_type_union.kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( + kv_type_union.k_cache_type, + kv_type_union.v_cache_type, + kv_type_union.v_cache_type); + break; + } + + cur_type = cosyvoice_kv_cache_type_to_ggml(kv_type_union.kv_cache_fallback); + } + else + cur_type = cosyvoice_kv_cache_type_to_ggml(kv_type_union.kv_cache_type); + + do + { + if (fattn_check(cur_type, cur_type)) + { + k_type = cur_type; + v_type = cur_type; + kv_type_union.kv_cache_type = cosyvoice_ggml_to_kv_cache_type(cur_type); + break; + } + + cur_type = cosyvoice_get_kv_fallback_type(cur_type); + } while (cur_type != GGML_TYPE_F32); + + if (kv_type_union.kv_cache_separate_buffers) + kv_type_union.kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( + cosyvoice_ggml_to_kv_cache_type(k_type), + cosyvoice_ggml_to_kv_cache_type(v_type), + cosyvoice_ggml_to_kv_cache_type(v_type)); + } while (false); + } + else if (kv_type_union.kv_cache_separate_buffers) + { + k_type = cosyvoice_kv_cache_type_to_ggml(kv_type_union.k_cache_type); + v_type = cosyvoice_kv_cache_type_to_ggml(kv_type_union.v_cache_type); + fattn = fattn_check(k_type, v_type); + if (fattn) + kv_type_union.kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( + kv_type_union.k_cache_type, + kv_type_union.v_cache_type, + kv_type_union.v_cache_type); + } + else + { + k_type = cosyvoice_kv_cache_type_to_ggml(kv_type_union.kv_cache_type); + v_type = k_type; + fattn = fattn_check(k_type, v_type); + } + } + } + + if (!fattn) + { + ggml_type cur_type; + auto attn_check = [&](ggml_type check_k, ggml_type check_v) -> bool + { + if (ggml_is_quantized(check_v)) return false; + auto q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, q_proj.weight->ne[1] / num_attn_heads, 1, num_attn_heads); + auto k = ggml_new_tensor_3d(ctx, check_k, k_proj.weight->ne[1] / num_kv_heads, 1, num_kv_heads); + auto v = ggml_new_tensor_3d(ctx, check_v, 1, v_proj.weight->ne[1] / num_kv_heads, num_kv_heads); + auto s = ggml_mul_mat(ctx, k, q); + auto o = ggml_mul_mat(ctx, v, s); + return ggml_backend_supports_op(backend, o) && ggml_backend_supports_op(backend, s); + }; + + do + { + if (kv_type_union.kv_cache_separate_buffers) + { + if (k_type = cosyvoice_kv_cache_type_to_ggml(kv_type_union.k_cache_type), + v_type = cosyvoice_kv_cache_type_to_ggml(kv_type_union.v_cache_type); + attn_check(k_type, v_type)) + { + kv_type_union.kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( + kv_type_union.k_cache_type, + kv_type_union.v_cache_type, + kv_type_union.v_cache_type); + break; + } + else if (kv_fallback && attn_check(k_type, GGML_TYPE_F16)) + { + k_type = cosyvoice_kv_cache_type_to_ggml(kv_type_union.k_cache_type); + v_type = GGML_TYPE_F16; + kv_type_union.kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( + kv_type_union.k_cache_type, + COSYVOICE_KV_CACHE_TYPE_F16, + COSYVOICE_KV_CACHE_TYPE_F16); + break; + } + else + cur_type = cosyvoice_kv_cache_type_to_ggml(kv_type_union.kv_cache_fallback); + } + else + cur_type = cosyvoice_kv_cache_type_to_ggml(kv_type_union.kv_cache_type); + + if (kv_type_union.kv_cache_separate_buffers) + { + v_type = ggml_is_quantized(cur_type) ? GGML_TYPE_F16 : cur_type; + do + { + if (attn_check(cur_type, v_type)) + { + k_type = cur_type; + break; + } + + GGML_ASSERT(kv_fallback); + cur_type = cosyvoice_get_kv_fallback_type(cur_type); + } while (cur_type != GGML_TYPE_F32); + + kv_type_union.kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( + cosyvoice_ggml_to_kv_cache_type(cur_type), + cosyvoice_ggml_to_kv_cache_type(v_type), + cosyvoice_ggml_to_kv_cache_type(v_type)); + + k_type = cur_type; + } + else + { + if (attn_check(cur_type, cur_type)) + { + k_type = cur_type; + v_type = cur_type; + kv_type_union.kv_cache_type = cosyvoice_ggml_to_kv_cache_type(cur_type); + break; + } + else + { + GGML_ASSERT(kv_fallback); + cur_type = GGML_TYPE_F16; + } + k_type = cur_type; + v_type = cur_type; + kv_type_union.kv_cache_type = cosyvoice_ggml_to_kv_cache_type(cur_type); + } + } while (false); + } + + return { k_type, v_type }; +} + void cosyvoice_model_3::load(gguf_loader& loader) { auto& flow = cv3_shared->flow; auto& hift = cv3_shared->hift; auto& llm = cv3_shared->llm; + { + constexpr auto diffusion_steps = CausalConditionalCFM::diffusion_steps; + auto& n_fixed_slots = shared->params.dit_kv_fixed_slots; + auto& n_offloadable_slots = shared->params.dit_kv_offloadable_slots; + if (n_fixed_slots > diffusion_steps) + n_fixed_slots = diffusion_steps; + if (n_offloadable_slots + n_fixed_slots > diffusion_steps) + n_offloadable_slots = diffusion_steps - n_fixed_slots; + } + + flow.OnLoad(loader, {}); hift.OnLoad(loader, {}); llm.OnLoad(loader, {}, shared->params); @@ -688,186 +897,15 @@ void cosyvoice_model_3::load(gguf_loader& loader) for (uint32_t i = 0; i != shared->params.n_workers; ++i) workers[i].config = shared->config; - if (shared->params.llm_use_flash_attn) - { - auto fattn_check = [&](ggml_type check_k, ggml_type check_v) -> bool - { - auto q = ggml_new_tensor_3d(worker->ctx0.get(), GGML_TYPE_F32, llm.layers[0].self_attn.q_proj.weight->ne[1] / llm.num_attention_heads, 1, llm.num_attention_heads); - auto k = ggml_new_tensor_3d(worker->ctx0.get(), check_k, llm.layers[0].self_attn.k_proj.weight->ne[1] / llm.num_key_value_heads, 1, llm.num_key_value_heads); - auto v = ggml_new_tensor_3d(worker->ctx0.get(), check_v, llm.layers[0].self_attn.v_proj.weight->ne[1] / llm.num_key_value_heads, 1, llm.num_key_value_heads); - auto o = ggml_flash_attn_ext(worker->ctx0.get(), q, k, v, nullptr, 1.f / std::sqrt(static_cast(k->ne[0])), 0.f, 0.f); - return ggml_backend_supports_op(backend, o); - }; - - shared->params.llm_use_flash_attn = fattn_check(GGML_TYPE_F32, GGML_TYPE_F32); - if (shared->params.llm_use_flash_attn) - { - if (shared->params.llm_allow_kv_cache_fallback) - { - ggml_type cur_type; - - do - { - if (shared->params.llm_kv_cache_separate_buffers) - { - if (auto k_type = cosyvoice_llm_kv_cache_type_to_ggml(shared->params.llm_k_cache_type), - v_type = cosyvoice_llm_kv_cache_type_to_ggml(shared->params.llm_v_cache_type); - fattn_check(k_type, v_type)) - { - cv3_shared->k_type = k_type; - cv3_shared->v_type = v_type; - shared->params.llm_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( - shared->params.llm_k_cache_type, - shared->params.llm_v_cache_type, - shared->params.llm_v_cache_type); - break; - } - - cur_type = cosyvoice_llm_kv_cache_type_to_ggml(shared->params.llm_kv_cache_fallback); - } - else - cur_type = cosyvoice_llm_kv_cache_type_to_ggml(shared->params.llm_kv_cache_type); - - do - { - if (fattn_check(cur_type, cur_type)) - { - cv3_shared->k_type = cur_type; - cv3_shared->v_type = cur_type; - shared->params.llm_kv_cache_type = cosyvoice_ggml_to_llm_kv_cache_type(cur_type); - break; - } - - cur_type = cosyvoice_get_kv_fallback_type(cur_type); - } while (cur_type != GGML_TYPE_F32); - - if (shared->params.llm_kv_cache_separate_buffers) - shared->params.llm_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( - cosyvoice_ggml_to_llm_kv_cache_type(cv3_shared->k_type), - cosyvoice_ggml_to_llm_kv_cache_type(cv3_shared->v_type), - cosyvoice_ggml_to_llm_kv_cache_type(cv3_shared->v_type)); - } while (false); - } - else if (shared->params.llm_kv_cache_separate_buffers) - { - auto k_type = cosyvoice_llm_kv_cache_type_to_ggml(shared->params.llm_k_cache_type); - auto v_type = cosyvoice_llm_kv_cache_type_to_ggml(shared->params.llm_v_cache_type); - shared->params.llm_use_flash_attn = fattn_check(k_type, v_type); - if (shared->params.llm_use_flash_attn) - { - cv3_shared->k_type = k_type; - cv3_shared->v_type = v_type; - shared->params.llm_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( - shared->params.llm_k_cache_type, - shared->params.llm_v_cache_type, - shared->params.llm_v_cache_type); - } - } - else - { - auto kv_type = cosyvoice_llm_kv_cache_type_to_ggml(shared->params.llm_kv_cache_type); - shared->params.llm_use_flash_attn = fattn_check(kv_type, kv_type); - } - } - } - - if (!shared->params.llm_use_flash_attn) - { - ggml_type cur_type; - auto attn_check = [&](ggml_type check_k, ggml_type check_v) -> bool - { - if (ggml_is_quantized(check_v)) return false; - auto q = ggml_new_tensor_3d(worker->ctx0.get(), GGML_TYPE_F32, llm.layers[0].self_attn.q_proj.weight->ne[1] / llm.num_attention_heads, 1, llm.num_attention_heads); - auto k = ggml_new_tensor_3d(worker->ctx0.get(), check_k, llm.layers[0].self_attn.k_proj.weight->ne[1] / llm.num_key_value_heads, 1, llm.num_key_value_heads); - auto v = ggml_new_tensor_3d(worker->ctx0.get(), check_v, 1, llm.layers[0].self_attn.v_proj.weight->ne[1] / llm.num_key_value_heads, llm.num_key_value_heads); - auto s = ggml_mul_mat(worker->ctx0.get(), k, q); - auto o = ggml_mul_mat(worker->ctx0.get(), v, s); - return ggml_backend_supports_op(backend, o) && ggml_backend_supports_op(backend, s); - }; - - do - { - if (shared->params.llm_kv_cache_separate_buffers) - { - if (auto k_type = cosyvoice_llm_kv_cache_type_to_ggml(shared->params.llm_k_cache_type), - v_type = cosyvoice_llm_kv_cache_type_to_ggml(shared->params.llm_v_cache_type); - attn_check(k_type, v_type)) - { - cv3_shared->k_type = k_type; - cv3_shared->v_type = v_type; - shared->params.llm_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( - shared->params.llm_k_cache_type, - shared->params.llm_v_cache_type, - shared->params.llm_v_cache_type); - break; - } - else if (shared->params.llm_allow_kv_cache_fallback && attn_check(k_type, GGML_TYPE_F16)) - { - cv3_shared->k_type = cosyvoice_llm_kv_cache_type_to_ggml(shared->params.llm_k_cache_type); - cv3_shared->v_type = GGML_TYPE_F16; - shared->params.llm_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( - shared->params.llm_k_cache_type, - COSYVOICE_LLM_KV_CACHE_TYPE_F16, - COSYVOICE_LLM_KV_CACHE_TYPE_F16); - break; - } - else - cur_type = cosyvoice_llm_kv_cache_type_to_ggml(shared->params.llm_kv_cache_fallback); - } - else - cur_type = cosyvoice_llm_kv_cache_type_to_ggml(shared->params.llm_kv_cache_type); - - if (shared->params.llm_kv_cache_separate_buffers) - { - auto v_type = ggml_is_quantized(cur_type) ? GGML_TYPE_F16 : cur_type; - do - { - if (attn_check(cur_type, v_type)) - { - cv3_shared->k_type = cur_type; - cv3_shared->v_type = v_type; - break; - } - - GGML_ASSERT(shared->params.llm_allow_kv_cache_fallback); - cur_type = cosyvoice_get_kv_fallback_type(cur_type); - } while (cur_type != GGML_TYPE_F32); + auto [llm_k_type, llm_v_type] = cosyvoice_check_kv_cache_types(worker->ctx0.get(), worker->backend.get(), + shared->params.flow_use_flash_attn, cv3_shared->llm.num_attention_heads, cv3_shared->llm.num_key_value_heads, + cv3_shared->llm.layers[0].self_attn.q_proj, cv3_shared->llm.layers[0].self_attn.k_proj, cv3_shared->llm.layers[0].self_attn.v_proj, + shared->params.llm_allow_kv_cache_fallback, reinterpret_cast(shared->params.llm_kv_cache_type)); - shared->params.llm_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( - cosyvoice_ggml_to_llm_kv_cache_type(cur_type), - cosyvoice_ggml_to_llm_kv_cache_type(v_type), - cosyvoice_ggml_to_llm_kv_cache_type(v_type)); - } - else - { - if (attn_check(cur_type, cur_type)) - { - cv3_shared->k_type = cur_type; - cv3_shared->v_type = cur_type; - shared->params.llm_kv_cache_type = cosyvoice_ggml_to_llm_kv_cache_type(cur_type); - break; - } - else - { - GGML_ASSERT(shared->params.llm_allow_kv_cache_fallback); - cur_type = GGML_TYPE_F16; - } - cv3_shared->k_type = cur_type; - cv3_shared->v_type = cur_type; - shared->params.llm_kv_cache_type = cosyvoice_ggml_to_llm_kv_cache_type(cur_type); - } - } while (false); - } - - if (shared->params.flow_use_flash_attn) - { - int heads = flow.decoder.estimator.transformer_blocks[0].attn.heads; - auto q = ggml_new_tensor_4d(worker->ctx0.get(), GGML_TYPE_F32, flow.decoder.estimator.transformer_blocks[0].attn.to_q.weight->ne[1] / heads, 1, heads, 2); - auto k = ggml_new_tensor_4d(worker->ctx0.get(), GGML_TYPE_F32, flow.decoder.estimator.transformer_blocks[0].attn.to_k.weight->ne[1] / heads, 1, heads, 2); - auto v = ggml_new_tensor_4d(worker->ctx0.get(), GGML_TYPE_F32, flow.decoder.estimator.transformer_blocks[0].attn.to_v.weight->ne[1] / heads, 1, heads, 2); - auto o = ggml_flash_attn_ext(worker->ctx0.get(), q, k, v, nullptr, 1.f / std::sqrt(static_cast(k->ne[0])), 0.f, 0.f); - shared->params.flow_use_flash_attn = ggml_backend_supports_op(backend, o); - } + auto [dit_k_type, dit_v_type] = cosyvoice_check_kv_cache_types(worker->ctx0.get(), worker->backend.get(), + shared->params.flow_use_flash_attn, cv3_shared->flow.decoder.estimator.transformer_blocks[0].attn.heads, cv3_shared->flow.decoder.estimator.transformer_blocks[0].attn.heads, + cv3_shared->flow.decoder.estimator.transformer_blocks[0].attn.to_q, cv3_shared->flow.decoder.estimator.transformer_blocks[0].attn.to_k, cv3_shared->flow.decoder.estimator.transformer_blocks[0].attn.to_v, + shared->params.dit_allow_kv_cache_fallback, reinterpret_cast(shared->params.dit_kv_cache_type)); for (auto& worker : std::span(workers, shared->params.n_workers)) { @@ -885,24 +923,28 @@ void cosyvoice_model_3::load(gguf_loader& loader) static_cast(llm.layers[0].self_attn.v_proj.weight->ne[1] / llm.num_key_value_heads), llm.num_key_value_heads, shared->params.n_max_seq, - cv3_shared->k_type, - cv3_shared->v_type, + llm_k_type, + llm_v_type, 1, + 1, + shared->params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_DEDICATED ? 0 : 1, shared->params.llm_use_flash_attn ); - + const auto dit_blocks = flow.decoder.estimator.transformer_blocks; worker.dit_kv_cache.build_kv_cache( backend, worker.dit_kv_buffer, - static_cast(dit_blocks.size()* (flow.decoder.t_span.size() - 1)), + static_cast(dit_blocks.size()), static_cast(dit_blocks[0].attn.to_k.weight->ne[1] / dit_blocks[0].attn.heads), static_cast(dit_blocks[0].attn.to_v.weight->ne[1] / dit_blocks[0].attn.heads), dit_blocks[0].attn.heads, - shared->params.n_max_seq, - GGML_TYPE_Q8_0, - GGML_TYPE_Q4_0, + shared->params.dit_kv_cache_length, + dit_k_type, + dit_v_type, 2, + shared->params.dit_kv_fixed_slots, + shared->params.dit_kv_offloadable_slots, shared->params.flow_use_flash_attn ); } diff --git a/src/cosyvoice-model.cpp b/src/cosyvoice-model.cpp index acab562..24c73ef 100644 --- a/src/cosyvoice-model.cpp +++ b/src/cosyvoice-model.cpp @@ -46,7 +46,7 @@ void cosyvoice_init_default_context_params(cosyvoice_context_params_t* params) params->flow_use_flash_attn = true; params->llm_use_flash_attn = true; - params->llm_kv_cache_type = COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0; + params->llm_kv_cache_type = COSYVOICE_KV_CACHE_TYPE_Q8_0; params->llm_allow_kv_cache_fallback = true; params->inference_buffer_policy = COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED; @@ -60,7 +60,7 @@ void cosyvoice_init_default_context_params(cosyvoice_context_params_t* params) params->sampler_ctx = nullptr; } -cosyvoice_model_shared::cosyvoice_model_shared(const cosyvoice_context_params_v2_cpp& params) +cosyvoice_model_shared::cosyvoice_model_shared(const cosyvoice_context_params_v3_cpp& params) : params(params), ctx(nullptr), backend_uma(false), rand_noise_len(0), noise_callback(nullptr), noise_callback_ctx(nullptr) {} cosyvoice_worker_context::cosyvoice_worker_context(ggml_backend_t backend) @@ -69,7 +69,7 @@ cosyvoice_worker_context::cosyvoice_worker_context(ggml_backend_t backend) gf(nullptr), llm_input(nullptr), llm_probs(nullptr), position_ids(nullptr), causal_mask(nullptr), llm_kv_cache(), status(GGML_STATUS_SUCCESS), prompt_crc32(0), sampler_seed(0), sampler(nullptr), sampler_ctx(nullptr), builtin_sampler_rng_policy(COSYVOICE_BUILTIN_SAMPLER_RNG_POLICY_RESET_PER_SESSION), nucleus_probs_capacity(0), nucleus_probs_len(0) {} -cosyvoice_model::cosyvoice_model(ggml_backend_t backend, const cosyvoice_context_params_v2_cpp& params) +cosyvoice_model::cosyvoice_model(ggml_backend_t backend, const cosyvoice_context_params_v3_cpp& params) : shared(new cosyvoice_model_shared(params)), workers(reinterpret_cast(malloc(sizeof(cosyvoice_worker_context) * params.n_workers))) { auto dev = ggml_backend_get_device(backend); @@ -267,12 +267,8 @@ void cosyvoice_model_3::empty_buffer_cache() worker->backend.get(), static_cast(cv3_shared->llm.layers[0].self_attn.k_proj.weight->ne[1] / cv3_shared->llm.num_key_value_heads), static_cast(cv3_shared->llm.layers[0].self_attn.v_proj.weight->ne[1] / cv3_shared->llm.num_key_value_heads), - cv3_shared->llm.num_key_value_heads, shared->params.n_max_seq, - cv3_shared->k_type, - cv3_shared->v_type, - 1, - shared->params.llm_use_flash_attn)); + 1)); switch (shared->params.inference_buffer_policy) { @@ -334,22 +330,10 @@ uint32_t cosyvoice_model_3::get_chunk_tokens() return 25 + cv3_shared->flow.pre_lookahead_layer.pre_lookahead_len; } -uint32_t cosyvoice_model_3::get_flow_overlap_tokens() -{ - // TODO: calculate the actual overlap tokens - return 0; -} - -uint32_t cosyvoice_model_3::get_hift_overlap_tokens() -{ - // TODO: calculate the actual overlap tokens - return 0; -} - cosyvoice_3_worker_context::cosyvoice_3_worker_context() : ctx1(ggml_init(ggml_init_params{ .mem_size = ggml_tensor_overhead() * 4, .no_alloc = true })) {} -cosyvoice_model_3::cosyvoice_model_3(ggml_backend_t backend, const cosyvoice_context_params_v2_cpp& params) +cosyvoice_model_3::cosyvoice_model_3(ggml_backend_t backend, const cosyvoice_context_params_v3_cpp& params) : cosyvoice_model(backend, params), cv3_shared(new cosyvoice_model_3_shared), cv3_workers(new cosyvoice_3_worker_context[params.n_workers]()) { cv3_worker = cv3_workers; diff --git a/src/cosyvoice-model.h b/src/cosyvoice-model.h index 4cc4c87..064d655 100644 --- a/src/cosyvoice-model.h +++ b/src/cosyvoice-model.h @@ -14,14 +14,14 @@ struct cosyvoice_model_shared { - cosyvoice_model_shared(const cosyvoice_context_params_v2_cpp& params); + cosyvoice_model_shared(const cosyvoice_context_params_v3_cpp& params); ggml_context_ptr ctx; ggml_backend_buffer_ptr buffer; ggml_backend_buffer_ptr cpu_buffer; - cosyvoice_context_params_v2_cpp params; + cosyvoice_context_params_v3_cpp params; ggml_backend_op_capabilities op_caps; bool backend_uma; @@ -83,7 +83,7 @@ struct cosyvoice_worker_context struct cosyvoice_model : virtual cosyvoice_model_context, virtual cosyvoice_object_ref_counter { - cosyvoice_model(ggml_backend_t backend, const cosyvoice_context_params_v2_cpp& params); + cosyvoice_model(ggml_backend_t backend, const cosyvoice_context_params_v3_cpp& params); ~cosyvoice_model(); virtual void load(gguf_loader& loader) = 0; @@ -143,9 +143,6 @@ struct cosyvoice_model_3_shared CausalHiFTGenerator hift; CosyVoice3LM llm; - ggml_type k_type; - ggml_type v_type; - std::set stop_tokens; std::set silent_tokens; }; @@ -163,7 +160,7 @@ struct cosyvoice_3_worker_context struct cosyvoice_model_3 : cosyvoice_model { - cosyvoice_model_3(ggml_backend_t backend, const cosyvoice_context_params_v2_cpp& params); + cosyvoice_model_3(ggml_backend_t backend, const cosyvoice_context_params_v3_cpp& params); ~cosyvoice_model_3(); void load(gguf_loader& loader); @@ -189,8 +186,6 @@ struct cosyvoice_model_3 : cosyvoice_model bool token2wav(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, cosyvoice_generated_speech_ptr result); bool token2wav_ext(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* offset, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result); uint32_t get_chunk_tokens(); - uint32_t get_flow_overlap_tokens(); - uint32_t get_hift_overlap_tokens(); void empty_buffer_cache(); void get_memory_usage(cosyvoice_memory_usage_t* usage); diff --git a/src/cosyvoice-modules.h b/src/cosyvoice-modules.h index 110d1be..ea43c68 100644 --- a/src/cosyvoice-modules.h +++ b/src/cosyvoice-modules.h @@ -131,7 +131,7 @@ struct Attention : Module void OnLoad(const gguf_loader& loader, const std::string& prefix); - ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* position_ids, int64_t cut_len, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, int layer_idx) const; + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* position_ids, int64_t cut_len, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, ggml_tensor* attn_mask, int layer_idx) const; }; struct FeedForward : Module @@ -154,7 +154,7 @@ struct DiTBlock : Module void OnLoad(const gguf_loader& loader, const std::string& prefix); - ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* time_emb, ggml_tensor* position_ids, int64_t cut_len, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, int layer_idx) const; + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* time_emb, ggml_tensor* position_ids, int64_t cut_len, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, ggml_tensor* attn_mask, int layer_idx) const; }; struct AdaLayerNorm_Final : Module @@ -183,12 +183,13 @@ struct DiT : Module void OnLoad(const gguf_loader& loader, const std::string& prefix); - ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* mu, ggml_tensor* t, ggml_tensor* spks, ggml_tensor* cond, int64_t cut_len, ggml_tensor*& position_ids, ggml_backend_op_capabilities capabilities, cosyvoice_kv_cache* kv_cache, ggml_cgraph* gf, int step) const; + ggml_tensor* build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_tensor* mu, ggml_tensor* t, ggml_tensor* spks, ggml_tensor* cond, int64_t cut_len, ggml_tensor*& position_ids, ggml_backend_op_capabilities capabilities, cosyvoice_kv_cache* kv_cache, ggml_tensor** ref_attn_mask, ggml_cgraph* gf) const; }; struct CausalConditionalCFM : Module { - std::array t_span; + constexpr static int diffusion_steps = 10; + std::array t_span; float inference_cfg_rate; DiT estimator; @@ -206,7 +207,7 @@ struct CausalConditionalCFM : Module DiTContext prepare_context(ggml_context* ctx, ggml_tensor* mu, ggml_tensor* spks, ggml_tensor* cond) const; std::array get_t_and_dt(ggml_context* ctx, int step) const; - ggml_tensor* build_cgraph_one_step(ggml_context* ctx, const DiTContext& ditctx, int step, ggml_backend_op_capabilities capabilities, int64_t cut_len, ggml_tensor*& t_tensor, ggml_tensor*& position_ids, ggml_cgraph* gf, cosyvoice_kv_cache* kv_cache) const; + ggml_tensor* build_cgraph_one_step(ggml_context* ctx, const DiTContext& ditctx, int step, ggml_backend_op_capabilities capabilities, int64_t cut_len, ggml_tensor*& t_tensor, ggml_tensor*& position_ids, ggml_cgraph* gf, cosyvoice_kv_cache* kv_cache, ggml_tensor** attn_mask) const; }; struct PreLookaheadLayer : Module diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index 143500b..2cd0c23 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -311,6 +311,75 @@ bool cosyvoice_model_3::token2wav(const int* token_ids, uint32_t n_tokens, float return token2wav_ext(token_ids, n_tokens, speed, prompt, nullptr, false, true, result); } +template +struct dit_sched_config +{ + struct graph_config_t + { + bool rebuild; + bool cache_kv; + bool load; + bool offload; + bool slide; + bool slice; + bool mask; + int64_t cut_len; + }; + + graph_config_t graph_config[n]; + + const auto& operator[](int i) const { return graph_config[i]; } + + dit_sched_config(const cosyvoice_context_params_v3_cpp& params, int64_t cut_len, uint32_t offset, bool streaming) + { + if (!streaming) + { + for (int i = 0; i != n; ++i) + { + graph_config[i].rebuild = false; + graph_config[i].cache_kv = false; + graph_config[i].load = false; + graph_config[i].offload = false; + graph_config[i].slide = false; + graph_config[i].slice = false; + graph_config[i].mask = false; + graph_config[i].cut_len = 0; + } + + graph_config[0].rebuild = true; + graph_config[1].rebuild = true; + graph_config[n - 1].rebuild = true; + graph_config[n - 1].cut_len = cut_len; + } + else + { + if (offset != 0) cut_len = 0; + + const auto n_offloadable_steps = params.dit_kv_offloadable_slots; + const auto n_fixed_steps = params.dit_kv_fixed_slots; + const auto n_no_cache_steps = static_cast(n) - n_offloadable_steps - n_fixed_steps; + + for (int i = 0; i != n; ++i) + { + graph_config[i].rebuild = i == 0 + || i == 1 + || offset != 0 && i == n_no_cache_steps - 1 + || i == n_no_cache_steps + || i == n - 1 && cut_len != 0; + graph_config[i].cache_kv = i >= n_no_cache_steps; + graph_config[i].offload = i >= n_no_cache_steps && i < n_no_cache_steps + n_offloadable_steps; + graph_config[i].load = graph_config[i].offload && offset != 0; + graph_config[i].mask = i < n_no_cache_steps && offset != 0; + graph_config[i].cut_len = i == n - 1 && offset == 0 ? cut_len : 0; + if (i == n_no_cache_steps - 1 && offset != 0) + graph_config[i].cut_len = offset; + graph_config[i].slice = offset != 0 && i == n_no_cache_steps; + graph_config[i].slide = graph_config[i].cache_kv && !graph_config[i].offload && i != n - 1; + } + } + } +}; + bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* offset, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result) { const auto& params = shared->params; @@ -334,12 +403,12 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f ggml_tensor* prompt_token = ggml_new_tensor_1d(ctx0.get(), GGML_TYPE_I32, prompt->flow_prompt_speech_tokens.second); ggml_tensor* prompt_feat = ggml_new_tensor_2d(ctx0.get(), GGML_TYPE_F32, prompt->prompt_speech_feat.shape[1], prompt->prompt_speech_feat.shape[0]); ggml_tensor* embedding = ggml_new_tensor_2d(ctx0.get(), GGML_TYPE_F32, prompt->flow_embedding.shape[1], prompt->flow_embedding.shape[0]); - auto kv_cache = streaming ? &worker->dit_kv_cache : nullptr; + auto kv_cache = &worker->dit_kv_cache; // Phase 1: Flow encode ggml_cgraph* gf = new_cgraph(ctx0.get()); - auto [mu, spks, conds, cut_len] = flow.build_cgraph_encode(ctx0.get(), token, prompt_token, prompt_feat, embedding, op_caps, offset ? *offset : 0, streaming); - if (offset) *offset += mu->ne[1] / flow.token_mel_ratio; + auto [mu, spks, conds, cut_len] = flow.build_cgraph_encode(ctx0.get(), token, prompt_token, prompt_feat, embedding, op_caps, (flow.decoder.diffusion_steps - params.dit_kv_fixed_slots - params.dit_kv_offloadable_slots) == 0 && offset ? *offset : 0, streaming); + const auto chunk_len = mu->ne[1]; auto ditctx = flow.decoder.prepare_context(ctx1.get(), mu, spks, conds); do { @@ -363,6 +432,13 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f } } while (false); + dit_sched_config config(params, cut_len, offset ? *offset : 0, streaming); + if (offset) + if (flow.decoder.diffusion_steps == params.dit_kv_fixed_slots + params.dit_kv_offloadable_slots) + *offset += chunk_len; + else + *offset = chunk_len; + int position_offset = kv_cache ? static_cast(kv_cache->cur_len) : 0; uint32_t noise_len = static_cast(ggml_nelements(ditctx.x)); float* noise_buffer = shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_BEFORE_FLOW, noise_len, nullptr, shared->noise_callback_ctx); @@ -371,19 +447,37 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f // Phase 2: Flow decode steps ggml_tensor* t_leaf; ggml_tensor* position_ids; - auto feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, 1, op_caps, cut_len, t_leaf, position_ids, gf, kv_cache); + ggml_tensor* attn_mask; + if (config[0].cache_kv) kv_cache->bind_slot(0); + auto feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, 1, op_caps, config[0].cut_len, t_leaf, position_ids, gf, config[0].cache_kv ? kv_cache : nullptr, config[0].mask ? &attn_mask : nullptr); ggml_build_forward_expand(gf, feat); set_graph_backends(gf, sched.get(), backend.get(), cpu_backend.get(), op_caps); ggml_backend_sched_synchronize(sched.get()); ggml_backend_sched_alloc_graph(sched.get(), gf); - if (!op_caps.fill) ggml_set_zero(t_leaf); - for (int64_t i = 0; i < position_ids->ne[1]; ++i) + auto post_process = [&](int step) { - auto cur_row = reinterpret_cast(position_ids->data) + i * position_ids->ne[0]; - for (int32_t j = 0; j < position_ids->ne[0]; ++j) - cur_row[j] = j + position_offset; - } + if (!op_caps.fill) ggml_set_zero(t_leaf); + for (int64_t i = 0; i < position_ids->ne[1]; ++i) + { + auto cur_row = reinterpret_cast(position_ids->data) + i * position_ids->ne[0]; + for (int32_t j = 0; j < position_ids->ne[0]; ++j) + cur_row[j] = j + position_offset; + } + if (config[step].mask) + { + const auto cs = flow.decoder.estimator.static_chunk_size; + for (int64_t i = 0; i < attn_mask->ne[0]; ++i) + { + int64_t block_end = ((i / cs) + 1) * cs; + auto row = reinterpret_cast(attn_mask->data) + i * attn_mask->ne[1]; + for (int64_t j = 0; j < attn_mask->ne[1]; ++j) + row[j] = j < block_end ? 0 : 0xFC00; + } + } + }; + + post_process(0); ggml_backend_tensor_set_async(backend.get(), token, token_ids, 0, token->nb[1]); ggml_backend_tensor_set_async(backend.get(), prompt_token, prompt->flow_prompt_speech_tokens.first.get(), 0, prompt_token->nb[1]); @@ -405,64 +499,57 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f memset(tensor->src, 0, sizeof(tensor->src)); } - ggml_backend_tensor_copy_async(backend.get(), backend.get(), feat, ditctx.x); - - ggml_reset(ctx0.get()); - ggml_backend_sched_reset(sched.get()); - gf = new_cgraph(ctx0.get()); - feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, 2, op_caps, cut_len, t_leaf, position_ids, gf, kv_cache); - - ggml_build_forward_expand(gf, feat); - set_graph_backends(gf, sched.get(), backend.get(), cpu_backend.get(), op_caps); - - ggml_backend_sched_alloc_graph(sched.get(), gf); - if (!op_caps.fill) ggml_set_zero(t_leaf); - for (int64_t i = 0; i < position_ids->ne[1]; ++i) - { - auto cur_row = reinterpret_cast(position_ids->data) + i * position_ids->ne[0]; - for (int32_t j = 0; j < position_ids->ne[0]; ++j) - cur_row[j] = j + position_offset; - } - worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); - if (worker->status != GGML_STATUS_SUCCESS) return false; - - auto scale_node = feat->src[1]; - GGML_ASSERT(scale_node->op == GGML_OP_SCALE); - int num_blocks = static_cast(flow.decoder.estimator.transformer_blocks.size()); - for (int step = 3; step != flow.decoder.t_span.size() - 1; ++step) + int offload_slot = 0; + int kv_slot = config[0].cache_kv && !config[0].offload; + for (int step = 1; step != flow.decoder.diffusion_steps; ++step) { + if (config[step].slice) + { + const auto cut_len = config[step - 1].cut_len; + reinterpret_cast(ditctx.cond_in->data) += static_cast(ditctx.cond_in->nb[1]) * cut_len; + ditctx.cond_in->ne[1] -= cut_len; + reinterpret_cast(ditctx.mu_in->data) += static_cast(ditctx.mu_in->nb[1]) * cut_len; + ditctx.mu_in->ne[1] -= cut_len; + + memcpy(ditctx.x->ne, feat->ne, sizeof(ditctx.x->ne)); + memcpy(ditctx.x->nb, feat->nb, sizeof(ditctx.x->nb)); + } ggml_backend_tensor_copy_async(backend.get(), backend.get(), feat, ditctx.x); - if (kv_cache) - kv_cache->slide_kv_layers((step - 2) * num_blocks, num_blocks); - auto [t, dt] = flow.decoder.get_t_and_dt(ctx0.get(), step); - reinterpret_cast(t_leaf->op_params)[op_caps.fill ? 0 : 1] = t; - reinterpret_cast(scale_node->op_params)[0] = dt; + if (config[step].rebuild) + { + if (config[step].cache_kv && !config[step].offload) + kv_cache->bind_slot(kv_slot++); + ggml_reset(ctx0.get()); + ggml_backend_sched_reset(sched.get()); + gf = new_cgraph(ctx0.get()); + feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, step + 1, op_caps, config[step].cut_len, t_leaf, position_ids, gf, config[step].cache_kv ? kv_cache : nullptr, config[step].mask ? &attn_mask : nullptr); + ggml_build_forward_expand(gf, feat); + set_graph_backends(gf, sched.get(), backend.get(), cpu_backend.get(), op_caps); + ggml_backend_sched_alloc_graph(sched.get(), gf); + post_process(step); + } + else + { + auto scale_node = feat->src[1]; + GGML_ASSERT(scale_node->op == GGML_OP_SCALE); + auto [t, dt] = flow.decoder.get_t_and_dt(ctx0.get(), step); + reinterpret_cast(t_leaf->op_params)[op_caps.fill ? 0 : 1] = t; + reinterpret_cast(scale_node->op_params)[0] = dt; + } + + if (config[step].load) + kv_cache->load_slot(backend.get(), sched.get(), offload_slot); worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); if (worker->status != GGML_STATUS_SUCCESS) return false; - } - - ggml_backend_tensor_copy_async(backend.get(), backend.get(), feat, ditctx.x); - ggml_reset(ctx0.get()); - ggml_backend_sched_reset(sched.get()); - gf = new_cgraph(ctx0.get()); - feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, static_cast(flow.decoder.t_span.size() - 1), op_caps, cut_len, t_leaf, position_ids, gf, kv_cache); - - ggml_build_forward_expand(gf, feat); - set_graph_backends(gf, sched.get(), backend.get(), cpu_backend.get(), op_caps); - ggml_backend_sched_alloc_graph(sched.get(), gf); - if (!op_caps.fill) ggml_set_zero(t_leaf); - for (int64_t i = 0; i < position_ids->ne[1]; ++i) - { - auto cur_row = reinterpret_cast(position_ids->data) + i * position_ids->ne[0]; - for (int32_t j = 0; j < position_ids->ne[0]; ++j) - cur_row[j] = j + position_offset; + if (config[step].offload) + kv_cache->offload_slot(backend.get(), sched.get(), offload_slot + 1); + if (config[step].slide) + kv_cache->slide_kv_slot(); } - worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); - if (worker->status != GGML_STATUS_SUCCESS) return false; // Phase 3: Copy flow output to speech_feat (persistent buffer) ggml_reset(ctx1.get()); @@ -478,7 +565,7 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f ggml_backend_tensor_alloc(token2wav_buffer.get(), speech_feat, ggml_backend_buffer_get_base(token2wav_buffer.get())); - ggml_backend_tensor_set_async(backend.get(), speech_feat, worker->flow_cache.data(), 0, speech_feat->nb[2]); + ggml_backend_tensor_set_async(backend.get(), speech_feat, worker->flow_cache.data(), 0, speech_feat->nb[2] - feat->nb[2]); ggml_backend_tensor_copy_async(backend.get(), backend.get(), feat, ggml_view_2d(ctx0.get(), speech_feat, speech_feat->ne[0], feat->ne[1], speech_feat->nb[1], speech_feat->nb[1] * cache_length)); } @@ -496,7 +583,7 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f ggml_backend_tensor_get_async(backend.get(), feat, worker->flow_cache.data() + n_cached_elements, 0, feat->nb[2]); } - if (kv_cache) + if (config[flow.decoder.diffusion_steps - 1].cache_kv) if (finalize) kv_cache->cur_len = 0; else diff --git a/src/cosyvoice.cpp b/src/cosyvoice.cpp index bc8922a..6bc78c7 100644 --- a/src/cosyvoice.cpp +++ b/src/cosyvoice.cpp @@ -60,7 +60,7 @@ struct cosyvoice_internal_context : cosyvoice_context struct cosyvoice_context_3 : cosyvoice_internal_context, cosyvoice_tokenizer, cosyvoice_model_3 { - cosyvoice_context_3(const cosyvoice_context_params_v2_cpp& params, ggml_backend_t backend) : + cosyvoice_context_3(const cosyvoice_context_params_v3_cpp& params, ggml_backend_t backend) : cosyvoice_model_3(backend, params) {} }; @@ -126,13 +126,41 @@ cosyvoice_context_t cosyvoice_load_ext(const void* data, size_t size, const cosy gguf_loader loader(parser, data, size); if (!loader) return nullptr; - cosyvoice_context_params_v2_t params_v2{ .base_params = *params }; - if (version == COSYVOICE_CONTEXT_PARAMS_V2_VERSION) - params_v2.n_workers = std::max(1u, reinterpret_cast(params)->n_workers); + cosyvoice_context_params_v3_cpp params_v3 = {}; + params_v3.cosyvoice_context_params_t::operator=(*params); + + if (version >= COSYVOICE_CONTEXT_PARAMS_V2_VERSION) + params_v3.n_workers = std::max(1u, reinterpret_cast(params)->n_workers); + else + params_v3.n_workers = 1; + + if (version >= COSYVOICE_CONTEXT_PARAMS_V3_VERSION) + { + auto* p_v3 = reinterpret_cast(params); + if (params_v3.dit_kv_fixed_slots == 0 && params_v3.dit_kv_offloadable_slots != 0) + { + params_v3.dit_kv_fixed_slots = 1; + --params_v3.dit_kv_offloadable_slots; + } + + params_v3.dit_kv_cache_type = p_v3->dit_kv_cache_type; + params_v3.dit_kv_fixed_slots = p_v3->dit_kv_fixed_slots; + params_v3.dit_kv_offloadable_slots = p_v3->dit_kv_offloadable_slots; + params_v3.dit_allow_kv_cache_fallback = p_v3->dit_allow_kv_cache_fallback; + params_v3.dit_kv_cache_length = p_v3->dit_kv_cache_length; + if (params_v3.dit_kv_cache_length == 0) + params_v3.dit_kv_cache_length = params->n_max_seq * 10; + } else - params_v2.n_workers = 1; + { + params_v3.dit_kv_fixed_slots = 0; + params_v3.dit_kv_offloadable_slots = 0; + params_v3.dit_kv_cache_length = params->n_max_seq * 10; + params_v3.dit_allow_kv_cache_fallback = true; + params_v3.dit_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE(COSYVOICE_KV_CACHE_TYPE_Q8_0, COSYVOICE_KV_CACHE_TYPE_Q4_0, COSYVOICE_KV_CACHE_TYPE_Q8_0); + } - auto ctx = new cosyvoice_context_3(reinterpret_cast(params_v2), + auto ctx = new cosyvoice_context_3(params_v3, backend ? backend : ggml_backend_init_best() ); ctx->cosyvoice_model_3::load(loader); @@ -140,7 +168,7 @@ cosyvoice_context_t cosyvoice_load_ext(const void* data, size_t size, const cosy auto ggml_backend_set_n_threads = reinterpret_cast(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(ggml_backend_get_device(ctx->worker->cpu_backend.get())), "ggml_backend_set_n_threads")); if (n_threads == 0) - n_threads = std::max(1, std::thread::hardware_concurrency() / params_v2.n_workers); + n_threads = std::max(1, std::thread::hardware_concurrency() / params_v3.n_workers); if (n_threads != 0) ggml_backend_set_n_threads(ctx->worker->cpu_backend.get(), n_threads); @@ -175,6 +203,11 @@ cosyvoice_context_t cosyvoice_load_from_file_with_params_v2(const char* filename return cosyvoice_load_from_file_ext(filename, params, nullptr, 0); } +cosyvoice_context_t cosyvoice_load_from_file_with_params_v3(const char* filename, const cosyvoice_context_params_v3_t* params) +{ + return cosyvoice_load_from_file_ext(filename, params, nullptr, 0); +} + cosyvoice_context_t cosyvoice_duplicate_context(cosyvoice_context_t ctx) { if (auto cv3_ctx = dynamic_cast(ctx); cv3_ctx) @@ -392,16 +425,6 @@ uint32_t cosyvoice_get_chunk_tokens(cosyvoice_context_t ctx) return ctx->get_chunk_tokens(); } -uint32_t cosyvoice_get_flow_overlap_tokens(cosyvoice_context_t ctx) -{ - return ctx->get_flow_overlap_tokens(); -} - -uint32_t cosyvoice_get_hift_overlap_tokens(cosyvoice_context_t ctx) -{ - return ctx->get_hift_overlap_tokens(); -} - ggml_status cosyvoice_get_last_status(cosyvoice_context_t ctx) { return ctx->get_last_status(); From 36f5edfa913c7967327b813e39642ffa5857065e Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Mon, 13 Jul 2026 22:25:38 +0800 Subject: [PATCH 07/34] cli: add DiT KV cache configuration and --stream flag for interactive mode --- tools/cli/cosyvoice-cli.cpp | 117 +++++++++++++++++++++++++++++++----- 1 file changed, 101 insertions(+), 16 deletions(-) diff --git a/tools/cli/cosyvoice-cli.cpp b/tools/cli/cosyvoice-cli.cpp index 38ae62b..366da86 100644 --- a/tools/cli/cosyvoice-cli.cpp +++ b/tools/cli/cosyvoice-cli.cpp @@ -39,6 +39,8 @@ #include #endif +constexpr auto parse_kv_cache_type_arg = parse_llm_kv_cache_type_arg; + struct cli_options { uint32_t max_llm_len = COSYVOICE_DEFAULT_LLM_MAX_SEQ_LEN; @@ -77,12 +79,21 @@ struct cli_options seed_policy_mode seed_policy = seed_policy_mode::auto_mode; uint32_t n_threads = 0; bool has_llm_kv_cache_type = false; - cosyvoice_llm_kv_cache_type_t llm_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( - COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0, - COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0, - COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0); + cosyvoice_kv_cache_type_t llm_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( + COSYVOICE_KV_CACHE_TYPE_Q8_0, + COSYVOICE_KV_CACHE_TYPE_Q4_0, + COSYVOICE_KV_CACHE_TYPE_Q8_0); + bool has_dit_kv_cache_type = false; + cosyvoice_kv_cache_type_t dit_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( + COSYVOICE_KV_CACHE_TYPE_Q8_0, + COSYVOICE_KV_CACHE_TYPE_Q4_0, + COSYVOICE_KV_CACHE_TYPE_Q8_0); + uint32_t dit_kv_fixed_slots = 0; + uint32_t dit_kv_offloadable_slots = 0; + uint32_t dit_kv_cache_length = 0; bool has_inference_buffer_policy = false; cosyvoice_inference_buffer_policy_t inference_buffer_policy = COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED; + bool stream = false; bool verbose = false; bool quiet = false; bool has_temperature = false; @@ -305,6 +316,13 @@ static void print_usage(const char* argv0) printf(" Default: k=q8_0,v=q4_0,fallback=q8_0.\n"); printf(" --inference-buffer-policy \n"); printf(" Inference buffer policy (interactive only). Default: balanced.\n"); + printf(" --dit-kv-cache-type ,v=[,fallback=]>\n"); + printf(" DiT KV cache type (interactive only).\n"); + printf(" Default: k=q8_0,v=q4_0,fallback=q8_0.\n"); + printf(" --dit-kv-fixed-slots Number of fixed (non-offloadable) DiT KV slots (interactive only). Default: 0.\n"); + printf(" --dit-kv-offloadable-slots Number of offloadable DiT KV slots (interactive only). Default: 0.\n"); + printf(" --dit-kv-cache-length DiT KV cache max seq length (interactive only). Default: max-llm-len * 10.\n"); + printf(" --stream Enable streaming playback in interactive mode.\n"); printf(" --seed Fixed seed for sampling.\n"); printf(" --seed-policy Seed strategy. Default: auto (fixed if --seed is set).\n"); @@ -851,7 +869,7 @@ static void run_interactive_loop( uint32_t sample_rate) { audio_cache cache; - bool streaming = false; + bool streaming = options.stream; std::string line; if (seed_state && !seed_state->has_next_seed) { @@ -1353,12 +1371,23 @@ static bool validate_options(cli_options& options) if (options.has_llm_kv_cache_type && !COSYVOICE_IS_SEPARATE_KV_CACHE(options.llm_kv_cache_type) - && (options.llm_kv_cache_type < 0 || options.llm_kv_cache_type >= COSYVOICE_LLM_KV_CACHE_TYPE_COUNT)) + && (options.llm_kv_cache_type < 0 || options.llm_kv_cache_type >= COSYVOICE_KV_CACHE_TYPE_COUNT)) { print_error_log("Error: invalid --llm-kv-cache-type. Allowed values: f32, f16, q8_0, q5_1, q5_0, q4_1, q4_0 or k=,v=.\n"); ok = false; } + if (options.has_dit_kv_cache_type + && !COSYVOICE_IS_SEPARATE_KV_CACHE(options.dit_kv_cache_type) + && (options.dit_kv_cache_type < 0 || options.dit_kv_cache_type >= COSYVOICE_KV_CACHE_TYPE_COUNT)) + { + print_error_log("Error: invalid --dit-kv-cache-type. Allowed values: f32, f16, q8_0, q5_1, q5_0, q4_1, q4_0 or k=,v=.\n"); + ok = false; + } + + if (options.dit_kv_cache_length == 0) + options.dit_kv_cache_length = options.max_llm_len * 10; + if (options.has_win_size && options.win_size <= 0) { print_error_log("Error: --win-size must be > 0.\n"); @@ -1483,8 +1512,8 @@ int tool_entry(int argc, char** argv) else if (str_casecmp(arg, "--llm-kv-cache-type") == 0) { auto value = get_arg_value(); - cosyvoice_llm_kv_cache_type_t type; - if (!parse_llm_kv_cache_type_arg(value, &type)) + cosyvoice_kv_cache_type_t type; + if (!parse_kv_cache_type_arg(value, &type)) { print_error_log("Error: invalid --llm-kv-cache-type value \"%s\".\n", value); return 1; @@ -1492,6 +1521,53 @@ int tool_entry(int argc, char** argv) options.llm_kv_cache_type = type; options.has_llm_kv_cache_type = true; } + else if (str_casecmp(arg, "--dit-kv-cache-type") == 0) + { + auto value = get_arg_value(); + cosyvoice_kv_cache_type_t type; + if (!parse_kv_cache_type_arg(value, &type)) + { + print_error_log("Error: invalid --dit-kv-cache-type value \"%s\".\n", value); + return 1; + } + options.dit_kv_cache_type = type; + options.has_dit_kv_cache_type = true; + } + else if (str_casecmp(arg, "--dit-kv-fixed-slots") == 0) + { + auto value = get_arg_value(); + uint32_t v; + if (!parse_uint32_arg(value, &v)) + { + print_error_log("Error: invalid --dit-kv-fixed-slots value \"%s\".\n", value); + return 1; + } + options.dit_kv_fixed_slots = v; + } + else if (str_casecmp(arg, "--dit-kv-offloadable-slots") == 0) + { + auto value = get_arg_value(); + uint32_t v; + if (!parse_uint32_arg(value, &v)) + { + print_error_log("Error: invalid --dit-kv-offloadable-slots value \"%s\".\n", value); + return 1; + } + options.dit_kv_offloadable_slots = v; + } + else if (str_casecmp(arg, "--dit-kv-cache-length") == 0) + { + auto value = get_arg_value(); + uint32_t v; + if (!parse_uint32_arg(value, &v)) + { + print_error_log("Error: invalid --dit-kv-cache-length value \"%s\".\n", value); + return 1; + } + options.dit_kv_cache_length = v; + } + else if (str_casecmp(arg, "--stream") == 0) + options.stream = true; else if (str_casecmp(arg, "--inference-buffer-policy") == 0) { auto value = get_arg_value(); @@ -1828,17 +1904,26 @@ int tool_entry(int argc, char** argv) cosyvoice_init_backend_from_path(options.backend_path.empty() ? nullptr : options.backend_path.c_str()); timing.backend_init_ms = elapsed_ms(stage_start, std::chrono::steady_clock::now()); - cosyvoice_context_params_t params; - cosyvoice_init_default_context_params(¶ms); + cosyvoice_context_params_v3_t params = {}; + cosyvoice_init_default_context_params(¶ms.base_params.base_params); + params.base_params.base_params.n_max_seq = options.max_llm_len; + params.base_params.base_params.llm_kv_cache_type = options.llm_kv_cache_type; + params.base_params.n_workers = 1; if (options.interactive) { if (options.has_inference_buffer_policy) - params.inference_buffer_policy = options.inference_buffer_policy; + params.base_params.base_params.inference_buffer_policy = options.inference_buffer_policy; } else - params.inference_buffer_policy = COSYVOICE_INFERENCE_BUFFER_POLICY_SHARED; - params.n_max_seq = options.max_llm_len; - params.llm_kv_cache_type = options.llm_kv_cache_type; + params.base_params.base_params.inference_buffer_policy = COSYVOICE_INFERENCE_BUFFER_POLICY_SHARED; + if (options.interactive) + { + params.dit_kv_cache_type = options.dit_kv_cache_type; + params.dit_kv_fixed_slots = options.dit_kv_fixed_slots; + params.dit_kv_offloadable_slots = options.dit_kv_offloadable_slots; + params.dit_kv_cache_length = options.dit_kv_cache_length; + params.dit_allow_kv_cache_fallback = true; + } tts_seed_state seed_state; const bool has_seed_value = !options.seed.empty(); const cli_options::seed_policy_mode policy = resolve_seed_policy_mode(options); @@ -1851,14 +1936,14 @@ int tool_entry(int argc, char** argv) print_error_log("Error: invalid --seed value \"%s\". It should be a non-negative integer between 0 and %u.\n", options.seed.c_str(), UINT32_MAX); return 1; } - params.seed = seed_value; + params.base_params.base_params.seed = seed_value; seed_state.next_seed = seed_value; seed_state.has_next_seed = true; } else if (seed_state.fixed) { const uint32_t seed_value = cosyvoice_generate_random_seed(); - params.seed = seed_value; + params.base_params.base_params.seed = seed_value; seed_state.next_seed = seed_value; seed_state.has_next_seed = true; } From 94ebef59b800ebe81d2d2be7490262985bacf0e2 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Tue, 14 Jul 2026 00:07:56 +0800 Subject: [PATCH 08/34] cosyvoice: non-flash attention support and KV cache slot management fixes --- src/cosyvoice-graph.cpp | 4 ++-- src/cosyvoice-kv-cache.cpp | 10 +++------- src/cosyvoice-kv-cache.h | 2 +- src/cosyvoice-llm.cpp | 4 ++-- src/cosyvoice-tts.cpp | 28 ++++++++++++++++++++-------- 5 files changed, 28 insertions(+), 20 deletions(-) diff --git a/src/cosyvoice-graph.cpp b/src/cosyvoice-graph.cpp index df9236b..0ac8442 100644 --- a/src/cosyvoice-graph.cpp +++ b/src/cosyvoice-graph.cpp @@ -377,11 +377,11 @@ ggml_tensor* Attention::build_cgraph(ggml_context* ctx, ggml_tensor* x, ggml_ten attn_output = ggml_flash_attn_ext(ctx, query, key, value, attn_mask, 1.f / std::sqrt(static_cast(head_dim)), 0.f, 0.f); else { - value = ggml_permute(ctx, value, 0, 2, 1, 3); + value = ggml_permute(ctx, value, 1, 0, 2, 3); value = ggml_cont(ctx, value); auto attn_scores = ggml_mul_mat(ctx, key, query); auto attn_weights = ggml_soft_max_ext_inplace(ctx, attn_scores, attn_mask, 1.f / std::sqrt(static_cast(head_dim)), 0.f); - auto attn_output = ggml_mul_mat(ctx, value, attn_weights); + attn_output = ggml_mul_mat(ctx, value, attn_weights); attn_output = ggml_permute(ctx, attn_output, 0, 2, 1, 3); attn_output = ggml_cont(ctx, attn_output); } diff --git a/src/cosyvoice-kv-cache.cpp b/src/cosyvoice-kv-cache.cpp index 1e66906..51e11c8 100644 --- a/src/cosyvoice-kv-cache.cpp +++ b/src/cosyvoice-kv-cache.cpp @@ -376,13 +376,9 @@ void cosyvoice_kv_cache::shift_kv_node_pos(uint32_t shift_pos) } } -bool cosyvoice_kv_cache::can_reuse(bool prefill) const +bool cosyvoice_kv_cache::can_reuse() const { - if (!fattn) return false; - for (auto& layer : std::span(kv_cache_layers + cur_slot_idx * layers, layers)) - if (!layer.k_view || !layer.v_view) - return false; - return true; + return fattn; } bool cosyvoice_kv_cache::bind_slot(int slot_idx) @@ -424,5 +420,5 @@ void cosyvoice_kv_cache::slide_kv_slot() } } else - GGML_ABORT("slide_kv_layers is not supported for non-flash attention"); + GGML_ABORT("slide_kv_slot is not supported for non-flash attention"); } diff --git a/src/cosyvoice-kv-cache.h b/src/cosyvoice-kv-cache.h index 667c318..054f0b8 100644 --- a/src/cosyvoice-kv-cache.h +++ b/src/cosyvoice-kv-cache.h @@ -31,7 +31,7 @@ class cosyvoice_kv_cache ggml_tensor* attention_forward(ggml_context* ctx0, ggml_tensor* query_states, ggml_tensor* key_states, ggml_tensor* value_states, ggml_tensor* attention_mask) const; void shift_kv_node_pos(uint32_t shift_pos); - bool can_reuse(bool prefill) const; + bool can_reuse() const; bool bind_slot(int slot_idx); void slide_kv_slot(); diff --git a/src/cosyvoice-llm.cpp b/src/cosyvoice-llm.cpp index 1a8ba4b..d666e09 100644 --- a/src/cosyvoice-llm.cpp +++ b/src/cosyvoice-llm.cpp @@ -130,7 +130,7 @@ bool cosyvoice_model_3::llm_prefill( auto& llm_probs = worker->llm_probs; if (gf && llm_input && !llm_probs && llm_input->type == type && n_tokens == llm_input->ne[1] - && kv_cache->can_reuse(true)) + && kv_cache->can_reuse()) { ggml_backend_tensor_set_async(worker->backend.get(), llm_input, data, 0, ggml_nbytes(llm_input)); kv_cache->shift_kv_node_pos(n_tokens); @@ -215,7 +215,7 @@ bool cosyvoice_model_3::llm_decode(ggml_type type, const void* data) auto& llm_probs = worker->llm_probs; if (gf && llm_input && llm_probs && llm_input->type == type && 1 == llm_input->ne[1] - && kv_cache->can_reuse(false)) + && kv_cache->can_reuse()) { ggml_backend_tensor_set_async(shared->op_caps.emb_cast_f32 ? worker->backend.get() : worker->cpu_backend.get(), llm_input, data, 0, ggml_nbytes(llm_input)); kv_cache->shift_kv_node_pos(1); diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index 2cd0c23..4ff5a34 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -330,7 +330,7 @@ struct dit_sched_config const auto& operator[](int i) const { return graph_config[i]; } - dit_sched_config(const cosyvoice_context_params_v3_cpp& params, int64_t cut_len, uint32_t offset, bool streaming) + dit_sched_config(const cosyvoice_context_params_v3_cpp& params, int64_t cut_len, uint32_t offset, bool streaming, bool kv_slidable) { if (!streaming) { @@ -365,7 +365,8 @@ struct dit_sched_config || i == 1 || offset != 0 && i == n_no_cache_steps - 1 || i == n_no_cache_steps - || i == n - 1 && cut_len != 0; + || i == n - 1 && cut_len != 0 + || !kv_slidable && i >= n_no_cache_steps; graph_config[i].cache_kv = i >= n_no_cache_steps; graph_config[i].offload = i >= n_no_cache_steps && i < n_no_cache_steps + n_offloadable_steps; graph_config[i].load = graph_config[i].offload && offset != 0; @@ -374,7 +375,7 @@ struct dit_sched_config if (i == n_no_cache_steps - 1 && offset != 0) graph_config[i].cut_len = offset; graph_config[i].slice = offset != 0 && i == n_no_cache_steps; - graph_config[i].slide = graph_config[i].cache_kv && !graph_config[i].offload && i != n - 1; + graph_config[i].slide = kv_slidable && graph_config[i].cache_kv && !graph_config[i].offload && i != n - 1; } } } @@ -432,7 +433,7 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f } } while (false); - dit_sched_config config(params, cut_len, offset ? *offset : 0, streaming); + dit_sched_config config(params, cut_len, offset ? *offset : 0, streaming, kv_cache->can_reuse()); if (offset) if (flow.decoder.diffusion_steps == params.dit_kv_fixed_slots + params.dit_kv_offloadable_slots) *offset += chunk_len; @@ -522,6 +523,13 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f kv_cache->bind_slot(kv_slot++); ggml_reset(ctx0.get()); ggml_backend_sched_reset(sched.get()); + + if (config[step].load) + { + kv_cache->load_slot(backend.get(), sched.get(), offload_slot); + ggml_backend_sched_reset(sched.get()); + } + gf = new_cgraph(ctx0.get()); feat = flow.decoder.build_cgraph_one_step(ctx0.get(), ditctx, step + 1, op_caps, config[step].cut_len, t_leaf, position_ids, gf, config[step].cache_kv ? kv_cache : nullptr, config[step].mask ? &attn_mask : nullptr); ggml_build_forward_expand(gf, feat); @@ -536,17 +544,21 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f auto [t, dt] = flow.decoder.get_t_and_dt(ctx0.get(), step); reinterpret_cast(t_leaf->op_params)[op_caps.fill ? 0 : 1] = t; reinterpret_cast(scale_node->op_params)[0] = dt; - } - if (config[step].load) - kv_cache->load_slot(backend.get(), sched.get(), offload_slot); + if (config[step].load) + kv_cache->load_slot(backend.get(), sched.get(), offload_slot); + } worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); if (worker->status != GGML_STATUS_SUCCESS) return false; if (config[step].offload) - kv_cache->offload_slot(backend.get(), sched.get(), offload_slot + 1); + { + if (step != flow.decoder.diffusion_steps - 1 && config[step + 1].rebuild) + ggml_backend_sched_reset(sched.get()); + kv_cache->offload_slot(backend.get(), sched.get(), offload_slot++); + } if (config[step].slide) kv_cache->slide_kv_slot(); } From c616f36d55a7cbb5bb83b45ca6a3fa18ad530ce1 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Tue, 14 Jul 2026 23:12:38 +0800 Subject: [PATCH 09/34] cosyvoice: accurate causal mask and position_ids for streaming DiT --- src/cosyvoice-model.h | 1 + src/cosyvoice-tts.cpp | 19 +++++++++++++++---- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/src/cosyvoice-model.h b/src/cosyvoice-model.h index 064d655..22e30cb 100644 --- a/src/cosyvoice-model.h +++ b/src/cosyvoice-model.h @@ -72,6 +72,7 @@ struct cosyvoice_worker_context cosyvoice_builtin_sampler_rng_policy_t builtin_sampler_rng_policy; std::vector flow_cache; + std::vector chunk_boundaries; std::unique_ptr batch_buffer; std::unique_ptr nucleus_probs; uint32_t nucleus_probs_capacity; diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index 4ff5a34..66332fc 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -440,7 +440,9 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f else *offset = chunk_len; - int position_offset = kv_cache ? static_cast(kv_cache->cur_len) : 0; + if (streaming && !finalize && offset) + worker->chunk_boundaries.push_back(*offset); + uint32_t noise_len = static_cast(ggml_nelements(ditctx.x)); float* noise_buffer = shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_BEFORE_FLOW, noise_len, nullptr, shared->noise_callback_ctx); ggml_backend_tensor_set_async(backend.get(), ditctx.x, noise_buffer, 0, ggml_nbytes(ditctx.x)); @@ -459,18 +461,24 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f auto post_process = [&](int step) { if (!op_caps.fill) ggml_set_zero(t_leaf); + + const auto base = config[step].cache_kv ? kv_cache->cur_len : 0; for (int64_t i = 0; i < position_ids->ne[1]; ++i) { auto cur_row = reinterpret_cast(position_ids->data) + i * position_ids->ne[0]; for (int32_t j = 0; j < position_ids->ne[0]; ++j) - cur_row[j] = j + position_offset; + cur_row[j] = j + base; } + if (config[step].mask) { - const auto cs = flow.decoder.estimator.static_chunk_size; + const auto& bounds = worker->chunk_boundaries; for (int64_t i = 0; i < attn_mask->ne[0]; ++i) { - int64_t block_end = ((i / cs) + 1) * cs; + auto it = std::upper_bound(bounds.begin(), bounds.end(), static_cast(i)); + int64_t block_end = it != bounds.end() + ? static_cast(*it) + : attn_mask->ne[0]; auto row = reinterpret_cast(attn_mask->data) + i * attn_mask->ne[1]; for (int64_t j = 0; j < attn_mask->ne[1]; ++j) row[j] = j < block_end ? 0 : 0xFC00; @@ -632,7 +640,10 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_AFTER_HIFT, noise_len, noise_buffer, shared->noise_callback_ctx); if (finalize) + { worker->flow_cache.clear(); + worker->chunk_boundaries.clear(); + } if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) { From e7af70de2be5e5812b6a7ada0614486a4e45f607 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Wed, 15 Jul 2026 00:04:41 +0800 Subject: [PATCH 10/34] cosyvoice: make chunk_tokens runtime-configurable via setter API --- include/cosyvoice-interface.h | 1 + include/cosyvoice-lowlevel.h | 5 +++++ src/cosyvoice-loader.cpp | 1 + src/cosyvoice-model.cpp | 15 ++++++++++----- src/cosyvoice-model.h | 4 +++- src/cosyvoice.cpp | 5 +++++ 6 files changed, 25 insertions(+), 6 deletions(-) diff --git a/include/cosyvoice-interface.h b/include/cosyvoice-interface.h index ca570cd..3570d78 100644 --- a/include/cosyvoice-interface.h +++ b/include/cosyvoice-interface.h @@ -97,6 +97,7 @@ struct cosyvoice_model_context ) = 0; ///< Convert speech tokens into waveform samples with additional options. virtual uint32_t get_chunk_tokens() = 0; ///< Get the number of tokens processed in each chunk during streaming inference. + virtual void set_chunk_tokens(uint32_t n_tokens) = 0; ///< Set the number of tokens processed in each chunk during streaming inference. // Status virtual ggml_status get_last_status() = 0; ///< Get the status of the most recent backend operation. diff --git a/include/cosyvoice-lowlevel.h b/include/cosyvoice-lowlevel.h index 4860380..ae4d1e4 100644 --- a/include/cosyvoice-lowlevel.h +++ b/include/cosyvoice-lowlevel.h @@ -355,6 +355,11 @@ COSYVOICE_API bool cosyvoice_tts_stream( */ COSYVOICE_API uint32_t cosyvoice_get_chunk_tokens(cosyvoice_context_t ctx); +/** +* @brief Set the number of tokens processed in each chunk during streaming inference. + */ +COSYVOICE_API void cosyvoice_set_chunk_tokens(cosyvoice_context_t ctx, uint32_t n_tokens); + // ---------------------------------------------------------------------------- // Tokenizer Operations // ---------------------------------------------------------------------------- diff --git a/src/cosyvoice-loader.cpp b/src/cosyvoice-loader.cpp index 0cd14f2..533a455 100644 --- a/src/cosyvoice-loader.cpp +++ b/src/cosyvoice-loader.cpp @@ -909,6 +909,7 @@ void cosyvoice_model_3::load(gguf_loader& loader) for (auto& worker : std::span(workers, shared->params.n_workers)) { + worker.chunk_size = 25 + cv3_shared->flow.pre_lookahead_layer.pre_lookahead_len; worker.nucleus_probs_capacity = static_cast(sampling.top_k * 2); worker.nucleus_probs.reset(new float[worker.nucleus_probs_capacity]); worker.nucleus_probs_len = 0; diff --git a/src/cosyvoice-model.cpp b/src/cosyvoice-model.cpp index 24c73ef..40fd79c 100644 --- a/src/cosyvoice-model.cpp +++ b/src/cosyvoice-model.cpp @@ -183,6 +183,16 @@ cosyvoice_builtin_sampler_rng_policy_t cosyvoice_model::get_builtin_sampler_rng_ return worker->builtin_sampler_rng_policy; } +uint32_t cosyvoice_model::get_chunk_tokens() +{ + return worker->chunk_size; +} + +void cosyvoice_model::set_chunk_tokens(uint32_t n_tokens) +{ + worker->chunk_size = n_tokens; +} + bool cosyvoice_model::set_sampler_seed(uint32_t seed) { worker->sampler_seed = seed; @@ -325,11 +335,6 @@ void cosyvoice_model_3::reset_shared_buffer(ggml_backend_buffer* new_buffer) } } -uint32_t cosyvoice_model_3::get_chunk_tokens() -{ - return 25 + cv3_shared->flow.pre_lookahead_layer.pre_lookahead_len; -} - cosyvoice_3_worker_context::cosyvoice_3_worker_context() : ctx1(ggml_init(ggml_init_params{ .mem_size = ggml_tensor_overhead() * 4, .no_alloc = true })) {} diff --git a/src/cosyvoice-model.h b/src/cosyvoice-model.h index 22e30cb..fc7f359 100644 --- a/src/cosyvoice-model.h +++ b/src/cosyvoice-model.h @@ -71,6 +71,7 @@ struct cosyvoice_worker_context void* sampler_ctx; cosyvoice_builtin_sampler_rng_policy_t builtin_sampler_rng_policy; + uint32_t chunk_size; std::vector flow_cache; std::vector chunk_boundaries; std::unique_ptr batch_buffer; @@ -129,6 +130,8 @@ struct cosyvoice_model : virtual cosyvoice_model_context, virtual cosyvoice_obje bool set_builtin_sampler_rng_policy(cosyvoice_builtin_sampler_rng_policy_t policy); bool set_sampler_seed(uint32_t seed); uint32_t get_sampler_seed(); + uint32_t get_chunk_tokens(); + void set_chunk_tokens(uint32_t n_tokens); void set_noise_callback(cosyvoice_noise_callback_t callback, void* callback_ctx); void get_noise_callback(cosyvoice_noise_callback_t* callback, void** callback_ctx); @@ -186,7 +189,6 @@ struct cosyvoice_model_3 : cosyvoice_model bool llm_job_ext(const int* text, uint32_t text_len, cosyvoice_prompt_t prompt, uint32_t max_new_tokens, bool* final); bool token2wav(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, cosyvoice_generated_speech_ptr result); bool token2wav_ext(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* offset, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result); - uint32_t get_chunk_tokens(); void empty_buffer_cache(); void get_memory_usage(cosyvoice_memory_usage_t* usage); diff --git a/src/cosyvoice.cpp b/src/cosyvoice.cpp index 6bc78c7..a4cb79c 100644 --- a/src/cosyvoice.cpp +++ b/src/cosyvoice.cpp @@ -425,6 +425,11 @@ uint32_t cosyvoice_get_chunk_tokens(cosyvoice_context_t ctx) return ctx->get_chunk_tokens(); } +void cosyvoice_set_chunk_tokens(cosyvoice_context_t ctx, uint32_t n_tokens) +{ + return ctx->set_chunk_tokens(n_tokens); +} + ggml_status cosyvoice_get_last_status(cosyvoice_context_t ctx) { return ctx->get_last_status(); From 1ac1e405c894ae41ff18d81844c2f31febe1bc7e Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Wed, 15 Jul 2026 00:22:55 +0800 Subject: [PATCH 11/34] cosyvoice: correct KV cache slot offload/load state management --- src/cosyvoice-kv-cache.cpp | 6 ++---- src/cosyvoice-kv-cache.h | 2 +- src/cosyvoice-tts.cpp | 2 +- 3 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/cosyvoice-kv-cache.cpp b/src/cosyvoice-kv-cache.cpp index 51e11c8..5a892b9 100644 --- a/src/cosyvoice-kv-cache.cpp +++ b/src/cosyvoice-kv-cache.cpp @@ -206,7 +206,6 @@ void cosyvoice_kv_cache::offload_cache(ggml_backend_t backend, ggml_backend_sche ggml_backend_tensor_get_async(backend, layer.v_tensor, layer.v, 0, v_nbytes); } - cur_len = 0; ggml_backend_sched_synchronize(sched); } @@ -253,7 +252,6 @@ void cosyvoice_kv_cache::load_cache(ggml_backend_t backend, ggml_backend_sched* } } - offloaded_cache->len = 0; ggml_backend_sched_synchronize(sched); } @@ -285,11 +283,11 @@ void cosyvoice_kv_cache::clear_offloaded_cache() } } -void cosyvoice_kv_cache::offload_slot(ggml_backend_t backend, ggml_backend_sched* sched, int offloaded_slot_idx) +void cosyvoice_kv_cache::offload_slot(ggml_backend_t backend, ggml_backend_sched* sched, int offloaded_slot_idx, uint32_t n_tokens) { auto offset = static_cast(get_offloaded_kv_cache_struct_size(layers) * offloaded_slot_idx); offloaded_cache = advance_ptr(offloaded_cache, offset); - offload_cache(backend, sched, cur_len); + offload_cache(backend, sched, n_tokens); offloaded_cache = advance_ptr(offloaded_cache, -offset); } diff --git a/src/cosyvoice-kv-cache.h b/src/cosyvoice-kv-cache.h index 054f0b8..909f6b5 100644 --- a/src/cosyvoice-kv-cache.h +++ b/src/cosyvoice-kv-cache.h @@ -44,7 +44,7 @@ class cosyvoice_kv_cache size_t get_offloaded_cache_size() const; void clear_offloaded_cache(); - void offload_slot(ggml_backend_t backend, ggml_backend_sched* sched, int offloaded_slot_idx); + void offload_slot(ggml_backend_t backend, ggml_backend_sched* sched, int offloaded_slot_idx, uint32_t n_tokens); void load_slot(ggml_backend_t backend, ggml_backend_sched* sched, int offloaded_slot_idx); uint32_t cur_len; diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index 66332fc..35bb2ee 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -565,7 +565,7 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f { if (step != flow.decoder.diffusion_steps - 1 && config[step + 1].rebuild) ggml_backend_sched_reset(sched.get()); - kv_cache->offload_slot(backend.get(), sched.get(), offload_slot++); + kv_cache->offload_slot(backend.get(), sched.get(), offload_slot++, kv_cache->cur_len + static_cast(position_ids->ne[0])); } if (config[step].slide) kv_cache->slide_kv_slot(); From 2269c8c09aa6250e9b499143b8fdc0197ab725b0 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Wed, 15 Jul 2026 14:13:02 +0800 Subject: [PATCH 12/34] cosyvoice: truncate KV cache when exceeding dit_kv_cache_length --- src/cosyvoice-tts.cpp | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index 35bb2ee..e90d6a9 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -409,7 +409,7 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f // Phase 1: Flow encode ggml_cgraph* gf = new_cgraph(ctx0.get()); auto [mu, spks, conds, cut_len] = flow.build_cgraph_encode(ctx0.get(), token, prompt_token, prompt_feat, embedding, op_caps, (flow.decoder.diffusion_steps - params.dit_kv_fixed_slots - params.dit_kv_offloadable_slots) == 0 && offset ? *offset : 0, streaming); - const auto chunk_len = mu->ne[1]; + const auto seq_len = mu->ne[1]; auto ditctx = flow.decoder.prepare_context(ctx1.get(), mu, spks, conds); do { @@ -436,13 +436,20 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f dit_sched_config config(params, cut_len, offset ? *offset : 0, streaming, kv_cache->can_reuse()); if (offset) if (flow.decoder.diffusion_steps == params.dit_kv_fixed_slots + params.dit_kv_offloadable_slots) - *offset += chunk_len; + *offset += seq_len; else - *offset = chunk_len; + *offset = seq_len; - if (streaming && !finalize && offset) + if (streaming && offset) worker->chunk_boundaries.push_back(*offset); + if (offset) + { + const auto chunk_len = *offset - (worker->chunk_boundaries.size() > 1 ? worker->chunk_boundaries.rbegin()[1] : 0); + if (kv_cache->cur_len + chunk_len >= params.dit_kv_cache_length) + kv_cache->cur_len = params.dit_kv_cache_length - chunk_len; + } + uint32_t noise_len = static_cast(ggml_nelements(ditctx.x)); float* noise_buffer = shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_BEFORE_FLOW, noise_len, nullptr, shared->noise_callback_ctx); ggml_backend_tensor_set_async(backend.get(), ditctx.x, noise_buffer, 0, ggml_nbytes(ditctx.x)); From c11761b4ad818546ee6fd823e5eab6bf76bd76c2 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Wed, 15 Jul 2026 14:57:37 +0800 Subject: [PATCH 13/34] common: rename LLM KV cache type functions and types for consistency --- tools/common/tool_common_cosyvoice.h | 48 ++++++++++++++-------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/tools/common/tool_common_cosyvoice.h b/tools/common/tool_common_cosyvoice.h index 69aab68..01c8f42 100644 --- a/tools/common/tool_common_cosyvoice.h +++ b/tools/common/tool_common_cosyvoice.h @@ -28,7 +28,7 @@ using audio_buffer_handle = std::unique_ptr>; #endif -inline bool parse_llm_kv_cache_type_arg(const std::string& value, cosyvoice_llm_kv_cache_type_t* result) +inline bool parse_kv_cache_type_arg(const std::string& value, cosyvoice_kv_cache_type_t* result) { // Check for separate K/V format: "k=,v=" or "k=,v=,fallback=" auto k_pos = value.find("k="); @@ -42,21 +42,21 @@ inline bool parse_llm_kv_cache_type_arg(const std::string& value, cosyvoice_llm_ auto v_end = value.find_first_of(",", v_start); auto v_str = value.substr(v_start, v_end - v_start); - cosyvoice_llm_kv_cache_type_t k_type, v_type; + cosyvoice_kv_cache_type_t k_type, v_type; - if (!parse_llm_kv_cache_type_arg(k_str, &k_type) || - !parse_llm_kv_cache_type_arg(v_str, &v_type)) + if (!parse_kv_cache_type_arg(k_str, &k_type) || + !parse_kv_cache_type_arg(v_str, &v_type)) return false; // Parse optional fallback auto f_pos = value.find("fallback="); - cosyvoice_llm_kv_cache_type_t fallback_type; + cosyvoice_kv_cache_type_t fallback_type; if (f_pos != std::string::npos) { auto f_start = f_pos + 9; auto f_end = value.find_first_of(",", f_start); auto f_str = value.substr(f_start, f_end - f_start); - if (!parse_llm_kv_cache_type_arg(f_str, &fallback_type)) + if (!parse_kv_cache_type_arg(f_str, &fallback_type)) return false; } else @@ -71,19 +71,19 @@ inline bool parse_llm_kv_cache_type_arg(const std::string& value, cosyvoice_llm_ // Unified type. const auto lowered = to_lower(value); if (lowered == "f32") - *result = COSYVOICE_LLM_KV_CACHE_TYPE_F32; + *result = COSYVOICE_KV_CACHE_TYPE_F32; else if (lowered == "f16") - *result = COSYVOICE_LLM_KV_CACHE_TYPE_F16; + *result = COSYVOICE_KV_CACHE_TYPE_F16; else if (lowered == "q8_0") - *result = COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0; + *result = COSYVOICE_KV_CACHE_TYPE_Q8_0; else if (lowered == "q5_1") - *result = COSYVOICE_LLM_KV_CACHE_TYPE_Q5_1; + *result = COSYVOICE_KV_CACHE_TYPE_Q5_1; else if (lowered == "q5_0") - *result = COSYVOICE_LLM_KV_CACHE_TYPE_Q5_0; + *result = COSYVOICE_KV_CACHE_TYPE_Q5_0; else if (lowered == "q4_1") - *result = COSYVOICE_LLM_KV_CACHE_TYPE_Q4_1; + *result = COSYVOICE_KV_CACHE_TYPE_Q4_1; else if (lowered == "q4_0") - *result = COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0; + *result = COSYVOICE_KV_CACHE_TYPE_Q4_0; else return false; return true; @@ -118,33 +118,33 @@ inline const char* inference_buffer_policy_to_string(cosyvoice_inference_buffer_ } } -inline std::string llm_kv_cache_type_to_string(cosyvoice_llm_kv_cache_type_t type) +inline std::string kv_cache_type_to_string(cosyvoice_kv_cache_type_t type) { if (COSYVOICE_IS_SEPARATE_KV_CACHE(type)) { auto fallback = COSYVOICE_KV_CACHE_FALLBACK(type); - auto result = "k=" + std::string(llm_kv_cache_type_to_string(COSYVOICE_K_CACHE_TYPE(type))) - + ",v=" + llm_kv_cache_type_to_string(COSYVOICE_V_CACHE_TYPE(type)); + auto result = "k=" + std::string(kv_cache_type_to_string(COSYVOICE_K_CACHE_TYPE(type))) + + ",v=" + kv_cache_type_to_string(COSYVOICE_V_CACHE_TYPE(type)); if (fallback != COSYVOICE_V_CACHE_TYPE(type)) - result += ",fallback=" + llm_kv_cache_type_to_string(fallback); + result += ",fallback=" + kv_cache_type_to_string(fallback); return result; } switch (type) { - case COSYVOICE_LLM_KV_CACHE_TYPE_F32: + case COSYVOICE_KV_CACHE_TYPE_F32: return "f32"; - case COSYVOICE_LLM_KV_CACHE_TYPE_F16: + case COSYVOICE_KV_CACHE_TYPE_F16: return "f16"; - case COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0: + case COSYVOICE_KV_CACHE_TYPE_Q8_0: return "q8_0"; - case COSYVOICE_LLM_KV_CACHE_TYPE_Q5_1: + case COSYVOICE_KV_CACHE_TYPE_Q5_1: return "q5_1"; - case COSYVOICE_LLM_KV_CACHE_TYPE_Q5_0: + case COSYVOICE_KV_CACHE_TYPE_Q5_0: return "q5_0"; - case COSYVOICE_LLM_KV_CACHE_TYPE_Q4_1: + case COSYVOICE_KV_CACHE_TYPE_Q4_1: return "q4_1"; - case COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0: + case COSYVOICE_KV_CACHE_TYPE_Q4_0: return "q4_0"; default: return "unknown"; From 24e631708dda1ea89d6562fe687567862381dcc0 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Wed, 15 Jul 2026 14:57:50 +0800 Subject: [PATCH 14/34] cli: add chunk_tokens option for streaming playback configuration --- tools/cli/cosyvoice-cli.cpp | 34 ++++++++++++++++++++++++++++++---- 1 file changed, 30 insertions(+), 4 deletions(-) diff --git a/tools/cli/cosyvoice-cli.cpp b/tools/cli/cosyvoice-cli.cpp index 366da86..f5defa9 100644 --- a/tools/cli/cosyvoice-cli.cpp +++ b/tools/cli/cosyvoice-cli.cpp @@ -39,8 +39,6 @@ #include #endif -constexpr auto parse_kv_cache_type_arg = parse_llm_kv_cache_type_arg; - struct cli_options { uint32_t max_llm_len = COSYVOICE_DEFAULT_LLM_MAX_SEQ_LEN; @@ -94,6 +92,8 @@ struct cli_options bool has_inference_buffer_policy = false; cosyvoice_inference_buffer_policy_t inference_buffer_policy = COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED; bool stream = false; + uint32_t chunk_tokens = 0; + bool has_chunk_tokens = false; bool verbose = false; bool quiet = false; bool has_temperature = false; @@ -323,6 +323,7 @@ static void print_usage(const char* argv0) printf(" --dit-kv-offloadable-slots Number of offloadable DiT KV slots (interactive only). Default: 0.\n"); printf(" --dit-kv-cache-length DiT KV cache max seq length (interactive only). Default: max-llm-len * 10.\n"); printf(" --stream Enable streaming playback in interactive mode.\n"); + printf(" --chunk-tokens Tokens per streaming chunk (interactive only). Default: model-defined.\n"); printf(" --seed Fixed seed for sampling.\n"); printf(" --seed-policy Seed strategy. Default: auto (fixed if --seed is set).\n"); @@ -771,11 +772,22 @@ static void print_tts_runtime_info( { char buf[256]; snprintf(buf, sizeof(buf), "requested: %s, actual: %s (%s)", - llm_kv_cache_type_to_string(options.llm_kv_cache_type).c_str(), - llm_kv_cache_type_to_string(context_params.llm_kv_cache_type).c_str(), + kv_cache_type_to_string(options.llm_kv_cache_type).c_str(), + kv_cache_type_to_string(context_params.llm_kv_cache_type).c_str(), options.has_llm_kv_cache_type ? "cli override" : "default"); print_kv_line_string("llm_kv_cache_type", buf); } + if (options.interactive) + { + { + char buf[256]; + snprintf(buf, sizeof(buf), "requested: %s (%s)", + kv_cache_type_to_string(options.dit_kv_cache_type).c_str(), + options.has_dit_kv_cache_type ? "cli override" : "default"); + print_kv_line_string("dit_kv_cache_type", buf); + } + print_kv_line_u32("chunk_tokens", cosyvoice_get_chunk_tokens(ctx)); + } print_kv_line_string("buffer_policy", inference_buffer_policy_to_string(context_params.inference_buffer_policy)); if (log_level == cli_log_level::verbose) { @@ -1568,6 +1580,18 @@ int tool_entry(int argc, char** argv) } else if (str_casecmp(arg, "--stream") == 0) options.stream = true; + else if (str_casecmp(arg, "--chunk-tokens") == 0) + { + auto value = get_arg_value(); + uint32_t v; + if (!parse_uint32_arg(value, &v)) + { + print_error_log("Error: invalid --chunk-tokens value \"%s\".\n", value); + return 1; + } + options.chunk_tokens = v; + options.has_chunk_tokens = true; + } else if (str_casecmp(arg, "--inference-buffer-policy") == 0) { auto value = get_arg_value(); @@ -1991,6 +2015,8 @@ int tool_entry(int argc, char** argv) cosyvoice_tts_context_set_split_text_enabled(tts_ctx.get(), options.split_text_enabled); cosyvoice_tts_context_set_fast_split_text_enabled(tts_ctx.get(), options.fast_split_text_enabled); cosyvoice_tts_context_set_fade_in_enabled(tts_ctx.get(), options.fade_in_enabled); + if (options.interactive && options.has_chunk_tokens) + cosyvoice_set_chunk_tokens(ctx.get(), options.chunk_tokens); cosyvoice_context_params_t effective_params; cosyvoice_get_context_params(ctx.get(), &effective_params); cosyvoice_generation_config_t generation_config; From b44aa9c1e943ecc73309fea13d781f9233cc34eb Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Wed, 15 Jul 2026 23:58:01 +0800 Subject: [PATCH 15/34] server: add streaming TTS with DiT KV cache support --- tools/server/cosyvoice-server-backend.cpp | 242 ++++++++++++++--- tools/server/cosyvoice-server-webui.cpp | 313 ++++++++++++++++------ tools/server/cosyvoice-server.cpp | 113 +++++++- tools/server/cosyvoice-server.h | 48 +++- tools/server/httplib.ixx | 9 +- tools/server/server_common.cpp | 117 ++++++++ tools/server/server_common.h | 64 ++++- tools/server/synthesize_via_api.py | 121 ++++++++- tools/server/webui/cosyvoice-webui.css | 9 + tools/server/webui/cosyvoice-webui.html | 72 +++-- tools/server/webui/cosyvoice-webui.js | 238 ++++++++++++++-- 11 files changed, 1140 insertions(+), 206 deletions(-) diff --git a/tools/server/cosyvoice-server-backend.cpp b/tools/server/cosyvoice-server-backend.cpp index 24c8d8a..5ac33f3 100644 --- a/tools/server/cosyvoice-server-backend.cpp +++ b/tools/server/cosyvoice-server-backend.cpp @@ -1,4 +1,4 @@ -#include "cosyvoice-server.h" +#include "cosyvoice-server.h" #include "server_common.h" #ifndef COSYVOICE_NO_AUDIO @@ -47,9 +47,13 @@ struct speech_request float min_token_text_ratio = 0.0f; bool has_max_token_text_ratio = false; float max_token_text_ratio = 0.0f; + + bool stream = false; + bool has_chunk_tokens = false; + uint32_t chunk_tokens = 0; }; -static void set_openai_error(Response& res, int status, const std::string& message, const std::string& type, const char* param, const char* code) +void set_openai_error(Response& res, int status, const std::string& message, const std::string& type, const char* param, const char* code) { json payload; payload["error"] = { @@ -63,7 +67,7 @@ static void set_openai_error(Response& res, int status, const std::string& messa res.set_content(payload.dump(), "application/json"); } -static bool require_auth(const Request& req, const server_runtime& runtime, Response& res) +bool require_auth(const Request& req, const server_runtime& runtime, Response& res) { if (runtime.api_key.empty()) return true; @@ -299,6 +303,33 @@ static bool parse_speech_request_json(const json& body, speech_request* request, return false; } + if (const auto it = body.find("stream"); it != body.end()) + { + if (!it->is_boolean()) + { + set_openai_error(res, 400, "Field must be a boolean: stream", "invalid_request_error", "stream", "invalid_type"); + return false; + } + request->stream = it->get(); + } + + if (const auto it = body.find("chunk_tokens"); it != body.end()) + { + if (!it->is_number_integer()) + { + set_openai_error(res, 400, "Field must be an integer: chunk_tokens", "invalid_request_error", "chunk_tokens", "invalid_type"); + return false; + } + const auto value = it->get(); + if (value < 0 || static_cast(value) > UINT32_MAX) + { + set_openai_error(res, 400, "Field must be in [0, 4294967295]: chunk_tokens", "invalid_request_error", "chunk_tokens", "invalid_value"); + return false; + } + request->chunk_tokens = static_cast(value); + request->has_chunk_tokens = true; + } + const auto it_instructions = body.find("instructions"); const auto it_instruction = body.find("instruction"); if (it_instructions != body.end()) @@ -349,28 +380,19 @@ static void fill_request_log_context_from_request(request_log_context* ctx, cons ctx->has_instructions = request.has_instructions; ctx->has_seed = request.has_seed; ctx->seed = request.seed; + ctx->stream = request.stream; } -int cosyvoice_server_backend_run(server_runtime& runtime) +void cosyvoice_server_register_api_routes(Server& server, server_runtime& runtime) { - Server server; - server.new_task_queue = [&runtime]() { - return new ThreadPool(runtime.concurrency, runtime.concurrency); - }; - - server.set_exception_handler([&](const Request& req, Response& res, std::exception_ptr ep) - { - (void)ep; - if (req.path.rfind("/v1/", 0) == 0) - set_openai_error(res, 500, "Internal server error.", "server_error", nullptr, "internal_error"); - else - res.status = 500; - }); - - server.Get("/", [](const Request&, Response& res) + // ---- CORS preflight ---- + server.Options(R"(/.*)", [](const Request&, Response& res) { - res.status = 200; - res.set_content("CosyVoice OpenAI Speech API Server is running.\n", "text/plain"); + res.set_header("Access-Control-Allow-Origin", "*"); + res.set_header("Access-Control-Allow-Methods", "GET, POST, DELETE, PUT, OPTIONS"); + res.set_header("Access-Control-Allow-Headers", "Content-Type, Authorization"); + res.set_header("Access-Control-Max-Age", "86400"); + res.status = 204; }); server.Get("/healthz", [&](const Request&, Response& res) @@ -508,8 +530,8 @@ int cosyvoice_server_backend_run(server_runtime& runtime) return; } - const uint32_t slot = get_or_assign_thread_slot(runtime); - if (slot >= runtime.concurrency) + const uint32_t slot = acquire_thread_slot(runtime); + if (slot == UINT32_MAX) { set_openai_error(res, 503, "Server is overloaded.", "server_error", nullptr, "server_overloaded"); log_request_done(runtime.log_level, log_ctx, request_log_status::failed, res.status, 0, res.body.size(), "server_overloaded"); @@ -520,12 +542,12 @@ int cosyvoice_server_backend_run(server_runtime& runtime) auto voice_ctx = get_slot_voice_session(runtime, slot, request.voice); if (!voice_ctx) { + release_thread_slot(runtime, slot); set_openai_error(res, 400, "Unknown voice for current slot.", "invalid_request_error", "voice", "invalid_voice"); log_request_done(runtime.log_level, log_ctx, request_log_status::bad_request, res.status, 0, res.body.size(), "voice_mismatch"); return; } - cosyvoice_generated_speech generated = {}; std::string payload; std::string encoding_error; std::string generation_error; @@ -543,36 +565,170 @@ int cosyvoice_server_backend_run(server_runtime& runtime) if (!apply_generation_overrides(go, runtime, model_ctx, &applied_seed, &generation_error)) { + release_thread_slot(runtime, slot); set_openai_error(res, 400, generation_error, "invalid_request_error", nullptr, "invalid_generation_params"); log_request_done(runtime.log_level, log_ctx, request_log_status::bad_request, res.status, 0, res.body.size(), "generation_override"); return; } - bool ok = false; - if (request.has_instructions) - ok = cosyvoice_tts_instruct(voice_ctx, request.input.c_str(), request.instructions.c_str(), request.speed, &generated); - else - ok = cosyvoice_tts_zero_shot(voice_ctx, request.input.c_str(), request.speed, &generated); + const bool is_streaming = request.stream || runtime.stream; + + // ---- Streaming path ---- + if (is_streaming) + { + // Apply chunk_tokens (from request or server default) + uint32_t ct = request.has_chunk_tokens ? request.chunk_tokens : + (runtime.has_chunk_tokens ? runtime.chunk_tokens : 0); + if (ct > 0) + cosyvoice_set_chunk_tokens(model_ctx, ct); + + // Build WAV header ahead of time (static header, sent before PCM data) + std::string wav_header; + if (format == response_audio_format::wav) + build_wav_header(&wav_header, runtime.sample_rate); + + auto content_type = response_format_to_content_type(format); + bool has_wav = (format == response_audio_format::wav); + + // Capture copies of resources needed inside the provider lambda. + // IMPORTANT: the provider runs AFTER the route handler returns, so + // any value captured by reference would be a dangling reference. + // Capture by value or via raw pointers that outlive the handler. + auto& rt = runtime; + auto* vctx = voice_ctx; + auto log_ctx_copy = log_ctx; + auto req_copy = request; // value copy for provider use + + res.set_chunked_content_provider(content_type, + [&rt, slot, req_copy, vctx, format, wav_header, has_wav, log_ctx_copy, applied_seed] + (size_t /*offset*/, DataSink& sink) -> bool + { + // Write WAV header first (for WAV output format) + if (has_wav && !wav_header.empty()) + { + if (!sink.write(wav_header.data(), wav_header.size())) + { + release_thread_slot(rt, slot); + sink.done(); + return true; + } + } - if (!ok || !generated.data || generated.length == 0) - { - set_openai_error(res, 500, "TTS generation failed.", "server_error", nullptr, "generation_failed"); - log_request_done(runtime.log_level, log_ctx, request_log_status::failed, res.status, applied_seed, res.body.size(), "generation_failed"); + std::atomic aborted{false}; + std::string stream_error; + + streaming_callback_context cb_ctx; + cb_ctx.sink = &sink; + cb_ctx.encoder = nullptr; +#ifndef COSYVOICE_NO_AUDIO + cb_ctx.encoder = rt.audio_encoder.get(); +#endif + cb_ctx.format = format; + cb_ctx.sample_rate = rt.sample_rate; + cb_ctx.wav_header_written = has_wav; + cb_ctx.error_out = &stream_error; + cb_ctx.aborted = &aborted; + + // Streaming callback: encode PCM chunk and write to HTTP sink + auto stream_cb = [](const float* audio, uint32_t n, void* ud) -> bool + { + return stream_audio_chunk(static_cast(ud), audio, n); + }; + + bool ok = false; + if (req_copy.has_instructions) + ok = cosyvoice_tts_instruct_stream(vctx, req_copy.input.c_str(), + req_copy.instructions.c_str(), req_copy.speed, stream_cb, &cb_ctx); + else + ok = cosyvoice_tts_zero_shot_stream(vctx, req_copy.input.c_str(), + req_copy.speed, stream_cb, &cb_ctx); + + release_thread_slot(rt, slot); + sink.done(); + + if (!ok && !aborted) + { + log_request_done(rt.log_level, log_ctx_copy, + request_log_status::failed, 200, applied_seed, 0, "stream_generation_failed"); + } + else + { + log_request_done(rt.log_level, log_ctx_copy, + request_log_status::ok, 200, applied_seed, 0, "stream_speech"); + } + + return true; + }, + [](bool /*success*/) { /* resource releaser — no-op */ }); + + log_request_details(runtime.log_level, log_ctx); return; } - if (!build_audio_payload(format, generated, runtime, &payload, &encoding_error)) + // ---- Non-streaming (blocking) path ---- { - set_openai_error(res, 500, encoding_error, "server_error", nullptr, "audio_encode_failed"); - log_request_done(runtime.log_level, log_ctx, request_log_status::failed, res.status, applied_seed, res.body.size(), "audio_encode_failed"); - return; + cosyvoice_generated_speech generated = {}; + + bool ok = false; + if (request.has_instructions) + ok = cosyvoice_tts_instruct(voice_ctx, request.input.c_str(), request.instructions.c_str(), request.speed, &generated); + else + ok = cosyvoice_tts_zero_shot(voice_ctx, request.input.c_str(), request.speed, &generated); + + if (!ok || !generated.data || generated.length == 0) + { + release_thread_slot(runtime, slot); + set_openai_error(res, 500, "TTS generation failed.", "server_error", nullptr, "generation_failed"); + log_request_done(runtime.log_level, log_ctx, request_log_status::failed, res.status, applied_seed, res.body.size(), "generation_failed"); + return; + } + + if (!build_audio_payload(format, generated, runtime, &payload, &encoding_error)) + { + release_thread_slot(runtime, slot); + set_openai_error(res, 500, encoding_error, "server_error", nullptr, "audio_encode_failed"); + log_request_done(runtime.log_level, log_ctx, request_log_status::failed, res.status, applied_seed, res.body.size(), "audio_encode_failed"); + return; + } + + release_thread_slot(runtime, slot); + res.status = 200; + res.set_content(std::move(payload), response_format_to_content_type(format)); + log_request_done(runtime.log_level, log_ctx, request_log_status::ok, res.status, applied_seed, res.body.size(), "speech"); } + }); +} + +int cosyvoice_server_backend_run(server_runtime& runtime) +{ + Server server; + server.new_task_queue = [&runtime]() { + return new ThreadPool(runtime.concurrency, runtime.concurrency); + }; + + server.set_exception_handler([&](const Request& req, Response& res, std::exception_ptr ep) + { + (void)ep; + if (req.path.rfind("/v1/", 0) == 0) + set_openai_error(res, 500, "Internal server error.", "server_error", nullptr, "internal_error"); + else + res.status = 500; + }); + + server.set_pre_routing_handler([](const Request&, Response& res) -> HandlerResponse + { + res.set_header("Access-Control-Allow-Origin", "*"); + return HandlerResponse::Unhandled; + }); + server.Get("/", [](const Request&, Response& res) + { res.status = 200; - res.set_content(std::move(payload), response_format_to_content_type(format)); - log_request_done(runtime.log_level, log_ctx, request_log_status::ok, res.status, applied_seed, res.body.size(), "speech"); + res.set_content("CosyVoice OpenAI Speech API Server is running.\n", "text/plain"); }); + cosyvoice_server_register_api_routes(server, runtime); + server.set_error_handler([&](const Request& req, Response& res) { if (!res.body.empty()) @@ -599,8 +755,8 @@ int cosyvoice_server_backend_run(server_runtime& runtime) { char kv_buf[256]; snprintf(kv_buf, sizeof(kv_buf), "requested: %s, actual: %s (%s)", - llm_kv_cache_type_to_string(runtime.requested_llm_kv_cache_type).c_str(), - llm_kv_cache_type_to_string(runtime.actual_llm_kv_cache_type).c_str(), + kv_cache_type_to_string(runtime.requested_llm_kv_cache_type).c_str(), + kv_cache_type_to_string(runtime.actual_llm_kv_cache_type).c_str(), runtime.has_llm_kv_cache_override ? "user override" : "default"); print_info_log(runtime.log_level, " llm_kv_cache_type : %s\n", kv_buf); } @@ -616,6 +772,9 @@ int cosyvoice_server_backend_run(server_runtime& runtime) print_info_log(runtime.log_level, " text_splitting : %s\n", runtime.split_text_enabled ? "enabled" : "disabled"); print_info_log(runtime.log_level, " fast_split : %s\n", runtime.fast_split_text_enabled ? "enabled" : "disabled"); print_info_log(runtime.log_level, " fade_in : %s\n", runtime.fade_in_enabled ? "enabled" : "disabled"); + print_info_log(runtime.log_level, " stream : %s\n", runtime.stream ? "enabled" : "disabled"); + if (runtime.has_chunk_tokens) + print_info_log(runtime.log_level, " chunk_tokens : %u\n", runtime.chunk_tokens); #ifndef COSYVOICE_NO_AUDIO print_info_log(runtime.log_level, " audio_encoder : %s\n", runtime.audio_encoder ? "available" : "unavailable"); #else @@ -633,4 +792,3 @@ int cosyvoice_server_backend_run(server_runtime& runtime) return 0; } - diff --git a/tools/server/cosyvoice-server-webui.cpp b/tools/server/cosyvoice-server-webui.cpp index 57423e6..255c0a4 100644 --- a/tools/server/cosyvoice-server-webui.cpp +++ b/tools/server/cosyvoice-server-webui.cpp @@ -1,4 +1,4 @@ -#include "cosyvoice-server.h" +#include "cosyvoice-server.h" #include "server_common.h" #include "resource.h" #include "tool_common_cosyvoice.h" @@ -6,6 +6,7 @@ #include +#include #include #include #include @@ -186,7 +187,7 @@ static bool register_speaker_from_audio( #endif // !COSYVOICE_NO_AUDIO // --------------------------------------------------------------------------- -// Route handlers — WebUI mode +// Route handlers - WebUI mode // --------------------------------------------------------------------------- int cosyvoice_server_webui_run(server_runtime& runtime) @@ -203,10 +204,10 @@ int cosyvoice_server_webui_run(server_runtime& runtime) res.set_content(json, "application/json"); }); - // ---- Pre-routing auth handler (checks cookie for every request) ---- + // ---- Pre-routing auth handler (checks cookie for WebUI, Bearer for API) ---- server.set_pre_routing_handler([&runtime](const Request& req, Response& res) -> HandlerResponse { - // No API key configured → no auth needed + // No API key configured - no auth needed if (runtime.api_key.empty()) return HandlerResponse::Unhandled; @@ -218,7 +219,15 @@ int cosyvoice_server_webui_run(server_runtime& runtime) if (req.method == "GET" && req.path == "/ping") return HandlerResponse::Unhandled; - // Check auth cookie + // For OpenAI-compatible API paths (/v1/*), use Bearer token auth + if (req.path.rfind("/v1/", 0) == 0 || req.path == "/healthz") + { + if (require_auth(req, runtime, res)) + return HandlerResponse::Unhandled; + return HandlerResponse::Handled; + } + + // Check auth cookie (WebUI) const std::string token = get_cookie_value(req, "cosyvoice_auth_token"); if (!token.empty() && token == runtime.api_key) return HandlerResponse::Unhandled; @@ -237,19 +246,19 @@ int cosyvoice_server_webui_run(server_runtime& runtime) else { res.status = 200; - res.set_content("CosyVoice — Login page resource not found.", "text/plain"); + res.set_content("CosyVoice - Login page resource not found.", "text/plain"); } return HandlerResponse::Handled; } - // All other routes → 401 + // All other routes - 401 res.status = 401; nlohmann::json err = {{"error", "Authentication required."}}; res.set_content(err.dump(), "application/json"); return HandlerResponse::Handled; }); - // ---- POST /api/auth/login — authenticate and set session cookie ---- + // ---- POST /api/auth/login - authenticate and set session cookie ---- server.Post("/api/auth/login", [&runtime](const Request& req, Response& res) { nlohmann::json body; @@ -282,7 +291,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) "Login successful from %s", req.remote_addr.c_str()); }); - // ---- GET / — serve the WebUI HTML ---- + // ---- GET / - serve the WebUI HTML ---- server.Get("/", [&runtime](const Request&, Response& res) { res.status = 200; @@ -317,7 +326,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) res.set_content(std::move(html), "text/html"); }); - // ---- GET /cosyvoice-webui.css — serve embedded CSS ---- + // ---- GET /cosyvoice-webui.css - serve embedded CSS ---- server.Get("/cosyvoice-webui.css", [](const Request&, Response& res) { size_t size = 0; @@ -332,7 +341,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) res.set_content(std::string(static_cast(data), size), "text/css"); }); - // ---- GET /cosyvoice-webui.js — serve embedded JavaScript ---- + // ---- GET /cosyvoice-webui.js - serve embedded JavaScript ---- server.Get("/cosyvoice-webui.js", [](const Request&, Response& res) { size_t size = 0; @@ -353,7 +362,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) res.set_content("pong", "text/plain"); }); - // ---- GET /backends — list available GGML backends ---- + // ---- GET /backends - list available GGML backends ---- server.Get("/backends", [](const Request&, Response& res) { nlohmann::json list = nlohmann::json::array(); @@ -382,7 +391,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) res.set_content(list.dump(), "application/json"); }); - // ---- GET /formats — list supported audio output formats ---- + // ---- GET /formats - list supported audio output formats ---- server.Get("/formats", [](const Request&, Response& res) { nlohmann::json list = nlohmann::json::array(); @@ -405,10 +414,10 @@ int cosyvoice_server_webui_run(server_runtime& runtime) res.set_content(list.dump(), "application/json"); }); - // ---- GET /status — server status ---- + // ---- GET /status - server status ---- server.Get("/status", [&runtime](const Request&, Response& res) { - // Use manual JSON building — avoids an issue with nlohmann::json in this handler + // Use manual JSON building - avoids an issue with nlohmann::json in this handler std::string json = "{\"status\":\"ok\""; json += ",\"model_loaded\":"; json += runtime.model_slots.empty() ? "false" : "true"; @@ -423,8 +432,8 @@ int cosyvoice_server_webui_run(server_runtime& runtime) cosyvoice_context_params_t actual_params; cosyvoice_get_context_params(runtime.model_slots.front().get(), &actual_params); json += ",\"max_llm_len\":" + std::to_string(actual_params.n_max_seq); - json += ",\"k_cache_type\":\"" + llm_kv_cache_type_to_string(actual_params.llm_k_cache_type) + "\""; - json += ",\"v_cache_type\":\"" + llm_kv_cache_type_to_string(actual_params.llm_v_cache_type) + "\""; + json += ",\"k_cache_type\":\"" + kv_cache_type_to_string(actual_params.llm_k_cache_type) + "\""; + json += ",\"v_cache_type\":\"" + kv_cache_type_to_string(actual_params.llm_v_cache_type) + "\""; json += ",\"buffer_policy\":\"" + std::string(inference_buffer_policy_to_string(actual_params.inference_buffer_policy)) + "\""; auto arch = cosyvoice_get_architecture(runtime.model_slots.front().get()); if (arch && *arch) @@ -448,7 +457,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) res.set_content(std::move(json), "application/json"); }); - // ---- GET /speaker — list registered speakers ---- + // ---- GET /speaker - list registered speakers ---- server.Get("/speaker", [&runtime](const Request&, Response& res) { nlohmann::json payload = runtime.voice_names; @@ -456,7 +465,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) res.set_content(payload.dump(), "application/json"); }); - // ---- POST /speaker — register a new speaker ---- + // ---- POST /speaker - register a new speaker ---- server.Post("/speaker", [&runtime](const Request& req, Response& res) { // Detect multipart vs JSON @@ -584,7 +593,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) "Speaker registered: %s (type=%s, total=%zu)", name.c_str(), type.c_str(), runtime.voice_names.size()); }); - // ---- DELETE /speaker/ — remove a speaker ---- + // ---- DELETE /speaker/ - remove a speaker ---- server.Delete(R"(/speaker/(.*))", [&runtime](const Request& req, Response& res) { std::string name = req.matches[1]; @@ -631,7 +640,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) "Speaker removed: %s (remaining=%zu)", name.c_str(), runtime.voice_names.size()); }); - // ---- POST /speaker/save — save speaker prompt speech to a server-side path ---- + // ---- POST /speaker/save - save speaker prompt speech to a server-side path ---- server.Post("/speaker/save", [&runtime](const Request& req, Response& res) { nlohmann::json body; @@ -760,12 +769,23 @@ int cosyvoice_server_webui_run(server_runtime& runtime) log_ctx.has_seed = body.contains("seed"); if (log_ctx.has_seed) log_ctx.seed = body["seed"].get(); + log_ctx.stream = body.value("stream", false); log_request_details(runtime.log_level, log_ctx); // Speed float speed = body.value("speed", 1.0f); if (speed <= 0.0f) speed = 1.0f; + // Chunk tokens (applied before streaming/blocking either way) + if (body.contains("chunk_tokens")) + { + uint32_t ct = body["chunk_tokens"].get(); + if (ct > 0) + cosyvoice_set_chunk_tokens(runtime.model_slots[0].get(), ct); + } + else if (runtime.has_chunk_tokens) + cosyvoice_set_chunk_tokens(runtime.model_slots[0].get(), runtime.chunk_tokens); + // Apply TTS context flags (optional overrides per request) if (body.contains("fade_in")) cosyvoice_tts_context_set_fade_in_enabled(tts_ctx, body["fade_in"].get()); @@ -800,62 +820,146 @@ int cosyvoice_server_webui_run(server_runtime& runtime) } } - // Generate speech - cosyvoice_generated_speech generated = {}; - bool ok = false; + // ---- Determine streaming vs blocking ---- + const bool stream = body.value("stream", false); + const bool is_streaming = stream || runtime.stream; + + // Determine output format + const std::string format_str = body.value("format", "wav"); + auto fmt = parse_response_format(format_str); + if (fmt == response_audio_format::unknown) + fmt = response_audio_format::wav; const std::string mode = body.value("mode", "auto"); const std::string instructions = body.value("instructions", body.value("instruction", "")); - if (mode == "cross_lingual") - { - ok = cosyvoice_tts_cross_lingual(tts_ctx, text.c_str(), speed, &generated); - } - else if (mode == "instruct" || (mode == "auto" && !instructions.empty())) - { - ok = cosyvoice_tts_instruct(tts_ctx, text.c_str(), instructions.c_str(), speed, &generated); - } - else // "zero_shot" or "auto" + // ---- Streaming path ---- + if (is_streaming) { - ok = cosyvoice_tts_zero_shot(tts_ctx, text.c_str(), speed, &generated); - } + // Build WAV header ahead of time + std::string wav_header; + if (fmt == response_audio_format::wav) + build_wav_header(&wav_header, runtime.sample_rate); + + auto content_type = response_format_to_content_type(fmt); + bool has_wav = (fmt == response_audio_format::wav); + // Copy shared state needed inside the provider. + // The provider runs AFTER the route handler returns, so any + // reference to local variables is dangling — capture by value only. + auto& rt = runtime; + auto* vctx = tts_ctx; + auto log_ctx_copy = log_ctx; + auto text_copy = text; // value copy for provider + auto instr_copy = instructions; // value copy for provider + + res.set_chunked_content_provider(content_type, + [&rt, vctx, text_copy, instr_copy, speed, mode, fmt, wav_header, has_wav, log_ctx_copy, applied_seed] + (size_t /*offset*/, DataSink& sink) -> bool + { + // Write WAV header first + if (has_wav && !wav_header.empty()) + { + if (!sink.write(wav_header.data(), wav_header.size())) + { + sink.done(); + return true; + } + } - if (!ok || !generated.data || generated.length == 0) - { - res.status = 500; - nlohmann::json err = {{"error", "TTS generation failed"}}; - res.set_content(err.dump(), "application/json"); - log_request_done(runtime.log_level, log_ctx, request_log_status::failed, res.status, applied_seed, res.body.size(), "generation_failed"); - return; - } + std::atomic aborted{false}; + std::string stream_error; - // Encode output using shared utilities - std::string audio_payload; - const std::string format_str = body.value("format", "wav"); - auto fmt = parse_response_format(format_str); - if (fmt == response_audio_format::unknown) - fmt = response_audio_format::wav; + streaming_callback_context cb_ctx; + cb_ctx.sink = &sink; + cb_ctx.encoder = nullptr; +#ifndef COSYVOICE_NO_AUDIO + cb_ctx.encoder = rt.audio_encoder.get(); +#endif + cb_ctx.format = fmt; + cb_ctx.sample_rate = rt.sample_rate; + cb_ctx.wav_header_written = has_wav; + cb_ctx.error_out = &stream_error; + cb_ctx.aborted = &aborted; + + auto stream_cb = [](const float* audio, uint32_t n, void* ud) -> bool + { + return stream_audio_chunk(static_cast(ud), audio, n); + }; + + bool ok = false; + if (mode == "cross_lingual") + ok = cosyvoice_tts_cross_lingual_stream(vctx, text_copy.c_str(), speed, stream_cb, &cb_ctx); + else if (mode == "instruct" || (mode == "auto" && !instr_copy.empty())) + ok = cosyvoice_tts_instruct_stream(vctx, text_copy.c_str(), instr_copy.c_str(), speed, stream_cb, &cb_ctx); + else + ok = cosyvoice_tts_zero_shot_stream(vctx, text_copy.c_str(), speed, stream_cb, &cb_ctx); + + sink.done(); + + if (!ok && !aborted) + log_request_done(rt.log_level, log_ctx_copy, + request_log_status::failed, 200, applied_seed, 0, "stream_tts_failed"); + else + log_request_done(rt.log_level, log_ctx_copy, + request_log_status::ok, 200, applied_seed, 0, "stream_tts"); + + return true; + }, + [](bool /*success*/) { /* resource releaser - no-op */ }); - std::string encoding_error; - if (!build_audio_payload(fmt, generated, runtime, &audio_payload, &encoding_error)) - { - res.status = 500; - nlohmann::json err = {{"error", "Audio encoding failed: " + encoding_error}}; - res.set_content(err.dump(), "application/json"); - log_request_done(runtime.log_level, log_ctx, request_log_status::failed, res.status, applied_seed, res.body.size(), "audio_encode_failed"); return; } - log_message(runtime.log_level, server_log_level::concise, "WEBUI", - "TTS generated: voice=%s, text_len=%zu, format=%s, samples=%u", - voice.c_str(), text.size(), format_str.c_str(), generated.length); + // ---- Non-streaming (blocking) path ---- + { + cosyvoice_generated_speech generated = {}; + bool ok = false; - res.status = 200; - res.set_content(std::move(audio_payload), response_format_to_content_type(fmt)); - log_request_done(runtime.log_level, log_ctx, request_log_status::ok, res.status, applied_seed, audio_payload.size(), "tts"); + if (mode == "cross_lingual") + { + ok = cosyvoice_tts_cross_lingual(tts_ctx, text.c_str(), speed, &generated); + } + else if (mode == "instruct" || (mode == "auto" && !instructions.empty())) + { + ok = cosyvoice_tts_instruct(tts_ctx, text.c_str(), instructions.c_str(), speed, &generated); + } + else // "zero_shot" or "auto" + { + ok = cosyvoice_tts_zero_shot(tts_ctx, text.c_str(), speed, &generated); + } + + if (!ok || !generated.data || generated.length == 0) + { + res.status = 500; + nlohmann::json err = {{"error", "TTS generation failed"}}; + res.set_content(err.dump(), "application/json"); + log_request_done(runtime.log_level, log_ctx, request_log_status::failed, res.status, applied_seed, res.body.size(), "generation_failed"); + return; + } + + // Encode output using shared utilities + std::string audio_payload; + std::string encoding_error; + if (!build_audio_payload(fmt, generated, runtime, &audio_payload, &encoding_error)) + { + res.status = 500; + nlohmann::json err = {{"error", "Audio encoding failed: " + encoding_error}}; + res.set_content(err.dump(), "application/json"); + log_request_done(runtime.log_level, log_ctx, request_log_status::failed, res.status, applied_seed, res.body.size(), "audio_encode_failed"); + return; + } + + log_message(runtime.log_level, server_log_level::concise, "WEBUI", + "TTS generated: voice=%s, text_len=%zu, format=%s, samples=%u", + voice.c_str(), text.size(), format_str.c_str(), generated.length); + + res.status = 200; + res.set_content(std::move(audio_payload), response_format_to_content_type(fmt)); + log_request_done(runtime.log_level, log_ctx, request_log_status::ok, res.status, applied_seed, audio_payload.size(), "tts"); + } }); - // ---- GET /frontend/model — return frontend model paths ---- + // ---- GET /frontend/model - return frontend model paths ---- server.Get("/frontend/model", [&runtime](const Request&, Response& res) { nlohmann::json payload; @@ -865,7 +969,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) res.set_content(payload.dump(), "application/json"); }); - // ---- PUT /frontend/model — update frontend model paths ---- + // ---- PUT /frontend/model - update frontend model paths ---- server.Put("/frontend/model", [&runtime](const Request& req, Response& res) { nlohmann::json body; @@ -888,7 +992,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) res.set_content(ok.dump(), "application/json"); }); - // ---- POST /frontend/model/load — load ONNX frontend models into memory ---- + // ---- POST /frontend/model/load - load ONNX frontend models into memory ---- #if !defined(COSYVOICE_NO_FRONTEND) server.Post("/frontend/model/load", [&runtime](const Request& req, Response& res) { @@ -943,7 +1047,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) }); #endif - // ---- POST /frontend/model/unload — unload ONNX frontend models ---- + // ---- POST /frontend/model/unload - unload ONNX frontend models ---- #if !defined(COSYVOICE_NO_FRONTEND) server.Post("/frontend/model/unload", [&runtime](const Request&, Response& res) { @@ -972,7 +1076,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) }); #endif - // ---- POST /model/load — dynamically load a GGUF model ---- + // ---- POST /model/load - dynamically load a GGUF model ---- server.Post("/model/load", [&runtime](const Request& req, Response& res) { request_log_context log_ctx = make_request_log_context(req, "/model/load"); @@ -1050,7 +1154,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) } // Build context params with defaults - cosyvoice_context_params_v2_cpp context_params; + cosyvoice_context_params_v3_cpp context_params; cosyvoice_init_default_context_params(&context_params); if (runtime.has_seed) @@ -1060,10 +1164,31 @@ int cosyvoice_server_webui_run(server_runtime& runtime) // Apply optional advanced config from request if (body.contains("llm_kv_cache_type")) { - cosyvoice_llm_kv_cache_type_t kv_type; - if (parse_llm_kv_cache_type_arg(body["llm_kv_cache_type"].get(), &kv_type)) + cosyvoice_kv_cache_type_t kv_type; + if (parse_kv_cache_type_arg(body["llm_kv_cache_type"].get(), &kv_type)) context_params.llm_kv_cache_type = kv_type; } + if (body.contains("dit_kv_cache_type")) + { + cosyvoice_kv_cache_type_t kv_type; + if (parse_kv_cache_type_arg(body["dit_kv_cache_type"].get(), &kv_type)) + context_params.dit_kv_cache_type = kv_type; + } + if (body.contains("dit_kv_fixed_slots")) + { + uint32_t v = body["dit_kv_fixed_slots"].get(); + context_params.dit_kv_fixed_slots = v; + } + if (body.contains("dit_kv_offloadable_slots")) + { + uint32_t v = body["dit_kv_offloadable_slots"].get(); + context_params.dit_kv_offloadable_slots = v; + } + if (body.contains("dit_kv_cache_length")) + { + uint32_t v = body["dit_kv_cache_length"].get(); + context_params.dit_kv_cache_length = v; + } if (body.contains("inference_buffer_policy")) { cosyvoice_inference_buffer_policy_t policy; @@ -1076,7 +1201,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) if (v > 0) context_params.n_max_seq = v; } - // Load the model — check result FIRST, then store + // Load the model - check result FIRST, then store auto loaded_ctx = cosyvoice_load_from_file_ext( model_path.c_str(), &context_params, @@ -1104,7 +1229,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) cosyvoice_get_architecture(loaded_ctx) ? cosyvoice_get_architecture(loaded_ctx) : "?", backend_type.c_str(), n_threads ? n_threads : (uint32_t)0, - llm_kv_cache_type_to_string(p.llm_kv_cache_type).c_str(), + kv_cache_type_to_string(p.llm_kv_cache_type).c_str(), inference_buffer_policy_to_string(p.inference_buffer_policy), p.n_max_seq); } @@ -1127,6 +1252,16 @@ int cosyvoice_server_webui_run(server_runtime& runtime) // Sample rate runtime.sample_rate = cosyvoice_get_sample_rate(runtime.model_slots.front().get()); + // Save effective DiT KV cache params (v3-specific, not available via get_context_params) + { + // context_params is a cosyvoice_context_params_v3_cpp from the outer scope + // (defined above when building the load request) + const auto& cp = context_params; + runtime.dit_kv_fixed_slots = cp.dit_kv_fixed_slots; + runtime.dit_kv_offloadable_slots = cp.dit_kv_offloadable_slots; + runtime.dit_kv_cache_length = cp.dit_kv_cache_length; + } + // Get default generation config cosyvoice_get_default_generation_config(runtime.model_slots.front().get(), &runtime.default_generation_config); for (auto& slot : runtime.model_slots) @@ -1152,7 +1287,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) log_request_done(runtime.log_level, log_ctx, request_log_status::ok, res.status, 0, res.body.size(), "model_loaded"); }); - // ---- POST /model/unload — unload the current model ---- + // ---- POST /model/unload - unload the current model ---- server.Post("/model/unload", [&runtime](const Request& req, Response& res) { request_log_context log_ctx = make_request_log_context(req, "/model/unload"); @@ -1167,7 +1302,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) return; } - // Order matters: TTS sessions → voices → model context + // Order matters: TTS sessions -> voices -> model context runtime.voice_sessions.clear(); runtime.voices.clear(); runtime.voice_names.clear(); @@ -1194,7 +1329,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) log_request_done(runtime.log_level, log_ctx, request_log_status::ok, res.status, 0, res.body.size(), "model_unloaded"); }); - // ---- GET /model/defaults — return defaults for all configurable parameters ---- + // ---- GET /model/defaults - return defaults for all configurable parameters ---- // Works even without a loaded model (sensible fallbacks for generation params). server.Get("/model/defaults", [&runtime](const Request&, Response& res) { @@ -1208,14 +1343,29 @@ int cosyvoice_server_webui_run(server_runtime& runtime) d["default_backend"] = "auto"; d["default_n_threads"] = 0; + // Default DiT KV types (overridden below when model is loaded) + d["default_dit_k_cache_type"] = "q8_0"; + d["default_dit_v_cache_type"] = "q4_0"; + d["default_dit_kv_fixed_slots"] = 0; + d["default_dit_kv_offloadable_slots"] = 0; + d["default_dit_kv_cache_length"] = 0; + // Generation defaults (model-dependent or sensible fallbacks) if (!runtime.model_slots.empty()) { + auto ctx = runtime.model_slots.front().get(); + d["temperature"] = runtime.default_generation_config.temperature; d["top_k"] = runtime.default_generation_config.sampling.top_k; d["top_p"] = runtime.default_generation_config.sampling.top_p; d["win_size"] = runtime.default_generation_config.sampling.win_size; d["tau_r"] = runtime.default_generation_config.sampling.tau_r; + d["chunk_tokens"] = cosyvoice_get_chunk_tokens(ctx); + + // Use effective DiT params saved in runtime after model load + d["default_dit_kv_fixed_slots"] = runtime.dit_kv_fixed_slots; + d["default_dit_kv_offloadable_slots"] = runtime.dit_kv_offloadable_slots; + d["default_dit_kv_cache_length"] = runtime.dit_kv_cache_length; } else { @@ -1224,6 +1374,7 @@ int cosyvoice_server_webui_run(server_runtime& runtime) d["top_p"] = 0.8; d["win_size"] = 100; d["tau_r"] = 0.0; + d["chunk_tokens"] = 0; } #ifndef COSYVOICE_NO_ICU d["text_normalization"] = runtime.text_normalization_enabled; @@ -1237,6 +1388,10 @@ int cosyvoice_server_webui_run(server_runtime& runtime) res.set_content(d.dump(), "application/json"); }); + // ---- Register OpenAI-compatible API routes (if model is loaded) ---- + if (!runtime.model_slots.empty()) + cosyvoice_server_register_api_routes(server, runtime); + // ---- Error handler (404) ---- server.set_error_handler([](const Request& req, Response& res) { @@ -1264,7 +1419,13 @@ int cosyvoice_server_webui_run(server_runtime& runtime) runtime.api_key.empty() ? "no" : "yes"); if (runtime.sample_rate > 0) print_info_log(runtime.log_level, " sample_rate : %u\n", runtime.sample_rate); - print_info_log(runtime.log_level, " speakers : %zu\n", runtime.voice_names.size()); + { + const auto speakers_str = join_strings(runtime.voice_names, ", "); + print_info_log(runtime.log_level, " speakers : %s\n", speakers_str.empty() ? "-" : speakers_str.c_str()); + } + print_info_log(runtime.log_level, " stream : %s\n", runtime.stream ? "enabled" : "disabled"); + if (runtime.has_chunk_tokens) + print_info_log(runtime.log_level, " chunk_tokens : %u\n", runtime.chunk_tokens); #if !defined(COSYVOICE_NO_FRONTEND) print_info_log(runtime.log_level, " frontend_available : %s\n", runtime.frontend_ctx ? "yes" : "no"); diff --git a/tools/server/cosyvoice-server.cpp b/tools/server/cosyvoice-server.cpp index b90eb96..3222054 100644 --- a/tools/server/cosyvoice-server.cpp +++ b/tools/server/cosyvoice-server.cpp @@ -50,10 +50,22 @@ struct server_options uint32_t seed = 0; bool has_llm_kv_cache_type = false; - cosyvoice_llm_kv_cache_type_t llm_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( - COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0, - COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0, - COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0); + cosyvoice_kv_cache_type_t llm_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( + COSYVOICE_KV_CACHE_TYPE_Q8_0, + COSYVOICE_KV_CACHE_TYPE_Q4_0, + COSYVOICE_KV_CACHE_TYPE_Q8_0); + + bool has_dit_kv_cache_type = false; + cosyvoice_kv_cache_type_t dit_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( + COSYVOICE_KV_CACHE_TYPE_Q8_0, + COSYVOICE_KV_CACHE_TYPE_Q4_0, + COSYVOICE_KV_CACHE_TYPE_Q8_0); + uint32_t dit_kv_fixed_slots = 0; + uint32_t dit_kv_offloadable_slots = 0; + uint32_t dit_kv_cache_length = 0; + bool stream = false; + bool has_chunk_tokens = false; + uint32_t chunk_tokens = 0; bool has_temperature = false; float temperature = 0.0f; @@ -115,6 +127,13 @@ static void print_usage(const char* argv0) printf(" KV cache type. Single type (e.g. q8_0) uses the same format for K and V.\n"); printf(" Default: k=q8_0,v=q4_0,fallback=q8_0.\n"); printf(" --seed Default random seed for built-in sampler.\n"); + printf(" --dit-kv-cache-type ,v=[,fallback=]>\n"); + printf(" DiT KV cache type. Default: k=q8_0,v=q4_0,fallback=q8_0.\n"); + printf(" --dit-kv-fixed-slots DiT KV fixed slots (0 = auto).\n"); + printf(" --dit-kv-offloadable-slots DiT KV offloadable slots (0 = auto).\n"); + printf(" --dit-kv-cache-length DiT KV cache length (0 = auto).\n"); + printf(" --stream Enable streaming for TTS requests.\n"); + printf(" --chunk-tokens Tokens per streaming chunk. Default: model-defined.\n"); printf("\nSampling defaults (server-level):\n"); printf(" --temperature Sampling temperature (> 0).\n"); @@ -224,11 +243,15 @@ static std::string derive_served_model_name(const std::string& model_path) static bool init_model_context(const server_options& options, ggml_backend_t backend, server_runtime* runtime) { - cosyvoice_context_params_v2_cpp context_params; + cosyvoice_context_params_v3_cpp context_params; cosyvoice_init_default_context_params(&context_params); context_params.inference_buffer_policy = options.inference_buffer_policy; context_params.n_max_seq = options.max_llm_len; context_params.llm_kv_cache_type = options.llm_kv_cache_type; + context_params.dit_kv_cache_type = options.dit_kv_cache_type; + context_params.dit_kv_fixed_slots = options.dit_kv_fixed_slots; + context_params.dit_kv_offloadable_slots = options.dit_kv_offloadable_slots; + context_params.dit_kv_cache_length = options.dit_kv_cache_length; if (options.has_seed) context_params.seed = options.seed; context_params.n_workers = options.concurrency; @@ -246,6 +269,11 @@ static bool init_model_context(const server_options& options, ggml_backend_t bac return false; } + // Save effective DiT KV cache params from the context used to load + runtime->dit_kv_fixed_slots = context_params.dit_kv_fixed_slots; + runtime->dit_kv_offloadable_slots = context_params.dit_kv_offloadable_slots; + runtime->dit_kv_cache_length = context_params.dit_kv_cache_length; + return true; } @@ -363,9 +391,13 @@ static bool build_runtime(const server_options& options, server_runtime* runtime runtime->split_text_enabled = options.split_text_enabled; runtime->fast_split_text_enabled = options.fast_split_text_enabled; runtime->fade_in_enabled = options.fade_in_enabled; + runtime->stream = options.stream; + runtime->has_chunk_tokens = options.has_chunk_tokens; + runtime->chunk_tokens = options.chunk_tokens; runtime->log_level = get_log_level(options); runtime->seed_rng.seed(cosyvoice_generate_random_seed()); + runtime->slot_in_use.resize(runtime->concurrency, false); cosyvoice_init_backend_from_path(options.backend_path.empty() ? nullptr : options.backend_path.c_str()); @@ -393,8 +425,8 @@ static bool build_runtime(const server_options& options, server_runtime* runtime runtime->actual_llm_kv_cache_type = actual_params.llm_kv_cache_type; if (!options.quiet) fprintf(stderr, "llm_kv_cache_type: requested: %s, actual: %s\n", - llm_kv_cache_type_to_string(options.llm_kv_cache_type).c_str(), - llm_kv_cache_type_to_string(actual_params.llm_kv_cache_type).c_str()); + kv_cache_type_to_string(options.llm_kv_cache_type).c_str(), + kv_cache_type_to_string(actual_params.llm_kv_cache_type).c_str()); } if (!apply_generation_defaults(options, runtime)) @@ -446,9 +478,13 @@ static bool webui_build_runtime(const server_options& options, server_runtime* r runtime->split_text_enabled = options.split_text_enabled; runtime->fast_split_text_enabled = options.fast_split_text_enabled; runtime->fade_in_enabled = options.fade_in_enabled; + runtime->stream = options.stream; + runtime->has_chunk_tokens = options.has_chunk_tokens; + runtime->chunk_tokens = options.chunk_tokens; runtime->log_level = get_log_level(options); runtime->seed_rng.seed(cosyvoice_generate_random_seed()); + runtime->slot_in_use.resize(runtime->concurrency, false); cosyvoice_init_backend_from_path(options.backend_path.empty() ? nullptr : options.backend_path.c_str()); @@ -810,8 +846,8 @@ int tool_entry(int argc, char** argv) else if (str_casecmp(arg, "--llm-kv-cache-type") == 0) { const auto value = get_arg_value(); - cosyvoice_llm_kv_cache_type_t type; - if (!parse_llm_kv_cache_type_arg(value, &type)) + cosyvoice_kv_cache_type_t type; + if (!parse_kv_cache_type_arg(value, &type)) { fprintf(stderr, "Error: invalid --llm-kv-cache-type value \"%s\".\n", value); return 1; @@ -819,6 +855,65 @@ int tool_entry(int argc, char** argv) options.llm_kv_cache_type = type; options.has_llm_kv_cache_type = true; } + else if (str_casecmp(arg, "--dit-kv-cache-type") == 0) + { + const auto value = get_arg_value(); + cosyvoice_kv_cache_type_t type; + if (!parse_kv_cache_type_arg(value, &type)) + { + fprintf(stderr, "Error: invalid --dit-kv-cache-type value \"%s\".\n", value); + return 1; + } + options.dit_kv_cache_type = type; + options.has_dit_kv_cache_type = true; + } + else if (str_casecmp(arg, "--dit-kv-fixed-slots") == 0) + { + const auto value = get_arg_value(); + uint32_t v; + if (!parse_uint32_arg(value, &v)) + { + fprintf(stderr, "Error: invalid --dit-kv-fixed-slots value \"%s\".\n", value); + return 1; + } + options.dit_kv_fixed_slots = v; + } + else if (str_casecmp(arg, "--dit-kv-offloadable-slots") == 0) + { + const auto value = get_arg_value(); + uint32_t v; + if (!parse_uint32_arg(value, &v)) + { + fprintf(stderr, "Error: invalid --dit-kv-offloadable-slots value \"%s\".\n", value); + return 1; + } + options.dit_kv_offloadable_slots = v; + } + else if (str_casecmp(arg, "--dit-kv-cache-length") == 0) + { + const auto value = get_arg_value(); + uint32_t v; + if (!parse_uint32_arg(value, &v)) + { + fprintf(stderr, "Error: invalid --dit-kv-cache-length value \"%s\".\n", value); + return 1; + } + options.dit_kv_cache_length = v; + } + else if (str_casecmp(arg, "--stream") == 0) + options.stream = true; + else if (str_casecmp(arg, "--chunk-tokens") == 0) + { + const auto value = get_arg_value(); + uint32_t v; + if (!parse_uint32_arg(value, &v)) + { + fprintf(stderr, "Error: invalid --chunk-tokens value \"%s\".\n", value); + return 1; + } + options.chunk_tokens = v; + options.has_chunk_tokens = true; + } else if (str_casecmp(arg, "--verbose") == 0 || str_casecmp(arg, "-v") == 0) options.verbose = true; else if (str_casecmp(arg, "--quiet") == 0 || str_casecmp(arg, "-q") == 0) diff --git a/tools/server/cosyvoice-server.h b/tools/server/cosyvoice-server.h index 5332488..0a18ef0 100644 --- a/tools/server/cosyvoice-server.h +++ b/tools/server/cosyvoice-server.h @@ -3,7 +3,7 @@ #include "tool_common_cosyvoice.h" #include -#include +#include #include #include #include @@ -36,8 +36,8 @@ struct server_runtime uint32_t concurrency = 1; cosyvoice_inference_buffer_policy_t inference_buffer_policy = COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED; bool has_llm_kv_cache_override = false; - cosyvoice_llm_kv_cache_type_t requested_llm_kv_cache_type = static_cast(0); - cosyvoice_llm_kv_cache_type_t actual_llm_kv_cache_type = static_cast(0); + cosyvoice_kv_cache_type_t requested_llm_kv_cache_type = static_cast(0); + cosyvoice_kv_cache_type_t actual_llm_kv_cache_type = static_cast(0); #ifndef COSYVOICE_NO_ICU bool text_normalization_enabled = true; #endif @@ -46,6 +46,15 @@ struct server_runtime bool fade_in_enabled = true; server_log_level log_level = server_log_level::concise; + bool stream = false; + bool has_chunk_tokens = false; + uint32_t chunk_tokens = 0; + + // Effective DiT KV cache params (populated after model load) + uint32_t dit_kv_fixed_slots = 0; + uint32_t dit_kv_offloadable_slots = 0; + uint32_t dit_kv_cache_length = 0; + // Frontend model paths (ONNX, for feature extraction) std::string frontend_model; std::string speech_tokenizer; @@ -67,7 +76,10 @@ struct server_runtime std::vector>> voice_sessions; std::mt19937 seed_rng; - std::atomic_uint32_t thread_slot_counter{0}; + + // Concurrency slot management (request-scoped, not thread-local) + std::mutex slot_mutex; + std::vector slot_in_use; }; inline cosyvoice_context_t get_slot_model_context(server_runtime& runtime, uint32_t slot) @@ -84,12 +96,32 @@ inline cosyvoice_tts_context_t get_slot_voice_session(server_runtime& runtime, u return nullptr; } +// Acquire a concurrency slot for this request. Returns UINT32_MAX if all +// slots are in use (caller should return 503 / overloaded). +inline uint32_t acquire_thread_slot(server_runtime& runtime) +{ + std::lock_guard lock(runtime.slot_mutex); + for (uint32_t i = 0; i < runtime.concurrency; ++i) + if (!runtime.slot_in_use[i]) + { + runtime.slot_in_use[i] = true; + return i; + } + return UINT32_MAX; +} + +// Release a concurrency slot previously acquired via acquire_thread_slot. +inline void release_thread_slot(server_runtime& runtime, uint32_t slot) +{ + if (slot >= runtime.concurrency) + return; + std::lock_guard lock(runtime.slot_mutex); + runtime.slot_in_use[slot] = false; +} + inline uint32_t get_or_assign_thread_slot(server_runtime& runtime) { - thread_local uint32_t slot = UINT32_MAX; - if (slot == UINT32_MAX) - slot = runtime.thread_slot_counter.fetch_add(1, std::memory_order_relaxed); - return slot; + return acquire_thread_slot(runtime); } int cosyvoice_server_backend_run(server_runtime& runtime); diff --git a/tools/server/httplib.ixx b/tools/server/httplib.ixx index 0caeeb4..abaa939 100644 --- a/tools/server/httplib.ixx +++ b/tools/server/httplib.ixx @@ -7,6 +7,7 @@ export using Request = httplib::Request; export using Response = httplib::Response; export using ThreadPool = httplib::ThreadPool; export using HandlerResponse = httplib::Server::HandlerResponse; +export using DataSink = httplib::DataSink; // NOTE: Do NOT export `using Server = httplib::Server;` directly! // Tested on MSVC: this triggers unresolved externals (e.g. `__tlregdtor`) when using C++20 modules, @@ -21,7 +22,7 @@ public: Server() : new_task_queue(server.new_task_queue) {} decltype(httplib::Server::new_task_queue)& new_task_queue; - + template void Get(const std::string& pattern, Handler handler) { @@ -45,6 +46,12 @@ public: { server.Put(pattern, std::move(handler)); } + + template + void Options(const std::string& pattern, Handler handler) + { + server.Options(pattern, std::move(handler)); + } template void set_pre_routing_handler(Handler handler) diff --git a/tools/server/server_common.cpp b/tools/server/server_common.cpp index e660e9a..5cbf11c 100644 --- a/tools/server/server_common.cpp +++ b/tools/server/server_common.cpp @@ -486,6 +486,123 @@ std::string supported_response_formats_to_string(const server_runtime& runtime) return result; } +// --------------------------------------------------------------------------- +// Streaming audio helpers +// --------------------------------------------------------------------------- + +bool build_wav_header(std::string* output, uint32_t sample_rate) +{ + constexpr uint16_t num_channels = 1; + constexpr uint16_t bits_per_sample = 16; + const uint32_t byte_rate = sample_rate * num_channels * bits_per_sample / 8; + const uint16_t block_align = num_channels * bits_per_sample / 8; + + // Use 0x7FFFFFFF as the data size placeholder (most players treat + // this as "until EOF", which is correct for streaming). + const uint32_t subchunk2_size = 0x7FFFFFFF; + const uint32_t chunk_size = 4 + (8 + 16) + (8 + subchunk2_size); + + output->clear(); + output->reserve(44); + + // RIFF header + output->append("RIFF", 4); + output->append(reinterpret_cast(&chunk_size), 4); + output->append("WAVE", 4); + + // fmt subchunk + output->append("fmt ", 4); + const uint32_t fmt_subchunk_size = 16; + output->append(reinterpret_cast(&fmt_subchunk_size), 4); + const uint16_t audio_format = 1; // PCM + output->append(reinterpret_cast(&audio_format), 2); + output->append(reinterpret_cast(&num_channels), 2); + output->append(reinterpret_cast(&sample_rate), 4); + output->append(reinterpret_cast(&byte_rate), 4); + output->append(reinterpret_cast(&block_align), 2); + output->append(reinterpret_cast(&bits_per_sample), 2); + + // data subchunk header + output->append("data", 4); + output->append(reinterpret_cast(&subchunk2_size), 4); + + return true; +} + +bool stream_audio_chunk(streaming_callback_context* ctx, const float* pcm, uint32_t n_samples) +{ + if (!ctx || !ctx->sink || !pcm || n_samples == 0) + return false; + + constexpr size_t kBytesPerSample = sizeof(int16_t); + const size_t chunk_bytes = static_cast(n_samples) * kBytesPerSample; + + // PCM or WAV: convert float -> int16 and write raw PCM16 + if (ctx->format == response_audio_format::pcm || + ctx->format == response_audio_format::wav) + { + // WAV format: PCM16 data written after the header already sent + std::string buf; + buf.resize(chunk_bytes); + for (uint32_t i = 0; i < n_samples; ++i) + { + const auto sample = float_to_pcm16(pcm[i]); + buf[2 * i] = static_cast(sample & 0xFF); + buf[2 * i + 1] = static_cast((sample >> 8) & 0xFF); + } + + if (!ctx->sink->write(buf.data(), buf.size())) + { + if (ctx->aborted) *ctx->aborted = true; + return false; + } + return true; + } + + // MP3, OPUS, AAC, FLAC, M4A: use the audio encoder +#ifndef COSYVOICE_NO_AUDIO + if (ctx->encoder) + { + auto enc = static_cast(ctx->encoder); + cosyvoice_audio_encoding_format_t enc_format; + switch (ctx->format) + { + case response_audio_format::mp3: enc_format = COSYVOICE_AUDIO_ENCODING_FORMAT_MP3; break; + case response_audio_format::opus: enc_format = COSYVOICE_AUDIO_ENCODING_FORMAT_OPUS; break; + case response_audio_format::aac: enc_format = COSYVOICE_AUDIO_ENCODING_FORMAT_AAC; break; + case response_audio_format::m4a: enc_format = COSYVOICE_AUDIO_ENCODING_FORMAT_M4A; break; + case response_audio_format::flac: enc_format = COSYVOICE_AUDIO_ENCODING_FORMAT_FLAC; break; + default: + if (ctx->error_out) *ctx->error_out = "Unsupported format for streaming."; + return false; + } + + if (!cosyvoice_audio_encoder_encode(enc, pcm, n_samples, enc_format)) + { + if (ctx->error_out) *ctx->error_out = "Audio encoder failed on streaming chunk."; + return false; + } + + const uint8_t* encoded = nullptr; + uint32_t encoded_length = 0; + cosyvoice_audio_encoder_get_encoded_data(enc, &encoded, &encoded_length); + + if (encoded && encoded_length > 0) + { + if (!ctx->sink->write(reinterpret_cast(encoded), encoded_length)) + { + if (ctx->aborted) *ctx->aborted = true; + return false; + } + } + return true; + } +#endif + + if (ctx->error_out) *ctx->error_out = "No encoder available for this format."; + return false; +} + // --------------------------------------------------------------------------- // String join // --------------------------------------------------------------------------- diff --git a/tools/server/server_common.h b/tools/server/server_common.h index b654434..b8da74c 100644 --- a/tools/server/server_common.h +++ b/tools/server/server_common.h @@ -1,13 +1,21 @@ -#pragma once +#pragma once #include "cosyvoice-server.h" +#include #include #include #include #include -namespace httplib { class Request; } +namespace httplib { class Request; class Response; } + +#ifdef COSYVOICE_SERVER_USE_PCH +namespace httplib { class Server; } +using Server = httplib::Server; +#else +class Server; +#endif // --------------------------------------------------------------------------- // Audio output format @@ -38,7 +46,7 @@ void log_message(server_log_level current, server_log_level required, // Timestamped error message to stderr. void print_error_log(const char* format, ...); -// Level-gated info — no timestamp, no trailing newline (for startup banners). +// Level-gated info - no timestamp, no trailing newline (for startup banners). void print_info_log(server_log_level level, const char* format, ...); // --------------------------------------------------------------------------- @@ -95,7 +103,7 @@ bool apply_generation_overrides( std::string* error_message); // --------------------------------------------------------------------------- -// Request logging (request ID, timing, status) — shared between API and WebUI +// Request logging (request ID, timing, status) - shared between API and WebUI // --------------------------------------------------------------------------- enum class request_log_status @@ -118,6 +126,7 @@ struct request_log_context bool has_instructions = false; bool has_seed = false; uint32_t seed = 0; + bool stream = false; std::chrono::steady_clock::time_point start; }; @@ -146,13 +155,13 @@ const char* request_log_status_to_string(request_log_status status); bool build_pcm16_bytes(const float* data, uint32_t length, std::string* output, std::string* error); -// Native WAV builder (no audio encoder) — only compiled when COSYVOICE_NO_AUDIO. +// Native WAV builder (no audio encoder) - only compiled when COSYVOICE_NO_AUDIO. #ifdef COSYVOICE_NO_AUDIO bool build_wav_bytes(const float* data, uint32_t length, uint32_t sample_rate, std::string* output, std::string* error); #endif -// Audio-encoder-based builder — only compiled when !COSYVOICE_NO_AUDIO. +// Audio-encoder-based builder - only compiled when !COSYVOICE_NO_AUDIO. #ifndef COSYVOICE_NO_AUDIO bool build_wav_bytes_with_audio_encoder(cosyvoice_audio_encoder_t encoder, const float* data, uint32_t length, @@ -160,7 +169,7 @@ bool build_wav_bytes_with_audio_encoder(cosyvoice_audio_encoder_t encoder, std::string* output, std::string* error); #endif -// Main audio payload builder — dispatches to the correct encoder based on format. +// Main audio payload builder - dispatches to the correct encoder based on format. bool build_audio_payload(response_audio_format format, const cosyvoice_generated_speech& generated, const server_runtime& runtime, @@ -177,8 +186,49 @@ bool is_response_format_supported(const server_runtime& runtime response_audio_format format); std::string supported_response_formats_to_string(const server_runtime& runtime); +// --------------------------------------------------------------------------- +// Streaming audio helpers +// --------------------------------------------------------------------------- + +// Forward declarations from httplib +namespace httplib { class DataSink; } + +// Context passed through the cosyvoice streaming TTS callback. +// encoder is a cosyvoice_audio_encoder_t stored as void* to avoid +// conditional-compilation complexity in this header. +struct streaming_callback_context +{ + httplib::DataSink* sink; + void* encoder; // cosyvoice_audio_encoder_t or NULL + response_audio_format format; + uint32_t sample_rate; + bool wav_header_written; + std::string* error_out; + std::atomic* aborted; +}; + +// Build a standard 44-byte RIFF/WAV header with a placeholder data size +// (0x7FFFFFFF, meaning "until EOF" — accepted by most WAV players). +bool build_wav_header(std::string* output, uint32_t sample_rate); + +// Encode a single PCM chunk and write it to the DataSink. +// Called from within a cosyvoice streaming TTS callback. +// Returns false to abort TTS if the sink write fails (sets *aborted). +bool stream_audio_chunk(streaming_callback_context* ctx, const float* pcm, uint32_t n_samples); + // --------------------------------------------------------------------------- // Generic string join // --------------------------------------------------------------------------- std::string join_strings(const std::vector& items, const char* sep); + +// OpenAI-compatible error response helper (shared between API and WebUI modes) +void set_openai_error(httplib::Response& res, int status, const std::string& message, + const std::string& type, const char* param, const char* code); + +// Bearer token authentication (shared between API and WebUI modes) +bool require_auth(const httplib::Request& req, const server_runtime& runtime, httplib::Response& res); + +// Register OpenAI-compatible API routes on an existing server instance. +// Called by both API mode and WebUI mode (when a model is loaded). +void cosyvoice_server_register_api_routes(Server& server, server_runtime& runtime); diff --git a/tools/server/synthesize_via_api.py b/tools/server/synthesize_via_api.py index 7f4bb8d..8d9fe6d 100644 --- a/tools/server/synthesize_via_api.py +++ b/tools/server/synthesize_via_api.py @@ -24,7 +24,7 @@ def guess_output_path(response_format: str) -> str: def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( - description="Call cosyvoice-server via OpenAI SDK and save speech audio." + description="Call cosyvoice-server via OpenAI API and save/play speech audio." ) parser.add_argument( "--base-url", @@ -60,6 +60,9 @@ def parse_args() -> argparse.Namespace: default=None, help="Optional max_token_text_ratio extension (>= 0)", ) + parser.add_argument("--stream", action="store_true", help="Enable streaming TTS (play while receiving if format is wav)") + parser.add_argument("--chunk-tokens", type=int, default=None, help="Tokens per streaming chunk") + parser.add_argument("--no-play", action="store_true", help="Disable real-time playback even with --stream") parser.add_argument( "--output", default="", @@ -107,6 +110,9 @@ def validate_args(args: argparse.Namespace) -> int: if args.max_token_text_ratio is not None and args.max_token_text_ratio < 0: print("Error: --max-token-text-ratio must be >= 0", file=sys.stderr) return 2 + if args.chunk_tokens is not None and args.chunk_tokens < 0: + print("Error: --chunk-tokens must be >= 0", file=sys.stderr) + return 2 if ( args.min_token_text_ratio is not None and args.max_token_text_ratio is not None @@ -118,6 +124,76 @@ def validate_args(args: argparse.Namespace) -> int: return 0 +def request_and_stream(base_url, api_key, body, response_format, output_path, play, timeout): + """Send streaming TTS request, play chunks in real-time, and save to file.""" + import requests + + url = base_url + "/audio/speech" + headers = { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + } + + resp = requests.post(url, headers=headers, json=body, stream=True, timeout=timeout) + resp.raise_for_status() + + _play_stream = None + _pyaudio_instance = None + _all_data = bytearray() + _wav_header_buf = bytearray() + + try: + _use_pyaudio = play and response_format == "wav" + if _use_pyaudio: + try: + import pyaudio as _pa + except ImportError: + print("Warning: pyaudio not installed. Install with: pip install pyaudio", file=sys.stderr) + _use_pyaudio = False + else: + _pyaudio_instance = _pa.PyAudio() + + for chunk in resp.iter_content(chunk_size=None): + if not chunk: + continue + _all_data.extend(chunk) + + if _use_pyaudio: + if _play_stream is None: + _wav_header_buf.extend(chunk) + if len(_wav_header_buf) >= 44: + import struct + channels = struct.unpack_from(' 44: + _play_stream.write(bytes(_wav_header_buf[44:])) + else: + _play_stream.write(chunk) + + output_path.write_bytes(bytes(_all_data)) + print(f"Saved {len(_all_data)} bytes to {output_path}") + finally: + if _play_stream: + _play_stream.stop_stream() + _play_stream.close() + if _pyaudio_instance: + _pyaudio_instance.terminate() + + +def request_and_save(client, kwargs, response_format, output_path): + """Non-streaming request via OpenAI SDK.""" + audio = client.audio.speech.create(**kwargs) + audio_bytes = audio.read() + output_path.write_bytes(audio_bytes) + return len(audio_bytes) + + def main() -> int: args = parse_args() code = validate_args(args) @@ -165,9 +241,13 @@ def main() -> int: extension_body["min_token_text_ratio"] = args.min_token_text_ratio if args.max_token_text_ratio is not None: extension_body["max_token_text_ratio"] = args.max_token_text_ratio + if args.stream: + extension_body["stream"] = True + if args.chunk_tokens is not None: + extension_body["chunk_tokens"] = args.chunk_tokens if extension_body: request_kwargs["extra_body"] = extension_body - # Determine response_format: explicit arg > infer from output extension > default 'wav' + response_format = args.response_format if response_format is None: if args.output: @@ -185,16 +265,11 @@ def main() -> int: output_path = pathlib.Path(output) output_path.parent.mkdir(parents=True, exist_ok=True) + play = args.stream and not args.no_play + total_bytes = 0 saved_files = [] for i in range(args.requests): - try: - audio = client.audio.speech.create(**request_kwargs) - audio_bytes = audio.read() - except Exception as exc: - print(f"Request failed (#{i + 1}): {exc}", file=sys.stderr) - return 1 - if args.requests == 1: curr_path = output_path else: @@ -202,11 +277,29 @@ def main() -> int: f"{output_path.stem}_{i + 1:03d}{output_path.suffix}" ) - curr_path.write_bytes(audio_bytes) - total_bytes += len(audio_bytes) - saved_files.append(curr_path) - - print(f"Saved {len(audio_bytes)} bytes to {curr_path}") + try: + if args.stream: + # Build the full request body (kwargs + extra_body) + body = {k: v for k, v in request_kwargs.items() if k != "extra_body"} + if "extra_body" in request_kwargs: + body.update(request_kwargs["extra_body"]) + request_and_stream( + base_url=base_url, + api_key=api_key, + body=body, + response_format=response_format, + output_path=curr_path, + play=play, + timeout=args.timeout, + ) + else: + total = request_and_save(client, request_kwargs, response_format, curr_path) + total_bytes += total + saved_files.append(curr_path) + print(f"Saved {total} bytes to {curr_path}") + except Exception as exc: + print(f"Request failed (#{i + 1}): {exc}", file=sys.stderr) + return 1 if args.requests > 1: print(f"Saved {len(saved_files)} files, total {total_bytes} bytes") diff --git a/tools/server/webui/cosyvoice-webui.css b/tools/server/webui/cosyvoice-webui.css index 5c84268..59876da 100644 --- a/tools/server/webui/cosyvoice-webui.css +++ b/tools/server/webui/cosyvoice-webui.css @@ -701,6 +701,15 @@ details .params > * { min-width: 120px; } +.config-section-title { + font-size: 11px; + font-weight: 600; + text-transform: uppercase; + letter-spacing: 0.6px; + color: var(--text-muted); + margin-bottom: 6px; +} + /* ---- Spinner ---- */ .spinner { display: inline-block; diff --git a/tools/server/webui/cosyvoice-webui.html b/tools/server/webui/cosyvoice-webui.html index 41d7bf4..77c5273 100644 --- a/tools/server/webui/cosyvoice-webui.html +++ b/tools/server/webui/cosyvoice-webui.html @@ -102,25 +102,48 @@

Model Management

Advanced Model Config -
-
-
-
-
-
-
-
+
+
LLM
+
+
+
+
+
+
+
+
+
+
+
+
DiT
+
+
+
+
+
+
+
+
+
@@ -282,7 +305,16 @@

TTS Synthesis

-
+ +
+ +
+
+
+
+
+
+
diff --git a/tools/server/webui/cosyvoice-webui.js b/tools/server/webui/cosyvoice-webui.js index c61d795..f377fa5 100644 --- a/tools/server/webui/cosyvoice-webui.js +++ b/tools/server/webui/cosyvoice-webui.js @@ -29,7 +29,9 @@ const ADV_PARAM_IDS = [ 'tts-fadein','tts-textnorm','tts-split','tts-fastsplit', 'tts-mode','tts-format', 'model-max-llm','model-kv-k','model-kv-v','model-buffer-policy', - 'model-threads','model-backend' + 'model-threads','model-backend', + 'model-dit-kv-k','model-dit-kv-v','model-dit-fixed-slots','model-dit-offloadable-slots','model-dit-cache-length', + 'tts-stream','tts-chunk-tokens' ]; function saveAdvParam(id) { @@ -137,6 +139,7 @@ function initEls() { // Model 'model-path', 'model-backend', 'model-threads', 'model-kv-k', 'model-kv-v', 'model-buffer-policy', 'model-max-llm', + 'model-dit-kv-k', 'model-dit-kv-v', 'model-dit-fixed-slots', 'model-dit-offloadable-slots', 'model-dit-cache-length', 'btn-reset-model-config', 'btn-load-model', 'btn-unload-model', 'model-load-area', 'model-loaded-area', 'model-loaded-info', @@ -164,6 +167,7 @@ function initEls() { 'tts-speed', 'speed-val', 'tts-seed', 'tts-temp', 'tts-topk', 'tts-topp', 'tts-winsize', 'tts-taur', 'tts-fadein', 'tts-textnorm', 'tts-split', 'tts-fastsplit', + 'tts-stream', 'tts-chunk-tokens', 'seed-dice', 'seed-lock', 'btn-tts', 'btn-reset-gen-config', 'player-area', 'audio-player', 'download-link', @@ -485,6 +489,16 @@ function initModelLoad() { const ml = parseInt(els['model-max-llm'].value, 10); if (ml > 0) body.max_llm_len = ml; + const dkt = els['model-dit-kv-k'].value; + const dvt = els['model-dit-kv-v'].value; + body.dit_kv_cache_type = (dkt === dvt) ? dkt : 'k=' + dkt + ',v=' + dvt; + const dfs = parseInt(els['model-dit-fixed-slots'].value, 10) || 0; + body.dit_kv_fixed_slots = dfs; + const dos = parseInt(els['model-dit-offloadable-slots'].value, 10) || 0; + body.dit_kv_offloadable_slots = dos; + const dcl = parseInt(els['model-dit-cache-length'].value, 10) || 0; + body.dit_kv_cache_length = dcl; + await apiFetch('/model/load', { method: 'POST', ...jsonBody(body) }); showSuccess(els['model-success'], 'Model loaded successfully'); showToast('Model loaded: ' + model, 'success'); @@ -1139,6 +1153,25 @@ function initTts() { } } +// After streaming finishes: set up download link, history, reset button. +// The audio player already has the data via MediaSource or blob URL. +function onStreamDone(chunks, contentType, els, voice, ext, text) { + if (chunks.length === 0) return; + const blob = new Blob(chunks, { type: contentType }); + const url = URL.createObjectURL(blob); + if (els['download-link']) { + els['download-link'].href = url; + els['download-link'].download = 'cosyvoice_' + voice + '.' + ext; + els['download-link'].innerHTML = ICONS.download + ' Download ' + ext.toUpperCase() + ' (' + formatFileSize(blob.size) + ')'; + } + addHistory({ voice, text, blob, url, format: ext, mode: els['tts-mode'] ? els['tts-mode'].value : '', timestamp: new Date() }); + showToast('Stream complete: ' + formatFileSize(blob.size), 'success'); + if (els['btn-tts']) { + els['btn-tts'].disabled = false; + els['btn-tts'].innerHTML = ICONS.speaker + ' Generate Speech'; + } +} + async function generateTts() { hideError(els['tts-error']); const voice = els['tts-voice'].value; @@ -1158,6 +1191,12 @@ async function generateTts() { fast_split: els['tts-fastsplit'] ? els['tts-fastsplit'].checked : true, }; + if (els['tts-stream']) body.stream = els['tts-stream'].checked; + if (els['tts-chunk-tokens']) { + const ct = parseInt(els['tts-chunk-tokens'].value, 10); + if (ct > 0) body.chunk_tokens = ct; + } + const instr = els['tts-instructions'] ? els['tts-instructions'].value.trim() : ''; if (instr) body.instructions = instr; @@ -1171,10 +1210,7 @@ async function generateTts() { ]; for (const [k, el] of adv) { const v = el.value.trim(); - if (v === '') { - showError(els['tts-error'], 'Please fill in all advanced parameters (empty: "' + k + '").'); - return; - } + if (v === '') continue; // skip empty — server uses its defaults body[k] = k === 'seed' || k === 'top_k' || k === 'win_size' ? parseInt(v, 10) : parseFloat(v); } @@ -1185,35 +1221,119 @@ async function generateTts() { btn.innerHTML = ' Generating...'; try { - const res = await apiFetch('/tts', { method: 'POST', ...jsonBody(body) }); - const blob = await res.blob(); - const url = URL.createObjectURL(blob); - if (els['audio-player']) els['audio-player'].src = url; - if (els['player-area']) els['player-area'].style.display = 'block'; + const isStream = body.stream === true; + let ext = body.format || 'wav'; + + if (isStream) { + const res = await fetch('/tts', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(body), + credentials: 'same-origin', + }); + if (!res.ok) { + let msg = 'HTTP ' + res.status; + try { const j = await res.json(); msg = j.error || j.message || msg; } catch(e) { /* ignore */ } + throw new Error(typeof msg === 'string' ? msg : JSON.stringify(msg)); + } - const ext = body.format || 'wav'; - if (els['download-link']) { - els['download-link'].href = url; - els['download-link'].download = 'cosyvoice_' + voice + '.' + ext; - els['download-link'].innerHTML = ICONS.download + ' Download ' + ext.toUpperCase() + ' (' + formatFileSize(blob.size) + ')'; - } + const contentType = res.headers.get('Content-Type') || 'audio/mpeg'; + const reader = res.body.getReader(); + const chunks = []; + + // Set up audio player immediately + const audio = els['audio-player']; + if (els['player-area']) els['player-area'].style.display = 'block'; + btn.innerHTML = ' Streaming...'; + + // Use MediaSource for true progressive playback when supported + if (window.MediaSource && MediaSource.isTypeSupported(contentType)) { + const mediaSource = new MediaSource(); + const audioUrl = URL.createObjectURL(mediaSource); + if (audio) audio.src = audioUrl; + + mediaSource.onsourceopen = () => { + const sb = mediaSource.addSourceBuffer(contentType); + const pump = async () => { + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + chunks.push(value); + // Wait for previous append to complete + if (sb.updating) + await new Promise(r => sb.addEventListener('updateend', r, { once: true })); + sb.appendBuffer(value); + } + // Wait for final append to finish + if (sb.updating) + await new Promise(r => sb.addEventListener('updateend', r, { once: true })); + if (mediaSource.readyState === 'open') + mediaSource.endOfStream(); + onStreamDone(chunks, contentType, els, voice, ext, text); + } catch(e) { + console.error('Stream error:', e); + try { if (mediaSource.readyState === 'open') mediaSource.endOfStream(); } catch(_) {} + onStreamDone(chunks, contentType, els, voice, ext, text); + } + }; + pump(); + }; + + // Auto-play once data arrives + audio.oncanplaythrough = () => audio.play().catch(() => {}); + audio.oncanplay = () => audio.play().catch(() => {}); + } else { + // Fallback: accumulate all chunks then play + while (true) { + const { done, value } = await reader.read(); + if (done) break; + chunks.push(value); + } + const blob = new Blob(chunks, { type: contentType }); + const url = URL.createObjectURL(blob); + if (audio) audio.src = url; + audio.oncanplay = () => audio.play().catch(() => {}); + onStreamDone(chunks, contentType, els, voice, ext, text); + } + } else { + // Non-streaming path: wait for full blob + const res = await apiFetch('/tts', { method: 'POST', ...jsonBody(body) }); + const blob = await res.blob(); + const url = URL.createObjectURL(blob); + if (els['audio-player']) els['audio-player'].src = url; + if (els['player-area']) els['player-area'].style.display = 'block'; + + if (els['download-link']) { + els['download-link'].href = url; + els['download-link'].download = 'cosyvoice_' + voice + '.' + ext; + els['download-link'].innerHTML = ICONS.download + ' Download ' + ext.toUpperCase() + ' (' + formatFileSize(blob.size) + ')'; + } - // Auto-play - if (els['audio-player']) els['audio-player'].play().catch(() => {}); - - // Add to history - addHistory({ - voice: voice, - text: text, - blob: blob, - url: url, - format: ext, - mode: els['tts-mode'].value, - timestamp: new Date(), - }); + // Auto-play + if (els['audio-player']) els['audio-player'].play().catch(() => {}); + + // Add to history + addHistory({ + voice: voice, + text: text, + blob: blob, + url: url, + format: ext, + mode: els['tts-mode'].value, + timestamp: new Date(), + }); + } // Auto-roll seed if (!seedLocked && els['seed-dice']) els['seed-dice'].click(); + + // For streaming, the MediaSource pump runs asynchronously — don't + // reset the button here; it stays "Streaming..." until onStreamDone. + if (isStream) { + isGenerating = false; + return; + } } catch(e) { showError(els['tts-error'], 'Synthesis failed: ' + e.message); showToast('TTS generation failed', 'error'); @@ -1580,6 +1700,30 @@ function initResetButtons() { const splitCheck = document.getElementById('tts-split'); if (splitCheck) splitCheck.addEventListener('change', updateFastSplitState); + + // Streaming ↔ format auto-switch: enabling streaming forces mp3/opus/aac; + // switching format away from a streamable one disables streaming. + const streamCheck = document.getElementById('tts-stream'); + const formatSelect = document.getElementById('tts-format'); + const STREAMABLE_FORMATS = ['mp3', 'opus', 'm4a', 'aac', 'flac']; + if (streamCheck && formatSelect) { + const syncStreamable = () => { + if (streamCheck.checked && !STREAMABLE_FORMATS.includes(formatSelect.value)) { + if (!confirm('Streaming requires MP3, Opus, AAC or FLAC. Switch to MP3?')) { + streamCheck.checked = false; + } else { + formatSelect.value = 'mp3'; + } + } + }; + streamCheck.addEventListener('change', syncStreamable); + formatSelect.addEventListener('change', () => { + if (!STREAMABLE_FORMATS.includes(formatSelect.value)) { + streamCheck.checked = false; + showToast('Streaming disabled — format not compatible with progressive playback', 'warning'); + } + }); + } } // ---- Fetch Defaults ---- @@ -1606,6 +1750,42 @@ async function fetchDefaults() { } els['model-threads'].value = d.default_n_threads !== undefined ? d.default_n_threads : 0; + if (d.default_dit_k_cache_type && els['model-dit-kv-k']) { + const kSel = els['model-dit-kv-k']; + for (let i = 0; i < kSel.options.length; i++) { + if (kSel.options[i].value === d.default_dit_k_cache_type) { kSel.selectedIndex = i; break; } + } + } + if (d.default_dit_v_cache_type && els['model-dit-kv-v']) { + const vSel = els['model-dit-kv-v']; + for (let i = 0; i < vSel.options.length; i++) { + if (vSel.options[i].value === d.default_dit_v_cache_type) { vSel.selectedIndex = i; break; } + } + } + // DiT defaults: 0 = auto for slots, 10× LLM max seq for cache length + if (d.default_dit_kv_fixed_slots !== undefined && els['model-dit-fixed-slots']) + els['model-dit-fixed-slots'].value = d.default_dit_kv_fixed_slots; + else if (els['model-dit-fixed-slots'] && !els['model-dit-fixed-slots'].value) + els['model-dit-fixed-slots'].value = '0'; + + if (d.default_dit_kv_offloadable_slots !== undefined && els['model-dit-offloadable-slots']) + els['model-dit-offloadable-slots'].value = d.default_dit_kv_offloadable_slots; + else if (els['model-dit-offloadable-slots'] && !els['model-dit-offloadable-slots'].value) + els['model-dit-offloadable-slots'].value = '0'; + + // DiT KV Cache Length: server value > 0 → use it; otherwise 10× LLM max seq + if (d.default_dit_kv_cache_length && els['model-dit-cache-length']) + els['model-dit-cache-length'].value = d.default_dit_kv_cache_length; + else if (els['model-dit-cache-length'] && els['model-max-llm']) + { + const llm = parseInt(els['model-max-llm'].value, 10); + els['model-dit-cache-length'].value = llm > 0 ? String(llm * 10) : '0'; + } + + // Chunk tokens (0 = model default) + if (d.chunk_tokens !== undefined && els['tts-chunk-tokens']) + els['tts-chunk-tokens'].value = d.chunk_tokens; + if (d.temperature !== undefined) els['tts-temp'].value = r(d.temperature, 6); if (d.top_k !== undefined) els['tts-topk'].value = d.top_k; if (d.top_p !== undefined) els['tts-topp'].value = r(d.top_p, 6); From 167ab8054e1f13fc90f888412f10f149ab3a7cf3 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Thu, 16 Jul 2026 11:59:44 +0800 Subject: [PATCH 16/34] cli: add support for LLM and Flow flash attention options --- tools/cli/cosyvoice-cli.cpp | 78 +++++++++++++++++++++++++++++++++++-- 1 file changed, 75 insertions(+), 3 deletions(-) diff --git a/tools/cli/cosyvoice-cli.cpp b/tools/cli/cosyvoice-cli.cpp index f5defa9..ee1f71f 100644 --- a/tools/cli/cosyvoice-cli.cpp +++ b/tools/cli/cosyvoice-cli.cpp @@ -110,6 +110,10 @@ struct cli_options float min_token_text_ratio = 0.0f; bool has_max_token_text_ratio = false; float max_token_text_ratio = 0.0f; + bool has_llm_flash_attn = false; + bool llm_flash_attn = true; + bool has_flow_flash_attn = false; + bool flow_flash_attn = true; }; enum class cli_log_level @@ -277,6 +281,7 @@ static void print_interactive_commands() printf(" /seed-policy Show or set seed policy.\n"); printf(" /help Show command list.\n"); printf(" /stream Toggle streaming playback.\n"); + printf(" /chunk-tokens [value] Show or set tokens per streaming chunk.\n"); printf(" /exit Exit interactive mode. Ctrl+C also exits.\n"); } @@ -324,6 +329,8 @@ static void print_usage(const char* argv0) printf(" --dit-kv-cache-length DiT KV cache max seq length (interactive only). Default: max-llm-len * 10.\n"); printf(" --stream Enable streaming playback in interactive mode.\n"); printf(" --chunk-tokens Tokens per streaming chunk (interactive only). Default: model-defined.\n"); + printf(" --llm-flash-attn <0|1> Enable/disable LLM flash attention. Default: 1.\n"); + printf(" --flow-flash-attn <0|1> Enable/disable Flow/DiT flash attention. Default: 1.\n"); printf(" --seed Fixed seed for sampling.\n"); printf(" --seed-policy Seed strategy. Default: auto (fixed if --seed is set).\n"); @@ -789,12 +796,12 @@ static void print_tts_runtime_info( print_kv_line_u32("chunk_tokens", cosyvoice_get_chunk_tokens(ctx)); } print_kv_line_string("buffer_policy", inference_buffer_policy_to_string(context_params.inference_buffer_policy)); + print_kv_line_string("llm_use_flash_attn", enabled_to_string(context_params.llm_use_flash_attn)); + print_kv_line_string("flow_use_flash_attn", enabled_to_string(context_params.flow_use_flash_attn)); if (log_level == cli_log_level::verbose) { print_kv_line_u32("n_batch", context_params.n_batch); print_kv_line_u32("n_max_seq", context_params.n_max_seq); - print_kv_line_string("llm_use_flash_attn", enabled_to_string(context_params.llm_use_flash_attn)); - print_kv_line_string("flow_use_flash_attn", enabled_to_string(context_params.flow_use_flash_attn)); print_kv_line_string("llm_kv_fallback", enabled_to_string(context_params.llm_allow_kv_cache_fallback)); } @@ -1049,6 +1056,20 @@ static void run_interactive_loop( streaming = !streaming; printf("Streaming playback %s.\n", streaming ? "enabled" : "disabled"); } + else if (cmd == "/chunk-tokens") + { + if (tokens.size() > 1) + { + uint32_t v; + if (!parse_uint32_arg(tokens[1], &v)) + { + print_error_log("Error: invalid chunk-tokens value.\n"); + continue; + } + cosyvoice_set_chunk_tokens(ctx, v); + } + printf("chunk_tokens: %u\n", cosyvoice_get_chunk_tokens(ctx)); + } else if (cmd == "/play") { bool id_ok = false; @@ -1089,6 +1110,9 @@ static void run_interactive_loop( entry.seed = used_seed; entry.text = trimmed; + std::chrono::steady_clock::time_point gen_start; + std::chrono::steady_clock::time_point gen_end; + if (streaming) { struct stream_play_state @@ -1131,6 +1155,7 @@ static void run_interactive_loop( return true; })); + gen_start = std::chrono::steady_clock::now(); bool ok; if (options.mode == "cross-lingual") ok = cosyvoice_tts_cross_lingual_stream(tts_ctx, trimmed.c_str(), options.speed, stream_callback, &play_state); @@ -1140,6 +1165,7 @@ static void run_interactive_loop( ok = cosyvoice_tts_instruct_stream(tts_ctx, trimmed.c_str(), options.instruction.c_str(), options.speed, stream_callback, &play_state); inferencing_done = true; + gen_end = std::chrono::steady_clock::now(); playback_thread.join(); ctrl_c_guard._register(); @@ -1151,6 +1177,7 @@ static void run_interactive_loop( } else { + gen_start = std::chrono::steady_clock::now(); std::string error; auto pcm = generate_tts_audio(tts_ctx, options, trimmed, &error); if (!pcm.data) @@ -1160,13 +1187,16 @@ static void run_interactive_loop( continue; } entry.pcm.insert(entry.pcm.end(), pcm.data, pcm.data + pcm.length); + gen_end = std::chrono::steady_clock::now(); } const auto id = cache.next_id++; cache.last_success_id = id; const auto duration = pcm_length_seconds(sample_rate, entry.pcm.size()); cache.items.emplace(id, std::move(entry)); - printf("Generated audio %u (%.2fs), seed=%u.\n", cache.last_success_id, duration, used_seed); + const double elapsed_sec = elapsed_ms(gen_start, gen_end) / 1000.0; + const double rtf = (duration > 0.0) ? (elapsed_sec / duration) : 0.0; + printf("Generated audio %u (%.2fs, rtf=%.2f), seed=%u.\n", cache.last_success_id, duration, rtf, used_seed); } } @@ -1649,6 +1679,44 @@ int tool_entry(int argc, char** argv) options.fast_split_text_enabled = false; else if (str_casecmp(arg, "--disable-fade-in") == 0) options.fade_in_enabled = false; + else if (str_casecmp(arg, "--llm-flash-attn") == 0) + { + const auto v = to_lower(get_arg_value()); + if (v == "1" || v == "yes" || v == "true" || v == "on") + { + options.llm_flash_attn = true; + options.has_llm_flash_attn = true; + } + else if (v == "0" || v == "no" || v == "false" || v == "off") + { + options.llm_flash_attn = false; + options.has_llm_flash_attn = true; + } + else + { + print_error_log("Error: invalid --llm-flash-attn value \"%s\". Use 0/1, yes/no, true/false, on/off.\n", v.c_str()); + return 1; + } + } + else if (str_casecmp(arg, "--flow-flash-attn") == 0) + { + const auto v = to_lower(get_arg_value()); + if (v == "1" || v == "yes" || v == "true" || v == "on") + { + options.flow_flash_attn = true; + options.has_flow_flash_attn = true; + } + else if (v == "0" || v == "no" || v == "false" || v == "off") + { + options.flow_flash_attn = false; + options.has_flow_flash_attn = true; + } + else + { + print_error_log("Error: invalid --flow-flash-attn value \"%s\". Use 0/1, yes/no, true/false, on/off.\n", v.c_str()); + return 1; + } + } else if (str_casecmp(arg, "--speed") == 0 || str_casecmp(arg, "-s") == 0) { auto value = get_arg_value(); @@ -1932,6 +2000,10 @@ int tool_entry(int argc, char** argv) cosyvoice_init_default_context_params(¶ms.base_params.base_params); params.base_params.base_params.n_max_seq = options.max_llm_len; params.base_params.base_params.llm_kv_cache_type = options.llm_kv_cache_type; + if (options.has_llm_flash_attn) + params.base_params.base_params.llm_use_flash_attn = options.llm_flash_attn; + if (options.has_flow_flash_attn) + params.base_params.base_params.flow_use_flash_attn = options.flow_flash_attn; params.base_params.n_workers = 1; if (options.interactive) { From 708662d99c1d100ced0092ae7cbb5cb721e828e2 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Thu, 16 Jul 2026 12:00:10 +0800 Subject: [PATCH 17/34] server: add LLM and Flow flash attention options to server and web UI --- tools/server/cosyvoice-server-backend.cpp | 6 +++++ tools/server/cosyvoice-server-webui.cpp | 23 ++++++++++++++++ tools/server/cosyvoice-server.cpp | 32 ++++++++++++++++++++++- tools/server/webui/cosyvoice-webui.html | 6 +++++ tools/server/webui/cosyvoice-webui.js | 8 ++++++ 5 files changed, 74 insertions(+), 1 deletion(-) diff --git a/tools/server/cosyvoice-server-backend.cpp b/tools/server/cosyvoice-server-backend.cpp index 5ac33f3..5cf7018 100644 --- a/tools/server/cosyvoice-server-backend.cpp +++ b/tools/server/cosyvoice-server-backend.cpp @@ -760,6 +760,12 @@ int cosyvoice_server_backend_run(server_runtime& runtime) runtime.has_llm_kv_cache_override ? "user override" : "default"); print_info_log(runtime.log_level, " llm_kv_cache_type : %s\n", kv_buf); } + { + cosyvoice_context_params_t ap; + cosyvoice_get_context_params(runtime.model_slots.front().get(), &ap); + print_info_log(runtime.log_level, " llm_use_flash_attn : %s\n", ap.llm_use_flash_attn ? "yes" : "no"); + print_info_log(runtime.log_level, " flow_use_flash_attn: %s\n", ap.flow_use_flash_attn ? "yes" : "no"); + } if (runtime.has_seed) print_info_log(runtime.log_level, " seed_strategy : server default=%u, request override allowed\n", runtime.seed); else diff --git a/tools/server/cosyvoice-server-webui.cpp b/tools/server/cosyvoice-server-webui.cpp index 255c0a4..30e762a 100644 --- a/tools/server/cosyvoice-server-webui.cpp +++ b/tools/server/cosyvoice-server-webui.cpp @@ -435,6 +435,8 @@ int cosyvoice_server_webui_run(server_runtime& runtime) json += ",\"k_cache_type\":\"" + kv_cache_type_to_string(actual_params.llm_k_cache_type) + "\""; json += ",\"v_cache_type\":\"" + kv_cache_type_to_string(actual_params.llm_v_cache_type) + "\""; json += ",\"buffer_policy\":\"" + std::string(inference_buffer_policy_to_string(actual_params.inference_buffer_policy)) + "\""; + json += ",\"llm_use_flash_attn\":" + std::string(actual_params.llm_use_flash_attn ? "true" : "false"); + json += ",\"flow_use_flash_attn\":" + std::string(actual_params.flow_use_flash_attn ? "true" : "false"); auto arch = cosyvoice_get_architecture(runtime.model_slots.front().get()); if (arch && *arch) json += ",\"model_arch\":\"" + std::string(arch) + "\""; @@ -1200,8 +1202,15 @@ int cosyvoice_server_webui_run(server_runtime& runtime) uint32_t v = body["max_llm_len"].get(); if (v > 0) context_params.n_max_seq = v; } + if (body.contains("llm_use_flash_attn")) + context_params.llm_use_flash_attn = body["llm_use_flash_attn"].get(); + if (body.contains("flow_use_flash_attn")) + context_params.flow_use_flash_attn = body["flow_use_flash_attn"].get(); // Load the model - check result FIRST, then store + uint32_t chunk_tokens_val = 0; + if (body.contains("chunk_tokens")) + chunk_tokens_val = body["chunk_tokens"].get(); auto loaded_ctx = cosyvoice_load_from_file_ext( model_path.c_str(), &context_params, @@ -1262,6 +1271,13 @@ int cosyvoice_server_webui_run(server_runtime& runtime) runtime.dit_kv_cache_length = cp.dit_kv_cache_length; } + // Apply chunk_tokens if specified + if (chunk_tokens_val > 0) + { + for (auto& slot : runtime.model_slots) + cosyvoice_set_chunk_tokens(slot.get(), chunk_tokens_val); + } + // Get default generation config cosyvoice_get_default_generation_config(runtime.model_slots.front().get(), &runtime.default_generation_config); for (auto& slot : runtime.model_slots) @@ -1423,6 +1439,13 @@ int cosyvoice_server_webui_run(server_runtime& runtime) const auto speakers_str = join_strings(runtime.voice_names, ", "); print_info_log(runtime.log_level, " speakers : %s\n", speakers_str.empty() ? "-" : speakers_str.c_str()); } + if (!runtime.model_slots.empty()) + { + cosyvoice_context_params_t ap; + cosyvoice_get_context_params(runtime.model_slots.front().get(), &ap); + print_info_log(runtime.log_level, " llm_use_flash_attn : %s\n", ap.llm_use_flash_attn ? "yes" : "no"); + print_info_log(runtime.log_level, " flow_use_flash_attn: %s\n", ap.flow_use_flash_attn ? "yes" : "no"); + } print_info_log(runtime.log_level, " stream : %s\n", runtime.stream ? "enabled" : "disabled"); if (runtime.has_chunk_tokens) print_info_log(runtime.log_level, " chunk_tokens : %u\n", runtime.chunk_tokens); diff --git a/tools/server/cosyvoice-server.cpp b/tools/server/cosyvoice-server.cpp index 3222054..4f94f39 100644 --- a/tools/server/cosyvoice-server.cpp +++ b/tools/server/cosyvoice-server.cpp @@ -90,6 +90,10 @@ struct server_options bool fade_in_enabled = true; bool verbose = false; bool quiet = false; + bool has_llm_flash_attn = false; + bool llm_flash_attn = true; + bool has_flow_flash_attn = false; + bool flow_flash_attn = true; }; static void print_usage(const char* argv0) @@ -144,7 +148,9 @@ static void print_usage(const char* argv0) printf(" --min-token-text-ratio Minimum token/text ratio (>= 0).\n"); printf(" --max-token-text-ratio Maximum token/text ratio (>= 0).\n"); printf(" --verbose, -v Verbose logs.\n"); - printf(" --quiet, -q Suppress non-error logs.\n"); + printf(" --quiet, -q Suppress non-error logs.\n"); + printf(" --llm-flash-attn <0|1> Enable/disable LLM flash attention. Default: 1.\n"); + printf(" --flow-flash-attn <0|1> Enable/disable Flow/DiT flash attention. Default: 1.\n"); #ifndef COSYVOICE_NO_ICU printf("\nText normalization:\n"); @@ -248,6 +254,10 @@ static bool init_model_context(const server_options& options, ggml_backend_t bac context_params.inference_buffer_policy = options.inference_buffer_policy; context_params.n_max_seq = options.max_llm_len; context_params.llm_kv_cache_type = options.llm_kv_cache_type; + if (options.has_llm_flash_attn) + context_params.llm_use_flash_attn = options.llm_flash_attn; + if (options.has_flow_flash_attn) + context_params.flow_use_flash_attn = options.flow_flash_attn; context_params.dit_kv_cache_type = options.dit_kv_cache_type; context_params.dit_kv_fixed_slots = options.dit_kv_fixed_slots; context_params.dit_kv_offloadable_slots = options.dit_kv_offloadable_slots; @@ -928,6 +938,26 @@ int tool_entry(int argc, char** argv) options.fast_split_text_enabled = false; else if (str_casecmp(arg, "--disable-fade-in") == 0) options.fade_in_enabled = false; + else if (str_casecmp(arg, "--llm-flash-attn") == 0) + { + const auto v = to_lower(get_arg_value()); + if (v == "1" || v == "yes" || v == "true" || v == "on") + { options.llm_flash_attn = true; options.has_llm_flash_attn = true; } + else if (v == "0" || v == "no" || v == "false" || v == "off") + { options.llm_flash_attn = false; options.has_llm_flash_attn = true; } + else + { fprintf(stderr, "Error: invalid --llm-flash-attn value \"%s\". Use 0/1, yes/no, true/false, on/off.\n", v.c_str()); return 1; } + } + else if (str_casecmp(arg, "--flow-flash-attn") == 0) + { + const auto v = to_lower(get_arg_value()); + if (v == "1" || v == "yes" || v == "true" || v == "on") + { options.flow_flash_attn = true; options.has_flow_flash_attn = true; } + else if (v == "0" || v == "no" || v == "false" || v == "off") + { options.flow_flash_attn = false; options.has_flow_flash_attn = true; } + else + { fprintf(stderr, "Error: invalid --flow-flash-attn value \"%s\". Use 0/1, yes/no, true/false, on/off.\n", v.c_str()); return 1; } + } else { const auto option = arg; diff --git a/tools/server/webui/cosyvoice-webui.html b/tools/server/webui/cosyvoice-webui.html index 77c5273..e2d5600 100644 --- a/tools/server/webui/cosyvoice-webui.html +++ b/tools/server/webui/cosyvoice-webui.html @@ -124,6 +124,9 @@

Model Management

+
+ +
DiT
@@ -144,6 +147,9 @@

Model Management

+
+ +
diff --git a/tools/server/webui/cosyvoice-webui.js b/tools/server/webui/cosyvoice-webui.js index f377fa5..a1b66e0 100644 --- a/tools/server/webui/cosyvoice-webui.js +++ b/tools/server/webui/cosyvoice-webui.js @@ -31,6 +31,7 @@ const ADV_PARAM_IDS = [ 'model-max-llm','model-kv-k','model-kv-v','model-buffer-policy', 'model-threads','model-backend', 'model-dit-kv-k','model-dit-kv-v','model-dit-fixed-slots','model-dit-offloadable-slots','model-dit-cache-length', + 'model-llm-flash-attn','model-flow-flash-attn', 'tts-stream','tts-chunk-tokens' ]; @@ -342,10 +343,14 @@ function updateModelUI() { const kv_v = statusData.v_cache_type || '?'; const buf = statusData.buffer_policy || '?'; const mll = statusData.max_llm_len || '?'; + const llmFattn = statusData.llm_use_flash_attn !== undefined ? (statusData.llm_use_flash_attn ? 'yes' : 'no') : '?'; + const flowFattn = statusData.flow_use_flash_attn !== undefined ? (statusData.flow_use_flash_attn ? 'yes' : 'no') : '?'; els['model-loaded-info'].innerHTML = '' + escapeHtml(arch) + '
' + 'KV Cache: K=' + kv_k + ', V=' + kv_v + '
' + 'Buffer Policy: ' + buf + '
' + 'Max LLM Length: ' + mll + '
' + + 'LLM Flash Attn: ' + llmFattn + '
' + + 'Flow Flash Attn: ' + flowFattn + '
' + 'Sample Rate: ' + (statusData.sample_rate || '?') + ' Hz'; } } @@ -489,6 +494,9 @@ function initModelLoad() { const ml = parseInt(els['model-max-llm'].value, 10); if (ml > 0) body.max_llm_len = ml; + if (els['model-llm-flash-attn']) body.llm_use_flash_attn = els['model-llm-flash-attn'].checked; + if (els['model-flow-flash-attn']) body.flow_use_flash_attn = els['model-flow-flash-attn'].checked; + const dkt = els['model-dit-kv-k'].value; const dvt = els['model-dit-kv-v'].value; body.dit_kv_cache_type = (dkt === dvt) ? dkt : 'k=' + dkt + ',v=' + dvt; From b0df7144dbc3e6f5f1b0690aad021e8db0efb809 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Thu, 16 Jul 2026 12:10:02 +0800 Subject: [PATCH 18/34] cosyvoice: fix null deref in non-split and single-fragment streaming paths --- src/cosyvoice-tts.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index e90d6a9..d58dc3a 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -704,14 +704,16 @@ struct cosyvoice_tts_context : cosyvoice_tokenization_result_impl, cosyvoice_pro if (!(flags & COSYVOICE_TTS_FLAG_SPLIT_TEXT)) { ctx->tokenize(effective_text, this, true); - return cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, result); + return result ? cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, result) + : cosyvoice_tts_stream_with_postprocess(get_tokens(), get_n_tokens(), speed, callback, user_data); } const auto fragments = cosyvoice_internal::split_into_fragments(effective_text); if (fragments.size() <= 1) { ctx->tokenize(effective_text, this, true); - return cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, result); + return result ? cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, result) + : cosyvoice_tts_stream_with_postprocess(get_tokens(), get_n_tokens(), speed, callback, user_data); } // Compute the maximum text-token budget per chunk: From 727ef31e364e9ea3b996706756c37852740bf85d Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Thu, 16 Jul 2026 14:37:52 +0800 Subject: [PATCH 19/34] docs: add streaming TTS and DiT KV cache API documentation --- README.md | 41 ++++ README_zh.md | 41 ++++ docs/API_cosyvoice.md | 329 ++++++++++++++++++++++++++++---- docs/API_cosyvoice_interface.md | 180 +++++++++++++++++ docs/API_cosyvoice_lowlevel.md | 173 +++++++++++++++++ docs/TOOLS.md | 117 ++++++++++-- docs/TOOLS_zh.md | 196 +++++++++++++------ 7 files changed, 974 insertions(+), 103 deletions(-) diff --git a/README.md b/README.md index 07f115a..fda4219 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,8 @@ C++/GGML port of the Python CosyVoice inference pipeline released by the origina This repository ships independent engineering work and does not contain official support commitments. +Supports **zero-shot**, **instruct**, and **cross-lingual** TTS modes with both synchronous and **streaming** output. The frontend pipeline (speech tokenizer + speaker embedding) handles reference audio processing, or pre-encoded `prompt_speech` can be reused across sessions to skip the ONNX frontend entirely. + This project provides: - A core C/C++ inference library (`cosyvoice`) - A CLI synthesis tool (`cosyvoice-cli`) @@ -51,6 +53,10 @@ This project provides: | **Interactive REPL** | CLI interactive mode with slash commands for play, save, list, query, and seed control | | **Concurrent Serving** | Server `--concurrency` for parallel request handling | | **Model Quantization** | Quantize GGUF models to smaller formats (Q2_K through F16) with the built-in `quantize` tool | +| **Streaming TTS** | Real-time speech generation with low-latency audio delivery via callback — delivers audio chunks as they are synthesized, before the full utterance completes | +| **DiT KV Cache** | Avoid redundant attention recomputation across diffusion steps during streaming — configurable with fixed (device), offloadable (CPU), and uncached slot categories to trade memory vs. speed | +| **Flash Attention** | LLM and Flow flash attention support (`--llm-flash-attn`, `--flow-flash-attn`) for reduced memory and faster inference when the backend supports it | +| **Chunk Tokens Control** | Tune streaming latency vs. overhead tradeoff via `--chunk-tokens` — smaller chunks reduce first-chunk latency, larger chunks reduce RTF | | **KV Cache Quantization** | Reduce LLM memory usage via `--llm-kv-cache-type` (f32 / f16 / q8_0 / q5_1 / q4_0 / ...). Supports asymmetric quantization with separate K/V types (e.g. `k=q8_0,v=q4_0`). | | **Prompt Speech Reuse** | Pre-encode reference voice once, reuse across multiple synthesis runs — no ONNX overhead | | **Audio Backend Plugins** | Choose MINIAUDIO (default) or FFMPEG for multi-format encoding (WAV, MP3, AAC, FLAC, OPUS, M4A) | @@ -60,6 +66,41 @@ This project provides: | **Multiple Backends** | CPU, CUDA, Metal, SYCL (see [Backend Test Status](#backend-test-status)) | | **Cross-Platform** | Windows (x64), Linux (x86_64), macOS (arm64) — all tested in CI | +## Streaming TTS & DiT KV Cache + +Streaming TTS delivers audio chunks incrementally via a callback function as they are synthesized, without waiting for the full utterance to complete. This enables real-time playback and lower perceived latency. + +The streaming pipeline introduces a **DiT KV cache** to avoid redundant computation. During non-streaming inference, the DiT module runs 10 diffusion steps, each computing self-attention over the full audio sequence — resulting in 10× attention recomputation. The KV cache stores intermediate key/value tensors across steps so that each position is computed only once. + +### Slot Organization + +The DiT KV cache is organized into **slots**, where each slot holds the KV cache for one diffusion step. With the default 10 steps, there can be at most 10 slots. + +Slots fall into three categories: + +| Category | Memory | Behavior | +|----------|--------|----------| +| **Fixed** | Stays on device (GPU) | Fastest; never offloaded | +| **Offloadable** | Offloaded to CPU when not in use | Saves device memory at the cost of transfer latency | +| **Uncached** | Not stored at all | Full attention recomputation every step, no extra memory | + +Total slots = `fixed + offloadable`. Remaining steps (10 − total) use full recomputation. + +The cache is large, so the default is 0 slots (all 10 steps fully recomputed). When enabled and the sequence exceeds the configured cache length, some positions are discarded — inference continues normally but output quality may degrade. + +> The DiT KV cache is **only used during streaming TTS**; non-streaming calls ignore it and always perform full recomputation. + +### Configuration + +DiT KV cache parameters are configured via `cosyvoice_context_params_v3_t` (C API) or the respective CLI/server `--dit-kv-*` flags: + +- `--dit-kv-type` / `dit_kv_cache_type`: Storage format (f32/f16/q8_0/...) for the DiT KV cache. +- `--dit-kv-fixed-slots` / `dit_kv_fixed_slots`: Number of device-resident slots. +- `--dit-kv-offloadable-slots` / `dit_kv_offloadable_slots`: Number of CPU-offloadable slots. +- `--dit-kv-cache-length` / `dit_kv_cache_length`: Maximum sequence positions kept in the cache. + +Streaming is enabled via `--stream` flag on CLI/server. Chunk granularity is controlled by `--chunk-tokens` (default: 128 tokens per chunk). + ## Pre-converted Models Download ready-to-use GGUF models (no conversion needed): diff --git a/README_zh.md b/README_zh.md index 8d99807..873d333 100644 --- a/README_zh.md +++ b/README_zh.md @@ -15,6 +15,8 @@ 本仓库仅提供独立社区实现，不包含任何官方支持承诺。 +支持 **zero-shot**、**instruct** 和 **cross-lingual** 三种 TTS 模式，同时提供同步输出与**流式输出**（streaming）。前端流水线（speech tokenizer + speaker embedding）可处理参考音频，也可复用预编码的 `prompt_speech` 跳过 ONNX 前端以提升性能。 + 本项目提供： - 核心 C/C++ 推理库（`cosyvoice`） - 命令行合成工具（`cosyvoice-cli`） @@ -51,6 +53,10 @@ | **交互式 REPL** | CLI 交互模式，支持 /play、/save、/list、/query、/seed 等斜杠命令 | | **并发服务** | Server 的 `--concurrency` 参数，支持并行请求处理 | | **模型量化** | 内置 `quantize` 工具，支持 Q2_K 到 F16 多种量化格式 | +| **流式 TTS** | 实时语音生成，通过回调逐段交付音频——在完整语句尚未合成完毕前即可开始播放 | +| **DiT KV 缓存** | 在流式推理过程中，跨扩散步复用注意力 KV 缓存，避免冗余计算——可配置固定（设备内存）、可卸载（CPU）、不缓存三种槽位类型 | +| **Flash Attention** | LLM 与 Flow 模块均支持 flash attention（`--llm-flash-attn`、`--flow-flash-attn`），后端支持时可降低显存并加速推理 | +| **分块 Token 控制** | 通过 `--chunk-tokens` 调节流式推理的延迟与开销平衡——较小分块降低首块延迟，较大分块降低 RTF | | **KV Cache 量化** | 通过 `--llm-kv-cache-type` 降低 LLM 内存占用（f32 / f16 / q8_0 / q5_1 / q4_0 / ...）。支持非对称量化，K 和 V 可独立指定类型（如 `k=q8_0,v=q4_0`）。 | | **Prompt Speech 复用** | 一次编码参考音色，后续合成直接复用，无需再跑 ONNX | | **音频后端可切换** | 可选 MINIAUDIO（默认）或 FFMPEG，支持 WAV、MP3、AAC、FLAC、OPUS、M4A | @@ -60,6 +66,41 @@ | **多后端支持** | CPU、CUDA、Metal、SYCL（见[后端测试情况](#后端测试情况)） | | **跨平台** | Windows (x64)、Linux (x86_64)、macOS (arm64) — 均在 CI 中测试 | +## 流式 TTS 与 DiT KV 缓存 + +流式 TTS 在合成过程中通过回调函数逐段交付音频，无需等待完整语句生成完毕即可开始播放，从而实现实时播放与更低的主观延迟。 + +流式流水线引入了 **DiT KV 缓存** 以避免冗余计算。非流式推理时，DiT 模块运行 10 步扩散，每步对完整音频序列计算自注意力——总共会执行 10 次注意力重算。KV 缓存跨扩散步存储中间 key/value 张量，使每个位置只需计算一次。 + +### 槽位组织 + +DiT KV 缓存按 **槽位（slot）** 组织，每个槽位对应一个扩散步的 KV 缓存。默认 10 个扩散步意味着最多 10 个槽位。 + +槽位分为三类： + +| 类别 | 内存位置 | 行为 | +|----------|--------|------| +| **固定** | 常驻设备（GPU） | 最快，从不卸载 | +| **可卸载** | 不使用时卸载到 CPU | 节省设备显存，但增加传输开销 | +| **不缓存** | 不存储 | 每步全量重算注意力，无额外内存开销 | + +总槽位数 = `固定 + 可卸载`。剩余步（10 − 总槽位）使用全量重算。 + +KV 缓存占用较大，因此默认 0 个槽位（全部 10 步全量重算）。启用缓存后，若序列长度超过配置的缓存长度，会丢弃部分位置——推理可正常继续，但输出质量可能下降。 + +> DiT KV 缓存**仅在流式 TTS 时使用**；非流式调用始终执行全量重算，忽略此缓存。 + +### 配置 + +DiT KV 缓存参数通过 `cosyvoice_context_params_v3_t`（C API）或 CLI/server 的 `--dit-kv-*` 参数配置： + +- `--dit-kv-type` / `dit_kv_cache_type`：DiT KV 缓存存储格式（f32/f16/q8_0/...）。 +- `--dit-kv-fixed-slots` / `dit_kv_fixed_slots`：常驻设备内存的槽位数。 +- `--dit-kv-offloadable-slots` / `dit_kv_offloadable_slots`：可卸载到 CPU 的槽位数。 +- `--dit-kv-cache-length` / `dit_kv_cache_length`：缓存保留的最大序列位置数。 + +流式输出通过 `--stream` 标志启用。分块粒度由 `--chunk-tokens` 控制（默认 128 tokens/块）。 + ## 预转换模型 下载即用的 GGUF 模型（无需自行转换）： diff --git a/docs/API_cosyvoice.md b/docs/API_cosyvoice.md index 08b1d08..d8578ea 100644 --- a/docs/API_cosyvoice.md +++ b/docs/API_cosyvoice.md @@ -29,42 +29,58 @@ Controls symbol import/export visibility for the C API. Define `COSYVOICE_STATIC On Windows, the macro expands to `__declspec(dllimport)` by default. On non-Windows targets, default symbol visibility is used. -## cosyvoice_llm_kv_cache_type_t +## cosyvoice_kv_cache_type_t ### Syntax ```c -typedef enum cosyvoice_llm_kv_cache_type +typedef enum cosyvoice_kv_cache_type { - COSYVOICE_LLM_KV_CACHE_TYPE_F32, - COSYVOICE_LLM_KV_CACHE_TYPE_F16, - COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0, - COSYVOICE_LLM_KV_CACHE_TYPE_Q5_1, - COSYVOICE_LLM_KV_CACHE_TYPE_Q5_0, - COSYVOICE_LLM_KV_CACHE_TYPE_Q4_1, - COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0, - COSYVOICE_LLM_KV_CACHE_TYPE_COUNT -} cosyvoice_llm_kv_cache_type_t; + COSYVOICE_KV_CACHE_TYPE_F32, + COSYVOICE_KV_CACHE_TYPE_F16, + COSYVOICE_KV_CACHE_TYPE_Q8_0, + COSYVOICE_KV_CACHE_TYPE_Q5_1, + COSYVOICE_KV_CACHE_TYPE_Q5_0, + COSYVOICE_KV_CACHE_TYPE_Q4_1, + COSYVOICE_KV_CACHE_TYPE_Q4_0, + COSYVOICE_KV_CACHE_TYPE_COUNT +} cosyvoice_kv_cache_type_t, cosyvoice_llm_kv_cache_type_t; +``` + +#### Backward Compatibility Aliases + +The old `cosyvoice_llm_kv_cache_type_t` and `COSYVOICE_LLM_KV_CACHE_TYPE_*` constants +are provided as backward-compatibility aliases: + +```c +#define COSYVOICE_LLM_KV_CACHE_TYPE_F32 COSYVOICE_KV_CACHE_TYPE_F32 +#define COSYVOICE_LLM_KV_CACHE_TYPE_F16 COSYVOICE_KV_CACHE_TYPE_F16 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0 COSYVOICE_KV_CACHE_TYPE_Q8_0 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q5_1 COSYVOICE_KV_CACHE_TYPE_Q5_1 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q5_0 COSYVOICE_KV_CACHE_TYPE_Q5_0 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q4_1 COSYVOICE_KV_CACHE_TYPE_Q4_1 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0 COSYVOICE_KV_CACHE_TYPE_Q4_0 +#define COSYVOICE_LLM_KV_CACHE_TYPE_COUNT COSYVOICE_KV_CACHE_TYPE_COUNT ``` ### Description -Specifies supported KV-cache storage formats for the LLM module. +Specifies supported KV-cache storage formats. Used by both the LLM and DiT modules. ### Values -- `COSYVOICE_LLM_KV_CACHE_TYPE_F32`: 32-bit floating-point cache. -- `COSYVOICE_LLM_KV_CACHE_TYPE_F16`: 16-bit floating-point cache. -- `COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0`: GGML `Q8_0` quantized cache. -- `COSYVOICE_LLM_KV_CACHE_TYPE_Q5_1`: GGML `Q5_1` quantized cache. -- `COSYVOICE_LLM_KV_CACHE_TYPE_Q5_0`: GGML `Q5_0` quantized cache. -- `COSYVOICE_LLM_KV_CACHE_TYPE_Q4_1`: GGML `Q4_1` quantized cache. -- `COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0`: GGML `Q4_0` quantized cache. -- `COSYVOICE_LLM_KV_CACHE_TYPE_COUNT`: Sentinel value, not a runtime mode. +- `COSYVOICE_KV_CACHE_TYPE_F32`: 32-bit floating-point cache. +- `COSYVOICE_KV_CACHE_TYPE_F16`: 16-bit floating-point cache. +- `COSYVOICE_KV_CACHE_TYPE_Q8_0`: GGML `Q8_0` quantized cache. +- `COSYVOICE_KV_CACHE_TYPE_Q5_1`: GGML `Q5_1` quantized cache. +- `COSYVOICE_KV_CACHE_TYPE_Q5_0`: GGML `Q5_0` quantized cache. +- `COSYVOICE_KV_CACHE_TYPE_Q4_1`: GGML `Q4_1` quantized cache. +- `COSYVOICE_KV_CACHE_TYPE_Q4_0`: GGML `Q4_0` quantized cache. +- `COSYVOICE_KV_CACHE_TYPE_COUNT`: Sentinel value, not a runtime mode. ### Separate K/V Cache Macros -The `cosyvoice_llm_kv_cache_type_t` value can encode separate types for the K and V cache +The `cosyvoice_kv_cache_type_t` value can encode separate types for the K and V cache buffers via bit-packing. This allows using different quantization formats for the K and V tensors (e.g. K in `Q8_0` and V in `Q4_0`) to trade off quality vs. memory. @@ -78,7 +94,7 @@ tensors (e.g. K in `Q8_0` and V in `Q4_0`) to trade off quality vs. memory. | 15–30 | Reserved | | 31 | Separate-K/V flag | -When bit 31 is 0, the value is a plain `cosyvoice_llm_kv_cache_type_t` that applies +When bit 31 is 0, the value is a plain `cosyvoice_kv_cache_type_t` that applies to both K and V (backward compatible). ```c @@ -100,12 +116,13 @@ tried first; if that also fails, auto-fallback applies. #### Application in parameters -Use the bitfield members of `cosyvoice_context_params_t` directly, or pack a value -with `COSYVOICE_MAKE_SEPARATE_KV_CACHE` and assign it to `llm_kv_cache_type`: +Use the bitfield members of `cosyvoice_context_params_t` or `cosyvoice_context_params_v3_t` +directly, or pack a value with `COSYVOICE_MAKE_SEPARATE_KV_CACHE` and assign it to +`llm_kv_cache_type` / `dit_kv_cache_type`: ```c -params.llm_k_cache_type = COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0; -params.llm_v_cache_type = COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0; +params.llm_k_cache_type = COSYVOICE_KV_CACHE_TYPE_Q8_0; +params.llm_v_cache_type = COSYVOICE_KV_CACHE_TYPE_Q4_0; params.llm_kv_cache_separate_buffers = true; ``` @@ -113,9 +130,9 @@ or equivalently: ```c params.llm_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( - COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0, - COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0, - COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0); + COSYVOICE_KV_CACHE_TYPE_Q8_0, + COSYVOICE_KV_CACHE_TYPE_Q4_0, + COSYVOICE_KV_CACHE_TYPE_Q8_0); ``` ## cosyvoice_inference_buffer_policy_t @@ -356,13 +373,13 @@ typedef struct cosyvoice_context_params { struct { - cosyvoice_llm_kv_cache_type_t llm_k_cache_type : 5; ///< K cache data type. - cosyvoice_llm_kv_cache_type_t llm_v_cache_type : 5; ///< V cache data type. - cosyvoice_llm_kv_cache_type_t llm_kv_cache_fallback : 5; ///< Fallback type when preferred type unsupported. - cosyvoice_llm_kv_cache_type_t : 16; ///< Reserved. - cosyvoice_llm_kv_cache_type_t llm_kv_cache_separate_buffers : 1; ///< Allocate separate K & V buffers. + cosyvoice_kv_cache_type_t llm_k_cache_type : 5; ///< K cache data type. + cosyvoice_kv_cache_type_t llm_v_cache_type : 5; ///< V cache data type. + cosyvoice_kv_cache_type_t llm_kv_cache_fallback : 5; ///< Fallback type when preferred type unsupported. + cosyvoice_kv_cache_type_t : 16; ///< Reserved. + uint32_t llm_kv_cache_separate_buffers : 1; ///< Allocate separate K & V buffers. }; - cosyvoice_llm_kv_cache_type_t llm_kv_cache_type; ///< Backward-compatible unified type. + cosyvoice_kv_cache_type_t llm_kv_cache_type; ///< Backward-compatible unified type. }; bool llm_allow_kv_cache_fallback; cosyvoice_inference_buffer_policy_t inference_buffer_policy; @@ -549,6 +566,117 @@ Loads a model context with extended parameters, including worker-count configura Loaded context handle on success; `NULL` on failure. +## cosyvoice_context_params_v3_t + +### Syntax + +```c +typedef struct cosyvoice_context_params_v3 +{ + cosyvoice_context_params_v2_t base_params; + + union + { + struct + { +#if defined(__BYTE_ORDER__) && (__BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) || defined(_BYTE_ORDER) && (_BYTE_ORDER == _BIG_ENDIAN) || defined(__BIG_ENDIAN__) && !defined(__LITTLE_ENDIAN__) || defined(__ARMEB__) || defined(__MIPSEB__) || defined(__sparc__) + uint32_t dit_kv_cache_separate_buffers : 1; + cosyvoice_kv_cache_type_t : 16; + cosyvoice_kv_cache_type_t dit_kv_cache_fallback : 5; + cosyvoice_kv_cache_type_t dit_v_cache_type : 5; + cosyvoice_kv_cache_type_t dit_k_cache_type : 5; +#else + cosyvoice_kv_cache_type_t dit_k_cache_type : 5; + cosyvoice_kv_cache_type_t dit_v_cache_type : 5; + cosyvoice_kv_cache_type_t dit_kv_cache_fallback : 5; + cosyvoice_kv_cache_type_t : 16; + uint32_t dit_kv_cache_separate_buffers : 1; +#endif + }; + cosyvoice_kv_cache_type_t dit_kv_cache_type; + }; + bool dit_allow_kv_cache_fallback; + uint32_t dit_kv_fixed_slots; + uint32_t dit_kv_offloadable_slots; + uint32_t dit_kv_cache_length; +} cosyvoice_context_params_v3_t; + +#ifdef __cplusplus +struct cosyvoice_context_params_v3_cpp : cosyvoice_context_params_v2_cpp +{ + union + { + struct + { +#if defined(__BYTE_ORDER__) && (__BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) || defined(_BYTE_ORDER) && (_BYTE_ORDER == _BIG_ENDIAN) || defined(__BIG_ENDIAN__) && !defined(__LITTLE_ENDIAN__) || defined(__ARMEB__) || defined(__MIPSEB__) || defined(__sparc__) + uint32_t dit_kv_cache_separate_buffers : 1; + cosyvoice_kv_cache_type_t : 16; + cosyvoice_kv_cache_type_t dit_kv_cache_fallback : 5; + cosyvoice_kv_cache_type_t dit_v_cache_type : 5; + cosyvoice_kv_cache_type_t dit_k_cache_type : 5; +#else + cosyvoice_kv_cache_type_t dit_k_cache_type : 5; + cosyvoice_kv_cache_type_t dit_v_cache_type : 5; + cosyvoice_kv_cache_type_t dit_kv_cache_fallback : 5; + cosyvoice_kv_cache_type_t : 16; + uint32_t dit_kv_cache_separate_buffers : 1; +#endif + }; + cosyvoice_kv_cache_type_t dit_kv_cache_type; + }; + bool dit_allow_kv_cache_fallback; + uint32_t dit_kv_fixed_slots; + uint32_t dit_kv_offloadable_slots; + uint32_t dit_kv_cache_length; +}; +#endif +``` + +### Description + +Extends `cosyvoice_context_params_v2_t` with DiT (diffusion) KV cache configuration. The DiT module runs multiple diffusion steps during streaming TTS — each step computes self-attention over the full audio sequence. A KV cache can avoid redundant attention recomputation across steps, but the cache is very large (up to `sequence_length × n_diffusion_steps` key-value pairs). + +### Fields + +- `base_params`: V2 base parameters. +- `dit_k_cache_type`: Data type of the K cache in the DiT module. +- `dit_v_cache_type`: Data type of the V cache in the DiT module. +- `dit_kv_cache_separate_buffers`: If true, allocate separate buffers for K and V. +- `dit_kv_cache_fallback`: Fallback type when preferred K/V type is unsupported. +- `dit_kv_cache_type`: Shorthand — assigns a unified type (no separate K/V). +- `dit_allow_kv_cache_fallback`: If true, fall back to a flash-attention-compatible type. +- `dit_kv_fixed_slots`: Number of fixed (device memory, never offloaded) DiT KV slots. Each slot holds the KV cache for one diffusion step. +- `dit_kv_offloadable_slots`: Number of offloadable (CPU offload) DiT KV slots. +- `dit_kv_cache_length`: Maximum sequence length for the DiT KV cache. 0 to use default (`n_max_seq × 10`). + +### DiT KV Cache Concept + +See [README.md — Streaming TTS](#streaming-tts--dit-kv-cache) for a detailed explanation of the DiT KV cache concept. + +## cosyvoice_load_from_file_with_params_v3 + +### Syntax + +```c +COSYVOICE_API cosyvoice_context_t cosyvoice_load_from_file_with_params_v3( + const char* filename, + const cosyvoice_context_params_v3_t* params +); +``` + +### Description + +Loads a model context with V3 extended parameters, including DiT KV cache configuration. + +### Parameters + +- `filename`: Path to the model file. +- `params`: V3 context-parameter block. + +### Returns + +Loaded context handle on success; `NULL` on failure. + ## cosyvoice_duplicate_context ### Syntax @@ -1466,6 +1594,120 @@ Generates speech in cross-lingual mode. `true` on success; otherwise `false`. +## cosyvoice_tts_audio_callback_t + +### Syntax + +```c +typedef bool (*cosyvoice_tts_audio_callback_t)(const float* audio, uint32_t n_samples, void* user_data); +``` + +### Description + +Callback invoked by the streaming TTS API for each generated audio chunk. + +### Parameters + +- `audio`: PCM samples in 32-bit floating point format, 1 channel. +- `n_samples`: Number of samples in this chunk. +- `user_data`: Opaque context passed when registering the callback. + +### Returns + +`true` to continue streaming, `false` to abort synthesis. + +## cosyvoice_tts_zero_shot_stream + +### Syntax + +```c +COSYVOICE_API bool cosyvoice_tts_zero_shot_stream( + cosyvoice_tts_context_t ctx, + const char* text, + float speed, + cosyvoice_tts_audio_callback_t callback, + void* user_data +); +``` + +### Description + +Generates speech with streaming output in zero-shot mode. Audio chunks are delivered incrementally via the callback. + +### Parameters + +- `ctx`: TTS context handle. +- `text`: Input text. +- `speed`: Speech speed multiplier. +- `callback`: Callback receiving each audio chunk. +- `user_data`: Opaque context passed to the callback. + +### Returns + +`true` on success; otherwise `false`. + +## cosyvoice_tts_instruct_stream + +### Syntax + +```c +COSYVOICE_API bool cosyvoice_tts_instruct_stream( + cosyvoice_tts_context_t ctx, + const char* text, + const char* instruction, + float speed, + cosyvoice_tts_audio_callback_t callback, + void* user_data +); +``` + +### Description + +Generates speech with streaming output in instruct mode. + +### Parameters + +- `ctx`: TTS context handle. +- `text`: Input text. +- `instruction`: Instruction text guiding style or behavior. +- `speed`: Speech speed multiplier. +- `callback`: Callback receiving each audio chunk. +- `user_data`: Opaque context passed to the callback. + +### Returns + +`true` on success; otherwise `false`. + +## cosyvoice_tts_cross_lingual_stream + +### Syntax + +```c +COSYVOICE_API bool cosyvoice_tts_cross_lingual_stream( + cosyvoice_tts_context_t ctx, + const char* text, + float speed, + cosyvoice_tts_audio_callback_t callback, + void* user_data +); +``` + +### Description + +Generates speech with streaming output in cross-lingual mode. + +### Parameters + +- `ctx`: TTS context handle. +- `text`: Input text. +- `speed`: Speech speed multiplier. +- `callback`: Callback receiving each audio chunk. +- `user_data`: Opaque context passed to the callback. + +### Returns + +`true` on success; otherwise `false`. + ## cosyvoice_save_wav ### Syntax @@ -1541,6 +1783,23 @@ Retrieves a snapshot of current memory usage for the context. - `ctx`: Context handle. - `usage`: Output structure receiving usage values. +## cosyvoice_get_total_memory_usage + +### Syntax + +```c +COSYVOICE_API void cosyvoice_get_total_memory_usage(cosyvoice_context_t ctx, cosyvoice_memory_usage_t* usage); +``` + +### Description + +Retrieves a snapshot of total memory usage across all workers in the context. + +### Parameters + +- `ctx`: Context handle. +- `usage`: Output structure receiving usage values. + ## cosyvoice_empty_buffer_cache ### Syntax diff --git a/docs/API_cosyvoice_interface.md b/docs/API_cosyvoice_interface.md index 7a99ab5..a080bf3 100644 --- a/docs/API_cosyvoice_interface.md +++ b/docs/API_cosyvoice_interface.md @@ -28,6 +28,7 @@ struct cosyvoice_model_context virtual cosyvoice_builtin_sampler_rng_policy_t get_builtin_sampler_rng_policy() = 0; virtual bool set_builtin_sampler_rng_policy(cosyvoice_builtin_sampler_rng_policy_t policy) = 0; virtual bool set_sampler_seed(uint32_t seed) = 0; + virtual uint32_t get_sampler_seed() = 0; virtual bool llm_prefill(ggml_type type, const void* data, uint32_t seq_len) = 0; virtual bool llm_decode(ggml_type type, const void* data) = 0; @@ -38,6 +39,8 @@ struct cosyvoice_model_context virtual uint32_t llm_get_kv_cache_len() = 0; virtual bool llm_set_kv_cache_len(uint32_t len) = 0; + virtual void llm_offload_kv_cache() = 0; + virtual void llm_load_kv_cache() = 0; virtual int llm_sample_token() = 0; virtual bool llm_is_stop_token(int token_id) = 0; @@ -55,6 +58,28 @@ struct cosyvoice_model_context cosyvoice_generated_speech_ptr result ) = 0; + virtual bool llm_job_ext( + const int* text, + uint32_t text_len, + cosyvoice_prompt_t prompt, + uint32_t max_new_tokens, + bool* final + ) = 0; + + virtual bool token2wav_ext( + const int* token_ids, + uint32_t n_tokens, + float speed, + cosyvoice_prompt_t prompt, + uint32_t* offset, + bool streaming, + bool finalize, + cosyvoice_generated_speech_ptr result + ) = 0; + + virtual uint32_t get_chunk_tokens() = 0; + virtual void set_chunk_tokens(uint32_t n_tokens) = 0; + virtual ggml_status get_last_status() = 0; virtual void set_prompt( @@ -65,6 +90,7 @@ struct cosyvoice_model_context ) = 0; virtual void get_memory_usage(cosyvoice_memory_usage_t* usage) = 0; + virtual void get_total_memory_usage(cosyvoice_memory_usage_t* usage) = 0; virtual void empty_buffer_cache() = 0; virtual void set_noise_callback(cosyvoice_noise_callback_t callback, void* callback_ctx) = 0; @@ -533,6 +559,38 @@ Trims KV-cache length. `true` on success; otherwise `false`. +## cosyvoice_model_context::llm_offload_kv_cache + +### Syntax + +```cpp +virtual void llm_offload_kv_cache() = 0; +``` + +### Description + +Offloads the KV cache to CPU memory. + +### Returns + +No return value. + +## cosyvoice_model_context::llm_load_kv_cache + +### Syntax + +```cpp +virtual void llm_load_kv_cache() = 0; +``` + +### Description + +Loads the KV cache from CPU memory back to the backend. + +### Returns + +No return value. + ## cosyvoice_model_context::llm_sample_token ### Syntax @@ -659,6 +717,36 @@ Runs low-level LLM generation. `true` on success; otherwise `false`. +## cosyvoice_model_context::llm_job_ext + +### Syntax + +```cpp +virtual bool llm_job_ext( + const int* text, + uint32_t text_len, + cosyvoice_prompt_t prompt, + uint32_t max_new_tokens, + bool* final +) = 0; +``` + +### Description + +Runs low-level LLM generation with additional options. + +### Parameters + +- `text`: Input text tokens. +- `text_len`: Number of tokens. +- `prompt`: Prompt handle. +- `max_new_tokens`: Maximum number of new tokens to generate. If 0, no new tokens are generated. +- `final`: Output parameter indicating whether generation is complete (`true`) or more tokens can be generated (`false`). + +### Returns + +`true` on success; otherwise `false`. + ## cosyvoice_model_context::token2wav ### Syntax @@ -689,6 +777,78 @@ Converts speech tokens to waveform. `true` on success; otherwise `false`. +## cosyvoice_model_context::token2wav_ext + +### Syntax + +```cpp +virtual bool token2wav_ext( + const int* token_ids, + uint32_t n_tokens, + float speed, + cosyvoice_prompt_t prompt, + uint32_t* offset, + bool streaming, + bool finalize, + cosyvoice_generated_speech_ptr result +) = 0; +``` + +### Description + +Converts speech tokens to waveform with additional options. + +### Parameters + +- `token_ids`: Speech token ids. +- `n_tokens`: Number of tokens. +- `speed`: Speed multiplier. +- `prompt`: Prompt handle. +- `offset`: Input/output token offset for incremental conversion. +- `streaming`: If true, convert incrementally. +- `finalize`: If true, flush and finalize output. +- `result`: Output waveform container. + +### Returns + +`true` on success; otherwise `false`. + +## cosyvoice_model_context::get_chunk_tokens + +### Syntax + +```cpp +virtual uint32_t get_chunk_tokens() = 0; +``` + +### Description + +Gets the number of tokens processed in each chunk during streaming inference. + +### Returns + +Current chunk token count. + +## cosyvoice_model_context::set_chunk_tokens + +### Syntax + +```cpp +virtual void set_chunk_tokens(uint32_t n_tokens) = 0; +``` + +### Description + +Sets the number of tokens processed in each chunk during streaming inference. + +### Parameters + +- `n_tokens`: Number of tokens per chunk. + +### Returns + +No return value. + ## cosyvoice_model_context::get_last_status ### Syntax @@ -753,6 +913,26 @@ Gets memory-usage snapshot. No return value. +## cosyvoice_model_context::get_total_memory_usage + +### Syntax + +```cpp +virtual void get_total_memory_usage(cosyvoice_memory_usage_t* usage) = 0; +``` + +### Description + +Gets total memory usage across all workers. + +### Parameters + +- `usage`: Output usage structure. + +### Returns + +No return value. + ## cosyvoice_model_context::empty_buffer_cache ### Syntax diff --git a/docs/API_cosyvoice_lowlevel.md b/docs/API_cosyvoice_lowlevel.md index 316da69..31026d5 100644 --- a/docs/API_cosyvoice_lowlevel.md +++ b/docs/API_cosyvoice_lowlevel.md @@ -395,6 +395,38 @@ Trims current KV-cache length. `len` must be less than or equal to the current length. +## cosyvoice_llm_offload_kv_cache + +### Syntax + +```c +COSYVOICE_API void cosyvoice_llm_offload_kv_cache(cosyvoice_context_t ctx); +``` + +### Description + +Offloads the LLM KV cache from device memory to CPU memory. + +### Parameters + +- `ctx`: Context handle. + +## cosyvoice_llm_load_kv_cache + +### Syntax + +```c +COSYVOICE_API void cosyvoice_llm_load_kv_cache(cosyvoice_context_t ctx); +``` + +### Description + +Loads the LLM KV cache from CPU memory back to the backend device. + +### Parameters + +- `ctx`: Context handle. + ## cosyvoice_llm_sample_token ### Syntax @@ -545,6 +577,38 @@ Runs low-level LLM generation for tokenized text and prompt. `true` on success; otherwise `false`. +## cosyvoice_llm_job_ext + +### Syntax + +```c +COSYVOICE_API bool cosyvoice_llm_job_ext( + cosyvoice_context_t ctx, + const int* text, + uint32_t text_len, + cosyvoice_prompt_t prompt, + uint32_t max_new_tokens, + bool* final +); +``` + +### Description + +Runs low-level LLM generation with additional options. + +### Parameters + +- `ctx`: Context handle. +- `text`: Text token ids. +- `text_len`: Number of tokens in `text`. +- `prompt`: Prompt handle. +- `max_new_tokens`: Maximum number of new tokens to generate. If 0, no new tokens are generated. +- `final`: Output parameter indicating whether the generation is complete (`true`) or more tokens can be generated (`false`). + +### Returns + +`true` on success; otherwise `false`. + ## cosyvoice_token2wav ### Syntax @@ -577,6 +641,44 @@ Converts speech tokens to waveform. `true` on success; otherwise `false`. +## cosyvoice_token2wav_ext + +### Syntax + +```c +COSYVOICE_API bool cosyvoice_token2wav_ext( + cosyvoice_context_t ctx, + const int* token_ids, + uint32_t n_tokens, + float speed, + cosyvoice_prompt_t prompt, + uint32_t* offset, + bool streaming, + bool finalize, + cosyvoice_generated_speech_ptr result +); +``` + +### Description + +Converts speech tokens to waveform with additional options. + +### Parameters + +- `ctx`: Context handle. +- `token_ids`: Speech token ids. +- `n_tokens`: Number of speech tokens. +- `speed`: Speech speed multiplier. +- `prompt`: Prompt handle. +- `offset`: Input/output token offset for incremental conversion. +- `streaming`: If true, convert incrementally (partial chunks). +- `finalize`: If true, flush and finalize the output. +- `result`: Output waveform container. + +### Returns + +`true` on success; otherwise `false`. + ## cosyvoice_tts ### Syntax @@ -613,6 +715,40 @@ Runs full low-level text-to-speech pipeline. Runs LLM generation and waveform conversion as a single convenience pipeline. +## cosyvoice_tts_stream + +### Syntax + +```c +COSYVOICE_API bool cosyvoice_tts_stream( + cosyvoice_context_t ctx, + const int* text, + uint32_t text_len, + float speed, + cosyvoice_prompt_t prompt, + cosyvoice_tts_audio_callback_t callback, + void* user_data +); +``` + +### Description + +Runs full low-level text-to-speech pipeline with streaming output. Audio chunks are delivered incrementally via the callback. + +### Parameters + +- `ctx`: Context handle. +- `text`: Text token ids. +- `text_len`: Number of text tokens. +- `speed`: Speech speed multiplier. +- `prompt`: Prompt handle. +- `callback`: Callback receiving each audio chunk. +- `user_data`: Opaque context passed to the callback. + +### Returns + +`true` on success; otherwise `false`. + ## cosyvoice_get_tokenizer ### Syntax @@ -884,6 +1020,26 @@ Gets currently registered noise callback and context. No return value. +## cosyvoice_get_chunk_tokens + +### Syntax + +```c +COSYVOICE_API uint32_t cosyvoice_get_chunk_tokens(cosyvoice_context_t ctx); +``` + +### Description + +Gets the number of tokens processed in each chunk during streaming inference. + +### Parameters + +- `ctx`: Context handle. + +### Returns + +Current chunk token count. + ## cosyvoice_get_hift_rand_ini_len ### Syntax @@ -904,6 +1060,23 @@ Gets required length of HiFT initialization noise buffer. Required sample count. +## cosyvoice_set_chunk_tokens + +### Syntax + +```c +COSYVOICE_API void cosyvoice_set_chunk_tokens(cosyvoice_context_t ctx, uint32_t n_tokens); +``` + +### Description + +Sets the number of tokens processed in each chunk during streaming inference. + +### Parameters + +- `ctx`: Context handle. +- `n_tokens`: Number of tokens per chunk. Smaller values reduce first-chunk latency but increase overhead; larger values reduce RTF but increase first-chunk latency. + ## cosyvoice_set_hift_rand_ini ### Syntax diff --git a/docs/TOOLS.md b/docs/TOOLS.md index 7535cc1..8014303 100644 --- a/docs/TOOLS.md +++ b/docs/TOOLS.md @@ -129,7 +129,7 @@ The WebUI provides a modern single-page application with the following capabilit #### Model Management - **Load a model** at runtime: enter a `.gguf` file path, select backend and threads, then click "Load Model". - **Unload a model**: releases the model from memory without restarting the server. -- **Parameter configuration**: KV cache types, buffer policy, max LLM length — configurable before loading. +- **Parameter configuration**: LLM/DiT KV cache types, buffer policy, max LLM length, flash attention toggles, DiT KV cache slots — configurable before loading. - **Dynamic backend selection**: queries available GGML backends at startup (Auto / CPU / CUDA / Vulkan / Metal). #### Speaker (Voice) Management @@ -143,12 +143,13 @@ The WebUI provides a modern single-page application with the following capabilit #### TTS Generation - **Text input** with configurable voice, mode (zero-shot / instruct / cross-lingual), and output format. +- **Streaming TTS**: enable streaming mode for progressive playback as audio chunks arrive. Configurable chunk token count. - **Advanced sampling controls**: temperature, top-k, top-p, repetition window, tau-r, seed (with "lock for session"). - **Instruct mode**: enter instructions to control speaking style (requires an instruct-capable model). -- **Audio playback**: generated audio plays directly in the browser with a live waveform visualizer. +- **Audio playback**: generated audio plays directly in the browser with a live waveform visualizer. Streaming mode uses MediaSource for true progressive playback. - **Download**: save generated audio as a file. - **History panel**: recent generations are cached up to 20 entries; click to replay or re-download. -- **Responsive output formats**: `wav`, `mp3`, `flac`, `opus`, `aac`, `m4a`, `pcm` (subset depends on FFmpeg runtime). +- **Responsive output formats**: `wav`, `mp3`, `flac`, `opus`, `aac`, `m4a`, `pcm` (subset depends on FFmpeg runtime). Streaming via the WebUI browser player requires MP3, Opus, AAC or FLAC (the MediaSource API does not support WAV progressive playback; WAV streaming via the API works through chunked transfer with a placeholder header). - **Theme toggle**: light / dark mode, persisted in `localStorage`. #### Drag & Drop @@ -206,7 +207,15 @@ If provided at startup, the WebUI pre-populates the frontend configuration field | `--max-llm-len ` | Maximum LLM sequence length. Default: `2048`. | | `--threads, -j ` | CPU thread count. Default: `0` (hardware concurrency). | | `--inference-buffer-policy ` | Inference buffer policy. Default: `balanced`. | -| `--llm-kv-cache-type ,v=[,fallback=]>` | KV cache type. Single type (e.g. `q8_0`) uses the same format for K and V. Default: `k=q8_0,v=q4_0,fallback=q8_0`. | +| `--llm-kv-cache-type ,v=[,fallback=]>` | LLM KV cache type. Single type (e.g. `q8_0`) uses the same format for K and V. Default: `k=q8_0,v=q4_0,fallback=q8_0`. | +| `--dit-kv-cache-type ,v=[,fallback=]>` | DiT (flow matching) KV cache type. Same format as LLM KV cache. Default: `k=q8_0,v=q4_0,fallback=q8_0`. | +| `--dit-kv-fixed-slots ` | Number of fixed (non-offloadable) DiT KV cache slots. Default: `0` (auto). | +| `--dit-kv-offloadable-slots ` | Number of CPU-offloadable DiT KV cache slots. Default: `0` (auto). | +| `--dit-kv-cache-length ` | Maximum sequence length for the DiT KV cache. Default: `0` (auto, `max-llm-len * 10`). | +| `--llm-flash-attn <0\|1>` | Enable/disable LLM flash attention. Default: `1` (enabled). | +| `--flow-flash-attn <0\|1>` | Enable/disable Flow/DiT flash attention. Default: `1` (enabled). | +| `--stream` | Enable streaming TTS for all requests by default (WebUI and API). | +| `--chunk-tokens ` | Number of tokens per streaming chunk. Smaller chunks → lower first-chunk latency but higher overhead and RTF; larger chunks → lower RTF but higher first-chunk latency. Default: model-defined (varies by model). | | `--seed ` | Default random seed for the built-in sampler (when request does not provide a seed). | ### Sampling Default Overrides (Server-level) @@ -298,6 +307,8 @@ Returns server status as JSON: "k_cache_type": "Q8_0", "v_cache_type": "Q4_0", "buffer_policy": "balanced", + "llm_use_flash_attn": true, + "flow_use_flash_attn": true, "model_arch": "cosyvoice3-2512", "frontend_available": false, "speakers": ["alloy", "nova"] @@ -362,25 +373,37 @@ Generates speech. JSON body: "min_token_text_ratio": 1.0, "max_token_text_ratio": 5.0, "instruction": "Speak with a warm and calm style.", - "fadein": true + "fadein": true, + "stream": false, + "chunk_tokens": 0 } ``` Returns raw audio bytes with the appropriate `Content-Type` header (`audio/wav`, `audio/mpeg`, etc.). +When `stream` is `true` (or the server was started with `--stream`), the response is delivered as a chunked HTTP transfer. Audio chunks are streamed incrementally as they are generated. All output formats support streaming at the server level. When using the WebUI, the browser player requires MP3, Opus, AAC or FLAC (the MediaSource API does not support WAV progressive playback). When `stream` is `false` (default), the full audio is buffered and returned as a single blob. + #### `GET /model/defaults` -Returns default model parameters as JSON: +Returns default model parameters as JSON (field set varies based on loaded model): ```json { "max_llm_len": 2048, "temperature": 1.0, "top_k": 50, "top_p": 0.9, + "chunk_tokens": 75, + "default_dit_k_cache_type": "q8_0", + "default_dit_v_cache_type": "q4_0", + "default_dit_kv_fixed_slots": 1, + "default_dit_kv_offloadable_slots": 0, + "default_dit_kv_cache_length": 20480, ... } ``` +DiT KV cache defaults are populated from the effective configuration after model load. `chunk_tokens` reflects the currently configured value (0 = model default). + #### `POST /model/load` Loads a model into the server at runtime. JSON body: @@ -392,10 +415,25 @@ Loads a model into the server at runtime. JSON body: "max_llm_len": 2048, "buffer_policy": "balanced", "k_cache_type": "q8_0", - "v_cache_type": "q4_0" + "v_cache_type": "q4_0", + "llm_use_flash_attn": true, + "flow_use_flash_attn": true, + "dit_kv_cache_type": "k=q8_0,v=q4_0,fallback=q8_0", + "dit_kv_fixed_slots": 0, + "dit_kv_offloadable_slots": 0, + "dit_kv_cache_length": 0, + "chunk_tokens": 0 } ``` +Additional fields: +- `llm_use_flash_attn`, `flow_use_flash_attn`: `true`/`false` toggles for flash attention. +- `dit_kv_cache_type`: DiT KV cache type (same format as `k_cache_type`). +- `dit_kv_fixed_slots`: Number of fixed DiT KV cache slots (0 = auto). +- `dit_kv_offloadable_slots`: Number of CPU-offloadable DiT KV cache slots (0 = auto). +- `dit_kv_cache_length`: Maximum DiT KV cache sequence length (0 = auto, `max_llm_len * 10`). +- `chunk_tokens`: Tokens per streaming chunk (0 = model default). + #### `POST /model/unload` Unloads the current model. JSON body not required. @@ -433,7 +471,7 @@ Authentication behavior: #### Speech request behavior (`POST /v1/audio/speech`): - Required fields: `model`, `input`, `voice`. - Optional standard fields: `instructions`, `speed`, `response_format`. -- Optional extension fields: `seed`, `temperature`, `top_k`, `top_p`, `win_size`, `tau_r`, `min_token_text_ratio`, `max_token_text_ratio`. +- Optional extension fields: `seed`, `temperature`, `top_k`, `top_p`, `win_size`, `tau_r`, `min_token_text_ratio`, `max_token_text_ratio`, `stream` (boolean), `chunk_tokens` (uint32). - Seed precedence: request `seed` extension > server `--seed` > random seed per request. - Mode selection is automatic: non-empty `instructions` -> instruct mode; otherwise zero-shot mode. - Supported `response_format`: `wav`, `pcm`, plus FFmpeg-backed formats `mp3`, `aac`, `flac`, `m4a`, and `opus` when the linked FFmpeg runtime actually provides the required encoders. The server prints the runtime-supported subset in its help text. @@ -441,6 +479,7 @@ Authentication behavior: - `wav` is supported through the audio encoder in normal builds, and through an internal PCM16 WAV fallback when audio helpers are disabled. - `pcm` returns raw 16-bit little-endian mono PCM payload. - If a requested format is not available at runtime the server will return an error explaining the unsupported format; clients should fall back to `wav` or `pcm`. +- **Streaming**: When `stream` is `true` (or the server was started with `--stream`), the API returns a chunked transfer response. Each chunk contains raw audio bytes — no SSE framing. The client must read the stream until exhaustion. All response formats support streaming — there is no format restriction on the API side. WAV sends a placeholder header first, followed by raw PCM16 chunks. OpenAI Python client example: ```python @@ -462,15 +501,18 @@ with open("out.mp3", "wb") as f: Non-standard extension fields (OpenAI SDK `extra_body`): - Standard OpenAI Speech requests usually use: `model`, `input`, `voice`, optional `instructions/speed/response_format`. -- `cosyvoice-server` additionally supports request-level sampling extensions: +- `cosyvoice-server` additionally supports request-level extensions: - `seed`, `temperature`, `top_k`, `top_p`, `win_size`, `tau_r`, `min_token_text_ratio`, `max_token_text_ratio` -- With the OpenAI Python SDK, pass these extensions via `extra_body`: + - `stream` (boolean): enable streaming TTS. + - `chunk_tokens` (uint32): tokens per streaming chunk (0 = model default). +- With the OpenAI Python SDK, pass these via `extra_body`: ```python from openai import OpenAI client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="sk-local-demo") +# Non-streaming request with sampling overrides audio = client.audio.speech.create( model="cosyvoice-3", voice="alloy", @@ -487,9 +529,25 @@ audio = client.audio.speech.create( "max_token_text_ratio": 5.0, }, ) - with open("out.wav", "wb") as f: f.write(audio.read()) + +# Streaming request (read raw stream manually) +import httpx +response = httpx.post( + "http://127.0.0.1:8080/v1/audio/speech", + json={ + "model": "cosyvoice-3", + "voice": "alloy", + "input": "Hello from streaming TTS", + "response_format": "mp3", + "stream": True, + }, + headers={"Authorization": "Bearer sk-local-demo"}, +) +with open("out.mp3", "wb") as f: + for chunk in response.iter_bytes(): + f.write(chunk) ``` Validation rules for these extensions: @@ -501,6 +559,8 @@ Validation rules for these extensions: - `tau_r`: `>= 0` - `min_token_text_ratio`: `>= 0` - `max_token_text_ratio`: `>= 0` and `>= min_token_text_ratio` (when both are set) +- `stream`: boolean +- `chunk_tokens`: uint32 (`[0, 4294967295]`) If validation fails, the server returns OpenAI-style `400` with an error object. @@ -571,6 +631,23 @@ python tools/server/synthesize_via_api.py \ --max-token-text-ratio 5.0 ``` +- Streaming request with real-time playback (requires `pyaudio`): + +```bash +python tools/server/synthesize_via_api.py \ + --base-url http://127.0.0.1:8080 \ + --model cosyvoice-3 \ + --voice alloy \ + --text "Streaming TTS demo" \ + --response-format wav \ + --stream +``` + +Additional options for streaming: +- `--stream`: Enable streaming TTS. Audio plays in real-time via pyaudio (WAV format) or saves progressively. +- `--chunk-tokens `: Tokens per streaming chunk (0 = model default). +- `--no-play`: Disable real-time playback even with `--stream` (useful for saving streaming output to file). + #### `tools/server/batch_tts_stress_test.py` Concurrent stress test helper that sends multiple concurrent requests and keeps all generated audio files. @@ -645,6 +722,8 @@ Type text at the `> ` prompt to synthesize, or use slash commands: - `/clear`: Clear cached audio. - `/seed [value]`: Show or set next seed. - `/seed-policy `: Show or set seed policy. +- `/stream`: Toggle streaming playback (audio plays progressively during generation). +- `/chunk-tokens [value]`: Show or set tokens per streaming chunk. - `/help`: Show command list. - `/exit`: Exit interactive mode. `Ctrl+C` also exits. @@ -672,6 +751,14 @@ Core options: - `--inference-buffer-policy `: Inference buffer policy. Only effective in interactive mode; non-interactive mode always uses `shared` for minimal memory footprint and fastest single-shot synthesis. Default: `balanced`. - `--mode `: TTS mode. Default: auto-detect from `--instruction`. - `--instruction, -i `: Instruction text for instruct mode. +- `--dit-kv-cache-type ,v=[,fallback=]>`: DiT KV cache type (interactive only). Same format as `--llm-kv-cache-type`. Default: `k=q8_0,v=q4_0,fallback=q8_0`. +- `--dit-kv-fixed-slots `: Number of fixed (non-offloadable) DiT KV cache slots (interactive only). Default: `0` (auto). +- `--dit-kv-offloadable-slots `: Number of CPU-offloadable DiT KV cache slots (interactive only). Default: `0` (auto). +- `--dit-kv-cache-length `: Maximum DiT KV cache sequence length (interactive only). Default: `0` (auto, `max-llm-len * 10`). +- `--stream`: Enable streaming playback in interactive mode (audio plays progressively during generation). +- `--chunk-tokens `: Tokens per streaming chunk (interactive only). Smaller chunks → lower first-chunk latency but higher overhead and RTF; larger chunks → lower RTF but higher first-chunk latency. Default: model-defined. +- `--llm-flash-attn <0|1>`: Enable/disable LLM flash attention. Default: `1` (enabled). +- `--flow-flash-attn <0|1>`: Enable/disable Flow/DiT flash attention. Default: `1` (enabled). Sampling override options: - `--temperature `: Sampling temperature, must be `> 0`. @@ -753,6 +840,14 @@ Required vs optional: | `--llm-kv-cache-type` | `k=q8_0,v=q4_0,fallback=q8_0` | CLI | | `--inference-buffer-policy` | `balanced` (interactive) / `shared` (non-interactive) | CLI | | `--seed` | random | runtime | +| `--stream` | `false` | CLI | +| `--chunk-tokens` | model-defined | model | +| `--dit-kv-cache-type` | `k=q8_0,v=q4_0,fallback=q8_0` | CLI | +| `--dit-kv-fixed-slots` | `0` (auto) | CLI | +| `--dit-kv-offloadable-slots` | `0` (auto) | CLI | +| `--dit-kv-cache-length` | `0` (auto, `max-llm-len * 10`) | CLI | +| `--llm-flash-attn` | `1` (enabled) | CLI | +| `--flow-flash-attn` | `1` (enabled) | CLI | | `temperature`, `top_k`, `top_p`, `win_size`, `tau_r`, `min/max_token_text_ratio` | model metadata | model | ### Typical command templates diff --git a/docs/TOOLS_zh.md b/docs/TOOLS_zh.md index 99bd8c1..f17e6aa 100644 --- a/docs/TOOLS_zh.md +++ b/docs/TOOLS_zh.md @@ -129,7 +129,7 @@ WebUI 是一个现代化的单页应用，提供以下功能： #### 模型管理 - **运行时加载模型**：输入 `.gguf` 文件路径，选择后端和线程数，点击「Load Model」加载。 - **卸载模型**：释放模型内存，无需重启服务端。 -- **参数配置**：KV cache 类型、buffer 策略、最大 LLM 长度——可在加载前配置。 +- **参数配置**：LLM/DiT KV cache 类型、buffer 策略、最大 LLM 长度、Flash Attention 开关、DiT KV cache 槽位数——可在加载前配置。 - **动态后端选择**：启动时自动探测可用的 GGML 后端（Auto / CPU / CUDA / Vulkan / Metal）。 #### 音色管理 @@ -143,12 +143,13 @@ WebUI 是一个现代化的单页应用，提供以下功能： #### TTS 生成 - **文本输入**：可选择音色、模式（zero-shot / instruct / cross-lingual）和输出格式。 +- **流式 TTS**：开启流式模式后，音频块到达时即可渐进播放。可配置每个块的 token 数。 - **高级采样控制**：temperature、top-k、top-p、重复惩罚窗口、tau-r、随机种子（支持锁定）。 - **Instruct 模式**：输入指令控制说话风格（需要模型支持 instruct）。 -- **音频播放**：生成的音频直接在浏览器中播放，配有波形可视化。 +- **音频播放**：生成的音频直接在浏览器中播放，配有波形可视化。流式模式使用 MediaSource 实现真正的渐进式播放。 - **下载**：保存生成的音频文件。 - **历史记录**：最近最多 20 条生成记录，可点击回放或重新下载。 -- **多格式输出**：`wav`、`mp3`、`flac`、`opus`、`aac`、`m4a`、`pcm`（可用格式取决于 FFmpeg 运行时）。 +- **多格式输出**：`wav`、`mp3`、`flac`、`opus`、`aac`、`m4a`、`pcm`（可用格式取决于 FFmpeg 运行时）。通过 WebUI 浏览器播放流式音频需要 MP3、Opus、AAC 或 FLAC（MediaSource API 不支持 WAV 渐进播放；通过 API 的 WAV 流式可通过分块传输加占位头实现）。 - **主题切换**：亮色/暗色模式，通过 `localStorage` 持久化。 #### 拖拽支持 @@ -206,7 +207,15 @@ WebUI 是一个现代化的单页应用，提供以下功能： | `--threads, -j ` | CPU 线程数，默认 `0`（硬件并发数）。 | | `--concurrency, -c ` | 并发请求槽数，默认 `1`（仅 API 模式；WebUI 模式始终为单槽）。 | | `--inference-buffer-policy ` | 推理缓冲区策略，默认 `balanced`。 | -| `--llm-kv-cache-type ,v=[,fallback=]>` | KV cache 类型。单一类型（如 `q8_0`）为 K 和 V 使用相同格式。默认 `k=q8_0,v=q4_0,fallback=q8_0`。 | +| `--llm-kv-cache-type ,v=[,fallback=]>` | LLM KV cache 类型。单一类型（如 `q8_0`）为 K 和 V 使用相同格式。默认 `k=q8_0,v=q4_0,fallback=q8_0`。 | +| `--dit-kv-cache-type ,v=[,fallback=]>` | DiT（flow matching）KV cache 类型。格式同 LLM KV cache。默认 `k=q8_0,v=q4_0,fallback=q8_0`。 | +| `--dit-kv-fixed-slots ` | 固定（不可卸载）DiT KV cache 槽位数，默认 `0`（自动）。 | +| `--dit-kv-offloadable-slots ` | 可 CPU 卸载的 DiT KV cache 槽位数，默认 `0`（自动）。 | +| `--dit-kv-cache-length ` | DiT KV cache 最大序列长度，默认 `0`（自动，为 `max-llm-len * 10`）。 | +| `--llm-flash-attn <0\|1>` | 启用/禁用 LLM Flash Attention。默认 `1`（启用）。 | +| `--flow-flash-attn <0\|1>` | 启用/禁用 Flow/DiT Flash Attention。默认 `1`（启用）。 | +| `--stream` | 默认对所有请求启用流式 TTS（WebUI 和 API 模式均生效）。 | +| `--chunk-tokens ` | 每个流式块的 token 数。chunk 越小，首包延迟越低，但上下文调度开销越大，RTF 越高；chunk 越大，RTF 越低，但首包延迟越高。默认：模型定义（因模型而异）。 | | `--seed ` | 默认随机种子（当请求未传 seed 时使用）。 | ### 采样默认值覆盖（服务级） @@ -290,17 +299,19 @@ WebUI 暴露以下 REST 接口，由前端 JavaScript 调用： 返回服务端状态 JSON： ```json { - “status”: “ok”, - “model_loaded”: true, - “model”: “cosyvoice-3”, - “sample_rate”: 24000, - “max_llm_len”: 2048, - “k_cache_type”: “Q8_0”, - “v_cache_type”: “Q4_0”, - “buffer_policy”: “balanced”, - “model_arch”: “cosyvoice3-2512”, - “frontend_available”: false, - “speakers”: [“alloy”, “nova”] + "status": "ok", + "model_loaded": true, + "model": "cosyvoice-3", + "sample_rate": 24000, + "max_llm_len": 2048, + "k_cache_type": "Q8_0", + "v_cache_type": "Q4_0", + "buffer_policy": "balanced", + "llm_use_flash_attn": true, + "flow_use_flash_attn": true, + "model_arch": "cosyvoice3-2512", + "frontend_available": false, + "speakers": ["alloy", "nova"] } ``` @@ -348,54 +359,81 @@ JSON 请求体： ```json { - “text”: “你好，这里是 CosyVoice”, - “voice”: “alloy”, - “mode”: “zero-shot”, - “response_format”: “wav”, - “speed”: 1.0, - “seed”: 12345, - “temperature”: 0.8, - “top_k”: 32, - “top_p”: 0.9, - “win_size”: 10, - “tau_r”: 0.2, - “min_token_text_ratio”: 1.0, - “max_token_text_ratio”: 5.0, - “instruction”: “语气温和、平稳。”, - “fadein”: true + "text": "你好，这里是 CosyVoice", + "voice": "alloy", + "mode": "zero-shot", + "response_format": "wav", + "speed": 1.0, + "seed": 12345, + "temperature": 0.8, + "top_k": 32, + "top_p": 0.9, + "win_size": 10, + "tau_r": 0.2, + "min_token_text_ratio": 1.0, + "max_token_text_ratio": 5.0, + "instruction": "语气温和、平稳。", + "fadein": true, + "stream": false, + "chunk_tokens": 0 } ``` 返回原始音频字节，带对应 `Content-Type` 头（`audio/wav`、`audio/mpeg` 等）。 +当 `stream` 为 `true`（或服务端以 `--stream` 启动）时，响应以 HTTP 分块传输方式返回。音频块在生成过程中逐步传输。服务端层面对所有输出格式均支持流式。通过 WebUI 使用浏览器播放时，需要 MP3、Opus、AAC 或 FLAC（MediaSource API 不支持 WAV 渐进播放）。当 `stream` 为 `false`（默认）时，完整音频缓存后一次性返回。 + #### `GET /model/defaults` -返回模型默认参数 JSON： +返回模型默认参数 JSON（字段集因已加载模型而异）： ```json { - “max_llm_len”: 2048, - “temperature”: 1.0, - “top_k”: 50, - “top_p”: 0.9, + "max_llm_len": 2048, + "temperature": 1.0, + "top_k": 50, + "top_p": 0.9, + "chunk_tokens": 75, + "default_dit_k_cache_type": "q8_0", + "default_dit_v_cache_type": "q4_0", + "default_dit_kv_fixed_slots": 1, + "default_dit_kv_offloadable_slots": 0, + "default_dit_kv_cache_length": 20480, … } ``` +DiT KV cache 默认值从加载后的生效配置中获取。`chunk_tokens` 反映当前配置值（0 = 模型默认）。 + #### `POST /model/load` 运行时加载模型。JSON 请求体： ```json { - “model_path”: “/data/models/cosyvoice-3.gguf”, - “backend”: “auto”, - “n_threads”: 8, - “max_llm_len”: 2048, - “buffer_policy”: “balanced”, - “k_cache_type”: “q8_0”, - “v_cache_type”: “q4_0” + "model_path": "/data/models/cosyvoice-3.gguf", + "backend": "auto", + "n_threads": 8, + "max_llm_len": 2048, + "buffer_policy": "balanced", + "k_cache_type": "q8_0", + "v_cache_type": "q4_0", + "llm_use_flash_attn": true, + "flow_use_flash_attn": true, + "dit_kv_cache_type": "k=q8_0,v=q4_0,fallback=q8_0", + "dit_kv_fixed_slots": 0, + "dit_kv_offloadable_slots": 0, + "dit_kv_cache_length": 0, + "chunk_tokens": 0 } ``` +其他字段： +- `llm_use_flash_attn`、`flow_use_flash_attn`：`true`/`false`，Flash Attention 开关。 +- `dit_kv_cache_type`：DiT KV cache 类型（格式同 `k_cache_type`）。 +- `dit_kv_fixed_slots`：固定 DiT KV cache 槽位数（0 = 自动）。 +- `dit_kv_offloadable_slots`：可 CPU 卸载的 DiT KV cache 槽位数（0 = 自动）。 +- `dit_kv_cache_length`：DiT KV cache 最大序列长度（0 = 自动，为 `max_llm_len * 10`）。 +- `chunk_tokens`：每个流式块的 token 数（0 = 模型默认）。 + #### `POST /model/unload` 卸载当前模型（无需 JSON 请求体）。 @@ -433,7 +471,7 @@ JSON 请求体： #### 语音请求行为（`POST /v1/audio/speech`）： - 必填字段：`model`、`input`、`voice`。 - 标准可选字段：`instructions`、`speed`、`response_format`。 -- 扩展可选字段：`seed`、`temperature`、`top_k`、`top_p`、`win_size`、`tau_r`、`min_token_text_ratio`、`max_token_text_ratio`。 +- 扩展可选字段：`seed`、`temperature`、`top_k`、`top_p`、`win_size`、`tau_r`、`min_token_text_ratio`、`max_token_text_ratio`、`stream`（布尔值）、`chunk_tokens`（uint32）。 - seed 优先级：请求扩展字段 `seed` > 服务端 `--seed` > 随机 seed。 - 模式自动判定：`instructions` 非空时走 instruct；否则走 zero-shot。 - 支持 `response_format`：`wav`、`pcm`，以及 FFmpeg 后端提供的 `mp3`、`aac`、`flac`、`m4a`、`opus`。 @@ -441,6 +479,7 @@ JSON 请求体： - `wav` 通过音频编码器生成；关闭音频辅助 API 时回退到内部 PCM16 WAV 实现。 - `pcm` 返回原始 16 位小端单声道 PCM 负载。 - 格式不可用时返回错误。 +- **流式**：当 `stream` 为 `true`（或服务端以 `--stream` 启动）时，API 返回分块传输响应。每个块包含原始音频字节——无 SSE 帧格式。客户端需读取流直至结束。**API 流式无格式限制**——所有 `response_format` 均可用于流式。WAV 先发送占位头，之后发送原始 PCM16 块。 OpenAI Python 客户端示例： ```python @@ -461,29 +500,36 @@ with open(“out.mp3”, “wb”) as f: ``` 非标准扩展字段（OpenAI SDK `extra_body`）： +- 标准 OpenAI Speech 请求通常使用：`model`、`input`、`voice`，可选 `instructions/speed/response_format`。 +- `cosyvoice-server` 额外支持请求级扩展字段： + - `seed`、`temperature`、`top_k`、`top_p`、`win_size`、`tau_r`、`min_token_text_ratio`、`max_token_text_ratio` + - `stream`（布尔值）：启用流式 TTS。 + - `chunk_tokens`（uint32）：每个流式块的 token 数（0 = 模型默认）。 + ```python from openai import OpenAI -client = OpenAI(base_url=”http://127.0.0.1:8080/v1”, api_key=”sk-local-demo”) +client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="sk-local-demo") +# 非流式请求，带采样覆盖 audio = client.audio.speech.create( - model=”cosyvoice-3”, - voice=”alloy”, - input=”你好，这里是 cosyvoice-server”, - response_format=”wav”, + model="cosyvoice-3", + voice="alloy", + input="你好，这里是 cosyvoice-server", + response_format="wav", extra_body={ - “seed”: 123, - “temperature”: 0.8, - “top_k”: 32, - “top_p”: 0.9, - “win_size”: 10, - “tau_r”: 0.2, - “min_token_text_ratio”: 1.0, - “max_token_text_ratio”: 5.0, + "seed": 123, + "temperature": 0.8, + "top_k": 32, + "top_p": 0.9, + "win_size": 10, + "tau_r": 0.2, + "min_token_text_ratio": 1.0, + "max_token_text_ratio": 5.0, }, ) -with open(“out.wav”, “wb”) as f: +with open("out.wav", "wb") as f: f.write(audio.read()) ``` @@ -496,6 +542,8 @@ with open(“out.wav”, “wb”) as f: - `tau_r`：`>= 0` - `min_token_text_ratio`：`>= 0` - `max_token_text_ratio`：`>= 0` 且同时设置时 `>= min_token_text_ratio` +- `stream`：布尔值 +- `chunk_tokens`：uint32（`[0, 4294967295]`） 校验失败时返回 OpenAI 风格 `400` 错误对象。 @@ -562,6 +610,22 @@ python tools/server/synthesize_via_api.py \ --max-token-text-ratio 5.0 ``` +流式请求（实时播放，需安装 pyaudio）： +```bash +python tools/server/synthesize_via_api.py \ + --base-url http://127.0.0.1:8080 \ + --model cosyvoice-3 \ + --voice alloy \ + --text "流式 TTS 演示" \ + --response-format wav \ + --stream +``` + +流式相关选项： +- `--stream`：启用流式 TTS。WAV 格式通过 pyaudio 实时播放，其他格式渐进保存。 +- `--chunk-tokens `：每个流式块的 token 数（0 = 模型默认）。 +- `--no-play`：即使启用 `--stream` 也禁用实时播放（用于将流式输出保存到文件）。 + #### `tools/server/batch_tts_stress_test.py` 并发压测脚本，会并发发送多次请求并保留全部输出音频文件。 @@ -636,6 +700,8 @@ cosyvoice-cli \ - `/clear`：清空缓存。 - `/seed [value]`：显示或设置下一次 seed。 - `/seed-policy `：显示或设置 seed 策略。 +- `/stream`：切换流式播放（生成过程中渐进播放音频）。 +- `/chunk-tokens [value]`：显示或设置每个流式块的 token 数。 - `/help`：显示命令列表。 - `/exit`：退出交互模式，`Ctrl+C` 也可以退出。 @@ -663,6 +729,14 @@ cosyvoice-cli \ - `--inference-buffer-policy `：推理缓冲区策略。仅在交互模式下生效；非交互模式始终使用 `shared` 以最小化内存占用并获得最快的单次合成速度。默认 `balanced`。 - `--mode `：TTS 模式。默认按 `--instruction` 自动判定。 - `--instruction, -i `：instruct 模式指令文本。 +- `--dit-kv-cache-type ,v=[,fallback=]>`：DiT KV cache 类型（仅交互模式）。格式同 `--llm-kv-cache-type`。默认 `k=q8_0,v=q4_0,fallback=q8_0`。 +- `--dit-kv-fixed-slots `：固定（不可卸载）DiT KV cache 槽位数（仅交互模式）。默认 `0`（自动）。 +- `--dit-kv-offloadable-slots `：可 CPU 卸载的 DiT KV cache 槽位数（仅交互模式）。默认 `0`（自动）。 +- `--dit-kv-cache-length `：DiT KV cache 最大序列长度（仅交互模式）。默认 `0`（自动，为 `max-llm-len * 10`）。 +- `--stream`：在交互模式下启用流式播放（生成过程中渐进播放音频）。 +- `--chunk-tokens `：每个流式块的 token 数（仅交互模式）。chunk 越小，首包延迟越低，但上下文调度开销越大，RTF 越高；chunk 越大，RTF 越低，但首包延迟越高。默认：模型定义。 +- `--llm-flash-attn <0|1>`：启用/禁用 LLM Flash Attention。默认 `1`（启用）。 +- `--flow-flash-attn <0|1>`：启用/禁用 Flow/DiT Flash Attention。默认 `1`（启用）。 采样覆盖参数： - `--temperature `：采样温度，必须 `> 0`。 @@ -744,6 +818,14 @@ cosyvoice-cli \ | `--llm-kv-cache-type` | `k=q8_0,v=q4_0,fallback=q8_0` | CLI | | `--inference-buffer-policy` | `balanced`（交互模式）/ `shared`（非交互模式） | CLI | | `--seed` | 随机 | 运行时 | +| `--stream` | `false` | CLI | +| `--chunk-tokens` | 模型定义 | 模型 | +| `--dit-kv-cache-type` | `k=q8_0,v=q4_0,fallback=q8_0` | CLI | +| `--dit-kv-fixed-slots` | `0`（自动） | CLI | +| `--dit-kv-offloadable-slots` | `0`（自动） | CLI | +| `--dit-kv-cache-length` | `0`（自动，`max-llm-len * 10`） | CLI | +| `--llm-flash-attn` | `1`（启用） | CLI | +| `--flow-flash-attn` | `1`（启用） | CLI | | `temperature/top_k/top_p/win_size/tau_r/min/max_token_text_ratio` | 模型元数据 | 模型 | ### 常用命令模板 From 87f44c9d677706eec0d63dc36ea3250147fb8ad3 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Thu, 16 Jul 2026 14:39:44 +0800 Subject: [PATCH 20/34] server: correct preprocessor directive for PCH usage --- tools/server/server_common.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/server/server_common.h b/tools/server/server_common.h index b8da74c..82fa1ed 100644 --- a/tools/server/server_common.h +++ b/tools/server/server_common.h @@ -10,7 +10,7 @@ namespace httplib { class Request; class Response; } -#ifdef COSYVOICE_SERVER_USE_PCH +#ifdef COSYVOICE_USE_PCH namespace httplib { class Server; } using Server = httplib::Server; #else From 807149579d85984dac01751201f87d8e0c952fa0 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Thu, 16 Jul 2026 16:57:06 +0800 Subject: [PATCH 21/34] cosyvoice: correct layer index calculation in slide_kv_slot function --- src/cosyvoice-kv-cache.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cosyvoice-kv-cache.cpp b/src/cosyvoice-kv-cache.cpp index 5a892b9..8c405d6 100644 --- a/src/cosyvoice-kv-cache.cpp +++ b/src/cosyvoice-kv-cache.cpp @@ -391,7 +391,7 @@ void cosyvoice_kv_cache::slide_kv_slot() GGML_ASSERT(cur_slot_idx < n_slots); if (fattn) { - auto layer_idx = cur_slot_idx * layers; + auto layer_idx = cur_slot_idx++ * layers; const auto end = layer_idx + layers; for (int cur = layer_idx; cur != end; ++cur) { From e7ea6554c6b5bed4908f6e77f3c4d35a98a9f28e Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Thu, 16 Jul 2026 17:03:28 +0800 Subject: [PATCH 22/34] cosyvoice: enhance token2wav_ext for streaming support and cache management --- src/cosyvoice-tts.cpp | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index d58dc3a..ea4e151 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -383,6 +383,12 @@ struct dit_sched_config bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* offset, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result) { + if (streaming && offset && *offset == 0) + { + worker->flow_cache.clear(); + worker->chunk_boundaries.clear(); + } + const auto& params = shared->params; auto& sched = worker->sched; auto& flow = cv3_shared->flow; @@ -580,7 +586,7 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f // Phase 3: Copy flow output to speech_feat (persistent buffer) ggml_reset(ctx1.get()); - auto cache_length = static_cast(worker->flow_cache.size() / feat->ne[0]); + const auto cache_length = streaming && offset ? static_cast(worker->flow_cache.size() / feat->ne[0]) : 0; ggml_tensor* speech_feat = ggml_new_tensor_2d(ctx1.get(), feat->type, feat->ne[0], feat->ne[1] + cache_length); if (cache_length != 0) { @@ -646,12 +652,6 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_AFTER_HIFT, noise_len, noise_buffer, shared->noise_callback_ctx); - if (finalize) - { - worker->flow_cache.clear(); - worker->chunk_boundaries.clear(); - } - if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) { ggml_backend_sched_reset(sched.get()); From 49b2ca931991a1c6483e82d914cef47f14745566 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Tue, 21 Jul 2026 23:48:08 +0800 Subject: [PATCH 23/34] cosyvoice: add stop-request API for graceful TTS interruption --- include/cosyvoice-interface.h | 4 +++ include/cosyvoice-lowlevel.h | 20 +++++++++++++++ src/cosyvoice-internal.h | 18 ++++++++++++-- src/cosyvoice-kv-cache.cpp | 2 ++ src/cosyvoice-model.cpp | 23 ++++++++++++++++++ src/cosyvoice-model.h | 9 +++++++ src/cosyvoice-tts.cpp | 46 +++++++++++++++++++++++++++++++---- src/cosyvoice.cpp | 42 ++++++++++++++++++++++++++++++++ 8 files changed, 157 insertions(+), 7 deletions(-) diff --git a/include/cosyvoice-interface.h b/include/cosyvoice-interface.h index 3570d78..f630c60 100644 --- a/include/cosyvoice-interface.h +++ b/include/cosyvoice-interface.h @@ -120,6 +120,10 @@ struct cosyvoice_model_context virtual void get_noise_callback(cosyvoice_noise_callback_t* callback, void** callback_ctx) = 0; ///< Query the shared noise callback. virtual uint32_t get_hift_rand_ini_len() = 0; ///< Get the required shared HiFT initialization-noise length. virtual void set_hift_rand_ini(const float* data) = 0; ///< Set the shared HiFT initialization-noise buffer. + + // Stop-request API + virtual void request_stop() = 0; ///< Request that the active worker's current job stop as soon as possible. + virtual bool stop_requested() = 0; ///< Check and atomically clear the stop-requested flag. Returns true if a stop was requested. }; // ---------------------------------------------------------------------------- diff --git a/include/cosyvoice-lowlevel.h b/include/cosyvoice-lowlevel.h index ae4d1e4..d130a87 100644 --- a/include/cosyvoice-lowlevel.h +++ b/include/cosyvoice-lowlevel.h @@ -457,6 +457,26 @@ COSYVOICE_API uint32_t cosyvoice_get_hift_rand_ini_len(cosyvoice_context_t ctx); */ COSYVOICE_API void cosyvoice_set_hift_rand_ini(cosyvoice_context_t ctx, const float* data); +// ---------------------------------------------------------------------------- +// Stop-Request API +// ---------------------------------------------------------------------------- + +/** + * @brief Request that the current worker's job stop as soon as possible. + * @param ctx Model context. + */ +COSYVOICE_API void cosyvoice_request_stop(cosyvoice_context_t ctx); + +/** + * @brief Check and atomically clear the stop-requested flag for the current worker. + * @param ctx Model context. + * @return True if a stop was requested, false otherwise. + * @note Used internally by hot paths (llm_job_ext, token2wav_ext, tts) to + * detect stop requests. The flag is atomically reset so subsequent calls + * return false. + */ +COSYVOICE_API bool cosyvoice_stop_requested(cosyvoice_context_t ctx); + // ---------------------------------------------------------------------------- // Prompt Utilities // ---------------------------------------------------------------------------- diff --git a/src/cosyvoice-internal.h b/src/cosyvoice-internal.h index 191a77f..684c21b 100644 --- a/src/cosyvoice-internal.h +++ b/src/cosyvoice-internal.h @@ -145,6 +145,22 @@ struct cosyvoice_prompt void calculate_crc32(); }; +struct cosyvoice_worker_context; +struct cosyvoice_model; + +struct use_count_guard +{ + cosyvoice_worker_context* worker; + + use_count_guard(cosyvoice_context* ctx); + use_count_guard(cosyvoice_model* model); + + ~use_count_guard(); + + use_count_guard(const use_count_guard&) = delete; + use_count_guard& operator=(const use_count_guard&) = delete; +}; + struct cosyvoice_tokenization_result_impl : cosyvoice_tokenization_result { cosyvoice_tokenization_result_impl() {} @@ -161,8 +177,6 @@ bool cosyvoice_frontend_util_text_normalize( const char* locale ); -struct cosyvoice_worker_context; - int cosyvoice_llm_sampler( cosyvoice_llm_token_prob_t* nucleus_probs, int k, diff --git a/src/cosyvoice-kv-cache.cpp b/src/cosyvoice-kv-cache.cpp index 8c405d6..b3cc2cb 100644 --- a/src/cosyvoice-kv-cache.cpp +++ b/src/cosyvoice-kv-cache.cpp @@ -217,6 +217,8 @@ void cosyvoice_kv_cache::load_cache(ggml_backend_t backend, ggml_backend_sched* const size_t v_head_nbytes = fattn ? kv_cache_layers[0].v->nb[1] * offloaded_cache->len : ggml_row_size(kv_cache_layers[0].v->type, offloaded_cache->len) * kv_cache_layers[0].v->ne[1]; cur_len = offloaded_cache->len; + if (cur_len == 0) return; + if (!fattn) { ggml_reset(offloaded_cache->ctx); diff --git a/src/cosyvoice-model.cpp b/src/cosyvoice-model.cpp index 40fd79c..860a0d4 100644 --- a/src/cosyvoice-model.cpp +++ b/src/cosyvoice-model.cpp @@ -559,6 +559,29 @@ uint32_t cosyvoice_model::get_n_workers() return shared->params.n_workers; } +void cosyvoice_model::request_stop() +{ + if (worker->use_count.load(std::memory_order_acquire) == 0) + { + worker->stop_flag.store(false, std::memory_order_release); + return; + } + + worker->stop_flag.store(true, std::memory_order_release); + + std::unique_lock lock(worker->cv_mutex); + worker->cv.wait(lock, [this]() { + return worker->use_count.load(std::memory_order_acquire) == 0; + }); + + worker->stop_flag.store(false, std::memory_order_release); +} + +bool cosyvoice_model::stop_requested() +{ + return worker->stop_flag.exchange(false, std::memory_order_acq_rel); +} + uint32_t cosyvoice_model::llm_get_kv_cache_len() { return worker->llm_kv_cache.cur_len; diff --git a/src/cosyvoice-model.h b/src/cosyvoice-model.h index fc7f359..d964b31 100644 --- a/src/cosyvoice-model.h +++ b/src/cosyvoice-model.h @@ -8,6 +8,8 @@ #include #include +#include +#include #include #include #include @@ -81,6 +83,11 @@ struct cosyvoice_worker_context std::unique_ptr probs; ggml_backend_buffer_ptr llm_kv_buffer; ggml_backend_buffer_ptr dit_kv_buffer; + + std::atomic stop_flag{false}; + std::atomic use_count{0}; + std::mutex cv_mutex; + std::condition_variable cv; }; struct cosyvoice_model : virtual cosyvoice_model_context, virtual cosyvoice_object_ref_counter @@ -92,6 +99,8 @@ struct cosyvoice_model : virtual cosyvoice_model_context, virtual cosyvoice_obje uint32_t get_worker_no(); uint32_t get_n_workers(); + void request_stop(); + bool stop_requested(); uint32_t llm_get_kv_cache_len(); bool llm_set_kv_cache_len(uint32_t len); diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index ea4e151..9a36ee3 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -23,6 +23,7 @@ bool cosyvoice_model_3::llm_job(const int* text, uint32_t text_len, cosyvoice_pr bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoice_prompt_t prompt, uint32_t max_new_tokens, bool* final) { + use_count_guard _guard(this); const auto& params = shared->params; auto& sched = worker->sched; const auto& llm = cv3_shared->llm; @@ -31,6 +32,12 @@ bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoic auto last_prompt_crc32 = prompt_crc32; bool stop_reached = false; + if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) + { + ggml_backend_sched_reset(sched.get()); + worker->llm_kv_cache.load_cache(worker->backend.get(), sched.get()); + } + try { const auto n_batch = params.n_batch; @@ -161,6 +168,16 @@ bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoic for (uint32_t n = llm_get_n_accepted_tokens(); n != limit; ++n) { + if (worker->stop_flag.load(std::memory_order_acquire)) + { + worker->stop_flag.store(false, std::memory_order_release); + worker->llm_input = nullptr; + cosyvoice_call_ggml_log_callback(GGML_LOG_LEVEL_INFO, "LLM generation stopped by user.\n"); + if (params.builtin_sampler_rng_policy == COSYVOICE_BUILTIN_SAMPLER_RNG_POLICY_RESET_PER_SESSION) + reset_builtin_sampler_rng(); + return false; + } + if (!llm_decode(speech_type, cur)) throw std::runtime_error("Failed to decode LLM output.\n"); @@ -383,6 +400,7 @@ struct dit_sched_config bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* offset, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result) { + use_count_guard _guard(this); if (streaming && offset && *offset == 0) { worker->flow_cache.clear(); @@ -525,6 +543,15 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f int kv_slot = config[0].cache_kv && !config[0].offload; for (int step = 1; step != flow.decoder.diffusion_steps; ++step) { + if (worker->stop_flag.load(std::memory_order_acquire)) + { + worker->stop_flag.store(false, std::memory_order_release); + ggml_reset(ctx0.get()); + ggml_backend_sched_reset(sched.get()); + cosyvoice_call_ggml_log_callback(GGML_LOG_LEVEL_INFO, "token2wav stopped during DiT steps.\n"); + return false; + } + if (config[step].slice) { const auto cut_len = config[step - 1].cut_len; @@ -652,11 +679,6 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_AFTER_HIFT, noise_len, noise_buffer, shared->noise_callback_ctx); - if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) - { - ggml_backend_sched_reset(sched.get()); - worker->llm_kv_cache.load_cache(backend.get(), sched.get()); - } return worker->status == GGML_STATUS_SUCCESS; } @@ -781,6 +803,12 @@ struct cosyvoice_tts_context : cosyvoice_tokenization_result_impl, cosyvoice_pro combined_pcm.clear(); for (auto& chunk_tokens : chunk_token_list) { + if (ctx->stop_requested()) + { + result->data = nullptr; + result->length = 0; + return false; + } tokens = std::move(chunk_tokens); cosyvoice_generated_speech part = {}; if (!cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, &part) @@ -799,6 +827,8 @@ struct cosyvoice_tts_context : cosyvoice_tokenization_result_impl, cosyvoice_pro for (auto& chunk_tokens : chunk_token_list) { + if (ctx->stop_requested()) + return false; tokens = std::move(chunk_tokens); if (!cosyvoice_tts_stream_with_postprocess(get_tokens(), get_n_tokens(), speed, callback, user_data)) return false; @@ -829,6 +859,12 @@ struct cosyvoice_tts_context : cosyvoice_tokenization_result_impl, cosyvoice_pro combined_pcm.clear(); for (const auto& chunk : chunks) { + if (ctx->stop_requested()) + { + result->data = nullptr; + result->length = 0; + return false; + } ctx->tokenize(chunk.c_str(), this, true); cosyvoice_generated_speech part = {}; if (!cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, &part) diff --git a/src/cosyvoice.cpp b/src/cosyvoice.cpp index a4cb79c..2b1c72a 100644 --- a/src/cosyvoice.cpp +++ b/src/cosyvoice.cpp @@ -25,6 +25,20 @@ uint32_t cosyvoice_generate_random_seed() return dist(gen); } +use_count_guard::use_count_guard(cosyvoice_model* model) : worker(model->worker) +{ + worker->use_count.fetch_add(1, std::memory_order_relaxed); +} + +use_count_guard::use_count_guard(cosyvoice_context* ctx) + : use_count_guard(dynamic_cast(ctx)) {} + +use_count_guard::~use_count_guard() +{ + if (worker->use_count.fetch_sub(1, std::memory_order_release) == 1) + worker->cv.notify_all(); +} + #ifdef _WIN32 #define NOMINMAX #include @@ -363,6 +377,8 @@ static bool check_length(cosyvoice_prompt_t prompt, uint32_t text_len, uint32_t bool cosyvoice_tts(cosyvoice_context_t ctx, const int* text, uint32_t text_len, float speed, cosyvoice_prompt_t prompt, cosyvoice_generated_speech_ptr result) { + use_count_guard guard(ctx); + { cosyvoice_context_params_t params; ctx->get_context_params(¶ms); @@ -376,7 +392,18 @@ bool cosyvoice_tts(cosyvoice_context_t ctx, const int* text, uint32_t text_len, if (ctx->llm_job(text, text_len, prompt) && ctx->token2wav(ctx->llm_get_accepted_tokens(), ctx->llm_get_n_accepted_tokens(), speed, prompt, result)) + { + // Check if stop was requested during token2wav + if (ctx->stop_requested()) + { + result->data = nullptr; + result->length = 0; + return false; + } return true; + } + // Clear stop flag if generation failed due to stop request + ctx->stop_requested(); result->data = nullptr; result->length = 0; return false; @@ -384,6 +411,8 @@ bool cosyvoice_tts(cosyvoice_context_t ctx, const int* text, uint32_t text_len, bool cosyvoice_tts_stream(cosyvoice_context_t ctx, const int* text, uint32_t text_len, float speed, cosyvoice_prompt_t prompt, cosyvoice_tts_audio_callback_t callback, void* user_data) { + use_count_guard guard(ctx); + cosyvoice_context_params_t params; ctx->get_context_params(¶ms); if (!check_length(prompt, text_len, params.n_max_seq)) @@ -398,6 +427,9 @@ bool cosyvoice_tts_stream(cosyvoice_context_t ctx, const int* text, uint32_t tex bool final = false; do { + if (ctx->stop_requested()) + return false; + if (!final) if (!ctx->llm_job_ext(text, text_len, prompt, chunk_tokens + 1, &final)) return false; @@ -450,6 +482,16 @@ void cosyvoice_empty_buffer_cache(cosyvoice_context_t ctx) ctx->empty_buffer_cache(); } +void cosyvoice_request_stop(cosyvoice_context_t ctx) +{ + ctx->request_stop(); +} + +bool cosyvoice_stop_requested(cosyvoice_context_t ctx) +{ + return ctx->stop_requested(); +} + void cosyvoice_set_noise_callback(cosyvoice_context_t ctx, cosyvoice_noise_callback_t callback, void* callback_ctx) { ctx->set_noise_callback(callback, callback_ctx); From 0b9b8df519cdcb17f5e2b88448591fed8d00c0ae Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Wed, 22 Jul 2026 00:09:47 +0800 Subject: [PATCH 24/34] cli: add Ctrl+C stop support --- tools/cli/cosyvoice-cli.cpp | 94 +++++++++++++++++++++++++++++++++++-- 1 file changed, 89 insertions(+), 5 deletions(-) diff --git a/tools/cli/cosyvoice-cli.cpp b/tools/cli/cosyvoice-cli.cpp index ee1f71f..2859d60 100644 --- a/tools/cli/cosyvoice-cli.cpp +++ b/tools/cli/cosyvoice-cli.cpp @@ -124,6 +124,7 @@ enum class cli_log_level }; static bool g_quiet_logs = false; +static std::atomic_bool g_stop_requested = false; inline std::atomic_bool g_interactive_exit_requested = false; static constexpr const char* ANSI_RESET = "\033[0m"; static constexpr const char* ANSI_RED = "\033[31m"; @@ -139,6 +140,7 @@ static BOOL WINAPI handle_console_ctrl(DWORD ctrl_type) { g_interactive_exit_requested.store(true, std::memory_order_relaxed); interrupt_playback(); + g_stop_requested.store(true, std::memory_order_release); // Unregister self: next Ctrl+C triggers default handler (process termination) SetConsoleCtrlHandler(handle_console_ctrl, FALSE); @@ -153,6 +155,7 @@ static void handle_sigint(int) { g_interactive_exit_requested.store(true, std::memory_order_relaxed); interrupt_playback(); + g_stop_requested.store(true, std::memory_order_release); // Restore default handler: next Ctrl+C terminates the process immediately signal(SIGINT, SIG_DFL); @@ -179,6 +182,7 @@ class interactive_ctrl_c_guard void _register() { g_interactive_exit_requested.store(false, std::memory_order_relaxed); + g_stop_requested.store(false, std::memory_order_release); #ifndef COSYVOICE_CLI_NO_PLAYBACK g_playback_interrupted.store(false, std::memory_order_relaxed); #endif @@ -197,6 +201,55 @@ class interactive_ctrl_c_guard #endif }; +// RAII: spawns a monitoring thread during TTS generation that calls +// cosyvoice_request_stop() when Ctrl+C is detected (g_stop_requested). +// The destructor signals the thread to exit and joins it. +class stop_monitor +{ +public: + stop_monitor(cosyvoice_context_t ctx) + { + worker = std::thread([this, ctx]() + { + for (;;) + { + if (stop_thread_flag.load(std::memory_order_acquire)) + break; + if (g_stop_requested.load(std::memory_order_acquire)) + { + if (ctx) + { + cosyvoice_request_stop(ctx); + triggered = true; + } + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + } + }); + } + + ~stop_monitor() + { + stop_thread_flag.store(true, std::memory_order_release); + if (worker.joinable()) + worker.join(); + } + + // True if cosyvoice_request_stop() was actually called. + bool was_stop_requested() + { + if (worker.joinable()) + worker.join(); + return triggered; + } + +private: + bool triggered{false}; + std::atomic stop_thread_flag{false}; + std::thread worker; +}; + static bool has_generation_overrides(const cli_options& options) { return options.has_temperature @@ -1157,6 +1210,7 @@ static void run_interactive_loop( gen_start = std::chrono::steady_clock::now(); bool ok; + stop_monitor sm(ctx); if (options.mode == "cross-lingual") ok = cosyvoice_tts_cross_lingual_stream(tts_ctx, trimmed.c_str(), options.speed, stream_callback, &play_state); else if (options.mode == "zero-shot") @@ -1171,7 +1225,10 @@ static void run_interactive_loop( if (!ok) { - print_error_log("Error: TTS streaming generation failed.\n"); + if (sm.was_stop_requested()) + printf("Generation cancelled.\n"); + else + print_error_log("Error: TTS streaming generation failed.\n"); continue; } } @@ -1179,10 +1236,21 @@ static void run_interactive_loop( { gen_start = std::chrono::steady_clock::now(); std::string error; - auto pcm = generate_tts_audio(tts_ctx, options, trimmed, &error); + cosyvoice_generated_speech pcm{}; + { + stop_monitor sm(ctx); + pcm = generate_tts_audio(tts_ctx, options, trimmed, &error); + if (!pcm.data && !sm.was_stop_requested() && error.empty()) + error = "TTS generation failed."; + } if (!pcm.data) { - if (!error.empty()) + if (g_stop_requested.load(std::memory_order_acquire)) + { + ctrl_c_guard._register(); + printf("Generation cancelled.\n"); + } + else if (!error.empty()) print_error_log("Error: %s\n", error.c_str()); continue; } @@ -2140,8 +2208,24 @@ int tool_entry(int argc, char** argv) stage_start = std::chrono::steady_clock::now(); std::string tts_error; - auto pcm = generate_tts_audio(tts_ctx.get(), options, options.text, &tts_error); - timing.tts_generate_ms = elapsed_ms(stage_start, std::chrono::steady_clock::now()); + cosyvoice_generated_speech pcm = {}; + { + stop_monitor sm(ctx.get()); + // Install a signal handler for non-interactive mode so Ctrl+C stops generation + interactive_ctrl_c_guard ctrl_c_guard; + pcm = generate_tts_audio(tts_ctx.get(), options, options.text, &tts_error); + timing.tts_generate_ms = elapsed_ms(stage_start, std::chrono::steady_clock::now()); + if (!pcm.data) + { + if (sm.was_stop_requested()) + { + if (tts_error.empty()) + tts_error = "TTS generation cancelled."; + } + else if (tts_error.empty()) + tts_error = "TTS generation failed."; + } + } if (!pcm.data) { From 4f520f980884fb488d3448edda477fc41738f953 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Wed, 22 Jul 2026 00:47:39 +0800 Subject: [PATCH 25/34] docs: add streaming TTS and DiT KV cache API documentation(zh) --- docs/API_zh_cosyvoice.md | 329 ++++++++++++++++++++++++++--- docs/API_zh_cosyvoice_interface.md | 215 +++++++++++++++++++ docs/API_zh_cosyvoice_lowlevel.md | 209 ++++++++++++++++++ 3 files changed, 719 insertions(+), 34 deletions(-) diff --git a/docs/API_zh_cosyvoice.md b/docs/API_zh_cosyvoice.md index fa88a81..4d50cb8 100644 --- a/docs/API_zh_cosyvoice.md +++ b/docs/API_zh_cosyvoice.md @@ -29,42 +29,57 @@ 静态链接时定义 `COSYVOICE_STATIC` 可关闭导入属性；Windows 默认使用 `__declspec(dllimport)`。 -## cosyvoice_llm_kv_cache_type_t +## cosyvoice_kv_cache_type_t ### 语法 ```c -typedef enum cosyvoice_llm_kv_cache_type +typedef enum cosyvoice_kv_cache_type { - COSYVOICE_LLM_KV_CACHE_TYPE_F32, - COSYVOICE_LLM_KV_CACHE_TYPE_F16, - COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0, - COSYVOICE_LLM_KV_CACHE_TYPE_Q5_1, - COSYVOICE_LLM_KV_CACHE_TYPE_Q5_0, - COSYVOICE_LLM_KV_CACHE_TYPE_Q4_1, - COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0, - COSYVOICE_LLM_KV_CACHE_TYPE_COUNT -} cosyvoice_llm_kv_cache_type_t; + COSYVOICE_KV_CACHE_TYPE_F32, + COSYVOICE_KV_CACHE_TYPE_F16, + COSYVOICE_KV_CACHE_TYPE_Q8_0, + COSYVOICE_KV_CACHE_TYPE_Q5_1, + COSYVOICE_KV_CACHE_TYPE_Q5_0, + COSYVOICE_KV_CACHE_TYPE_Q4_1, + COSYVOICE_KV_CACHE_TYPE_Q4_0, + COSYVOICE_KV_CACHE_TYPE_COUNT +} cosyvoice_kv_cache_type_t, cosyvoice_llm_kv_cache_type_t; +``` + +#### 向后兼容别名 + +旧的 `cosyvoice_llm_kv_cache_type_t` 和 `COSYVOICE_LLM_KV_CACHE_TYPE_*` 常量以后向兼容宏提供： + +```c +#define COSYVOICE_LLM_KV_CACHE_TYPE_F32 COSYVOICE_KV_CACHE_TYPE_F32 +#define COSYVOICE_LLM_KV_CACHE_TYPE_F16 COSYVOICE_KV_CACHE_TYPE_F16 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0 COSYVOICE_KV_CACHE_TYPE_Q8_0 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q5_1 COSYVOICE_KV_CACHE_TYPE_Q5_1 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q5_0 COSYVOICE_KV_CACHE_TYPE_Q5_0 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q4_1 COSYVOICE_KV_CACHE_TYPE_Q4_1 +#define COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0 COSYVOICE_KV_CACHE_TYPE_Q4_0 +#define COSYVOICE_LLM_KV_CACHE_TYPE_COUNT COSYVOICE_KV_CACHE_TYPE_COUNT ``` ### 说明 -指定 LLM 的 KV 缓存存储格式。 +指定 KV 缓存存储格式。同时用于 LLM 和 DiT 模块。 ### 枚举值 -- `COSYVOICE_LLM_KV_CACHE_TYPE_F32`：32 位浮点格式。 -- `COSYVOICE_LLM_KV_CACHE_TYPE_F16`：16 位浮点格式。 -- `COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0`：GGML `Q8_0` 量化格式。 -- `COSYVOICE_LLM_KV_CACHE_TYPE_Q5_1`：GGML `Q5_1` 量化格式。 -- `COSYVOICE_LLM_KV_CACHE_TYPE_Q5_0`：GGML `Q5_0` 量化格式。 -- `COSYVOICE_LLM_KV_CACHE_TYPE_Q4_1`：GGML `Q4_1` 量化格式。 -- `COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0`：GGML `Q4_0` 量化格式。 -- `COSYVOICE_LLM_KV_CACHE_TYPE_COUNT`：哨兵值，不用于运行配置。 +- `COSYVOICE_KV_CACHE_TYPE_F32`：32 位浮点格式。 +- `COSYVOICE_KV_CACHE_TYPE_F16`：16 位浮点格式。 +- `COSYVOICE_KV_CACHE_TYPE_Q8_0`：GGML `Q8_0` 量化格式。 +- `COSYVOICE_KV_CACHE_TYPE_Q5_1`：GGML `Q5_1` 量化格式。 +- `COSYVOICE_KV_CACHE_TYPE_Q5_0`：GGML `Q5_0` 量化格式。 +- `COSYVOICE_KV_CACHE_TYPE_Q4_1`：GGML `Q4_1` 量化格式。 +- `COSYVOICE_KV_CACHE_TYPE_Q4_0`：GGML `Q4_0` 量化格式。 +- `COSYVOICE_KV_CACHE_TYPE_COUNT`：哨兵值，不用于运行配置。 ### 分离 K/V 缓存宏 -`cosyvoice_llm_kv_cache_type_t` 可以通过位打包（bit-packing）编码 K 和 V 缓存的不同类型， +`cosyvoice_kv_cache_type_t` 可以通过位打包（bit-packing）编码 K 和 V 缓存的不同类型， 从而为 K 和 V 张量使用不同的量化格式（例如 K 用 `Q8_0`、V 用 `Q4_0`），在质量与内存之间灵活取舍。 #### 打包格式（bit 31 = 1 表示分离模式） @@ -77,7 +92,7 @@ typedef enum cosyvoice_llm_kv_cache_type | 15–30 | 保留 | | 31 | 分离标志位 | -当 bit 31 为 0 时，值被当作普通的 `cosyvoice_llm_kv_cache_type_t`，K 和 V 使用同一类型（向后兼容）。 +当 bit 31 为 0 时，值被当作普通的 `cosyvoice_kv_cache_type_t`，K 和 V 使用同一类型（向后兼容）。 ```c #define COSYVOICE_MAKE_SEPARATE_KV_CACHE(k_type, v_type, fallback_type) @@ -97,11 +112,11 @@ typedef enum cosyvoice_llm_kv_cache_type #### 在参数中使用 -可以直接使用 `cosyvoice_context_params_t` 中的位域成员，或通过 `COSYVOICE_MAKE_SEPARATE_KV_CACHE` 打包后赋值给 `llm_kv_cache_type`： +可以直接使用 `cosyvoice_context_params_t` 或 `cosyvoice_context_params_v3_t` 中的位域成员，或通过 `COSYVOICE_MAKE_SEPARATE_KV_CACHE` 打包后赋值给 `llm_kv_cache_type` / `dit_kv_cache_type`： ```c -params.llm_k_cache_type = COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0; -params.llm_v_cache_type = COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0; +params.llm_k_cache_type = COSYVOICE_KV_CACHE_TYPE_Q8_0; +params.llm_v_cache_type = COSYVOICE_KV_CACHE_TYPE_Q4_0; params.llm_kv_cache_separate_buffers = true; ``` @@ -109,9 +124,9 @@ params.llm_kv_cache_separate_buffers = true; ```c params.llm_kv_cache_type = COSYVOICE_MAKE_SEPARATE_KV_CACHE( - COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0, - COSYVOICE_LLM_KV_CACHE_TYPE_Q4_0, - COSYVOICE_LLM_KV_CACHE_TYPE_Q8_0); + COSYVOICE_KV_CACHE_TYPE_Q8_0, + COSYVOICE_KV_CACHE_TYPE_Q4_0, + COSYVOICE_KV_CACHE_TYPE_Q8_0); ``` ## cosyvoice_inference_buffer_policy_t @@ -352,13 +367,13 @@ typedef struct cosyvoice_context_params { struct { - cosyvoice_llm_kv_cache_type_t llm_k_cache_type : 5; ///< K 缓存数据类型。 - cosyvoice_llm_kv_cache_type_t llm_v_cache_type : 5; ///< V 缓存数据类型。 - cosyvoice_llm_kv_cache_type_t llm_kv_cache_fallback : 5; ///< 首选类型不受支持时的回退类型。 - cosyvoice_llm_kv_cache_type_t : 16; ///< 保留位。 - cosyvoice_llm_kv_cache_type_t llm_kv_cache_separate_buffers : 1; ///< 是否分别为 K 和 V 分配独立缓存。 + cosyvoice_kv_cache_type_t llm_k_cache_type : 5; ///< K 缓存数据类型。 + cosyvoice_kv_cache_type_t llm_v_cache_type : 5; ///< V 缓存数据类型。 + cosyvoice_kv_cache_type_t llm_kv_cache_fallback : 5; ///< 首选类型不受支持时的回退类型。 + cosyvoice_kv_cache_type_t : 16; ///< 保留位。 + uint32_t llm_kv_cache_separate_buffers : 1; ///< 是否分别为 K 和 V 分配独立缓存。 }; - cosyvoice_llm_kv_cache_type_t llm_kv_cache_type; ///< 向后兼容的统一类型。 + cosyvoice_kv_cache_type_t llm_kv_cache_type; ///< 向后兼容的统一类型。 }; bool llm_allow_kv_cache_fallback; cosyvoice_inference_buffer_policy_t inference_buffer_policy; @@ -557,6 +572,117 @@ COSYVOICE_API cosyvoice_context_t cosyvoice_load_from_file_with_params_v2( 成功返回上下文句柄，失败返回 `NULL`。 +## cosyvoice_context_params_v3_t + +### 语法 + +```c +typedef struct cosyvoice_context_params_v3 +{ + cosyvoice_context_params_v2_t base_params; + + union + { + struct + { +#if defined(__BYTE_ORDER__) && (__BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) || defined(_BYTE_ORDER) && (_BYTE_ORDER == _BIG_ENDIAN) || defined(__BIG_ENDIAN__) && !defined(__LITTLE_ENDIAN__) || defined(__ARMEB__) || defined(__MIPSEB__) || defined(__sparc__) + uint32_t dit_kv_cache_separate_buffers : 1; + cosyvoice_kv_cache_type_t : 16; + cosyvoice_kv_cache_type_t dit_kv_cache_fallback : 5; + cosyvoice_kv_cache_type_t dit_v_cache_type : 5; + cosyvoice_kv_cache_type_t dit_k_cache_type : 5; +#else + cosyvoice_kv_cache_type_t dit_k_cache_type : 5; + cosyvoice_kv_cache_type_t dit_v_cache_type : 5; + cosyvoice_kv_cache_type_t dit_kv_cache_fallback : 5; + cosyvoice_kv_cache_type_t : 16; + uint32_t dit_kv_cache_separate_buffers : 1; +#endif + }; + cosyvoice_kv_cache_type_t dit_kv_cache_type; + }; + bool dit_allow_kv_cache_fallback; + uint32_t dit_kv_fixed_slots; + uint32_t dit_kv_offloadable_slots; + uint32_t dit_kv_cache_length; +} cosyvoice_context_params_v3_t; + +#ifdef __cplusplus +struct cosyvoice_context_params_v3_cpp : cosyvoice_context_params_v2_cpp +{ + union + { + struct + { +#if defined(__BYTE_ORDER__) && (__BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) || defined(_BYTE_ORDER) && (_BYTE_ORDER == _BIG_ENDIAN) || defined(__BIG_ENDIAN__) && !defined(__LITTLE_ENDIAN__) || defined(__ARMEB__) || defined(__MIPSEB__) || defined(__sparc__) + uint32_t dit_kv_cache_separate_buffers : 1; + cosyvoice_kv_cache_type_t : 16; + cosyvoice_kv_cache_type_t dit_kv_cache_fallback : 5; + cosyvoice_kv_cache_type_t dit_v_cache_type : 5; + cosyvoice_kv_cache_type_t dit_k_cache_type : 5; +#else + cosyvoice_kv_cache_type_t dit_k_cache_type : 5; + cosyvoice_kv_cache_type_t dit_v_cache_type : 5; + cosyvoice_kv_cache_type_t dit_kv_cache_fallback : 5; + cosyvoice_kv_cache_type_t : 16; + uint32_t dit_kv_cache_separate_buffers : 1; +#endif + }; + cosyvoice_kv_cache_type_t dit_kv_cache_type; + }; + bool dit_allow_kv_cache_fallback; + uint32_t dit_kv_fixed_slots; + uint32_t dit_kv_offloadable_slots; + uint32_t dit_kv_cache_length; +}; +#endif +``` + +### 说明 + +在 `cosyvoice_context_params_v2_t` 基础上增加 DiT（扩散模型）KV 缓存配置。流式 TTS 时 DiT 模块会运行多个扩散步，每步计算完整音频序列的自注意力——KV 缓存可以跨步复用注意力计算结果，但缓存本身非常大（最多 `sequence_length × n_diffusion_steps` 个 key-value 对）。 + +### 成员 + +- `base_params`：V2 基础参数。 +- `dit_k_cache_type`：DiT 模块 K 缓存数据类型。 +- `dit_v_cache_type`：DiT 模块 V 缓存数据类型。 +- `dit_kv_cache_separate_buffers`：若为 true，则分别为 K 和 V 分配独立缓存。 +- `dit_kv_cache_fallback`：首选类型不受支持时的回退类型。 +- `dit_kv_cache_type`：快捷方式——指定统一类型（不分离 K/V）。 +- `dit_allow_kv_cache_fallback`：若为 true，不支持时回退到 flash attention 兼容类型。 +- `dit_kv_fixed_slots`：固定（设备内存，从不卸载）DiT KV 槽位数。每个槽位对应一个扩散步的 KV 缓存。 +- `dit_kv_offloadable_slots`：可卸载（CPU 卸载）DiT KV 槽位数。 +- `dit_kv_cache_length`：DiT KV 缓存最大序列长度。0 表示使用默认值（`n_max_seq × 10`）。 + +### DiT KV 缓存概念 + +详见 [README_zh.md — 流式 TTS 与 DiT KV 缓存](#) 的详细说明。 + +## cosyvoice_load_from_file_with_params_v3 + +### 语法 + +```c +COSYVOICE_API cosyvoice_context_t cosyvoice_load_from_file_with_params_v3( + const char* filename, + const cosyvoice_context_params_v3_t* params +); +``` + +### 说明 + +使用 V3 扩展参数加载模型上下文，包含 DiT KV 缓存配置。 + +### 参数 + +- `filename`：模型文件路径。 +- `params`：V3 上下文参数。 + +### 返回值 + +成功返回上下文句柄，失败返回 `NULL`。 + ## cosyvoice_duplicate_context ### 语法 @@ -1506,6 +1632,120 @@ COSYVOICE_API bool cosyvoice_tts_cross_lingual( 成功返回 `true`，失败返回 `false`。 +## cosyvoice_tts_audio_callback_t + +### 语法 + +```c +typedef bool (*cosyvoice_tts_audio_callback_t)(const float* audio, uint32_t n_samples, void* user_data); +``` + +### 说明 + +流式 TTS API 用于逐段交付音频的回调函数类型。 + +### 参数 + +- `audio`：32 位浮点 PCM 采样数据，单声道。 +- `n_samples`：本次回调的采样点数。 +- `user_data`：注册回调时传入的不透明上下文。 + +### 返回值 + +返回 `true` 继续流式合成，返回 `false` 中止合成。 + +## cosyvoice_tts_zero_shot_stream + +### 语法 + +```c +COSYVOICE_API bool cosyvoice_tts_zero_shot_stream( + cosyvoice_tts_context_t ctx, + const char* text, + float speed, + cosyvoice_tts_audio_callback_t callback, + void* user_data +); +``` + +### 说明 + +以 zero-shot 模式流式生成语音，通过回调逐段交付音频。 + +### 参数 + +- `ctx`：TTS 会话。 +- `text`：输入文本。 +- `speed`：语速系数。 +- `callback`：接收每段音频的回调函数。 +- `user_data`：传入回调的不透明上下文。 + +### 返回值 + +成功返回 `true`，失败返回 `false`。 + +## cosyvoice_tts_instruct_stream + +### 语法 + +```c +COSYVOICE_API bool cosyvoice_tts_instruct_stream( + cosyvoice_tts_context_t ctx, + const char* text, + const char* instruction, + float speed, + cosyvoice_tts_audio_callback_t callback, + void* user_data +); +``` + +### 说明 + +以 instruct 模式流式生成语音。 + +### 参数 + +- `ctx`：TTS 会话。 +- `text`：输入文本。 +- `instruction`：指令文本。 +- `speed`：语速系数。 +- `callback`：接收每段音频的回调函数。 +- `user_data`：传入回调的不透明上下文。 + +### 返回值 + +成功返回 `true`，失败返回 `false`。 + +## cosyvoice_tts_cross_lingual_stream + +### 语法 + +```c +COSYVOICE_API bool cosyvoice_tts_cross_lingual_stream( + cosyvoice_tts_context_t ctx, + const char* text, + float speed, + cosyvoice_tts_audio_callback_t callback, + void* user_data +); +``` + +### 说明 + +以 cross-lingual 模式流式生成语音。 + +### 参数 + +- `ctx`：TTS 会话。 +- `text`：输入文本。 +- `speed`：语速系数。 +- `callback`：接收每段音频的回调函数。 +- `user_data`：传入回调的不透明上下文。 + +### 返回值 + +成功返回 `true`，失败返回 `false`。 + ## cosyvoice_save_wav ### 语法 @@ -1585,6 +1825,27 @@ COSYVOICE_API void cosyvoice_get_memory_usage(cosyvoice_context_t ctx, cosyvoice 无返回值。 +## cosyvoice_get_total_memory_usage + +### 语法 + +```c +COSYVOICE_API void cosyvoice_get_total_memory_usage(cosyvoice_context_t ctx, cosyvoice_memory_usage_t* usage); +``` + +### 说明 + +获取所有 worker 的总内存占用快照。 + +### 参数 + +- `ctx`：模型上下文。 +- `usage`：输出内存明细结构体。 + +### 返回值 + +无返回值。 + ## cosyvoice_empty_buffer_cache ### 语法 diff --git a/docs/API_zh_cosyvoice_interface.md b/docs/API_zh_cosyvoice_interface.md index e05f1bb..2ab6ffb 100644 --- a/docs/API_zh_cosyvoice_interface.md +++ b/docs/API_zh_cosyvoice_interface.md @@ -28,6 +28,7 @@ struct cosyvoice_model_context virtual cosyvoice_builtin_sampler_rng_policy_t get_builtin_sampler_rng_policy() = 0; virtual bool set_builtin_sampler_rng_policy(cosyvoice_builtin_sampler_rng_policy_t policy) = 0; virtual bool set_sampler_seed(uint32_t seed) = 0; + virtual uint32_t get_sampler_seed() = 0; virtual bool llm_prefill(ggml_type type, const void* data, uint32_t seq_len) = 0; virtual bool llm_decode(ggml_type type, const void* data) = 0; @@ -38,6 +39,8 @@ struct cosyvoice_model_context virtual uint32_t llm_get_kv_cache_len() = 0; virtual bool llm_set_kv_cache_len(uint32_t len) = 0; + virtual void llm_offload_kv_cache() = 0; + virtual void llm_load_kv_cache() = 0; virtual int llm_sample_token() = 0; virtual bool llm_is_stop_token(int token_id) = 0; @@ -55,6 +58,28 @@ struct cosyvoice_model_context cosyvoice_generated_speech_ptr result ) = 0; + virtual bool llm_job_ext( + const int* text, + uint32_t text_len, + cosyvoice_prompt_t prompt, + uint32_t max_new_tokens, + bool* final + ) = 0; + + virtual bool token2wav_ext( + const int* token_ids, + uint32_t n_tokens, + float speed, + cosyvoice_prompt_t prompt, + uint32_t* offset, + bool streaming, + bool finalize, + cosyvoice_generated_speech_ptr result + ) = 0; + + virtual uint32_t get_chunk_tokens() = 0; + virtual void set_chunk_tokens(uint32_t n_tokens) = 0; + virtual ggml_status get_last_status() = 0; virtual void set_prompt( @@ -65,12 +90,16 @@ struct cosyvoice_model_context ) = 0; virtual void get_memory_usage(cosyvoice_memory_usage_t* usage) = 0; + virtual void get_total_memory_usage(cosyvoice_memory_usage_t* usage) = 0; virtual void empty_buffer_cache() = 0; virtual void set_noise_callback(cosyvoice_noise_callback_t callback, void* callback_ctx) = 0; virtual void get_noise_callback(cosyvoice_noise_callback_t* callback, void** callback_ctx) = 0; virtual uint32_t get_hift_rand_ini_len() = 0; virtual void set_hift_rand_ini(const float* data) = 0; + + virtual void request_stop() = 0; + virtual bool stop_requested() = 0; }; ``` @@ -557,6 +586,38 @@ virtual bool llm_set_kv_cache_len(uint32_t len) = 0; 成功返回 `true`，失败返回 `false`。 +## cosyvoice_model_context::llm_offload_kv_cache + +### 语法 + +```cpp +virtual void llm_offload_kv_cache() = 0; +``` + +### 说明 + +将 KV 缓存卸载到 CPU 内存。 + +### 返回值 + +无返回值。 + +## cosyvoice_model_context::llm_load_kv_cache + +### 语法 + +```cpp +virtual void llm_load_kv_cache() = 0; +``` + +### 说明 + +将 KV 缓存从 CPU 内存加载回后端设备。 + +### 返回值 + +无返回值。 + ## cosyvoice_model_context::llm_sample_token ### 语法 @@ -699,6 +760,36 @@ virtual bool llm_job(const int* text, uint32_t text_len, cosyvoice_prompt_t prom 成功返回 `true`，失败返回 `false`。 +## cosyvoice_model_context::llm_job_ext + +### 语法 + +```cpp +virtual bool llm_job_ext( + const int* text, + uint32_t text_len, + cosyvoice_prompt_t prompt, + uint32_t max_new_tokens, + bool* final +) = 0; +``` + +### 说明 + +运行低层级 LLM 生成，附增额外选项。 + +### 参数 + +- `text`：输入文本 token。 +- `text_len`：token 数量。 +- `prompt`：提示句柄。 +- `max_new_tokens`：最大新生成 token 数。0 表示不生成新 token。 +- `final`：输出参数，指示生成是否已完成（`true`）或可继续生成更多 token（`false`）。 + +### 返回值 + +成功返回 `true`，失败返回 `false`。 + ## cosyvoice_model_context::token2wav ### 语法 @@ -729,6 +820,78 @@ virtual bool token2wav( 成功返回 `true`，失败返回 `false`。 +## cosyvoice_model_context::token2wav_ext + +### 语法 + +```cpp +virtual bool token2wav_ext( + const int* token_ids, + uint32_t n_tokens, + float speed, + cosyvoice_prompt_t prompt, + uint32_t* offset, + bool streaming, + bool finalize, + cosyvoice_generated_speech_ptr result +) = 0; +``` + +### 说明 + +将语音 token 转换为波形数据，附增额外选项。 + +### 参数 + +- `token_ids`：语音 token ID。 +- `n_tokens`：token 数量。 +- `speed`：语速系数。 +- `prompt`：提示句柄。 +- `offset`：输入/输出 token 偏移量，用于增量转换。 +- `streaming`：若为 true，增量转换。 +- `finalize`：若为 true，刷新并完成输出。 +- `result`：输出波形容器。 + +### 返回值 + +成功返回 `true`，失败返回 `false`。 + +## cosyvoice_model_context::get_chunk_tokens + +### 语法 + +```cpp +virtual uint32_t get_chunk_tokens() = 0; +``` + +### 说明 + +获取流式推理时每个分块处理的 token 数。 + +### 返回值 + +当前分块 token 数。 + +## cosyvoice_model_context::set_chunk_tokens + +### 语法 + +```cpp +virtual void set_chunk_tokens(uint32_t n_tokens) = 0; +``` + +### 说明 + +设置流式推理时每个分块处理的 token 数。 + +### 参数 + +- `n_tokens`：每个分块的 token 数。 + +### 返回值 + +无返回值。 + ## cosyvoice_model_context::get_last_status ### 语法 @@ -797,6 +960,26 @@ virtual void get_memory_usage(cosyvoice_memory_usage_t* usage) = 0; 无返回值。 +## cosyvoice_model_context::get_total_memory_usage + +### 语法 + +```cpp +virtual void get_total_memory_usage(cosyvoice_memory_usage_t* usage) = 0; +``` + +### 说明 + +获取所有 worker 的总内存占用。 + +### 参数 + +- `usage`：输出内存占用结构体。 + +### 返回值 + +无返回值。 + ## cosyvoice_model_context::empty_buffer_cache ### 语法 @@ -899,6 +1082,38 @@ virtual void set_hift_rand_ini(const float* data) = 0; 无返回值。 +## cosyvoice_model_context::request_stop + +### 语法 + +```cpp +virtual void request_stop() = 0; +``` + +### 说明 + +请求当前 worker 的当前任务尽快停止。worker 完成当前操作后才响应，截止点之前的输出仍然有效。 + +### 返回值 + +无返回值。 + +## cosyvoice_model_context::stop_requested + +### 语法 + +```cpp +virtual bool stop_requested() = 0; +``` + +### 说明 + +检查并原子清除停止请求标志。标志被原子重置，后续调用返回 `false`。 + +### 返回值 + +若自上次检查以来有停止请求则返回 `true`，否则返回 `false`。 + ## cosyvoice_tokenization_result ### 语法 diff --git a/docs/API_zh_cosyvoice_lowlevel.md b/docs/API_zh_cosyvoice_lowlevel.md index 4778ad3..679b7ab 100644 --- a/docs/API_zh_cosyvoice_lowlevel.md +++ b/docs/API_zh_cosyvoice_lowlevel.md @@ -395,6 +395,38 @@ COSYVOICE_API bool cosyvoice_llm_set_kv_cache_len(cosyvoice_context_t ctx, uint3 `len` 不得大于当前长度。 +## cosyvoice_llm_offload_kv_cache + +### 语法 + +```c +COSYVOICE_API void cosyvoice_llm_offload_kv_cache(cosyvoice_context_t ctx); +``` + +### 说明 + +将 LLM KV 缓存从设备内存卸载到 CPU 内存。 + +### 参数 + +- `ctx`：模型上下文。 + +## cosyvoice_llm_load_kv_cache + +### 语法 + +```c +COSYVOICE_API void cosyvoice_llm_load_kv_cache(cosyvoice_context_t ctx); +``` + +### 说明 + +将 LLM KV 缓存从 CPU 内存加载回后端设备。 + +### 参数 + +- `ctx`：模型上下文。 + ## cosyvoice_llm_sample_token ### 语法 @@ -545,6 +577,38 @@ COSYVOICE_API bool cosyvoice_llm_job( 成功返回 `true`，失败返回 `false`。 +## cosyvoice_llm_job_ext + +### 语法 + +```c +COSYVOICE_API bool cosyvoice_llm_job_ext( + cosyvoice_context_t ctx, + const int* text, + uint32_t text_len, + cosyvoice_prompt_t prompt, + uint32_t max_new_tokens, + bool* final +); +``` + +### 说明 + +运行低层级 LLM 生成，附增额外选项。 + +### 参数 + +- `ctx`：模型上下文。 +- `text`：文本 token ID。 +- `text_len`：`text` 中的 token 数量。 +- `prompt`：提示句柄。 +- `max_new_tokens`：最大新生成 token 数。0 表示不生成新 token。 +- `final`：输出参数，指示生成是否已完成（`true`）或可继续生成更多 token（`false`）。 + +### 返回值 + +成功返回 `true`，失败返回 `false`。 + ## cosyvoice_token2wav ### 语法 @@ -577,6 +641,44 @@ COSYVOICE_API bool cosyvoice_token2wav( 成功返回 `true`，失败返回 `false`。 +## cosyvoice_token2wav_ext + +### 语法 + +```c +COSYVOICE_API bool cosyvoice_token2wav_ext( + cosyvoice_context_t ctx, + const int* token_ids, + uint32_t n_tokens, + float speed, + cosyvoice_prompt_t prompt, + uint32_t* offset, + bool streaming, + bool finalize, + cosyvoice_generated_speech_ptr result +); +``` + +### 说明 + +将语音 token 转换为波形数据，附增额外选项。 + +### 参数 + +- `ctx`：模型上下文。 +- `token_ids`：语音 token ID。 +- `n_tokens`：语音 token 数量。 +- `speed`：语速系数。 +- `prompt`：提示句柄。 +- `offset`：输入/输出 token 偏移量，用于增量转换。 +- `streaming`：若为 true，增量转换（部分分块）。 +- `finalize`：若为 true，刷新并完成输出。 +- `result`：输出波形容器。 + +### 返回值 + +成功返回 `true`，失败返回 `false`。 + ## cosyvoice_tts ### 语法 @@ -613,6 +715,40 @@ COSYVOICE_API bool cosyvoice_tts( 该接口将 LLM 生成与波形转换封装为一次调用。 +## cosyvoice_tts_stream + +### 语法 + +```c +COSYVOICE_API bool cosyvoice_tts_stream( + cosyvoice_context_t ctx, + const int* text, + uint32_t text_len, + float speed, + cosyvoice_prompt_t prompt, + cosyvoice_tts_audio_callback_t callback, + void* user_data +); +``` + +### 说明 + +运行完整低层级 TTS 流水线并流式输出。音频块通过回调逐段交付。 + +### 参数 + +- `ctx`：模型上下文。 +- `text`：文本 token ID。 +- `text_len`：文本 token 数量。 +- `speed`：语速系数。 +- `prompt`：提示句柄。 +- `callback`：接收每段音频的回调函数。 +- `user_data`：传入回调的不透明上下文。 + +### 返回值 + +成功返回 `true`，失败返回 `false`。 + ## cosyvoice_get_tokenizer ### 语法 @@ -888,6 +1024,26 @@ COSYVOICE_API void cosyvoice_get_noise_callback(cosyvoice_context_t ctx, cosyvoi 无返回值。 +## cosyvoice_get_chunk_tokens + +### 语法 + +```c +COSYVOICE_API uint32_t cosyvoice_get_chunk_tokens(cosyvoice_context_t ctx); +``` + +### 说明 + +获取流式推理时每个分块处理的 token 数。 + +### 参数 + +- `ctx`：模型上下文。 + +### 返回值 + +当前分块 token 数。 + ## cosyvoice_get_hift_rand_ini_len ### 语法 @@ -908,6 +1064,23 @@ COSYVOICE_API uint32_t cosyvoice_get_hift_rand_ini_len(cosyvoice_context_t ctx); 所需采样点数量。 +## cosyvoice_set_chunk_tokens + +### 语法 + +```c +COSYVOICE_API void cosyvoice_set_chunk_tokens(cosyvoice_context_t ctx, uint32_t n_tokens); +``` + +### 说明 + +设置流式推理时每个分块处理的 token 数。 + +### 参数 + +- `ctx`：模型上下文。 +- `n_tokens`：每个分块的 token 数。越小首块延迟越低但开销越大；越大 RTF 越低但首块延迟越高。 + ## cosyvoice_set_hift_rand_ini ### 语法 @@ -929,6 +1102,42 @@ COSYVOICE_API void cosyvoice_set_hift_rand_ini(cosyvoice_context_t ctx, const fl 无返回值。 +## cosyvoice_request_stop + +### 语法 + +```c +COSYVOICE_API void cosyvoice_request_stop(cosyvoice_context_t ctx); +``` + +### 说明 + +请求当前 worker 的当前任务尽快停止。停止是优雅的——worker 会完成当前操作后才响应请求，因此截止点之前的输出仍然有效。 + +### 参数 + +- `ctx`：模型上下文。 + +## cosyvoice_stop_requested + +### 语法 + +```c +COSYVOICE_API bool cosyvoice_stop_requested(cosyvoice_context_t ctx); +``` + +### 说明 + +检查并原子清除当前 worker 的停止请求标志。内部热路径（`llm_job_ext`、`token2wav_ext`、`tts`）使用此函数检测停止请求。标志被原子重置，后续调用返回 `false`。 + +### 参数 + +- `ctx`：模型上下文。 + +### 返回值 + +若自上次检查以来有停止请求则返回 `true`，否则返回 `false`。 + ## cosyvoice_prompt_speech_get_crc32 ### 语法 From a63535279ef048a2f76deaa74abbbf30c284dcc2 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Wed, 22 Jul 2026 00:48:19 +0800 Subject: [PATCH 26/34] docs: add stop-request API docs --- docs/API_cosyvoice_interface.md | 35 ++++++++++++++++++++++++++++++++ docs/API_cosyvoice_lowlevel.md | 36 +++++++++++++++++++++++++++++++++ docs/TOOLS.md | 1 + docs/TOOLS_zh.md | 1 + 4 files changed, 73 insertions(+) diff --git a/docs/API_cosyvoice_interface.md b/docs/API_cosyvoice_interface.md index a080bf3..d2758f5 100644 --- a/docs/API_cosyvoice_interface.md +++ b/docs/API_cosyvoice_interface.md @@ -97,6 +97,9 @@ struct cosyvoice_model_context virtual void get_noise_callback(cosyvoice_noise_callback_t* callback, void** callback_ctx) = 0; virtual uint32_t get_hift_rand_ini_len() = 0; virtual void set_hift_rand_ini(const float* data) = 0; + + virtual void request_stop() = 0; + virtual bool stop_requested() = 0; }; ``` @@ -1027,6 +1030,38 @@ Sets HiFT initialization-noise buffer. No return value. +## cosyvoice_model_context::request_stop + +### Syntax + +```cpp +virtual void request_stop() = 0; +``` + +### Description + +Requests that the active worker's current job stop as soon as possible. The worker finishes the current operation before honouring the request, so output up to that point remains valid. + +### Returns + +No return value. + +## cosyvoice_model_context::stop_requested + +### Syntax + +```cpp +virtual bool stop_requested() = 0; +``` + +### Description + +Checks and atomically clears the stop-requested flag for the active worker. Used internally by hot paths to detect stop requests. The flag is atomically reset so subsequent calls return `false`. + +### Returns + +`true` if a stop was requested since the last check; `false` otherwise. + ## cosyvoice_tokenization_result ### Syntax diff --git a/docs/API_cosyvoice_lowlevel.md b/docs/API_cosyvoice_lowlevel.md index 31026d5..63fc45d 100644 --- a/docs/API_cosyvoice_lowlevel.md +++ b/docs/API_cosyvoice_lowlevel.md @@ -1098,6 +1098,42 @@ Overrides HiFT initialization noise buffer. No return value. +## cosyvoice_request_stop + +### Syntax + +```c +COSYVOICE_API void cosyvoice_request_stop(cosyvoice_context_t ctx); +``` + +### Description + +Requests that the active worker's current job stop as soon as possible. The stop is graceful — the worker will finish the current operation before honouring the request, so the output up to that point remains valid. + +### Parameters + +- `ctx`: Context handle. + +## cosyvoice_stop_requested + +### Syntax + +```c +COSYVOICE_API bool cosyvoice_stop_requested(cosyvoice_context_t ctx); +``` + +### Description + +Checks and atomically clears the stop-requested flag for the active worker. Used internally by hot paths (`llm_job_ext`, `token2wav_ext`, `tts`) to detect stop requests. The flag is atomically reset so subsequent calls return `false`. + +### Parameters + +- `ctx`: Context handle. + +### Returns + +`true` if a stop was requested since the last check; `false` otherwise. + ## cosyvoice_prompt_speech_get_crc32 ### Syntax diff --git a/docs/TOOLS.md b/docs/TOOLS.md index 8014303..765bc20 100644 --- a/docs/TOOLS.md +++ b/docs/TOOLS.md @@ -715,6 +715,7 @@ When neither `--text` nor `--output` is provided, interactive mode is enabled au Type text at the `> ` prompt to synthesize, or use slash commands: - `/play [code]`: Play cached audio (requires audio support; disabled with `COSYVOICE_CLI_NO_PLAYBACK=ON`). Press `Ctrl+C` to stop playback. +- Press `Ctrl+C` during TTS generation to gracefully stop synthesis mid-generation via the stop-request API — output up to that point remains valid. - `/save [code]`: Save cached audio to file. - `/list`: List cached audio. - `/query [code]`: Show audio details. diff --git a/docs/TOOLS_zh.md b/docs/TOOLS_zh.md index f17e6aa..d657e91 100644 --- a/docs/TOOLS_zh.md +++ b/docs/TOOLS_zh.md @@ -693,6 +693,7 @@ cosyvoice-cli \ 在 `> ` 提示符下输入文本即可合成，或使用斜杠命令： - `/play [code]`：播放缓存音频（需要音频支持；`COSYVOICE_CLI_NO_PLAYBACK=ON` 时不可用）。按 `Ctrl+C` 可停止播放。 +- 在 TTS 生成过程中按 `Ctrl+C` 可通过 stop-request API 优雅地停止合成——终止点之前的输出仍然有效。 - `/save [code]`：保存缓存音频到文件。 - `/list`：列出缓存音频。 - `/query [code]`：查看音频详情。 From 95812e0e67778fadbc93392708f816e8ab068ee2 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Thu, 23 Jul 2026 16:14:03 +0800 Subject: [PATCH 27/34] cosyvoice: fix SIGSEGV in tts_job when using streaming with long text --- src/cosyvoice-tts.cpp | 58 +++++++++++++++++++++++++++---------------- 1 file changed, 36 insertions(+), 22 deletions(-) diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index 9a36ee3..936f98c 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -833,51 +833,65 @@ struct cosyvoice_tts_context : cosyvoice_tokenization_result_impl, cosyvoice_pro if (!cosyvoice_tts_stream_with_postprocess(get_tokens(), get_n_tokens(), speed, callback, user_data)) return false; } + return true; } // Slow-split path: tokenize each fragment only for counting, // reassemble text chunks, then re-tokenize each chunk. cosyvoice_tokenization_result_impl tokens_scratch; const auto token_count = [&](std::string_view fragment) -> std::size_t - { - tokens_scratch.tokens.clear(); - ctx->tokenize(fragment.data(), static_cast(fragment.size()), &tokens_scratch, true); - return tokens_scratch.get_n_tokens(); - }; + { + tokens_scratch.tokens.clear(); + ctx->tokenize(fragment.data(), static_cast(fragment.size()), &tokens_scratch, true); + return tokens_scratch.get_n_tokens(); + }; const auto chunks = cosyvoice_internal::reassemble_by_token_budget( fragments, max_text_tokens, token_count); if (chunks.size() <= 1) { ctx->tokenize(effective_text, this, true); - return cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, result); + return result ? cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, result) + : cosyvoice_tts_stream_with_postprocess(get_tokens(), get_n_tokens(), speed, callback, user_data); } // Multi-chunk path: synthesize each chunk and copy its PCM into a context-owned buffer. // cosyvoice_tts() points result->data at an internal token2wav buffer that is overwritten // on the next call, so the copy must happen before the following synthesis begins. combined_pcm.clear(); - for (const auto& chunk : chunks) + if (result) { - if (ctx->stop_requested()) + for (const auto& chunk : chunks) { - result->data = nullptr; - result->length = 0; - return false; + if (ctx->stop_requested()) + { + result->data = nullptr; + result->length = 0; + return false; + } + ctx->tokenize(chunk.c_str(), this, true); + cosyvoice_generated_speech part = {}; + if (!cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, &part) + || !part.data || part.length == 0) + { + result->data = nullptr; + result->length = 0; + return false; + } + combined_pcm.insert(combined_pcm.end(), part.data, part.data + part.length); } - ctx->tokenize(chunk.c_str(), this, true); - cosyvoice_generated_speech part = {}; - if (!cosyvoice_tts_with_postprocess(get_tokens(), get_n_tokens(), speed, &part) - || !part.data || part.length == 0) + result->data = combined_pcm.data(); + result->length = static_cast(combined_pcm.size()); + } + else + for (const auto& chunk : chunks) { - result->data = nullptr; - result->length = 0; - return false; + if (ctx->stop_requested()) + return false; + ctx->tokenize(chunk.c_str(), this, true); + if (!cosyvoice_tts_stream_with_postprocess(get_tokens(), get_n_tokens(), speed, callback, user_data)) + return false; } - combined_pcm.insert(combined_pcm.end(), part.data, part.data + part.length); - } - result->data = combined_pcm.data(); - result->length = static_cast(combined_pcm.size()); return true; } From 9cd8778f09860871f2fccef43a66393a679051ba Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Fri, 24 Jul 2026 18:08:47 +0800 Subject: [PATCH 28/34] cosyvoice: use HiFT overlap for flow cache slicing and streaming callback --- src/cosyvoice-graph.cpp | 8 ++++++++ src/cosyvoice-loader.cpp | 1 + src/cosyvoice-model.h | 1 + src/cosyvoice-modules.h | 1 + src/cosyvoice-tts.cpp | 15 ++++++++++++--- src/cosyvoice.cpp | 6 ++---- 6 files changed, 25 insertions(+), 7 deletions(-) diff --git a/src/cosyvoice-graph.cpp b/src/cosyvoice-graph.cpp index 0ac8442..cfc7eab 100644 --- a/src/cosyvoice-graph.cpp +++ b/src/cosyvoice-graph.cpp @@ -805,6 +805,14 @@ std::array CausalHiFTGenerator::build_cgraph(ggml_context* ctx, } } +int CausalHiFTGenerator::overlap_length() const +{ + const int c0 = f0_predictor.condnet_0.causal_padding(); + const int cp = conv_pre.causal_padding(); + const int stft_overlap = (nfft + scale_factor - 1) / scale_factor; + return std::max(c0 + cp + 1, stft_overlap); +} + ggml_tensor* Qwen2MLP::build_cgraph(ggml_context* ctx, ggml_tensor* x) const { auto gate = gate_proj.build_cgraph(ctx, x); diff --git a/src/cosyvoice-loader.cpp b/src/cosyvoice-loader.cpp index 533a455..2dd85d4 100644 --- a/src/cosyvoice-loader.cpp +++ b/src/cosyvoice-loader.cpp @@ -980,6 +980,7 @@ void cosyvoice_model_3::load(gguf_loader& loader) auto arch = loader.get_string("general.architecture"); shared->architecture.reset(new char[arch.size() + 1]); memcpy(shared->architecture.get(), arch.data(), arch.size() + 1); + cv3_shared->hift_overlap = hift.overlap_length(); ggml_backend_synchronize(backend); } diff --git a/src/cosyvoice-model.h b/src/cosyvoice-model.h index d964b31..c02b587 100644 --- a/src/cosyvoice-model.h +++ b/src/cosyvoice-model.h @@ -158,6 +158,7 @@ struct cosyvoice_model_3_shared std::set stop_tokens; std::set silent_tokens; + int hift_overlap; }; struct cosyvoice_3_worker_context diff --git a/src/cosyvoice-modules.h b/src/cosyvoice-modules.h index ea43c68..4b25c28 100644 --- a/src/cosyvoice-modules.h +++ b/src/cosyvoice-modules.h @@ -345,6 +345,7 @@ struct CausalHiFTGenerator : Module std::array build_cgraph(ggml_context* ctx, ggml_tensor* speech_feat, bool finalize) const; + int overlap_length() const; void set_rand_ini(const float* data) const; float lrelu_slope; diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index 936f98c..aba17a4 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -638,9 +638,18 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f if (!finalize) { - auto n_cached_elements = cache_length * feat->ne[0]; - worker->flow_cache.resize(n_cached_elements + static_cast(ggml_nelements(feat))); - ggml_backend_tensor_get_async(backend.get(), feat, worker->flow_cache.data() + n_cached_elements, 0, feat->nb[2]); + const auto overlap = static_cast(cv3_shared->hift_overlap); + const auto n_feat_frames = feat->ne[1]; + const auto frames_to_keep = std::min(overlap, n_feat_frames); + const auto elements_to_keep = static_cast(frames_to_keep * feat->ne[0]); + + worker->flow_cache.resize(elements_to_keep); + if (frames_to_keep > 0) + { + const auto byte_offset = static_cast((n_feat_frames - frames_to_keep) * feat->nb[1]); + const auto byte_size = static_cast(frames_to_keep * feat->nb[1]); + ggml_backend_tensor_get_async(backend.get(), feat, worker->flow_cache.data(), byte_offset, byte_size); + } } if (config[flow.decoder.diffusion_steps - 1].cache_kv) diff --git a/src/cosyvoice.cpp b/src/cosyvoice.cpp index 2b1c72a..074285d 100644 --- a/src/cosyvoice.cpp +++ b/src/cosyvoice.cpp @@ -423,7 +423,6 @@ bool cosyvoice_tts_stream(cosyvoice_context_t ctx, const int* text, uint32_t tex ctx->llm_clear_accepted_tokens(); uint32_t offset = 0; uint32_t n_tokens = 0; - uint32_t last_length = 0; bool final = false; do { @@ -443,10 +442,9 @@ bool cosyvoice_tts_stream(cosyvoice_context_t ctx, const int* text, uint32_t tex if (!ctx->token2wav_ext(ctx->llm_get_accepted_tokens(), n_tokens, speed, prompt, &offset, true, ctx->llm_get_n_accepted_tokens() == n_tokens, &result)) return false; - if (result.data && result.length > last_length - && !callback(result.data + last_length, result.length - last_length, user_data)) + if (result.data && result.length > 0 + && !callback(result.data, result.length, user_data)) return false; - last_length = result.length; } while (n_tokens != ctx->llm_get_n_accepted_tokens()); return true; From f883fe26ba7bc3a3870014b28c58faf35fcbf71e Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Fri, 24 Jul 2026 22:06:01 +0800 Subject: [PATCH 29/34] cli: fix stop_monitor race conditions --- tools/cli/cosyvoice-cli.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tools/cli/cosyvoice-cli.cpp b/tools/cli/cosyvoice-cli.cpp index 2859d60..4c2abb1 100644 --- a/tools/cli/cosyvoice-cli.cpp +++ b/tools/cli/cosyvoice-cli.cpp @@ -213,8 +213,6 @@ class stop_monitor { for (;;) { - if (stop_thread_flag.load(std::memory_order_acquire)) - break; if (g_stop_requested.load(std::memory_order_acquire)) { if (ctx) @@ -224,6 +222,8 @@ class stop_monitor } break; } + if (stop_thread_flag.load(std::memory_order_acquire)) + break; std::this_thread::sleep_for(std::chrono::milliseconds(50)); } }); @@ -240,7 +240,10 @@ class stop_monitor bool was_stop_requested() { if (worker.joinable()) + { + stop_thread_flag.store(true, std::memory_order_release); worker.join(); + } return triggered; } From c5a3b9fd8d9c6183a412ca9f8376ca246381f0f3 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Sat, 25 Jul 2026 10:34:40 +0800 Subject: [PATCH 30/34] cosyvoice: remove unused code and refactor Snake epsilon handling --- src/cosyvoice-kv-cache.cpp | 4 +--- src/cosyvoice-kv-cache.h | 2 +- src/cosyvoice-loader.cpp | 49 ++++++++++++++++++++++++++++++++------ src/cosyvoice-model.cpp | 32 ------------------------- src/cosyvoice-model.h | 5 +--- src/cosyvoice-tts.cpp | 2 +- 6 files changed, 46 insertions(+), 48 deletions(-) diff --git a/src/cosyvoice-kv-cache.cpp b/src/cosyvoice-kv-cache.cpp index b3cc2cb..6ffbd26 100644 --- a/src/cosyvoice-kv-cache.cpp +++ b/src/cosyvoice-kv-cache.cpp @@ -118,11 +118,10 @@ ggml_backend_buffer* cosyvoice_kv_cache::initialize_buffer(ggml_backend_t backen return ggml_backend_alloc_ctx_tensors(ctx, backend); } -uint32_t cosyvoice_kv_cache::reset_buffer(ggml_backend_buffer* buffer) +void cosyvoice_kv_cache::reset_buffer(ggml_backend_buffer* buffer) { cur_len = 0; auto alignment = ggml_backend_buffer_get_alignment(buffer); - auto max_seq_len = static_cast(kv_cache_layers[0].k->ne[1]); auto k_ne = *reinterpret_cast*>(&kv_cache_layers[0].k->ne); auto v_ne = *reinterpret_cast*>(&kv_cache_layers[0].v->ne); @@ -151,7 +150,6 @@ uint32_t cosyvoice_kv_cache::reset_buffer(ggml_backend_buffer* buffer) ggml_backend_tensor_alloc(buffer, v, buffer_base); buffer_base += get_aligned_size(ggml_backend_buffer_get_alloc_size(buffer, v), alignment); } - return max_seq_len; } void cosyvoice_kv_cache::offload_cache(ggml_backend_t backend, ggml_backend_sched* sched, uint32_t n_tokens) diff --git a/src/cosyvoice-kv-cache.h b/src/cosyvoice-kv-cache.h index 909f6b5..43c472a 100644 --- a/src/cosyvoice-kv-cache.h +++ b/src/cosyvoice-kv-cache.h @@ -37,7 +37,7 @@ class cosyvoice_kv_cache void slide_kv_slot(); ggml_backend_buffer* initialize_buffer(ggml_backend_t backend, int k_head_dim, int v_head_dim, uint32_t max_seq, int batch_size); - uint32_t reset_buffer(ggml_backend_buffer* buffer); + void reset_buffer(ggml_backend_buffer* buffer); void offload_cache(ggml_backend_t backend, ggml_backend_sched* sched, uint32_t n_tokens); void load_cache(ggml_backend_t backend, ggml_backend_sched* sched); diff --git a/src/cosyvoice-loader.cpp b/src/cosyvoice-loader.cpp index 2dd85d4..c102a85 100644 --- a/src/cosyvoice-loader.cpp +++ b/src/cosyvoice-loader.cpp @@ -155,6 +155,24 @@ bool backend_looks_uma(ggml_backend_t backend, ggml_backend_buffer* buffer) #define LOAD_METADATA(name) GGML_ASSERT(loader.get_metadata(prefix, #name, name)) #define LOAD_METADATA_NOPREFIX(name) GGML_ASSERT(loader.get_metadata(#name, name)) +constexpr uint32_t GGML_TENSOR_FLAG_NEED_EPS_ADJUST = 28u << 0; + +static inline void set_tensor_eps(ggml_tensor* tensor, float eps) +{ + GGML_ASSERT(tensor && tensor->type == GGML_TYPE_F32); + + tensor->flags |= GGML_TENSOR_FLAG_NEED_EPS_ADJUST; + *reinterpret_cast(tensor->op_params) = eps; +} + +static inline float get_tensor_eps(const ggml_tensor* tensor) +{ + return (tensor->flags & GGML_TENSOR_FLAG_NEED_EPS_ADJUST) + ? *reinterpret_cast(tensor->op_params) + : 0.0f; +} + + void Module::OnLoad(const gguf_loader& loader, const std::string& prefix) {} void BasicModule::OnLoad(const gguf_loader& loader, const std::string& prefix) @@ -293,10 +311,7 @@ void Snake::OnLoad(const gguf_loader& loader, const std::string& prefix) { LOAD_TENSOR(alpha); - constexpr float epsilon = 0.000000001f; - for (auto& i : std::span(reinterpret_cast(ggml_get_data(alpha)), ggml_nelements(alpha))) - if (std::abs(i) < epsilon) - i = epsilon; + set_tensor_eps(alpha, 1e-6f); } void ResBlock::OnLoad(const gguf_loader& loader, const std::string& prefix) @@ -797,7 +812,28 @@ void cosyvoice_model_3::load(gguf_loader& loader) else { ggml_backend_tensor_alloc(shared->buffer.get(), new_tensor, buffer_base); - set_tensor(new_tensor, ggml_get_data(*tensor), tensor_size); + + const auto eps = get_tensor_eps(*tensor); + if (eps == 0.f) + set_tensor(new_tensor, ggml_get_data(*tensor), tensor_size); + else + { + auto data = reinterpret_cast(ggml_get_data(*tensor)); + const auto nelements = ggml_nelements(*tensor); + std::unique_ptr buffer; + for (int64_t i = 0; i != nelements; ++i) + if (std::abs(data[i]) < eps) + { + if (!buffer) + { + buffer.reset(new float[nelements]); + memcpy(buffer.get(), data, sizeof(float) * nelements); + } + + buffer[i] = eps; + } + set_tensor(new_tensor, buffer ? buffer.get() : data, tensor_size); + } } *tensor = new_tensor; @@ -974,13 +1010,12 @@ void cosyvoice_model_3::load(gguf_loader& loader) default: throw std::invalid_argument("unexpected policy"); } - cv3_worker->orig_max_seq_len = shared->params.n_max_seq; } auto arch = loader.get_string("general.architecture"); shared->architecture.reset(new char[arch.size() + 1]); memcpy(shared->architecture.get(), arch.data(), arch.size() + 1); - cv3_shared->hift_overlap = hift.overlap_length(); + shared->hift_overlap = hift.overlap_length(); ggml_backend_synchronize(backend); } diff --git a/src/cosyvoice-model.cpp b/src/cosyvoice-model.cpp index 860a0d4..d4bf4eb 100644 --- a/src/cosyvoice-model.cpp +++ b/src/cosyvoice-model.cpp @@ -265,37 +265,6 @@ void cosyvoice_model::empty_buffer_cache() worker->prompt_crc32 = 0; } -void cosyvoice_model_3::empty_buffer_cache() -{ - cosyvoice_model::empty_buffer_cache(); - - if (cv3_worker->orig_max_seq_len != shared->params.n_max_seq) - { - shared->params.n_max_seq = cv3_worker->orig_max_seq_len; - worker->llm_kv_buffer.reset( - worker->llm_kv_cache.initialize_buffer( - worker->backend.get(), - static_cast(cv3_shared->llm.layers[0].self_attn.k_proj.weight->ne[1] / cv3_shared->llm.num_key_value_heads), - static_cast(cv3_shared->llm.layers[0].self_attn.v_proj.weight->ne[1] / cv3_shared->llm.num_key_value_heads), - shared->params.n_max_seq, - 1)); - - switch (shared->params.inference_buffer_policy) - { - case COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED: - case COSYVOICE_INFERENCE_BUFFER_POLICY_SHARED: - cv3_worker->token2wav_buffer.release(); - cv3_worker->token2wav_buffer.reset(worker->llm_kv_buffer.get()); - break; - case COSYVOICE_INFERENCE_BUFFER_POLICY_DEDICATED: - cv3_worker->token2wav_buffer.reset(); - break; - default: - GGML_ABORT("unexpected policy"); - } - } -} - void cosyvoice_model_3::get_memory_usage(cosyvoice_memory_usage_t* usage) { usage->parameters = ggml_backend_buffer_get_size(shared->buffer.get()); @@ -331,7 +300,6 @@ void cosyvoice_model_3::reset_shared_buffer(ggml_backend_buffer* new_buffer) { worker->llm_kv_buffer.release(); worker->llm_kv_buffer.reset(new_buffer); - shared->params.n_max_seq = worker->llm_kv_cache.reset_buffer(new_buffer); } } diff --git a/src/cosyvoice-model.h b/src/cosyvoice-model.h index c02b587..c4677ae 100644 --- a/src/cosyvoice-model.h +++ b/src/cosyvoice-model.h @@ -26,6 +26,7 @@ struct cosyvoice_model_shared cosyvoice_context_params_v3_cpp params; ggml_backend_op_capabilities op_caps; bool backend_uma; + int hift_overlap; std::mt19937 noise_rng; std::shared_mutex noise_mutex; @@ -158,7 +159,6 @@ struct cosyvoice_model_3_shared std::set stop_tokens; std::set silent_tokens; - int hift_overlap; }; struct cosyvoice_3_worker_context @@ -168,8 +168,6 @@ struct cosyvoice_3_worker_context ggml_context_ptr ctx1; ggml_backend_buffer_ptr token2wav_buffer; - - uint32_t orig_max_seq_len; }; struct cosyvoice_model_3 : cosyvoice_model @@ -200,7 +198,6 @@ struct cosyvoice_model_3 : cosyvoice_model bool token2wav(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, cosyvoice_generated_speech_ptr result); bool token2wav_ext(const int* token_ids, uint32_t n_tokens, float speed, cosyvoice_prompt_t prompt, uint32_t* offset, bool streaming, bool finalize, cosyvoice_generated_speech_ptr result); - void empty_buffer_cache(); void get_memory_usage(cosyvoice_memory_usage_t* usage); void get_total_memory_usage(cosyvoice_memory_usage_t* usage); void reset_shared_buffer(ggml_backend_buffer* new_buffer); diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index aba17a4..d407b2c 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -638,7 +638,7 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f if (!finalize) { - const auto overlap = static_cast(cv3_shared->hift_overlap); + const auto overlap = static_cast(shared->hift_overlap); const auto n_feat_frames = feat->ne[1]; const auto frames_to_keep = std::min(overlap, n_feat_frames); const auto elements_to_keep = static_cast(frames_to_keep * feat->ne[0]); From 6ca298a31642a8391bdacbb1cf5a76873f8922dc Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Sat, 25 Jul 2026 11:54:39 +0800 Subject: [PATCH 31/34] cosyvoice: support shared inference buffer policy in LLM streaming --- src/cosyvoice-tts.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index d407b2c..ba9f30d 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -32,7 +32,8 @@ bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoic auto last_prompt_crc32 = prompt_crc32; bool stop_reached = false; - if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) + if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED + || params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_SHARED && max_new_tokens != UINT32_MAX) { ggml_backend_sched_reset(sched.get()); worker->llm_kv_cache.load_cache(worker->backend.get(), sched.get()); @@ -230,6 +231,8 @@ bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoic llm_set_kv_cache_len(1 + static_cast(prompt->prompt_text.size())); llm_offload_kv_cache(); } + else if (params.inference_buffer_policy != COSYVOICE_INFERENCE_BUFFER_POLICY_DEDICATED && max_new_tokens != UINT32_MAX && !stop_reached) + llm_offload_kv_cache(); // Reset RNG on stop (matches non-streaming behavior) if (stop_reached && params.builtin_sampler_rng_policy == COSYVOICE_BUILTIN_SAMPLER_RNG_POLICY_RESET_PER_SESSION) From 71157b62df1af78e5c146f244e2f41d4e5196bd5 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Sat, 25 Jul 2026 14:35:30 +0800 Subject: [PATCH 32/34] cosyvoice: correct LLM flash attention arg and refactor get_total_memory_usage --- src/cosyvoice-loader.cpp | 12 +----------- src/cosyvoice-model.cpp | 35 +++++++++++++++++++++++++---------- 2 files changed, 26 insertions(+), 21 deletions(-) diff --git a/src/cosyvoice-loader.cpp b/src/cosyvoice-loader.cpp index c102a85..fb0721c 100644 --- a/src/cosyvoice-loader.cpp +++ b/src/cosyvoice-loader.cpp @@ -934,7 +934,7 @@ void cosyvoice_model_3::load(gguf_loader& loader) workers[i].config = shared->config; auto [llm_k_type, llm_v_type] = cosyvoice_check_kv_cache_types(worker->ctx0.get(), worker->backend.get(), - shared->params.flow_use_flash_attn, cv3_shared->llm.num_attention_heads, cv3_shared->llm.num_key_value_heads, + shared->params.llm_use_flash_attn, cv3_shared->llm.num_attention_heads, cv3_shared->llm.num_key_value_heads, cv3_shared->llm.layers[0].self_attn.q_proj, cv3_shared->llm.layers[0].self_attn.k_proj, cv3_shared->llm.layers[0].self_attn.v_proj, shared->params.llm_allow_kv_cache_fallback, reinterpret_cast(shared->params.llm_kv_cache_type)); @@ -1019,13 +1019,3 @@ void cosyvoice_model_3::load(gguf_loader& loader) ggml_backend_synchronize(backend); } - -bool cosyvoice_model_3::set_worker_no(uint32_t worker_no) -{ - if (worker_no >= shared->params.n_workers) - return false; - - worker = workers + worker_no; - cv3_worker = cv3_workers + worker_no; - return true; -} diff --git a/src/cosyvoice-model.cpp b/src/cosyvoice-model.cpp index d4bf4eb..e28f8d7 100644 --- a/src/cosyvoice-model.cpp +++ b/src/cosyvoice-model.cpp @@ -278,19 +278,24 @@ void cosyvoice_model_3::get_memory_usage(cosyvoice_memory_usage_t* usage) void cosyvoice_model_3::get_total_memory_usage(cosyvoice_memory_usage_t* usage) { - usage->parameters = ggml_backend_buffer_get_size(shared->buffer.get()); + cosyvoice_memory_usage_t cur; + const auto worker_no = get_worker_no(); + + *usage = {}; for (uint32_t i = 0; i != shared->params.n_workers; ++i) { - auto worker = workers + i; - auto cv3_worker = cv3_workers + i; - - usage->buffers = ggml_backend_sched_get_buffer_size(worker->sched.get(), worker->backend.get()); - usage->cpu_buffers = ggml_backend_sched_get_buffer_size(worker->sched.get(), worker->cpu_backend.get()); - usage->offloaded_kv_cache = worker->llm_kv_cache.get_offloaded_cache_size(); - usage->random_noise = sizeof(float) * shared->rand_noise_len; - usage->kv_cache = worker->llm_kv_buffer.get() ? ggml_backend_buffer_get_size(worker->llm_kv_buffer.get()) : 0; - usage->token2wav = cv3_worker->token2wav_buffer.get() ? ggml_backend_buffer_get_size(cv3_worker->token2wav_buffer.get()) : 0; + set_worker_no(i); + get_memory_usage(&cur); + + usage->buffers += cur.buffers; + usage->cpu_buffers += cur.cpu_buffers; + usage->offloaded_kv_cache += cur.offloaded_kv_cache; + usage->random_noise += cur.random_noise; + usage->kv_cache += cur.kv_cache; + usage->token2wav += cur.token2wav; } + usage->parameters = cur.parameters; + set_worker_no(worker_no); } void cosyvoice_model_3::reset_shared_buffer(ggml_backend_buffer* new_buffer) @@ -522,6 +527,16 @@ uint32_t cosyvoice_model::get_worker_no() return static_cast(worker - workers); } +bool cosyvoice_model_3::set_worker_no(uint32_t worker_no) +{ + if (worker_no >= shared->params.n_workers) + return false; + + worker = workers + worker_no; + cv3_worker = cv3_workers + worker_no; + return true; +} + uint32_t cosyvoice_model::get_n_workers() { return shared->params.n_workers; From 3009cc5cebda37ef7dfb8685790a606c98cd51da Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Sat, 25 Jul 2026 15:25:54 +0800 Subject: [PATCH 33/34] cosyvoice: guard uninitialized kv cache destructor and skip dit build when slots are zero --- src/cosyvoice-kv-cache.cpp | 2 ++ src/cosyvoice-loader.cpp | 35 +++++++++++++++++++---------------- src/cosyvoice-model.cpp | 2 +- 3 files changed, 22 insertions(+), 17 deletions(-) diff --git a/src/cosyvoice-kv-cache.cpp b/src/cosyvoice-kv-cache.cpp index 6ffbd26..c6fbf0f 100644 --- a/src/cosyvoice-kv-cache.cpp +++ b/src/cosyvoice-kv-cache.cpp @@ -301,6 +301,8 @@ void cosyvoice_kv_cache::load_slot(ggml_backend_t backend, ggml_backend_sched* s cosyvoice_kv_cache::~cosyvoice_kv_cache() { + if (!ctx) return; + ggml_free(ctx); delete[] kv_cache_layers; auto object_size = get_offloaded_kv_cache_struct_size(layers); diff --git a/src/cosyvoice-loader.cpp b/src/cosyvoice-loader.cpp index fb0721c..72aee11 100644 --- a/src/cosyvoice-loader.cpp +++ b/src/cosyvoice-loader.cpp @@ -968,22 +968,25 @@ void cosyvoice_model_3::load(gguf_loader& loader) shared->params.llm_use_flash_attn ); - const auto dit_blocks = flow.decoder.estimator.transformer_blocks; - worker.dit_kv_cache.build_kv_cache( - backend, - worker.dit_kv_buffer, - static_cast(dit_blocks.size()), - static_cast(dit_blocks[0].attn.to_k.weight->ne[1] / dit_blocks[0].attn.heads), - static_cast(dit_blocks[0].attn.to_v.weight->ne[1] / dit_blocks[0].attn.heads), - dit_blocks[0].attn.heads, - shared->params.dit_kv_cache_length, - dit_k_type, - dit_v_type, - 2, - shared->params.dit_kv_fixed_slots, - shared->params.dit_kv_offloadable_slots, - shared->params.flow_use_flash_attn - ); + if (shared->params.dit_kv_fixed_slots != 0) + { + const auto dit_blocks = flow.decoder.estimator.transformer_blocks; + worker.dit_kv_cache.build_kv_cache( + backend, + worker.dit_kv_buffer, + static_cast(dit_blocks.size()), + static_cast(dit_blocks[0].attn.to_k.weight->ne[1] / dit_blocks[0].attn.heads), + static_cast(dit_blocks[0].attn.to_v.weight->ne[1] / dit_blocks[0].attn.heads), + dit_blocks[0].attn.heads, + shared->params.dit_kv_cache_length, + dit_k_type, + dit_v_type, + 2, + shared->params.dit_kv_fixed_slots, + shared->params.dit_kv_offloadable_slots, + shared->params.flow_use_flash_attn + ); + } } { diff --git a/src/cosyvoice-model.cpp b/src/cosyvoice-model.cpp index e28f8d7..db624c5 100644 --- a/src/cosyvoice-model.cpp +++ b/src/cosyvoice-model.cpp @@ -66,7 +66,7 @@ cosyvoice_model_shared::cosyvoice_model_shared(const cosyvoice_context_params_v3 cosyvoice_worker_context::cosyvoice_worker_context(ggml_backend_t backend) : backend(backend), cpu_backend(backend), ctx0(ggml_init(ggml_init_params{ .mem_size = ggml_graph_overhead() * kCosyVoiceGraphSize, .no_alloc = true })), - gf(nullptr), llm_input(nullptr), llm_probs(nullptr), position_ids(nullptr), causal_mask(nullptr), llm_kv_cache(), + gf(nullptr), llm_input(nullptr), llm_probs(nullptr), position_ids(nullptr), causal_mask(nullptr), llm_kv_cache(), dit_kv_cache(), status(GGML_STATUS_SUCCESS), prompt_crc32(0), sampler_seed(0), sampler(nullptr), sampler_ctx(nullptr), builtin_sampler_rng_policy(COSYVOICE_BUILTIN_SAMPLER_RNG_POLICY_RESET_PER_SESSION), nucleus_probs_capacity(0), nucleus_probs_len(0) {} cosyvoice_model::cosyvoice_model(ggml_backend_t backend, const cosyvoice_context_params_v3_cpp& params) From 056072ed2faf648c146af80f752a996f2c4b43a8 Mon Sep 17 00:00:00 2001 From: Lourdle Duan Date: Sun, 26 Jul 2026 13:44:57 +0800 Subject: [PATCH 34/34] cosyvoice: unify llm_job_ext prefill paths and fix buffer policy handling --- include/cosyvoice-lowlevel.h | 1 + src/cosyvoice-model.cpp | 1 + src/cosyvoice-tts.cpp | 138 ++++++++++++++--------------------- src/cosyvoice.cpp | 3 - 4 files changed, 58 insertions(+), 85 deletions(-) diff --git a/include/cosyvoice-lowlevel.h b/include/cosyvoice-lowlevel.h index d130a87..4ef0f2b 100644 --- a/include/cosyvoice-lowlevel.h +++ b/include/cosyvoice-lowlevel.h @@ -286,6 +286,7 @@ COSYVOICE_API bool cosyvoice_llm_job( * @brief Run the LLM with the given input tokens and prompt, with additional options. * @param max_new_tokens Maximum number of new tokens to generate. If 0, no new tokens are generated. * @param final Output parameter indicating whether the generation is complete (true) or more tokens can be generated (false). + * @note @p text must not be nullptr. Use @ref cosyvoice_llm_get_n_accepted_tokens to check whether this is a fresh run or a continuation. */ COSYVOICE_API bool cosyvoice_llm_job_ext( cosyvoice_context_t ctx, diff --git a/src/cosyvoice-model.cpp b/src/cosyvoice-model.cpp index db624c5..d95f313 100644 --- a/src/cosyvoice-model.cpp +++ b/src/cosyvoice-model.cpp @@ -305,6 +305,7 @@ void cosyvoice_model_3::reset_shared_buffer(ggml_backend_buffer* new_buffer) { worker->llm_kv_buffer.release(); worker->llm_kv_buffer.reset(new_buffer); + worker->llm_kv_cache.reset_buffer(worker->llm_kv_buffer.get()); } } diff --git a/src/cosyvoice-tts.cpp b/src/cosyvoice-tts.cpp index ba9f30d..2114030 100644 --- a/src/cosyvoice-tts.cpp +++ b/src/cosyvoice-tts.cpp @@ -32,8 +32,7 @@ bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoic auto last_prompt_crc32 = prompt_crc32; bool stop_reached = false; - if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED - || params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_SHARED && max_new_tokens != UINT32_MAX) + if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) { ggml_backend_sched_reset(sched.get()); worker->llm_kv_cache.load_cache(worker->backend.get(), sched.get()); @@ -49,97 +48,74 @@ bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoic const char* cur; uint32_t offset = 0; - if (text != nullptr) + if (llm_get_n_accepted_tokens() == 0 || llm_get_kv_cache_len() == 0) { - if (speech_type == llm.embed_tokens_weight->type) - { - auto prefill_embedding = [&](const char* data, int token_id) - { - if (offset == n_batch) - { - if (!llm_prefill(speech_type, batch_buffer.get(), n_batch)) - throw std::runtime_error("Failed to prefill LLM KV cache.\n"); - offset = 0; - } + GGML_ASSERT(text); - memcpy(batch_buffer.get() + offset++ * speech_row_size, data + token_id * speech_row_size, speech_row_size); - }; + const auto token_type = llm.embed_tokens_weight->type; + const auto token_row_size = static_cast(llm.embed_tokens_weight->nb[1]); - if (llm_get_kv_cache_len() == 0) - { - prefill_embedding(speech_emb, llm.sos_token_id); - prompt_crc32 = 0; - } - // The first token is assumed to be the SOS token already stored in the KV cache. - if (prompt_crc32 != prompt->prompt_crc32) + auto prefill_embedding = [&](const char* data, int token_id, uint32_t row_size, ggml_type type) + { + if (offset == n_batch) { - llm_set_kv_cache_len(1); - for (const auto& i : prompt->prompt_text) - prefill_embedding(token_emb, i); - prompt_crc32 = prompt->prompt_crc32; + if (!llm_prefill(type, batch_buffer.get(), n_batch)) + throw std::runtime_error("Failed to prefill LLM KV cache.\n"); + offset = 0; } - else llm_set_kv_cache_len(1 + static_cast(prompt->prompt_text.size())); - - for (uint32_t i = 0; i != text_len; ++i) - prefill_embedding(token_emb, text[i]); - if (prompt->llm_prompt_speech_tokens.second != 0) - { - prefill_embedding(speech_emb, llm.task_token_id); + memcpy(batch_buffer.get() + offset++ * row_size, data + token_id * row_size, row_size); + }; - const auto end = prompt->llm_prompt_speech_tokens.second - 1; - for (uint32_t i = 0; i != end; ++i) - prefill_embedding(speech_emb, prompt->llm_prompt_speech_tokens.first[i]); - cur = speech_emb + prompt->llm_prompt_speech_tokens.first[end] * speech_row_size; - } - else cur = speech_emb + llm.task_token_id * speech_row_size; - } - else + if (llm_get_kv_cache_len() == 0) { - const auto token_type = llm.embed_tokens_weight->type; - const auto token_row_size = static_cast(llm.embed_tokens_weight->nb[1]); - - auto prefill_embedding = [&](const char* data, int token_id, uint32_t row_size, ggml_type type) - { - if (offset == n_batch) - { - if (!llm_prefill(type, batch_buffer.get(), n_batch)) - throw std::runtime_error("Failed to prefill LLM KV cache.\n"); - offset = 0; - } - - memcpy(batch_buffer.get() + offset++ * row_size, data + token_id * row_size, row_size); - }; - - if (llm_get_kv_cache_len() == 0) + if (speech_type == token_type) + prefill_embedding(speech_emb, llm.sos_token_id, speech_row_size, speech_type); + else llm_prefill(speech_type, speech_emb + llm.sos_token_id * speech_row_size, 1); + prefill_prompt: + for (const auto& i : prompt->prompt_text) + prefill_embedding(token_emb, i, token_row_size, token_type); + prompt_crc32 = prompt->prompt_crc32; + } + else if (prompt_crc32 != prompt->prompt_crc32 + || !llm_set_kv_cache_len(1 + static_cast(prompt->prompt_text.size()))) + { // The first token is assumed to be the SOS token already stored in the KV cache. - if (prompt_crc32 != prompt->prompt_crc32) - { - llm_set_kv_cache_len(1); - for (const auto& i : prompt->prompt_text) - prefill_embedding(token_emb, i, token_row_size, token_type); - prompt_crc32 = prompt->prompt_crc32; - } - else llm_set_kv_cache_len(1 + static_cast(prompt->prompt_text.size())); + llm_set_kv_cache_len(1); + goto prefill_prompt; + } + ; - for (uint32_t i = 0; i != text_len; ++i) - prefill_embedding(token_emb, text[i], token_row_size, token_type); + for (uint32_t i = 0; i != text_len; ++i) + prefill_embedding(token_emb, text[i], token_row_size, token_type); + if (token_type != speech_type) + { if (offset != 0 && !llm_prefill(token_type, batch_buffer.get(), offset)) throw std::runtime_error("Failed to prefill LLM KV cache.\n"); offset = 0; + } - if (prompt->llm_prompt_speech_tokens.second != 0) - { - prefill_embedding(speech_emb, llm.task_token_id, speech_row_size, speech_type); + if (prompt->llm_prompt_speech_tokens.second != 0) + { + prefill_embedding(speech_emb, llm.task_token_id, speech_row_size, speech_type); - const auto end = prompt->llm_prompt_speech_tokens.second - 1; - for (uint32_t i = 0; i != end; ++i) - prefill_embedding(speech_emb, prompt->llm_prompt_speech_tokens.first[i], speech_row_size, speech_type); - cur = speech_emb + prompt->llm_prompt_speech_tokens.first[end] * speech_row_size; - } - else cur = speech_emb + llm.task_token_id * speech_row_size; + const auto end = prompt->llm_prompt_speech_tokens.second - 1; + for (uint32_t i = 0; i != end; ++i) + prefill_embedding(speech_emb, prompt->llm_prompt_speech_tokens.first[i], speech_row_size, speech_type); + cur = speech_emb + prompt->llm_prompt_speech_tokens.first[end] * speech_row_size; + } + else cur = speech_emb + llm.task_token_id * speech_row_size; + + if (const auto n_acc = llm_get_n_accepted_tokens(); n_acc > 0) + { + const auto* acc_tokens = llm_get_accepted_tokens(); + const auto end = n_acc - 1; + prefill_embedding(cur, 0, speech_row_size, speech_type); + for (uint32_t i = 0; i < end; ++i) + prefill_embedding(speech_emb, acc_tokens[i], speech_row_size, speech_type); + cur = speech_emb + acc_tokens[end] * speech_row_size; } if (offset != 0 && !llm_prefill(speech_type, batch_buffer.get(), offset)) @@ -148,10 +124,7 @@ bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoic else { // Continue from existing KV cache — get last accepted token's embedding - const auto n_acc = llm_get_n_accepted_tokens(); - if (n_acc == 0) - throw std::runtime_error("llm_job_ext: cannot continue generation with empty token history.\n"); - const auto last_token_id = llm_get_accepted_tokens()[n_acc - 1]; + const auto last_token_id = llm_get_accepted_tokens()[llm_get_n_accepted_tokens() - 1]; cur = speech_emb + last_token_id * speech_row_size; } @@ -216,6 +189,7 @@ bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoic catch (const std::exception& e) { worker->llm_input = nullptr; + llm_set_kv_cache_len(0); cosyvoice_call_ggml_log_callback(GGML_LOG_LEVEL_ERROR, e.what()); if (params.builtin_sampler_rng_policy == COSYVOICE_BUILTIN_SAMPLER_RNG_POLICY_RESET_PER_SESSION) reset_builtin_sampler_rng(); @@ -231,7 +205,7 @@ bool cosyvoice_model_3::llm_job_ext(const int* text, uint32_t text_len, cosyvoic llm_set_kv_cache_len(1 + static_cast(prompt->prompt_text.size())); llm_offload_kv_cache(); } - else if (params.inference_buffer_policy != COSYVOICE_INFERENCE_BUFFER_POLICY_DEDICATED && max_new_tokens != UINT32_MAX && !stop_reached) + else if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED && max_new_tokens != UINT32_MAX && !stop_reached) llm_offload_kv_cache(); // Reset RNG on stop (matches non-streaming behavior) @@ -529,7 +503,7 @@ bool cosyvoice_model_3::token2wav_ext(const int* token_ids, uint32_t n_tokens, f worker->status = ggml_backend_sched_graph_compute(sched.get(), gf); shared->noise_callback(COSYVOICE_NOISE_CALLBACK_STAGE_AFTER_FLOW, noise_len, noise_buffer, shared->noise_callback_ctx); - if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_SHARED) + if (params.inference_buffer_policy != COSYVOICE_INFERENCE_BUFFER_POLICY_DEDICATED) llm_set_kv_cache_len(0); if (worker->status != GGML_STATUS_SUCCESS) return false; diff --git a/src/cosyvoice.cpp b/src/cosyvoice.cpp index 074285d..90a4cc4 100644 --- a/src/cosyvoice.cpp +++ b/src/cosyvoice.cpp @@ -434,9 +434,6 @@ bool cosyvoice_tts_stream(cosyvoice_context_t ctx, const int* text, uint32_t tex return false; n_tokens = std::min(n_tokens + chunk_tokens, ctx->llm_get_n_accepted_tokens()); - text = nullptr; - if (params.inference_buffer_policy == COSYVOICE_INFERENCE_BUFFER_POLICY_BALANCED) - ctx->llm_offload_kv_cache(); cosyvoice_generated_speech result = {}; if (!ctx->token2wav_ext(ctx->llm_get_accepted_tokens(), n_tokens, speed, prompt, &offset, true, ctx->llm_get_n_accepted_tokens() == n_tokens, &result))