From 5eb0271a41d1b9a4d9f5ee399c743b8b363743b7 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:42:57 -0400 Subject: [PATCH] Add PocketTTS streaming support --- README.md | 2 +- .../engine/models/pocket_tts/acoustic_model.h | 23 + .../engine/models/pocket_tts/audio_decoder.h | 12 + .../engine/models/pocket_tts/mimi_decoder.h | 14 + include/engine/models/pocket_tts/session.h | 25 +- model_specs/pocket_tts.json | 3 +- src/models/pocket_tts/acoustic_model.cpp | 176 +++++--- src/models/pocket_tts/audio_decoder.cpp | 33 ++ src/models/pocket_tts/loader.cpp | 4 +- src/models/pocket_tts/mimi_decoder.cpp | 402 ++++++++++++++++++ src/models/pocket_tts/session.cpp | 172 +++++++- 11 files changed, 782 insertions(+), 84 deletions(-) diff --git a/README.md b/README.md index 8c1e25673..7907c4647 100644 --- a/README.md +++ b/README.md @@ -88,7 +88,7 @@ Runtime tags summarize the supported loading paths. GGUF package precision varie | **neutts** | TTS, Ctrl | en | NeuTTS 2E with built-in speaker prompts and emotion control | GGUF original precision, Stream | | **omnivoice** | TTS, Clone, Design, Ctrl | 646+ langs | OmniVoice, Qwen3-0.6B based | GGUF 16/Q8, Stream | | **personaplex** | Dialogue, S2S | en | PersonaPlex 7B v1 speech-to-speech conversational model with packaged voice/persona prompts | GGUF Q4/Q8, Stream | -| **pocket_tts** | TTS, Clone | en, de, it, pt, es | PocketTTS-100M | GGUF 16/Q8 | +| **pocket_tts** | TTS, Clone | en, de, it, pt, es | PocketTTS-100M | GGUF 16/Q8, Stream | | **qwen3_tts** | TTS, Clone, Design, Ctrl | zh, en, fr, de, it, ja, ko, pt, ru, es | Qwen3-TTS-12Hz-0.6B-Base, Qwen3-TTS-12Hz-1.7B-Base, Qwen3-TTS-12Hz-1.7B-CustomVoice, Qwen3-TTS-12Hz-1.7B-VoiceDesign | GGUF 16/Q8 | | **supertonic** | TTS | en, ko, ja, ar, bg, cs, da, de, el, es, et, fi, fr, hi, hr, hu, id, it, lt, lv, nl, pl, pt, ro, ru, sk, sl, sv, tr, uk, vi, na | Supertonic 3 | GGUF F32, Stream | | **vibevoice** | TTS, Dialogue | en, zh | VibeVoice-1.5B, VibeVoice-7B | GGUF 16/Q8 | diff --git a/include/engine/models/pocket_tts/acoustic_model.h b/include/engine/models/pocket_tts/acoustic_model.h index 3503b13e7..3f8ad0d3e 100644 --- a/include/engine/models/pocket_tts/acoustic_model.h +++ b/include/engine/models/pocket_tts/acoustic_model.h @@ -7,6 +7,8 @@ #include #include #include +#include +#include #include namespace engine::models::pocket_tts { @@ -35,6 +37,17 @@ struct AcousticPreparedRuntime { std::shared_ptr step_runtime; }; +struct AcousticStreamState { + AcousticPreparedRuntime runtime; + AcousticGenerationConfig config; + std::vector current_input; + std::mt19937 rng; + int step = 0; + int eos_step = -1; + int generated_steps = 0; + bool done = false; +}; + class AcousticModel { public: explicit AcousticModel(FlowLMConfig config = {}); @@ -63,6 +76,16 @@ class AcousticModel { const FlowLMState & initial_state, const AcousticGenerationConfig & config) const; + AcousticStreamState start_stream( + const AcousticPreparedRuntime & runtime, + const PocketTTSAssets & manifest, + const PocketTTSBackendWeights & weights, + const std::vector & text_embeddings, + const FlowLMState & initial_state, + const AcousticGenerationConfig & config) const; + + std::optional next_stream_step(AcousticStreamState & state) const; + void clear_runtime_cache() const noexcept; int64_t prepared_prompt_capacity() const noexcept; int prepared_max_steps_capacity() const noexcept; diff --git a/include/engine/models/pocket_tts/audio_decoder.h b/include/engine/models/pocket_tts/audio_decoder.h index dc8620181..c5bd3408c 100644 --- a/include/engine/models/pocket_tts/audio_decoder.h +++ b/include/engine/models/pocket_tts/audio_decoder.h @@ -31,6 +31,18 @@ class AudioDecoder { int64_t stage2_chunk_frames, bool use_full_sequence_path) const; + void reset_streaming_state() const; + + std::vector decode_streaming_step( + ggml_backend_t backend, + int threads, + const PocketTTSAssets & manifest, + const PocketTTSBackendWeights & weights, + const std::vector & normalized_latent, + size_t conv_graph_context_bytes, + size_t transformer_graph_context_bytes, + size_t tail_graph_context_bytes) const; + void clear_runtime_cache() const noexcept; private: diff --git a/include/engine/models/pocket_tts/mimi_decoder.h b/include/engine/models/pocket_tts/mimi_decoder.h index 22a540666..eb8c7eaeb 100644 --- a/include/engine/models/pocket_tts/mimi_decoder.h +++ b/include/engine/models/pocket_tts/mimi_decoder.h @@ -42,12 +42,26 @@ class MimiDecoder { int64_t stage2_chunk_frames, bool use_full_sequence_path) const; + void reset_streaming_state() const; + + std::vector decode_streaming_step( + ggml_backend_t backend, + int threads, + const PocketTTSAssets & manifest, + const PocketTTSBackendWeights & weights, + const std::vector & latent, + size_t conv_graph_context_bytes, + size_t transformer_graph_context_bytes, + size_t tail_graph_context_bytes) const; + void clear_runtime_cache() const noexcept; private: struct RuntimeCache; + struct StreamingState; MimiDecoderConfig config_; mutable std::unique_ptr runtime_cache_; + mutable std::unique_ptr streaming_state_; }; } // namespace engine::models::pocket_tts diff --git a/include/engine/models/pocket_tts/session.h b/include/engine/models/pocket_tts/session.h index 165da8657..87841942c 100644 --- a/include/engine/models/pocket_tts/session.h +++ b/include/engine/models/pocket_tts/session.h @@ -9,6 +9,7 @@ #include "engine/models/pocket_tts/text_conditioner.h" #include "engine/models/pocket_tts/voice_conditioner.h" +#include #include #include #include @@ -38,7 +39,8 @@ struct PocketTTSGraphCapacityConfig { class PocketTTSSession final : public runtime::RuntimeSessionBase - , public runtime::IOfflineVoiceTaskSession { + , public runtime::IOfflineVoiceTaskSession + , public runtime::IStreamingVoiceTaskSession { public: PocketTTSSession( runtime::TaskSpec task, @@ -55,6 +57,14 @@ class PocketTTSSession final runtime::RunMode run_mode() const override; void prepare(const runtime::SessionPreparationRequest & request) override; runtime::TaskResult run(const runtime::TaskRequest & request) override; + runtime::StreamingPolicy streaming_policy() const override; + void start_stream(const runtime::TaskRequest & request) override; + std::optional next_stream_event() override; + void set_stream_event_sink(runtime::StreamEventCallback sink) override; + runtime::TaskResult finish_stream() override; + void reset() override; + runtime::StreamEvent process_audio_chunk(const runtime::AudioChunk & chunk) override; + runtime::TaskResult finalize() override; void prepare_generation(const GenerationRequest & request); GenerationResult generate(const GenerationRequest & request); @@ -72,6 +82,8 @@ class PocketTTSSession final runtime::MappedGraphCapacityAdapter make_prompt_capacity_adapter() const; runtime::MappedGraphCapacityAdapter make_generation_capacity_adapter() const; AcousticCapacitySelection select_acoustic_capacities(int64_t prompt_steps, int max_steps) const; + GenerationRequest effective_request_for_run(const runtime::TaskRequest & request) const; + bool start_next_stream_text_chunk(); std::vector prepared_prompt_capacities() const; std::vector prepared_generation_capacities() const; @@ -88,6 +100,17 @@ class PocketTTSSession final GenerationRequest prepared_session_request_; runtime::GraphCapacityController prompt_capacity_controller_; runtime::GraphCapacityController generation_capacity_controller_; + + GenerationRequest stream_request_; + FlowLMState stream_voice_state_; + std::vector stream_text_chunks_; + size_t stream_text_chunk_index_ = 0; + std::optional stream_acoustic_state_; + runtime::AudioBuffer stream_merged_audio_; + runtime::StreamEventCallback stream_event_sink_; + std::chrono::steady_clock::time_point stream_started_at_; + size_t stream_audio_chunk_index_ = 0; + bool stream_started_ = false; }; } // namespace engine::models::pocket_tts diff --git a/model_specs/pocket_tts.json b/model_specs/pocket_tts.json index 522467eb5..cffc2f6c3 100644 --- a/model_specs/pocket_tts.json +++ b/model_specs/pocket_tts.json @@ -9,7 +9,8 @@ "clone" ], "modes": [ - "offline" + "offline", + "streaming" ], "languages": [ "en", diff --git a/src/models/pocket_tts/acoustic_model.cpp b/src/models/pocket_tts/acoustic_model.cpp index 062da351e..8fcf8ffb4 100644 --- a/src/models/pocket_tts/acoustic_model.cpp +++ b/src/models/pocket_tts/acoustic_model.cpp @@ -37,6 +37,51 @@ std::vector sample_normal(std::mt19937 & rng, int64_t count, float stddev return values; } +void validate_generation_inputs( + const FlowLMConfig & flow_config, + const std::vector & text_embeddings, + const AcousticGenerationConfig & config) { + if (config.max_steps <= 0) { + throw std::runtime_error("PocketTTS acoustic max_steps must be positive"); + } + if (config.temperature <= 0.0F) { + throw std::runtime_error("PocketTTS acoustic temperature must be positive"); + } + if (!config.noise_schedule.empty() && config.noise_schedule.size() % static_cast(flow_config.latent_size) != 0) { + throw std::runtime_error("PocketTTS acoustic noise_schedule must be a multiple of latent_size"); + } + if (!config.noise_schedule.empty()) { + const size_t scheduled_steps = + config.noise_schedule.size() / static_cast(flow_config.latent_size); + if (scheduled_steps < static_cast(config.max_steps)) { + throw std::runtime_error("PocketTTS acoustic noise_schedule must provide at least max_steps latent noise vectors"); + } + } + if (text_embeddings.size() % static_cast(flow_config.hidden_size) != 0) { + throw std::runtime_error("PocketTTS acoustic text embeddings must be a multiple of hidden_size"); + } +} + +std::vector sample_noise_for_step(const FlowLMConfig & flow_config, AcousticStreamState & state) { + if (!state.config.noise_schedule.empty()) { + const size_t start = static_cast(state.step) * static_cast(flow_config.latent_size); + return std::vector( + state.config.noise_schedule.begin() + static_cast(start), + state.config.noise_schedule.begin() + static_cast(start + static_cast(flow_config.latent_size))); + } + if (state.config.noise_clamp > 0.0F) { + return sample_trunc_normal( + state.rng, + flow_config.latent_size, + std::sqrt(state.config.temperature), + state.config.noise_clamp); + } + return sample_normal( + state.rng, + flow_config.latent_size, + std::sqrt(state.config.temperature)); +} + } // namespace AcousticModel::AcousticModel(FlowLMConfig config) : flow_lm_(std::move(config)) {} @@ -133,80 +178,15 @@ AcousticModelResult AcousticModel::generate( const std::vector & text_embeddings, const FlowLMState & initial_state, const AcousticGenerationConfig & config) const { - (void) manifest; - (void) weights; - if (config.max_steps <= 0) { - throw std::runtime_error("PocketTTS acoustic max_steps must be positive"); - } - if (config.temperature <= 0.0F) { - throw std::runtime_error("PocketTTS acoustic temperature must be positive"); - } - if (!config.noise_schedule.empty() && config.noise_schedule.size() % static_cast(flow_lm_.config().latent_size) != 0) { - throw std::runtime_error("PocketTTS acoustic noise_schedule must be a multiple of latent_size"); - } - if (!config.noise_schedule.empty()) { - const size_t scheduled_steps = - config.noise_schedule.size() / static_cast(flow_lm_.config().latent_size); - if (scheduled_steps < static_cast(config.max_steps)) { - throw std::runtime_error("PocketTTS acoustic noise_schedule must provide at least max_steps latent noise vectors"); - } - } - if (text_embeddings.size() % static_cast(flow_lm_.config().hidden_size) != 0) { - throw std::runtime_error("PocketTTS acoustic text embeddings must be a multiple of hidden_size"); - } - - const int64_t prompt_steps = runtime.prompt_steps; - if (runtime.step_runtime == nullptr) { - throw std::runtime_error("PocketTTS acoustic runtime is not initialized"); - } AcousticModelResult result; const double generate_ms = engine::debug::measure_ms([&]() { - flow_lm_.apply_prompt(*runtime.step_runtime, text_embeddings, prompt_steps, initial_state); - - std::vector current_input( - static_cast(flow_lm_.config().latent_size), - std::numeric_limits::quiet_NaN()); + auto state = start_stream(runtime, manifest, weights, text_embeddings, initial_state, config); result.latents.reserve(static_cast(config.max_steps) * static_cast(flow_lm_.config().latent_size)); result.eos_logits.reserve(static_cast(config.max_steps)); - std::mt19937 rng(config.seed); - int eos_step = -1; - for (int step = 0; step < config.max_steps; ++step) { - std::vector noise; - if (!config.noise_schedule.empty()) { - const size_t start = static_cast(step) * static_cast(flow_lm_.config().latent_size); - noise.assign( - config.noise_schedule.begin() + static_cast(start), - config.noise_schedule.begin() + static_cast(start + static_cast(flow_lm_.config().latent_size))); - } else if (config.noise_clamp > 0.0F) { - noise = sample_trunc_normal( - rng, - flow_lm_.config().latent_size, - std::sqrt(config.temperature), - config.noise_clamp); - } else { - noise = sample_normal( - rng, - flow_lm_.config().latent_size, - std::sqrt(config.temperature)); - } - - const auto step_result = flow_lm_.run_step_in_place( - *runtime.step_runtime, - current_input, - noise); - - const bool is_eos = step_result.eos_logit > config.eos_threshold; - if (is_eos && eos_step < 0) { - eos_step = step; - } - if (eos_step >= 0 && step >= eos_step + config.frames_after_eos) { - break; - } - - result.eos_logits.push_back(step_result.eos_logit); - result.latents.insert(result.latents.end(), step_result.next_latent.begin(), step_result.next_latent.end()); - current_input = step_result.next_latent; + while (auto step_result = next_stream_step(state)) { + result.eos_logits.push_back(step_result->eos_logit); + result.latents.insert(result.latents.end(), step_result->next_latent.begin(), step_result->next_latent.end()); result.generated_steps += 1; } const auto flow_timing = flow_lm_.runtime_timing(*runtime.step_runtime); @@ -233,6 +213,64 @@ AcousticModelResult AcousticModel::generate( return result; } +AcousticStreamState AcousticModel::start_stream( + const AcousticPreparedRuntime & runtime, + const models::pocket_tts::PocketTTSAssets & manifest, + const models::pocket_tts::PocketTTSBackendWeights & weights, + const std::vector & text_embeddings, + const FlowLMState & initial_state, + const AcousticGenerationConfig & config) const { + (void) manifest; + (void) weights; + validate_generation_inputs(flow_lm_.config(), text_embeddings, config); + if (runtime.step_runtime == nullptr) { + throw std::runtime_error("PocketTTS acoustic runtime is not initialized"); + } + flow_lm_.apply_prompt(*runtime.step_runtime, text_embeddings, runtime.prompt_steps, initial_state); + + AcousticStreamState state; + state.runtime = runtime; + state.config = config; + state.current_input.assign( + static_cast(flow_lm_.config().latent_size), + std::numeric_limits::quiet_NaN()); + state.rng.seed(config.seed); + return state; +} + +std::optional AcousticModel::next_stream_step(AcousticStreamState & state) const { + if (state.done) { + return std::nullopt; + } + if (state.runtime.step_runtime == nullptr) { + throw std::runtime_error("PocketTTS acoustic runtime is not initialized"); + } + if (state.step >= state.config.max_steps) { + state.done = true; + return std::nullopt; + } + + auto noise = sample_noise_for_step(flow_lm_.config(), state); + auto step_result = flow_lm_.run_step_in_place( + *state.runtime.step_runtime, + state.current_input, + noise); + + const bool is_eos = step_result.eos_logit > state.config.eos_threshold; + if (is_eos && state.eos_step < 0) { + state.eos_step = state.step; + } + if (state.eos_step >= 0 && state.step >= state.eos_step + state.config.frames_after_eos) { + state.done = true; + return std::nullopt; + } + + state.current_input = step_result.next_latent; + ++state.step; + ++state.generated_steps; + return step_result; +} + void AcousticModel::clear_runtime_cache() const noexcept { runtime_cache_ = {}; } diff --git a/src/models/pocket_tts/audio_decoder.cpp b/src/models/pocket_tts/audio_decoder.cpp index 691a45cd0..a55d892bb 100644 --- a/src/models/pocket_tts/audio_decoder.cpp +++ b/src/models/pocket_tts/audio_decoder.cpp @@ -66,6 +66,39 @@ std::vector AudioDecoder::decode( use_full_sequence_path); } +void AudioDecoder::reset_streaming_state() const { + decoder_.reset_streaming_state(); +} + +std::vector AudioDecoder::decode_streaming_step( + ggml_backend_t backend, + int threads, + const models::pocket_tts::PocketTTSAssets & manifest, + const models::pocket_tts::PocketTTSBackendWeights & weights, + const std::vector & normalized_latent, + size_t conv_graph_context_bytes, + size_t transformer_graph_context_bytes, + size_t tail_graph_context_bytes) const { + if (normalized_latent.size() != static_cast(decoder_.config().latent_size)) { + throw std::runtime_error("PocketTTS streaming audio decoder expects one latent step"); + } + const auto & emb_mean = weights.host.emb_mean; + const auto & emb_std = weights.host.emb_std; + if (emb_mean.size() != emb_std.size() || emb_mean.size() != static_cast(decoder_.config().latent_size)) { + throw std::runtime_error("PocketTTS latent normalization stats must match Mimi latent_size"); + } + auto denormalized = denormalize_latents(normalized_latent, emb_mean, emb_std); + return decoder_.decode_streaming_step( + backend, + threads, + manifest, + weights, + denormalized, + conv_graph_context_bytes, + transformer_graph_context_bytes, + tail_graph_context_bytes); +} + void AudioDecoder::clear_runtime_cache() const noexcept { decoder_.clear_runtime_cache(); } diff --git a/src/models/pocket_tts/loader.cpp b/src/models/pocket_tts/loader.cpp index 0098a1d35..3b769c003 100644 --- a/src/models/pocket_tts/loader.cpp +++ b/src/models/pocket_tts/loader.cpp @@ -22,7 +22,7 @@ std::string requested_language(const runtime::ModelLoadRequest & request) { runtime::CapabilitySet capabilities(const PocketTTSAssets & assets) { runtime::CapabilitySet out; out.supported_tasks = { - {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline, runtime::RunMode::Streaming}}, }; out.languages = {assets.language}; out.supports_speaker_reference = true; @@ -63,7 +63,7 @@ class PocketTTSLoader final : public runtime::IVoiceModelLoader { runtime::CapabilitySet advertised_capabilities() const override { runtime::CapabilitySet out; out.supported_tasks = { - {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline, runtime::RunMode::Streaming}}, }; out.supports_speaker_reference = true; out.supports_style_condition = true; diff --git a/src/models/pocket_tts/mimi_decoder.cpp b/src/models/pocket_tts/mimi_decoder.cpp index 0ad1b52cf..6e1cec483 100644 --- a/src/models/pocket_tts/mimi_decoder.cpp +++ b/src/models/pocket_tts/mimi_decoder.cpp @@ -882,6 +882,21 @@ class MimiTransformerRuntime { } core::set_backend_threads(backend_, threads_); graph_ = ggml_new_graph_custom(ggml_ctx_, 32768, false); + for (const auto & mask : attention_masks_) { + ggml_build_forward_expand(graph_, mask.tensor); + } + for (const auto & tensor : work_prefix_keys_) { + ggml_build_forward_expand(graph_, tensor.tensor); + } + for (const auto & tensor : work_prefix_values_) { + ggml_build_forward_expand(graph_, tensor.tensor); + } + for (const auto & tensor : zero_prefix_keys_) { + ggml_build_forward_expand(graph_, tensor.tensor); + } + for (const auto & tensor : zero_prefix_values_) { + ggml_build_forward_expand(graph_, tensor.tensor); + } ggml_build_forward_expand(graph_, output_bct_.tensor); galloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); if (galloc_ == nullptr || @@ -890,6 +905,20 @@ class MimiTransformerRuntime { release_partial_graph_runtime(galloc_, params_buffer_, ggml_ctx_); throw std::runtime_error("Mimi transformer graph allocation failed"); } + for (size_t keep = 0; keep < carry_key_sources_.size(); ++keep) { + for (size_t layer = 0; layer < carry_key_sources_[keep].size(); ++layer) { + if (carry_key_sources_[keep][layer] != nullptr) { + ggml_backend_view_init(carry_key_sources_[keep][layer]); + ggml_backend_view_init(carry_value_sources_[keep][layer]); + ggml_backend_view_init(carry_key_destinations_[keep][layer]); + ggml_backend_view_init(carry_value_destinations_[keep][layer]); + } + ggml_backend_view_init(append_key_sources_[keep][layer]); + ggml_backend_view_init(append_value_sources_[keep][layer]); + ggml_backend_view_init(append_key_destinations_[keep][layer]); + ggml_backend_view_init(append_value_destinations_[keep][layer]); + } + } core::write_tensor_f32(input_bct_, std::vector(static_cast(config_.hidden_size * frames_), 0.0F)); core::write_tensor_i32(positions_, std::vector(static_cast(frames_), 0)); core::write_tensor_f32(attention_mask_, std::vector(static_cast(frames_ * (cache_steps_ + frames_)), -INFINITY)); @@ -1340,6 +1369,14 @@ struct MimiDecoder::RuntimeCache { int64_t full_output_runtime_frames = -1; }; +struct MimiDecoder::StreamingState { + DecoderState decoder_state; + bool transformer_sequence_initialized = false; + + explicit StreamingState(const MimiDecoderConfig & config) + : decoder_state(make_decoder_state(config)) {} +}; + MimiDecoder::MimiDecoder(MimiDecoderConfig config) : config_(std::move(config)) {} MimiDecoder::~MimiDecoder() { @@ -2059,8 +2096,373 @@ std::vector MimiDecoder::decode( return audio; } +void MimiDecoder::reset_streaming_state() const { + streaming_state_ = std::make_unique(config_); +} + +std::vector MimiDecoder::decode_streaming_step( + ggml_backend_t backend, + int threads, + const models::pocket_tts::PocketTTSAssets & manifest, + const models::pocket_tts::PocketTTSBackendWeights & weights, + const std::vector & latent, + size_t conv_graph_context_bytes, + size_t transformer_graph_context_bytes, + size_t tail_graph_context_bytes) const { + const auto decode_started = std::chrono::steady_clock::now(); + if (latent.size() != static_cast(config_.latent_size)) { + throw std::runtime_error("PocketTTS Mimi streaming decoder expects one latent step"); + } + auto & runtime_cache = runtime_cache_; + if (!runtime_cache || runtime_cache->manifest != &manifest || runtime_cache->backend != backend || runtime_cache->threads != threads + || runtime_cache->conv_graph_context_bytes != conv_graph_context_bytes + || runtime_cache->transformer_graph_context_bytes != transformer_graph_context_bytes + || runtime_cache->tail_graph_context_bytes != tail_graph_context_bytes) { + runtime_cache = std::make_unique(); + runtime_cache->manifest = &manifest; + runtime_cache->backend = backend; + runtime_cache->threads = threads; + runtime_cache->conv_graph_context_bytes = conv_graph_context_bytes; + runtime_cache->transformer_graph_context_bytes = transformer_graph_context_bytes; + runtime_cache->tail_graph_context_bytes = tail_graph_context_bytes; + } + if (!streaming_state_) { + reset_streaming_state(); + } + auto & cache = *runtime_cache; + auto & state = streaming_state_->decoder_state; + const auto & decoder_weights = weights.mimi_decoder; + const auto & quantizer_weight = decoder_weights.quantizer_output_proj_weight; + const auto & encoder_upsample_weight = decoder_weights.encoder_upsample_weight; + const auto & input_projection = decoder_weights.input_projection; + const auto & stage0_upsample = decoder_weights.stage0_upsample; + const auto & stage1_upsample = decoder_weights.stage1_upsample; + const auto & stage2_upsample = decoder_weights.stage2_upsample; + const auto & output_projection = decoder_weights.output_projection; + auto & quantizer_runtime = cache.quantizer_runtime; + auto & transformer_runtime = cache.transformer_runtime; + auto & input_projection_runtime = cache.input_projection_runtime; + auto & stage0_upsample_runtime = cache.stage0_upsample_runtime; + auto & stage0_conv1_runtime = cache.stage0_conv1_runtime; + auto & stage0_conv2_runtime = cache.stage0_conv2_runtime; + auto & stage1_upsample_runtime = cache.stage1_upsample_runtime; + auto & stage1_conv1_runtime = cache.stage1_conv1_runtime; + auto & stage1_conv2_runtime = cache.stage1_conv2_runtime; + auto & stage2_upsample_runtime = cache.stage2_upsample_runtime; + auto & stage2_conv1_runtime = cache.stage2_conv1_runtime; + auto & stage2_conv2_runtime = cache.stage2_conv2_runtime; + auto & output_projection_runtime = cache.output_projection_runtime; + auto & resblock_conv1_frames = cache.resblock_conv1_frames; + auto & resblock_conv2_frames = cache.resblock_conv2_frames; + + auto run_resblock = [&](DecoderState & state_ref, + const std::vector & input_bct, + int64_t channels, + int64_t hidden_channels, + int stage_index, + std::unique_ptr & conv1_runtime, + std::unique_ptr & conv2_runtime, + const PocketTTSBackendResidualBlockWeights & block_weights) { + const int64_t frames_bct = static_cast(input_bct.size()) / channels; + const int64_t conv1_needed_frames = + frames_bct + state_ref.stage_residual_convs[static_cast(stage_index)][0].history_frames; + if (!conv1_runtime || resblock_conv1_frames[static_cast(stage_index)] != conv1_needed_frames) { + conv1_runtime = std::make_unique( + backend, + weights.backend_type, + threads, + conv_graph_context_bytes, + block_weights.conv1.weight, + block_weights.conv1.bias, + channels, + conv1_needed_frames, + hidden_channels, + 3, + 1, + 1); + resblock_conv1_frames[static_cast(stage_index)] = conv1_needed_frames; + } + const int64_t conv2_needed_frames = + frames_bct + state_ref.stage_residual_convs[static_cast(stage_index)][1].history_frames; + if (!conv2_runtime || resblock_conv2_frames[static_cast(stage_index)] != conv2_needed_frames) { + conv2_runtime = std::make_unique( + backend, + weights.backend_type, + threads, + conv_graph_context_bytes, + block_weights.conv2.weight, + block_weights.conv2.bias, + hidden_channels, + conv2_needed_frames, + channels, + 1, + 1, + 1); + resblock_conv2_frames[static_cast(stage_index)] = conv2_needed_frames; + } + auto x = elu(input_bct); + x = run_streaming_conv1d_step( + *conv1_runtime, + x, + channels, + frames_bct, + hidden_channels, + 3, + 1, + 1, + modules::StreamingPadMode::Constant, + state_ref.stage_residual_convs[static_cast(stage_index)][0]); + x = elu(x); + x = run_streaming_conv1d_step( + *conv2_runtime, + x, + hidden_channels, + frames_bct, + channels, + 1, + 1, + 1, + modules::StreamingPadMode::Constant, + state_ref.stage_residual_convs[static_cast(stage_index)][1]); + return add_bct(input_bct, x); + }; + + if (!quantizer_runtime || cache.quantizer_steps != 1) { + quantizer_runtime = std::make_unique( + backend, + weights.backend_type, + threads, + conv_graph_context_bytes, + quantizer_weight, + std::nullopt, + config_.latent_size, + 1, + config_.hidden_size, + 1, + 1, + 1); + cache.quantizer_steps = 1; + } + if (!cache.encoder_rate_upsample_runtime || cache.encoder_rate_upsample_steps != 1) { + cache.encoder_rate_upsample_runtime = std::make_unique( + backend, + weights.backend_type, + threads, + conv_graph_context_bytes, + encoder_upsample_weight, + 1, + config_.hidden_size, + config_.encoder_upsample_stride * 2, + static_cast(config_.encoder_upsample_stride)); + cache.encoder_rate_upsample_steps = 1; + } + auto & encoder_rate_upsample_runtime = *cache.encoder_rate_upsample_runtime; + + auto x = quantizer_runtime->run(latent); + x = run_depthwise_convtranspose1d_step( + encoder_rate_upsample_runtime, + x, + config_.hidden_size, + 1, + config_.encoder_upsample_stride * 2, + static_cast(config_.encoder_upsample_stride), + state.encoder_rate_upsample); + const int64_t encoder_frames = static_cast(x.size()) / config_.hidden_size; + if (!transformer_runtime || cache.transformer_frames != encoder_frames) { + transformer_runtime = std::make_unique( + backend, + threads, + transformer_graph_context_bytes, + weights, + config_, + encoder_frames, + 250); + cache.transformer_frames = encoder_frames; + } + if (!streaming_state_->transformer_sequence_initialized) { + transformer_runtime->reset_sequence(state.transformer.current_end); + streaming_state_->transformer_sequence_initialized = true; + } + x = transformer_runtime->run(x).output_bct; + const int64_t needed_input_frames = encoder_frames + state.input_projection.history_frames; + if (!input_projection_runtime || cache.input_projection_frames != needed_input_frames) { + input_projection_runtime = std::make_unique( + backend, + weights.backend_type, + threads, + conv_graph_context_bytes, + input_projection.weight, + input_projection.bias, + config_.hidden_size, + needed_input_frames, + config_.hidden_size, + 7, + 1, + 1); + cache.input_projection_frames = needed_input_frames; + } + x = run_streaming_conv1d_step( + *input_projection_runtime, + x, + config_.hidden_size, + encoder_frames, + config_.hidden_size, + 7, + 1, + 1, + modules::StreamingPadMode::Constant, + state.input_projection); + + x = elu(x); + if (!stage0_upsample_runtime || cache.stage0_upsample_frames != encoder_frames) { + stage0_upsample_runtime = std::make_unique( + backend, + weights.backend_type, + threads, + conv_graph_context_bytes, + stage0_upsample.weight, + stage0_upsample.bias, + config_.hidden_size, + encoder_frames, + 256, + 12, + 6); + cache.stage0_upsample_frames = encoder_frames; + } + x = run_streaming_convtranspose1d_step( + *stage0_upsample_runtime, + x, + config_.hidden_size, + encoder_frames, + 256, + 12, + 6, + decoder_weights.stage0_upsample_bias_values, + state.stage_upsamples[0]); + x = run_resblock( + state, + x, + 256, + 128, + 0, + stage0_conv1_runtime, + stage0_conv2_runtime, + decoder_weights.stage0_block); + + const int64_t stage1_frames = static_cast(x.size()) / 256; + x = elu(x); + if (!stage1_upsample_runtime || cache.stage1_upsample_frames != stage1_frames) { + stage1_upsample_runtime = std::make_unique( + backend, + weights.backend_type, + threads, + conv_graph_context_bytes, + stage1_upsample.weight, + stage1_upsample.bias, + 256, + stage1_frames, + 128, + 10, + 5); + cache.stage1_upsample_frames = stage1_frames; + } + x = run_streaming_convtranspose1d_step( + *stage1_upsample_runtime, + x, + 256, + stage1_frames, + 128, + 10, + 5, + decoder_weights.stage1_upsample_bias_values, + state.stage_upsamples[1]); + x = run_resblock( + state, + x, + 128, + 64, + 1, + stage1_conv1_runtime, + stage1_conv2_runtime, + decoder_weights.stage1_block); + + const int64_t stage2_frames = static_cast(x.size()) / 128; + x = elu(x); + if (!stage2_upsample_runtime || cache.stage2_upsample_frames != stage2_frames) { + stage2_upsample_runtime = std::make_unique( + backend, + weights.backend_type, + threads, + conv_graph_context_bytes, + stage2_upsample.weight, + stage2_upsample.bias, + 128, + stage2_frames, + 64, + 8, + 4); + cache.stage2_upsample_frames = stage2_frames; + } + x = run_streaming_convtranspose1d_step( + *stage2_upsample_runtime, + x, + 128, + stage2_frames, + 64, + 8, + 4, + decoder_weights.stage2_upsample_bias_values, + state.stage_upsamples[2]); + x = run_resblock( + state, + x, + 64, + 32, + 2, + stage2_conv1_runtime, + stage2_conv2_runtime, + decoder_weights.stage2_block); + + const int64_t output_frames = static_cast(x.size()) / 64; + x = elu(x); + const int64_t needed_output_frames = output_frames + state.output_projection.history_frames; + if (!output_projection_runtime || cache.output_projection_frames != needed_output_frames) { + output_projection_runtime = std::make_unique( + backend, + weights.backend_type, + threads, + conv_graph_context_bytes, + output_projection.weight, + output_projection.bias, + 64, + needed_output_frames, + 1, + 3, + 1, + 1); + cache.output_projection_frames = needed_output_frames; + } + x = run_streaming_conv1d_step( + *output_projection_runtime, + x, + 64, + output_frames, + 1, + 3, + 1, + 1, + modules::StreamingPadMode::Constant, + state.output_projection); + + engine::debug::timing_log_scalar( + "pocket_tts.mimi.streaming_decoder_step_ms", + engine::debug::elapsed_ms(decode_started)); + return x; +} + void MimiDecoder::clear_runtime_cache() const noexcept { runtime_cache_.reset(); + streaming_state_.reset(); } } // namespace engine::models::pocket_tts diff --git a/src/models/pocket_tts/session.cpp b/src/models/pocket_tts/session.cpp index 4a7bc2288..5f4fe48e0 100644 --- a/src/models/pocket_tts/session.cpp +++ b/src/models/pocket_tts/session.cpp @@ -474,8 +474,8 @@ PocketTTSSession::PocketTTSSession( if (task_.task != runtime::VoiceTaskKind::Tts) { throw std::runtime_error("PocketTTS only supports VoiceTaskKind::Tts"); } - if (task_.mode != runtime::RunMode::Offline) { - throw std::runtime_error("PocketTTS only supports offline mode"); + if (task_.mode != runtime::RunMode::Offline && task_.mode != runtime::RunMode::Streaming) { + throw std::runtime_error("PocketTTS only supports offline and streaming mode"); } if (graph_capacity_.prompt_mode == runtime::GraphCapacityMode::Unsupported || graph_capacity_.generation_mode == runtime::GraphCapacityMode::Unsupported) { @@ -761,6 +761,128 @@ void PocketTTSSession::prepare(const runtime::SessionPreparationRequest & reques runtime::TaskResult PocketTTSSession::run(const runtime::TaskRequest & request) { require_prepared("PocketTTS run()"); audio_decoder_.clear_runtime_cache(); + const GenerationRequest generation_request = effective_request_for_run(request); + const GenerationResult generated = generate(generation_request); + runtime::TaskResult result; + result.audio_output = runtime::AudioBuffer{ + generated.sample_rate, + 1, + generated.audio, + }; + return result; +} + +runtime::StreamingPolicy PocketTTSSession::streaming_policy() const { + runtime::StreamingPolicy policy; + policy.input = runtime::StreamingInputKind::None; + policy.output = runtime::StreamingOutputKind::PullEvents; + return policy; +} + +void PocketTTSSession::start_stream(const runtime::TaskRequest & request) { + require_prepared("PocketTTS streaming"); + if (task_.mode != runtime::RunMode::Streaming) { + throw std::runtime_error("PocketTTS start_stream requires a streaming session"); + } + reset(); + stream_request_ = effective_request_for_run(request); + validate_generation_request(stream_request_); + const int64_t streaming_chunk_size = + engine::text::parse_text_chunk_size_override(request.options).value_or(stream_request_.max_tokens); + stream_text_chunks_ = engine::text::split_text_chunks(stream_request_.text, streaming_chunk_size); + if (stream_text_chunks_.empty()) { + throw std::runtime_error("PocketTTS streaming text chunking produced no segments"); + } + const auto voice_plan = resolve_voice_conditioning_plan(model_dir_, stream_request_); + stream_voice_state_ = resolve_prepared_voice_state(voice_plan); + stream_merged_audio_ = runtime::AudioBuffer{manifest_->model_config.sample_rate, 1, {}}; + audio_decoder_.reset_streaming_state(); + stream_started_at_ = std::chrono::steady_clock::now(); + stream_started_ = true; + engine::debug::trace_log_scalar("pocket_tts.streaming.text_chunk_size", streaming_chunk_size); + engine::debug::trace_log_scalar("pocket_tts.streaming.text_chunk_count", static_cast(stream_text_chunks_.size())); +} + +std::optional PocketTTSSession::next_stream_event() { + if (!stream_started_) { + throw std::runtime_error("PocketTTS streaming has not been started"); + } + while (true) { + if (!stream_acoustic_state_.has_value()) { + if (!start_next_stream_text_chunk()) { + return std::nullopt; + } + } + auto acoustic_step = acoustic_model_.next_stream_step(*stream_acoustic_state_); + if (!acoustic_step.has_value()) { + stream_acoustic_state_.reset(); + continue; + } + auto audio = audio_decoder_.decode_streaming_step( + execution_context().backend(), + options().backend.threads, + *manifest_, + *weights_, + acoustic_step->next_latent, + graph_capacity_.mimi_conv_graph_context_bytes, + graph_capacity_.mimi_transformer_graph_context_bytes, + graph_capacity_.mimi_tail_graph_context_bytes); + runtime::AudioBuffer chunk{ + manifest_->model_config.sample_rate, + 1, + std::move(audio), + }; + runtime::append_audio_buffer(stream_merged_audio_, chunk); + + runtime::StreamEvent event; + event.named_audio_outputs.push_back({ + "chunk_" + std::to_string(stream_audio_chunk_index_++), + std::move(chunk), + {}, + }); + return event; + } +} + +void PocketTTSSession::set_stream_event_sink(runtime::StreamEventCallback sink) { + stream_event_sink_ = std::move(sink); +} + +runtime::TaskResult PocketTTSSession::finish_stream() { + if (!stream_started_) { + throw std::runtime_error("PocketTTS streaming has not been started"); + } + while (next_stream_event().has_value()) { + } + runtime::TaskResult result; + result.audio_output = std::move(stream_merged_audio_); + engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(stream_started_at_)); + reset(); + return result; +} + +void PocketTTSSession::reset() { + stream_request_ = {}; + stream_voice_state_ = {}; + stream_text_chunks_.clear(); + stream_text_chunk_index_ = 0; + stream_acoustic_state_.reset(); + audio_decoder_.reset_streaming_state(); + stream_merged_audio_ = {}; + stream_audio_chunk_index_ = 0; + stream_started_ = false; +} + +runtime::StreamEvent PocketTTSSession::process_audio_chunk(const runtime::AudioChunk & chunk) { + (void) chunk; + throw std::runtime_error("PocketTTS streaming does not accept audio chunks"); +} + +runtime::TaskResult PocketTTSSession::finalize() { + return finish_stream(); +} + +GenerationRequest PocketTTSSession::effective_request_for_run(const runtime::TaskRequest & request) const { runtime::TaskRequest effective_request = request; effective_request.voice.reset(); GenerationRequest generation_request = @@ -774,14 +896,44 @@ runtime::TaskResult PocketTTSSession::run(const runtime::TaskRequest & request) generation_request.noise_schedule = prepared_session_request_.noise_schedule; generation_request.noise_schedule_path = prepared_session_request_.noise_schedule_path; generation_request.voice = prepared_session_request_.voice; - const GenerationResult generated = generate(generation_request); - runtime::TaskResult result; - result.audio_output = runtime::AudioBuffer{ - generated.sample_rate, - 1, - generated.audio, - }; - return result; + return generation_request; +} + +bool PocketTTSSession::start_next_stream_text_chunk() { + if (stream_text_chunk_index_ >= stream_text_chunks_.size()) { + return false; + } + const auto & chunk = stream_text_chunks_[stream_text_chunk_index_++]; + const TextConditioningResult text_state = text_conditioner_.prepare(*manifest_, weights_->host, chunk); + const AcousticGenerationConfig acoustic_config = resolve_acoustic_generation_config( + *manifest_, + text_state, + stream_request_, + acoustic_model_.config().latent_size); + const int64_t prompt_steps = static_cast( + text_state.text_embeddings.size() / static_cast(acoustic_model_.config().hidden_size)); + const AcousticCapacitySelection capacities = select_acoustic_capacities(prompt_steps, acoustic_config.max_steps); + AcousticPreparedRuntime acoustic_runtime = acoustic_model_.prepare_runtime( + execution_context().backend(), + options().backend.threads, + *manifest_, + *weights_, + text_state.text_embeddings, + stream_voice_state_, + acoustic_config, + capacities.prompt_capacity, + stream_voice_state_.current_end, + capacities.generation_capacity, + graph_capacity_.flow_weights_view_context_bytes, + graph_capacity_.flow_step_graph_context_bytes); + stream_acoustic_state_ = acoustic_model_.start_stream( + acoustic_runtime, + *manifest_, + *weights_, + text_state.text_embeddings, + stream_voice_state_, + acoustic_config); + return true; } void PocketTTSSession::prepare_generation(const GenerationRequest & request) {