feat(kokoro): add multilingual GGUF synthesis and optimized CPU inference

CMakeLists.txt

@@ -583,6 +583,28 @@ audiocpp_add_model(demucs
         htdemucs
 )
 
+audiocpp_add_model(kokoro_tts
+    SOURCES
+        src/models/kokoro_tts/assets.cpp
+        src/models/kokoro_tts/package.cpp
+        src/models/kokoro_tts/g2p_multilingual.cpp
+        src/models/kokoro_tts/loader.cpp
+        src/models/kokoro_tts/frontend.cpp
+        src/models/kokoro_tts/g2p_en.cpp
+        src/models/kokoro_tts/predictor.cpp
+        src/models/kokoro_tts/decoder.cpp
+        src/models/kokoro_tts/plbert.cpp
+        src/models/kokoro_tts/session.cpp
+    INCLUDES engine/models/kokoro_tts/loader.h
+    LOADERS engine::models::kokoro_tts::make_kokoro_tts_loader
+)
+if(MSVC)
+    target_compile_options(engine_model_kokoro_tts PRIVATE /utf-8)
+endif()
+if(AUDIOCPP_STATIC_ESPEAK)
+    target_compile_definitions(engine_model_kokoro_tts PRIVATE AUDIOCPP_STATIC_ESPEAK=1)
+endif()
+
 audiocpp_add_model(omnivoice
     SOURCES
         src/models/omnivoice/assets.cpp
@@ -2109,6 +2131,9 @@ if (ENGINE_BUILD_WARMBENCH)
 
     add_engine_warmbench(campplus_shared_default_probe tests/glm_tts/campplus_shared_default_probe.cpp)
     add_engine_warmbench(chatterbox_warm_bench tests/chatterbox/chatterbox_warm_bench.cpp)
+    if(kokoro_tts IN_LIST AUDIOCPP_LINKED_MODELS)
+        add_engine_warmbench(kokoro_tts_warm_bench tests/kokoro_tts/kokoro_tts_warm_bench.cpp)
+    endif()
     add_engine_warmbench(citrinet_asr_warm_bench tests/citrinet_asr/citrinet_asr_warm_bench.cpp)
     add_engine_warmbench(confucius4_tts_warm_bench tests/confucius4_tts/confucius4_tts_warm_bench.cpp)
     add_engine_warmbench(controlfoley_warm_bench tests/controlfoley/controlfoley_warm_bench.cpp)
@@ -2184,6 +2209,14 @@ if (ENGINE_BUILD_TESTS OR ENGINE_BUILD_EXTENDED_TESTS OR ENGINE_BUILD_MODEL_TEST
     add_test(NAME sortformer_v2_schedule_test COMMAND sortformer_v2_schedule_test)
 
     if (ENGINE_BUILD_TESTS)
+        if(kokoro_tts IN_LIST AUDIOCPP_LINKED_MODELS)
+            add_engine_unittest(kokoro_cpu_kernel_test tests/kokoro_tts/kokoro_cpu_kernel_test.cpp)
+            add_engine_unittest(kokoro_g2p_probe tests/kokoro_tts/kokoro_g2p_probe.cpp)
+            add_test(NAME kokoro_cpu_kernel_test COMMAND kokoro_cpu_kernel_test)
+            if(AUDIOCPP_STATIC_ESPEAK)
+                audiocpp_stage_espeak(kokoro_g2p_probe)
+            endif()
+        endif()
         add_engine_unittest(espeak_data_test tests/unittests/test_espeak_data.cpp)
         target_link_libraries(espeak_data_test PRIVATE Threads::Threads)
         add_test(NAME espeak_data_test COMMAND espeak_data_test)

README.md

@@ -84,6 +84,7 @@ Runtime tags summarize the supported loading paths. GGUF package precision varie
 | **fireredtts3** | TTS, Clone, Design, Ctrl | 24 langs + 21 zh dialects | FireRedTTS3 Base and Instruct packages for voice cloning, design, semantic edit, and acoustic edit | GGUF original/Q8 |
 | **higgs_audio_tts** | TTS, Clone, Ctrl | auto | Higgs Audio v3 TTS 4B | GGUF 16/Q8 |
 | **index_tts2** | TTS, Clone, Ctrl | zh, en, ja, es, ar | IndexTTS-2, IndexTTS-2.5 (variant) | GGUF 16/Q8 |
+| **kokoro_tts** | TTS | en-us, en-gb, es, fr, hi, it, ja, pt-br, zh | [Kokoro 82M](tests/kokoro_tts/MULTILINGUAL_GGUF.md), 54 preset voices | Safetensors, local GGUF BF16/Q8 |
 | **irodori_tts** | TTS, Clone, Design, Ctrl | ja | Irodori-TTS-v4-Small, Irodori-TTS-500M-v3, Irodori-TTS-600M-v3-VoiceDesign | GGUF 16/Q8 |
 | **magpie_tts** | TTS | ar-AE, ar-MSA, ar-SA, de, en, es, fr, hi, it, ko, pt-BR, vi, zh | NVIDIA MagpieTTS Multilingual 357M (v2607) with baked speaker prompts and NanoCodec decode | GGUF original/Q8 |
 | **miotts** | TTS, Clone | en, ja | MioTTS-1.7B | GGUF 16/Q8 |

include/engine/models/kokoro_tts/assets.h

@@ -0,0 +1,285 @@
+#pragma once
+
+#include "engine/framework/assets/tensor_source.h"
+#include "engine/framework/core/backend_weight_store.h"
+#include "engine/framework/core/module.h"
+#include "engine/framework/io/json.h"
+
+#include <ggml.h>
+
+#include <cstdint>
+#include <filesystem>
+#include <memory>
+#include <optional>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace kokoro_ggml {
+struct KokoroWeights {
+    struct LinearWeights {
+        engine::core::TensorValue weight;
+        std::optional<engine::core::TensorValue> bias;
+        int64_t out_features = 0;
+        int64_t in_features = 0;
+        bool use_bias = false;
+    };
+
+    struct HostAffineWeights {
+        std::vector<float> weight;
+        std::vector<float> bias;
+        int64_t out_features = 0;
+        int64_t in_features = 0;
+        bool use_bias = false;
+    };
+
+    struct EmbeddingWeights {
+        engine::core::TensorValue weight;
+        int64_t num_embeddings = 0;
+        int64_t embedding_dim = 0;
+    };
+
+    struct LayerNormWeights {
+        engine::core::TensorValue weight;
+        engine::core::TensorValue bias;
+        int64_t channels = 0;
+        float eps = 1.0e-5f;
+    };
+
+    struct WeightNormConv1dWeights {
+        engine::core::TensorValue weight;
+        std::optional<engine::core::TensorValue> bias;
+        int64_t out_channels = 0;
+        int64_t in_channels = 0;
+        int64_t kernel = 0;
+        int64_t stride = 1;
+        int64_t padding = 0;
+        int64_t dilation = 1;
+        int64_t groups = 1;
+        bool use_bias = false;
+    };
+
+    struct Conv1dWeights {
+        engine::core::TensorValue weight;
+        std::optional<engine::core::TensorValue> bias;
+        int64_t out_channels = 0;
+        int64_t in_channels = 0;
+        int64_t kernel = 0;
+        int64_t stride = 1;
+        int64_t padding = 0;
+        int64_t dilation = 1;
+        int64_t groups = 1;
+        bool use_bias = false;
+    };
+
+    struct WeightNormConvTranspose1dWeights {
+        engine::core::TensorValue weight;
+        engine::core::TensorValue dense_weight;
+        std::optional<engine::core::TensorValue> bias;
+        std::shared_ptr<Conv1dWeights> phase_shuffle_conv;
+        int64_t in_channels = 0;
+        int64_t out_channels = 0;
+        int64_t kernel = 0;
+        int64_t stride = 1;
+        int64_t padding = 0;
+        int64_t output_padding = 0;
+        int64_t groups = 1;
+        bool use_bias = false;
+    };
+
+    struct LstmWeights {
+        engine::core::TensorValue weight_ih_l0;
+        engine::core::TensorValue weight_hh_l0;
+        engine::core::TensorValue bias_ih_l0;
+        engine::core::TensorValue bias_hh_l0;
+        engine::core::TensorValue combined_bias_l0;
+        engine::core::TensorValue weight_ih_l0_reverse;
+        engine::core::TensorValue weight_hh_l0_reverse;
+        engine::core::TensorValue bias_ih_l0_reverse;
+        engine::core::TensorValue bias_hh_l0_reverse;
+        engine::core::TensorValue combined_bias_l0_reverse;
+        int64_t input_size = 0;
+        int64_t hidden_size = 0;
+    };
+
+    struct AdaLayerNormWeights {
+        LinearWeights fc;
+        int64_t channels = 0;
+        float eps = 1.0e-5f;
+    };
+
+    struct AdaIn1dWeights {
+        LinearWeights fc;
+        int64_t channels = 0;
+        float eps = 1.0e-5f;
+    };
+
+    struct AlbertEmbeddingsWeights {
+        EmbeddingWeights word_embeddings;
+        EmbeddingWeights position_embeddings;
+        EmbeddingWeights token_type_embeddings;
+        LayerNormWeights layer_norm;
+    };
+
+    struct AlbertAttentionWeights {
+        LinearWeights query;
+        LinearWeights key;
+        LinearWeights value;
+        LinearWeights dense;
+        LayerNormWeights layer_norm;
+    };
+
+    struct AlbertLayerWeights {
+        AlbertAttentionWeights attention;
+        LinearWeights ffn;
+        LinearWeights ffn_output;
+        LayerNormWeights full_layer_layer_norm;
+    };
+
+    struct AlbertWeights {
+        AlbertEmbeddingsWeights embeddings;
+        LinearWeights embedding_hidden_mapping_in;
+        AlbertLayerWeights shared_layer;
+        LinearWeights pooler;
+        int64_t hidden_size = 768;
+        int64_t embedding_size = 128;
+        int64_t intermediate_size = 2048;
+        int64_t max_position_embeddings = 512;
+        int64_t num_hidden_layers = 12;
+        int64_t num_attention_heads = 12;
+        float layer_norm_eps = 1.0e-12f;
+    };
+
+    struct TextEncoderBlockWeights {
+        WeightNormConv1dWeights conv;
+        LayerNormWeights layer_norm;
+    };
+
+    struct TextEncoderWeights {
+        EmbeddingWeights embedding;
+        std::vector<TextEncoderBlockWeights> cnn;
+        LstmWeights lstm;
+    };
+
+    struct DurationEncoderWeights {
+        std::vector<LstmWeights> lstms;
+        std::vector<AdaLayerNormWeights> ada_layer_norms;
+    };
+
+    struct AdainResBlock1dWeights {
+        WeightNormConv1dWeights conv1;
+        WeightNormConv1dWeights conv2;
+        WeightNormConv1dWeights conv1x1;
+        WeightNormConvTranspose1dWeights pool;
+        AdaIn1dWeights norm1;
+        AdaIn1dWeights norm2;
+        bool learned_sc = false;
+        bool use_pool = false;
+        bool upsample = false;
+    };
+
+    struct GeneratorResBlockWeights {
+        std::vector<WeightNormConv1dWeights> convs1;
+        std::vector<WeightNormConv1dWeights> convs2;
+        std::vector<AdaIn1dWeights> adain1;
+        std::vector<AdaIn1dWeights> adain2;
+        std::vector<engine::core::TensorValue> alpha1;
+        std::vector<engine::core::TensorValue> alpha2;
+    };
+
+    struct GeneratorWeights {
+        std::vector<WeightNormConvTranspose1dWeights> ups;
+        std::vector<Conv1dWeights> noise_convs;
+        std::vector<GeneratorResBlockWeights> noise_res;
+        std::vector<GeneratorResBlockWeights> resblocks;
+        WeightNormConv1dWeights conv_post;
+        HostAffineWeights source_linear;
+        int64_t harmonic_num = 8;
+        int64_t sampling_rate = 24000;
+        float sine_amp = 0.1f;
+        float noise_std = 0.003f;
+        float voiced_threshold = 10.0f;
+        int64_t gen_istft_n_fft = 20;
+        int64_t gen_istft_hop_size = 5;
+    };
+
+    struct ProsodyPredictorWeights {
+        DurationEncoderWeights duration_encoder;
+        LstmWeights lstm;
+        LinearWeights duration_proj;
+        LstmWeights shared;
+        std::vector<AdainResBlock1dWeights> f0_blocks;
+        std::vector<AdainResBlock1dWeights> n_blocks;
+        Conv1dWeights f0_proj;
+        Conv1dWeights n_proj;
+    };
+
+    struct DecoderWeights {
+        AdainResBlock1dWeights encode;
+        std::vector<AdainResBlock1dWeights> decode;
+        WeightNormConv1dWeights f0_conv;
+        WeightNormConv1dWeights n_conv;
+        WeightNormConv1dWeights asr_res;
+        GeneratorWeights generator;
+    };
+
+    AlbertWeights bert;
+    LinearWeights bert_encoder;
+    ProsodyPredictorWeights predictor;
+    TextEncoderWeights text_encoder;
+    DecoderWeights decoder;
+
+    int64_t n_token = 178;
+    int64_t hidden_dim = 512;
+    int64_t style_dim = 128;
+    int64_t n_layer = 3;
+    int64_t max_dur = 50;
+    int64_t n_mels = 80;
+    int64_t context_length = 512;
+    int64_t max_output_tokens = 512 * 50;
+    float dropout = 0.2f;
+    float lrelu_slope = 0.1f;
+    float post_lrelu_slope = 0.01f;
+
+    std::shared_ptr<engine::core::BackendWeightStore> store;
+};
+namespace g2p_en {
+class EnglishG2P;
+}
+}
+
+namespace engine::models::kokoro_tts {
+
+struct KokoroVoicePack {
+    std::string id;
+    std::string language_code;
+    int64_t rows = 0;
+    int64_t cols = 0;
+    std::vector<float> values;
+};
+
+struct KokoroAssets {
+    std::shared_ptr<struct KokoroPackage> package;
+    std::shared_ptr<class MultilingualG2P> multilingual_g2p;
+    std::filesystem::path model_root;
+    engine::io::json::Value config;
+    std::shared_ptr<const engine::assets::TensorSource> model_weights;
+    int64_t context_length = 512;
+    std::unordered_map<std::string, int32_t> vocab;
+    std::unordered_map<std::string, KokoroVoicePack> voices;
+    std::filesystem::path english_lexicon_dir;
+    std::shared_ptr<const kokoro_ggml::g2p_en::EnglishG2P> english_g2p_us;
+    std::shared_ptr<const kokoro_ggml::g2p_en::EnglishG2P> english_g2p_gb;
+};
+
+std::shared_ptr<const KokoroAssets> load_kokoro_assets(const std::filesystem::path & model_root);
+
+std::shared_ptr<const kokoro_ggml::KokoroWeights> load_kokoro_backend_weights(
+    const KokoroAssets & assets,
+    ggml_backend_t backend,
+    engine::core::BackendType backend_type,
+    engine::assets::TensorStorageType matmul_storage_type,
+    engine::assets::TensorStorageType conv_storage_type,
+    size_t weight_context_bytes);
+
+}  // namespace engine::models::kokoro_tts

include/engine/models/kokoro_tts/decoder.h

@@ -0,0 +1,42 @@
+#pragma once
+
+#include "engine/models/kokoro_tts/predictor.h"
+
+#include <cstdint>
+#include <vector>
+
+typedef struct ggml_backend * ggml_backend_t;
+
+namespace kokoro_ggml {
+
+struct KokoroDecoderCapacityContract {
+    int64_t decoder_frames = 0;
+    int64_t conditioning_frames = 0;
+};
+
+class KokoroDecoderRuntime {
+public:
+    KokoroDecoderRuntime(
+        std::shared_ptr<const KokoroWeights> weights,
+        ggml_backend_t backend,
+        int n_threads,
+        bool use_device_backend,
+        uint64_t rng_seed,
+        KokoroDecoderCapacityContract contract);
+    ~KokoroDecoderRuntime();
+
+    KokoroDecoderRuntime(const KokoroDecoderRuntime &) = delete;
+    KokoroDecoderRuntime & operator=(const KokoroDecoderRuntime &) = delete;
+
+    void prepare(KokoroDecoderCapacityContract contract);
+
+    std::vector<float> decode(
+        const PredictorOutputs & predictor,
+        const std::vector<float> & ref_s);
+
+private:
+    struct Impl;
+    std::unique_ptr<Impl> impl_;
+};
+
+}  // namespace kokoro_ggml