Add PocketTTS streaming support

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 |

include/engine/models/pocket_tts/acoustic_model.h

@@ -7,6 +7,8 @@
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <optional>
+#include <random>
 #include <vector>
 
 namespace engine::models::pocket_tts {
@@ -35,6 +37,17 @@ struct AcousticPreparedRuntime {
     std::shared_ptr<FlowLMStepRuntime> step_runtime;
 };
 
+struct AcousticStreamState {
+    AcousticPreparedRuntime runtime;
+    AcousticGenerationConfig config;
+    std::vector<float> 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<float> & text_embeddings,
+        const FlowLMState & initial_state,
+        const AcousticGenerationConfig & config) const;
+
+    std::optional<FlowLMStepResult> 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;

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<float> decode_streaming_step(
+        ggml_backend_t backend,
+        int threads,
+        const PocketTTSAssets & manifest,
+        const PocketTTSBackendWeights & weights,
+        const std::vector<float> & 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:

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<float> decode_streaming_step(
+        ggml_backend_t backend,
+        int threads,
+        const PocketTTSAssets & manifest,
+        const PocketTTSBackendWeights & weights,
+        const std::vector<float> & 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<RuntimeCache> runtime_cache_;
+    mutable std::unique_ptr<StreamingState> streaming_state_;
 };
 
 }  // namespace engine::models::pocket_tts

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 <chrono>
 #include <cstddef>
 #include <memory>
 #include <string>
@@ -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<runtime::StreamEvent> 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<int64_t> prepared_prompt_capacities() const;
     std::vector<int64_t> 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<std::string> stream_text_chunks_;
+    size_t stream_text_chunk_index_ = 0;
+    std::optional<AcousticStreamState> 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

model_specs/pocket_tts.json

@@ -9,7 +9,8 @@
     "clone"
   ],
   "modes": [
-    "offline"
+    "offline",
+    "streaming"
   ],
   "languages": [
     "en",

src/models/pocket_tts/acoustic_model.cpp

@@ -37,6 +37,51 @@ std::vector<float> sample_normal(std::mt19937 & rng, int64_t count, float stddev
     return values;
 }
 
+void validate_generation_inputs(
+    const FlowLMConfig & flow_config,
+    const std::vector<float> & 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<size_t>(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<size_t>(flow_config.latent_size);
+        if (scheduled_steps < static_cast<size_t>(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<size_t>(flow_config.hidden_size) != 0) {
+        throw std::runtime_error("PocketTTS acoustic text embeddings must be a multiple of hidden_size");
+    }
+}
+
+std::vector<float> sample_noise_for_step(const FlowLMConfig & flow_config, AcousticStreamState & state) {
+    if (!state.config.noise_schedule.empty()) {
+        const size_t start = static_cast<size_t>(state.step) * static_cast<size_t>(flow_config.latent_size);
+        return std::vector<float>(
+            state.config.noise_schedule.begin() + static_cast<ptrdiff_t>(start),
+            state.config.noise_schedule.begin() + static_cast<ptrdiff_t>(start + static_cast<size_t>(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<float> & 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<size_t>(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<size_t>(flow_lm_.config().latent_size);
-        if (scheduled_steps < static_cast<size_t>(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<size_t>(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<float> current_input(
-            static_cast<size_t>(flow_lm_.config().latent_size),
-            std::numeric_limits<float>::quiet_NaN());
+        auto state = start_stream(runtime, manifest, weights, text_embeddings, initial_state, config);
         result.latents.reserve(static_cast<size_t>(config.max_steps) * static_cast<size_t>(flow_lm_.config().latent_size));
         result.eos_logits.reserve(static_cast<size_t>(config.max_steps));
 
-        std::mt19937 rng(config.seed);
-        int eos_step = -1;
-        for (int step = 0; step < config.max_steps; ++step) {
-            std::vector<float> noise;
-            if (!config.noise_schedule.empty()) {
-                const size_t start = static_cast<size_t>(step) * static_cast<size_t>(flow_lm_.config().latent_size);
-                noise.assign(
-                    config.noise_schedule.begin() + static_cast<ptrdiff_t>(start),
-                    config.noise_schedule.begin() + static_cast<ptrdiff_t>(start + static_cast<size_t>(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<float> & 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<size_t>(flow_lm_.config().latent_size),
+        std::numeric_limits<float>::quiet_NaN());
+    state.rng.seed(config.seed);
+    return state;
+}
+
+std::optional<FlowLMStepResult> 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_ = {};
 }