From 8e92aec1080ff9a18ecf8af71910a43f11606eae Mon Sep 17 00:00:00 2001 From: justinjohn0306 Date: Wed, 1 Jul 2026 19:55:50 +0530 Subject: [PATCH 1/4] Add MOSS-TTS-Local model: Qwen3 backbone, depth transformer, generation Introduce the moss_tts_local family (OpenMOSS MOSS-TTS-Local-Transformer-v1.5): downloader package, config/asset loading, and family registration; the Qwen3 backbone (per-head QK-norm, GQA, NEOX RoPE); the 1-layer GPT-2-J depth transformer (fused c_attn, interleaved RoPE) with the 12-codebook generation loop and binary end gate; and the text processor that builds the generation prefix. Verified numerically: the C++ greedy first-frame codes match the fp32 Python reference exactly on all 12 codebooks (harness: moss_tts_local_smoke). Codec decoder and session/CLI wiring (audio output, voice cloning) still to come. Co-Authored-By: Claude Opus 4.8 (1M context) --- CMakeLists.txt | 15 + include/engine/models/moss_tts_local/assets.h | 78 ++++ .../engine/models/moss_tts_local/backbone.h | 50 +++ .../models/moss_tts_local/depth_transformer.h | 40 ++ .../engine/models/moss_tts_local/generator.h | 61 +++ include/engine/models/moss_tts_local/loader.h | 11 + .../engine/models/moss_tts_local/processor.h | 42 ++ src/framework/runtime/registry.cpp | 2 + src/models/moss_tts_local/assets.cpp | 153 +++++++ src/models/moss_tts_local/backbone.cpp | 422 ++++++++++++++++++ .../moss_tts_local/depth_transformer.cpp | 294 ++++++++++++ src/models/moss_tts_local/generator.cpp | 254 +++++++++++ src/models/moss_tts_local/loader.cpp | 118 +++++ src/models/moss_tts_local/processor.cpp | 112 +++++ tests/moss_tts_local/moss_tts_local_smoke.cpp | 149 +++++++ tools/model_manager.py | 51 +++ 16 files changed, 1852 insertions(+) create mode 100644 include/engine/models/moss_tts_local/assets.h create mode 100644 include/engine/models/moss_tts_local/backbone.h create mode 100644 include/engine/models/moss_tts_local/depth_transformer.h create mode 100644 include/engine/models/moss_tts_local/generator.h create mode 100644 include/engine/models/moss_tts_local/loader.h create mode 100644 include/engine/models/moss_tts_local/processor.h create mode 100644 src/models/moss_tts_local/assets.cpp create mode 100644 src/models/moss_tts_local/backbone.cpp create mode 100644 src/models/moss_tts_local/depth_transformer.cpp create mode 100644 src/models/moss_tts_local/generator.cpp create mode 100644 src/models/moss_tts_local/loader.cpp create mode 100644 src/models/moss_tts_local/processor.cpp create mode 100644 tests/moss_tts_local/moss_tts_local_smoke.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 9d35a92..0b660a9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -217,6 +217,12 @@ add_library(engine_runtime STATIC src/models/miotts/causal_lm.cpp src/models/miotts/session.cpp src/models/miotts/loader.cpp + src/models/moss_tts_local/assets.cpp + src/models/moss_tts_local/backbone.cpp + src/models/moss_tts_local/depth_transformer.cpp + src/models/moss_tts_local/generator.cpp + src/models/moss_tts_local/loader.cpp + src/models/moss_tts_local/processor.cpp src/models/voxcpm2/assets.cpp src/models/voxcpm2/audiovae.cpp src/models/voxcpm2/generator.cpp @@ -488,6 +494,15 @@ if (ENGINE_ENABLE_OPENMP) target_link_libraries(vibevoice_finetune_overlay_check PRIVATE OpenMP::OpenMP_CXX) endif() +add_executable(moss_tts_local_smoke + tests/moss_tts_local/moss_tts_local_smoke.cpp +) + +target_link_libraries(moss_tts_local_smoke PRIVATE engine_runtime ggml) +if (ENGINE_ENABLE_OPENMP) + target_link_libraries(moss_tts_local_smoke PRIVATE OpenMP::OpenMP_CXX) +endif() + if (ENGINE_BUILD_WARMBENCH) function(add_engine_warmbench target_name source_file) add_executable(${target_name} diff --git a/include/engine/models/moss_tts_local/assets.h b/include/engine/models/moss_tts_local/assets.h new file mode 100644 index 0000000..0ea2420 --- /dev/null +++ b/include/engine/models/moss_tts_local/assets.h @@ -0,0 +1,78 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" + +#include +#include +#include +#include +#include +#include + +namespace engine::models::moss_tts_local { + +struct MossBackboneConfig { + int64_t hidden_size = 0; + int64_t intermediate_size = 0; + int64_t num_hidden_layers = 0; + int64_t num_attention_heads = 0; + int64_t num_key_value_heads = 0; + int64_t head_dim = 0; + int64_t max_position_embeddings = 0; + int64_t vocab_size = 0; + float rms_norm_eps = 1.0e-6F; + float rope_theta = 1000000.0F; + bool tie_word_embeddings = true; +}; + +struct MossLocalTransformerConfig { + int64_t hidden_size = 0; + int64_t intermediate_size = 0; + int64_t num_layers = 0; + int64_t num_heads = 0; + int64_t max_positions = 0; + float layer_norm_eps = 1.0e-6F; + float rope_base = 1000000.0F; +}; + +struct MossTTSLocalConfig { + MossBackboneConfig backbone; + MossLocalTransformerConfig local; + int64_t num_codebooks = 0; + int64_t audio_vocab_size = 0; + std::vector audio_codebook_sizes; + int64_t audio_pad_token_id = 0; + int64_t pad_token_id = 0; + int64_t im_start_token_id = 0; + int64_t im_end_token_id = 0; + int64_t audio_start_token_id = 0; + int64_t audio_end_token_id = 0; + int64_t audio_user_slot_token_id = 0; + int64_t audio_assistant_slot_token_id = 0; + int64_t sampling_rate = 0; + std::string local_text_head_mode; + std::string audio_tokenizer_name_or_path; +}; + +struct MossTTSLocalAssetPaths { + std::filesystem::path model_root; + std::filesystem::path config_path; + std::filesystem::path model_weights_path; + std::optional tokenizer_json_path; + std::optional tokenizer_config_path; + std::optional tokenizer_vocab_path; + std::optional tokenizer_merges_path; + std::optional audio_tokenizer_root; +}; + +struct MossTTSLocalAssets { + MossTTSLocalAssetPaths paths; + MossTTSLocalConfig config; + std::shared_ptr model_weights; +}; + +MossTTSLocalAssetPaths resolve_moss_tts_local_assets(const std::filesystem::path & model_path); +MossTTSLocalConfig load_moss_tts_local_config(const std::filesystem::path & model_path); +std::shared_ptr load_moss_tts_local_assets(const std::filesystem::path & model_path); + +} // namespace engine::models::moss_tts_local diff --git a/include/engine/models/moss_tts_local/backbone.h b/include/engine/models/moss_tts_local/backbone.h new file mode 100644 index 0000000..217375c --- /dev/null +++ b/include/engine/models/moss_tts_local/backbone.h @@ -0,0 +1,50 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" +#include "engine/models/moss_tts_local/assets.h" + +#include +#include +#include + +namespace engine::models::moss_tts_local { + +// Qwen3 backbone (transformer.*) runtime: loads the language-model weights and runs +// a prefill forward that returns the final hidden states. Text tokens are embedded in +// the graph; the summed audio-codebook contribution is supplied as a precomputed bias +// so the depth transformer and the generator can share a single embedding table. +class MossBackboneRuntime { +public: + MossBackboneRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution_context, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + ~MossBackboneRuntime(); + + MossBackboneRuntime(const MossBackboneRuntime &) = delete; + MossBackboneRuntime & operator=(const MossBackboneRuntime &) = delete; + + int64_t hidden_size() const noexcept; + + // Runs the backbone over a prefill sequence of text token ids and returns the + // final hidden states as [steps, hidden_size] row-major float. + std::vector forward_prefill(const std::vector & token_ids) const; + + // Same as forward_prefill, but adds the per-position audio-codebook embedding sum + // (audio_bias, [steps * hidden_size] row-major) to the text embedding before the + // decoder stack. This mirrors MossTTSLocalModel._build_inputs_embeds. + std::vector forward_prefill_fused( + const std::vector & token_ids, + const std::vector & audio_bias) const; + +private: + std::vector run_prefill(const std::vector & token_ids, const float * audio_bias) const; + + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::moss_tts_local diff --git a/include/engine/models/moss_tts_local/depth_transformer.h b/include/engine/models/moss_tts_local/depth_transformer.h new file mode 100644 index 0000000..9f53d7d --- /dev/null +++ b/include/engine/models/moss_tts_local/depth_transformer.h @@ -0,0 +1,40 @@ +#pragma once + +#include "engine/framework/core/execution_context.h" +#include "engine/models/moss_tts_local/assets.h" + +#include +#include +#include + +namespace engine::models::moss_tts_local { + +// Single-layer GPT-2-J depth transformer (local_transformer.*). Each generation frame +// seeds it with the backbone hidden state and then feeds back one audio-codebook +// embedding at a time; the returned final hidden state drives the RVQ code heads and +// the assistant/end gate. The layer count is tiny, so the weights are kept in f32 and +// the prefix is recomputed each step instead of maintaining a KV cache. +class MossDepthTransformer { +public: + MossDepthTransformer( + std::shared_ptr assets, + core::ExecutionContext & execution_context, + size_t graph_arena_bytes, + size_t weight_context_bytes); + ~MossDepthTransformer(); + + MossDepthTransformer(const MossDepthTransformer &) = delete; + MossDepthTransformer & operator=(const MossDepthTransformer &) = delete; + + int64_t hidden_size() const noexcept; + + // Runs the local transformer over inputs_embeds ([steps, hidden_size] row-major + // float) and returns the final-step hidden state as [hidden_size] float. + std::vector forward(const std::vector & inputs_embeds, int64_t steps) const; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::moss_tts_local diff --git a/include/engine/models/moss_tts_local/generator.h b/include/engine/models/moss_tts_local/generator.h new file mode 100644 index 0000000..75ea2a6 --- /dev/null +++ b/include/engine/models/moss_tts_local/generator.h @@ -0,0 +1,61 @@ +#pragma once + +#include "engine/models/moss_tts_local/assets.h" +#include "engine/models/moss_tts_local/backbone.h" +#include "engine/models/moss_tts_local/depth_transformer.h" + +#include +#include +#include + +namespace engine::models::moss_tts_local { + +// Sampling controls for the autoregressive audio-token loop. The audio_* fields drive +// the RVQ codebook heads and the text_* fields drive the binary assistant/end gate, +// mirroring the arguments of MossTTSLocalModel.generate. +struct MossGenerationOptions { + int64_t max_new_frames = 4096; + bool do_sample = true; + float audio_temperature = 1.0F; + float audio_top_p = 0.95F; + int audio_top_k = 50; + float audio_repetition_penalty = 1.0F; + float text_temperature = 1.0F; + float text_top_p = 1.0F; + int text_top_k = 50; + uint64_t seed = 0; +}; + +// Drives the MOSS-TTS-Local generation loop: for every frame it runs the Qwen3 backbone +// over the decoder sequence, then walks the depth transformer to emit one RVQ code per +// codebook. Generation stops when the binary gate selects the audio-end token. +class MossGenerator { +public: + MossGenerator( + std::shared_ptr assets, + const MossBackboneRuntime & backbone, + const MossDepthTransformer & depth); + + MossGenerator(const MossGenerator &) = delete; + MossGenerator & operator=(const MossGenerator &) = delete; + + // text_tokens holds the text channel (input_ids[..., 0]) and audio_codes holds the + // n_vq audio channels flattened row-major as [seq, n_vq] (input_ids[..., 1:], with + // audio_pad_token_id where a position carries no code). Returns the generated frames + // as [num_frames][n_vq] audio codes. + std::vector> generate( + const std::vector & text_tokens, + const std::vector & audio_codes, + const MossGenerationOptions & options) const; + +private: + std::shared_ptr assets_; + const MossBackboneRuntime & backbone_; + const MossDepthTransformer & depth_; + int64_t hidden_size_ = 0; + int64_t num_codebooks_ = 0; + std::vector> audio_embeddings_; // [n_vq][codebook_size * hidden] + std::vector local_text_head_; // [2 * hidden] +}; + +} // namespace engine::models::moss_tts_local diff --git a/include/engine/models/moss_tts_local/loader.h b/include/engine/models/moss_tts_local/loader.h new file mode 100644 index 0000000..5813f5b --- /dev/null +++ b/include/engine/models/moss_tts_local/loader.h @@ -0,0 +1,11 @@ +#pragma once + +#include "engine/framework/runtime/model.h" + +#include + +namespace engine::models::moss_tts_local { + +std::shared_ptr make_moss_tts_local_loader(); + +} // namespace engine::models::moss_tts_local diff --git a/include/engine/models/moss_tts_local/processor.h b/include/engine/models/moss_tts_local/processor.h new file mode 100644 index 0000000..707b52e --- /dev/null +++ b/include/engine/models/moss_tts_local/processor.h @@ -0,0 +1,42 @@ +#pragma once + +#include "engine/models/moss_tts_local/assets.h" + +#include +#include +#include +#include +#include + +namespace engine::models::moss_tts_local { + +// Decoder input for a generation request: the text channel (input_ids[..., 0]) plus the +// n_vq audio channels flattened row-major as [seq, n_vq] (input_ids[..., 1:]). Every audio +// slot of the prompt carries audio_pad_token_id, matching MossTTSLocalProcessor._build_text_rows. +struct MossGenerationPrefix { + std::vector text_tokens; + std::vector audio_codes; +}; + +// Reproduces the text-only branch of MossTTSLocalProcessor: it renders the +// template, byte-level BPE encodes each piece with the Qwen tokenizer, and splices in the +// im_start/im_end/audio_start ids to open the assistant audio turn. Reference-audio voice +// cloning is a later phase and is intentionally not handled here. +class MossTextProcessor { +public: + explicit MossTextProcessor(std::shared_ptr assets); + ~MossTextProcessor(); + + MossTextProcessor(const MossTextProcessor &) = delete; + MossTextProcessor & operator=(const MossTextProcessor &) = delete; + + MossGenerationPrefix build_generation_prefix( + const std::string & text, + const std::optional & language = std::nullopt) const; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::moss_tts_local diff --git a/src/framework/runtime/registry.cpp b/src/framework/runtime/registry.cpp index c2b1414..77b9753 100644 --- a/src/framework/runtime/registry.cpp +++ b/src/framework/runtime/registry.cpp @@ -16,6 +16,7 @@ #include "engine/models/marblenet_vad/session.h" #include "engine/models/miocodec/loader.h" #include "engine/models/miotts/loader.h" +#include "engine/models/moss_tts_local/loader.h" #include "engine/models/omnivoice/loader.h" #include "engine/models/pocket_tts/loader.h" #include "engine/models/qwen3_asr/loader.h" @@ -216,6 +217,7 @@ ModelRegistry make_default_registry(const std::optional & engine::models::omnivoice::make_omnivoice_loader(), engine::models::miocodec::make_miocodec_loader(), engine::models::miotts::make_miotts_loader(), + engine::models::moss_tts_local::make_moss_tts_local_loader(), engine::models::voxcpm2::make_voxcpm2_loader(), engine::models::vibevoice::make_vibevoice_loader(), engine::models::heartmula::make_heartmula_loader(), diff --git a/src/models/moss_tts_local/assets.cpp b/src/models/moss_tts_local/assets.cpp new file mode 100644 index 0000000..e956f76 --- /dev/null +++ b/src/models/moss_tts_local/assets.cpp @@ -0,0 +1,153 @@ +#include "engine/models/moss_tts_local/assets.h" + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/io/filesystem.h" +#include "engine/framework/io/json.h" + +#include +#include +#include +#include + +namespace engine::models::moss_tts_local { +namespace json = engine::io::json; +namespace { + +std::filesystem::path resolve_model_root(const std::filesystem::path & model_path) { + if (engine::io::is_existing_directory(model_path)) { + return std::filesystem::weakly_canonical(model_path); + } + if (engine::io::is_existing_file(model_path)) { + return std::filesystem::weakly_canonical(model_path.parent_path()); + } + throw std::runtime_error("MOSS-TTS-Local model path does not exist: " + model_path.string()); +} + +assets::ResourceBundle make_resource_bundle(const std::filesystem::path & model_path) { + assets::ResourceBundle resources(resolve_model_root(model_path)); + resources.add_model_files({ + {"config", "config.json", true}, + {"tokenizer_json", "tokenizer.json", false}, + {"tokenizer_config", "tokenizer_config.json", false}, + {"tokenizer_vocab", "vocab.json", false}, + {"tokenizer_merges", "merges.txt", false}, + }); + return resources; +} + +void require_positive(int64_t value, const char * label) { + if (value <= 0) { + throw std::runtime_error(std::string("MOSS-TTS-Local config has non-positive ") + label); + } +} + +MossBackboneConfig parse_backbone_config(const engine::io::json::Value & value) { + MossBackboneConfig config; + config.hidden_size = json::require_i64(value, "hidden_size"); + config.intermediate_size = json::require_i64(value, "intermediate_size"); + config.num_hidden_layers = json::require_i64(value, "num_hidden_layers"); + config.num_attention_heads = json::require_i64(value, "num_attention_heads"); + config.num_key_value_heads = json::require_i64(value, "num_key_value_heads"); + require_positive(config.num_attention_heads, "backbone num_attention_heads"); + config.head_dim = + json::optional_i64(value, "head_dim", config.hidden_size / config.num_attention_heads); + config.max_position_embeddings = json::require_i64(value, "max_position_embeddings"); + config.vocab_size = json::require_i64(value, "vocab_size"); + config.rms_norm_eps = json::optional_f32(value, "rms_norm_eps", config.rms_norm_eps); + config.rope_theta = json::optional_f32(value, "rope_theta", config.rope_theta); + config.tie_word_embeddings = json::optional_bool(value, "tie_word_embeddings", config.tie_word_embeddings); + require_positive(config.hidden_size, "backbone hidden_size"); + require_positive(config.intermediate_size, "backbone intermediate_size"); + require_positive(config.num_hidden_layers, "backbone num_hidden_layers"); + require_positive(config.num_key_value_heads, "backbone num_key_value_heads"); + require_positive(config.head_dim, "backbone head_dim"); + require_positive(config.vocab_size, "backbone vocab_size"); + return config; +} + +MossLocalTransformerConfig parse_local_config(const engine::io::json::Value & value) { + MossLocalTransformerConfig config; + config.hidden_size = json::require_i64(value, "n_embd"); + config.intermediate_size = json::optional_i64(value, "n_inner", config.hidden_size * 4); + config.num_layers = json::require_i64(value, "n_layer"); + config.num_heads = json::require_i64(value, "n_head"); + config.max_positions = json::optional_i64(value, "n_positions", config.max_positions); + config.layer_norm_eps = json::optional_f32(value, "layer_norm_epsilon", config.layer_norm_eps); + config.rope_base = json::optional_f32(value, "rope_base", config.rope_base); + require_positive(config.hidden_size, "local n_embd"); + require_positive(config.num_layers, "local n_layer"); + require_positive(config.num_heads, "local n_head"); + return config; +} + +} // namespace + +MossTTSLocalAssetPaths resolve_moss_tts_local_assets(const std::filesystem::path & model_path) { + const auto root = resolve_model_root(model_path); + MossTTSLocalAssetPaths paths; + paths.model_root = root; + paths.config_path = root / "config.json"; + paths.model_weights_path = root / "model.safetensors"; + const auto add_optional_file = [&](std::optional & slot, const char * name) { + const auto path = root / name; + if (engine::io::is_existing_file(path)) { + slot = path; + } + }; + add_optional_file(paths.tokenizer_json_path, "tokenizer.json"); + add_optional_file(paths.tokenizer_config_path, "tokenizer_config.json"); + add_optional_file(paths.tokenizer_vocab_path, "vocab.json"); + add_optional_file(paths.tokenizer_merges_path, "merges.txt"); + const auto audio_tokenizer_root = root / "audio_tokenizer"; + if (engine::io::is_existing_directory(audio_tokenizer_root)) { + paths.audio_tokenizer_root = audio_tokenizer_root; + } + return paths; +} + +MossTTSLocalConfig load_moss_tts_local_config(const std::filesystem::path & model_path) { + auto resources = make_resource_bundle(model_path); + const auto root = resources.parse_json("config"); + const auto model_type = json::optional_string(root, "model_type", ""); + if (model_type != "moss_tts_local") { + throw std::runtime_error( + "MOSS-TTS-Local config model_type mismatch: expected moss_tts_local, got " + model_type); + } + MossTTSLocalConfig config; + config.backbone = parse_backbone_config(root.require("qwen3_config")); + config.local = parse_local_config(root.require("gpt2_config")); + config.num_codebooks = json::require_i64(root, "n_vq"); + config.audio_vocab_size = json::optional_i64(root, "audio_vocab_size", 0); + config.audio_codebook_sizes = json::optional_i64_array(root, "audio_codebook_sizes"); + config.audio_pad_token_id = json::optional_i64(root, "audio_pad_token_id", 0); + config.pad_token_id = json::optional_i64(root, "pad_token_id", 0); + config.im_start_token_id = json::optional_i64(root, "im_start_token_id", 0); + config.im_end_token_id = json::optional_i64(root, "im_end_token_id", 0); + config.audio_start_token_id = json::optional_i64(root, "audio_start_token_id", 0); + config.audio_end_token_id = json::optional_i64(root, "audio_end_token_id", 0); + config.audio_user_slot_token_id = json::optional_i64(root, "audio_user_slot_token_id", 0); + config.audio_assistant_slot_token_id = json::optional_i64(root, "audio_assistant_slot_token_id", 0); + config.sampling_rate = json::optional_i64(root, "sampling_rate", 0); + config.local_text_head_mode = json::optional_string(root, "local_text_head_mode", ""); + config.audio_tokenizer_name_or_path = json::optional_string(root, "audio_tokenizer_name_or_path", ""); + require_positive(config.num_codebooks, "n_vq"); + if (!config.audio_codebook_sizes.empty() && + static_cast(config.audio_codebook_sizes.size()) != config.num_codebooks) { + throw std::runtime_error("MOSS-TTS-Local audio_codebook_sizes length does not match n_vq"); + } + return config; +} + +std::shared_ptr load_moss_tts_local_assets(const std::filesystem::path & model_path) { + MossTTSLocalAssets assets; + assets.paths = resolve_moss_tts_local_assets(model_path); + assets.config = load_moss_tts_local_config(model_path); + if (!engine::io::is_existing_file(assets.paths.model_weights_path)) { + throw std::runtime_error( + "MOSS-TTS-Local weights not found: " + assets.paths.model_weights_path.string()); + } + assets.model_weights = assets::open_tensor_source(assets.paths.model_weights_path); + return std::make_shared(std::move(assets)); +} + +} // namespace engine::models::moss_tts_local diff --git a/src/models/moss_tts_local/backbone.cpp b/src/models/moss_tts_local/backbone.cpp new file mode 100644 index 0000000..83993ef --- /dev/null +++ b/src/models/moss_tts_local/backbone.cpp @@ -0,0 +1,422 @@ +#include "engine/models/moss_tts_local/backbone.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/module.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::moss_tts_local { +namespace { + +namespace modules = engine::modules; +namespace binding = engine::modules::binding; + +constexpr float kMaskedAttentionBias = std::numeric_limits::lowest(); + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct BackboneLayerWeights { + core::TensorValue input_norm; + core::TensorValue q_proj; + core::TensorValue k_proj; + core::TensorValue v_proj; + core::TensorValue o_proj; + core::TensorValue q_norm; + core::TensorValue k_norm; + core::TensorValue post_norm; + core::TensorValue gate_proj; + core::TensorValue up_proj; + core::TensorValue down_proj; +}; + +struct BackboneWeights { + std::shared_ptr store; + core::TensorValue embed_tokens; + std::vector layers; + core::TensorValue norm; +}; + +void validate_weight_storage_type(assets::TensorStorageType storage_type) { + switch (storage_type) { + case assets::TensorStorageType::Native: + case assets::TensorStorageType::F32: + case assets::TensorStorageType::F16: + case assets::TensorStorageType::BF16: + case assets::TensorStorageType::Q8_0: + return; + default: + throw std::runtime_error( + "MOSS-TTS-Local backbone weight_type supports only native, f32, f16, bf16, and q8_0"); + } +} + +BackboneWeights load_backbone_weights( + const MossTTSLocalAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) { + validate_weight_storage_type(storage_type); + const auto & config = assets.config.backbone; + const auto & source = *assets.model_weights; + BackboneWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "moss_tts_local.backbone.weights", + weight_context_bytes); + weights.embed_tokens = weights.store->load_tensor( + source, + "transformer.embed_tokens.weight", + storage_type, + {config.vocab_size, config.hidden_size}); + const int64_t dim = config.head_dim; + weights.layers.reserve(static_cast(config.num_hidden_layers)); + for (int64_t layer = 0; layer < config.num_hidden_layers; ++layer) { + const std::string prefix = "transformer.layers." + std::to_string(layer); + BackboneLayerWeights w; + w.input_norm = weights.store->load_f32_tensor(source, prefix + ".input_layernorm.weight", {config.hidden_size}); + w.q_proj = weights.store->load_tensor( + source, + prefix + ".self_attn.q_proj.weight", + storage_type, + {config.num_attention_heads * dim, config.hidden_size}); + w.k_proj = weights.store->load_tensor( + source, + prefix + ".self_attn.k_proj.weight", + storage_type, + {config.num_key_value_heads * dim, config.hidden_size}); + w.v_proj = weights.store->load_tensor( + source, + prefix + ".self_attn.v_proj.weight", + storage_type, + {config.num_key_value_heads * dim, config.hidden_size}); + w.o_proj = weights.store->load_tensor( + source, + prefix + ".self_attn.o_proj.weight", + storage_type, + {config.hidden_size, config.num_attention_heads * dim}); + w.q_norm = weights.store->load_f32_tensor(source, prefix + ".self_attn.q_norm.weight", {dim}); + w.k_norm = weights.store->load_f32_tensor(source, prefix + ".self_attn.k_norm.weight", {dim}); + w.post_norm = weights.store->load_f32_tensor( + source, + prefix + ".post_attention_layernorm.weight", + {config.hidden_size}); + w.gate_proj = weights.store->load_tensor( + source, + prefix + ".mlp.gate_proj.weight", + storage_type, + {config.intermediate_size, config.hidden_size}); + w.up_proj = weights.store->load_tensor( + source, + prefix + ".mlp.up_proj.weight", + storage_type, + {config.intermediate_size, config.hidden_size}); + w.down_proj = weights.store->load_tensor( + source, + prefix + ".mlp.down_proj.weight", + storage_type, + {config.hidden_size, config.intermediate_size}); + weights.layers.push_back(std::move(w)); + } + weights.norm = weights.store->load_f32_tensor(source, "transformer.norm.weight", {config.hidden_size}); + weights.store->upload(); + return weights; +} + +core::TensorValue ensure_contiguous(core::ModuleBuildContext & ctx, const core::TensorValue & value) { + return core::ensure_backend_addressable_layout(ctx, value); +} + +core::TensorValue reshape_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t heads, + int64_t dim) { + auto contiguous = ensure_contiguous(ctx, input); + return core::reshape_tensor( + ctx, + contiguous, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim})); +} + +core::TensorValue repeat_kv_heads(core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t repeats) { + if (repeats == 1) { + return input; + } + auto contiguous = ensure_contiguous(ctx, input); + const int64_t batch = contiguous.shape.dims[0]; + const int64_t kv_heads = contiguous.shape.dims[1]; + const int64_t steps = contiguous.shape.dims[2]; + const int64_t dim = contiguous.shape.dims[3]; + auto expanded = core::reshape_tensor( + ctx, + contiguous, + core::TensorShape::from_dims({batch, kv_heads, 1, steps * dim})); + expanded = modules::RepeatModule({core::TensorShape::from_dims({batch, kv_heads, repeats, steps * dim})}) + .build(ctx, expanded); + expanded = ensure_contiguous(ctx, expanded); + return core::reshape_tensor( + ctx, + expanded, + core::TensorShape::from_dims({batch, kv_heads * repeats, steps, dim})); +} + +core::TensorValue attention_from_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & q_heads, + const core::TensorValue & k_heads, + const core::TensorValue & v_heads, + int64_t dim, + const core::TensorValue & attention_mask) { + const modules::MatMulModule matmul; + auto scores = matmul.build( + ctx, + q_heads, + modules::TransposeModule({{0, 1, 3, 2}, k_heads.shape.rank}).build(ctx, k_heads)); + scores = core::ensure_backend_addressable_layout(ctx, scores); + auto attn = core::wrap_tensor( + ggml_soft_max_ext( + ctx.ggml, + scores.tensor, + attention_mask.tensor, + 1.0F / std::sqrt(static_cast(dim)), + 0.0F), + scores.shape, + GGML_TYPE_F32); + return matmul.build(ctx, attn, v_heads); +} + +core::TensorValue decoder_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const BackboneLayerWeights & weights, + const MossBackboneConfig & config, + const core::TensorValue & attention_mask) { + const int64_t dim = config.head_dim; + const int64_t kv_repeats = config.num_attention_heads / config.num_key_value_heads; + const modules::LinearModule q_proj( + binding::linear_config(config.hidden_size, config.num_attention_heads * dim, false)); + const modules::LinearModule k_proj( + binding::linear_config(config.hidden_size, config.num_key_value_heads * dim, false)); + const modules::LinearModule v_proj( + binding::linear_config(config.hidden_size, config.num_key_value_heads * dim, false)); + const modules::LinearModule o_proj( + binding::linear_config(config.num_attention_heads * dim, config.hidden_size, false)); + const modules::RMSNormModule hidden_norm({config.hidden_size, config.rms_norm_eps, true, false}); + const modules::RMSNormModule head_norm({dim, config.rms_norm_eps, true, false}); + auto x_norm = hidden_norm.build(ctx, input, binding::norm_data(ctx, weights.input_norm)); + auto q = q_proj.build(ctx, x_norm, binding::linear_data(ctx, weights.q_proj)); + auto k = k_proj.build(ctx, x_norm, binding::linear_data(ctx, weights.k_proj)); + auto v = v_proj.build(ctx, x_norm, binding::linear_data(ctx, weights.v_proj)); + q = head_norm.build(ctx, reshape_heads(ctx, q, config.num_attention_heads, dim), binding::norm_data(ctx, weights.q_norm)); + k = head_norm.build(ctx, reshape_heads(ctx, k, config.num_key_value_heads, dim), binding::norm_data(ctx, weights.k_norm)); + v = reshape_heads(ctx, v, config.num_key_value_heads, dim); + q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, q, positions); + k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, k, positions); + + auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); + auto k_heads = repeat_kv_heads(ctx, modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k), kv_repeats); + auto v_heads = repeat_kv_heads(ctx, modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v), kv_repeats); + auto context = attention_from_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask); + context = modules::TransposeModule({{0, 2, 1, 3}, context.shape.rank}).build(ctx, context); + context = ensure_contiguous(ctx, context); + context = core::reshape_tensor( + ctx, + context, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], config.num_attention_heads * dim})); + auto x = modules::AddModule{}.build(ctx, input, o_proj.build(ctx, context, binding::linear_data(ctx, weights.o_proj))); + + auto ff_in = hidden_norm.build(ctx, x, binding::norm_data(ctx, weights.post_norm)); + auto gate = modules::LinearModule( + binding::linear_config(config.hidden_size, config.intermediate_size, false)) + .build(ctx, ff_in, binding::linear_data(ctx, weights.gate_proj)); + gate = modules::SiluModule{}.build(ctx, gate); + auto up = modules::LinearModule( + binding::linear_config(config.hidden_size, config.intermediate_size, false)) + .build(ctx, ff_in, binding::linear_data(ctx, weights.up_proj)); + auto ff = modules::LinearModule( + binding::linear_config(config.intermediate_size, config.hidden_size, false)) + .build(ctx, modules::MulModule{}.build(ctx, gate, up), binding::linear_data(ctx, weights.down_proj)); + return modules::AddModule{}.build(ctx, x, ff); +} + +} // namespace + +struct MossBackboneRuntime::Impl { + std::shared_ptr assets; + ggml_backend_t backend = nullptr; + core::BackendType backend_type = core::BackendType::Cpu; + size_t graph_arena_bytes = 0; + BackboneWeights weights; +}; + +MossBackboneRuntime::MossBackboneRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution_context, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : impl_(std::make_unique()) { + if (assets == nullptr) { + throw std::runtime_error("MOSS-TTS-Local backbone requires assets"); + } + if (assets->model_weights == nullptr) { + throw std::runtime_error("MOSS-TTS-Local backbone requires model weights"); + } + impl_->backend = execution_context.backend(); + if (impl_->backend == nullptr) { + throw std::runtime_error("MOSS-TTS-Local backbone backend is not initialized"); + } + impl_->backend_type = execution_context.backend_type(); + impl_->graph_arena_bytes = graph_arena_bytes; + impl_->weights = load_backbone_weights( + *assets, + impl_->backend, + impl_->backend_type, + weight_context_bytes, + weight_storage_type); + impl_->assets = std::move(assets); +} + +MossBackboneRuntime::~MossBackboneRuntime() = default; + +int64_t MossBackboneRuntime::hidden_size() const noexcept { + return impl_->assets->config.backbone.hidden_size; +} + +std::vector MossBackboneRuntime::forward_prefill(const std::vector & token_ids) const { + return run_prefill(token_ids, nullptr); +} + +std::vector MossBackboneRuntime::forward_prefill_fused( + const std::vector & token_ids, + const std::vector & audio_bias) const { + const int64_t steps = static_cast(token_ids.size()); + const int64_t hidden = impl_->assets->config.backbone.hidden_size; + if (static_cast(audio_bias.size()) != steps * hidden) { + throw std::runtime_error("MOSS-TTS-Local backbone audio bias size does not match [steps, hidden]"); + } + return run_prefill(token_ids, audio_bias.data()); +} + +std::vector MossBackboneRuntime::run_prefill( + const std::vector & token_ids, + const float * audio_bias) const { + const int64_t steps = static_cast(token_ids.size()); + if (steps <= 0) { + throw std::runtime_error("MOSS-TTS-Local backbone prefill requires a non-empty token sequence"); + } + const auto & config = impl_->assets->config.backbone; + const auto & weights = impl_->weights; + + ggml_init_params params{impl_->graph_arena_bytes, nullptr, true}; + std::unique_ptr graph_ctx(ggml_init(params)); + if (graph_ctx == nullptr) { + throw std::runtime_error("failed to initialize MOSS-TTS-Local backbone graph context"); + } + core::ModuleBuildContext ctx{graph_ctx.get(), "moss_tts_local.backbone.forward", impl_->backend_type}; + + auto token_input = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, steps})); + ggml_set_input(token_input.tensor); + auto positions = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({steps})); + ggml_set_input(positions.tensor); + auto attention_mask = + core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, steps, steps})); + ggml_set_input(attention_mask.tensor); + core::TensorValue audio_bias_input; + if (audio_bias != nullptr) { + audio_bias_input = + core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, steps, config.hidden_size})); + ggml_set_input(audio_bias_input.tensor); + } + + auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size}) + .build(ctx, token_input, weights.embed_tokens); + if (audio_bias != nullptr) { + x = modules::AddModule{}.build(ctx, x, audio_bias_input); + } + for (const auto & layer : weights.layers) { + x = decoder_layer(ctx, x, positions, layer, config, attention_mask); + } + x = modules::RMSNormModule({config.hidden_size, config.rms_norm_eps, true, false}) + .build(ctx, x, binding::norm_data(ctx, weights.norm)); + x = ensure_contiguous(ctx, x); + auto hidden = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({steps, config.hidden_size})); + hidden = ensure_contiguous(ctx, hidden); + ggml_set_output(hidden.tensor); + + ggml_cgraph * graph = ggml_new_graph_custom(graph_ctx.get(), 65536, false); + ggml_build_forward_expand(graph, hidden.tensor); + + ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(impl_->backend)); + if (gallocr == nullptr || !ggml_gallocr_reserve(gallocr, graph) || !ggml_gallocr_alloc_graph(gallocr, graph)) { + if (gallocr != nullptr) { + ggml_gallocr_free(gallocr); + } + throw std::runtime_error("failed to allocate MOSS-TTS-Local backbone forward graph"); + } + + ggml_backend_tensor_set(token_input.tensor, token_ids.data(), 0, token_ids.size() * sizeof(int32_t)); + + std::vector position_host(static_cast(steps)); + for (int64_t i = 0; i < steps; ++i) { + position_host[static_cast(i)] = static_cast(i); + } + ggml_backend_tensor_set(positions.tensor, position_host.data(), 0, position_host.size() * sizeof(int32_t)); + + std::vector mask_host(static_cast(steps * steps), kMaskedAttentionBias); + for (int64_t q = 0; q < steps; ++q) { + for (int64_t k = 0; k <= q; ++k) { + mask_host[static_cast(q * steps + k)] = 0.0F; + } + } + ggml_backend_tensor_set(attention_mask.tensor, mask_host.data(), 0, mask_host.size() * sizeof(float)); + + if (audio_bias != nullptr) { + ggml_backend_tensor_set( + audio_bias_input.tensor, + audio_bias, + 0, + static_cast(steps * config.hidden_size) * sizeof(float)); + } + + const ggml_status status = ggml_backend_graph_compute(impl_->backend, graph); + ggml_backend_synchronize(impl_->backend); + if (status != GGML_STATUS_SUCCESS) { + ggml_gallocr_free(gallocr); + throw std::runtime_error("MOSS-TTS-Local backbone forward graph compute failed"); + } + + std::vector hidden_states(static_cast(steps * config.hidden_size)); + ggml_backend_tensor_get(hidden.tensor, hidden_states.data(), 0, hidden_states.size() * sizeof(float)); + ggml_gallocr_free(gallocr); + return hidden_states; +} + +} // namespace engine::models::moss_tts_local diff --git a/src/models/moss_tts_local/depth_transformer.cpp b/src/models/moss_tts_local/depth_transformer.cpp new file mode 100644 index 0000000..94f50d3 --- /dev/null +++ b/src/models/moss_tts_local/depth_transformer.cpp @@ -0,0 +1,294 @@ +#include "engine/models/moss_tts_local/depth_transformer.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/module.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::moss_tts_local { +namespace { + +namespace modules = engine::modules; +namespace binding = engine::modules::binding; + +constexpr float kMaskedAttentionBias = std::numeric_limits::lowest(); + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct DepthWeights { + std::shared_ptr store; + modules::NormWeights ln_1; + modules::LinearWeights c_attn; + modules::LinearWeights c_proj; + modules::NormWeights ln_2; + modules::LinearWeights fc_in; + modules::LinearWeights fc_out; + modules::NormWeights ln_f; +}; + +DepthWeights load_depth_weights( + const MossTTSLocalAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes) { + const auto & config = assets.config.local; + const auto & source = *assets.model_weights; + const int64_t hidden = config.hidden_size; + const int64_t inner = config.intermediate_size; + DepthWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "moss_tts_local.depth.weights", + weight_context_bytes); + auto & store = *weights.store; + const std::string prefix = "local_transformer.h.0"; + weights.ln_1 = binding::norm_from_source(store, source, prefix + ".ln_1", hidden); + weights.c_attn = binding::linear_from_source( + store, source, prefix + ".attn.c_attn", assets::TensorStorageType::F32, 3 * hidden, hidden, true); + weights.c_proj = binding::linear_from_source( + store, source, prefix + ".attn.c_proj", assets::TensorStorageType::F32, hidden, hidden, true); + weights.ln_2 = binding::norm_from_source(store, source, prefix + ".ln_2", hidden); + weights.fc_in = binding::linear_from_source( + store, source, prefix + ".mlp.fc_in", assets::TensorStorageType::F32, inner, hidden, true); + weights.fc_out = binding::linear_from_source( + store, source, prefix + ".mlp.fc_out", assets::TensorStorageType::F32, hidden, inner, true); + weights.ln_f = binding::norm_from_source(store, source, "local_transformer.ln_f", hidden); + store.upload(); + return weights; +} + +core::TensorValue ensure_contiguous(core::ModuleBuildContext & ctx, const core::TensorValue & value) { + return core::ensure_backend_addressable_layout(ctx, value); +} + +core::TensorValue split_projection( + core::ModuleBuildContext & ctx, + const core::TensorValue & fused, + int64_t index, + int64_t width) { + return modules::SliceModule({2, index * width, width}).build(ctx, fused); +} + +core::TensorValue reshape_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t heads, + int64_t dim) { + auto contiguous = ensure_contiguous(ctx, input); + return core::reshape_tensor( + ctx, + contiguous, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim})); +} + +core::TensorValue attention_from_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & q_heads, + const core::TensorValue & k_heads, + const core::TensorValue & v_heads, + int64_t dim, + const core::TensorValue & attention_mask) { + const modules::MatMulModule matmul; + auto scores = matmul.build( + ctx, + q_heads, + modules::TransposeModule({{0, 1, 3, 2}, k_heads.shape.rank}).build(ctx, k_heads)); + scores = core::ensure_backend_addressable_layout(ctx, scores); + auto attn = core::wrap_tensor( + ggml_soft_max_ext( + ctx.ggml, + scores.tensor, + attention_mask.tensor, + 1.0F / std::sqrt(static_cast(dim)), + 0.0F), + scores.shape, + GGML_TYPE_F32); + return matmul.build(ctx, attn, v_heads); +} + +core::TensorValue local_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const DepthWeights & weights, + const MossLocalTransformerConfig & config, + const core::TensorValue & attention_mask) { + const int64_t hidden = config.hidden_size; + const int64_t heads = config.num_heads; + const int64_t dim = hidden / heads; + const modules::LayerNormModule attn_norm({hidden, config.layer_norm_eps, true, true}); + const modules::LinearModule c_attn(binding::linear_config(hidden, 3 * hidden, true)); + const modules::LinearModule c_proj(binding::linear_config(hidden, hidden, true)); + + auto x_norm = attn_norm.build(ctx, input, weights.ln_1); + auto qkv = c_attn.build(ctx, x_norm, weights.c_attn); + auto q = reshape_heads(ctx, split_projection(ctx, qkv, 0, hidden), heads, dim); + auto k = reshape_heads(ctx, split_projection(ctx, qkv, 1, hidden), heads, dim); + auto v = reshape_heads(ctx, split_projection(ctx, qkv, 2, hidden), heads, dim); + q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NORMAL, config.rope_base}).build(ctx, q, positions); + k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NORMAL, config.rope_base}).build(ctx, k, positions); + + auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); + auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k); + auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v); + auto context = attention_from_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask); + context = modules::TransposeModule({{0, 2, 1, 3}, context.shape.rank}).build(ctx, context); + context = ensure_contiguous(ctx, context); + context = core::reshape_tensor( + ctx, context, core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], hidden})); + auto x = modules::AddModule{}.build(ctx, input, c_proj.build(ctx, context, weights.c_proj)); + + auto ff_in = modules::LayerNormModule({hidden, config.layer_norm_eps, true, true}).build(ctx, x, weights.ln_2); + auto ff = modules::LinearModule(binding::linear_config(hidden, config.intermediate_size, true)) + .build(ctx, ff_in, weights.fc_in); + ff = modules::SiluModule{}.build(ctx, ff); + ff = modules::LinearModule(binding::linear_config(config.intermediate_size, hidden, true)) + .build(ctx, ff, weights.fc_out); + return modules::AddModule{}.build(ctx, x, ff); +} + +} // namespace + +struct MossDepthTransformer::Impl { + std::shared_ptr assets; + ggml_backend_t backend = nullptr; + core::BackendType backend_type = core::BackendType::Cpu; + size_t graph_arena_bytes = 0; + DepthWeights weights; +}; + +MossDepthTransformer::MossDepthTransformer( + std::shared_ptr assets, + core::ExecutionContext & execution_context, + size_t graph_arena_bytes, + size_t weight_context_bytes) + : impl_(std::make_unique()) { + if (assets == nullptr) { + throw std::runtime_error("MOSS-TTS-Local depth transformer requires assets"); + } + if (assets->model_weights == nullptr) { + throw std::runtime_error("MOSS-TTS-Local depth transformer requires model weights"); + } + if (assets->config.local.num_layers != 1) { + throw std::runtime_error("MOSS-TTS-Local depth transformer expects a single local_transformer layer"); + } + impl_->backend = execution_context.backend(); + if (impl_->backend == nullptr) { + throw std::runtime_error("MOSS-TTS-Local depth transformer backend is not initialized"); + } + impl_->backend_type = execution_context.backend_type(); + impl_->graph_arena_bytes = graph_arena_bytes; + impl_->weights = load_depth_weights(*assets, impl_->backend, impl_->backend_type, weight_context_bytes); + impl_->assets = std::move(assets); +} + +MossDepthTransformer::~MossDepthTransformer() = default; + +int64_t MossDepthTransformer::hidden_size() const noexcept { + return impl_->assets->config.local.hidden_size; +} + +std::vector MossDepthTransformer::forward(const std::vector & inputs_embeds, int64_t steps) const { + const auto & config = impl_->assets->config.local; + const int64_t hidden = config.hidden_size; + if (steps <= 0) { + throw std::runtime_error("MOSS-TTS-Local depth transformer forward requires a non-empty prefix"); + } + if (static_cast(inputs_embeds.size()) != steps * hidden) { + throw std::runtime_error("MOSS-TTS-Local depth transformer input size does not match [steps, hidden]"); + } + const auto & weights = impl_->weights; + + ggml_init_params params{impl_->graph_arena_bytes, nullptr, true}; + std::unique_ptr graph_ctx(ggml_init(params)); + if (graph_ctx == nullptr) { + throw std::runtime_error("failed to initialize MOSS-TTS-Local depth graph context"); + } + core::ModuleBuildContext ctx{graph_ctx.get(), "moss_tts_local.depth.forward", impl_->backend_type}; + + auto embeds_input = + core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, steps, hidden})); + ggml_set_input(embeds_input.tensor); + auto positions = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({steps})); + ggml_set_input(positions.tensor); + auto attention_mask = + core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, steps, steps})); + ggml_set_input(attention_mask.tensor); + + auto x = local_layer(ctx, embeds_input, positions, weights, config, attention_mask); + x = modules::LayerNormModule({hidden, config.layer_norm_eps, true, true}).build(ctx, x, weights.ln_f); + x = ensure_contiguous(ctx, x); + auto hidden_states = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({steps, hidden})); + hidden_states = ensure_contiguous(ctx, hidden_states); + ggml_set_output(hidden_states.tensor); + + ggml_cgraph * graph = ggml_new_graph_custom(graph_ctx.get(), 8192, false); + ggml_build_forward_expand(graph, hidden_states.tensor); + + ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(impl_->backend)); + if (gallocr == nullptr || !ggml_gallocr_reserve(gallocr, graph) || !ggml_gallocr_alloc_graph(gallocr, graph)) { + if (gallocr != nullptr) { + ggml_gallocr_free(gallocr); + } + throw std::runtime_error("failed to allocate MOSS-TTS-Local depth forward graph"); + } + + ggml_backend_tensor_set( + embeds_input.tensor, inputs_embeds.data(), 0, inputs_embeds.size() * sizeof(float)); + + std::vector position_host(static_cast(steps)); + for (int64_t i = 0; i < steps; ++i) { + position_host[static_cast(i)] = static_cast(i); + } + ggml_backend_tensor_set(positions.tensor, position_host.data(), 0, position_host.size() * sizeof(int32_t)); + + std::vector mask_host(static_cast(steps * steps), kMaskedAttentionBias); + for (int64_t q = 0; q < steps; ++q) { + for (int64_t k = 0; k <= q; ++k) { + mask_host[static_cast(q * steps + k)] = 0.0F; + } + } + ggml_backend_tensor_set(attention_mask.tensor, mask_host.data(), 0, mask_host.size() * sizeof(float)); + + const ggml_status status = ggml_backend_graph_compute(impl_->backend, graph); + ggml_backend_synchronize(impl_->backend); + if (status != GGML_STATUS_SUCCESS) { + ggml_gallocr_free(gallocr); + throw std::runtime_error("MOSS-TTS-Local depth forward graph compute failed"); + } + + std::vector last_hidden(static_cast(hidden)); + ggml_backend_tensor_get( + hidden_states.tensor, + last_hidden.data(), + static_cast((steps - 1) * hidden) * sizeof(float), + static_cast(hidden) * sizeof(float)); + ggml_gallocr_free(gallocr); + return last_hidden; +} + +} // namespace engine::models::moss_tts_local diff --git a/src/models/moss_tts_local/generator.cpp b/src/models/moss_tts_local/generator.cpp new file mode 100644 index 0000000..0f20378 --- /dev/null +++ b/src/models/moss_tts_local/generator.cpp @@ -0,0 +1,254 @@ +#include "engine/models/moss_tts_local/generator.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::moss_tts_local { +namespace { + +// Projects a hidden-state vector through a [rows, hidden] row-major weight matrix. The +// audio codebook heads are tied to the audio embeddings and the gate head is tiny, so +// these products are cheap to run on the host. +std::vector project( + const std::vector & weight, + const std::vector & hidden, + int64_t rows, + int64_t hidden_size) { + std::vector logits(static_cast(rows)); + for (int64_t row = 0; row < rows; ++row) { + const float * weight_row = weight.data() + static_cast(row) * hidden_size; + double accumulator = 0.0; + for (int64_t index = 0; index < hidden_size; ++index) { + accumulator += static_cast(weight_row[index]) * static_cast(hidden[static_cast(index)]); + } + logits[static_cast(row)] = static_cast(accumulator); + } + return logits; +} + +int32_t argmax_index(const std::vector & logits) { + if (logits.empty()) { + throw std::runtime_error("MOSS-TTS-Local sampler received empty logits"); + } + size_t best = 0; + for (size_t index = 1; index < logits.size(); ++index) { + if (logits[index] > logits[best]) { + best = index; + } + } + return static_cast(best); +} + +void apply_repetition_penalty( + std::vector & logits, + const std::vector & previous_codes, + float penalty) { + if (penalty == 1.0F) { + return; + } + if (penalty <= 0.0F) { + throw std::runtime_error("MOSS-TTS-Local repetition penalty must be positive"); + } + std::unordered_set seen_codes; + for (const int32_t code : previous_codes) { + if (code < 0 || code >= static_cast(logits.size())) { + continue; + } + if (!seen_codes.insert(code).second) { + continue; + } + float & value = logits[static_cast(code)]; + value = value < 0.0F ? value * penalty : value / penalty; + } +} + +int32_t sample_index( + const std::vector & logits, + int top_k, + float top_p, + float temperature, + std::mt19937 & rng) { + if (temperature <= 0.0F) { + throw std::runtime_error("MOSS-TTS-Local sampler temperature must be positive"); + } + std::vector indices; + indices.reserve(logits.size()); + for (size_t index = 0; index < logits.size(); ++index) { + if (std::isfinite(logits[index])) { + indices.push_back(static_cast(index)); + } + } + if (indices.empty()) { + throw std::runtime_error("MOSS-TTS-Local sampler has no finite logits"); + } + std::sort(indices.begin(), indices.end(), [&](int32_t lhs, int32_t rhs) { + return logits[static_cast(lhs)] > logits[static_cast(rhs)]; + }); + if (top_k > 0 && static_cast(indices.size()) > top_k) { + indices.resize(static_cast(top_k)); + } + + const float max_logit = logits[static_cast(indices.front())] / temperature; + std::vector weights; + weights.reserve(indices.size()); + double total = 0.0; + for (const int32_t index : indices) { + const double weight = std::exp(static_cast(logits[static_cast(index)] / temperature - max_logit)); + weights.push_back(weight); + total += weight; + } + if (top_p > 0.0F && top_p < 1.0F) { + double cumulative = 0.0; + size_t keep = weights.size(); + for (size_t index = 0; index < weights.size(); ++index) { + cumulative += weights[index] / total; + if (cumulative >= top_p) { + keep = index + 1; + break; + } + } + indices.resize(keep); + weights.resize(keep); + } + std::discrete_distribution distribution(weights.begin(), weights.end()); + return indices[distribution(rng)]; +} + +} // namespace + +MossGenerator::MossGenerator( + std::shared_ptr assets, + const MossBackboneRuntime & backbone, + const MossDepthTransformer & depth) + : assets_(std::move(assets)), + backbone_(backbone), + depth_(depth) { + if (assets_ == nullptr) { + throw std::runtime_error("MOSS-TTS-Local generator requires assets"); + } + if (assets_->model_weights == nullptr) { + throw std::runtime_error("MOSS-TTS-Local generator requires model weights"); + } + const auto & config = assets_->config; + hidden_size_ = config.backbone.hidden_size; + num_codebooks_ = config.num_codebooks; + if (config.local_text_head_mode != "binary") { + throw std::runtime_error("MOSS-TTS-Local generator only supports the binary local text head"); + } + const auto & source = *assets_->model_weights; + audio_embeddings_.reserve(static_cast(num_codebooks_)); + for (int64_t codebook = 0; codebook < num_codebooks_; ++codebook) { + const int64_t codebook_size = config.audio_codebook_sizes.empty() + ? config.audio_vocab_size + : config.audio_codebook_sizes[static_cast(codebook)]; + if (codebook_size <= 0) { + throw std::runtime_error("MOSS-TTS-Local generator has an invalid audio codebook size"); + } + audio_embeddings_.push_back(source.require_f32( + "audio_embeddings." + std::to_string(codebook) + ".weight", {codebook_size, hidden_size_})); + } + local_text_head_ = source.require_f32("local_text_lm_head.weight", {2, hidden_size_}); +} + +std::vector> MossGenerator::generate( + const std::vector & text_tokens, + const std::vector & audio_codes, + const MossGenerationOptions & options) const { + const auto & config = assets_->config; + const int64_t hidden = hidden_size_; + const int64_t n_vq = num_codebooks_; + const int32_t pad_code = static_cast(config.audio_pad_token_id); + const int32_t assistant_slot = static_cast(config.audio_assistant_slot_token_id); + + if (text_tokens.empty()) { + throw std::runtime_error("MOSS-TTS-Local generation requires a non-empty prompt"); + } + if (static_cast(audio_codes.size()) != static_cast(text_tokens.size()) * n_vq) { + throw std::runtime_error("MOSS-TTS-Local generation audio codes must be [seq, n_vq]"); + } + + std::vector sequence_text = text_tokens; + std::vector sequence_audio = audio_codes; + std::vector> generated_frames; + std::vector> code_history(static_cast(n_vq)); + std::mt19937 rng(options.seed); + + for (int64_t frame = 0; frame < options.max_new_frames; ++frame) { + const int64_t steps = static_cast(sequence_text.size()); + + // Fuse the summed audio-codebook embeddings into an additive bias so the backbone + // only has to embed the text channel in the graph. + std::vector audio_bias(static_cast(steps * hidden), 0.0F); + for (int64_t position = 0; position < steps; ++position) { + float * bias_row = audio_bias.data() + static_cast(position) * hidden; + for (int64_t codebook = 0; codebook < n_vq; ++codebook) { + const int32_t code = sequence_audio[static_cast(position * n_vq + codebook)]; + if (code == pad_code) { + continue; + } + const float * embedding = + audio_embeddings_[static_cast(codebook)].data() + static_cast(code) * hidden; + for (int64_t index = 0; index < hidden; ++index) { + bias_row[index] += embedding[index]; + } + } + } + + const std::vector backbone_hidden = backbone_.forward_prefill_fused(sequence_text, audio_bias); + std::vector local_hidden( + backbone_hidden.end() - hidden, + backbone_hidden.end()); + + // Seed the depth transformer with the backbone hidden state (local position 0). + std::vector local_embeds = local_hidden; + local_hidden = depth_.forward(local_embeds, 1); + + // Binary gate: continue with the assistant slot or stop on the audio-end token. + std::vector gate_logits = project(local_text_head_, local_hidden, 2, hidden); + const int32_t gate_index = options.do_sample + ? sample_index(gate_logits, options.text_top_k, options.text_top_p, options.text_temperature, rng) + : argmax_index(gate_logits); + if (gate_index != 0) { + break; + } + + std::vector frame_codes(static_cast(n_vq)); + for (int64_t codebook = 0; codebook < n_vq; ++codebook) { + const int64_t codebook_size = + static_cast(audio_embeddings_[static_cast(codebook)].size()) / hidden; + std::vector logits = + project(audio_embeddings_[static_cast(codebook)], local_hidden, codebook_size, hidden); + apply_repetition_penalty( + logits, code_history[static_cast(codebook)], options.audio_repetition_penalty); + const int32_t code = options.do_sample + ? sample_index(logits, options.audio_top_k, options.audio_top_p, options.audio_temperature, rng) + : argmax_index(logits); + frame_codes[static_cast(codebook)] = code; + code_history[static_cast(codebook)].push_back(code); + + if (codebook + 1 < n_vq) { + const float * embedding = + audio_embeddings_[static_cast(codebook)].data() + static_cast(code) * hidden; + local_embeds.insert(local_embeds.end(), embedding, embedding + hidden); + local_hidden = depth_.forward(local_embeds, codebook + 2); + } + } + generated_frames.push_back(frame_codes); + + // Append the emitted frame as the next decoder row (assistant slot + codes). + sequence_text.push_back(assistant_slot); + sequence_audio.insert(sequence_audio.end(), frame_codes.begin(), frame_codes.end()); + } + + return generated_frames; +} + +} // namespace engine::models::moss_tts_local diff --git a/src/models/moss_tts_local/loader.cpp b/src/models/moss_tts_local/loader.cpp new file mode 100644 index 0000000..0d668c3 --- /dev/null +++ b/src/models/moss_tts_local/loader.cpp @@ -0,0 +1,118 @@ +#include "engine/models/moss_tts_local/loader.h" + +#include "engine/framework/io/filesystem.h" +#include "engine/models/moss_tts_local/assets.h" + +#include +#include +#include +#include + +namespace engine::models::moss_tts_local { +namespace { + +std::filesystem::path resolve_model_root(const std::filesystem::path & model_path) { + if (engine::io::is_existing_directory(model_path)) { + return std::filesystem::weakly_canonical(model_path); + } + if (engine::io::is_existing_file(model_path)) { + return std::filesystem::weakly_canonical(model_path.parent_path()); + } + throw std::runtime_error("MOSS-TTS-Local model path does not exist: " + model_path.string()); +} + +bool has_moss_tts_local_assets(const std::filesystem::path & root) { + return engine::io::is_existing_file(root / "config.json") && + engine::io::is_existing_file(root / "model.safetensors") && + engine::io::is_existing_file(root / "tokenizer.json"); +} + +runtime::CapabilitySet moss_tts_local_capabilities() { + runtime::CapabilitySet capabilities; + capabilities.supported_tasks = { + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + }; + capabilities.languages = {"Auto"}; + capabilities.supports_speaker_reference = true; + return capabilities; +} + +runtime::ModelMetadata moss_tts_local_metadata() { + runtime::ModelMetadata metadata; + metadata.family = "moss_tts_local"; + metadata.variant = "MOSS-TTS-Local-Transformer-v1.5"; + metadata.description = "MOSS-TTS Local Transformer with the MOSS-Audio-Tokenizer-v2 codec."; + metadata.config_candidates = { + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + }; + metadata.weight_candidates = {"model.safetensors"}; + return metadata; +} + +std::vector discover_config_assets(const std::filesystem::path & root) { + return runtime::discover_named_assets( + root, + { + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + }); +} + +std::vector discover_weight_assets(const std::filesystem::path & root) { + return runtime::discover_named_assets(root, {"model.safetensors"}); +} + +class MossTTSLocalLoader final : public runtime::IVoiceModelLoader { +public: + std::string family() const override { + return "moss_tts_local"; + } + + bool can_load(const runtime::ModelLoadRequest & request) const override { + try { + const auto root = resolve_model_root(request.model_path); + return has_moss_tts_local_assets(root) && + (!request.family_hint.has_value() || *request.family_hint == family()); + } catch (...) { + return false; + } + } + + runtime::ModelInspection inspect(const runtime::ModelLoadRequest & request) const override { + const auto root = resolve_model_root(request.model_path); + const auto config = load_moss_tts_local_config(root); + (void) config; + runtime::ModelInspection inspection; + inspection.model_root = root; + inspection.metadata = moss_tts_local_metadata(); + inspection.capabilities = moss_tts_local_capabilities(); + inspection.cli.request_options = { + {"voice_ref", "path", "Reference speaker WAV for zero-shot voice cloning."}, + {"reference_text", "text", "Transcript of the reference audio when known."}, + {"language", "code", "Optional language tag hint for multilingual synthesis."}, + }; + inspection.discovered_configs = discover_config_assets(root); + inspection.discovered_weights = discover_weight_assets(root); + return inspection; + } + + std::unique_ptr load(const runtime::ModelLoadRequest & request) const override { + (void) request; + throw std::runtime_error("MOSS-TTS-Local generation is not implemented yet"); + } +}; + +} // namespace + +std::shared_ptr make_moss_tts_local_loader() { + return std::make_shared(); +} + +} // namespace engine::models::moss_tts_local diff --git a/src/models/moss_tts_local/processor.cpp b/src/models/moss_tts_local/processor.cpp new file mode 100644 index 0000000..1143ae8 --- /dev/null +++ b/src/models/moss_tts_local/processor.cpp @@ -0,0 +1,112 @@ +#include "engine/models/moss_tts_local/processor.h" + +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include + +namespace engine::models::moss_tts_local { +namespace { + +// User-template fragments, copied verbatim from MossTTSLocalProcessor so the encoded +// prompt matches the reference token-for-token. +constexpr const char * kUserRolePrefix = "user\n"; +constexpr const char * kUserReferencePrefix = "\n- Reference(s):\n"; +constexpr const char * kUserTextSuffix = "\n- Text:\n"; +constexpr const char * kUserInstSuffix = "\n"; +constexpr const char * kAssistantTurnPrefix = "\n"; +constexpr const char * kAssistantRolePrefix = "assistant\n"; +constexpr const char * kNoneValue = "None"; + +std::string normalize_template_value(const std::optional & value) { + if (!value.has_value()) { + return kNoneValue; + } + std::string resolved = *value; + const auto first = resolved.find_first_not_of(" \t\r\n"); + if (first == std::string::npos) { + return kNoneValue; + } + const auto last = resolved.find_last_not_of(" \t\r\n"); + resolved = resolved.substr(first, last - first + 1); + return resolved.empty() ? kNoneValue : resolved; +} + +// Mirrors _render_user_prompt_after_reference for the text-only request: every optional +// field defaults to "None" and only the language slot can carry a caller value. +std::string render_after_reference(const std::optional & language) { + return std::string("\n- Instruction:\n") + kNoneValue + + "\n- Tokens:\n" + kNoneValue + + "\n- Quality:\n" + kNoneValue + + "\n- Sound Event:\n" + kNoneValue + + "\n- Ambient Sound:\n" + kNoneValue + + "\n- Language:\n" + normalize_template_value(language) + + kUserTextSuffix; +} + +std::filesystem::path require_path( + const std::optional & path, + const char * label) { + if (!path.has_value()) { + throw std::runtime_error(std::string("MOSS-TTS-Local processor requires ") + label); + } + return *path; +} + +} // namespace + +struct MossTextProcessor::Impl { + std::shared_ptr assets; + std::shared_ptr tokenizer; + + void append_encoded(std::vector & tokens, const std::string & text) const { + const auto ids = tokenizer->encode(text); + tokens.insert(tokens.end(), ids.begin(), ids.end()); + } +}; + +MossTextProcessor::MossTextProcessor(std::shared_ptr assets) + : impl_(std::make_unique()) { + if (assets == nullptr) { + throw std::runtime_error("MOSS-TTS-Local processor requires assets"); + } + engine::tokenizers::LlamaBpeTokenizerSpec spec; + spec.vocab_path = require_path(assets->paths.tokenizer_vocab_path, "vocab.json"); + spec.merges_path = require_path(assets->paths.tokenizer_merges_path, "merges.txt"); + spec.tokenizer_config_path = require_path(assets->paths.tokenizer_config_path, "tokenizer_config.json"); + spec.tokenizer_json_path = assets->paths.tokenizer_json_path; + spec.pre_type = engine::tokenizers::LlamaBpePreTokenizer::Qwen2; + impl_->tokenizer = engine::tokenizers::load_llama_bpe_tokenizer(spec); + impl_->assets = std::move(assets); +} + +MossTextProcessor::~MossTextProcessor() = default; + +MossGenerationPrefix MossTextProcessor::build_generation_prefix( + const std::string & text, + const std::optional & language) const { + const auto & config = impl_->assets->config; + + std::vector tokens; + tokens.push_back(static_cast(config.im_start_token_id)); + impl_->append_encoded(tokens, kUserRolePrefix); + impl_->append_encoded(tokens, kUserReferencePrefix); + impl_->append_encoded(tokens, kNoneValue); + impl_->append_encoded(tokens, render_after_reference(language)); + impl_->append_encoded(tokens, text); + impl_->append_encoded(tokens, kUserInstSuffix); + tokens.push_back(static_cast(config.im_end_token_id)); + impl_->append_encoded(tokens, kAssistantTurnPrefix); + tokens.push_back(static_cast(config.im_start_token_id)); + impl_->append_encoded(tokens, kAssistantRolePrefix); + tokens.push_back(static_cast(config.audio_start_token_id)); + + MossGenerationPrefix prefix; + prefix.text_tokens = std::move(tokens); + prefix.audio_codes.assign( + prefix.text_tokens.size() * static_cast(config.num_codebooks), + static_cast(config.audio_pad_token_id)); + return prefix; +} + +} // namespace engine::models::moss_tts_local diff --git a/tests/moss_tts_local/moss_tts_local_smoke.cpp b/tests/moss_tts_local/moss_tts_local_smoke.cpp new file mode 100644 index 0000000..c65775b --- /dev/null +++ b/tests/moss_tts_local/moss_tts_local_smoke.cpp @@ -0,0 +1,149 @@ +// Standalone verification for the MOSS-TTS-Local backbone + depth transformer + generator. +// It loads the real weights, builds a text-only generation prefix, and runs a short greedy +// generation so we can confirm the pipeline executes and emits in-range RVQ codes. No audio +// is produced here: the MOSS-Audio-Tokenizer-v2 decoder is a separate phase. + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" +#include "engine/models/moss_tts_local/assets.h" +#include "engine/models/moss_tts_local/backbone.h" +#include "engine/models/moss_tts_local/depth_transformer.h" +#include "engine/models/moss_tts_local/generator.h" +#include "engine/models/moss_tts_local/processor.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +std::string arg_value(int argc, char ** argv, const std::string & name, const std::string & fallback) { + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) { + return argv[i + 1]; + } + } + return fallback; +} + +int int_arg(int argc, char ** argv, const std::string & name, int fallback) { + return std::stoi(arg_value(argc, argv, name, std::to_string(fallback))); +} + +engine::assets::TensorStorageType parse_weight_type(const std::string & value) { + if (value == "native") { + return engine::assets::TensorStorageType::Native; + } + if (value == "f32") { + return engine::assets::TensorStorageType::F32; + } + if (value == "f16") { + return engine::assets::TensorStorageType::F16; + } + if (value == "bf16") { + return engine::assets::TensorStorageType::BF16; + } + if (value == "q8_0") { + return engine::assets::TensorStorageType::Q8_0; + } + throw std::runtime_error("unsupported weight type: " + value); +} + +} // namespace + +int main(int argc, char ** argv) { + try { + const std::filesystem::path model_path = + arg_value(argc, argv, "--model", "models/MOSS-TTS-Local-Transformer-v1.5"); + const std::string text = arg_value( + argc, argv, "--text", "Hello, this is a test of the MOSS text to speech system."); + const int threads = int_arg(argc, argv, "--threads", 8); + const int max_frames = int_arg(argc, argv, "--frames", 12); + const auto weight_type = parse_weight_type(arg_value(argc, argv, "--weight-type", "native")); + + constexpr size_t kBackboneWeightContextBytes = 64ull * 1024 * 1024; + constexpr size_t kBackboneGraphArenaBytes = 512ull * 1024 * 1024; + constexpr size_t kDepthWeightContextBytes = 16ull * 1024 * 1024; + constexpr size_t kDepthGraphArenaBytes = 64ull * 1024 * 1024; + + std::cout << "model=" << model_path.string() << "\n"; + std::cout << "loading assets...\n" << std::flush; + auto assets = engine::models::moss_tts_local::load_moss_tts_local_assets(model_path); + + engine::core::BackendConfig backend_config; + backend_config.type = engine::core::BackendType::Cpu; + backend_config.device = 0; + backend_config.threads = threads; + engine::core::ExecutionContext execution_context(backend_config); + + std::cout << "loading backbone weights (36 layers)...\n" << std::flush; + engine::models::moss_tts_local::MossBackboneRuntime backbone( + assets, execution_context, kBackboneGraphArenaBytes, kBackboneWeightContextBytes, weight_type); + + std::cout << "loading depth transformer weights...\n" << std::flush; + engine::models::moss_tts_local::MossDepthTransformer depth( + assets, execution_context, kDepthGraphArenaBytes, kDepthWeightContextBytes); + + engine::models::moss_tts_local::MossGenerator generator(assets, backbone, depth); + engine::models::moss_tts_local::MossTextProcessor processor(assets); + + const auto prefix = processor.build_generation_prefix(text); + std::cout << "prompt_tokens=" << prefix.text_tokens.size() << "\n"; + std::cout << "first_token_ids=["; + const size_t preview = std::min(20, prefix.text_tokens.size()); + for (size_t i = 0; i < preview; ++i) { + std::cout << (i == 0 ? "" : ",") << prefix.text_tokens[i]; + } + std::cout << "]\n" << std::flush; + + engine::models::moss_tts_local::MossGenerationOptions options; + options.do_sample = false; + options.max_new_frames = max_frames; + + std::cout << "generating up to " << max_frames << " frames (greedy)...\n" << std::flush; + const auto frames = generator.generate(prefix.text_tokens, prefix.audio_codes, options); + + const int64_t n_vq = assets->config.num_codebooks; + std::vector min_code(static_cast(n_vq), std::numeric_limits::max()); + std::vector max_code(static_cast(n_vq), std::numeric_limits::min()); + bool all_in_range = true; + for (const auto & frame : frames) { + for (int64_t codebook = 0; codebook < n_vq; ++codebook) { + const int32_t code = frame[static_cast(codebook)]; + min_code[static_cast(codebook)] = std::min(min_code[static_cast(codebook)], code); + max_code[static_cast(codebook)] = std::max(max_code[static_cast(codebook)], code); + if (code < 0 || code >= 1024) { + all_in_range = false; + } + } + } + + std::cout << "\n=== RESULT ===\n"; + std::cout << "ran_without_crash=true\n"; + std::cout << "frames_produced=" << frames.size() << "\n"; + std::cout << "stopped_on_audio_end=" << ((int)frames.size() < max_frames ? "true" : "false") + << " (false means it hit the frame cap)\n"; + std::cout << "all_codes_in_[0,1023]=" << (all_in_range ? "true" : "false") << "\n"; + if (!frames.empty()) { + std::cout << "per_codebook_min_max=\n"; + for (int64_t codebook = 0; codebook < n_vq; ++codebook) { + std::cout << " cb" << codebook << ": min=" << min_code[static_cast(codebook)] + << " max=" << max_code[static_cast(codebook)] << "\n"; + } + std::cout << "first_frame_codes=["; + for (int64_t codebook = 0; codebook < n_vq; ++codebook) { + std::cout << (codebook == 0 ? "" : ",") << frames.front()[static_cast(codebook)]; + } + std::cout << "]\n"; + } + return 0; + } catch (const std::exception & error) { + std::cerr << "moss_tts_local_smoke failed: " << error.what() << "\n"; + return 1; + } +} diff --git a/tools/model_manager.py b/tools/model_manager.py index b692eb7..e8bc56d 100644 --- a/tools/model_manager.py +++ b/tools/model_manager.py @@ -231,6 +231,57 @@ def package_usage_examples(package: ModelPackage) -> list[str]: required_files=("config.json", "model-00001-of-00001.safetensors", "model.safetensors.index.json"), description="Subcomponent only. Use moss_tts for the full framework-ready MOSS runtime layout.", ), + ModelPackage( + id="moss_tts_local_v1_5", + display_name="MOSS TTS Local Transformer v1.5", + target_directory="MOSS-TTS-Local-Transformer-v1.5", + source=CompositeSnapshotSource( + placements=( + SnapshotPlacement( + source=SnapshotSource(repo_id="OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5"), + required_files=( + "config.json", + "model.safetensors", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + "special_tokens_map.json", + "added_tokens.json", + "chat_template.jinja", + ), + ), + SnapshotPlacement( + source=SnapshotSource(repo_id="OpenMOSS-Team/MOSS-Audio-Tokenizer-v2"), + target_subdir="audio_tokenizer", + required_files=( + "config.json", + "model.safetensors.index.json", + "model-00001-of-00003.safetensors", + "model-00002-of-00003.safetensors", + "model-00003-of-00003.safetensors", + ), + ), + ), + ), + required_files=( + "config.json", + "model.safetensors", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + "special_tokens_map.json", + "added_tokens.json", + "chat_template.jinja", + "audio_tokenizer/config.json", + "audio_tokenizer/model.safetensors.index.json", + "audio_tokenizer/model-00001-of-00003.safetensors", + "audio_tokenizer/model-00002-of-00003.safetensors", + "audio_tokenizer/model-00003-of-00003.safetensors", + ), + description="MOSS TTS Local Transformer v1.5 with the MOSS-Audio-Tokenizer-v2 codec dependency.", + ), ModelPackage( id="omnivoice", display_name="OmniVoice", From 34db64342f401e78cee7f9f017769c3543496fa3 Mon Sep 17 00:00:00 2001 From: justinjohn0306 Date: Wed, 1 Jul 2026 20:33:14 +0530 Subject: [PATCH 2/4] Add MOSS-Audio-Tokenizer-v2 decoder: codes -> 48 kHz stereo Port the codec decode path: the RLFQ dequantizer (12-of-32 residual quantizers, weight-normed 1x1 projections) turning codes into the 768-dim latent, then the "CNN-free" decoder -- six causal ProjectedTransformers (fused-QKV attention, interleaved RoPE, LayerScale, erf-GELU MLP) interleaved with reshape-based patch upsampling -- and channel de-interleave to stereo. Verified numerically against the fp32 Python model.decode: cosine 1.0 / max-abs-diff 1.6e-5 per channel (dequant alone is exact, cosine 1.0). Harnesses: codec_dequant_parity, codec_decode_parity. Session/CLI wiring and the voice-clone encoder remain (Phase 5). Co-Authored-By: Claude Opus 4.8 (1M context) --- CMakeLists.txt | 20 + .../models/moss_tts_local/codec_decoder.h | 42 ++ .../models/moss_tts_local/codec_quantizer.h | 42 ++ src/models/moss_tts_local/codec_decoder.cpp | 472 ++++++++++++++++++ src/models/moss_tts_local/codec_quantizer.cpp | 142 ++++++ tests/moss_tts_local/codec_decode_parity.cpp | 139 ++++++ tests/moss_tts_local/codec_dequant_parity.cpp | 70 +++ 7 files changed, 927 insertions(+) create mode 100644 include/engine/models/moss_tts_local/codec_decoder.h create mode 100644 include/engine/models/moss_tts_local/codec_quantizer.h create mode 100644 src/models/moss_tts_local/codec_decoder.cpp create mode 100644 src/models/moss_tts_local/codec_quantizer.cpp create mode 100644 tests/moss_tts_local/codec_decode_parity.cpp create mode 100644 tests/moss_tts_local/codec_dequant_parity.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 0b660a9..3ad9489 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -219,6 +219,8 @@ add_library(engine_runtime STATIC src/models/miotts/loader.cpp src/models/moss_tts_local/assets.cpp src/models/moss_tts_local/backbone.cpp + src/models/moss_tts_local/codec_decoder.cpp + src/models/moss_tts_local/codec_quantizer.cpp src/models/moss_tts_local/depth_transformer.cpp src/models/moss_tts_local/generator.cpp src/models/moss_tts_local/loader.cpp @@ -503,6 +505,24 @@ if (ENGINE_ENABLE_OPENMP) target_link_libraries(moss_tts_local_smoke PRIVATE OpenMP::OpenMP_CXX) endif() +add_executable(codec_dequant_parity + tests/moss_tts_local/codec_dequant_parity.cpp +) + +target_link_libraries(codec_dequant_parity PRIVATE engine_runtime ggml) +if (ENGINE_ENABLE_OPENMP) + target_link_libraries(codec_dequant_parity PRIVATE OpenMP::OpenMP_CXX) +endif() + +add_executable(codec_decode_parity + tests/moss_tts_local/codec_decode_parity.cpp +) + +target_link_libraries(codec_decode_parity PRIVATE engine_runtime ggml) +if (ENGINE_ENABLE_OPENMP) + target_link_libraries(codec_decode_parity PRIVATE OpenMP::OpenMP_CXX) +endif() + if (ENGINE_BUILD_WARMBENCH) function(add_engine_warmbench target_name source_file) add_executable(${target_name} diff --git a/include/engine/models/moss_tts_local/codec_decoder.h b/include/engine/models/moss_tts_local/codec_decoder.h new file mode 100644 index 0000000..ed28445 --- /dev/null +++ b/include/engine/models/moss_tts_local/codec_decoder.h @@ -0,0 +1,42 @@ +#pragma once + +#include "engine/framework/core/execution_context.h" + +#include +#include +#include +#include + +namespace engine::models::moss_tts_local { + +// MOSS-Audio-Tokenizer-v2 decoder: turns generated RVQ codes into a 48 kHz +// stereo waveform. The codec is "CNN-free" -- the decoder is a stack of causal +// Transformer blocks (interleaved RoPE, LayerScale, GELU MLP) separated by +// reshape-based patch upsamples, ending in a channel de-interleave that splits +// the jointly-processed stream back into left/right. The RLFQ dequantizer +// (codes -> latent) is provided by MossCodecDequantizer. +class MossCodecDecoder { +public: + MossCodecDecoder( + const std::filesystem::path & codec_dir, + core::ExecutionContext & execution_context, + int64_t num_quantizers, + size_t weight_context_bytes, + size_t graph_arena_bytes); + ~MossCodecDecoder(); + + MossCodecDecoder(const MossCodecDecoder &) = delete; + MossCodecDecoder & operator=(const MossCodecDecoder &) = delete; + + int64_t sampling_rate() const noexcept; + + // Decodes [num_quantizers][steps] codes into a stereo waveform returned as + // {left, right}, each with steps * 3840 samples at 48 kHz. + std::vector> decode(const std::vector> & codes) const; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::moss_tts_local diff --git a/include/engine/models/moss_tts_local/codec_quantizer.h b/include/engine/models/moss_tts_local/codec_quantizer.h new file mode 100644 index 0000000..33a34bd --- /dev/null +++ b/include/engine/models/moss_tts_local/codec_quantizer.h @@ -0,0 +1,42 @@ +#pragma once + +#include +#include +#include + +namespace engine::models::moss_tts_local { + +// Dequantizes MOSS-Audio-Tokenizer-v2 codes (RLFQ) into the codec's continuous +// latent, i.e. the input to the codec decoder stack. Codes are the +// [num_quantizers, steps] matrix produced by generation; the returned latent is +// [code_dim, steps] row-major (channel-major), matching the Python +// quantizer.decode_codes output [1, code_dim, steps]. This is the plain-linear +// dequant path (per-codebook embedding lookup -> weight-normalized 1x1 conv -> +// residual sum -> output projection); the transformer decoder is a later phase. +class MossCodecDequantizer { +public: + MossCodecDequantizer(const std::filesystem::path & codec_dir, int64_t num_quantizers); + + int64_t code_dim() const noexcept { return code_dim_; } + int64_t num_quantizers() const noexcept { return num_quantizers_; } + + std::vector decode(const std::vector> & codes) const; + +private: + struct Codebook { + std::vector table; // [codebook_size, codebook_dim] row-major + std::vector out_weight; // [rvq_dim, codebook_dim] row-major + std::vector out_bias; // [rvq_dim] + }; + + int64_t codebook_size_ = 0; + int64_t codebook_dim_ = 0; + int64_t rvq_dim_ = 0; + int64_t code_dim_ = 0; + int64_t num_quantizers_ = 0; + std::vector codebooks_; + std::vector output_weight_; // [code_dim, rvq_dim] row-major + std::vector output_bias_; // [code_dim] +}; + +} // namespace engine::models::moss_tts_local diff --git a/src/models/moss_tts_local/codec_decoder.cpp b/src/models/moss_tts_local/codec_decoder.cpp new file mode 100644 index 0000000..117b241 --- /dev/null +++ b/src/models/moss_tts_local/codec_decoder.cpp @@ -0,0 +1,472 @@ +#include "engine/models/moss_tts_local/codec_decoder.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/module.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" +#include "engine/models/moss_tts_local/codec_quantizer.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::moss_tts_local { +namespace { + +namespace modules = engine::modules; +namespace binding = engine::modules::binding; + +constexpr float kMaskedAttentionBias = std::numeric_limits::lowest(); +constexpr int64_t kCodeDim = 768; +constexpr int64_t kSamplesPerFrame = 3840; // downsample_rate (per interleaved stream frame) +constexpr float kRopeTheta = 10000.0F; +constexpr float kLayerNormEps = 1.0e-5F; + +// One decoder sub-transformer (ProjectedTransformer) description, mirroring the +// v2 codec config.json decoder_kwargs (indices decoder.0/2/4/6/8/10). context +// is the local-attention window in tokens at that stage's frame rate, computed +// as round(frame_rate * context_duration); the patch that follows doubles the +// frame rate (the last one expands by 240). +struct TransformerSpec { + int64_t input_dim; + int64_t output_dim; + int64_t d_model; + int64_t num_heads; + int64_t num_layers; + int64_t intermediate_size; + int64_t context; + int64_t patch_after; +}; + +constexpr TransformerSpec kDecoderSpecs[] = { + {kCodeDim, 1280, 1280, 20, 32, 5120, 125, 2}, + {640, 768, 768, 12, 12, 3072, 250, 2}, + {384, 768, 768, 12, 12, 3072, 400, 2}, + {384, 768, 768, 12, 12, 3072, 400, 2}, + {384, 768, 768, 12, 12, 3072, 400, 2}, + {384, 240, 768, 12, 12, 3072, 400, 240}, +}; + +constexpr size_t kNumTransformers = sizeof(kDecoderSpecs) / sizeof(kDecoderSpecs[0]); + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct LayerWeights { + core::TensorValue norm1_w; + core::TensorValue norm1_b; + core::TensorValue in_proj; // fused qkv [3 * d_model, d_model] + core::TensorValue out_proj; // [d_model, d_model] + core::TensorValue norm2_w; + core::TensorValue norm2_b; + core::TensorValue fc1; // [intermediate_size, d_model] + core::TensorValue fc2; // [d_model, intermediate_size] + core::TensorValue layer_scale1; // [d_model] + core::TensorValue layer_scale2; // [d_model] +}; + +struct TransformerWeights { + TransformerSpec spec; + core::TensorValue input_proj; // [d_model, input_dim] + core::TensorValue output_proj; // [output_dim, d_model] + std::vector layers; +}; + +// Opens the codec safetensors shards and resolves a tensor to the shard that +// holds it (the codec ships model-0000N-of-00003.safetensors + an index). +class CodecShards { +public: + explicit CodecShards(const std::filesystem::path & codec_dir) { + for (const auto & entry : std::filesystem::directory_iterator(codec_dir)) { + if (entry.is_regular_file() && entry.path().extension() == ".safetensors") { + sources_.push_back(assets::open_tensor_source(entry.path())); + } + } + if (sources_.empty()) { + throw std::runtime_error("MOSS codec has no safetensors shards: " + codec_dir.string()); + } + } + + const assets::TensorSource & source_for(const std::string & name) const { + for (const auto & source : sources_) { + if (source->has_tensor(name)) { + return *source; + } + } + throw std::runtime_error("MOSS codec tensor not found: " + name); + } + +private: + std::vector> sources_; +}; + +TransformerWeights load_transformer( + core::BackendWeightStore & store, + const CodecShards & shards, + const TransformerSpec & spec, + int64_t decoder_index) { + const std::string prefix = "decoder." + std::to_string(decoder_index); + const auto load = [&](const std::string & name, std::initializer_list shape) { + return store.load_tensor(shards.source_for(name), name, assets::TensorStorageType::F32, shape); + }; + const auto load_f32 = [&](const std::string & name, std::initializer_list shape) { + return store.load_f32_tensor(shards.source_for(name), name, shape); + }; + + TransformerWeights weights; + weights.spec = spec; + weights.input_proj = load(prefix + ".input_proj.weight", {spec.d_model, spec.input_dim}); + weights.output_proj = load(prefix + ".output_proj.weight", {spec.output_dim, spec.d_model}); + weights.layers.reserve(static_cast(spec.num_layers)); + for (int64_t layer = 0; layer < spec.num_layers; ++layer) { + const std::string lp = prefix + ".transformer.layers." + std::to_string(layer); + LayerWeights w; + w.norm1_w = load_f32(lp + ".norm1.weight", {spec.d_model}); + w.norm1_b = load_f32(lp + ".norm1.bias", {spec.d_model}); + w.in_proj = load(lp + ".self_attn.in_proj.weight", {3 * spec.d_model, spec.d_model}); + w.out_proj = load(lp + ".self_attn.out_proj.weight", {spec.d_model, spec.d_model}); + w.norm2_w = load_f32(lp + ".norm2.weight", {spec.d_model}); + w.norm2_b = load_f32(lp + ".norm2.bias", {spec.d_model}); + w.fc1 = load(lp + ".ffn.0.weight", {spec.intermediate_size, spec.d_model}); + w.fc2 = load(lp + ".ffn.2.weight", {spec.d_model, spec.intermediate_size}); + w.layer_scale1 = load_f32(lp + ".layer_scale_1.scale", {spec.d_model}); + w.layer_scale2 = load_f32(lp + ".layer_scale_2.scale", {spec.d_model}); + weights.layers.push_back(std::move(w)); + } + return weights; +} + +core::TensorValue ensure_contiguous(core::ModuleBuildContext & ctx, const core::TensorValue & value) { + return core::ensure_backend_addressable_layout(ctx, value); +} + +core::TensorValue reshape_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t heads, + int64_t dim) { + auto contiguous = ensure_contiguous(ctx, input); + return core::reshape_tensor( + ctx, + contiguous, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim})); +} + +// Slices a fused qkv projection [1, T, 3*d_model] into its q/k/v thirds, each +// [1, T, d_model]. The 3*d_model axis is laid out as [3, heads, head_dim]. +core::TensorValue slice_projection( + core::ModuleBuildContext & ctx, + const core::TensorValue & qkv, + int64_t which, + int64_t d_model, + int64_t steps) { + ggml_tensor * source = qkv.tensor; + const size_t element_size = ggml_element_size(source); + ggml_tensor * view = ggml_view_3d( + ctx.ggml, + source, + d_model, + steps, + 1, + source->nb[1], + source->nb[2], + static_cast(which * d_model) * element_size); + return core::wrap_tensor( + ggml_cont(ctx.ggml, view), + core::TensorShape::from_dims({1, steps, d_model}), + GGML_TYPE_F32); +} + +core::TensorValue attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & q_heads, + const core::TensorValue & k_heads, + const core::TensorValue & v_heads, + int64_t dim, + const core::TensorValue & mask) { + const modules::MatMulModule matmul; + auto scores = matmul.build( + ctx, + q_heads, + modules::TransposeModule({{0, 1, 3, 2}, k_heads.shape.rank}).build(ctx, k_heads)); + scores = core::ensure_backend_addressable_layout(ctx, scores); + auto attn = core::wrap_tensor( + ggml_soft_max_ext( + ctx.ggml, + scores.tensor, + mask.tensor, + 1.0F / std::sqrt(static_cast(dim)), + 0.0F), + scores.shape, + GGML_TYPE_F32); + return matmul.build(ctx, attn, v_heads); +} + +core::TensorValue transformer_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const LayerWeights & weights, + const TransformerSpec & spec, + const core::TensorValue & positions, + const core::TensorValue & mask, + int64_t steps) { + const int64_t dim = spec.d_model / spec.num_heads; + const modules::LayerNormModule norm({spec.d_model, kLayerNormEps, true, true}); + + auto normed = norm.build(ctx, input, binding::norm_data(ctx, weights.norm1_w, weights.norm1_b)); + auto qkv = modules::LinearModule(binding::linear_config(spec.d_model, 3 * spec.d_model, false)) + .build(ctx, normed, binding::linear_data(ctx, weights.in_proj)); + + auto q = slice_projection(ctx, qkv, 0, spec.d_model, steps); + auto k = slice_projection(ctx, qkv, 1, spec.d_model, steps); + auto v = slice_projection(ctx, qkv, 2, spec.d_model, steps); + + q = reshape_heads(ctx, q, spec.num_heads, dim); + k = reshape_heads(ctx, k, spec.num_heads, dim); + v = reshape_heads(ctx, v, spec.num_heads, dim); + q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NORMAL, kRopeTheta}).build(ctx, q, positions); + k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NORMAL, kRopeTheta}).build(ctx, k, positions); + + auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); + auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k); + auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v); + auto context = attention(ctx, q_heads, k_heads, v_heads, dim, mask); + context = modules::TransposeModule({{0, 2, 1, 3}, context.shape.rank}).build(ctx, context); + context = ensure_contiguous(ctx, context); + context = core::reshape_tensor(ctx, context, core::TensorShape::from_dims({1, steps, spec.d_model})); + auto attn_out = modules::LinearModule(binding::linear_config(spec.d_model, spec.d_model, false)) + .build(ctx, context, binding::linear_data(ctx, weights.out_proj)); + attn_out = core::wrap_tensor( + ggml_mul(ctx.ggml, attn_out.tensor, weights.layer_scale1.tensor), attn_out.shape, GGML_TYPE_F32); + auto x = modules::AddModule{}.build(ctx, input, attn_out); + + auto ff_in = norm.build(ctx, x, binding::norm_data(ctx, weights.norm2_w, weights.norm2_b)); + auto ff = modules::LinearModule(binding::linear_config(spec.d_model, spec.intermediate_size, false)) + .build(ctx, ff_in, binding::linear_data(ctx, weights.fc1)); + ff = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, ff); + ff = modules::LinearModule(binding::linear_config(spec.intermediate_size, spec.d_model, false)) + .build(ctx, ff, binding::linear_data(ctx, weights.fc2)); + ff = core::wrap_tensor( + ggml_mul(ctx.ggml, ff.tensor, weights.layer_scale2.tensor), ff.shape, GGML_TYPE_F32); + return modules::AddModule{}.build(ctx, x, ff); +} + +// ProjectedTransformer: input projection -> transformer stack -> output +// projection. Input/output are [1, steps, channels] (feature-last). +core::TensorValue run_transformer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const TransformerWeights & weights, + const core::TensorValue & positions, + const core::TensorValue & mask, + int64_t steps) { + const auto & spec = weights.spec; + auto x = modules::LinearModule(binding::linear_config(spec.input_dim, spec.d_model, false)) + .build(ctx, input, binding::linear_data(ctx, weights.input_proj)); + for (const auto & layer : weights.layers) { + x = transformer_layer(ctx, x, layer, spec, positions, mask, steps); + } + return modules::LinearModule(binding::linear_config(spec.d_model, spec.output_dim, false)) + .build(ctx, x, binding::linear_data(ctx, weights.output_proj)); +} + +// PatchedPretransform (decode/upsample): [1, l, d*patch] -> [1, l*patch, d]. +// Each frame is unpacked into `patch` consecutive frames along time, matching +// x.reshape(b, d, h, l).permute(0, 1, 3, 2).reshape(b, d, l * h). +core::TensorValue patch_upsample( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t patch) { + auto contiguous = ensure_contiguous(ctx, input); + const int64_t length = contiguous.shape.dims[1]; + const int64_t packed = contiguous.shape.dims[2]; + const int64_t channels = packed / patch; + ggml_tensor * reshaped = ggml_reshape_3d(ctx.ggml, contiguous.tensor, patch, channels, length); + ggml_tensor * permuted = ggml_cont(ctx.ggml, ggml_permute(ctx.ggml, reshaped, 1, 0, 2, 3)); + ggml_tensor * unpacked = ggml_reshape_2d(ctx.ggml, permuted, channels, length * patch); + return core::wrap_tensor( + unpacked, core::TensorShape::from_dims({1, length * patch, channels}), GGML_TYPE_F32); +} + +std::vector causal_context_mask(int64_t steps, int64_t context) { + std::vector mask(static_cast(steps * steps), kMaskedAttentionBias); + for (int64_t query = 0; query < steps; ++query) { + for (int64_t key = 0; key <= query; ++key) { + if (query - key < context) { + mask[static_cast(query * steps + key)] = 0.0F; + } + } + } + return mask; +} + +} // namespace + +struct MossCodecDecoder::Impl { + ggml_backend_t backend = nullptr; + core::BackendType backend_type = core::BackendType::Cpu; + size_t graph_arena_bytes = 0; + std::unique_ptr dequantizer; + std::unique_ptr store; + std::vector transformers; +}; + +MossCodecDecoder::MossCodecDecoder( + const std::filesystem::path & codec_dir, + core::ExecutionContext & execution_context, + int64_t num_quantizers, + size_t weight_context_bytes, + size_t graph_arena_bytes) + : impl_(std::make_unique()) { + impl_->backend = execution_context.backend(); + if (impl_->backend == nullptr) { + throw std::runtime_error("MOSS codec decoder backend is not initialized"); + } + impl_->backend_type = execution_context.backend_type(); + impl_->graph_arena_bytes = graph_arena_bytes; + impl_->dequantizer = std::make_unique(codec_dir, num_quantizers); + + CodecShards shards(codec_dir); + impl_->store = std::make_unique( + impl_->backend, impl_->backend_type, "moss_tts_local.codec.decoder", weight_context_bytes); + impl_->transformers.reserve(kNumTransformers); + for (size_t index = 0; index < kNumTransformers; ++index) { + impl_->transformers.push_back( + load_transformer(*impl_->store, shards, kDecoderSpecs[index], static_cast(2 * index))); + } + impl_->store->upload(); +} + +MossCodecDecoder::~MossCodecDecoder() = default; + +int64_t MossCodecDecoder::sampling_rate() const noexcept { + return 48000; +} + +std::vector> MossCodecDecoder::decode( + const std::vector> & codes) const { + const int64_t frames = codes.empty() ? 0 : static_cast(codes.front().size()); + if (frames <= 0) { + throw std::runtime_error("MOSS codec decoder requires a non-empty code sequence"); + } + + // Codes -> continuous latent [code_dim, frames] (channel-major), transposed + // into the feature-last [1, frames, code_dim] layout the decoder expects. + const auto latent = impl_->dequantizer->decode(codes); + std::vector latent_input(static_cast(frames * kCodeDim)); + for (int64_t channel = 0; channel < kCodeDim; ++channel) { + for (int64_t step = 0; step < frames; ++step) { + latent_input[static_cast(step * kCodeDim + channel)] = + latent[static_cast(channel * frames + step)]; + } + } + + ggml_init_params params{impl_->graph_arena_bytes, nullptr, true}; + std::unique_ptr graph_ctx(ggml_init(params)); + if (graph_ctx == nullptr) { + throw std::runtime_error("failed to initialize MOSS codec decoder graph context"); + } + core::ModuleBuildContext ctx{graph_ctx.get(), "moss_tts_local.codec.decode", impl_->backend_type}; + + auto latent_tensor = + core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, frames, kCodeDim})); + ggml_set_input(latent_tensor.tensor); + + struct StageInput { + ggml_tensor * positions; + std::vector position_host; + ggml_tensor * mask; + std::vector mask_host; + }; + std::vector stage_inputs; + stage_inputs.reserve(impl_->transformers.size()); + + auto hidden = latent_tensor; + int64_t steps = frames; + for (const auto & transformer : impl_->transformers) { + auto positions = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({steps})); + ggml_set_input(positions.tensor); + auto mask = + core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, steps, steps})); + ggml_set_input(mask.tensor); + + StageInput stage; + stage.positions = positions.tensor; + stage.position_host.resize(static_cast(steps)); + for (int64_t i = 0; i < steps; ++i) { + stage.position_host[static_cast(i)] = static_cast(i); + } + stage.mask = mask.tensor; + stage.mask_host = causal_context_mask(steps, transformer.spec.context); + stage_inputs.push_back(std::move(stage)); + + hidden = run_transformer(ctx, hidden, transformer, positions, mask, steps); + hidden = patch_upsample(ctx, hidden, transformer.spec.patch_after); + steps *= transformer.spec.patch_after; + } + + hidden = ensure_contiguous(ctx, hidden); + ggml_set_output(hidden.tensor); + + ggml_cgraph * graph = ggml_new_graph_custom(graph_ctx.get(), 131072, false); + ggml_build_forward_expand(graph, hidden.tensor); + + ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(impl_->backend)); + if (gallocr == nullptr || !ggml_gallocr_reserve(gallocr, graph) || !ggml_gallocr_alloc_graph(gallocr, graph)) { + if (gallocr != nullptr) { + ggml_gallocr_free(gallocr); + } + throw std::runtime_error("failed to allocate MOSS codec decoder forward graph"); + } + + ggml_backend_tensor_set( + latent_tensor.tensor, latent_input.data(), 0, latent_input.size() * sizeof(float)); + for (const auto & stage : stage_inputs) { + ggml_backend_tensor_set( + stage.positions, stage.position_host.data(), 0, stage.position_host.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + stage.mask, stage.mask_host.data(), 0, stage.mask_host.size() * sizeof(float)); + } + + const ggml_status status = ggml_backend_graph_compute(impl_->backend, graph); + ggml_backend_synchronize(impl_->backend); + if (status != GGML_STATUS_SUCCESS) { + ggml_gallocr_free(gallocr); + throw std::runtime_error("MOSS codec decoder forward graph compute failed"); + } + + const int64_t interleaved = steps; // frames * 3840 * 2 (stereo interleaved) + std::vector flat(static_cast(interleaved)); + ggml_backend_tensor_get(hidden.tensor, flat.data(), 0, flat.size() * sizeof(float)); + ggml_gallocr_free(gallocr); + + // De-interleave the jointly-processed stream back into left/right channels + // (channel 0 = even samples, channel 1 = odd samples). + const int64_t per_channel = frames * kSamplesPerFrame; + std::vector> stereo(2, std::vector(static_cast(per_channel))); + for (int64_t i = 0; i < per_channel; ++i) { + stereo[0][static_cast(i)] = flat[static_cast(2 * i)]; + stereo[1][static_cast(i)] = flat[static_cast(2 * i + 1)]; + } + return stereo; +} + +} // namespace engine::models::moss_tts_local diff --git a/src/models/moss_tts_local/codec_quantizer.cpp b/src/models/moss_tts_local/codec_quantizer.cpp new file mode 100644 index 0000000..ecc6b56 --- /dev/null +++ b/src/models/moss_tts_local/codec_quantizer.cpp @@ -0,0 +1,142 @@ +#include "engine/models/moss_tts_local/codec_quantizer.h" + +#include "engine/framework/assets/tensor_source.h" + +#include +#include +#include +#include +#include + +namespace engine::models::moss_tts_local { +namespace { + +// Opens every safetensors shard in the codec directory and resolves tensors by +// name across them (the codec ships as model-0000N-of-00003.safetensors + an +// index; a name-scan avoids parsing the index for the handful of tensors the +// dequantizer needs). +class CodecShards { +public: + explicit CodecShards(const std::filesystem::path & codec_dir) { + for (const auto & entry : std::filesystem::directory_iterator(codec_dir)) { + if (entry.is_regular_file() && entry.path().extension() == ".safetensors") { + sources_.push_back(assets::open_tensor_source(entry.path())); + } + } + if (sources_.empty()) { + throw std::runtime_error("MOSS codec has no safetensors shards: " + codec_dir.string()); + } + } + + std::vector require_f32(const std::string & name) const { + for (const auto & source : sources_) { + if (source->has_tensor(name)) { + return source->require_f32(name); + } + } + throw std::runtime_error("MOSS codec tensor not found: " + name); + } + +private: + std::vector> sources_; +}; + +// Rebuilds a weight-normalized 1x1 conv weight from its parametrization +// (original0 = magnitude g per output channel, original1 = direction v), the +// PyTorch weight_norm(dim=0) reconstruction weight = g * v / ||v||. +std::vector reconstruct_weight_norm( + const std::vector & g, + const std::vector & v, + int64_t out_channels, + int64_t in_channels) { + std::vector weight(static_cast(out_channels * in_channels)); + for (int64_t o = 0; o < out_channels; ++o) { + double norm = 0.0; + for (int64_t k = 0; k < in_channels; ++k) { + const double value = v[static_cast(o * in_channels + k)]; + norm += value * value; + } + const float scale = static_cast(g[static_cast(o)] / std::sqrt(norm)); + for (int64_t k = 0; k < in_channels; ++k) { + weight[static_cast(o * in_channels + k)] = + v[static_cast(o * in_channels + k)] * scale; + } + } + return weight; +} + +std::vector load_wn_conv_weight( + const CodecShards & shards, + const std::string & prefix, + int64_t out_channels, + int64_t in_channels) { + const auto g = shards.require_f32(prefix + ".parametrizations.weight.original0"); + const auto v = shards.require_f32(prefix + ".parametrizations.weight.original1"); + return reconstruct_weight_norm(g, v, out_channels, in_channels); +} + +} // namespace + +MossCodecDequantizer::MossCodecDequantizer(const std::filesystem::path & codec_dir, int64_t num_quantizers) + : codebook_size_(1024), codebook_dim_(8), rvq_dim_(512), code_dim_(768), num_quantizers_(num_quantizers) { + if (num_quantizers_ <= 0) { + throw std::runtime_error("MOSS codec dequantizer requires a positive quantizer count"); + } + CodecShards shards(codec_dir); + + codebooks_.reserve(static_cast(num_quantizers_)); + for (int64_t index = 0; index < num_quantizers_; ++index) { + const std::string prefix = "quantizer.quantizers." + std::to_string(index); + Codebook codebook; + codebook.table = shards.require_f32(prefix + ".codebook.weight"); + codebook.out_weight = load_wn_conv_weight(shards, prefix + ".out_proj", rvq_dim_, codebook_dim_); + codebook.out_bias = shards.require_f32(prefix + ".out_proj.bias"); + codebooks_.push_back(std::move(codebook)); + } + + output_weight_ = load_wn_conv_weight(shards, "quantizer.output_proj", code_dim_, rvq_dim_); + output_bias_ = shards.require_f32("quantizer.output_proj.bias"); +} + +std::vector MossCodecDequantizer::decode(const std::vector> & codes) const { + if (static_cast(codes.size()) != num_quantizers_) { + throw std::runtime_error("MOSS codec dequantizer got the wrong number of codebooks"); + } + const int64_t steps = codes.empty() ? 0 : static_cast(codes.front().size()); + if (steps <= 0) { + throw std::runtime_error("MOSS codec dequantizer requires a non-empty code sequence"); + } + + std::vector latent(static_cast(code_dim_ * steps)); + std::vector residual(static_cast(rvq_dim_)); + for (int64_t step = 0; step < steps; ++step) { + std::fill(residual.begin(), residual.end(), 0.0); + for (int64_t index = 0; index < num_quantizers_; ++index) { + const auto & codebook = codebooks_[static_cast(index)]; + const int64_t code = codes[static_cast(index)][static_cast(step)]; + if (code < 0 || code >= codebook_size_) { + throw std::runtime_error("MOSS codec code index out of range"); + } + const float * embedding = &codebook.table[static_cast(code * codebook_dim_)]; + for (int64_t out = 0; out < rvq_dim_; ++out) { + double sum = codebook.out_bias[static_cast(out)]; + const float * row = &codebook.out_weight[static_cast(out * codebook_dim_)]; + for (int64_t k = 0; k < codebook_dim_; ++k) { + sum += static_cast(row[k]) * static_cast(embedding[k]); + } + residual[static_cast(out)] += sum; + } + } + for (int64_t out = 0; out < code_dim_; ++out) { + double sum = output_bias_[static_cast(out)]; + const float * row = &output_weight_[static_cast(out * rvq_dim_)]; + for (int64_t k = 0; k < rvq_dim_; ++k) { + sum += static_cast(row[k]) * residual[static_cast(k)]; + } + latent[static_cast(out * steps + step)] = static_cast(sum); + } + } + return latent; +} + +} // namespace engine::models::moss_tts_local diff --git a/tests/moss_tts_local/codec_decode_parity.cpp b/tests/moss_tts_local/codec_decode_parity.cpp new file mode 100644 index 0000000..9709256 --- /dev/null +++ b/tests/moss_tts_local/codec_decode_parity.cpp @@ -0,0 +1,139 @@ +// Standalone verification for the MOSS-Audio-Tokenizer-v2 decoder (codes -> +// 48 kHz stereo waveform). It decodes a fixed, deterministic code matrix and +// prints per-channel statistics + writes a WAV, so the output can be diffed +// against the Python reference (model.decode with fp32 / autocast disabled). + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" +#include "engine/models/moss_tts_local/codec_decoder.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +std::string arg_value(int argc, char ** argv, const std::string & name, const std::string & fallback) { + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) { + return argv[i + 1]; + } + } + return fallback; +} + +int int_arg(int argc, char ** argv, const std::string & name, int fallback) { + return std::stoi(arg_value(argc, argv, name, std::to_string(fallback))); +} + +void write_pcm16_wav( + const std::string & path, + const std::vector> & channels, + int32_t sample_rate) { + const int32_t num_channels = static_cast(channels.size()); + const int64_t frames = channels.empty() ? 0 : static_cast(channels.front().size()); + const int32_t bits = 16; + const int32_t block_align = num_channels * bits / 8; + const int32_t byte_rate = sample_rate * block_align; + const int32_t data_bytes = static_cast(frames * block_align); + + std::ofstream out(path, std::ios::binary); + const auto put32 = [&](int32_t value) { out.write(reinterpret_cast(&value), 4); }; + const auto put16 = [&](int16_t value) { out.write(reinterpret_cast(&value), 2); }; + out.write("RIFF", 4); + put32(36 + data_bytes); + out.write("WAVE", 4); + out.write("fmt ", 4); + put32(16); + put16(1); + put16(static_cast(num_channels)); + put32(sample_rate); + put32(byte_rate); + put16(static_cast(block_align)); + put16(static_cast(bits)); + out.write("data", 4); + put32(data_bytes); + for (int64_t frame = 0; frame < frames; ++frame) { + for (int32_t channel = 0; channel < num_channels; ++channel) { + float sample = channels[static_cast(channel)][static_cast(frame)]; + sample = std::max(-1.0F, std::min(1.0F, sample)); + put16(static_cast(std::lround(sample * 32767.0F))); + } + } +} + +} // namespace + +int main(int argc, char ** argv) { + try { + const std::string default_codec = + "C:/Users/justi/.cache/huggingface/hub/models--OpenMOSS-Team--MOSS-Audio-Tokenizer-v2/" + "snapshots/f6e20e543b33d2c252a7ef71bdf8aa71e5ff9169"; + const std::filesystem::path codec_dir = arg_value(argc, argv, "--codec", default_codec); + const int frames = int_arg(argc, argv, "--frames", 12); + const int threads = int_arg(argc, argv, "--threads", 16); + const int num_quantizers = int_arg(argc, argv, "--quantizers", 12); + const std::string wav_path = arg_value(argc, argv, "--out", "moss_codec_decode_cpp.wav"); + + constexpr size_t kWeightContextBytes = 256ull * 1024 * 1024; + constexpr size_t kGraphArenaBytes = 1024ull * 1024 * 1024; + + std::vector> codes( + static_cast(num_quantizers), std::vector(static_cast(frames))); + for (int q = 0; q < num_quantizers; ++q) { + for (int t = 0; t < frames; ++t) { + codes[static_cast(q)][static_cast(t)] = (q * 37 + t * 5) % 1024; + } + } + + engine::core::BackendConfig backend_config; + backend_config.type = engine::core::BackendType::Cpu; + backend_config.device = 0; + backend_config.threads = threads; + engine::core::ExecutionContext execution_context(backend_config); + + std::cout << "codec=" << codec_dir.string() << "\n"; + std::cout << "loading decoder weights...\n" << std::flush; + engine::models::moss_tts_local::MossCodecDecoder decoder( + codec_dir, execution_context, num_quantizers, kWeightContextBytes, kGraphArenaBytes); + + std::cout << "decoding " << frames << " frames...\n" << std::flush; + const auto stereo = decoder.decode(codes); + + std::cout << "\n=== RESULT ===\n"; + std::cout << "channels=" << stereo.size() << " samples_per_channel=" << stereo.front().size() << "\n"; + for (size_t channel = 0; channel < stereo.size(); ++channel) { + const auto & data = stereo[channel]; + double sum = 0.0; + double sq = 0.0; + for (float value : data) { + sum += value; + sq += static_cast(value) * value; + } + const double mean = sum / static_cast(data.size()); + const double rms = std::sqrt(sq / static_cast(data.size())); + const double variance = sq / static_cast(data.size()) - mean * mean; + std::cout << "channel " << channel << ": mean=" << mean << " std=" + << std::sqrt(std::max(0.0, variance)) << " rms=" << rms << "\n"; + std::cout << " first16=["; + for (int i = 0; i < 16 && i < static_cast(data.size()); ++i) { + std::cout << (i == 0 ? "" : ", ") << data[static_cast(i)]; + } + std::cout << "]\n"; + } + + write_pcm16_wav(wav_path, stereo, 48000); + std::cout << "wrote " << wav_path << "\n"; + return 0; + } catch (const std::exception & error) { + std::cerr << "codec_decode_parity failed: " << error.what() << "\n"; + return 1; + } +} diff --git a/tests/moss_tts_local/codec_dequant_parity.cpp b/tests/moss_tts_local/codec_dequant_parity.cpp new file mode 100644 index 0000000..9211087 --- /dev/null +++ b/tests/moss_tts_local/codec_dequant_parity.cpp @@ -0,0 +1,70 @@ +// Parity harness for the MOSS-Audio-Tokenizer-v2 RLFQ dequantizer: runs the +// dequant on a fixed code matrix and dumps the latent for comparison against +// the Python reference (scripts/codec_dequant_ref.py). + +#include "engine/models/moss_tts_local/codec_quantizer.h" + +#include +#include +#include +#include +#include +#include +#include + +int main(int argc, char ** argv) { + const std::filesystem::path codec_dir = + argc > 1 ? argv[1] + : "C:/Users/justi/.cache/huggingface/hub/models--OpenMOSS-Team--MOSS-Audio-Tokenizer-v2/" + "snapshots/f6e20e543b33d2c252a7ef71bdf8aa71e5ff9169"; + const std::filesystem::path out_path = + argc > 2 ? argv[2] + : "C:/Users/justi/AppData/Local/Temp/claude/E--REPOS-audio-cpp/" + "62af4e53-c9e0-4e66-ac0e-27e93cec72c9/scratchpad/cpp_latent.txt"; + + constexpr int64_t kNumQuantizers = 12; + constexpr int64_t kSteps = 8; + + std::vector> codes(kNumQuantizers, std::vector(kSteps)); + for (int64_t i = 0; i < kNumQuantizers; ++i) { + for (int64_t t = 0; t < kSteps; ++t) { + codes[static_cast(i)][static_cast(t)] = static_cast((i * 37 + t * 5) % 1024); + } + } + + try { + engine::models::moss_tts_local::MossCodecDequantizer dequantizer(codec_dir, kNumQuantizers); + const std::vector latent = dequantizer.decode(codes); // [code_dim, steps] + const int64_t code_dim = dequantizer.code_dim(); + + double mean = 0.0; + for (const float value : latent) { + mean += value; + } + mean /= static_cast(latent.size()); + double var = 0.0; + for (const float value : latent) { + var += (value - mean) * (value - mean); + } + var /= static_cast(latent.size()); + + std::printf("shape %lld %lld\n", static_cast(code_dim), static_cast(kSteps)); + std::printf("first16"); + for (int i = 0; i < 16; ++i) { + std::printf(" %.6f", latent[static_cast(i)]); + } + std::printf("\n"); + std::printf("mean %.6f std %.6f\n", mean, std::sqrt(var)); + + std::ofstream out(out_path); + out.precision(8); + for (const float value : latent) { + out << value << "\n"; + } + std::printf("wrote %lld values to %s\n", static_cast(latent.size()), out_path.string().c_str()); + } catch (const std::exception & error) { + std::fprintf(stderr, "error: %s\n", error.what()); + return 1; + } + return 0; +} From a968fc96a6584638cd0b815d020d63d87e5a0d5a Mon Sep 17 00:00:00 2001 From: justinjohn0306 Date: Thu, 2 Jul 2026 10:37:44 +0530 Subject: [PATCH 3/4] Add MOSS-TTS-Local Phase 5: voice cloning + session/CLI wiring Wire the MOSS-TTS-Local session end to end: the text processor builds the generation prefix, the generator emits RVQ codes, and the codec decoder renders 48 kHz stereo. Add voice cloning via the MOSS-Audio-Tokenizer-v2 encoder (audio -> RLFQ codes), the structural mirror of the decoder; extract the shared ProjectedTransformer machinery into codec_transformer.h so the encoder and decoder share one implementation. The processor's clone prefix embeds the reference speaker's codes under "- Reference(s):"; the session resamples and loudness-normalizes a --voice-ref clip, encodes it, and seeds generation. Parity-verified against the transformers reference: generation loop 96/96 codes over 8 frames, encoder 300/300 codes, clone input_ids 100x13 exact, decoder cosine 1.0. Co-Authored-By: Claude Opus 4.8 (1M context) --- CMakeLists.txt | 11 + include/engine/models/moss_tts_local/assets.h | 5 + .../models/moss_tts_local/codec_encoder.h | 40 +++ .../models/moss_tts_local/codec_quantizer.h | 18 +- .../models/moss_tts_local/codec_transformer.h | 299 +++++++++++++++++ .../engine/models/moss_tts_local/processor.h | 17 +- .../engine/models/moss_tts_local/session.h | 50 +++ src/models/moss_tts_local/assets.cpp | 52 +++ src/models/moss_tts_local/codec_decoder.cpp | 316 ++---------------- src/models/moss_tts_local/codec_encoder.cpp | 208 ++++++++++++ src/models/moss_tts_local/codec_quantizer.cpp | 97 ++++++ src/models/moss_tts_local/loader.cpp | 44 ++- src/models/moss_tts_local/processor.cpp | 113 +++++-- src/models/moss_tts_local/session.cpp | 223 ++++++++++++ tests/moss_tts_local/codec_encode_parity.cpp | 154 +++++++++ tests/moss_tts_local/moss_tts_local_smoke.cpp | 57 ++++ 16 files changed, 1380 insertions(+), 324 deletions(-) create mode 100644 include/engine/models/moss_tts_local/codec_encoder.h create mode 100644 include/engine/models/moss_tts_local/codec_transformer.h create mode 100644 include/engine/models/moss_tts_local/session.h create mode 100644 src/models/moss_tts_local/codec_encoder.cpp create mode 100644 src/models/moss_tts_local/session.cpp create mode 100644 tests/moss_tts_local/codec_encode_parity.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 3ad9489..45f1ae8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -220,11 +220,13 @@ add_library(engine_runtime STATIC src/models/moss_tts_local/assets.cpp src/models/moss_tts_local/backbone.cpp src/models/moss_tts_local/codec_decoder.cpp + src/models/moss_tts_local/codec_encoder.cpp src/models/moss_tts_local/codec_quantizer.cpp src/models/moss_tts_local/depth_transformer.cpp src/models/moss_tts_local/generator.cpp src/models/moss_tts_local/loader.cpp src/models/moss_tts_local/processor.cpp + src/models/moss_tts_local/session.cpp src/models/voxcpm2/assets.cpp src/models/voxcpm2/audiovae.cpp src/models/voxcpm2/generator.cpp @@ -523,6 +525,15 @@ if (ENGINE_ENABLE_OPENMP) target_link_libraries(codec_decode_parity PRIVATE OpenMP::OpenMP_CXX) endif() +add_executable(codec_encode_parity + tests/moss_tts_local/codec_encode_parity.cpp +) + +target_link_libraries(codec_encode_parity PRIVATE engine_runtime ggml) +if (ENGINE_ENABLE_OPENMP) + target_link_libraries(codec_encode_parity PRIVATE OpenMP::OpenMP_CXX) +endif() + if (ENGINE_BUILD_WARMBENCH) function(add_engine_warmbench target_name source_file) add_executable(${target_name} diff --git a/include/engine/models/moss_tts_local/assets.h b/include/engine/models/moss_tts_local/assets.h index 0ea2420..9c9bf26 100644 --- a/include/engine/models/moss_tts_local/assets.h +++ b/include/engine/models/moss_tts_local/assets.h @@ -75,4 +75,9 @@ MossTTSLocalAssetPaths resolve_moss_tts_local_assets(const std::filesystem::path MossTTSLocalConfig load_moss_tts_local_config(const std::filesystem::path & model_path); std::shared_ptr load_moss_tts_local_assets(const std::filesystem::path & model_path); +// Locates the MOSS-Audio-Tokenizer-v2 codec weights: prefers a local audio_tokenizer/ +// directory next to the model, otherwise resolves the Hugging Face hub cache snapshot for +// the repo id recorded in config (audio_tokenizer_name_or_path). +std::filesystem::path resolve_moss_codec_dir(const MossTTSLocalAssets & assets); + } // namespace engine::models::moss_tts_local diff --git a/include/engine/models/moss_tts_local/codec_encoder.h b/include/engine/models/moss_tts_local/codec_encoder.h new file mode 100644 index 0000000..79c6fca --- /dev/null +++ b/include/engine/models/moss_tts_local/codec_encoder.h @@ -0,0 +1,40 @@ +#pragma once + +#include "engine/framework/core/execution_context.h" + +#include +#include +#include +#include + +namespace engine::models::moss_tts_local { + +// MOSS-Audio-Tokenizer-v2 encoder: turns a reference waveform into RLFQ codes +// for zero-shot voice cloning. It is the structural mirror of MossCodecDecoder -- +// stereo is interleaved into one stream, patched down and run through a stack of +// causal Transformer blocks (interleaved RoPE, LayerScale, GELU MLP), then the +// RLFQ quantizer selects the nearest codes. Produces the same [num_quantizers, +// frames] code matrix the generator consumes. +class MossCodecEncoder { +public: + MossCodecEncoder( + const std::filesystem::path & codec_dir, + core::ExecutionContext & execution_context, + int64_t num_quantizers, + size_t weight_context_bytes, + size_t graph_arena_bytes); + ~MossCodecEncoder(); + + MossCodecEncoder(const MossCodecEncoder &) = delete; + MossCodecEncoder & operator=(const MossCodecEncoder &) = delete; + + // Encodes a waveform given as {left, right} channels (each with the same + // per-channel sample count, 48 kHz) into [num_quantizers][frames] codes. + std::vector> encode(const std::vector> & channels) const; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::moss_tts_local diff --git a/include/engine/models/moss_tts_local/codec_quantizer.h b/include/engine/models/moss_tts_local/codec_quantizer.h index 33a34bd..b3ff91c 100644 --- a/include/engine/models/moss_tts_local/codec_quantizer.h +++ b/include/engine/models/moss_tts_local/codec_quantizer.h @@ -22,11 +22,21 @@ class MossCodecDequantizer { std::vector decode(const std::vector> & codes) const; + // Quantizes the encoder latent into codes: the inverse of decode(). `hidden` + // is [frames, code_dim] feature-last (row-major: frame * code_dim + channel), + // matching the codec encoder's output. Mirrors the RLFQ forward pass + // (input_proj -> per-quantizer in_proj -> L2-normalized nearest code -> + // residual subtraction) and returns the [num_quantizers][frames] code matrix. + std::vector> encode(const std::vector & hidden, int64_t frames) const; + private: struct Codebook { - std::vector table; // [codebook_size, codebook_dim] row-major - std::vector out_weight; // [rvq_dim, codebook_dim] row-major - std::vector out_bias; // [rvq_dim] + std::vector table; // [codebook_size, codebook_dim] row-major + std::vector table_normalized; // [codebook_size, codebook_dim], L2-normalized rows (encode) + std::vector out_weight; // [rvq_dim, codebook_dim] row-major + std::vector out_bias; // [rvq_dim] + std::vector in_weight; // [codebook_dim, rvq_dim] row-major (encode) + std::vector in_bias; // [codebook_dim] (encode) }; int64_t codebook_size_ = 0; @@ -37,6 +47,8 @@ class MossCodecDequantizer { std::vector codebooks_; std::vector output_weight_; // [code_dim, rvq_dim] row-major std::vector output_bias_; // [code_dim] + std::vector input_weight_; // [rvq_dim, code_dim] row-major (encode) + std::vector input_bias_; // [rvq_dim] (encode) }; } // namespace engine::models::moss_tts_local diff --git a/include/engine/models/moss_tts_local/codec_transformer.h b/include/engine/models/moss_tts_local/codec_transformer.h new file mode 100644 index 0000000..46e7ec4 --- /dev/null +++ b/include/engine/models/moss_tts_local/codec_transformer.h @@ -0,0 +1,299 @@ +#pragma once + +// Shared building blocks for the MOSS-Audio-Tokenizer-v2 codec transformer +// stacks. The encoder and decoder are structural mirrors: both are stacks of +// causal ProjectedTransformers (fused qkv, interleaved RoPE, LayerScale, erf +// GELU MLP, pre-norm LayerNorm) separated by reshape-based patch transforms. +// The only differences are the module order (decoder: transformer -> upsample; +// encoder: downsample -> transformer), the per-stage specs, and the safetensors +// prefix ("decoder"/"encoder"). Everything below is prefix- and spec-agnostic so +// both stacks can share it. + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/module.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::moss_tts_local::codec_detail { + +namespace modules = engine::modules; +namespace binding = engine::modules::binding; + +inline constexpr float kMaskedAttentionBias = std::numeric_limits::lowest(); +inline constexpr int64_t kCodeDim = 768; +inline constexpr int64_t kSamplesPerFrame = 3840; // downsample_rate (per interleaved stream frame) +inline constexpr float kRopeTheta = 10000.0F; +inline constexpr float kLayerNormEps = 1.0e-5F; + +// One ProjectedTransformer stage. `patch` is the reshape factor applied to the +// stage (after the transformer for the decoder, before it for the encoder); +// `context` is the local-attention window in tokens at that stage's frame rate. +struct TransformerSpec { + int64_t input_dim; + int64_t output_dim; + int64_t d_model; + int64_t num_heads; + int64_t num_layers; + int64_t intermediate_size; + int64_t context; + int64_t patch; +}; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct LayerWeights { + core::TensorValue norm1_w; + core::TensorValue norm1_b; + core::TensorValue in_proj; // fused qkv [3 * d_model, d_model] + core::TensorValue out_proj; // [d_model, d_model] + core::TensorValue norm2_w; + core::TensorValue norm2_b; + core::TensorValue fc1; // [intermediate_size, d_model] + core::TensorValue fc2; // [d_model, intermediate_size] + core::TensorValue layer_scale1; // [d_model] + core::TensorValue layer_scale2; // [d_model] +}; + +struct TransformerWeights { + TransformerSpec spec; + core::TensorValue input_proj; // [d_model, input_dim] + core::TensorValue output_proj; // [output_dim, d_model] + std::vector layers; +}; + +// Opens the codec safetensors shards and resolves a tensor to the shard that +// holds it (the codec ships model-0000N-of-00003.safetensors + an index). +class CodecShards { +public: + explicit CodecShards(const std::filesystem::path & codec_dir) { + for (const auto & entry : std::filesystem::directory_iterator(codec_dir)) { + if (entry.is_regular_file() && entry.path().extension() == ".safetensors") { + sources_.push_back(assets::open_tensor_source(entry.path())); + } + } + if (sources_.empty()) { + throw std::runtime_error("MOSS codec has no safetensors shards: " + codec_dir.string()); + } + } + + const assets::TensorSource & source_for(const std::string & name) const { + for (const auto & source : sources_) { + if (source->has_tensor(name)) { + return *source; + } + } + throw std::runtime_error("MOSS codec tensor not found: " + name); + } + +private: + std::vector> sources_; +}; + +// Loads one ProjectedTransformer's weights. `stack_prefix` is "decoder" or +// "encoder"; `module_index` is the module's position in that ModuleList. +inline TransformerWeights load_transformer( + core::BackendWeightStore & store, + const CodecShards & shards, + const TransformerSpec & spec, + const std::string & stack_prefix, + int64_t module_index) { + const std::string prefix = stack_prefix + "." + std::to_string(module_index); + const auto load = [&](const std::string & name, std::initializer_list shape) { + return store.load_tensor(shards.source_for(name), name, assets::TensorStorageType::F32, shape); + }; + const auto load_f32 = [&](const std::string & name, std::initializer_list shape) { + return store.load_f32_tensor(shards.source_for(name), name, shape); + }; + + TransformerWeights weights; + weights.spec = spec; + weights.input_proj = load(prefix + ".input_proj.weight", {spec.d_model, spec.input_dim}); + weights.output_proj = load(prefix + ".output_proj.weight", {spec.output_dim, spec.d_model}); + weights.layers.reserve(static_cast(spec.num_layers)); + for (int64_t layer = 0; layer < spec.num_layers; ++layer) { + const std::string lp = prefix + ".transformer.layers." + std::to_string(layer); + LayerWeights w; + w.norm1_w = load_f32(lp + ".norm1.weight", {spec.d_model}); + w.norm1_b = load_f32(lp + ".norm1.bias", {spec.d_model}); + w.in_proj = load(lp + ".self_attn.in_proj.weight", {3 * spec.d_model, spec.d_model}); + w.out_proj = load(lp + ".self_attn.out_proj.weight", {spec.d_model, spec.d_model}); + w.norm2_w = load_f32(lp + ".norm2.weight", {spec.d_model}); + w.norm2_b = load_f32(lp + ".norm2.bias", {spec.d_model}); + w.fc1 = load(lp + ".ffn.0.weight", {spec.intermediate_size, spec.d_model}); + w.fc2 = load(lp + ".ffn.2.weight", {spec.d_model, spec.intermediate_size}); + w.layer_scale1 = load_f32(lp + ".layer_scale_1.scale", {spec.d_model}); + w.layer_scale2 = load_f32(lp + ".layer_scale_2.scale", {spec.d_model}); + weights.layers.push_back(std::move(w)); + } + return weights; +} + +inline core::TensorValue ensure_contiguous(core::ModuleBuildContext & ctx, const core::TensorValue & value) { + return core::ensure_backend_addressable_layout(ctx, value); +} + +inline core::TensorValue reshape_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t heads, + int64_t dim) { + auto contiguous = ensure_contiguous(ctx, input); + return core::reshape_tensor( + ctx, + contiguous, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim})); +} + +// Slices a fused qkv projection [1, T, 3*d_model] into its q/k/v thirds, each +// [1, T, d_model]. The 3*d_model axis is laid out as [3, heads, head_dim]. +inline core::TensorValue slice_projection( + core::ModuleBuildContext & ctx, + const core::TensorValue & qkv, + int64_t which, + int64_t d_model, + int64_t steps) { + ggml_tensor * source = qkv.tensor; + const size_t element_size = ggml_element_size(source); + ggml_tensor * view = ggml_view_3d( + ctx.ggml, + source, + d_model, + steps, + 1, + source->nb[1], + source->nb[2], + static_cast(which * d_model) * element_size); + return core::wrap_tensor( + ggml_cont(ctx.ggml, view), + core::TensorShape::from_dims({1, steps, d_model}), + GGML_TYPE_F32); +} + +inline core::TensorValue attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & q_heads, + const core::TensorValue & k_heads, + const core::TensorValue & v_heads, + int64_t dim, + const core::TensorValue & mask) { + const modules::MatMulModule matmul; + auto scores = matmul.build( + ctx, + q_heads, + modules::TransposeModule({{0, 1, 3, 2}, k_heads.shape.rank}).build(ctx, k_heads)); + scores = core::ensure_backend_addressable_layout(ctx, scores); + auto attn = core::wrap_tensor( + ggml_soft_max_ext( + ctx.ggml, + scores.tensor, + mask.tensor, + 1.0F / std::sqrt(static_cast(dim)), + 0.0F), + scores.shape, + GGML_TYPE_F32); + return matmul.build(ctx, attn, v_heads); +} + +inline core::TensorValue transformer_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const LayerWeights & weights, + const TransformerSpec & spec, + const core::TensorValue & positions, + const core::TensorValue & mask, + int64_t steps) { + const int64_t dim = spec.d_model / spec.num_heads; + const modules::LayerNormModule norm({spec.d_model, kLayerNormEps, true, true}); + + auto normed = norm.build(ctx, input, binding::norm_data(ctx, weights.norm1_w, weights.norm1_b)); + auto qkv = modules::LinearModule(binding::linear_config(spec.d_model, 3 * spec.d_model, false)) + .build(ctx, normed, binding::linear_data(ctx, weights.in_proj)); + + auto q = slice_projection(ctx, qkv, 0, spec.d_model, steps); + auto k = slice_projection(ctx, qkv, 1, spec.d_model, steps); + auto v = slice_projection(ctx, qkv, 2, spec.d_model, steps); + + q = reshape_heads(ctx, q, spec.num_heads, dim); + k = reshape_heads(ctx, k, spec.num_heads, dim); + v = reshape_heads(ctx, v, spec.num_heads, dim); + q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NORMAL, kRopeTheta}).build(ctx, q, positions); + k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NORMAL, kRopeTheta}).build(ctx, k, positions); + + auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); + auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k); + auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v); + auto context = attention(ctx, q_heads, k_heads, v_heads, dim, mask); + context = modules::TransposeModule({{0, 2, 1, 3}, context.shape.rank}).build(ctx, context); + context = ensure_contiguous(ctx, context); + context = core::reshape_tensor(ctx, context, core::TensorShape::from_dims({1, steps, spec.d_model})); + auto attn_out = modules::LinearModule(binding::linear_config(spec.d_model, spec.d_model, false)) + .build(ctx, context, binding::linear_data(ctx, weights.out_proj)); + attn_out = core::wrap_tensor( + ggml_mul(ctx.ggml, attn_out.tensor, weights.layer_scale1.tensor), attn_out.shape, GGML_TYPE_F32); + auto x = modules::AddModule{}.build(ctx, input, attn_out); + + auto ff_in = norm.build(ctx, x, binding::norm_data(ctx, weights.norm2_w, weights.norm2_b)); + auto ff = modules::LinearModule(binding::linear_config(spec.d_model, spec.intermediate_size, false)) + .build(ctx, ff_in, binding::linear_data(ctx, weights.fc1)); + ff = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, ff); + ff = modules::LinearModule(binding::linear_config(spec.intermediate_size, spec.d_model, false)) + .build(ctx, ff, binding::linear_data(ctx, weights.fc2)); + ff = core::wrap_tensor( + ggml_mul(ctx.ggml, ff.tensor, weights.layer_scale2.tensor), ff.shape, GGML_TYPE_F32); + return modules::AddModule{}.build(ctx, x, ff); +} + +// ProjectedTransformer: input projection -> transformer stack -> output +// projection. Input/output are [1, steps, channels] (feature-last). +inline core::TensorValue run_transformer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const TransformerWeights & weights, + const core::TensorValue & positions, + const core::TensorValue & mask, + int64_t steps) { + const auto & spec = weights.spec; + auto x = modules::LinearModule(binding::linear_config(spec.input_dim, spec.d_model, false)) + .build(ctx, input, binding::linear_data(ctx, weights.input_proj)); + for (const auto & layer : weights.layers) { + x = transformer_layer(ctx, x, layer, spec, positions, mask, steps); + } + return modules::LinearModule(binding::linear_config(spec.d_model, spec.output_dim, false)) + .build(ctx, x, binding::linear_data(ctx, weights.output_proj)); +} + +inline std::vector causal_context_mask(int64_t steps, int64_t context) { + std::vector mask(static_cast(steps * steps), kMaskedAttentionBias); + for (int64_t query = 0; query < steps; ++query) { + for (int64_t key = 0; key <= query; ++key) { + if (query - key < context) { + mask[static_cast(query * steps + key)] = 0.0F; + } + } + } + return mask; +} + +} // namespace engine::models::moss_tts_local::codec_detail diff --git a/include/engine/models/moss_tts_local/processor.h b/include/engine/models/moss_tts_local/processor.h index 707b52e..6a78bbb 100644 --- a/include/engine/models/moss_tts_local/processor.h +++ b/include/engine/models/moss_tts_local/processor.h @@ -18,10 +18,12 @@ struct MossGenerationPrefix { std::vector audio_codes; }; -// Reproduces the text-only branch of MossTTSLocalProcessor: it renders the -// template, byte-level BPE encodes each piece with the Qwen tokenizer, and splices in the -// im_start/im_end/audio_start ids to open the assistant audio turn. Reference-audio voice -// cloning is a later phase and is intentionally not handled here. +// Reproduces the direct-generation branch of MossTTSLocalProcessor: it renders the +// template, byte-level BPE encodes each piece with the Qwen tokenizer, and +// splices in the im_start/im_end/audio_start ids to open the assistant audio turn. The +// voice-clone variant additionally embeds a reference speaker's RLFQ codes under +// "- Reference(s):" (audio_start, audio_user_slot rows carrying the codes, audio_end), +// mirroring MossTTSLocalProcessor._build_generation_or_voice_clone_codes. class MossTextProcessor { public: explicit MossTextProcessor(std::shared_ptr assets); @@ -34,6 +36,13 @@ class MossTextProcessor { const std::string & text, const std::optional & language = std::nullopt) const; + // Builds a voice-clone prompt. reference_codes is [num_codebooks][frames] as produced + // by MossCodecEncoder for the reference speaker. + MossGenerationPrefix build_clone_prefix( + const std::string & text, + const std::vector> & reference_codes, + const std::optional & language = std::nullopt) const; + private: struct Impl; std::unique_ptr impl_; diff --git a/include/engine/models/moss_tts_local/session.h b/include/engine/models/moss_tts_local/session.h new file mode 100644 index 0000000..09fb322 --- /dev/null +++ b/include/engine/models/moss_tts_local/session.h @@ -0,0 +1,50 @@ +#pragma once + +#include "engine/framework/runtime/session_base.h" +#include "engine/models/moss_tts_local/assets.h" +#include "engine/models/moss_tts_local/backbone.h" +#include "engine/models/moss_tts_local/codec_decoder.h" +#include "engine/models/moss_tts_local/codec_encoder.h" +#include "engine/models/moss_tts_local/depth_transformer.h" +#include "engine/models/moss_tts_local/generator.h" +#include "engine/models/moss_tts_local/processor.h" + +#include + +namespace engine::models::moss_tts_local { + +// Offline TTS session: renders text into a 48 kHz stereo waveform by chaining the verified +// pieces -- the text processor builds the generation prefix, the generator (Qwen3 backbone + +// depth transformer) emits RVQ codes frame by frame, and the codec decoder turns those codes +// into audio. When a speaker reference is supplied, the codec encoder turns it into RLFQ +// codes that seed a voice-clone prompt. +class MossTTSLocalSession final + : public runtime::RuntimeSessionBase + , public runtime::IOfflineVoiceTaskSession { +public: + MossTTSLocalSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets); + + std::string family() const override; + runtime::VoiceTaskKind task_kind() const override; + runtime::RunMode run_mode() const override; + void prepare(const runtime::SessionPreparationRequest & request) override; + runtime::TaskResult run(const runtime::TaskRequest & request) override; + +private: + runtime::TaskSpec task_; + std::shared_ptr assets_; + // Declared before the generator so the generator (which holds references to them) is + // destroyed first. + std::unique_ptr backbone_; + std::unique_ptr depth_; + std::unique_ptr processor_; + std::unique_ptr codec_; + std::unique_ptr generator_; + // Lazily built the first time a speaker reference is provided (voice cloning). + std::unique_ptr encoder_; +}; + +} // namespace engine::models::moss_tts_local diff --git a/src/models/moss_tts_local/assets.cpp b/src/models/moss_tts_local/assets.cpp index e956f76..79a603c 100644 --- a/src/models/moss_tts_local/assets.cpp +++ b/src/models/moss_tts_local/assets.cpp @@ -5,6 +5,7 @@ #include "engine/framework/io/json.h" #include +#include #include #include #include @@ -80,6 +81,33 @@ MossLocalTransformerConfig parse_local_config(const engine::io::json::Value & va return config; } +std::filesystem::path huggingface_hub_root() { + if (const char * hf_home = std::getenv("HF_HOME")) { + return std::filesystem::path(hf_home) / "hub"; + } + const char * home = std::getenv("USERPROFILE"); + if (home == nullptr) { + home = std::getenv("HOME"); + } + if (home == nullptr) { + throw std::runtime_error("MOSS-TTS-Local cannot locate the Hugging Face cache: no HF_HOME/USERPROFILE/HOME"); + } + return std::filesystem::path(home) / ".cache" / "huggingface" / "hub"; +} + +std::string repo_id_to_cache_folder(const std::string & repo_id) { + // "Org/Name" -> "models--Org--Name" (Hugging Face hub cache naming). + std::string folder = "models--"; + for (const char ch : repo_id) { + if (ch == '/') { + folder += "--"; + } else { + folder += ch; + } + } + return folder; +} + } // namespace MossTTSLocalAssetPaths resolve_moss_tts_local_assets(const std::filesystem::path & model_path) { @@ -150,4 +178,28 @@ std::shared_ptr load_moss_tts_local_assets(const std:: return std::make_shared(std::move(assets)); } +std::filesystem::path resolve_moss_codec_dir(const MossTTSLocalAssets & assets) { + if (assets.paths.audio_tokenizer_root.has_value() && + engine::io::is_existing_file(*assets.paths.audio_tokenizer_root / "config.json")) { + return *assets.paths.audio_tokenizer_root; + } + const std::string & repo_id = assets.config.audio_tokenizer_name_or_path; + if (repo_id.empty()) { + throw std::runtime_error( + "MOSS-TTS-Local audio tokenizer not found: no local audio_tokenizer/ and no repo id in config"); + } + const auto snapshots = huggingface_hub_root() / repo_id_to_cache_folder(repo_id) / "snapshots"; + if (!engine::io::is_existing_directory(snapshots)) { + throw std::runtime_error( + "MOSS-TTS-Local codec not found in Hugging Face cache: " + snapshots.string() + + " (download " + repo_id + " or place it under the model's audio_tokenizer/ directory)"); + } + for (const auto & entry : std::filesystem::directory_iterator(snapshots)) { + if (entry.is_directory() && engine::io::is_existing_file(entry.path() / "config.json")) { + return entry.path(); + } + } + throw std::runtime_error("MOSS-TTS-Local codec snapshot with config.json not found under " + snapshots.string()); +} + } // namespace engine::models::moss_tts_local diff --git a/src/models/moss_tts_local/codec_decoder.cpp b/src/models/moss_tts_local/codec_decoder.cpp index 117b241..939e7f4 100644 --- a/src/models/moss_tts_local/codec_decoder.cpp +++ b/src/models/moss_tts_local/codec_decoder.cpp @@ -1,58 +1,28 @@ #include "engine/models/moss_tts_local/codec_decoder.h" -#include "engine/framework/assets/tensor_source.h" -#include "engine/framework/core/backend_weight_store.h" #include "engine/framework/core/module.h" -#include "engine/framework/modules/activation_modules.h" -#include "engine/framework/modules/linear_module.h" -#include "engine/framework/modules/norm_modules.h" -#include "engine/framework/modules/positional_modules.h" -#include "engine/framework/modules/primitive_modules.h" -#include "engine/framework/modules/structural_modules.h" -#include "engine/framework/modules/weight_binding.h" #include "engine/models/moss_tts_local/codec_quantizer.h" +#include "engine/models/moss_tts_local/codec_transformer.h" #include #include -#include -#include #include #include -#include #include #include namespace engine::models::moss_tts_local { namespace { -namespace modules = engine::modules; -namespace binding = engine::modules::binding; +namespace cd = codec_detail; -constexpr float kMaskedAttentionBias = std::numeric_limits::lowest(); -constexpr int64_t kCodeDim = 768; -constexpr int64_t kSamplesPerFrame = 3840; // downsample_rate (per interleaved stream frame) -constexpr float kRopeTheta = 10000.0F; -constexpr float kLayerNormEps = 1.0e-5F; - -// One decoder sub-transformer (ProjectedTransformer) description, mirroring the -// v2 codec config.json decoder_kwargs (indices decoder.0/2/4/6/8/10). context -// is the local-attention window in tokens at that stage's frame rate, computed -// as round(frame_rate * context_duration); the patch that follows doubles the -// frame rate (the last one expands by 240). -struct TransformerSpec { - int64_t input_dim; - int64_t output_dim; - int64_t d_model; - int64_t num_heads; - int64_t num_layers; - int64_t intermediate_size; - int64_t context; - int64_t patch_after; -}; - -constexpr TransformerSpec kDecoderSpecs[] = { - {kCodeDim, 1280, 1280, 20, 32, 5120, 125, 2}, +// Decoder ProjectedTransformer stages, mirroring the v2 codec config.json +// decoder_kwargs (modules decoder.0/2/4/6/8/10). context is the local-attention +// window at that stage's frame rate; the patch that follows expands the frame +// rate (the last one by 240). +constexpr cd::TransformerSpec kDecoderSpecs[] = { + {cd::kCodeDim, 1280, 1280, 20, 32, 5120, 125, 2}, {640, 768, 768, 12, 12, 3072, 250, 2}, {384, 768, 768, 12, 12, 3072, 400, 2}, {384, 768, 768, 12, 12, 3072, 400, 2}, @@ -62,232 +32,6 @@ constexpr TransformerSpec kDecoderSpecs[] = { constexpr size_t kNumTransformers = sizeof(kDecoderSpecs) / sizeof(kDecoderSpecs[0]); -struct GgmlContextDeleter { - void operator()(ggml_context * ctx) const noexcept { - if (ctx != nullptr) { - ggml_free(ctx); - } - } -}; - -struct LayerWeights { - core::TensorValue norm1_w; - core::TensorValue norm1_b; - core::TensorValue in_proj; // fused qkv [3 * d_model, d_model] - core::TensorValue out_proj; // [d_model, d_model] - core::TensorValue norm2_w; - core::TensorValue norm2_b; - core::TensorValue fc1; // [intermediate_size, d_model] - core::TensorValue fc2; // [d_model, intermediate_size] - core::TensorValue layer_scale1; // [d_model] - core::TensorValue layer_scale2; // [d_model] -}; - -struct TransformerWeights { - TransformerSpec spec; - core::TensorValue input_proj; // [d_model, input_dim] - core::TensorValue output_proj; // [output_dim, d_model] - std::vector layers; -}; - -// Opens the codec safetensors shards and resolves a tensor to the shard that -// holds it (the codec ships model-0000N-of-00003.safetensors + an index). -class CodecShards { -public: - explicit CodecShards(const std::filesystem::path & codec_dir) { - for (const auto & entry : std::filesystem::directory_iterator(codec_dir)) { - if (entry.is_regular_file() && entry.path().extension() == ".safetensors") { - sources_.push_back(assets::open_tensor_source(entry.path())); - } - } - if (sources_.empty()) { - throw std::runtime_error("MOSS codec has no safetensors shards: " + codec_dir.string()); - } - } - - const assets::TensorSource & source_for(const std::string & name) const { - for (const auto & source : sources_) { - if (source->has_tensor(name)) { - return *source; - } - } - throw std::runtime_error("MOSS codec tensor not found: " + name); - } - -private: - std::vector> sources_; -}; - -TransformerWeights load_transformer( - core::BackendWeightStore & store, - const CodecShards & shards, - const TransformerSpec & spec, - int64_t decoder_index) { - const std::string prefix = "decoder." + std::to_string(decoder_index); - const auto load = [&](const std::string & name, std::initializer_list shape) { - return store.load_tensor(shards.source_for(name), name, assets::TensorStorageType::F32, shape); - }; - const auto load_f32 = [&](const std::string & name, std::initializer_list shape) { - return store.load_f32_tensor(shards.source_for(name), name, shape); - }; - - TransformerWeights weights; - weights.spec = spec; - weights.input_proj = load(prefix + ".input_proj.weight", {spec.d_model, spec.input_dim}); - weights.output_proj = load(prefix + ".output_proj.weight", {spec.output_dim, spec.d_model}); - weights.layers.reserve(static_cast(spec.num_layers)); - for (int64_t layer = 0; layer < spec.num_layers; ++layer) { - const std::string lp = prefix + ".transformer.layers." + std::to_string(layer); - LayerWeights w; - w.norm1_w = load_f32(lp + ".norm1.weight", {spec.d_model}); - w.norm1_b = load_f32(lp + ".norm1.bias", {spec.d_model}); - w.in_proj = load(lp + ".self_attn.in_proj.weight", {3 * spec.d_model, spec.d_model}); - w.out_proj = load(lp + ".self_attn.out_proj.weight", {spec.d_model, spec.d_model}); - w.norm2_w = load_f32(lp + ".norm2.weight", {spec.d_model}); - w.norm2_b = load_f32(lp + ".norm2.bias", {spec.d_model}); - w.fc1 = load(lp + ".ffn.0.weight", {spec.intermediate_size, spec.d_model}); - w.fc2 = load(lp + ".ffn.2.weight", {spec.d_model, spec.intermediate_size}); - w.layer_scale1 = load_f32(lp + ".layer_scale_1.scale", {spec.d_model}); - w.layer_scale2 = load_f32(lp + ".layer_scale_2.scale", {spec.d_model}); - weights.layers.push_back(std::move(w)); - } - return weights; -} - -core::TensorValue ensure_contiguous(core::ModuleBuildContext & ctx, const core::TensorValue & value) { - return core::ensure_backend_addressable_layout(ctx, value); -} - -core::TensorValue reshape_heads( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - int64_t heads, - int64_t dim) { - auto contiguous = ensure_contiguous(ctx, input); - return core::reshape_tensor( - ctx, - contiguous, - core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim})); -} - -// Slices a fused qkv projection [1, T, 3*d_model] into its q/k/v thirds, each -// [1, T, d_model]. The 3*d_model axis is laid out as [3, heads, head_dim]. -core::TensorValue slice_projection( - core::ModuleBuildContext & ctx, - const core::TensorValue & qkv, - int64_t which, - int64_t d_model, - int64_t steps) { - ggml_tensor * source = qkv.tensor; - const size_t element_size = ggml_element_size(source); - ggml_tensor * view = ggml_view_3d( - ctx.ggml, - source, - d_model, - steps, - 1, - source->nb[1], - source->nb[2], - static_cast(which * d_model) * element_size); - return core::wrap_tensor( - ggml_cont(ctx.ggml, view), - core::TensorShape::from_dims({1, steps, d_model}), - GGML_TYPE_F32); -} - -core::TensorValue attention( - core::ModuleBuildContext & ctx, - const core::TensorValue & q_heads, - const core::TensorValue & k_heads, - const core::TensorValue & v_heads, - int64_t dim, - const core::TensorValue & mask) { - const modules::MatMulModule matmul; - auto scores = matmul.build( - ctx, - q_heads, - modules::TransposeModule({{0, 1, 3, 2}, k_heads.shape.rank}).build(ctx, k_heads)); - scores = core::ensure_backend_addressable_layout(ctx, scores); - auto attn = core::wrap_tensor( - ggml_soft_max_ext( - ctx.ggml, - scores.tensor, - mask.tensor, - 1.0F / std::sqrt(static_cast(dim)), - 0.0F), - scores.shape, - GGML_TYPE_F32); - return matmul.build(ctx, attn, v_heads); -} - -core::TensorValue transformer_layer( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const LayerWeights & weights, - const TransformerSpec & spec, - const core::TensorValue & positions, - const core::TensorValue & mask, - int64_t steps) { - const int64_t dim = spec.d_model / spec.num_heads; - const modules::LayerNormModule norm({spec.d_model, kLayerNormEps, true, true}); - - auto normed = norm.build(ctx, input, binding::norm_data(ctx, weights.norm1_w, weights.norm1_b)); - auto qkv = modules::LinearModule(binding::linear_config(spec.d_model, 3 * spec.d_model, false)) - .build(ctx, normed, binding::linear_data(ctx, weights.in_proj)); - - auto q = slice_projection(ctx, qkv, 0, spec.d_model, steps); - auto k = slice_projection(ctx, qkv, 1, spec.d_model, steps); - auto v = slice_projection(ctx, qkv, 2, spec.d_model, steps); - - q = reshape_heads(ctx, q, spec.num_heads, dim); - k = reshape_heads(ctx, k, spec.num_heads, dim); - v = reshape_heads(ctx, v, spec.num_heads, dim); - q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NORMAL, kRopeTheta}).build(ctx, q, positions); - k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NORMAL, kRopeTheta}).build(ctx, k, positions); - - auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); - auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k); - auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v); - auto context = attention(ctx, q_heads, k_heads, v_heads, dim, mask); - context = modules::TransposeModule({{0, 2, 1, 3}, context.shape.rank}).build(ctx, context); - context = ensure_contiguous(ctx, context); - context = core::reshape_tensor(ctx, context, core::TensorShape::from_dims({1, steps, spec.d_model})); - auto attn_out = modules::LinearModule(binding::linear_config(spec.d_model, spec.d_model, false)) - .build(ctx, context, binding::linear_data(ctx, weights.out_proj)); - attn_out = core::wrap_tensor( - ggml_mul(ctx.ggml, attn_out.tensor, weights.layer_scale1.tensor), attn_out.shape, GGML_TYPE_F32); - auto x = modules::AddModule{}.build(ctx, input, attn_out); - - auto ff_in = norm.build(ctx, x, binding::norm_data(ctx, weights.norm2_w, weights.norm2_b)); - auto ff = modules::LinearModule(binding::linear_config(spec.d_model, spec.intermediate_size, false)) - .build(ctx, ff_in, binding::linear_data(ctx, weights.fc1)); - ff = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, ff); - ff = modules::LinearModule(binding::linear_config(spec.intermediate_size, spec.d_model, false)) - .build(ctx, ff, binding::linear_data(ctx, weights.fc2)); - ff = core::wrap_tensor( - ggml_mul(ctx.ggml, ff.tensor, weights.layer_scale2.tensor), ff.shape, GGML_TYPE_F32); - return modules::AddModule{}.build(ctx, x, ff); -} - -// ProjectedTransformer: input projection -> transformer stack -> output -// projection. Input/output are [1, steps, channels] (feature-last). -core::TensorValue run_transformer( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const TransformerWeights & weights, - const core::TensorValue & positions, - const core::TensorValue & mask, - int64_t steps) { - const auto & spec = weights.spec; - auto x = modules::LinearModule(binding::linear_config(spec.input_dim, spec.d_model, false)) - .build(ctx, input, binding::linear_data(ctx, weights.input_proj)); - for (const auto & layer : weights.layers) { - x = transformer_layer(ctx, x, layer, spec, positions, mask, steps); - } - return modules::LinearModule(binding::linear_config(spec.d_model, spec.output_dim, false)) - .build(ctx, x, binding::linear_data(ctx, weights.output_proj)); -} - // PatchedPretransform (decode/upsample): [1, l, d*patch] -> [1, l*patch, d]. // Each frame is unpacked into `patch` consecutive frames along time, matching // x.reshape(b, d, h, l).permute(0, 1, 3, 2).reshape(b, d, l * h). @@ -295,7 +39,7 @@ core::TensorValue patch_upsample( core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t patch) { - auto contiguous = ensure_contiguous(ctx, input); + auto contiguous = cd::ensure_contiguous(ctx, input); const int64_t length = contiguous.shape.dims[1]; const int64_t packed = contiguous.shape.dims[2]; const int64_t channels = packed / patch; @@ -306,18 +50,6 @@ core::TensorValue patch_upsample( unpacked, core::TensorShape::from_dims({1, length * patch, channels}), GGML_TYPE_F32); } -std::vector causal_context_mask(int64_t steps, int64_t context) { - std::vector mask(static_cast(steps * steps), kMaskedAttentionBias); - for (int64_t query = 0; query < steps; ++query) { - for (int64_t key = 0; key <= query; ++key) { - if (query - key < context) { - mask[static_cast(query * steps + key)] = 0.0F; - } - } - } - return mask; -} - } // namespace struct MossCodecDecoder::Impl { @@ -326,7 +58,7 @@ struct MossCodecDecoder::Impl { size_t graph_arena_bytes = 0; std::unique_ptr dequantizer; std::unique_ptr store; - std::vector transformers; + std::vector transformers; }; MossCodecDecoder::MossCodecDecoder( @@ -344,13 +76,13 @@ MossCodecDecoder::MossCodecDecoder( impl_->graph_arena_bytes = graph_arena_bytes; impl_->dequantizer = std::make_unique(codec_dir, num_quantizers); - CodecShards shards(codec_dir); + cd::CodecShards shards(codec_dir); impl_->store = std::make_unique( impl_->backend, impl_->backend_type, "moss_tts_local.codec.decoder", weight_context_bytes); impl_->transformers.reserve(kNumTransformers); for (size_t index = 0; index < kNumTransformers; ++index) { - impl_->transformers.push_back( - load_transformer(*impl_->store, shards, kDecoderSpecs[index], static_cast(2 * index))); + impl_->transformers.push_back(cd::load_transformer( + *impl_->store, shards, kDecoderSpecs[index], "decoder", static_cast(2 * index))); } impl_->store->upload(); } @@ -371,23 +103,23 @@ std::vector> MossCodecDecoder::decode( // Codes -> continuous latent [code_dim, frames] (channel-major), transposed // into the feature-last [1, frames, code_dim] layout the decoder expects. const auto latent = impl_->dequantizer->decode(codes); - std::vector latent_input(static_cast(frames * kCodeDim)); - for (int64_t channel = 0; channel < kCodeDim; ++channel) { + std::vector latent_input(static_cast(frames * cd::kCodeDim)); + for (int64_t channel = 0; channel < cd::kCodeDim; ++channel) { for (int64_t step = 0; step < frames; ++step) { - latent_input[static_cast(step * kCodeDim + channel)] = + latent_input[static_cast(step * cd::kCodeDim + channel)] = latent[static_cast(channel * frames + step)]; } } ggml_init_params params{impl_->graph_arena_bytes, nullptr, true}; - std::unique_ptr graph_ctx(ggml_init(params)); + std::unique_ptr graph_ctx(ggml_init(params)); if (graph_ctx == nullptr) { throw std::runtime_error("failed to initialize MOSS codec decoder graph context"); } core::ModuleBuildContext ctx{graph_ctx.get(), "moss_tts_local.codec.decode", impl_->backend_type}; auto latent_tensor = - core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, frames, kCodeDim})); + core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, frames, cd::kCodeDim})); ggml_set_input(latent_tensor.tensor); struct StageInput { @@ -415,15 +147,15 @@ std::vector> MossCodecDecoder::decode( stage.position_host[static_cast(i)] = static_cast(i); } stage.mask = mask.tensor; - stage.mask_host = causal_context_mask(steps, transformer.spec.context); + stage.mask_host = cd::causal_context_mask(steps, transformer.spec.context); stage_inputs.push_back(std::move(stage)); - hidden = run_transformer(ctx, hidden, transformer, positions, mask, steps); - hidden = patch_upsample(ctx, hidden, transformer.spec.patch_after); - steps *= transformer.spec.patch_after; + hidden = cd::run_transformer(ctx, hidden, transformer, positions, mask, steps); + hidden = patch_upsample(ctx, hidden, transformer.spec.patch); + steps *= transformer.spec.patch; } - hidden = ensure_contiguous(ctx, hidden); + hidden = cd::ensure_contiguous(ctx, hidden); ggml_set_output(hidden.tensor); ggml_cgraph * graph = ggml_new_graph_custom(graph_ctx.get(), 131072, false); @@ -460,7 +192,7 @@ std::vector> MossCodecDecoder::decode( // De-interleave the jointly-processed stream back into left/right channels // (channel 0 = even samples, channel 1 = odd samples). - const int64_t per_channel = frames * kSamplesPerFrame; + const int64_t per_channel = frames * cd::kSamplesPerFrame; std::vector> stereo(2, std::vector(static_cast(per_channel))); for (int64_t i = 0; i < per_channel; ++i) { stereo[0][static_cast(i)] = flat[static_cast(2 * i)]; diff --git a/src/models/moss_tts_local/codec_encoder.cpp b/src/models/moss_tts_local/codec_encoder.cpp new file mode 100644 index 0000000..bb13c0b --- /dev/null +++ b/src/models/moss_tts_local/codec_encoder.cpp @@ -0,0 +1,208 @@ +#include "engine/models/moss_tts_local/codec_encoder.h" + +#include "engine/framework/core/module.h" +#include "engine/models/moss_tts_local/codec_quantizer.h" +#include "engine/models/moss_tts_local/codec_transformer.h" + +#include +#include + +#include +#include +#include +#include + +namespace engine::models::moss_tts_local { +namespace { + +namespace cd = codec_detail; + +// Encoder ProjectedTransformer stages, mirroring the v2 codec config.json +// encoder_kwargs (modules encoder.1/3/5/7/9/11). Each transformer is preceded by +// a PatchedPretransform downsample of factor `patch` (module 2*i); the first one +// (patch 240) turns the interleaved waveform into frames, the rest halve the +// frame count. context is the local-attention window at that stage's frame rate +// (mirror of the decoder's contexts, reversed). +constexpr cd::TransformerSpec kEncoderSpecs[] = { + {240, 384, 768, 12, 12, 3072, 400, 240}, + {768, 384, 768, 12, 12, 3072, 400, 2}, + {768, 384, 768, 12, 12, 3072, 400, 2}, + {768, 384, 768, 12, 12, 3072, 400, 2}, + {768, 640, 768, 12, 12, 3072, 250, 2}, + {1280, 768, 1280, 20, 32, 5120, 125, 2}, +}; + +constexpr size_t kNumTransformers = sizeof(kEncoderSpecs) / sizeof(kEncoderSpecs[0]); + +// PatchedPretransform (encode/downsample): [1, l, d] -> [1, l/patch, d*patch]. +// Packs `patch` consecutive frames into the feature dim, matching +// x.reshape(b, d, -1, h).permute(0, 1, 3, 2).reshape(b, d * h, -1) (conv layout), +// i.e. output feature (d_idx*patch + h_idx) at time lt = input feature d_idx at +// time lt*patch + h_idx. This is the exact inverse of the decoder's upsample. +core::TensorValue patch_downsample( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t patch) { + auto contiguous = cd::ensure_contiguous(ctx, input); + const int64_t total_length = contiguous.shape.dims[1]; + const int64_t channels = contiguous.shape.dims[2]; + const int64_t length = total_length / patch; + ggml_tensor * reshaped = ggml_reshape_3d(ctx.ggml, contiguous.tensor, channels, patch, length); + ggml_tensor * permuted = ggml_cont(ctx.ggml, ggml_permute(ctx.ggml, reshaped, 1, 0, 2, 3)); + ggml_tensor * packed = ggml_reshape_2d(ctx.ggml, permuted, channels * patch, length); + return core::wrap_tensor( + packed, core::TensorShape::from_dims({1, length, channels * patch}), GGML_TYPE_F32); +} + +} // namespace + +struct MossCodecEncoder::Impl { + ggml_backend_t backend = nullptr; + core::BackendType backend_type = core::BackendType::Cpu; + size_t graph_arena_bytes = 0; + std::unique_ptr quantizer; + std::unique_ptr store; + std::vector transformers; +}; + +MossCodecEncoder::MossCodecEncoder( + const std::filesystem::path & codec_dir, + core::ExecutionContext & execution_context, + int64_t num_quantizers, + size_t weight_context_bytes, + size_t graph_arena_bytes) + : impl_(std::make_unique()) { + impl_->backend = execution_context.backend(); + if (impl_->backend == nullptr) { + throw std::runtime_error("MOSS codec encoder backend is not initialized"); + } + impl_->backend_type = execution_context.backend_type(); + impl_->graph_arena_bytes = graph_arena_bytes; + impl_->quantizer = std::make_unique(codec_dir, num_quantizers); + + cd::CodecShards shards(codec_dir); + impl_->store = std::make_unique( + impl_->backend, impl_->backend_type, "moss_tts_local.codec.encoder", weight_context_bytes); + impl_->transformers.reserve(kNumTransformers); + for (size_t index = 0; index < kNumTransformers; ++index) { + impl_->transformers.push_back(cd::load_transformer( + *impl_->store, shards, kEncoderSpecs[index], "encoder", static_cast(2 * index + 1))); + } + impl_->store->upload(); +} + +MossCodecEncoder::~MossCodecEncoder() = default; + +std::vector> MossCodecEncoder::encode( + const std::vector> & channels) const { + if (channels.size() != 2) { + throw std::runtime_error("MOSS codec encoder requires stereo (2-channel) input"); + } + if (channels[0].size() != channels[1].size()) { + throw std::runtime_error("MOSS codec encoder channels must have equal length"); + } + const int64_t raw_per_channel = static_cast(channels[0].size()); + if (raw_per_channel <= 0) { + throw std::runtime_error("MOSS codec encoder requires a non-empty waveform"); + } + + // Pad each channel up to a multiple of the downsample rate, then interleave + // the two channels into one stream (even = left, odd = right), matching the + // codec's _flatten_channels_for_codec. + const int64_t frames = (raw_per_channel + cd::kSamplesPerFrame - 1) / cd::kSamplesPerFrame; + const int64_t per_channel = frames * cd::kSamplesPerFrame; + const int64_t interleaved = per_channel * 2; + std::vector waveform(static_cast(interleaved), 0.0F); + for (int64_t i = 0; i < raw_per_channel; ++i) { + waveform[static_cast(2 * i)] = channels[0][static_cast(i)]; + waveform[static_cast(2 * i + 1)] = channels[1][static_cast(i)]; + } + + ggml_init_params params{impl_->graph_arena_bytes, nullptr, true}; + std::unique_ptr graph_ctx(ggml_init(params)); + if (graph_ctx == nullptr) { + throw std::runtime_error("failed to initialize MOSS codec encoder graph context"); + } + core::ModuleBuildContext ctx{graph_ctx.get(), "moss_tts_local.codec.encode", impl_->backend_type}; + + // Input is the interleaved waveform as a single-feature stream [1, L, 1]. + auto input_tensor = + core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, interleaved, 1})); + ggml_set_input(input_tensor.tensor); + + struct StageInput { + ggml_tensor * positions; + std::vector position_host; + ggml_tensor * mask; + std::vector mask_host; + }; + std::vector stage_inputs; + stage_inputs.reserve(impl_->transformers.size()); + + auto hidden = input_tensor; + int64_t steps = interleaved; + for (const auto & transformer : impl_->transformers) { + // Downsample first (encoder order), then run the transformer at the + // reduced frame count. + hidden = patch_downsample(ctx, hidden, transformer.spec.patch); + steps /= transformer.spec.patch; + + auto positions = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({steps})); + ggml_set_input(positions.tensor); + auto mask = + core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, steps, steps})); + ggml_set_input(mask.tensor); + + StageInput stage; + stage.positions = positions.tensor; + stage.position_host.resize(static_cast(steps)); + for (int64_t i = 0; i < steps; ++i) { + stage.position_host[static_cast(i)] = static_cast(i); + } + stage.mask = mask.tensor; + stage.mask_host = cd::causal_context_mask(steps, transformer.spec.context); + stage_inputs.push_back(std::move(stage)); + + hidden = cd::run_transformer(ctx, hidden, transformer, positions, mask, steps); + } + + hidden = cd::ensure_contiguous(ctx, hidden); + ggml_set_output(hidden.tensor); + + ggml_cgraph * graph = ggml_new_graph_custom(graph_ctx.get(), 131072, false); + ggml_build_forward_expand(graph, hidden.tensor); + + ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(impl_->backend)); + if (gallocr == nullptr || !ggml_gallocr_reserve(gallocr, graph) || !ggml_gallocr_alloc_graph(gallocr, graph)) { + if (gallocr != nullptr) { + ggml_gallocr_free(gallocr); + } + throw std::runtime_error("failed to allocate MOSS codec encoder forward graph"); + } + + ggml_backend_tensor_set(input_tensor.tensor, waveform.data(), 0, waveform.size() * sizeof(float)); + for (const auto & stage : stage_inputs) { + ggml_backend_tensor_set( + stage.positions, stage.position_host.data(), 0, stage.position_host.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + stage.mask, stage.mask_host.data(), 0, stage.mask_host.size() * sizeof(float)); + } + + const ggml_status status = ggml_backend_graph_compute(impl_->backend, graph); + ggml_backend_synchronize(impl_->backend); + if (status != GGML_STATUS_SUCCESS) { + ggml_gallocr_free(gallocr); + throw std::runtime_error("MOSS codec encoder forward graph compute failed"); + } + + // hidden is [1, frames, code_dim] feature-last; ggml memory order is + // channel-fastest, i.e. flat[frame * code_dim + channel] -- exactly the + // layout MossCodecDequantizer::encode expects. + std::vector latent(static_cast(steps * cd::kCodeDim)); + ggml_backend_tensor_get(hidden.tensor, latent.data(), 0, latent.size() * sizeof(float)); + ggml_gallocr_free(gallocr); + + return impl_->quantizer->encode(latent, steps); +} + +} // namespace engine::models::moss_tts_local diff --git a/src/models/moss_tts_local/codec_quantizer.cpp b/src/models/moss_tts_local/codec_quantizer.cpp index ecc6b56..513c829 100644 --- a/src/models/moss_tts_local/codec_quantizer.cpp +++ b/src/models/moss_tts_local/codec_quantizer.cpp @@ -4,9 +4,11 @@ #include #include +#include #include #include #include +#include namespace engine::models::moss_tts_local { namespace { @@ -91,11 +93,29 @@ MossCodecDequantizer::MossCodecDequantizer(const std::filesystem::path & codec_d codebook.table = shards.require_f32(prefix + ".codebook.weight"); codebook.out_weight = load_wn_conv_weight(shards, prefix + ".out_proj", rvq_dim_, codebook_dim_); codebook.out_bias = shards.require_f32(prefix + ".out_proj.bias"); + codebook.in_weight = load_wn_conv_weight(shards, prefix + ".in_proj", codebook_dim_, rvq_dim_); + codebook.in_bias = shards.require_f32(prefix + ".in_proj.bias"); + // Pre-normalize the codebook rows once (encode does L2-normalized nearest + // search, matching the training LFQ; F.normalize uses eps=1e-12). + codebook.table_normalized = codebook.table; + for (int64_t code = 0; code < codebook_size_; ++code) { + float * row = &codebook.table_normalized[static_cast(code * codebook_dim_)]; + double norm = 0.0; + for (int64_t k = 0; k < codebook_dim_; ++k) { + norm += static_cast(row[k]) * static_cast(row[k]); + } + const double scale = 1.0 / std::max(std::sqrt(norm), 1.0e-12); + for (int64_t k = 0; k < codebook_dim_; ++k) { + row[k] = static_cast(row[k] * scale); + } + } codebooks_.push_back(std::move(codebook)); } output_weight_ = load_wn_conv_weight(shards, "quantizer.output_proj", code_dim_, rvq_dim_); output_bias_ = shards.require_f32("quantizer.output_proj.bias"); + input_weight_ = load_wn_conv_weight(shards, "quantizer.input_proj", rvq_dim_, code_dim_); + input_bias_ = shards.require_f32("quantizer.input_proj.bias"); } std::vector MossCodecDequantizer::decode(const std::vector> & codes) const { @@ -139,4 +159,81 @@ std::vector MossCodecDequantizer::decode(const std::vector> MossCodecDequantizer::encode( + const std::vector & hidden, int64_t frames) const { + if (frames <= 0) { + throw std::runtime_error("MOSS codec quantizer requires a non-empty encoder latent"); + } + if (static_cast(hidden.size()) != frames * code_dim_) { + throw std::runtime_error("MOSS codec quantizer got a mis-shaped encoder latent"); + } + + std::vector> codes( + static_cast(num_quantizers_), std::vector(static_cast(frames))); + + std::vector residual(static_cast(rvq_dim_)); + std::vector encoding(static_cast(codebook_dim_)); + for (int64_t step = 0; step < frames; ++step) { + // input_proj: encoder latent [code_dim] -> rvq_dim (WNConv1d 1x1). + const float * frame_hidden = &hidden[static_cast(step * code_dim_)]; + for (int64_t out = 0; out < rvq_dim_; ++out) { + double sum = input_bias_[static_cast(out)]; + const float * row = &input_weight_[static_cast(out * code_dim_)]; + for (int64_t k = 0; k < code_dim_; ++k) { + sum += static_cast(row[k]) * static_cast(frame_hidden[k]); + } + residual[static_cast(out)] = sum; + } + + for (int64_t index = 0; index < num_quantizers_; ++index) { + const auto & codebook = codebooks_[static_cast(index)]; + + // in_proj: residual [rvq_dim] -> codebook_dim, then L2-normalize. + double enc_norm = 0.0; + for (int64_t c = 0; c < codebook_dim_; ++c) { + double sum = codebook.in_bias[static_cast(c)]; + const float * row = &codebook.in_weight[static_cast(c * rvq_dim_)]; + for (int64_t k = 0; k < rvq_dim_; ++k) { + sum += static_cast(row[k]) * residual[static_cast(k)]; + } + encoding[static_cast(c)] = sum; + enc_norm += sum * sum; + } + const double enc_scale = 1.0 / std::max(std::sqrt(enc_norm), 1.0e-12); + for (int64_t c = 0; c < codebook_dim_; ++c) { + encoding[static_cast(c)] *= enc_scale; + } + + // Nearest code by cosine similarity (both sides L2-normalized), i.e. + // argmax dot == argmin squared distance on the unit sphere. + int32_t best_code = 0; + double best_dot = -std::numeric_limits::infinity(); + for (int64_t code = 0; code < codebook_size_; ++code) { + const float * row = &codebook.table_normalized[static_cast(code * codebook_dim_)]; + double dot = 0.0; + for (int64_t c = 0; c < codebook_dim_; ++c) { + dot += static_cast(row[c]) * encoding[static_cast(c)]; + } + if (dot > best_dot) { + best_dot = dot; + best_code = static_cast(code); + } + } + codes[static_cast(index)][static_cast(step)] = best_code; + + // Subtract the residual contribution: out_proj(raw codebook row). + const float * embedding = &codebook.table[static_cast(best_code * codebook_dim_)]; + for (int64_t out = 0; out < rvq_dim_; ++out) { + double sum = codebook.out_bias[static_cast(out)]; + const float * row = &codebook.out_weight[static_cast(out * codebook_dim_)]; + for (int64_t k = 0; k < codebook_dim_; ++k) { + sum += static_cast(row[k]) * static_cast(embedding[k]); + } + residual[static_cast(out)] -= sum; + } + } + } + return codes; +} + } // namespace engine::models::moss_tts_local diff --git a/src/models/moss_tts_local/loader.cpp b/src/models/moss_tts_local/loader.cpp index 0d668c3..df9376e 100644 --- a/src/models/moss_tts_local/loader.cpp +++ b/src/models/moss_tts_local/loader.cpp @@ -2,10 +2,13 @@ #include "engine/framework/io/filesystem.h" #include "engine/models/moss_tts_local/assets.h" +#include "engine/models/moss_tts_local/session.h" #include +#include #include #include +#include #include namespace engine::models::moss_tts_local { @@ -69,6 +72,42 @@ std::vector discover_weight_assets(const std::filesystem::p return runtime::discover_named_assets(root, {"model.safetensors"}); } +class MossTTSLocalLoadedModel final : public runtime::ILoadedVoiceModel { +public: + MossTTSLocalLoadedModel( + runtime::ModelMetadata metadata, + runtime::CapabilitySet capabilities, + std::shared_ptr assets) + : metadata_(std::move(metadata)), + capabilities_(std::move(capabilities)), + assets_(std::move(assets)) {} + + const runtime::ModelMetadata & metadata() const noexcept override { + return metadata_; + } + + const runtime::CapabilitySet & capabilities() const noexcept override { + return capabilities_; + } + + std::unique_ptr create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const override { + if (task.mode != runtime::RunMode::Offline) { + throw std::runtime_error("MOSS-TTS-Local only supports offline sessions"); + } + if (task.task != runtime::VoiceTaskKind::Tts) { + throw std::runtime_error("MOSS-TTS-Local only supports the Tts task"); + } + return std::make_unique(task, options, assets_); + } + +private: + runtime::ModelMetadata metadata_; + runtime::CapabilitySet capabilities_; + std::shared_ptr assets_; +}; + class MossTTSLocalLoader final : public runtime::IVoiceModelLoader { public: std::string family() const override { @@ -104,8 +143,9 @@ class MossTTSLocalLoader final : public runtime::IVoiceModelLoader { } std::unique_ptr load(const runtime::ModelLoadRequest & request) const override { - (void) request; - throw std::runtime_error("MOSS-TTS-Local generation is not implemented yet"); + const auto root = resolve_model_root(request.model_path); + return std::make_unique( + moss_tts_local_metadata(), moss_tts_local_capabilities(), load_moss_tts_local_assets(root)); } }; diff --git a/src/models/moss_tts_local/processor.cpp b/src/models/moss_tts_local/processor.cpp index 1143ae8..3b03d38 100644 --- a/src/models/moss_tts_local/processor.cpp +++ b/src/models/moss_tts_local/processor.cpp @@ -2,6 +2,7 @@ #include "engine/framework/tokenizers/llama_bpe.h" +#include #include #include @@ -55,13 +56,65 @@ std::filesystem::path require_path( } // namespace +// A decoder row carries a text-channel id plus n_vq audio-channel codes. Text rows pad the +// audio channels; audio rows carry the reference codes with the given text-channel slot. +struct RowBuilder { + std::vector text_tokens; + std::vector audio_codes; + int64_t n_vq; + int32_t audio_pad; + + void push_text_token(int32_t id) { + text_tokens.push_back(id); + audio_codes.insert(audio_codes.end(), static_cast(n_vq), audio_pad); + } + void push_text_ids(const std::vector & ids) { + for (const int32_t id : ids) { + push_text_token(id); + } + } +}; + struct MossTextProcessor::Impl { std::shared_ptr assets; std::shared_ptr tokenizer; - void append_encoded(std::vector & tokens, const std::string & text) const { - const auto ids = tokenizer->encode(text); - tokens.insert(tokens.end(), ids.begin(), ids.end()); + RowBuilder make_row_builder() const { + RowBuilder builder; + builder.n_vq = assets->config.num_codebooks; + builder.audio_pad = static_cast(assets->config.audio_pad_token_id); + return builder; + } + + void push_text(RowBuilder & builder, const std::string & text) const { + builder.push_text_ids(tokenizer->encode(text)); + } + + // Emits the shared scaffold. `emit_reference` fills the "- Reference(s):" + // slot: text-only passes "None"; the clone path emits the reference audio rows. + MossGenerationPrefix build_prefix( + const std::string & text, + const std::optional & language, + const std::function & emit_reference) const { + const auto & config = assets->config; + RowBuilder builder = make_row_builder(); + builder.push_text_token(static_cast(config.im_start_token_id)); + push_text(builder, kUserRolePrefix); + push_text(builder, kUserReferencePrefix); + emit_reference(builder); + push_text(builder, render_after_reference(language)); + push_text(builder, text); + push_text(builder, kUserInstSuffix); + builder.push_text_token(static_cast(config.im_end_token_id)); + push_text(builder, kAssistantTurnPrefix); + builder.push_text_token(static_cast(config.im_start_token_id)); + push_text(builder, kAssistantRolePrefix); + builder.push_text_token(static_cast(config.audio_start_token_id)); + + MossGenerationPrefix prefix; + prefix.text_tokens = std::move(builder.text_tokens); + prefix.audio_codes = std::move(builder.audio_codes); + return prefix; } }; @@ -85,28 +138,42 @@ MossTextProcessor::~MossTextProcessor() = default; MossGenerationPrefix MossTextProcessor::build_generation_prefix( const std::string & text, const std::optional & language) const { + return impl_->build_prefix(text, language, [this](RowBuilder & builder) { + impl_->push_text(builder, kNoneValue); + }); +} + +MossGenerationPrefix MossTextProcessor::build_clone_prefix( + const std::string & text, + const std::vector> & reference_codes, + const std::optional & language) const { const auto & config = impl_->assets->config; + const int64_t n_vq = config.num_codebooks; + if (static_cast(reference_codes.size()) != n_vq) { + throw std::runtime_error("MOSS-TTS-Local clone prefix expects num_codebooks reference rows"); + } + const size_t frames = reference_codes.empty() ? 0 : reference_codes.front().size(); + if (frames == 0) { + throw std::runtime_error("MOSS-TTS-Local clone prefix requires a non-empty reference"); + } + for (const auto & row : reference_codes) { + if (row.size() != frames) { + throw std::runtime_error("MOSS-TTS-Local clone reference codebooks must be equal length"); + } + } - std::vector tokens; - tokens.push_back(static_cast(config.im_start_token_id)); - impl_->append_encoded(tokens, kUserRolePrefix); - impl_->append_encoded(tokens, kUserReferencePrefix); - impl_->append_encoded(tokens, kNoneValue); - impl_->append_encoded(tokens, render_after_reference(language)); - impl_->append_encoded(tokens, text); - impl_->append_encoded(tokens, kUserInstSuffix); - tokens.push_back(static_cast(config.im_end_token_id)); - impl_->append_encoded(tokens, kAssistantTurnPrefix); - tokens.push_back(static_cast(config.im_start_token_id)); - impl_->append_encoded(tokens, kAssistantRolePrefix); - tokens.push_back(static_cast(config.audio_start_token_id)); - - MossGenerationPrefix prefix; - prefix.text_tokens = std::move(tokens); - prefix.audio_codes.assign( - prefix.text_tokens.size() * static_cast(config.num_codebooks), - static_cast(config.audio_pad_token_id)); - return prefix; + return impl_->build_prefix(text, language, [&](RowBuilder & builder) { + // "- Reference(s):" slot -> audio_start, one audio_user_slot row per reference + // frame carrying that frame's codes, then audio_end. + builder.push_text_token(static_cast(config.audio_start_token_id)); + for (size_t frame = 0; frame < frames; ++frame) { + builder.text_tokens.push_back(static_cast(config.audio_user_slot_token_id)); + for (int64_t codebook = 0; codebook < n_vq; ++codebook) { + builder.audio_codes.push_back(reference_codes[static_cast(codebook)][frame]); + } + } + builder.push_text_token(static_cast(config.audio_end_token_id)); + }); } } // namespace engine::models::moss_tts_local diff --git a/src/models/moss_tts_local/session.cpp b/src/models/moss_tts_local/session.cpp new file mode 100644 index 0000000..917c91f --- /dev/null +++ b/src/models/moss_tts_local/session.cpp @@ -0,0 +1,223 @@ +#include "engine/models/moss_tts_local/session.h" + +#include "engine/framework/audio/resampling.h" +#include "engine/framework/runtime/options.h" +#include "engine/models/moss_tts_local/assets.h" + +#include +#include +#include +#include +#include +#include + +namespace engine::models::moss_tts_local { +namespace { + +constexpr size_t kBackboneWeightContextBytes = 64ull * 1024 * 1024; +constexpr size_t kBackboneGraphArenaBytes = 1024ull * 1024 * 1024; +constexpr size_t kDepthWeightContextBytes = 16ull * 1024 * 1024; +constexpr size_t kDepthGraphArenaBytes = 64ull * 1024 * 1024; +constexpr size_t kCodecWeightContextBytes = 256ull * 1024 * 1024; +constexpr size_t kCodecGraphArenaBytes = 1536ull * 1024 * 1024; +constexpr size_t kEncoderWeightContextBytes = 256ull * 1024 * 1024; +constexpr size_t kEncoderGraphArenaBytes = 2048ull * 1024 * 1024; +constexpr int kCodecSampleRate = 48000; + +// Prepares a speaker reference for the codec encoder, mirroring the processor's +// encode_audios_from_wav front end: de-interleave, force stereo (mono duplicated), +// resample to 48 kHz with the torchaudio sinc-Hann kernel, then loudness-normalize the +// whole clip to -20 dBFS (power) with a +/-3 dB gain clamp. +std::vector> reference_to_codec_stereo(const runtime::AudioBuffer & audio) { + const int channels = std::max(1, audio.channels); + const int64_t frames = channels > 0 ? static_cast(audio.samples.size()) / channels : 0; + if (frames <= 0) { + throw std::runtime_error("MOSS-TTS-Local voice reference audio is empty"); + } + + // De-interleave into per-channel streams, then force exactly two channels + // (mono is duplicated; extra channels are dropped) before resampling. + std::vector> stereo(2, std::vector(static_cast(frames))); + for (int64_t t = 0; t < frames; ++t) { + const float left = audio.samples[static_cast(t * channels)]; + const float right = channels > 1 ? audio.samples[static_cast(t * channels + 1)] : left; + stereo[0][static_cast(t)] = left; + stereo[1][static_cast(t)] = right; + } + + if (audio.sample_rate != kCodecSampleRate) { + if (audio.sample_rate <= 0) { + throw std::runtime_error("MOSS-TTS-Local voice reference has an invalid sample rate"); + } + for (auto & channel : stereo) { + channel = engine::audio::resample_mono_torchaudio_sinc_hann( + channel, audio.sample_rate, kCodecSampleRate); + } + } + + double sum_sq = 0.0; + size_t count = 0; + for (const auto & channel : stereo) { + for (const float value : channel) { + sum_sq += static_cast(value) * value; + } + count += channel.size(); + } + if (count > 0) { + const double mean_sq = sum_sq / static_cast(count); + const double current_dbfs = 10.0 * std::log10(mean_sq + 1.0e-9); + const double gain = std::max(-3.0, std::min(-20.0 - current_dbfs, 3.0)); + const float factor = static_cast(std::pow(10.0, gain / 20.0)); + for (auto & channel : stereo) { + for (float & value : channel) { + value *= factor; + } + } + } + return stereo; +} + +std::shared_ptr require_assets(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("MOSS-TTS-Local session requires assets"); + } + return assets; +} + +MossGenerationOptions generation_options_from_request(const runtime::TaskRequest & request) { + MossGenerationOptions options; + if (const auto value = runtime::parse_i64_option(request.options, {"max_frames", "max_tokens"})) { + if (*value <= 0) { + throw std::runtime_error("MOSS-TTS-Local max_tokens must be positive"); + } + options.max_new_frames = *value; + } + if (const auto value = runtime::find_option(request.options, {"do_sample"})) { + options.do_sample = runtime::parse_bool_option(*value, "do_sample"); + } + if (const auto value = runtime::parse_float_option(request.options, {"audio_temperature", "temperature"})) { + options.audio_temperature = *value; + } + if (const auto value = runtime::parse_int_option(request.options, {"audio_top_k", "top_k"})) { + options.audio_top_k = *value; + } + if (const auto value = runtime::parse_float_option(request.options, {"audio_top_p", "top_p"})) { + options.audio_top_p = *value; + } + if (const auto value = + runtime::parse_float_option(request.options, {"audio_repetition_penalty", "repetition_penalty"})) { + options.audio_repetition_penalty = *value; + } + options.seed = runtime::parse_u32_option(request.options, {"seed"}).value_or(runtime::random_u32_seed()); + return options; +} + +} // namespace + +MossTTSLocalSession::MossTTSLocalSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets) + : RuntimeSessionBase(options), + task_(task), + assets_(require_assets(std::move(assets))) { + backbone_ = std::make_unique( + assets_, + execution_context(), + kBackboneGraphArenaBytes, + kBackboneWeightContextBytes, + assets::TensorStorageType::Native); + depth_ = std::make_unique( + assets_, execution_context(), kDepthGraphArenaBytes, kDepthWeightContextBytes); + processor_ = std::make_unique(assets_); + codec_ = std::make_unique( + resolve_moss_codec_dir(*assets_), + execution_context(), + assets_->config.num_codebooks, + kCodecWeightContextBytes, + kCodecGraphArenaBytes); + generator_ = std::make_unique(assets_, *backbone_, *depth_); +} + +std::string MossTTSLocalSession::family() const { + return "moss_tts_local"; +} + +runtime::VoiceTaskKind MossTTSLocalSession::task_kind() const { + return task_.task; +} + +runtime::RunMode MossTTSLocalSession::run_mode() const { + return task_.mode; +} + +void MossTTSLocalSession::prepare(const runtime::SessionPreparationRequest & request) { + (void) request; + mark_prepared(); +} + +runtime::TaskResult MossTTSLocalSession::run(const runtime::TaskRequest & request) { + if (!request.text_input.has_value() || request.text_input->text.empty()) { + throw std::runtime_error("MOSS-TTS-Local requires text input"); + } + const bool has_reference = request.voice.has_value() && request.voice->speaker.has_value() && + request.voice->speaker->audio.has_value(); + + std::optional language; + const std::string & language_tag = request.text_input->language; + if (!language_tag.empty() && language_tag != "Auto") { + language = language_tag; + } + + const auto options = generation_options_from_request(request); + MossGenerationPrefix prefix; + if (has_reference) { + const auto stereo = reference_to_codec_stereo(*request.voice->speaker->audio); + if (encoder_ == nullptr) { + encoder_ = std::make_unique( + resolve_moss_codec_dir(*assets_), + execution_context(), + assets_->config.num_codebooks, + kEncoderWeightContextBytes, + kEncoderGraphArenaBytes); + } + const auto reference_codes = encoder_->encode(stereo); + prefix = processor_->build_clone_prefix(request.text_input->text, reference_codes, language); + } else { + prefix = processor_->build_generation_prefix(request.text_input->text, language); + } + const auto frames = generator_->generate(prefix.text_tokens, prefix.audio_codes, options); + if (frames.empty()) { + throw std::runtime_error("MOSS-TTS-Local generated no audio frames"); + } + + const int64_t num_codebooks = assets_->config.num_codebooks; + std::vector> codes( + static_cast(num_codebooks), std::vector(frames.size())); + for (size_t frame = 0; frame < frames.size(); ++frame) { + for (int64_t codebook = 0; codebook < num_codebooks; ++codebook) { + codes[static_cast(codebook)][frame] = frames[frame][static_cast(codebook)]; + } + } + + const auto channels = codec_->decode(codes); + const int channel_count = static_cast(channels.size()); + const size_t samples_per_channel = channels.empty() ? 0 : channels.front().size(); + + runtime::AudioBuffer audio; + audio.sample_rate = static_cast(codec_->sampling_rate()); + audio.channels = channel_count; + audio.samples.resize(samples_per_channel * static_cast(channel_count)); + for (size_t sample = 0; sample < samples_per_channel; ++sample) { + for (int channel = 0; channel < channel_count; ++channel) { + audio.samples[sample * static_cast(channel_count) + static_cast(channel)] = + channels[static_cast(channel)][sample]; + } + } + + runtime::TaskResult result; + result.audio_output = std::move(audio); + return result; +} + +} // namespace engine::models::moss_tts_local diff --git a/tests/moss_tts_local/codec_encode_parity.cpp b/tests/moss_tts_local/codec_encode_parity.cpp new file mode 100644 index 0000000..756e42d --- /dev/null +++ b/tests/moss_tts_local/codec_encode_parity.cpp @@ -0,0 +1,154 @@ +// Parity check for the MOSS-Audio-Tokenizer-v2 encoder (voice-clone front end). +// Loads the prepared reference waveform captured by scratchpad/moss_encode_ref.py +// (enc_prepared.f32, [2, samples] channel-major @ 48 kHz), runs the C++ encoder, +// and compares the resulting [num_quantizers, frames] codes against the Python +// reference (enc_codes.csv, first `num_quantizers` rows). Exact code match is the +// gate: the RLFQ quantizer is robust to small latent differences, so matching +// codes confirms both the encoder stack and the quantize path. + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" +#include "engine/models/moss_tts_local/codec_encoder.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +std::string arg_value(int argc, char ** argv, const std::string & name, const std::string & fallback) { + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) { + return argv[i + 1]; + } + } + return fallback; +} + +int int_arg(int argc, char ** argv, const std::string & name, int fallback) { + return std::stoi(arg_value(argc, argv, name, std::to_string(fallback))); +} + +std::vector read_f32(const std::string & path, size_t expected) { + std::ifstream file(path, std::ios::binary); + if (!file) { + throw std::runtime_error("cannot open " + path); + } + std::vector data(expected); + file.read(reinterpret_cast(data.data()), static_cast(expected * sizeof(float))); + if (static_cast(file.gcount()) != expected * sizeof(float)) { + throw std::runtime_error("short read from " + path); + } + return data; +} + +std::vector> read_codes_csv(const std::string & path) { + std::ifstream file(path); + if (!file) { + throw std::runtime_error("cannot open " + path); + } + std::vector> rows; + std::string line; + while (std::getline(file, line)) { + if (line.empty()) { + continue; + } + std::vector row; + std::stringstream ss(line); + std::string cell; + while (std::getline(ss, cell, ',')) { + row.push_back(std::stoi(cell)); + } + rows.push_back(std::move(row)); + } + return rows; +} + +} // namespace + +int main(int argc, char ** argv) { + try { + const std::string codec_dir = arg_value( + argc, argv, "--codec", + "C:/Users/justi/.cache/huggingface/hub/models--OpenMOSS-Team--MOSS-Audio-Tokenizer-v2/" + "snapshots/f6e20e543b33d2c252a7ef71bdf8aa71e5ff9169"); + const std::string prepared = arg_value(argc, argv, "--prepared", "enc_prepared.f32"); + const std::string ref_codes_path = arg_value(argc, argv, "--ref-codes", "enc_codes.csv"); + const int channels = int_arg(argc, argv, "--channels", 2); + const int samples = int_arg(argc, argv, "--samples", 96000); + const int num_quantizers = int_arg(argc, argv, "--quantizers", 12); + const int threads = int_arg(argc, argv, "--threads", 16); + + std::cout << "loading prepared waveform [" << channels << "," << samples << "]...\n" << std::flush; + const auto wav = read_f32(prepared, static_cast(channels) * static_cast(samples)); + std::vector> stereo(2, std::vector(static_cast(samples))); + for (int c = 0; c < 2; ++c) { + for (int i = 0; i < samples; ++i) { + stereo[static_cast(c)][static_cast(i)] = + wav[static_cast(c) * static_cast(samples) + static_cast(i)]; + } + } + + constexpr size_t kWeightContextBytes = 256ull * 1024 * 1024; + constexpr size_t kGraphArenaBytes = 2048ull * 1024 * 1024; + + engine::core::BackendConfig backend_config; + backend_config.type = engine::core::BackendType::Cpu; + backend_config.device = 0; + backend_config.threads = threads; + engine::core::ExecutionContext execution_context(backend_config); + + std::cout << "loading codec encoder weights...\n" << std::flush; + engine::models::moss_tts_local::MossCodecEncoder encoder( + codec_dir, execution_context, num_quantizers, kWeightContextBytes, kGraphArenaBytes); + + std::cout << "encoding...\n" << std::flush; + const auto codes = encoder.encode(stereo); + const int64_t frames = codes.empty() ? 0 : static_cast(codes.front().size()); + std::cout << "produced codes [" << codes.size() << "," << frames << "]\n"; + + const auto ref = read_codes_csv(ref_codes_path); + std::cout << "reference codes rows=" << ref.size() + << " cols=" << (ref.empty() ? 0 : ref.front().size()) << "\n"; + + int64_t total = 0; + int64_t matched = 0; + int mismatched_rows = 0; + for (size_t q = 0; q < codes.size() && q < ref.size(); ++q) { + int64_t row_match = 0; + for (int64_t t = 0; t < frames && t < static_cast(ref[q].size()); ++t) { + ++total; + if (codes[q][static_cast(t)] == ref[q][static_cast(t)]) { + ++matched; + ++row_match; + } + } + if (row_match != frames) { + ++mismatched_rows; + if (mismatched_rows <= 4) { + std::cout << " cb" << q << ": " << row_match << "/" << frames << " match; cpp[0:8]=["; + for (int64_t t = 0; t < 8 && t < frames; ++t) { + std::cout << (t ? "," : "") << codes[q][static_cast(t)]; + } + std::cout << "] ref[0:8]=["; + for (int64_t t = 0; t < 8 && t < static_cast(ref[q].size()); ++t) { + std::cout << (t ? "," : "") << ref[q][static_cast(t)]; + } + std::cout << "]\n"; + } + } + } + + std::cout << "\n=== ENCODE PARITY: " << matched << "/" << total << " codes match ===\n"; + const bool ok = (total > 0 && matched == total); + std::cout << (ok ? "PASS" : "FAIL") << "\n"; + return ok ? 0 : 1; + } catch (const std::exception & error) { + std::cerr << "codec_encode_parity failed: " << error.what() << "\n"; + return 1; + } +} diff --git a/tests/moss_tts_local/moss_tts_local_smoke.cpp b/tests/moss_tts_local/moss_tts_local_smoke.cpp index c65775b..6ae5ce7 100644 --- a/tests/moss_tts_local/moss_tts_local_smoke.cpp +++ b/tests/moss_tts_local/moss_tts_local_smoke.cpp @@ -15,8 +15,10 @@ #include #include #include +#include #include #include +#include #include #include @@ -65,6 +67,7 @@ int main(int argc, char ** argv) { const int threads = int_arg(argc, argv, "--threads", 8); const int max_frames = int_arg(argc, argv, "--frames", 12); const auto weight_type = parse_weight_type(arg_value(argc, argv, "--weight-type", "native")); + const std::string dump_path = arg_value(argc, argv, "--dump", ""); constexpr size_t kBackboneWeightContextBytes = 64ull * 1024 * 1024; constexpr size_t kBackboneGraphArenaBytes = 512ull * 1024 * 1024; @@ -92,6 +95,46 @@ int main(int argc, char ** argv) { engine::models::moss_tts_local::MossGenerator generator(assets, backbone, depth); engine::models::moss_tts_local::MossTextProcessor processor(assets); + // Optional: build a voice-clone prefix from a reference-codes CSV ([nq, frames], + // e.g. the encoder's enc_codes.csv) and dump the [seq, 1+nq] input_ids so it can be + // diffed against the Python processor. Isolates the clone-prefix assembly. + const std::string clone_ref = arg_value(argc, argv, "--clone-ref", ""); + if (!clone_ref.empty()) { + const int64_t n_vq = assets->config.num_codebooks; + std::ifstream csv(clone_ref); + if (!csv) { + throw std::runtime_error("cannot open clone-ref csv: " + clone_ref); + } + std::vector> ref_rows; + std::string line; + while (std::getline(csv, line)) { + if (line.empty()) { + continue; + } + std::vector row; + std::stringstream ss(line); + std::string cell; + while (std::getline(ss, cell, ',')) { + row.push_back(std::stoi(cell)); + } + ref_rows.push_back(std::move(row)); + } + ref_rows.resize(static_cast(n_vq)); // first n_vq codebooks + const auto clone_prefix = processor.build_clone_prefix(text, ref_rows); + const std::string dump = arg_value(argc, argv, "--dump", "clone_prefix_cpp.csv"); + std::ofstream out(dump); + const size_t seq = clone_prefix.text_tokens.size(); + for (size_t row = 0; row < seq; ++row) { + out << clone_prefix.text_tokens[row]; + for (int64_t cb = 0; cb < n_vq; ++cb) { + out << "," << clone_prefix.audio_codes[row * static_cast(n_vq) + static_cast(cb)]; + } + out << "\n"; + } + std::cout << "clone_prefix rows=" << seq << " -> " << dump << "\n"; + return 0; + } + const auto prefix = processor.build_generation_prefix(text); std::cout << "prompt_tokens=" << prefix.text_tokens.size() << "\n"; std::cout << "first_token_ids=["; @@ -141,6 +184,20 @@ int main(int argc, char ** argv) { } std::cout << "]\n"; } + + if (!dump_path.empty()) { + std::ofstream dump(dump_path); + if (!dump) { + throw std::runtime_error("failed to open dump file: " + dump_path); + } + for (const auto & frame : frames) { + for (int64_t codebook = 0; codebook < n_vq; ++codebook) { + dump << (codebook == 0 ? "" : ",") << frame[static_cast(codebook)]; + } + dump << "\n"; + } + std::cout << "dumped " << frames.size() << " frames x " << n_vq << " codes to " << dump_path << "\n"; + } return 0; } catch (const std::exception & error) { std::cerr << "moss_tts_local_smoke failed: " << error.what() << "\n"; From 802b2b28926f0eab5759ef8a8d208acd803d1fe3 Mon Sep 17 00:00:00 2001 From: justinjohn0306 Date: Thu, 2 Jul 2026 21:46:26 +0530 Subject: [PATCH 4/4] Add MOSS-TTS-Local Phase 6: KV cache, batched prefill, auto dtype, test harness Backbone generation re-forwarded the whole sequence every frame (O(T^2)), which made long text slower than the Python reference and voice clones far slower. This adds an incremental KV-cached generation path plus the tests and tooling the PR review asked for. - backbone: per-layer KV cache with a reusable single-position step graph (begin_generation/step) and a single batched prefill that seeds the cache in one forward (prefill), replacing the per-frame re-forward. Removes the now-dead forward_prefill path. Greedy output is byte-identical to the previous path. - session: hardware-adaptive "auto" weight_type (CUDA bf16, CPU f32, other backends native), overridable via moss_tts_local.weight_type. - tools: audiocpp_cli path cases (text-only, voice clone, sampled, long-lived session) and a longform clone case; reference/parity/perf/session-probe scripts; per-request and session timing in the offline batch runner. Perf (CUDA bf16, RTF vs Python): text-only 0.59 vs 0.95, sampled 0.63 vs 0.92, long text 0.60 vs 0.74, long clone 0.61 vs 0.75 -- every generation-bound path now faster than Python. Short clones remain encoder-bound (codec encoder O(T^2), a separate path). Teacher-forced codec-decode parity: cosine 0.999999996. Co-Authored-By: Claude Opus 4.8 --- app/workflow/execution.cpp | 16 +- app/workflow/execution.h | 3 + app/workflow/file_sink.cpp | 4 + .../engine/models/moss_tts_local/backbone.h | 29 +- src/models/moss_tts_local/backbone.cpp | 352 +++++++++++++++--- src/models/moss_tts_local/generator.cpp | 61 +-- src/models/moss_tts_local/session.cpp | 41 +- ...audiocpp_cli_longform_tts_clone_cases.json | 21 ++ .../audiocpp_cli/audiocpp_cli_path_cases.json | 105 ++++++ .../moss_tts_local_codec_parity.py | 97 +++++ .../audiocpp_cli/moss_tts_local_reference.py | 296 +++++++++++++++ tools/audiocpp_cli/moss_tts_local_report.py | 169 +++++++++ .../moss_tts_local_session_probe.py | 218 +++++++++++ .../run_audiocpp_cli_path_tests.py | 8 + 14 files changed, 1328 insertions(+), 92 deletions(-) create mode 100644 tools/audiocpp_cli/moss_tts_local_codec_parity.py create mode 100644 tools/audiocpp_cli/moss_tts_local_reference.py create mode 100644 tools/audiocpp_cli/moss_tts_local_report.py create mode 100644 tools/audiocpp_cli/moss_tts_local_session_probe.py diff --git a/app/workflow/execution.cpp b/app/workflow/execution.cpp index 9efff56..0c873e0 100644 --- a/app/workflow/execution.cpp +++ b/app/workflow/execution.cpp @@ -1,5 +1,6 @@ #include "execution.h" +#include #include namespace minitts::app { @@ -53,18 +54,31 @@ AppBatchResult run_offline_batch( if (batch.requests.empty()) { throw std::runtime_error("offline batch requires at least one request"); } + using Clock = std::chrono::steady_clock; + const auto to_ms = [](Clock::duration d) { + return std::chrono::duration(d).count(); + }; + const auto session_start = Clock::now(); + + const auto prepare_start = Clock::now(); session.prepare(engine::runtime::build_preparation_request(batch.requests.front().request)); AppBatchResult out; + out.prepare_ms = to_ms(Clock::now() - prepare_start); out.results.reserve(batch.requests.size()); for (const auto & item : batch.requests) { + const auto run_start = Clock::now(); + auto result = offline.run(item.request); + const double wall_ms = to_ms(Clock::now() - run_start); out.results.push_back(AppRequestResult{ item.id, - offline.run(item.request), + std::move(result), + wall_ms, }); if (on_result) { on_result(out.results.size() - 1, out.results.back()); } } + out.session_wall_ms = to_ms(Clock::now() - session_start); if (audio_merge_mode == AudioMergeMode::Concat) { out.merged_audio = concat_audio_outputs(out.results, out.chapters); } diff --git a/app/workflow/execution.h b/app/workflow/execution.h index c92130d..fa9a743 100644 --- a/app/workflow/execution.h +++ b/app/workflow/execution.h @@ -19,6 +19,7 @@ struct AppRequest { struct AppRequestResult { std::string id; engine::runtime::TaskResult result; + double wall_ms = 0.0; }; struct AudioChapter { @@ -35,6 +36,8 @@ struct AppBatchResult { std::vector results; std::optional merged_audio; std::vector chapters; + double prepare_ms = 0.0; + double session_wall_ms = 0.0; }; enum class AudioMergeMode { diff --git a/app/workflow/file_sink.cpp b/app/workflow/file_sink.cpp index bfdb261..fb95e90 100644 --- a/app/workflow/file_sink.cpp +++ b/app/workflow/file_sink.cpp @@ -318,6 +318,9 @@ void emit_batch_result( for (size_t i = 0; i < batch.results.size(); ++i) { emit_batch_item_result(i, batch.results[i], policy); } + + std::cout << "[TIMING] session.prepare_ms " << batch.prepare_ms << "\n"; + std::cout << "[TIMING] session.wall_ms " << batch.session_wall_ms << "\n"; } void emit_batch_summary( @@ -351,6 +354,7 @@ void emit_batch_item_result( const std::string request_id = safe_output_name(item.id); std::cout << "request_index=" << index << "\n"; std::cout << "request_id=" << request_id << "\n"; + std::cout << "[TIMING] request." << request_id << ".wall_ms " << item.wall_ms << "\n"; std::optional audio_out; std::optional named_out_dir; diff --git a/include/engine/models/moss_tts_local/backbone.h b/include/engine/models/moss_tts_local/backbone.h index 217375c..74804ee 100644 --- a/include/engine/models/moss_tts_local/backbone.h +++ b/include/engine/models/moss_tts_local/backbone.h @@ -29,19 +29,26 @@ class MossBackboneRuntime { int64_t hidden_size() const noexcept; - // Runs the backbone over a prefill sequence of text token ids and returns the - // final hidden states as [steps, hidden_size] row-major float. - std::vector forward_prefill(const std::vector & token_ids) const; - - // Same as forward_prefill, but adds the per-position audio-codebook embedding sum - // (audio_bias, [steps * hidden_size] row-major) to the text embedding before the - // decoder stack. This mirrors MossTTSLocalModel._build_inputs_embeds. - std::vector forward_prefill_fused( - const std::vector & token_ids, - const std::vector & audio_bias) const; + // Incremental (KV-cached) generation. Call begin_generation once to size and allocate the + // per-layer cache, prefill the prompt in a single batched forward, then step one position + // at a time. This runs generation in O(T) work per step instead of the O(T^2) cost of + // re-forwarding the whole sequence every frame. + // + // begin_generation sizes the cache for max_positions and builds the reusable + // single-position step graph. + void begin_generation(int64_t max_positions) const; + // prefill forwards the whole prompt at once, writes every position's K/V into the cache, + // and returns the last position's hidden state ([hidden_size]) to seed the first frame. + // audio_bias is the summed audio-codebook embedding per position ([token_ids.size() * + // hidden_size] row-major), mirroring MossTTSLocalModel._build_inputs_embeds. + std::vector prefill(const std::vector & token_ids, const std::vector & audio_bias) const; + // step forwards one fused position (its text token embedded in-graph plus audio_bias_row), + // appends that position's K/V to the cache, and returns its hidden state ([hidden_size]). + std::vector step(int32_t token_id, const std::vector & audio_bias_row) const; + int64_t cached_positions() const noexcept; private: - std::vector run_prefill(const std::vector & token_ids, const float * audio_bias) const; + void build_step_graph(int64_t cache_steps) const; struct Impl; std::unique_ptr impl_; diff --git a/src/models/moss_tts_local/backbone.cpp b/src/models/moss_tts_local/backbone.cpp index 83993ef..2ab1b1e 100644 --- a/src/models/moss_tts_local/backbone.cpp +++ b/src/models/moss_tts_local/backbone.cpp @@ -7,6 +7,7 @@ #include "engine/framework/modules/linear_module.h" #include "engine/framework/modules/lookup_modules.h" #include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/optimizations/fast_kv_modules.h" #include "engine/framework/modules/positional_modules.h" #include "engine/framework/modules/primitive_modules.h" #include "engine/framework/modules/structural_modules.h" @@ -211,7 +212,17 @@ core::TensorValue attention_from_heads( return matmul.build(ctx, attn, v_heads); } -core::TensorValue decoder_layer( +// Prefill decoder layer: a standard batched causal-attention layer that also exposes this +// layer's post-RoPE key and pre-transpose value ([1, steps, kv_heads, dim]) so the caller can +// seed the generation cache from a single full-prompt forward instead of stepping position by +// position. key/value are laid out exactly as decoder_layer_cached writes them per step. +struct PrefillLayerOutput { + core::TensorValue hidden; + core::TensorValue key; + core::TensorValue value; +}; + +PrefillLayerOutput decoder_layer_prefill( core::ModuleBuildContext & ctx, const core::TensorValue & input, const core::TensorValue & positions, @@ -240,6 +251,10 @@ core::TensorValue decoder_layer( q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, q, positions); k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, k, positions); + // Capture k/v in the same [1, steps, kv_heads, dim] layout the per-step cache stores. + auto cache_key = ensure_contiguous(ctx, k); + auto cache_value = ensure_contiguous(ctx, v); + auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); auto k_heads = repeat_kv_heads(ctx, modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k), kv_repeats); auto v_heads = repeat_kv_heads(ctx, modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v), kv_repeats); @@ -252,6 +267,72 @@ core::TensorValue decoder_layer( core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], config.num_attention_heads * dim})); auto x = modules::AddModule{}.build(ctx, input, o_proj.build(ctx, context, binding::linear_data(ctx, weights.o_proj))); + auto ff_in = hidden_norm.build(ctx, x, binding::norm_data(ctx, weights.post_norm)); + auto gate = modules::LinearModule( + binding::linear_config(config.hidden_size, config.intermediate_size, false)) + .build(ctx, ff_in, binding::linear_data(ctx, weights.gate_proj)); + gate = modules::SiluModule{}.build(ctx, gate); + auto up = modules::LinearModule( + binding::linear_config(config.hidden_size, config.intermediate_size, false)) + .build(ctx, ff_in, binding::linear_data(ctx, weights.up_proj)); + auto ff = modules::LinearModule( + binding::linear_config(config.intermediate_size, config.hidden_size, false)) + .build(ctx, modules::MulModule{}.build(ctx, gate, up), binding::linear_data(ctx, weights.down_proj)); + return {modules::AddModule{}.build(ctx, x, ff), cache_key, cache_value}; +} + +// Cached (single-position) decoder layer: computes q/k/v for the one new row, writes k/v into +// the persistent per-layer cache at cache_slot, then attends the new query over the whole +// cache (invalid rows suppressed by attention_mask). Same math as the prefill layer, one row +// at a time, so cached generation matches a full re-forward. +core::TensorValue decoder_layer_cached( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, // [1, 1, hidden] + const core::TensorValue & positions, // i32 [1] + const BackboneLayerWeights & weights, + const MossBackboneConfig & config, + const core::TensorValue & cache_key, // [1, cache_steps, kv_heads, dim] + const core::TensorValue & cache_value, // [1, cache_steps, kv_heads, dim] + const core::TensorValue & cache_slot, // i32 [1] + const core::TensorValue & attention_mask) { // [1, 1, 1, cache_steps] + const int64_t dim = config.head_dim; + const int64_t kv_repeats = config.num_attention_heads / config.num_key_value_heads; + const modules::LinearModule q_proj( + binding::linear_config(config.hidden_size, config.num_attention_heads * dim, false)); + const modules::LinearModule k_proj( + binding::linear_config(config.hidden_size, config.num_key_value_heads * dim, false)); + const modules::LinearModule v_proj( + binding::linear_config(config.hidden_size, config.num_key_value_heads * dim, false)); + const modules::LinearModule o_proj( + binding::linear_config(config.num_attention_heads * dim, config.hidden_size, false)); + const modules::RMSNormModule hidden_norm({config.hidden_size, config.rms_norm_eps, true, false}); + const modules::RMSNormModule head_norm({dim, config.rms_norm_eps, true, false}); + auto x_norm = hidden_norm.build(ctx, input, binding::norm_data(ctx, weights.input_norm)); + auto q = q_proj.build(ctx, x_norm, binding::linear_data(ctx, weights.q_proj)); + auto k = k_proj.build(ctx, x_norm, binding::linear_data(ctx, weights.k_proj)); + auto v = v_proj.build(ctx, x_norm, binding::linear_data(ctx, weights.v_proj)); + q = head_norm.build(ctx, reshape_heads(ctx, q, config.num_attention_heads, dim), binding::norm_data(ctx, weights.q_norm)); + k = head_norm.build(ctx, reshape_heads(ctx, k, config.num_key_value_heads, dim), binding::norm_data(ctx, weights.k_norm)); + v = reshape_heads(ctx, v, config.num_key_value_heads, dim); + q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, q, positions); + k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, k, positions); + + const modules::FastKVSetRowsModule set_rows; + auto all_k = set_rows.build(ctx, cache_key, k, cache_slot); + auto all_v = set_rows.build(ctx, cache_value, v, cache_slot); + + auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); + auto k_heads = repeat_kv_heads(ctx, modules::TransposeModule({{0, 2, 1, 3}, all_k.shape.rank}).build(ctx, all_k), kv_repeats); + auto v_heads = repeat_kv_heads(ctx, modules::TransposeModule({{0, 2, 1, 3}, all_v.shape.rank}).build(ctx, all_v), kv_repeats); + auto context = attention_from_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask); + context = modules::TransposeModule({{0, 2, 1, 3}, context.shape.rank}).build(ctx, context); + context = ensure_contiguous(ctx, context); + context = core::reshape_tensor( + ctx, + context, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], config.num_attention_heads * dim})); + auto x = modules::AddModule{}.build(ctx, input, o_proj.build(ctx, context, binding::linear_data(ctx, weights.o_proj))); + auto ff_in = hidden_norm.build(ctx, x, binding::norm_data(ctx, weights.post_norm)); auto gate = modules::LinearModule( binding::linear_config(config.hidden_size, config.intermediate_size, false)) @@ -274,6 +355,31 @@ struct MossBackboneRuntime::Impl { core::BackendType backend_type = core::BackendType::Cpu; size_t graph_arena_bytes = 0; BackboneWeights weights; + + // Cached-generation step graph (built once by begin_generation, reused every step). + std::unique_ptr step_ctx; + ggml_cgraph * step_graph = nullptr; + ggml_backend_buffer_t step_buffer = nullptr; + ggml_tensor * step_token = nullptr; + ggml_tensor * step_bias = nullptr; + ggml_tensor * step_positions = nullptr; + ggml_tensor * step_cache_slot = nullptr; + ggml_tensor * step_mask = nullptr; + ggml_tensor * step_hidden = nullptr; + std::vector cache_keys; + std::vector cache_values; + int64_t cache_steps = 0; + int64_t valid_steps = 0; + std::vector mask_host; + + ~Impl() { + if (step_graph != nullptr && backend != nullptr) { + core::release_backend_graph_resources(backend, step_graph); + } + if (step_buffer != nullptr) { + ggml_backend_buffer_free(step_buffer); + } + } }; MossBackboneRuntime::MossBackboneRuntime( @@ -310,37 +416,166 @@ int64_t MossBackboneRuntime::hidden_size() const noexcept { return impl_->assets->config.backbone.hidden_size; } -std::vector MossBackboneRuntime::forward_prefill(const std::vector & token_ids) const { - return run_prefill(token_ids, nullptr); +void MossBackboneRuntime::build_step_graph(int64_t cache_steps) const { + auto & impl = *impl_; + const auto & config = impl.assets->config.backbone; + const auto & weights = impl.weights; + const int64_t dim = config.head_dim; + + ggml_init_params params{impl.graph_arena_bytes, nullptr, true}; + impl.step_ctx.reset(ggml_init(params)); + if (impl.step_ctx == nullptr) { + throw std::runtime_error("failed to initialize MOSS-TTS-Local backbone step graph context"); + } + ggml_context * gctx = impl.step_ctx.get(); + core::ModuleBuildContext ctx{gctx, "moss_tts_local.backbone.step", impl.backend_type}; + + auto token_input = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, 1})); + ggml_set_input(token_input.tensor); + auto bias_input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, config.hidden_size})); + ggml_set_input(bias_input.tensor); + auto positions = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1})); + ggml_set_input(positions.tensor); + auto cache_slot = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1})); + ggml_set_input(cache_slot.tensor); + auto attention_mask = + core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, 1, cache_steps})); + ggml_set_input(attention_mask.tensor); + + impl.cache_keys.clear(); + impl.cache_values.clear(); + impl.cache_keys.reserve(static_cast(config.num_hidden_layers)); + impl.cache_values.reserve(static_cast(config.num_hidden_layers)); + + auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size}) + .build(ctx, token_input, weights.embed_tokens); + x = modules::AddModule{}.build(ctx, x, bias_input); + for (const auto & layer : weights.layers) { + auto cache_key = core::make_tensor( + ctx, GGML_TYPE_F32, + core::TensorShape::from_dims({1, cache_steps, config.num_key_value_heads, dim})); + auto cache_value = core::make_tensor( + ctx, GGML_TYPE_F32, + core::TensorShape::from_dims({1, cache_steps, config.num_key_value_heads, dim})); + impl.cache_keys.push_back(cache_key.tensor); + impl.cache_values.push_back(cache_value.tensor); + x = decoder_layer_cached(ctx, x, positions, layer, config, cache_key, cache_value, cache_slot, attention_mask); + } + x = modules::RMSNormModule({config.hidden_size, config.rms_norm_eps, true, false}) + .build(ctx, x, binding::norm_data(ctx, weights.norm)); + x = ensure_contiguous(ctx, x); + auto hidden = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, config.hidden_size})); + hidden = ensure_contiguous(ctx, hidden); + ggml_set_output(hidden.tensor); + + impl.step_graph = ggml_new_graph_custom(gctx, 65536, false); + ggml_build_forward_expand(impl.step_graph, hidden.tensor); + + impl.step_buffer = ggml_backend_alloc_ctx_tensors(gctx, impl.backend); + if (impl.step_buffer == nullptr) { + throw std::runtime_error("failed to allocate MOSS-TTS-Local backbone step graph"); + } + + impl.step_token = token_input.tensor; + impl.step_bias = bias_input.tensor; + impl.step_positions = positions.tensor; + impl.step_cache_slot = cache_slot.tensor; + impl.step_mask = attention_mask.tensor; + impl.step_hidden = hidden.tensor; + impl.cache_steps = cache_steps; + impl.mask_host.assign(static_cast(cache_steps), kMaskedAttentionBias); } -std::vector MossBackboneRuntime::forward_prefill_fused( - const std::vector & token_ids, - const std::vector & audio_bias) const { - const int64_t steps = static_cast(token_ids.size()); - const int64_t hidden = impl_->assets->config.backbone.hidden_size; - if (static_cast(audio_bias.size()) != steps * hidden) { - throw std::runtime_error("MOSS-TTS-Local backbone audio bias size does not match [steps, hidden]"); +void MossBackboneRuntime::begin_generation(int64_t max_positions) const { + if (max_positions <= 0) { + throw std::runtime_error("MOSS-TTS-Local backbone begin_generation requires max_positions > 0"); + } + auto & impl = *impl_; + if (impl.step_graph == nullptr || impl.cache_steps < max_positions) { + if (impl.step_graph != nullptr) { + core::release_backend_graph_resources(impl.backend, impl.step_graph); + impl.step_graph = nullptr; + } + if (impl.step_buffer != nullptr) { + ggml_backend_buffer_free(impl.step_buffer); + impl.step_buffer = nullptr; + } + build_step_graph(max_positions); } - return run_prefill(token_ids, audio_bias.data()); + // Zero the caches so not-yet-written (masked) rows can never inject NaNs into the softmax. + const auto & config = impl.assets->config.backbone; + const size_t elems = + static_cast(impl.cache_steps * config.num_key_value_heads * config.head_dim); + const std::vector zeros(elems, 0.0F); + for (auto * tensor : impl.cache_keys) { + ggml_backend_tensor_set(tensor, zeros.data(), 0, elems * sizeof(float)); + } + for (auto * tensor : impl.cache_values) { + ggml_backend_tensor_set(tensor, zeros.data(), 0, elems * sizeof(float)); + } + impl.valid_steps = 0; +} + +std::vector MossBackboneRuntime::step(int32_t token_id, const std::vector & audio_bias_row) const { + auto & impl = *impl_; + if (impl.step_graph == nullptr) { + throw std::runtime_error("MOSS-TTS-Local backbone step called before begin_generation"); + } + const int64_t hidden = impl.assets->config.backbone.hidden_size; + if (static_cast(audio_bias_row.size()) != hidden) { + throw std::runtime_error("MOSS-TTS-Local backbone step audio bias row size does not match hidden_size"); + } + if (impl.valid_steps >= impl.cache_steps) { + throw std::runtime_error("MOSS-TTS-Local backbone step exceeds cache capacity"); + } + const int32_t position = static_cast(impl.valid_steps); + ggml_backend_tensor_set(impl.step_token, &token_id, 0, sizeof(int32_t)); + ggml_backend_tensor_set(impl.step_bias, audio_bias_row.data(), 0, audio_bias_row.size() * sizeof(float)); + ggml_backend_tensor_set(impl.step_positions, &position, 0, sizeof(int32_t)); + ggml_backend_tensor_set(impl.step_cache_slot, &position, 0, sizeof(int32_t)); + for (int64_t i = 0; i < impl.cache_steps; ++i) { + impl.mask_host[static_cast(i)] = (i <= impl.valid_steps) ? 0.0F : kMaskedAttentionBias; + } + ggml_backend_tensor_set(impl.step_mask, impl.mask_host.data(), 0, impl.mask_host.size() * sizeof(float)); + + const ggml_status status = ggml_backend_graph_compute(impl.backend, impl.step_graph); + ggml_backend_synchronize(impl.backend); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("MOSS-TTS-Local backbone step graph compute failed"); + } + std::vector hidden_state(static_cast(hidden)); + ggml_backend_tensor_get(impl.step_hidden, hidden_state.data(), 0, hidden_state.size() * sizeof(float)); + ++impl.valid_steps; + return hidden_state; } -std::vector MossBackboneRuntime::run_prefill( +std::vector MossBackboneRuntime::prefill( const std::vector & token_ids, - const float * audio_bias) const { + const std::vector & audio_bias) const { + auto & impl = *impl_; + if (impl.step_graph == nullptr) { + throw std::runtime_error("MOSS-TTS-Local backbone prefill called before begin_generation"); + } + const auto & config = impl.assets->config.backbone; const int64_t steps = static_cast(token_ids.size()); if (steps <= 0) { - throw std::runtime_error("MOSS-TTS-Local backbone prefill requires a non-empty token sequence"); + throw std::runtime_error("MOSS-TTS-Local backbone prefill requires a non-empty prompt"); + } + if (steps > impl.cache_steps) { + throw std::runtime_error("MOSS-TTS-Local backbone prefill prompt exceeds cache capacity"); } - const auto & config = impl_->assets->config.backbone; - const auto & weights = impl_->weights; + if (static_cast(audio_bias.size()) != steps * config.hidden_size) { + throw std::runtime_error("MOSS-TTS-Local backbone prefill audio bias size does not match [steps, hidden]"); + } + const int64_t dim = config.head_dim; + const int64_t kv_heads = config.num_key_value_heads; - ggml_init_params params{impl_->graph_arena_bytes, nullptr, true}; + ggml_init_params params{impl.graph_arena_bytes, nullptr, true}; std::unique_ptr graph_ctx(ggml_init(params)); if (graph_ctx == nullptr) { - throw std::runtime_error("failed to initialize MOSS-TTS-Local backbone graph context"); + throw std::runtime_error("failed to initialize MOSS-TTS-Local backbone prefill context"); } - core::ModuleBuildContext ctx{graph_ctx.get(), "moss_tts_local.backbone.forward", impl_->backend_type}; + core::ModuleBuildContext ctx{graph_ctx.get(), "moss_tts_local.backbone.prefill", impl.backend_type}; auto token_input = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, steps})); ggml_set_input(token_input.tensor); @@ -349,47 +584,62 @@ std::vector MossBackboneRuntime::run_prefill( auto attention_mask = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, steps, steps})); ggml_set_input(attention_mask.tensor); - core::TensorValue audio_bias_input; - if (audio_bias != nullptr) { - audio_bias_input = - core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, steps, config.hidden_size})); - ggml_set_input(audio_bias_input.tensor); - } + auto bias_input = + core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, steps, config.hidden_size})); + ggml_set_input(bias_input.tensor); auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size}) - .build(ctx, token_input, weights.embed_tokens); - if (audio_bias != nullptr) { - x = modules::AddModule{}.build(ctx, x, audio_bias_input); - } - for (const auto & layer : weights.layers) { - x = decoder_layer(ctx, x, positions, layer, config, attention_mask); + .build(ctx, token_input, impl.weights.embed_tokens); + x = modules::AddModule{}.build(ctx, x, bias_input); + std::vector layer_keys; + std::vector layer_values; + layer_keys.reserve(impl.weights.layers.size()); + layer_values.reserve(impl.weights.layers.size()); + for (const auto & layer : impl.weights.layers) { + auto out = decoder_layer_prefill(ctx, x, positions, layer, config, attention_mask); + x = out.hidden; + layer_keys.push_back(out.key); + layer_values.push_back(out.value); } x = modules::RMSNormModule({config.hidden_size, config.rms_norm_eps, true, false}) - .build(ctx, x, binding::norm_data(ctx, weights.norm)); + .build(ctx, x, binding::norm_data(ctx, impl.weights.norm)); x = ensure_contiguous(ctx, x); auto hidden = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({steps, config.hidden_size})); hidden = ensure_contiguous(ctx, hidden); ggml_set_output(hidden.tensor); + // Copy each layer's K/V straight into the generation cache rows [0, steps) as part of the + // graph. Doing it on-device (instead of reading the intermediates back to the host) is both + // faster and immune to the graph allocator reusing those tensors' storage after compute. ggml_cgraph * graph = ggml_new_graph_custom(graph_ctx.get(), 65536, false); ggml_build_forward_expand(graph, hidden.tensor); + for (size_t layer = 0; layer < impl.cache_keys.size(); ++layer) { + ggml_tensor * cache_key = impl.cache_keys[layer]; + ggml_tensor * cache_value = impl.cache_values[layer]; + ggml_tensor * key_view = ggml_view_4d( + graph_ctx.get(), cache_key, dim, kv_heads, steps, 1, + cache_key->nb[1], cache_key->nb[2], cache_key->nb[3], 0); + ggml_tensor * value_view = ggml_view_4d( + graph_ctx.get(), cache_value, dim, kv_heads, steps, 1, + cache_value->nb[1], cache_value->nb[2], cache_value->nb[3], 0); + ggml_build_forward_expand(graph, ggml_cpy(graph_ctx.get(), layer_keys[layer].tensor, key_view)); + ggml_build_forward_expand(graph, ggml_cpy(graph_ctx.get(), layer_values[layer].tensor, value_view)); + } - ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(impl_->backend)); + ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(impl.backend)); if (gallocr == nullptr || !ggml_gallocr_reserve(gallocr, graph) || !ggml_gallocr_alloc_graph(gallocr, graph)) { if (gallocr != nullptr) { ggml_gallocr_free(gallocr); } - throw std::runtime_error("failed to allocate MOSS-TTS-Local backbone forward graph"); + throw std::runtime_error("failed to allocate MOSS-TTS-Local backbone prefill graph"); } ggml_backend_tensor_set(token_input.tensor, token_ids.data(), 0, token_ids.size() * sizeof(int32_t)); - std::vector position_host(static_cast(steps)); for (int64_t i = 0; i < steps; ++i) { position_host[static_cast(i)] = static_cast(i); } ggml_backend_tensor_set(positions.tensor, position_host.data(), 0, position_host.size() * sizeof(int32_t)); - std::vector mask_host(static_cast(steps * steps), kMaskedAttentionBias); for (int64_t q = 0; q < steps; ++q) { for (int64_t k = 0; k <= q; ++k) { @@ -397,26 +647,28 @@ std::vector MossBackboneRuntime::run_prefill( } } ggml_backend_tensor_set(attention_mask.tensor, mask_host.data(), 0, mask_host.size() * sizeof(float)); + ggml_backend_tensor_set(bias_input.tensor, audio_bias.data(), 0, audio_bias.size() * sizeof(float)); - if (audio_bias != nullptr) { - ggml_backend_tensor_set( - audio_bias_input.tensor, - audio_bias, - 0, - static_cast(steps * config.hidden_size) * sizeof(float)); - } - - const ggml_status status = ggml_backend_graph_compute(impl_->backend, graph); - ggml_backend_synchronize(impl_->backend); + const ggml_status status = ggml_backend_graph_compute(impl.backend, graph); + ggml_backend_synchronize(impl.backend); if (status != GGML_STATUS_SUCCESS) { ggml_gallocr_free(gallocr); - throw std::runtime_error("MOSS-TTS-Local backbone forward graph compute failed"); + throw std::runtime_error("MOSS-TTS-Local backbone prefill graph compute failed"); } - std::vector hidden_states(static_cast(steps * config.hidden_size)); - ggml_backend_tensor_get(hidden.tensor, hidden_states.data(), 0, hidden_states.size() * sizeof(float)); + std::vector last_hidden(static_cast(config.hidden_size)); + ggml_backend_tensor_get( + hidden.tensor, + last_hidden.data(), + static_cast((steps - 1) * config.hidden_size) * sizeof(float), + last_hidden.size() * sizeof(float)); ggml_gallocr_free(gallocr); - return hidden_states; + impl.valid_steps = steps; + return last_hidden; +} + +int64_t MossBackboneRuntime::cached_positions() const noexcept { + return impl_->valid_steps; } } // namespace engine::models::moss_tts_local diff --git a/src/models/moss_tts_local/generator.cpp b/src/models/moss_tts_local/generator.cpp index 0f20378..afc5b71 100644 --- a/src/models/moss_tts_local/generator.cpp +++ b/src/models/moss_tts_local/generator.cpp @@ -175,41 +175,44 @@ std::vector> MossGenerator::generate( throw std::runtime_error("MOSS-TTS-Local generation audio codes must be [seq, n_vq]"); } - std::vector sequence_text = text_tokens; - std::vector sequence_audio = audio_codes; std::vector> generated_frames; std::vector> code_history(static_cast(n_vq)); std::mt19937 rng(options.seed); - for (int64_t frame = 0; frame < options.max_new_frames; ++frame) { - const int64_t steps = static_cast(sequence_text.size()); - - // Fuse the summed audio-codebook embeddings into an additive bias so the backbone - // only has to embed the text channel in the graph. - std::vector audio_bias(static_cast(steps * hidden), 0.0F); - for (int64_t position = 0; position < steps; ++position) { - float * bias_row = audio_bias.data() + static_cast(position) * hidden; - for (int64_t codebook = 0; codebook < n_vq; ++codebook) { - const int32_t code = sequence_audio[static_cast(position * n_vq + codebook)]; - if (code == pad_code) { - continue; - } - const float * embedding = - audio_embeddings_[static_cast(codebook)].data() + static_cast(code) * hidden; - for (int64_t index = 0; index < hidden; ++index) { - bias_row[index] += embedding[index]; - } + // Sums the audio-codebook embeddings for one decoder row into the additive bias the + // backbone expects (padding codes contribute nothing). + const auto bias_for = [&](const int32_t * codes) { + std::vector bias(static_cast(hidden), 0.0F); + for (int64_t codebook = 0; codebook < n_vq; ++codebook) { + const int32_t code = codes[codebook]; + if (code == pad_code) { + continue; + } + const float * embedding = + audio_embeddings_[static_cast(codebook)].data() + static_cast(code) * hidden; + for (int64_t index = 0; index < hidden; ++index) { + bias[static_cast(index)] += embedding[index]; } } + return bias; + }; - const std::vector backbone_hidden = backbone_.forward_prefill_fused(sequence_text, audio_bias); - std::vector local_hidden( - backbone_hidden.end() - hidden, - backbone_hidden.end()); + // Prefill the whole prompt in a single batched forward (fills the cache in one pass + // instead of one graph launch per token), then generate one row at a time from the cache. + // The last prompt position's hidden state seeds the first generated frame. + const int64_t prefix_len = static_cast(text_tokens.size()); + backbone_.begin_generation(prefix_len + options.max_new_frames); + std::vector prefix_bias(static_cast(prefix_len * hidden), 0.0F); + for (int64_t position = 0; position < prefix_len; ++position) { + const std::vector row = bias_for(audio_codes.data() + position * n_vq); + std::copy(row.begin(), row.end(), prefix_bias.begin() + position * hidden); + } + std::vector last_hidden = backbone_.prefill(text_tokens, prefix_bias); + for (int64_t frame = 0; frame < options.max_new_frames; ++frame) { // Seed the depth transformer with the backbone hidden state (local position 0). - std::vector local_embeds = local_hidden; - local_hidden = depth_.forward(local_embeds, 1); + std::vector local_embeds = last_hidden; + std::vector local_hidden = depth_.forward(local_embeds, 1); // Binary gate: continue with the assistant slot or stop on the audio-end token. std::vector gate_logits = project(local_text_head_, local_hidden, 2, hidden); @@ -243,9 +246,9 @@ std::vector> MossGenerator::generate( } generated_frames.push_back(frame_codes); - // Append the emitted frame as the next decoder row (assistant slot + codes). - sequence_text.push_back(assistant_slot); - sequence_audio.insert(sequence_audio.end(), frame_codes.begin(), frame_codes.end()); + // Append the emitted frame as the next decoder row (assistant slot + codes) and + // advance the backbone cache; the returned hidden seeds the next frame. + last_hidden = backbone_.step(assistant_slot, bias_for(frame_codes.data())); } return generated_frames; diff --git a/src/models/moss_tts_local/session.cpp b/src/models/moss_tts_local/session.cpp index 917c91f..6930f74 100644 --- a/src/models/moss_tts_local/session.cpp +++ b/src/models/moss_tts_local/session.cpp @@ -1,10 +1,13 @@ #include "engine/models/moss_tts_local/session.h" +#include "engine/framework/assets/tensor_source.h" #include "engine/framework/audio/resampling.h" +#include "engine/framework/core/backend.h" #include "engine/framework/runtime/options.h" #include "engine/models/moss_tts_local/assets.h" #include +#include #include #include #include @@ -77,6 +80,42 @@ std::vector> reference_to_codec_stereo(const runtime::AudioBu return stereo; } +// Hardware-adaptive backbone dtype for the "auto" mode: GPUs run the model's +// native bf16, while CPU uses f32 (bf16/f16 matmul is poorly accelerated on CPU, +// so f32 is both faster and more accurate there). Non-CUDA GPU backends keep the +// stored (native) dtype to stay on the known-good path. +engine::assets::TensorStorageType resolve_auto_weight_type(engine::core::BackendType backend) { + switch (backend) { + case engine::core::BackendType::Cpu: + return engine::assets::TensorStorageType::F32; + case engine::core::BackendType::Cuda: + return engine::assets::TensorStorageType::BF16; + default: // Metal, Vulkan, BestAvailable + return engine::assets::TensorStorageType::Native; + } +} + +// Resolves moss_tts_local.weight_type. Defaults to "auto" (hardware-adaptive per +// resolve_auto_weight_type); an explicit value (native/f32/f16/bf16/q8_0/...) +// overrides. Handled here because the generic parser maps "auto" -> Native. +engine::assets::TensorStorageType option_weight_type( + const runtime::SessionOptions & options, + const char * key) { + std::string value = "auto"; + const auto it = options.options.find(key); + if (it != options.options.end() && !it->second.empty()) { + value = it->second; + } + std::string lowered = value; + std::transform(lowered.begin(), lowered.end(), lowered.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + if (lowered == "auto") { + return resolve_auto_weight_type(options.backend.type); + } + return engine::assets::parse_tensor_storage_type(value); +} + std::shared_ptr require_assets(std::shared_ptr assets) { if (assets == nullptr) { throw std::runtime_error("MOSS-TTS-Local session requires assets"); @@ -126,7 +165,7 @@ MossTTSLocalSession::MossTTSLocalSession( execution_context(), kBackboneGraphArenaBytes, kBackboneWeightContextBytes, - assets::TensorStorageType::Native); + option_weight_type(options, "moss_tts_local.weight_type")); depth_ = std::make_unique( assets_, execution_context(), kDepthGraphArenaBytes, kDepthWeightContextBytes); processor_ = std::make_unique(assets_); diff --git a/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json b/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json index 0d2e8c0..5488ae2 100644 --- a/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json @@ -204,6 +204,27 @@ } ] }, + { + "id": "moss_tts_local_voice_clone_longform", + "coverage": "MOSS-TTS-Local voice clone with shared long-form text for single-prefix generation (no text chunking) and RTF measurement", + "family": "moss_tts_local", + "model": "models/MOSS-TTS-Local-Transformer-v1.5", + "task": "tts", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "clone_longform", + "text": "At dawn the harbor station opens its tall windows and the first clerk begins a careful report for the day. She notes the weather above the river, the slow cargo boats beyond the bridge, and the market voices arriving from the eastern road. A brass clock marks each quarter hour while porters stack wooden crates, bakers carry warm bread across the square, and a violinist practices the same bright phrase under the stone archway. By midmorning the keeper of the lighthouse sends a message about shifting currents, the museum guide unlocks a cabinet of maps, and a teacher leads a quiet line of students toward the ferry. In the afternoon a painter describes the silver color of the water, a mechanic jokes with the tram driver, and the station master reads an announcement that asks every traveler to keep close watch over letters, tickets, and parcels. After sunset the same clerk continues the report because new visitors keep arriving from the inland road. She explains that a florist carries pale roses past the fountain, two carpenters compare measurements beside the warehouse door, and the watchman checks each lock before the tide reaches its highest mark. A child laughs when the tram bell rings, a cook lowers a basket of fruit to the cellar, and three sailors unfold a chart that shows old channels, sandbars, and safe turning points for the morning crossing. Near midnight the lamps still glow on wet stone, the last cart rattles toward the market gate, and the report ends by saying that the harbor remains orderly, the wind has softened, the ferries are secure, and the town can rest until the next sunrise returns over the water. On the following morning the clerk resumes the record with even greater care because a week of inspections is about to begin. She writes that a ferry captain checks the mooring ropes one by one, a bookseller arranges travel guides beside the station cafe, and a pair of gardeners lift wet soil into bright clay pots near the west entrance. The bakery sends out trays of seed bread, the telegraph operator copies three official notices, and a tailor unfolds navy cloth across a polished wooden counter while customers wait in a line that bends toward the fountain. Before noon a surveyor compares bridge numbers against an old ledger, two cousins argue cheerfully about the best route to the fish market, and a choir director rehearses a patient scale that echoes against the warehouse wall. The lighthouse keeper reports that the northern channel is calmer than expected, the harbor pilot recommends a slower turn near the sandbar, and the customs officer stamps a packet of forms before waving a cart through the side gate. Later the schoolteacher returns with another group of students, asking them to observe the colors of rope, paint, stone, and water so they can write more exact descriptions in the classroom. A photographer kneels beside a rain barrel to capture the reflection of the clock tower, a mechanic tightens a brass hinge on the tram door, and an elderly traveler asks the clerk whether the evening ferry still stops at the orchard village beyond the marsh. As dusk arrives, lamps are trimmed again, shutters are tested against the wind, and the station kitchen sends bowls of soup to workers who remain on the late shift. The report continues with notes about a carpenter measuring floorboards in the east hall, a florist tying silver ribbon around the last stems of the day, and a violin case resting open on a bench beside the ticket window while its owner copies melody marks into a notebook. Long after the market gate closes, the clerk still writes that the harbor road stays busy, the river glints beneath scattered lamps, and the town maintains its patient rhythm of signals, footsteps, voices, bells, and distant engines. On the third day the clerk decides the record should be more precise, so she marks each event by the quarter hour and notes which sounds carry farthest through the station concourse. At first light she hears broom bristles on the stone steps, kettle lids in the cafe kitchen, and the slow scrape of crates being nudged across a loading cart beside the river wall. A messenger in a green coat delivers two canvas pouches, the ticket agent counts rolled coins into a brass tray, and a mother reads directions aloud while her son traces the painted ferry schedule with one curious finger. Midmorning brings a burst of sunlight across the waiting hall, making every brass handle shine while the museum guide escorts visitors toward the gallery of maps and navigational instruments. A porter pauses to describe the oldest compass in the display, a student sketches the harbor outline in graphite, and an apprentice clockmaker compares the station bell to a pocket watch that once belonged to his grandfather. By noon the fish market sends salt and seaweed scents through the open doors, tram wheels hiss at the curb, and the baker from the square exchanges a laugh with the florist who is carrying fresh lilies to the hotel veranda. The clerk writes that a cooper rolls three narrow barrels toward the cellar ramp, a translator copies weather bulletins for inland travelers, and a painter in a blue scarf studies the changing color of the tide as if each small wave might explain a different part of the sky. In the late afternoon the station master reviews freight tags, the customs officer checks a parcel of glassware, and a choir of children crosses the square singing a phrase so soft that the watchman removes his cap to listen. Evening settles slowly; lamps brighten in sequence, a cook inventories apples and onions in the pantry, and two sailors spread a faded chart on a crate so they can debate whether the shoals have shifted since the previous autumn. Before sleep the clerk closes the day with a final note that every vessel is accounted for, every platform has been swept, every lock has been tested twice, and the harbor seems ready to welcome another tide, another market, and another patient stream of voices at sunrise.", + "voice_ref": "resources/a.wav", + "reference_text": "This little work was finished in the year eighteen o three, and intended for immediate publication.", + "seed": 1234, + "do_sample": true + } + ] + }, { "id": "higgs_tts_voice_clone_longform", "coverage": "Higgs Audio v3 voice clone with framework long-form text chunking, bounded AR generation, and codec decode", diff --git a/tools/audiocpp_cli/audiocpp_cli_path_cases.json b/tools/audiocpp_cli/audiocpp_cli_path_cases.json index 67c283f..2572697 100644 --- a/tools/audiocpp_cli/audiocpp_cli_path_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_path_cases.json @@ -1255,6 +1255,111 @@ } ] }, + { + "id": "moss_tts_local_text_only_greedy", + "coverage": "MOSS-TTS-Local text-only generation: Qwen3 backbone prefill, GPT-2 depth transformer, binary audio-end gate, 12-codebook sampling, MOSS-Audio-Tokenizer-v2 decode to 48 kHz stereo", + "family": "moss_tts_local", + "model": "models/MOSS-TTS-Local-Transformer-v1.5", + "task": "tts", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "text_only_greedy", + "text": "The studio lights clicked on before the first phrase began, and the narrator kept a steady, natural tone.", + "seed": 1234, + "max_tokens": 256, + "do_sample": false + } + ] + }, + { + "id": "moss_tts_local_voice_clone_greedy", + "coverage": "MOSS-TTS-Local voice clone: reference resample and loudness normalization, MOSS-Audio-Tokenizer-v2 encode to codes, clone-prefix build, generation, codec decode", + "family": "moss_tts_local", + "model": "models/MOSS-TTS-Local-Transformer-v1.5", + "task": "tts", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "clone_greedy", + "text": "The cloned voice should stay steady and clear through this medium length request.", + "voice_ref": "resources/a.wav", + "reference_text": "This little work was finished in the year eighteen o three, and intended for immediate publication.", + "seed": 1234, + "max_tokens": 256, + "do_sample": false + } + ] + }, + { + "id": "moss_tts_local_sampled_language", + "coverage": "MOSS-TTS-Local sampled generation path: temperature, top-k, and top-p sampling with audio repetition penalty and an explicit language template slot", + "family": "moss_tts_local", + "model": "models/MOSS-TTS-Local-Transformer-v1.5", + "task": "tts", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "sampled_en", + "text": "This sampled request exercises temperature, top-k, top-p, and repetition penalty together.", + "language": "English", + "seed": 1234, + "max_tokens": 256, + "do_sample": true, + "temperature": 0.9, + "top_k": 50, + "top_p": 0.95, + "repetition_penalty": 1.1 + } + ] + }, + { + "id": "moss_tts_local_long_lived_session", + "coverage": "MOSS-TTS-Local long-lived session: long text-only, short clone, then long clone requests in one session to exercise backbone/depth/codec graph reuse, lazy encoder build, and stable memory after warmup", + "family": "moss_tts_local", + "model": "models/MOSS-TTS-Local-Transformer-v1.5", + "task": "tts", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "long_text_only", + "text": "The harbor clerk opened the tall windows and began a careful morning report about the tide, the lantern, and the ships crossing the horizon.", + "seed": 1234, + "max_tokens": 256, + "do_sample": false + }, + { + "id": "short_clone", + "text": "The shorter line should keep the same cloned speaker.", + "voice_ref": "resources/a.wav", + "reference_text": "This little work was finished in the year eighteen o three, and intended for immediate publication.", + "seed": 1234, + "max_tokens": 128, + "do_sample": false + }, + { + "id": "long_clone_again", + "text": "After the short request, the same cloned voice describes the quiet station again in a longer passage so graph reuse and memory behavior stay visible.", + "voice_ref": "resources/a.wav", + "reference_text": "This little work was finished in the year eighteen o three, and intended for immediate publication.", + "seed": 1234, + "max_tokens": 256, + "do_sample": false + } + ] + }, { "id": "vibevoice_two_speaker_long", "coverage": "VibeVoice 1.5B short two-speaker TTS smoke path with reference voices", diff --git a/tools/audiocpp_cli/moss_tts_local_codec_parity.py b/tools/audiocpp_cli/moss_tts_local_codec_parity.py new file mode 100644 index 0000000..8f8dd9a --- /dev/null +++ b/tools/audiocpp_cli/moss_tts_local_codec_parity.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python3 +"""Teacher-forced codec-decode parity: C++ vs Python reference. + +The C++ target ``codec_decode_parity`` decodes a FIXED, deterministic code +matrix (codes[q][t] = (q*37 + t*5) % 1024) to a 48 kHz stereo WAV. This script +decodes the *identical* codes through the Python MOSS-Audio-Tokenizer-v2 in fp32 +and reports the cosine / max-abs difference against the C++ WAV. + +Unlike end-to-end audio cosine (meaningless for a free-running AR model, whose +greedy rollout diverges), this is deterministic and teacher-forced: identical +codes in, so it directly measures codec-decoder numerical parity. + +Run in the ``moss_tts_local`` conda env. Example: + # C++ side first: + build/windows-cuda-release/bin/codec_decode_parity.exe --frames 64 --out cpp.wav + # then: + python moss_tts_local_codec_parity.py --frames 64 --cpp-wav cpp.wav +""" + +from __future__ import annotations + +import argparse +import wave +from pathlib import Path + +import numpy as np + +CODEC_REPO = "OpenMOSS-Team/MOSS-Audio-Tokenizer-v2" + + +def read_wav(path: Path) -> tuple[int, np.ndarray]: + with wave.open(str(path), "rb") as h: + ch, sr, sw, n = h.getnchannels(), h.getframerate(), h.getsampwidth(), h.getnframes() + raw = h.readframes(n) + assert sw == 2, f"expected pcm16, got sample width {sw}" + audio = np.frombuffer(raw, dtype=" float: + a, b = a.reshape(-1), b.reshape(-1) + n = min(a.size, b.size) + a, b = a[:n], b[:n] + denom = np.linalg.norm(a) * np.linalg.norm(b) + return 1.0 if denom == 0 else float(np.dot(a, b) / denom) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--codec", default=CODEC_REPO, help="codec repo id or local snapshot path") + parser.add_argument("--frames", type=int, default=12) + parser.add_argument("--num-quantizers", type=int, default=12) + parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu") + parser.add_argument("--out", type=Path, default=Path("moss_codec_decode_py.wav")) + parser.add_argument("--cpp-wav", type=Path, help="C++ codec_decode_parity WAV to compare against") + args = parser.parse_args() + + import torch + import torchaudio + from transformers import AutoModel + + codec = AutoModel.from_pretrained(str(args.codec), trust_remote_code=True) + codec = codec.to(device=args.device, dtype=torch.float32) + if hasattr(codec, "set_compute_dtype"): + codec.set_compute_dtype("fp32") + codec.eval() + + nq, frames = args.num_quantizers, args.frames + codes = torch.tensor( + [[(q * 37 + t * 5) % 1024 for t in range(frames)] for q in range(nq)], + dtype=torch.long, + device=args.device, + ) # [nq, frames], identical to the C++ target + + with torch.no_grad(): + out = codec.decode(codes, num_quantizers=nq) + audio = getattr(out, "audio", out[0] if isinstance(out, tuple) else out) + audio = audio.squeeze().to(torch.float32).cpu() + if audio.dim() == 1: + audio = audio.unsqueeze(0) + torchaudio.save(str(args.out), audio, 48000) + print(f"[PY] decoded codes[{nq},{frames}] -> {tuple(audio.shape)} @48000, wrote {args.out}") + + if args.cpp_wav: + py_sr, py = read_wav(args.out) + cpp_sr, cpp = read_wav(args.cpp_wav) + n = min(py.shape[0], cpp.shape[0]) + c = cosine(cpp[:n], py[:n]) + max_abs = float(np.max(np.abs(cpp[:n].reshape(-1) - py[:n].reshape(-1)))) + print(f"\n=== codec-decode parity (teacher-forced, fp32) ===") + print(f"cpp={args.cpp_wav} ({cpp_sr} Hz, {cpp.shape}), py={args.out} ({py_sr} Hz, {py.shape})") + print(f"cosine={c:.9f} max_abs_diff={max_abs:.6e} (16-bit WAV LSB ~= {1/32768:.2e})") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/audiocpp_cli/moss_tts_local_reference.py b/tools/audiocpp_cli/moss_tts_local_reference.py new file mode 100644 index 0000000..fc97cd4 --- /dev/null +++ b/tools/audiocpp_cli/moss_tts_local_reference.py @@ -0,0 +1,296 @@ +#!/usr/bin/env python3 +"""Run MOSS-TTS-Local test cases through the Python reference (transformers). + +Produces, for every case, a result directory that mirrors the ``audiocpp_cli`` +path-test layout so ``compare_audiocpp_cli_path_results.py`` and +``moss_tts_local_report.py`` can diff C++ vs Python symmetrically: + + //command.json (marks the dir as a case) + //outputs/.wav + //stdout.log (request_id= + [TIMING] lines) + //memory.json (per-request RSS / CUDA VRAM) + +Each case loads the model + codec once and runs its requests in sequence, so a +multi-request case is a genuine long-lived session (mirroring the C++ session). + +Run inside the ``moss_tts_local`` conda env (transformers>=5.0, torch, torchaudio). + +Examples +-------- + # fp32 parity reference (matches the C++ f32 path; no KV cache for exactness) + python moss_tts_local_reference.py --only moss_tts_local_text_only_greedy \ + --dtype fp32 --device cuda --out-root build/logs/moss_ref_fp32 + + # bf16 CUDA reference (for realistic performance comparison) + python moss_tts_local_reference.py --family moss_tts_local \ + --dtype bf16 --device cuda --out-root build/logs/moss_ref_bf16 +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path +from typing import Any, Optional + +REPO_ROOT = Path(__file__).resolve().parents[2] +DEFAULT_CASES = REPO_ROOT / "tools" / "audiocpp_cli" / "audiocpp_cli_path_cases.json" +DEFAULT_MODEL = REPO_ROOT / "models" / "MOSS-TTS-Local-Transformer-v1.5" +FAMILY = "moss_tts_local" + + +def load_cases(path: Path) -> list[dict[str, Any]]: + catalog = json.loads(path.read_text(encoding="utf-8")) + return catalog.get("cases", []) + + +def select(cases: list[dict[str, Any]], only: set[str]) -> list[dict[str, Any]]: + out = [] + for case in cases: + if case.get("family") != FAMILY: + continue + if only and case["id"] not in only: + continue + out.append(case) + if only: + missing = sorted(only - {c["id"] for c in out}) + if missing: + raise SystemExit(f"unknown/non-moss case id(s): {', '.join(missing)}") + return out + + +def resolve_resource(value: str) -> str: + path = Path(value) + if path.is_absolute(): + return str(path) + return str(REPO_ROOT / path) + + +def opt_float(req: dict[str, Any], key: str) -> Optional[float]: + return float(req[key]) if key in req and req[key] is not None else None + + +def opt_int(req: dict[str, Any], key: str) -> Optional[int]: + return int(req[key]) if key in req and req[key] is not None else None + + +def opt_bool(req: dict[str, Any], key: str, default: bool) -> bool: + if key not in req or req[key] is None: + return default + value = req[key] + if isinstance(value, bool): + return value + return str(value).strip().lower() in {"1", "true", "yes", "on"} + + +class MemoryProbe: + """Per-request process RSS and CUDA VRAM snapshots.""" + + def __init__(self, torch_mod: Any, device: str) -> None: + self.torch = torch_mod + self.device = device + try: + import psutil # noqa: WPS433 - optional dependency + + self.proc = psutil.Process() + except Exception: # noqa: BLE001 - psutil is optional + self.proc = None + + def reset_peak(self) -> None: + if self.device == "cuda" and self.torch.cuda.is_available(): + self.torch.cuda.reset_peak_memory_stats() + + def snapshot(self) -> dict[str, float]: + out: dict[str, float] = {} + if self.proc is not None: + out["rss_mb"] = self.proc.memory_info().rss / (1024.0 * 1024.0) + if self.device == "cuda" and self.torch.cuda.is_available(): + out["cuda_alloc_mb"] = self.torch.cuda.memory_allocated() / (1024.0 * 1024.0) + out["cuda_peak_mb"] = self.torch.cuda.max_memory_allocated() / (1024.0 * 1024.0) + out["cuda_reserved_mb"] = self.torch.cuda.memory_reserved() / (1024.0 * 1024.0) + return out + + +def run_case( + case: dict[str, Any], + processor: Any, + model: Any, + torch_mod: Any, + torchaudio_mod: Any, + device: str, + use_kv_cache: bool, + probe: MemoryProbe, + out_root: Path, +) -> None: + case_dir = out_root / case["id"] + outputs = case_dir / "outputs" + outputs.mkdir(parents=True, exist_ok=True) + (case_dir / "command.json").write_text( + json.dumps({"reference": "moss_tts_local_reference.py", "case": case["id"]}, indent=2) + "\n", + encoding="utf-8", + ) + + sampling_rate = int(processor.model_config.sampling_rate) + stdout_lines: list[str] = [f"family={FAMILY}", f"task={case.get('task', 'tts')}", "mode=offline"] + memory: dict[str, Any] = {"case": case["id"], "dtype": str(model.dtype), "requests": []} + + session_start = time.perf_counter() + for index, req in enumerate(case["requests"]): + request_id = req["id"] + language = req.get("language") + if language in (None, "", "Auto"): + language = None + reference = None + if "voice_ref" in req and req["voice_ref"]: + reference = [resolve_resource(req["voice_ref"])] + + message_kwargs: dict[str, Any] = {"text": req["text"], "language": language} + if reference is not None: + message_kwargs["reference"] = reference + if opt_int(req, "tokens") is not None: + message_kwargs["tokens"] = opt_int(req, "tokens") + conversation = [processor.build_user_message(**message_kwargs)] + + seed = opt_int(req, "seed") + if seed is not None: + torch_mod.manual_seed(seed) + if device == "cuda": + torch_mod.cuda.manual_seed_all(seed) + + gen_kwargs: dict[str, Any] = { + "max_new_tokens": opt_int(req, "max_tokens") or 4096, + "do_sample": opt_bool(req, "do_sample", True), + "use_kv_cache": use_kv_cache, + } + if opt_float(req, "temperature") is not None: + gen_kwargs["audio_temperature"] = opt_float(req, "temperature") + if opt_int(req, "top_k") is not None: + gen_kwargs["audio_top_k"] = opt_int(req, "top_k") + if opt_float(req, "top_p") is not None: + gen_kwargs["audio_top_p"] = opt_float(req, "top_p") + if opt_float(req, "repetition_penalty") is not None: + gen_kwargs["audio_repetition_penalty"] = opt_float(req, "repetition_penalty") + + probe.reset_peak() + run_start = time.perf_counter() + with torch_mod.no_grad(): + batch = processor([conversation], mode="generation") + input_ids = batch["input_ids"].to(model.device) + attention_mask = batch["attention_mask"].to(model.device) + gen_out = model.generate(input_ids=input_ids, attention_mask=attention_mask, **gen_kwargs) + audio = None + for message in processor.decode(gen_out): + if message is None: + continue + audio = message.audio_codes_list[0] + break + if device == "cuda": + torch_mod.cuda.synchronize() + wall_ms = (time.perf_counter() - run_start) * 1000.0 + + if audio is None: + raise RuntimeError(f"{case['id']}/{request_id}: reference produced no audio") + wav_path = outputs / f"{request_id}.wav" + torchaudio_mod.save(str(wav_path), audio.to(torch_mod.float32).cpu(), sampling_rate) + + snap = probe.snapshot() + snap.update({"request_id": request_id, "wall_ms": wall_ms}) + audio_seconds = float(audio.shape[-1]) / float(sampling_rate) + snap["audio_seconds"] = audio_seconds + snap["rtf"] = (wall_ms / 1000.0) / audio_seconds if audio_seconds > 0 else 0.0 + memory["requests"].append(snap) + + stdout_lines.append(f"request_index={index}") + stdout_lines.append(f"request_id={request_id}") + stdout_lines.append(f"[TIMING] request.{request_id}.wall_ms {wall_ms}") + stdout_lines.append(f"audio_out={wav_path}") + print( + f"[REF] {case['id']}/{request_id}: {wall_ms:.1f} ms, " + f"{audio_seconds:.2f} s audio, rtf={snap['rtf']:.3f}", + flush=True, + ) + + session_wall_ms = (time.perf_counter() - session_start) * 1000.0 + stdout_lines.append(f"[TIMING] session.wall_ms {session_wall_ms}") + (case_dir / "stdout.log").write_text("\n".join(stdout_lines) + "\n", encoding="utf-8") + memory["session_wall_ms"] = session_wall_ms + (case_dir / "memory.json").write_text(json.dumps(memory, indent=2) + "\n", encoding="utf-8") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--cases", type=Path, default=DEFAULT_CASES) + parser.add_argument("--model", type=Path, default=DEFAULT_MODEL) + parser.add_argument("--only", action="append", default=[], help="case id(s), comma-separated allowed") + parser.add_argument("--family", action="store_true", help="run all moss_tts_local cases (default if --only omitted)") + parser.add_argument("--dtype", choices=["fp32", "bf16", "fp16"], default="bf16") + parser.add_argument("--device", choices=["cuda", "cpu"], default="cuda") + parser.add_argument( + "--use-kv-cache", + choices=["auto", "true", "false"], + default="auto", + help="auto: off for fp32 (exact parity with C++), on otherwise", + ) + parser.add_argument("--out-root", type=Path, required=True) + args = parser.parse_args() + + try: + import torch + import torchaudio + from transformers import AutoModel, AutoProcessor + except Exception as exc: # noqa: BLE001 - guide the user to the right env + print(f"[FAIL] import error ({exc}). Activate the moss_tts_local conda env.", file=sys.stderr) + return 1 + + only = {item.strip() for raw in args.only for item in raw.split(",") if item.strip()} + cases = select(load_cases(args.cases), only) + if not cases: + print("[FAIL] no matching moss_tts_local cases", file=sys.stderr) + return 1 + + dtype = {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16}[args.dtype] + device = args.device if (args.device != "cuda" or torch.cuda.is_available()) else "cpu" + if args.use_kv_cache == "auto": + use_kv_cache = args.dtype != "fp32" + else: + use_kv_cache = args.use_kv_cache == "true" + attn = "sdpa" if device == "cuda" else "eager" + if device == "cuda" and hasattr(torch.backends.cuda, "enable_cudnn_sdp"): + torch.backends.cuda.enable_cudnn_sdp(False) + + print( + f"[INFO] model={args.model} dtype={args.dtype} device={device} " + f"attn={attn} use_kv_cache={use_kv_cache}", + flush=True, + ) + processor = AutoProcessor.from_pretrained(str(args.model), trust_remote_code=True) + if args.dtype == "fp32": + # The codec loads in bf16 by default; cast its weights to fp32 (and set + # its compute dtype to fp32) so decode stays fp32 for parity with the + # C++ f32 path. Casting params avoids the fp32-input/bf16-weight mismatch. + processor.audio_tokenizer = processor.audio_tokenizer.to(device=device, dtype=torch.float32) + if hasattr(processor.audio_tokenizer, "set_compute_dtype"): + processor.audio_tokenizer.set_compute_dtype("fp32") + else: + processor.audio_tokenizer = processor.audio_tokenizer.to(device) + model = AutoModel.from_pretrained( + str(args.model), + trust_remote_code=True, + attn_implementation=attn, + torch_dtype=dtype, + ).to(device) + model.eval() + + probe = MemoryProbe(torch, device) + args.out_root.mkdir(parents=True, exist_ok=True) + for case in cases: + print(f"[RUN] {case['id']}: {case.get('coverage', '')}", flush=True) + run_case(case, processor, model, torch, torchaudio, device, use_kv_cache, probe, args.out_root) + print(f"[DONE] {len(cases)} case(s), out_root={args.out_root}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/audiocpp_cli/moss_tts_local_report.py b/tools/audiocpp_cli/moss_tts_local_report.py new file mode 100644 index 0000000..bf6e4c7 --- /dev/null +++ b/tools/audiocpp_cli/moss_tts_local_report.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python3 +"""Build a PR-ready MOSS-TTS-Local report from a C++ result dir and a Python +reference result dir. + +Both dirs must follow the ``audiocpp_cli`` path-test layout (per-case subdirs +with ``outputs/.wav`` and a ``stdout.log`` carrying ``[TIMING]`` +lines). The C++ dir comes from ``run_audiocpp_cli_path_tests.py``; the Python +dir from ``moss_tts_local_reference.py``. + +Reuses the similarity metrics (cosine, log-mel) from +``compare_audiocpp_cli_path_results.py`` so the numbers match that tool exactly. + +Usage: + python moss_tts_local_report.py CPP_DIR PY_DIR [--title "..."] [--out report.md] +""" + +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path +from typing import Optional + +import compare_audiocpp_cli_path_results as cmp + +TIMING_RE = re.compile(r"^\[TIMING[^\]]*\]\s+(\S+)\s+([-+0-9.eE]+)\s*$") +LOG_MEL_RE = re.compile(r"log_mel_cos=([0-9.]+|n/a)") +SR_RE = re.compile(r"sr=(\d+)/(\d+)") + + +def case_dirs(root: Path) -> dict[str, Path]: + return {p.name: p for p in sorted(root.iterdir()) if p.is_dir() and (p / "command.json").exists()} + + +def timings(case_dir: Path) -> dict[str, float]: + out: dict[str, float] = {} + log = case_dir / "stdout.log" + if not log.exists(): + return out + for line in log.read_text(encoding="utf-8", errors="replace").splitlines(): + m = TIMING_RE.match(line) + if m: + out[m.group(1)] = float(m.group(2)) + return out + + +def request_ids(case: dict) -> list[str]: + return [r["id"] for r in case.get("requests", [])] + + +def load_case_meta(cases_files: list[Path]) -> dict[str, dict]: + meta: dict[str, dict] = {} + for path in cases_files: + if not path.exists(): + continue + for case in json.loads(path.read_text(encoding="utf-8")).get("cases", []): + meta[case["id"]] = case + return meta + + +def parse_detail(detail: str) -> tuple[Optional[str], Optional[str]]: + log_mel = LOG_MEL_RE.search(detail) + sr = SR_RE.search(detail) + sr_txt = f"{sr.group(1)}/{sr.group(2)}" if sr else "?" + return (log_mel.group(1) if log_mel else None), sr_txt + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("cpp_dir", type=Path) + parser.add_argument("py_dir", type=Path) + parser.add_argument("--title", default="MOSS-TTS-Local: audio.cpp vs Python reference") + parser.add_argument("--cases", type=Path, action="append", default=[]) + parser.add_argument("--out", type=Path) + args = parser.parse_args() + + repo = Path(__file__).resolve().parents[2] + cases_files = args.cases or [ + repo / "tools" / "audiocpp_cli" / "audiocpp_cli_path_cases.json", + repo / "tools" / "audiocpp_cli" / "audiocpp_cli_longform_tts_clone_cases.json", + ] + meta = load_case_meta(cases_files) + + cpp = case_dirs(args.cpp_dir) + py = case_dirs(args.py_dir) + shared = [cid for cid in sorted(set(cpp) & set(py))] + + lines: list[str] = [f"# {args.title}", ""] + lines.append(f"- C++ results: `{args.cpp_dir}`") + lines.append(f"- Python reference: `{args.py_dir}`") + lines.append("") + + sim_rows: list[str] = [] + perf_rows: list[str] = [] + mem_rows: list[str] = [] + + for cid in shared: + cpp_dir, py_dir = cpp[cid], py[cid] + cpp_t, py_t = timings(cpp_dir), timings(py_dir) + py_mem = {} + mem_path = py_dir / "memory.json" + if mem_path.exists(): + py_mem = {r["request_id"]: r for r in json.loads(mem_path.read_text(encoding="utf-8")).get("requests", [])} + + ids = request_ids(meta.get(cid, {})) or sorted( + p.stem for p in (cpp_dir / "outputs").glob("*.wav") + ) + for rid in ids: + cpp_wav = cpp_dir / "outputs" / f"{rid}.wav" + py_wav = py_dir / "outputs" / f"{rid}.wav" + if not cpp_wav.exists() or not py_wav.exists(): + sim_rows.append(f"| {cid} | {rid} | missing wav | | |") + continue + wav_cos, detail = cmp.wav_similarity_detail(cpp_wav, py_wav) + log_mel, sr = parse_detail(detail) + sim_rows.append(f"| {cid} | {rid} | {wav_cos:.6f} | {log_mel} | {sr} |") + + cpp_ms = cpp_t.get(f"request.{rid}.wall_ms") + py_ms = py_t.get(f"request.{rid}.wall_ms") + if cpp_ms and py_ms: + # Fair comparison uses RTF (time per second of audio): greedy can + # stop at different frame counts per backend, so raw ms is + # apples-to-oranges. Duration comes straight from each wav. + cpp_sr, cpp_audio = cmp.read_wav_f32(cpp_wav) + py_sr, py_audio = cmp.read_wav_f32(py_wav) + cpp_s = cpp_audio.shape[0] / cpp_sr if cpp_sr else 0.0 + py_s = py_audio.shape[0] / py_sr if py_sr else 0.0 + cpp_rtf = (cpp_ms / 1000.0) / cpp_s if cpp_s else 0.0 + py_rtf = (py_ms / 1000.0) / py_s if py_s else 0.0 + rtf_speedup = py_rtf / cpp_rtf if cpp_rtf else 0.0 + perf_rows.append( + f"| {cid} | {rid} | {cpp_s:.2f} | {py_s:.2f} | {cpp_ms:.0f} | {py_ms:.0f} | " + f"{cpp_rtf:.3f} | {py_rtf:.3f} | {rtf_speedup:.2f}x |" + ) + + m = py_mem.get(rid) + if m: + mem_rows.append( + f"| {cid} | {rid} | {m.get('rss_mb', 0):.0f} | " + f"{m.get('cuda_alloc_mb', 0):.0f} | {m.get('cuda_peak_mb', 0):.0f} |" + ) + + lines += ["## Similarity (C++ vs Python reference)", "", + "| case | request | wav_cos | log_mel_cos | sr Hz |", + "| --- | --- | ---: | ---: | ---: |", *sim_rows, ""] + lines += ["## Performance (per request, one loaded session; excludes model load)", "", + "_Greedy can stop at different frame counts per backend, so compare RTF " + "(wall seconds per second of generated audio), not raw ms. " + "RTF speedup = Python RTF / C++ RTF (>1 means C++ is faster per audio-second)._", + "", + "| case | request | C++ audio s | Py audio s | C++ ms | Py ms | C++ RTF | Py RTF | RTF speedup |", + "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |", *perf_rows, ""] + if mem_rows: + lines += ["## Python reference memory (per request; stable => no growth after warmup)", "", + "| case | request | RSS MB | CUDA alloc MB | CUDA peak MB |", + "| --- | --- | ---: | ---: | ---: |", *mem_rows, ""] + + report = "\n".join(lines) + "\n" + if args.out: + args.out.write_text(report, encoding="utf-8") + print(f"wrote {args.out}") + else: + print(report) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/audiocpp_cli/moss_tts_local_session_probe.py b/tools/audiocpp_cli/moss_tts_local_session_probe.py new file mode 100644 index 0000000..7efe87e --- /dev/null +++ b/tools/audiocpp_cli/moss_tts_local_session_probe.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python3 +"""Probe C++ audiocpp_cli memory across a long-lived MOSS-TTS-Local session. + +Runs one case (default: the multi-request ``moss_tts_local_long_lived_session``) +through ``audiocpp_cli`` while sampling the process RSS and its per-PID CUDA +memory (via ``nvidia-smi``), tagging each sample to the request that was active +when it was taken. Reports peak/final memory per request so you can show the +session is stable (no growth) after the first request or two. + +Reuses ``run_audiocpp_cli_path_tests`` to build the exact CLI command, so the +run matches the normal path-test harness. + +Usage: + python moss_tts_local_session_probe.py \ + --audiocpp-cli-bin build/windows-cuda-release/bin/audiocpp_cli.exe \ + --backend cuda --out-root build/logs/moss_session_probe +""" + +from __future__ import annotations + +import argparse +import re +import subprocess +import threading +import time +from pathlib import Path +from typing import Optional + +import run_audiocpp_cli_path_tests as runner + +REPO_ROOT = Path(__file__).resolve().parents[2] +TIMING_RE = re.compile(r"^\[TIMING[^\]]*\]\s+(\S+)\s+([-+0-9.eE]+)\s*$") + + +def gpu_used_mb(pid: int, device: int = 0) -> Optional[float]: + """Per-process VRAM if available, else device-total used memory. + + Per-process memory is not exposed under Windows WDDM, so we fall back to the + device total. On a dedicated run that total is dominated by this process, so + the curve shape (growth vs plateau) is still meaningful. + """ + def _nvsmi(query: str) -> str: + try: + return subprocess.run( + ["nvidia-smi", query, "--format=csv,noheader,nounits"], + text=True, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, check=False, + ).stdout + except FileNotFoundError: + return "" + + for line in _nvsmi("--query-compute-apps=pid,used_memory").splitlines(): + parts = [p.strip() for p in line.split(",")] + if len(parts) >= 2 and parts[0].isdigit() and int(parts[0]) == pid: + try: + return float(parts[1]) + except ValueError: + break + # Fallback: device-total used memory (index-th GPU). + rows = _nvsmi("--query-gpu=memory.used").splitlines() + if device < len(rows): + try: + return float(rows[device].strip()) + except ValueError: + return None + return None + + +class Sampler(threading.Thread): + def __init__(self, pid: int, interval: float, device: int = 0) -> None: + super().__init__(daemon=True) + self.pid = pid + self.interval = interval + self.device = device + self.samples: list[dict] = [] + self._stop = threading.Event() + try: + import psutil + + self.proc = psutil.Process(pid) + except Exception: # noqa: BLE001 - psutil optional / process may exit + self.proc = None + + def run(self) -> None: + while not self._stop.is_set(): + rss = None + if self.proc is not None: + try: + rss = self.proc.memory_info().rss / (1024.0 * 1024.0) + except Exception: # noqa: BLE001 - process ended + rss = None + self.samples.append({"t": time.perf_counter(), "rss_mb": rss, "gpu_mb": gpu_used_mb(self.pid, self.device)}) + self._stop.wait(self.interval) + + def stop(self) -> None: + self._stop.set() + + +def reconstruct_windows(stdout_lines: list[str], t0: float, t_end: float) -> list[tuple[str, float, float]]: + """Rebuild per-request time windows from the CLI's own [TIMING] lines. + + audiocpp_cli block-buffers stdout when piped, so real-time request markers + are unreliable. Instead we use the per-request wall_ms it reports: requests + run back-to-back at the tail of the process (after the big model load), so we + lay them out ending at t_end, oldest first. + """ + req: list[tuple[str, float]] = [] + for line in stdout_lines: + m = TIMING_RE.match(line) + if m and m.group(1).startswith("request.") and m.group(1).endswith(".wall_ms"): + rid = m.group(1)[len("request."):-len(".wall_ms")] + req.append((rid, float(m.group(2)) / 1000.0)) + total = sum(d for _, d in req) + windows: list[tuple[str, float, float]] = [] + # End of the last request ~= process end; walk backwards summing durations. + end = t_end + for rid, dur in reversed(req): + windows.append((rid, end - dur, end)) + end -= dur + windows.reverse() + return windows + + +def summarize(samples: list[dict], windows: list[tuple[str, float, float]]) -> list[dict]: + """Assign each memory sample to the request whose window contains it.""" + out: dict[str, dict] = {rid: {"rss": [], "gpu": []} for rid, _, _ in windows} + for s in samples: + for rid, start, end in windows: + if start <= s["t"] < end: + if s["rss_mb"] is not None: + out[rid]["rss"].append(s["rss_mb"]) + if s["gpu_mb"] is not None: + out[rid]["gpu"].append(s["gpu_mb"]) + break + rows = [] + for rid, _, _ in windows: + b = out[rid] + rows.append({ + "request_id": rid, + "rss_peak_mb": max(b["rss"]) if b["rss"] else None, + "rss_final_mb": b["rss"][-1] if b["rss"] else None, + "gpu_peak_mb": max(b["gpu"]) if b["gpu"] else None, + "gpu_final_mb": b["gpu"][-1] if b["gpu"] else None, + }) + return rows + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--cases", type=Path, default=runner.DEFAULT_CASES) + parser.add_argument("--case-id", default="moss_tts_local_long_lived_session") + parser.add_argument("--audiocpp-cli-bin", type=Path, default=runner.DEFAULT_AUDIOCPP_CLI_BIN) + parser.add_argument("--models-root", type=Path, default=runner.DEFAULT_MODELS_ROOT) + parser.add_argument("--backend", default="cuda", choices=["cpu", "cuda", "vulkan", "metal", "best"]) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--threads", type=int, default=1) + parser.add_argument("--session-option", action="append", default=[], help="key=value, e.g. moss_tts_local.weight_type=f32") + parser.add_argument("--interval", type=float, default=0.25) + parser.add_argument("--log", action="store_true", help="pass --log through to audiocpp_cli") + parser.add_argument("--out-root", type=Path, required=True) + args = parser.parse_args() + + if not args.models_root.is_absolute(): + args.models_root = REPO_ROOT / args.models_root + catalog = runner.load_cases(args.cases) + case = next((c for c in catalog.get("cases", []) if c["id"] == args.case_id), None) + if case is None: + raise SystemExit(f"case not found: {args.case_id}") + # Inject any extra session options (e.g. weight_type) without editing the shared JSON. + for kv in args.session_option: + key, _, value = kv.partition("=") + case.setdefault("session_options", {})[key] = value + + out_root = args.out_root + out_root.mkdir(parents=True, exist_ok=True) + case_dir = out_root / case["id"] + case_dir.mkdir(parents=True, exist_ok=True) + command = runner.build_command(args, case, case_dir) + (case_dir / "command.json").write_text(runner.json.dumps(command, indent=2) + "\n", encoding="utf-8") + + print(f"[PROBE] {args.case_id}: {' '.join(map(str, command))}", flush=True) + t0 = time.perf_counter() + proc = subprocess.Popen(command, cwd=REPO_ROOT, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + sampler = Sampler(proc.pid, args.interval, args.device) + sampler.start() + + # NOTE: audiocpp_cli block-buffers piped stdout, so these lines arrive in a + # burst at exit; we correlate memory to requests post-hoc via [TIMING] lines. + stdout_lines: list[str] = [] + assert proc.stdout is not None + for line in proc.stdout: + stdout_lines.append(line.rstrip("\n")) + proc.wait() + t_end = time.perf_counter() + sampler.stop() + sampler.join(timeout=2.0) + stderr = proc.stderr.read() if proc.stderr else "" + + (case_dir / "stdout.log").write_text("\n".join(stdout_lines) + "\n", encoding="utf-8") + (case_dir / "stderr.log").write_text(stderr, encoding="utf-8") + if proc.returncode != 0: + print(f"[FAIL] exit {proc.returncode}; see {case_dir}/stderr.log") + return 1 + + windows = reconstruct_windows(stdout_lines, t0, t_end) + rows = summarize(sampler.samples, windows) + report = {"case": args.case_id, "samples": len(sampler.samples), "per_request": rows} + (case_dir / "memory_cpp.json").write_text(runner.json.dumps(report, indent=2) + "\n", encoding="utf-8") + + print("\nrequest rss_peak rss_final gpu_peak gpu_final (MB)") + for r in rows: + print(f"{r['request_id']:<18} {r['rss_peak_mb'] or 0:8.0f} {r['rss_final_mb'] or 0:9.0f} " + f"{r['gpu_peak_mb'] or 0:8.0f} {r['gpu_final_mb'] or 0:9.0f}") + print(f"\n[DONE] memory report -> {case_dir/'memory_cpp.json'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/audiocpp_cli/run_audiocpp_cli_path_tests.py b/tools/audiocpp_cli/run_audiocpp_cli_path_tests.py index 29f126e..b239c43 100644 --- a/tools/audiocpp_cli/run_audiocpp_cli_path_tests.py +++ b/tools/audiocpp_cli/run_audiocpp_cli_path_tests.py @@ -325,6 +325,8 @@ def build_command(args: argparse.Namespace, case: dict[str, Any], case_dir: Path ] append_key_values(command, "--load-option", case.get("load_options", {})) append_key_values(command, "--session-option", case.get("session_options", {})) + for override in getattr(args, "session_option", []) or []: + command.extend(["--session-option", override]) if args.log: command.append("--log") @@ -452,6 +454,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--resource-sample-ms", type=int, default=DEFAULT_RESOURCE_SAMPLE_MS) parser.add_argument("--out-root", type=Path) parser.add_argument("--only", action="append", default=[], help="Case id or comma-separated case ids") + parser.add_argument( + "--session-option", + action="append", + default=[], + help="Extra key=value session option appended to every case (e.g. moss_tts_local.weight_type=f32)", + ) parser.add_argument("--family", help="Run only cases for one family") parser.add_argument("--log", action="store_true") parser.add_argument("--list", action="store_true")