From 04b9aa8afcccb477a1476ad318d3d153bfca24f9 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Sat, 29 Aug 2026 03:31:28 -0400 Subject: [PATCH 1/3] Promote MOSS codec runtime to framework --- CMakeLists.txt | 8 +- .../community_models/moss_voicegen/session.h | 7 +- .../community_models/vietneu_tts/session.h | 4 +- .../moss_audio_tokenizer_codec_runtime.h | 93 + .../modules/multi_codebook_embedding.h | 40 + .../models/moss/moss_tts_local/generator.h | 3 +- .../models/moss/moss_tts_local/session.h | 10 +- .../moss/moss_tts_local/tokenizer_text.h | 2 +- .../models/moss/moss_tts_nano/session.h | 8 +- .../moss/shared/audio_tokenizer_config.h | 48 - .../moss/shared/audio_tokenizer_decoder.h | 47 - .../moss/shared/audio_tokenizer_encoder.h | 45 - .../moss/shared/audio_tokenizer_quantizer.h | 63 - .../moss/shared/audio_tokenizer_transformer.h | 345 ---- .../engine/models/moss/shared/token_rows.h | 31 - .../moss_voicegen/session.cpp | 31 +- src/community_models/vietneu_tts/session.cpp | 29 +- .../moss_audio_tokenizer_codec_runtime.cpp | 1570 +++++++++++++++++ .../modules/multi_codebook_embedding.cpp | 64 + src/models/moss/moss_tts_local/generator.cpp | 10 +- src/models/moss/moss_tts_local/session.cpp | 39 +- src/models/moss/moss_tts_nano/session.cpp | 49 +- .../moss/shared/audio_tokenizer_config.cpp | 106 -- .../moss/shared/audio_tokenizer_decoder.cpp | 339 ---- .../moss/shared/audio_tokenizer_encoder.cpp | 230 --- .../moss/shared/audio_tokenizer_quantizer.cpp | 273 --- src/models/moss/shared/token_rows.cpp | 55 - tests/moss_tts_local/codec_decode_parity.cpp | 19 +- tests/moss_tts_local/codec_dequant_parity.cpp | 70 - tests/moss_tts_local/codec_encode_parity.cpp | 24 +- tests/moss_voicegen/backbone_parity.cpp | 9 +- tests/moss_voicegen/codec_decode_parity.cpp | 26 +- tests/moss_voicegen/generation_parity.cpp | 9 +- tests/moss_voicegen/voicegen_smoke.cpp | 33 +- .../audiocpp_cli/audiocpp_cli_path_cases.json | 38 +- 35 files changed, 1974 insertions(+), 1803 deletions(-) create mode 100644 include/engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h create mode 100644 include/engine/framework/modules/multi_codebook_embedding.h delete mode 100644 include/engine/models/moss/shared/audio_tokenizer_config.h delete mode 100644 include/engine/models/moss/shared/audio_tokenizer_decoder.h delete mode 100644 include/engine/models/moss/shared/audio_tokenizer_encoder.h delete mode 100644 include/engine/models/moss/shared/audio_tokenizer_quantizer.h delete mode 100644 include/engine/models/moss/shared/audio_tokenizer_transformer.h create mode 100644 src/framework/codecs/moss_audio_tokenizer_codec_runtime.cpp create mode 100644 src/framework/modules/multi_codebook_embedding.cpp delete mode 100644 src/models/moss/shared/audio_tokenizer_config.cpp delete mode 100644 src/models/moss/shared/audio_tokenizer_decoder.cpp delete mode 100644 src/models/moss/shared/audio_tokenizer_encoder.cpp delete mode 100644 src/models/moss/shared/audio_tokenizer_quantizer.cpp delete mode 100644 tests/moss_tts_local/codec_dequant_parity.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 37b153028..91060e17a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -394,6 +394,7 @@ add_library(engine_core OBJECT src/framework/modules/activation_modules.cpp src/framework/modules/norm_modules.cpp src/framework/modules/lookup_modules.cpp + src/framework/modules/multi_codebook_embedding.cpp src/framework/modules/structural_modules.cpp src/framework/modules/asr_helpers.cpp src/framework/modules/conv_modules.cpp @@ -472,6 +473,7 @@ add_library(engine_core OBJECT src/framework/codecs/fish_dac_codec_runtime.cpp src/framework/codecs/mel_latent_vae44k_runtime.cpp src/framework/codecs/mimi_codec_runtime.cpp + src/framework/codecs/moss_audio_tokenizer_codec_runtime.cpp src/framework/codecs/neural_audio.cpp src/framework/codecs/redae_codec_runtime.cpp src/framework/conditioners/clap_audio_conditioner_runtime.cpp @@ -799,10 +801,6 @@ audiocpp_add_model(moss src/models/moss/moss_tts_nano/prompt_builder.cpp src/models/moss/moss_tts_nano/session.cpp src/models/moss/moss_tts_nano/tokenizer_text.cpp - src/models/moss/shared/audio_tokenizer_decoder.cpp - src/models/moss/shared/audio_tokenizer_encoder.cpp - src/models/moss/shared/audio_tokenizer_config.cpp - src/models/moss/shared/audio_tokenizer_quantizer.cpp src/models/moss/shared/sampling.cpp src/models/moss/shared/token_rows.cpp src/models/moss/moss_tts_local/depth_transformer.cpp @@ -1195,8 +1193,6 @@ audiocpp_add_model(vietneu_tts engine/community_models/vietneu_tts/loader.h LOADERS engine::models::vietneu_tts::make_vietneu_tts_loader - DEPENDS - moss ) audiocpp_add_model(qwen3_asr diff --git a/include/engine/community_models/moss_voicegen/session.h b/include/engine/community_models/moss_voicegen/session.h index b38a07313..8c114a2e9 100644 --- a/include/engine/community_models/moss_voicegen/session.h +++ b/include/engine/community_models/moss_voicegen/session.h @@ -6,8 +6,9 @@ #include "engine/community_models/moss_voicegen/heads.h" #include "engine/community_models/moss_voicegen/tokenizer_text.h" #include "engine/framework/core/execution_context.h" +#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" +#include "engine/framework/modules/multi_codebook_embedding.h" #include "engine/framework/runtime/session_base.h" -#include "engine/models/moss/shared/audio_tokenizer_decoder.h" #include "engine/models/moss/shared/token_rows.h" #include @@ -63,10 +64,10 @@ class MossVoiceGenSession final // The execution context comes from RuntimeSessionBase; the runtimes below borrow it. std::unique_ptr text_processor_; - std::unique_ptr codebooks_; + std::unique_ptr codebooks_; std::unique_ptr backbone_; std::unique_ptr heads_; - std::unique_ptr codec_; + std::unique_ptr codec_; }; } // namespace engine::models::moss_voicegen diff --git a/include/engine/community_models/vietneu_tts/session.h b/include/engine/community_models/vietneu_tts/session.h index e6ff977e3..3c61de2d3 100644 --- a/include/engine/community_models/vietneu_tts/session.h +++ b/include/engine/community_models/vietneu_tts/session.h @@ -10,7 +10,7 @@ #include "engine/community_models/vietneu_tts/tokenizer_speech_encoder.h" #include "engine/community_models/vietneu_tts/tokenizer_text.h" -#include "engine/models/moss/shared/audio_tokenizer_decoder.h" +#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" #include #include @@ -80,7 +80,7 @@ class VietneuTTSSession final std::shared_ptr talker_weights_; std::shared_ptr talker_step_; core::ExecutionContext voice_prompt_context_; - std::unique_ptr moss_speech_decoder_; + std::unique_ptr moss_speech_decoder_; std::unique_ptr speech_encoder_; std::unique_ptr speaker_encoder_; runtime::CacheSlots voice_prompt_cache_; diff --git a/include/engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h b/include/engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h new file mode 100644 index 000000000..1fa685ea0 --- /dev/null +++ b/include/engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h @@ -0,0 +1,93 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" + +#include +#include +#include +#include + +namespace engine::codecs { + +struct MossAudioTokenizerTransformerStage { + int64_t input_dimension = 0; + int64_t output_dimension = 0; + int64_t model_dimension = 0; + int64_t num_heads = 0; + int64_t num_layers = 0; + int64_t feedforward_dimension = 0; + int64_t context_window = 0; + int64_t patch_size = 0; +}; + +struct MossAudioTokenizerQuantizerConfig { + int64_t codebook_size = 1024; + int64_t codebook_dim = 8; + int64_t rvq_dim = 512; + int64_t code_dim = 768; + int64_t num_quantizers = 12; +}; + +struct MossAudioTokenizerConfig { + int64_t sampling_rate = 48000; + int64_t samples_per_frame = 3840; + int64_t channels = 2; + MossAudioTokenizerQuantizerConfig quantizer; + std::vector encoder_stages; + std::vector decoder_stages; + int64_t encoder_final_patch = 1; + int64_t decoder_initial_patch = 1; + int64_t encoder_module_start = 1; + int64_t encoder_module_stride = 2; + int64_t decoder_module_start = 0; + int64_t decoder_module_stride = 2; +}; + +MossAudioTokenizerConfig moss_audio_tokenizer_v1_config(); +MossAudioTokenizerConfig moss_audio_tokenizer_v2_config(); +MossAudioTokenizerConfig moss_audio_tokenizer_nano_config(); + +struct MossAudioTokenizerAudio { + int64_t sampling_rate = 0; + std::vector> channels; +}; + +struct MossAudioTokenizerCodes { + int64_t frames = 0; + std::vector> codebooks; +}; + +struct MossAudioTokenizerCodecRuntimeOptions { + size_t weight_context_bytes = 256ull * 1024ull * 1024ull; + size_t encoder_graph_arena_bytes = 2048ull * 1024ull * 1024ull; + size_t decoder_graph_arena_bytes = 1536ull * 1024ull * 1024ull; + bool separate_encoder_context = false; +}; + +class MossAudioTokenizerCodecRuntime { +public: + MossAudioTokenizerCodecRuntime( + std::shared_ptr source, + core::ExecutionContext & execution_context, + int64_t num_quantizers, + MossAudioTokenizerCodecRuntimeOptions options, + MossAudioTokenizerConfig config = moss_audio_tokenizer_v2_config()); + ~MossAudioTokenizerCodecRuntime(); + + MossAudioTokenizerCodecRuntime(const MossAudioTokenizerCodecRuntime &) = delete; + MossAudioTokenizerCodecRuntime & operator=(const MossAudioTokenizerCodecRuntime &) = delete; + + int64_t sampling_rate() const noexcept; + void prepare_encoder(); + void prepare_decoder(); + MossAudioTokenizerCodes encode(const MossAudioTokenizerAudio & audio); + MossAudioTokenizerAudio decode(const MossAudioTokenizerCodes & codes); + void release_runtime_graphs(); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::codecs diff --git a/include/engine/framework/modules/multi_codebook_embedding.h b/include/engine/framework/modules/multi_codebook_embedding.h new file mode 100644 index 000000000..6865b3e97 --- /dev/null +++ b/include/engine/framework/modules/multi_codebook_embedding.h @@ -0,0 +1,40 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" + +#include +#include +#include +#include + +namespace engine::modules { + +struct MultiCodebookEmbeddingSpec { + int64_t hidden_size = 0; + int64_t num_codebooks = 0; + int64_t vocab_size = 0; + int64_t pad_token_id = 0; + std::vector codebook_sizes; + std::string tensor_prefix = "audio_embeddings"; +}; + +class MultiCodebookEmbedding { +public: + MultiCodebookEmbedding(const assets::TensorSource & source, MultiCodebookEmbeddingSpec spec); + + int64_t hidden_size() const noexcept { return hidden_size_; } + int64_t num_codebooks() const noexcept { return num_codebooks_; } + int32_t pad_token_id() const noexcept { return pad_token_id_; } + int64_t codebook_size(int64_t codebook) const; + const float * embedding(int64_t codebook, int32_t code) const; + void add_bias(const int32_t * codes, float * bias) const; + std::vector bias_for(const int32_t * codes) const; + +private: + int64_t hidden_size_ = 0; + int64_t num_codebooks_ = 0; + int32_t pad_token_id_ = 0; + std::vector> embeddings_; +}; + +} // namespace engine::modules diff --git a/include/engine/models/moss/moss_tts_local/generator.h b/include/engine/models/moss/moss_tts_local/generator.h index a10b8fdb2..03b120d23 100644 --- a/include/engine/models/moss/moss_tts_local/generator.h +++ b/include/engine/models/moss/moss_tts_local/generator.h @@ -3,6 +3,7 @@ #include "engine/models/moss/moss_tts_local/assets.h" #include "engine/models/moss/moss_tts_local/backbone.h" #include "engine/models/moss/moss_tts_local/depth_transformer.h" +#include "engine/framework/modules/multi_codebook_embedding.h" #include "engine/models/moss/shared/token_rows.h" #include "engine/framework/sampling/torch_random.h" @@ -65,7 +66,7 @@ class MossGenerator { const MossDepthTransformer & depth_; int64_t hidden_size_ = 0; int64_t num_codebooks_ = 0; - std::unique_ptr audio_codebooks_; + std::unique_ptr audio_codebooks_; std::vector local_text_head_; // [2 * hidden] engine::sampling::TorchCudaSamplingPolicy sampling_policy_; struct ProjectionRuntime; diff --git a/include/engine/models/moss/moss_tts_local/session.h b/include/engine/models/moss/moss_tts_local/session.h index 63a9fb52d..03738eda6 100644 --- a/include/engine/models/moss/moss_tts_local/session.h +++ b/include/engine/models/moss/moss_tts_local/session.h @@ -2,8 +2,7 @@ #include "engine/framework/runtime/cache_slots.h" #include "engine/framework/runtime/session_base.h" -#include "engine/models/moss/shared/audio_tokenizer_decoder.h" -#include "engine/models/moss/shared/audio_tokenizer_encoder.h" +#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" #include "engine/models/moss/moss_tts_local/assets.h" #include "engine/models/moss/moss_tts_local/backbone.h" #include "engine/models/moss/moss_tts_local/depth_transformer.h" @@ -38,8 +37,6 @@ class MossTTSLocalSession final runtime::TaskResult run(const runtime::TaskRequest & request) override; private: - moss::MossAudioTokenizerEncoder & encoder(); - struct ReferenceAudioCacheKey { uint64_t hash = 0; int sample_rate = 0; @@ -62,11 +59,8 @@ class MossTTSLocalSession final std::unique_ptr backbone_; std::unique_ptr depth_; std::unique_ptr processor_; - std::unique_ptr codec_; + std::unique_ptr codec_; std::unique_ptr generator_; - // Lazily built the first time a speaker reference is provided (voice cloning). - std::unique_ptr reference_encoder_execution_context_; - std::unique_ptr encoder_; runtime::CacheSlots reference_voice_cache_; }; diff --git a/include/engine/models/moss/moss_tts_local/tokenizer_text.h b/include/engine/models/moss/moss_tts_local/tokenizer_text.h index b94e52f06..da6d5be0a 100644 --- a/include/engine/models/moss/moss_tts_local/tokenizer_text.h +++ b/include/engine/models/moss/moss_tts_local/tokenizer_text.h @@ -35,7 +35,7 @@ class MossTextProcessor { const std::optional & language = std::nullopt) const; // Builds a voice-clone prompt. reference_codes is [num_codebooks][frames] as produced - // by MossAudioTokenizerEncoder for the reference speaker. + // by the MOSS audio tokenizer codec runtime for the reference speaker. MossGenerationPrefix build_clone_prefix( const std::string & text, const std::vector> & reference_codes, diff --git a/include/engine/models/moss/moss_tts_nano/session.h b/include/engine/models/moss/moss_tts_nano/session.h index e0c4b3e89..810453bb8 100644 --- a/include/engine/models/moss/moss_tts_nano/session.h +++ b/include/engine/models/moss/moss_tts_nano/session.h @@ -2,8 +2,7 @@ #include "engine/framework/core/execution_context.h" #include "engine/framework/runtime/session_base.h" -#include "engine/models/moss/shared/audio_tokenizer_decoder.h" -#include "engine/models/moss/shared/audio_tokenizer_encoder.h" +#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" #include "engine/models/moss/moss_tts_nano/assets.h" #include "engine/models/moss/moss_tts_nano/generator.h" #include "engine/models/moss/moss_tts_nano/global_transformer.h" @@ -33,7 +32,6 @@ class MossTTSNanoSession final runtime::TaskResult run(const runtime::TaskRequest & request) override; private: - moss::MossAudioTokenizerEncoder & encoder(); MossTTSNanoAudioCodes encode_reference_audio(const runtime::AudioBuffer & audio, int64_t active_codebooks); runtime::AudioBuffer decode_generated_audio(const MossTTSNanoAudioCodes & codes, int64_t active_codebooks); MossTTSNanoRequest make_request(const runtime::TaskRequest & request) const; @@ -56,9 +54,7 @@ class MossTTSNanoSession final MossTTSNanoGlobalTransformerRuntime global_transformer_; MossTTSNanoLocalFrameDecoderRuntime local_frame_decoder_; MossTTSNanoGenerator generator_; - moss::MossAudioTokenizerDecoder decoder_; - std::unique_ptr reference_encoder_execution_context_; - std::unique_ptr encoder_; + engine::codecs::MossAudioTokenizerCodecRuntime codec_; std::optional prepared_prompt_audio_; std::optional prepared_reference_codes_; }; diff --git a/include/engine/models/moss/shared/audio_tokenizer_config.h b/include/engine/models/moss/shared/audio_tokenizer_config.h deleted file mode 100644 index 2e9bb1e24..000000000 --- a/include/engine/models/moss/shared/audio_tokenizer_config.h +++ /dev/null @@ -1,48 +0,0 @@ -#pragma once - -#include -#include - -namespace engine::models::moss { - -struct AudioTokenizerTransformerStage { - int64_t input_dimension = 0; - int64_t output_dimension = 0; - int64_t model_dimension = 0; - int64_t num_heads = 0; - int64_t num_layers = 0; - int64_t feedforward_dimension = 0; - int64_t context_window = 0; - int64_t patch_size = 0; -}; - -struct AudioTokenizerQuantizerConfig { - int64_t codebook_size = 1024; - int64_t codebook_dim = 8; - int64_t rvq_dim = 512; - int64_t code_dim = 768; - int64_t num_quantizers = 12; -}; - -struct AudioTokenizerConfig { - int64_t sampling_rate = 48000; - int64_t samples_per_frame = 3840; - // v2 and Nano process left/right as one interleaved stream and split it at the end. - // v1 is mono, so its decoder output is already the waveform. - int64_t channels = 2; - AudioTokenizerQuantizerConfig quantizer; - std::vector encoder_stages; - std::vector decoder_stages; - int64_t encoder_final_patch = 1; - int64_t decoder_initial_patch = 1; - int64_t encoder_module_start = 1; - int64_t encoder_module_stride = 2; - int64_t decoder_module_start = 0; - int64_t decoder_module_stride = 2; -}; - -AudioTokenizerConfig moss_audio_tokenizer_v1_config(); -AudioTokenizerConfig moss_audio_tokenizer_v2_config(); -AudioTokenizerConfig moss_audio_tokenizer_nano_config(); - -} // namespace engine::models::moss diff --git a/include/engine/models/moss/shared/audio_tokenizer_decoder.h b/include/engine/models/moss/shared/audio_tokenizer_decoder.h deleted file mode 100644 index be1f4d921..000000000 --- a/include/engine/models/moss/shared/audio_tokenizer_decoder.h +++ /dev/null @@ -1,47 +0,0 @@ -#pragma once - -#include "engine/framework/core/execution_context.h" -#include "engine/models/moss/shared/audio_tokenizer_config.h" - -#include -#include -#include - -namespace engine::assets { -class TensorSource; -} - -namespace engine::models::moss { - -// 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 MossAudioTokenizerQuantizer. -class MossAudioTokenizerDecoder { -public: - MossAudioTokenizerDecoder( - const assets::TensorSource & source, - core::ExecutionContext & execution_context, - int64_t num_quantizers, - size_t weight_context_bytes, - size_t graph_arena_bytes, - AudioTokenizerConfig config = moss_audio_tokenizer_v2_config()); - ~MossAudioTokenizerDecoder(); - - MossAudioTokenizerDecoder(const MossAudioTokenizerDecoder &) = delete; - MossAudioTokenizerDecoder & operator=(const MossAudioTokenizerDecoder &) = 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 diff --git a/include/engine/models/moss/shared/audio_tokenizer_encoder.h b/include/engine/models/moss/shared/audio_tokenizer_encoder.h deleted file mode 100644 index 1e9937332..000000000 --- a/include/engine/models/moss/shared/audio_tokenizer_encoder.h +++ /dev/null @@ -1,45 +0,0 @@ -#pragma once - -#include "engine/framework/core/execution_context.h" -#include "engine/models/moss/shared/audio_tokenizer_config.h" - -#include -#include -#include - -namespace engine::assets { -class TensorSource; -} - -namespace engine::models::moss { - -// MOSS-Audio-Tokenizer encoder: turns a reference waveform into RLFQ codes -// for zero-shot voice cloning. It is the structural mirror of MossAudioTokenizerDecoder -- -// 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 MossAudioTokenizerEncoder { -public: - MossAudioTokenizerEncoder( - const assets::TensorSource & source, - core::ExecutionContext & execution_context, - int64_t num_quantizers, - size_t weight_context_bytes, - size_t graph_arena_bytes, - AudioTokenizerConfig config = moss_audio_tokenizer_v2_config()); - ~MossAudioTokenizerEncoder(); - - MossAudioTokenizerEncoder(const MossAudioTokenizerEncoder &) = delete; - MossAudioTokenizerEncoder & operator=(const MossAudioTokenizerEncoder &) = 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 diff --git a/include/engine/models/moss/shared/audio_tokenizer_quantizer.h b/include/engine/models/moss/shared/audio_tokenizer_quantizer.h deleted file mode 100644 index 0da9a4053..000000000 --- a/include/engine/models/moss/shared/audio_tokenizer_quantizer.h +++ /dev/null @@ -1,63 +0,0 @@ -#pragma once - -#include "engine/models/moss/shared/audio_tokenizer_config.h" - -#include -#include - -namespace engine::assets { -class TensorSource; -} - -namespace engine::models::moss { - -// 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 MossAudioTokenizerQuantizer { -public: - MossAudioTokenizerQuantizer( - const assets::TensorSource & source, - int64_t num_quantizers, - AudioTokenizerQuantizerConfig config = moss_audio_tokenizer_v2_config().quantizer); - - 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; - - // 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 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 latent_table; // [codebook_size, code_dim] row-major (decode) - std::vector in_weight; // [codebook_dim, rvq_dim] row-major (encode) - std::vector in_bias; // [codebook_dim] (encode) - }; - - 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] - std::vector input_weight_; // [rvq_dim, code_dim] row-major (encode) - std::vector input_bias_; // [rvq_dim] (encode) -}; - -} // namespace engine::models::moss diff --git a/include/engine/models/moss/shared/audio_tokenizer_transformer.h b/include/engine/models/moss/shared/audio_tokenizer_transformer.h deleted file mode 100644 index 230d0ab4e..000000000 --- a/include/engine/models/moss/shared/audio_tokenizer_transformer.h +++ /dev/null @@ -1,345 +0,0 @@ -#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 - -namespace engine::models::moss::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; -}; - -struct AttentionWindow { - int64_t query_start; - int64_t query_steps; - int64_t key_start; - int64_t key_steps; - core::TensorValue mask; -}; - -class CodecWeights { -public: - explicit CodecWeights(const assets::TensorSource & source) : source_(source) {} - - const assets::TensorSource & source_for(const std::string & name) const { - if (source_.has_tensor(name)) { - return source_; - } - throw std::runtime_error("MOSS codec tensor not found: " + name); - } - - bool has(const std::string & name) const noexcept { return source_.has_tensor(name); } - -private: - const assets::TensorSource & source_; -}; - -// 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 CodecWeights & codec_weights, - 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(codec_weights.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(codec_weights.source_for(name), name, shape); - }; - - TransformerWeights weights; - weights.spec = spec; - weights.input_proj = load(prefix + ".input_proj.weight", {spec.d_model, spec.input_dim}); - // Upstream's ProjectedTransformer only creates an output projection when the stage - // changes width. v2 ships one on every module; v1 leaves it out wherever - // output_dimension already equals d_model, so treat it as optional and fall through to - // the identity in that case. - const std::string output_proj_name = prefix + ".output_proj.weight"; - if (codec_weights.has(output_proj_name)) { - weights.output_proj = load(output_proj_name, {spec.output_dim, spec.d_model}); - } else if (spec.output_dim != spec.d_model) { - throw std::runtime_error( - "MOSS codec stage " + prefix + " changes width but carries no output projection"); - } - 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}); - // v2 stores one attention projection per layer; v1 keeps them in an indexed - // ModuleList (`in_projs.0`). Same tensor either way. - // v1 names the feed-forward layers linear1/linear2 where v2 uses an nn.Sequential. - const auto ffn_name = [&](const std::string & sequential, const std::string & named) { - return codec_weights.has(lp + sequential) ? lp + sequential : lp + named; - }; - const auto attention_name = [&](const std::string & single, const std::string & indexed) { - return codec_weights.has(lp + single) ? lp + single : lp + indexed; - }; - w.in_proj = load( - attention_name(".self_attn.in_proj.weight", ".self_attn.in_projs.0.weight"), - {3 * spec.d_model, spec.d_model}); - w.out_proj = load( - attention_name(".self_attn.out_proj.weight", ".self_attn.out_projs.0.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(ffn_name(".ffn.0.weight", ".linear1.weight"), {spec.intermediate_size, spec.d_model}); - w.fc2 = load(ffn_name(".ffn.2.weight", ".linear2.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 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 windowed_attention( - core::ModuleBuildContext & ctx, - const core::TensorValue & q_heads, - const core::TensorValue & k_heads, - const core::TensorValue & v_heads, - int64_t dim, - const std::vector & windows) { - if (windows.empty()) { - throw std::runtime_error("MOSS codec windowed attention requires at least one window"); - } - core::TensorValue merged; - for (const auto & window : windows) { - auto q_slice = modules::SliceModule({2, window.query_start, window.query_steps}).build(ctx, q_heads); - auto k_slice = modules::SliceModule({2, window.key_start, window.key_steps}).build(ctx, k_heads); - auto v_slice = modules::SliceModule({2, window.key_start, window.key_steps}).build(ctx, v_heads); - auto part = attention(ctx, q_slice, k_slice, v_slice, dim, window.mask); - merged = merged.valid() ? modules::ConcatModule({2}).build(ctx, merged, part) : part; - } - return merged; -} - -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, - const std::vector * windows, - 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 = core::ensure_backend_addressable_layout( - ctx, modules::SliceModule({2, 0, spec.d_model}).build(ctx, qkv)); - auto k = core::ensure_backend_addressable_layout( - ctx, modules::SliceModule({2, spec.d_model, spec.d_model}).build(ctx, qkv)); - auto v = core::ensure_backend_addressable_layout( - ctx, modules::SliceModule({2, 2 * spec.d_model, spec.d_model}).build(ctx, qkv)); - - q = modules::ReshapeModule({ - core::TensorShape::from_dims({q.shape.dims[0], q.shape.dims[1], spec.num_heads, dim}), - }).build(ctx, q); - k = modules::ReshapeModule({ - core::TensorShape::from_dims({k.shape.dims[0], k.shape.dims[1], spec.num_heads, dim}), - }).build(ctx, k); - v = modules::ReshapeModule({ - core::TensorShape::from_dims({v.shape.dims[0], v.shape.dims[1], spec.num_heads, dim}), - }).build(ctx, v); - 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 = windows == nullptr ? attention(ctx, q_heads, k_heads, v_heads, dim, mask) - : windowed_attention(ctx, q_heads, k_heads, v_heads, dim, *windows); - context = modules::TransposeModule({{0, 2, 1, 3}, context.shape.rank}).build(ctx, context); - context = core::ensure_backend_addressable_layout(ctx, context); - context = modules::ReshapeModule({ - core::TensorShape::from_dims({1, steps, spec.d_model}), - }).build(ctx, context); - 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)); - auto layer_scale1 = modules::ReshapeModule({ - core::TensorShape::from_dims({1, 1, spec.d_model}), - }).build(ctx, weights.layer_scale1); - layer_scale1 = modules::RepeatModule({attn_out.shape}).build(ctx, layer_scale1); - attn_out = modules::MulModule{}.build(ctx, attn_out, layer_scale1); - 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)); - auto layer_scale2 = modules::ReshapeModule({ - core::TensorShape::from_dims({1, 1, spec.d_model}), - }).build(ctx, weights.layer_scale2); - layer_scale2 = modules::RepeatModule({ff.shape}).build(ctx, layer_scale2); - ff = modules::MulModule{}.build(ctx, ff, layer_scale2); - 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 std::vector * windows = nullptr) { - 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, windows, steps); - } - if (!weights.output_proj.valid()) { - return x; - } - 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); -#ifdef _OPENMP -#pragma omp parallel for if(steps * steps >= 4096) -#endif - 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; -} - -inline std::vector causal_context_mask_window( - int64_t query_start, - int64_t query_steps, - int64_t key_start, - int64_t key_steps, - int64_t context) { - std::vector mask(static_cast(query_steps * key_steps), kMaskedAttentionBias); -#ifdef _OPENMP -#pragma omp parallel for if(query_steps * key_steps >= 4096) -#endif - for (int64_t q = 0; q < query_steps; ++q) { - const int64_t query = query_start + q; - for (int64_t k = 0; k < key_steps; ++k) { - const int64_t key = key_start + k; - if (key <= query && query - key < context) { - mask[static_cast(q * key_steps + k)] = 0.0F; - } - } - } - return mask; -} - -} // namespace engine::models::moss::codec_detail diff --git a/include/engine/models/moss/shared/token_rows.h b/include/engine/models/moss/shared/token_rows.h index 926ba7b93..76af2fdb5 100644 --- a/include/engine/models/moss/shared/token_rows.h +++ b/include/engine/models/moss/shared/token_rows.h @@ -1,9 +1,6 @@ #pragma once -#include "engine/framework/assets/tensor_source.h" - #include -#include #include namespace engine::models::moss { @@ -28,32 +25,4 @@ class TokenRowBuilder { TokenRows rows_; }; -struct AudioCodebookSpec { - int64_t hidden_size = 0; - int64_t num_codebooks = 0; - int64_t audio_vocab_size = 0; - int64_t audio_pad_token_id = 0; - std::vector audio_codebook_sizes; - std::string tensor_prefix = "audio_embeddings"; -}; - -class AudioCodebookEmbeddings { -public: - AudioCodebookEmbeddings(const assets::TensorSource & source, AudioCodebookSpec spec); - - int64_t hidden_size() const noexcept { return hidden_size_; } - int64_t num_codebooks() const noexcept { return num_codebooks_; } - int32_t audio_pad_token_id() const noexcept { return audio_pad_token_id_; } - int64_t codebook_size(int64_t codebook) const; - const float * embedding(int64_t codebook, int32_t code) const; - void add_bias(const int32_t * codes, float * bias) const; - std::vector bias_for(const int32_t * codes) const; - -private: - int64_t hidden_size_ = 0; - int64_t num_codebooks_ = 0; - int32_t audio_pad_token_id_ = 0; - std::vector> embeddings_; -}; - } // namespace engine::models::moss diff --git a/src/community_models/moss_voicegen/session.cpp b/src/community_models/moss_voicegen/session.cpp index 7380a7922..612aac1c7 100644 --- a/src/community_models/moss_voicegen/session.cpp +++ b/src/community_models/moss_voicegen/session.cpp @@ -123,15 +123,15 @@ void MossVoiceGenSession::prepare(const runtime::SessionPreparationRequest &) { const auto & config = assets_->config; text_processor_ = std::make_unique(assets_); - moss::AudioCodebookSpec codebook_spec; + engine::modules::MultiCodebookEmbeddingSpec codebook_spec; codebook_spec.hidden_size = config.backbone.hidden_size; codebook_spec.num_codebooks = config.num_codebooks; - codebook_spec.audio_vocab_size = config.audio_vocab_size + 1; - codebook_spec.audio_pad_token_id = config.audio_pad_code; + codebook_spec.vocab_size = config.audio_vocab_size + 1; + codebook_spec.pad_token_id = config.audio_pad_code; // The delay family stores its per-codebook input embeddings as emb_ext., not under // the shared default prefix. codebook_spec.tensor_prefix = "emb_ext"; - codebooks_ = std::make_unique(*assets_->model_weights, codebook_spec); + codebooks_ = std::make_unique(*assets_->model_weights, codebook_spec); backbone_ = std::make_unique( assets_, @@ -145,13 +145,18 @@ void MossVoiceGenSession::prepare(const runtime::SessionPreparationRequest &) { heads_graph_arena_bytes_, heads_weight_context_bytes_, weight_storage_type_); - codec_ = std::make_unique( - *assets_->audio_tokenizer_weights, + codec_ = std::make_unique( + assets_->audio_tokenizer_weights, execution_context(), config.num_codebooks, - codec_weight_context_bytes_, - codec_graph_arena_bytes_, - moss::moss_audio_tokenizer_v1_config()); + engine::codecs::MossAudioTokenizerCodecRuntimeOptions{ + codec_weight_context_bytes_, + codec_graph_arena_bytes_, + codec_graph_arena_bytes_, + false, + }, + engine::codecs::moss_audio_tokenizer_v1_config()); + codec_->prepare_decoder(); mark_prepared(); } @@ -234,11 +239,11 @@ std::vector MossVoiceGenSession::decode_codes(const GeneratedChunk & chun chunk.codes.begin() + static_cast(codebook * chunk.frames), chunk.codes.begin() + static_cast((codebook + 1) * chunk.frames)); } - auto channels = codec_->decode(codes); - if (channels.empty()) { + auto audio = codec_->decode(engine::codecs::MossAudioTokenizerCodes{chunk.frames, std::move(codes)}); + if (audio.channels.empty()) { throw std::runtime_error("MOSS-VoiceGenerator codec returned no audio"); } - return std::move(channels.front()); + return std::move(audio.channels.front()); } runtime::TaskResult MossVoiceGenSession::run(const runtime::TaskRequest & request) { @@ -312,7 +317,7 @@ runtime::TaskResult MossVoiceGenSession::run(const runtime::TaskRequest & reques debug::trace_log_scalar("moss_voicegen.chunk_count", static_cast(chunk_requests.size())); debug::trace_log_scalar("moss_voicegen.silent_chunks", silent_chunks); - debug::timing_log_scalar("moss_voicegen.run_ms", engine::debug::elapsed_ms(wall_start, Clock::now())); + debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start, Clock::now())); if (merged.samples.empty()) { throw std::runtime_error( diff --git a/src/community_models/vietneu_tts/session.cpp b/src/community_models/vietneu_tts/session.cpp index 8db7b11d5..6b9c18040 100644 --- a/src/community_models/vietneu_tts/session.cpp +++ b/src/community_models/vietneu_tts/session.cpp @@ -51,7 +51,7 @@ std::vector parse_speaker_embedding_file(const std::string & filepath) { runtime::AudioBuffer decode_moss_audio( const Qwen3SpeechCodes & codes, - const engine::models::moss::MossAudioTokenizerDecoder & decoder) { + engine::codecs::MossAudioTokenizerCodecRuntime & decoder) { const int64_t frames = codes.frames; const int64_t code_groups = codes.code_groups; std::vector> transposed(static_cast(code_groups), std::vector(static_cast(frames))); @@ -60,13 +60,13 @@ runtime::AudioBuffer decode_moss_audio( transposed[static_cast(g)][static_cast(f)] = codes.codes[static_cast(f * code_groups + g)]; } } - auto stereo = decoder.decode(transposed); + auto decoded = decoder.decode(engine::codecs::MossAudioTokenizerCodes{frames, std::move(transposed)}); runtime::AudioBuffer out; - out.sample_rate = 48000; - out.channels = 2; - if (stereo.size() >= 2) { - const auto & left = stereo[0]; - const auto & right = stereo[1]; + out.sample_rate = static_cast(decoded.sampling_rate); + out.channels = static_cast(decoded.channels.size()); + if (decoded.channels.size() >= 2) { + const auto & left = decoded.channels[0]; + const auto & right = decoded.channels[1]; out.samples.resize(left.size() * 2); for (size_t i = 0; i < left.size(); ++i) { out.samples[i * 2] = left[i]; @@ -307,13 +307,18 @@ VietneuTTSSession::VietneuTTSSession( talker_weights_, assets_->config.talker.max_position_embeddings, assets_->config.max_new_tokens); - moss_speech_decoder_ = std::make_unique( - *assets_->speech_tokenizer_weights, + moss_speech_decoder_ = std::make_unique( + assets_->speech_tokenizer_weights, execution_context(), assets_->config.speech_tokenizer.num_quantizers, - speech_decoder_constant_context_bytes_, - speech_decoder_graph_arena_bytes_, - engine::models::moss::moss_audio_tokenizer_nano_config()); + engine::codecs::MossAudioTokenizerCodecRuntimeOptions{ + speech_decoder_constant_context_bytes_, + speech_decoder_graph_arena_bytes_, + speech_decoder_graph_arena_bytes_, + false, + }, + engine::codecs::moss_audio_tokenizer_nano_config()); + moss_speech_decoder_->prepare_decoder(); if (task_.mode != runtime::RunMode::Offline) { throw std::runtime_error("Vietneu TTS currently supports offline sessions"); } diff --git a/src/framework/codecs/moss_audio_tokenizer_codec_runtime.cpp b/src/framework/codecs/moss_audio_tokenizer_codec_runtime.cpp new file mode 100644 index 000000000..1feffd1cf --- /dev/null +++ b/src/framework/codecs/moss_audio_tokenizer_codec_runtime.cpp @@ -0,0 +1,1570 @@ +#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/module.h" +#include "engine/framework/debug/profiler.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 +#include +#include +#include + +namespace engine::codecs::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 GgmlGallocrDeleter { + void operator()(ggml_gallocr_t alloc) const noexcept { + if (alloc != nullptr) { + ggml_gallocr_free(alloc); + } + } +}; + +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; +}; + +struct AttentionWindow { + int64_t query_start; + int64_t query_steps; + int64_t key_start; + int64_t key_steps; + core::TensorValue mask; +}; + +class CodecWeights { +public: + explicit CodecWeights(const assets::TensorSource & source) : source_(source) {} + + const assets::TensorSource & source_for(const std::string & name) const { + if (source_.has_tensor(name)) { + return source_; + } + throw std::runtime_error("MOSS codec tensor not found: " + name); + } + + bool has(const std::string & name) const noexcept { return source_.has_tensor(name); } + +private: + const assets::TensorSource & source_; +}; + +// 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 CodecWeights & codec_weights, + 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(codec_weights.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(codec_weights.source_for(name), name, shape); + }; + + TransformerWeights weights; + weights.spec = spec; + weights.input_proj = load(prefix + ".input_proj.weight", {spec.d_model, spec.input_dim}); + // Upstream's ProjectedTransformer only creates an output projection when the stage + // changes width. v2 ships one on every module; v1 leaves it out wherever + // output_dimension already equals d_model, so treat it as optional and fall through to + // the identity in that case. + const std::string output_proj_name = prefix + ".output_proj.weight"; + if (codec_weights.has(output_proj_name)) { + weights.output_proj = load(output_proj_name, {spec.output_dim, spec.d_model}); + } else if (spec.output_dim != spec.d_model) { + throw std::runtime_error( + "MOSS codec stage " + prefix + " changes width but carries no output projection"); + } + 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}); + // v2 stores one attention projection per layer; v1 keeps them in an indexed + // ModuleList (`in_projs.0`). Same tensor either way. + // v1 names the feed-forward layers linear1/linear2 where v2 uses an nn.Sequential. + const auto ffn_name = [&](const std::string & sequential, const std::string & named) { + return codec_weights.has(lp + sequential) ? lp + sequential : lp + named; + }; + const auto attention_name = [&](const std::string & single, const std::string & indexed) { + return codec_weights.has(lp + single) ? lp + single : lp + indexed; + }; + w.in_proj = load( + attention_name(".self_attn.in_proj.weight", ".self_attn.in_projs.0.weight"), + {3 * spec.d_model, spec.d_model}); + w.out_proj = load( + attention_name(".self_attn.out_proj.weight", ".self_attn.out_projs.0.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(ffn_name(".ffn.0.weight", ".linear1.weight"), {spec.intermediate_size, spec.d_model}); + w.fc2 = load(ffn_name(".ffn.2.weight", ".linear2.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 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 windowed_attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & q_heads, + const core::TensorValue & k_heads, + const core::TensorValue & v_heads, + int64_t dim, + const std::vector & windows) { + if (windows.empty()) { + throw std::runtime_error("MOSS codec windowed attention requires at least one window"); + } + core::TensorValue merged; + for (const auto & window : windows) { + auto q_slice = modules::SliceModule({2, window.query_start, window.query_steps}).build(ctx, q_heads); + auto k_slice = modules::SliceModule({2, window.key_start, window.key_steps}).build(ctx, k_heads); + auto v_slice = modules::SliceModule({2, window.key_start, window.key_steps}).build(ctx, v_heads); + auto part = attention(ctx, q_slice, k_slice, v_slice, dim, window.mask); + merged = merged.valid() ? modules::ConcatModule({2}).build(ctx, merged, part) : part; + } + return merged; +} + +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, + const std::vector * windows, + 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 = core::ensure_backend_addressable_layout( + ctx, modules::SliceModule({2, 0, spec.d_model}).build(ctx, qkv)); + auto k = core::ensure_backend_addressable_layout( + ctx, modules::SliceModule({2, spec.d_model, spec.d_model}).build(ctx, qkv)); + auto v = core::ensure_backend_addressable_layout( + ctx, modules::SliceModule({2, 2 * spec.d_model, spec.d_model}).build(ctx, qkv)); + + q = modules::ReshapeModule({ + core::TensorShape::from_dims({q.shape.dims[0], q.shape.dims[1], spec.num_heads, dim}), + }).build(ctx, q); + k = modules::ReshapeModule({ + core::TensorShape::from_dims({k.shape.dims[0], k.shape.dims[1], spec.num_heads, dim}), + }).build(ctx, k); + v = modules::ReshapeModule({ + core::TensorShape::from_dims({v.shape.dims[0], v.shape.dims[1], spec.num_heads, dim}), + }).build(ctx, v); + 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 = windows == nullptr ? attention(ctx, q_heads, k_heads, v_heads, dim, mask) + : windowed_attention(ctx, q_heads, k_heads, v_heads, dim, *windows); + context = modules::TransposeModule({{0, 2, 1, 3}, context.shape.rank}).build(ctx, context); + context = core::ensure_backend_addressable_layout(ctx, context); + context = modules::ReshapeModule({ + core::TensorShape::from_dims({1, steps, spec.d_model}), + }).build(ctx, context); + 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)); + auto layer_scale1 = modules::ReshapeModule({ + core::TensorShape::from_dims({1, 1, spec.d_model}), + }).build(ctx, weights.layer_scale1); + layer_scale1 = modules::RepeatModule({attn_out.shape}).build(ctx, layer_scale1); + attn_out = modules::MulModule{}.build(ctx, attn_out, layer_scale1); + 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)); + auto layer_scale2 = modules::ReshapeModule({ + core::TensorShape::from_dims({1, 1, spec.d_model}), + }).build(ctx, weights.layer_scale2); + layer_scale2 = modules::RepeatModule({ff.shape}).build(ctx, layer_scale2); + ff = modules::MulModule{}.build(ctx, ff, layer_scale2); + 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 std::vector * windows = nullptr) { + 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, windows, steps); + } + if (!weights.output_proj.valid()) { + return x; + } + 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); +#ifdef _OPENMP +#pragma omp parallel for if(steps * steps >= 4096) +#endif + 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; +} + +inline std::vector causal_context_mask_window( + int64_t query_start, + int64_t query_steps, + int64_t key_start, + int64_t key_steps, + int64_t context) { + std::vector mask(static_cast(query_steps * key_steps), kMaskedAttentionBias); +#ifdef _OPENMP +#pragma omp parallel for if(query_steps * key_steps >= 4096) +#endif + for (int64_t q = 0; q < query_steps; ++q) { + const int64_t query = query_start + q; + for (int64_t k = 0; k < key_steps; ++k) { + const int64_t key = key_start + k; + if (key <= query && query - key < context) { + mask[static_cast(q * key_steps + k)] = 0.0F; + } + } + } + return mask; +} + +} // namespace engine::codecs::codec_detail + +namespace engine::codecs { + +// 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 MossAudioTokenizerQuantizer { +public: + MossAudioTokenizerQuantizer( + const assets::TensorSource & source, + int64_t num_quantizers, + MossAudioTokenizerQuantizerConfig config = moss_audio_tokenizer_v2_config().quantizer); + + 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; + + // 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 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 latent_table; // [codebook_size, code_dim] row-major (decode) + std::vector in_weight; // [codebook_dim, rvq_dim] row-major (encode) + std::vector in_bias; // [codebook_dim] (encode) + }; + + 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] + std::vector input_weight_; // [rvq_dim, code_dim] row-major (encode) + std::vector input_bias_; // [rvq_dim] (encode) +}; + +// MOSS-Audio-Tokenizer encoder: turns a reference waveform into RLFQ codes +// for zero-shot voice cloning. It is the structural mirror of MossAudioTokenizerDecoder -- +// 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 MossAudioTokenizerEncoder { +public: + MossAudioTokenizerEncoder( + const assets::TensorSource & source, + std::shared_ptr quantizer, + core::ExecutionContext & execution_context, + size_t weight_context_bytes, + size_t graph_arena_bytes, + MossAudioTokenizerConfig config = moss_audio_tokenizer_v2_config()); + ~MossAudioTokenizerEncoder(); + + MossAudioTokenizerEncoder(const MossAudioTokenizerEncoder &) = delete; + MossAudioTokenizerEncoder & operator=(const MossAudioTokenizerEncoder &) = delete; + + // Encodes a waveform given as {left, right} channels (each with the same + // per-channel sample count, 48 kHz) into [num_quantizers][frames] codes. + MossAudioTokenizerCodes encode(const MossAudioTokenizerAudio & audio); + void release_runtime_graphs(); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +// 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 MossAudioTokenizerQuantizer. +class MossAudioTokenizerDecoder { +public: + MossAudioTokenizerDecoder( + const assets::TensorSource & source, + std::shared_ptr dequantizer, + core::ExecutionContext & execution_context, + size_t weight_context_bytes, + size_t graph_arena_bytes, + MossAudioTokenizerConfig config = moss_audio_tokenizer_v2_config()); + ~MossAudioTokenizerDecoder(); + + MossAudioTokenizerDecoder(const MossAudioTokenizerDecoder &) = delete; + MossAudioTokenizerDecoder & operator=(const MossAudioTokenizerDecoder &) = 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. + MossAudioTokenizerAudio decode(const MossAudioTokenizerCodes & codes); + void release_runtime_graphs(); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +namespace { + +// 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)); +#ifdef _OPENMP +#pragma omp parallel for if(out_channels * in_channels >= 4096) +#endif + 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 assets::TensorSource & source, + const std::string & prefix, + int64_t out_channels, + int64_t in_channels) { + const auto g = source.require_f32(prefix + ".parametrizations.weight.original0"); + const auto v = source.require_f32(prefix + ".parametrizations.weight.original1"); + return reconstruct_weight_norm(g, v, out_channels, in_channels); +} + +} // namespace + +MossAudioTokenizerQuantizer::MossAudioTokenizerQuantizer( + const assets::TensorSource & source, + int64_t num_quantizers, + MossAudioTokenizerQuantizerConfig config) + : codebook_size_(config.codebook_size), + codebook_dim_(config.codebook_dim), + rvq_dim_(config.rvq_dim), + code_dim_(config.code_dim), + num_quantizers_(num_quantizers) { + if (num_quantizers_ <= 0) { + throw std::runtime_error("MOSS codec dequantizer requires a positive quantizer count"); + } + + output_weight_ = load_wn_conv_weight(source, "quantizer.output_proj", code_dim_, rvq_dim_); + output_bias_ = source.require_f32("quantizer.output_proj.bias"); + + 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 = source.require_f32(prefix + ".codebook.weight"); + codebook.out_weight = load_wn_conv_weight(source, prefix + ".out_proj", rvq_dim_, codebook_dim_); + codebook.out_bias = source.require_f32(prefix + ".out_proj.bias"); + std::vector combined_bias(static_cast(code_dim_)); + std::vector combined_weight(static_cast(code_dim_ * codebook_dim_)); +#ifdef _OPENMP +#pragma omp parallel for if(code_dim_ >= 256) +#endif + for (int64_t out = 0; out < code_dim_; ++out) { + const float * output_row = &output_weight_[static_cast(out * rvq_dim_)]; + float bias_sum = 0.0F; + for (int64_t rvq = 0; rvq < rvq_dim_; ++rvq) { + bias_sum += output_row[rvq] * codebook.out_bias[static_cast(rvq)]; + } + combined_bias[static_cast(out)] = bias_sum; + for (int64_t k = 0; k < codebook_dim_; ++k) { + float sum = 0.0F; + for (int64_t rvq = 0; rvq < rvq_dim_; ++rvq) { + sum += output_row[rvq] * codebook.out_weight[static_cast(rvq * codebook_dim_ + k)]; + } + combined_weight[static_cast(out * codebook_dim_ + k)] = sum; + } + } + codebook.latent_table.resize(static_cast(codebook_size_ * code_dim_)); +#ifdef _OPENMP +#pragma omp parallel for collapse(2) if(codebook_size_ * code_dim_ >= 4096) +#endif + for (int64_t code = 0; code < codebook_size_; ++code) { + for (int64_t out = 0; out < code_dim_; ++out) { + const float * embedding = &codebook.table[static_cast(code * codebook_dim_)]; + float sum = combined_bias[static_cast(out)]; + const float * row = &combined_weight[static_cast(out * codebook_dim_)]; + for (int64_t k = 0; k < codebook_dim_; ++k) { + sum += row[k] * embedding[k]; + } + codebook.latent_table[static_cast(code * code_dim_ + out)] = sum; + } + } + codebook.in_weight = load_wn_conv_weight(source, prefix + ".in_proj", codebook_dim_, rvq_dim_); + codebook.in_bias = source.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; +#ifdef _OPENMP +#pragma omp parallel for if(codebook_size_ * codebook_dim_ >= 4096) +#endif + 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)); + } + + input_weight_ = load_wn_conv_weight(source, "quantizer.input_proj", rvq_dim_, code_dim_); + input_bias_ = source.require_f32("quantizer.input_proj.bias"); +} + +std::vector MossAudioTokenizerQuantizer::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"); + } + for (int64_t step = 0; step < steps; ++step) { + for (int64_t index = 0; index < num_quantizers_; ++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"); + } + } + } + + std::vector latent(static_cast(code_dim_ * steps)); +#ifdef _OPENMP +#pragma omp parallel for if(steps * code_dim_ >= 4096) +#endif + for (int64_t step = 0; step < steps; ++step) { + for (int64_t out = 0; out < code_dim_; ++out) { + float value = output_bias_[static_cast(out)]; + 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)]; + const float * decoded = &codebook.latent_table[static_cast(code * code_dim_)]; + value += decoded[static_cast(out)]; + } + latent[static_cast(out * steps + step)] = value; + } + } + return latent; +} + +std::vector> MossAudioTokenizerQuantizer::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))); + + const auto encode_frame = [&](int64_t step, std::vector & residual, std::vector & encoding) { + std::fill(residual.begin(), residual.end(), 0.0); + std::fill(encoding.begin(), encoding.end(), 0.0); + + // 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; + } + } + }; + +#ifdef _OPENMP + if (frames >= 8) { +#pragma omp parallel + { + std::vector residual(static_cast(rvq_dim_)); + std::vector encoding(static_cast(codebook_dim_)); +#pragma omp for + for (int64_t step = 0; step < frames; ++step) { + encode_frame(step, residual, encoding); + } + } + } else +#endif + { + std::vector residual(static_cast(rvq_dim_)); + std::vector encoding(static_cast(codebook_dim_)); + for (int64_t step = 0; step < frames; ++step) { + encode_frame(step, residual, encoding); + } + } + return codes; +} + +namespace { + +namespace cd = codec_detail; + +cd::TransformerSpec to_encoder_transformer_spec(const MossAudioTokenizerTransformerStage & stage) { + return { + stage.input_dimension, + stage.output_dimension, + stage.model_dimension, + stage.num_heads, + stage.num_layers, + stage.feedforward_dimension, + stage.context_window, + stage.patch_size, + }; +} + +// 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 = core::ensure_backend_addressable_layout(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; + auto reshaped = engine::modules::ReshapeModule({ + core::TensorShape::from_dims({1, length, patch, channels}), + }).build(ctx, contiguous); + auto transposed = engine::modules::TransposeModule({{0, 1, 3, 2}, reshaped.shape.rank}).build(ctx, reshaped); + return engine::modules::ReshapeModule({ + core::TensorShape::from_dims({1, length, channels * patch}), + }).build(ctx, core::ensure_backend_addressable_layout(ctx, transposed)); +} + +} // namespace + +struct MossAudioTokenizerEncoder::Impl { + struct StageInput { + ggml_tensor * positions = nullptr; + std::vector position_host; + ggml_tensor * mask = nullptr; + std::vector mask_host; + }; + + struct GraphCache { + int64_t interleaved = 0; + int64_t output_steps = 0; + std::unique_ptr graph_ctx; + ggml_cgraph * graph = nullptr; + ggml_tensor * input = nullptr; + ggml_tensor * output = nullptr; + std::vector stage_inputs; + std::unique_ptr, cd::GgmlGallocrDeleter> gallocr; + }; + + ggml_backend_t backend = nullptr; + core::BackendType backend_type = core::BackendType::Cpu; + int threads = 1; + int64_t samples_per_frame = 3840; + int64_t channels = 2; + size_t graph_arena_bytes = 0; + MossAudioTokenizerConfig config; + std::shared_ptr quantizer; + std::unique_ptr store; + std::vector transformers; + std::unique_ptr graph_cache; + + GraphCache & prepare_graph(int64_t interleaved) { + if (graph_cache != nullptr && graph_cache->interleaved == interleaved) { + return *graph_cache; + } + + auto cache = std::make_unique(); + cache->interleaved = interleaved; + + ggml_init_params params{graph_arena_bytes, nullptr, true}; + cache->graph_ctx.reset(ggml_init(params)); + if (cache->graph_ctx == nullptr) { + throw std::runtime_error("failed to initialize MOSS codec encoder graph context"); + } + core::ModuleBuildContext ctx{cache->graph_ctx.get(), "moss.audio_tokenizer.encode", backend_type}; + + auto input_tensor = + core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, interleaved, 1})); + ggml_set_input(input_tensor.tensor); + cache->input = input_tensor.tensor; + + auto hidden = input_tensor; + int64_t steps = interleaved; + cache->stage_inputs.reserve(transformers.size()); + for (const auto & transformer : transformers) { + 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); + cache->stage_inputs.push_back(std::move(stage)); + + hidden = cd::run_transformer(ctx, hidden, transformer, positions, mask, steps); + } + if (config.encoder_final_patch > 1) { + hidden = patch_downsample(ctx, hidden, config.encoder_final_patch); + steps /= config.encoder_final_patch; + } + + hidden = core::ensure_backend_addressable_layout(ctx, hidden); + ggml_set_output(hidden.tensor); + cache->output = hidden.tensor; + cache->output_steps = steps; + + cache->graph = ggml_new_graph_custom(cache->graph_ctx.get(), 131072, false); + ggml_build_forward_expand(cache->graph, hidden.tensor); + + cache->gallocr.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend))); + if (cache->gallocr == nullptr || + !ggml_gallocr_reserve(cache->gallocr.get(), cache->graph) || + !ggml_gallocr_alloc_graph(cache->gallocr.get(), cache->graph)) { + throw std::runtime_error("failed to allocate MOSS codec encoder forward graph"); + } + + graph_cache = std::move(cache); + return *graph_cache; + } + + void release_runtime_graphs() { + graph_cache.reset(); + } +}; + +MossAudioTokenizerEncoder::MossAudioTokenizerEncoder( + const assets::TensorSource & source, + std::shared_ptr quantizer, + core::ExecutionContext & execution_context, + size_t weight_context_bytes, + size_t graph_arena_bytes, + MossAudioTokenizerConfig config) + : 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_->threads = execution_context.config().threads; + impl_->config = config; + impl_->samples_per_frame = config.samples_per_frame; + impl_->channels = config.channels; + impl_->graph_arena_bytes = graph_arena_bytes; + impl_->quantizer = std::move(quantizer); + if (impl_->quantizer == nullptr) { + throw std::runtime_error("MOSS codec encoder requires a quantizer"); + } + + cd::CodecWeights weights(source); + impl_->store = std::make_unique( + impl_->backend, impl_->backend_type, "moss.audio_tokenizer.encoder", weight_context_bytes); + impl_->transformers.reserve(config.encoder_stages.size()); + for (size_t index = 0; index < config.encoder_stages.size(); ++index) { + const int64_t module_index = + config.encoder_module_start + static_cast(index) * config.encoder_module_stride; + impl_->transformers.push_back(cd::load_transformer( + *impl_->store, weights, to_encoder_transformer_spec(config.encoder_stages[index]), "encoder", module_index)); + } + impl_->store->upload(); +} + +MossAudioTokenizerEncoder::~MossAudioTokenizerEncoder() = default; + +MossAudioTokenizerCodes MossAudioTokenizerEncoder::encode(const MossAudioTokenizerAudio & audio) { + const auto & channels = audio.channels; + if (audio.sampling_rate != 0 && audio.sampling_rate != impl_->config.sampling_rate) { + throw std::runtime_error("MOSS codec encoder input sample rate does not match codec config"); + } + if (static_cast(channels.size()) != impl_->channels) { + throw std::runtime_error("MOSS codec encoder input channel count does not match codec config"); + } + for (int64_t channel = 1; channel < impl_->channels; ++channel) { + if (channels[static_cast(channel)].size() != channels.front().size()) { + throw std::runtime_error("MOSS codec encoder channels must have equal length"); + } + } + const int64_t raw_per_channel = static_cast(channels.front().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 for the encoder graph, + // but keep the official valid code length as floor(valid_samples / samples_per_frame). + // MossAudioTokenizerPatchedPretransform pads the tensor and propagates input_lengths + // with integer division, then slices audio_codes to audio_codes_lengths. + const int64_t frames = (raw_per_channel + impl_->samples_per_frame - 1) / impl_->samples_per_frame; + const int64_t valid_frames = raw_per_channel / impl_->samples_per_frame; + const int64_t per_channel = frames * impl_->samples_per_frame; + const int64_t interleaved = per_channel * impl_->channels; + std::vector waveform(static_cast(interleaved), 0.0F); +#ifdef _OPENMP +#pragma omp parallel for if(raw_per_channel >= 4096) +#endif + for (int64_t i = 0; i < raw_per_channel; ++i) { + for (int64_t channel = 0; channel < impl_->channels; ++channel) { + waveform[static_cast(impl_->channels * i + channel)] = + channels[static_cast(channel)][static_cast(i)]; + } + } + + auto & graph = impl_->prepare_graph(interleaved); + ggml_backend_tensor_set(graph.input, waveform.data(), 0, waveform.size() * sizeof(float)); + for (const auto & stage : graph.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)); + } + + core::set_backend_threads(impl_->backend, impl_->threads); + const ggml_status status = ggml_backend_graph_compute(impl_->backend, graph.graph); + ggml_backend_synchronize(impl_->backend); + if (status != GGML_STATUS_SUCCESS) { + 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 MossAudioTokenizerQuantizer::encode expects. + std::vector latent(static_cast(graph.output_steps * cd::kCodeDim)); + ggml_backend_tensor_get(graph.output, latent.data(), 0, latent.size() * sizeof(float)); + + if (valid_frames <= 0) { + throw std::runtime_error("MOSS codec encoder input is shorter than one codec frame"); + } + latent.resize(static_cast(valid_frames * cd::kCodeDim)); + return MossAudioTokenizerCodes{ + valid_frames, + impl_->quantizer->encode(latent, valid_frames), + }; +} + +void MossAudioTokenizerEncoder::release_runtime_graphs() { + impl_->release_runtime_graphs(); +} + +namespace { + +namespace cd = codec_detail; + +constexpr int64_t kAttentionQueryChunk = 1500; + +cd::TransformerSpec to_decoder_transformer_spec(const MossAudioTokenizerTransformerStage & stage) { + return { + stage.input_dimension, + stage.output_dimension, + stage.model_dimension, + stage.num_heads, + stage.num_layers, + stage.feedforward_dimension, + stage.context_window, + stage.patch_size, + }; +} + +// 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 = core::ensure_backend_addressable_layout(ctx, input); + const int64_t length = contiguous.shape.dims[1]; + const int64_t packed = contiguous.shape.dims[2]; + const int64_t channels = packed / patch; + auto reshaped = engine::modules::ReshapeModule({ + core::TensorShape::from_dims({1, length, channels, patch}), + }).build(ctx, contiguous); + auto transposed = engine::modules::TransposeModule({{0, 1, 3, 2}, reshaped.shape.rank}).build(ctx, reshaped); + return engine::modules::ReshapeModule({ + core::TensorShape::from_dims({1, length * patch, channels}), + }).build(ctx, core::ensure_backend_addressable_layout(ctx, transposed)); +} + +} // namespace + +struct MossAudioTokenizerDecoder::Impl { + struct StageInput { + struct MaskInput { + ggml_tensor * tensor = nullptr; + std::vector host; + }; + + ggml_tensor * positions = nullptr; + std::vector position_host; + std::vector masks; + }; + + struct GraphCache { + int64_t frames = 0; + int64_t interleaved = 0; + std::unique_ptr graph_ctx; + ggml_cgraph * graph = nullptr; + ggml_tensor * input = nullptr; + ggml_tensor * output = nullptr; + std::vector stage_inputs; + std::unique_ptr, cd::GgmlGallocrDeleter> gallocr; + }; + + ggml_backend_t backend = nullptr; + core::BackendType backend_type = core::BackendType::Cpu; + int threads = 1; + int64_t sampling_rate = 48000; + size_t graph_arena_bytes = 0; + MossAudioTokenizerConfig config; + std::shared_ptr dequantizer; + std::unique_ptr store; + std::vector transformers; + std::unique_ptr graph_cache; + + GraphCache & prepare_graph(int64_t frames) { + if (graph_cache != nullptr && graph_cache->frames == frames) { + return *graph_cache; + } + + auto cache = std::make_unique(); + cache->frames = frames; + + ggml_init_params params{graph_arena_bytes, nullptr, true}; + cache->graph_ctx.reset(ggml_init(params)); + if (cache->graph_ctx == nullptr) { + throw std::runtime_error("failed to initialize MOSS codec decoder graph context"); + } + core::ModuleBuildContext ctx{cache->graph_ctx.get(), "moss.audio_tokenizer.decode", backend_type}; + + auto latent_tensor = + core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, frames, cd::kCodeDim})); + ggml_set_input(latent_tensor.tensor); + cache->input = latent_tensor.tensor; + + auto hidden = latent_tensor; + int64_t steps = frames; + if (config.decoder_initial_patch > 1) { + hidden = patch_upsample(ctx, hidden, config.decoder_initial_patch); + steps *= config.decoder_initial_patch; + } + cache->stage_inputs.reserve(transformers.size()); + for (const auto & transformer : transformers) { + auto positions = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({steps})); + ggml_set_input(positions.tensor); + const bool use_windowed_attention = steps > kAttentionQueryChunk; + core::TensorValue mask; + if (!use_windowed_attention) { + 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); + } + std::vector windows; + if (use_windowed_attention) { + for (int64_t query_start = 0; query_start < steps; query_start += kAttentionQueryChunk) { + const int64_t query_steps = std::min(kAttentionQueryChunk, steps - query_start); + const int64_t key_start = std::max(0, query_start - transformer.spec.context + 1); + const int64_t key_steps = query_start + query_steps - key_start; + auto window_mask = core::make_tensor( + ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, 1, query_steps, key_steps})); + ggml_set_input(window_mask.tensor); + stage.masks.push_back(StageInput::MaskInput{ + window_mask.tensor, + cd::causal_context_mask_window( + query_start, + query_steps, + key_start, + key_steps, + transformer.spec.context), + }); + windows.push_back(cd::AttentionWindow{ + query_start, + query_steps, + key_start, + key_steps, + window_mask, + }); + } + } else { + stage.masks.push_back(StageInput::MaskInput{ + mask.tensor, + cd::causal_context_mask(steps, transformer.spec.context), + }); + } + cache->stage_inputs.push_back(std::move(stage)); + + hidden = windows.empty() + ? cd::run_transformer(ctx, hidden, transformer, positions, mask, steps) + : cd::run_transformer(ctx, hidden, transformer, positions, windows.front().mask, steps, &windows); + hidden = patch_upsample(ctx, hidden, transformer.spec.patch); + steps *= transformer.spec.patch; + } + + hidden = core::ensure_backend_addressable_layout(ctx, hidden); + ggml_set_output(hidden.tensor); + cache->output = hidden.tensor; + cache->interleaved = steps; + + cache->graph = ggml_new_graph_custom(cache->graph_ctx.get(), 131072, false); + ggml_build_forward_expand(cache->graph, hidden.tensor); + + cache->gallocr.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend))); + if (cache->gallocr == nullptr || + !ggml_gallocr_reserve(cache->gallocr.get(), cache->graph) || + !ggml_gallocr_alloc_graph(cache->gallocr.get(), cache->graph)) { + throw std::runtime_error("failed to allocate MOSS codec decoder forward graph"); + } + + graph_cache = std::move(cache); + return *graph_cache; + } + + void release_runtime_graphs() { + graph_cache.reset(); + } +}; + +MossAudioTokenizerDecoder::MossAudioTokenizerDecoder( + const assets::TensorSource & source, + std::shared_ptr dequantizer, + core::ExecutionContext & execution_context, + size_t weight_context_bytes, + size_t graph_arena_bytes, + MossAudioTokenizerConfig config) + : 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_->threads = execution_context.config().threads; + impl_->config = config; + impl_->sampling_rate = config.sampling_rate; + impl_->graph_arena_bytes = graph_arena_bytes; + impl_->dequantizer = std::move(dequantizer); + if (impl_->dequantizer == nullptr) { + throw std::runtime_error("MOSS codec decoder requires a quantizer"); + } + + cd::CodecWeights weights(source); + impl_->store = std::make_unique( + impl_->backend, impl_->backend_type, "moss.audio_tokenizer.decoder", weight_context_bytes); + impl_->transformers.reserve(config.decoder_stages.size()); + for (size_t index = 0; index < config.decoder_stages.size(); ++index) { + const int64_t module_index = + config.decoder_module_start + static_cast(index) * config.decoder_module_stride; + impl_->transformers.push_back(cd::load_transformer( + *impl_->store, weights, to_decoder_transformer_spec(config.decoder_stages[index]), "decoder", module_index)); + } + impl_->store->upload(); +} + +MossAudioTokenizerDecoder::~MossAudioTokenizerDecoder() = default; + +int64_t MossAudioTokenizerDecoder::sampling_rate() const noexcept { + return impl_->sampling_rate; +} + +MossAudioTokenizerAudio MossAudioTokenizerDecoder::decode(const MossAudioTokenizerCodes & codes) { + const int64_t frames = codes.frames; + if (frames <= 0) { + throw std::runtime_error("MOSS codec decoder requires a non-empty code sequence"); + } + if (codes.codebooks.empty()) { + throw std::runtime_error("MOSS codec decoder requires codebooks"); + } + for (const auto & codebook : codes.codebooks) { + if (static_cast(codebook.size()) != frames) { + throw std::runtime_error("MOSS codec decoder codebooks do not match frame count"); + } + } + + // Codes -> continuous latent [code_dim, frames] (channel-major), transposed + // into the feature-last [1, frames, code_dim] layout the decoder expects. + double dequant_ms = 0.0; + double latent_pack_ms = 0.0; + double graph_build_ms = 0.0; + double input_upload_ms = 0.0; + double graph_compute_ms = 0.0; + double output_read_ms = 0.0; + double deinterleave_ms = 0.0; + const bool collect_timing = engine::debug::timing_log_enabled(); + std::vector latent; + if (collect_timing) { + dequant_ms = engine::debug::measure_ms([&]() { + latent = impl_->dequantizer->decode(codes.codebooks); + }); + } else { + latent = impl_->dequantizer->decode(codes.codebooks); + } + std::vector latent_input; + if (collect_timing) { + latent_pack_ms = engine::debug::measure_ms([&]() { + latent_input.resize(static_cast(frames * cd::kCodeDim)); +#ifdef _OPENMP +#pragma omp parallel for collapse(2) if(frames * cd::kCodeDim >= 4096) +#endif + for (int64_t channel = 0; channel < cd::kCodeDim; ++channel) { + for (int64_t step = 0; step < frames; ++step) { + latent_input[static_cast(step * cd::kCodeDim + channel)] = + latent[static_cast(channel * frames + step)]; + } + } + }); + } else { + latent_input.resize(static_cast(frames * cd::kCodeDim)); +#ifdef _OPENMP +#pragma omp parallel for collapse(2) if(frames * cd::kCodeDim >= 4096) +#endif + for (int64_t channel = 0; channel < cd::kCodeDim; ++channel) { + for (int64_t step = 0; step < frames; ++step) { + latent_input[static_cast(step * cd::kCodeDim + channel)] = + latent[static_cast(channel * frames + step)]; + } + } + } + + const auto graph_build_start = std::chrono::steady_clock::now(); + auto & graph = impl_->prepare_graph(frames); + if (collect_timing) { + graph_build_ms = engine::debug::elapsed_ms(graph_build_start); + } + + const auto upload_start = std::chrono::steady_clock::now(); + ggml_backend_tensor_set( + graph.input, latent_input.data(), 0, latent_input.size() * sizeof(float)); + for (const auto & stage : graph.stage_inputs) { + ggml_backend_tensor_set( + stage.positions, stage.position_host.data(), 0, stage.position_host.size() * sizeof(int32_t)); + for (const auto & mask : stage.masks) { + ggml_backend_tensor_set(mask.tensor, mask.host.data(), 0, mask.host.size() * sizeof(float)); + } + } + if (collect_timing) { + input_upload_ms = engine::debug::elapsed_ms(upload_start); + } + + const auto compute_start = std::chrono::steady_clock::now(); + core::set_backend_threads(impl_->backend, impl_->threads); + const ggml_status status = ggml_backend_graph_compute(impl_->backend, graph.graph); + ggml_backend_synchronize(impl_->backend); + if (collect_timing) { + graph_compute_ms = engine::debug::elapsed_ms(compute_start); + } + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("MOSS codec decoder forward graph compute failed"); + } + + const int64_t interleaved = graph.interleaved; // frames * samples_per_frame * channels + std::vector flat(static_cast(interleaved)); + const auto read_start = std::chrono::steady_clock::now(); + ggml_backend_tensor_get(graph.output, flat.data(), 0, flat.size() * sizeof(float)); + if (collect_timing) { + output_read_ms = engine::debug::elapsed_ms(read_start); + } + + // 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 * impl_->config.samples_per_frame; + if (impl_->config.channels == 1) { + return MossAudioTokenizerAudio{impl_->sampling_rate, {std::move(flat)}}; + } + std::vector> channels( + static_cast(impl_->config.channels), + std::vector(static_cast(per_channel))); + if (collect_timing) { + deinterleave_ms = engine::debug::measure_ms([&]() { +#ifdef _OPENMP +#pragma omp parallel for if(per_channel >= 4096) +#endif + for (int64_t i = 0; i < per_channel; ++i) { + for (int64_t channel = 0; channel < impl_->config.channels; ++channel) { + channels[static_cast(channel)][static_cast(i)] = + flat[static_cast(impl_->config.channels * i + channel)]; + } + } + }); + engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.dequant_ms", dequant_ms); + engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.latent_pack_ms", latent_pack_ms); + engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.graph_build_ms", graph_build_ms); + engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.input_upload_ms", input_upload_ms); + engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.graph_compute_ms", graph_compute_ms); + engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.output_read_ms", output_read_ms); + engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.deinterleave_ms", deinterleave_ms); + } else { +#ifdef _OPENMP +#pragma omp parallel for if(per_channel >= 4096) +#endif + for (int64_t i = 0; i < per_channel; ++i) { + for (int64_t channel = 0; channel < impl_->config.channels; ++channel) { + channels[static_cast(channel)][static_cast(i)] = + flat[static_cast(impl_->config.channels * i + channel)]; + } + } + } + return MossAudioTokenizerAudio{impl_->sampling_rate, std::move(channels)}; +} + +void MossAudioTokenizerDecoder::release_runtime_graphs() { + impl_->release_runtime_graphs(); +} + +// MOSS-Audio-Tokenizer v1 - 24 kHz mono, hop 1920, four transformer stages per side. +// Same layer math as v2 with two fewer stages, and one 10 s attention window per stage +// rather than v2's per-stage durations, so the context lengths are just the stage frame +// rate times ten. +MossAudioTokenizerConfig moss_audio_tokenizer_v1_config() { + MossAudioTokenizerConfig config; + config.sampling_rate = 24000; + config.samples_per_frame = 1920; + config.channels = 1; + config.quantizer = MossAudioTokenizerQuantizerConfig{ + 1024, + 8, + 512, + 768, + 32, + }; + config.encoder_stages = { + {240, 384, 768, 12, 12, 3072, 1000, 240}, + {768, 384, 768, 12, 12, 3072, 500, 2}, + {768, 640, 768, 12, 12, 3072, 250, 2}, + {1280, 768, 1280, 20, 32, 5120, 125, 2}, + }; + config.decoder_stages = { + {768, 1280, 1280, 20, 32, 5120, 125, 2}, + {640, 768, 768, 12, 12, 3072, 250, 2}, + {384, 768, 768, 12, 12, 3072, 500, 2}, + {384, 240, 768, 12, 12, 3072, 1000, 240}, + }; + config.encoder_module_start = 1; + config.encoder_module_stride = 2; + config.decoder_module_start = 0; + config.decoder_module_stride = 2; + return config; +} + +MossAudioTokenizerConfig moss_audio_tokenizer_v2_config() { + MossAudioTokenizerConfig config; + config.sampling_rate = 48000; + config.samples_per_frame = 3840; + config.quantizer = MossAudioTokenizerQuantizerConfig{ + 1024, + 8, + 512, + 768, + 12, + }; + config.encoder_stages = { + {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}, + }; + config.decoder_stages = { + {768, 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}, + }; + config.encoder_module_start = 1; + config.encoder_module_stride = 2; + config.decoder_module_start = 0; + config.decoder_module_stride = 2; + return config; +} + +MossAudioTokenizerConfig moss_audio_tokenizer_nano_config() { + MossAudioTokenizerConfig config; + config.sampling_rate = 48000; + config.samples_per_frame = 3840; + config.quantizer = MossAudioTokenizerQuantizerConfig{ + 1024, + 8, + 512, + 768, + 16, + }; + config.encoder_stages = { + {240, 384, 256, 4, 4, 1024, 1600, 240}, + {768, 384, 256, 4, 2, 1024, 1200, 2}, + {768, 384, 256, 4, 2, 1024, 800, 2}, + {768, 192, 256, 4, 4, 1024, 500, 2}, + }; + config.decoder_stages = { + {192, 768, 256, 4, 4, 1024, 1000, 2}, + {384, 768, 256, 4, 2, 1024, 1600, 2}, + {384, 768, 256, 4, 2, 1024, 2400, 2}, + {384, 240, 256, 4, 4, 1024, 3200, 240}, + }; + config.encoder_final_patch = 4; + config.decoder_initial_patch = 4; + config.encoder_module_start = 1; + config.encoder_module_stride = 2; + config.decoder_module_start = 1; + config.decoder_module_stride = 2; + return config; +} + +struct MossAudioTokenizerCodecRuntime::Impl { + Impl( + std::shared_ptr source_in, + core::ExecutionContext & decode_context_in, + int64_t num_quantizers_in, + MossAudioTokenizerCodecRuntimeOptions options_in, + MossAudioTokenizerConfig config_in) + : source(std::move(source_in)), + decode_context(decode_context_in), + num_quantizers(num_quantizers_in), + options(options_in), + config(std::move(config_in)) {} + + std::shared_ptr source; + core::ExecutionContext & decode_context; + int64_t num_quantizers = 0; + MossAudioTokenizerCodecRuntimeOptions options; + MossAudioTokenizerConfig config; + std::shared_ptr quantizer; + std::unique_ptr encode_context; + std::unique_ptr encoder; + std::unique_ptr decoder; + + core::ExecutionContext & encoder_execution_context() { + if (!options.separate_encoder_context) { + return decode_context; + } + if (encode_context == nullptr) { + encode_context = std::make_unique(decode_context.config()); + } + return *encode_context; + } + + MossAudioTokenizerEncoder & require_encoder() { + if (encoder == nullptr) { + encoder = std::make_unique( + *source, + quantizer, + encoder_execution_context(), + options.weight_context_bytes, + options.encoder_graph_arena_bytes, + config); + } + return *encoder; + } + + MossAudioTokenizerDecoder & require_decoder() { + if (decoder == nullptr) { + decoder = std::make_unique( + *source, + quantizer, + decode_context, + options.weight_context_bytes, + options.decoder_graph_arena_bytes, + config); + } + return *decoder; + } +}; + +MossAudioTokenizerCodecRuntime::MossAudioTokenizerCodecRuntime( + std::shared_ptr source, + core::ExecutionContext & execution_context, + int64_t num_quantizers, + MossAudioTokenizerCodecRuntimeOptions options, + MossAudioTokenizerConfig config) + : impl_(std::make_unique( + std::move(source), + execution_context, + num_quantizers, + options, + std::move(config))) { + if (impl_->source == nullptr) { + throw std::runtime_error("MOSS audio tokenizer codec requires weights"); + } + if (impl_->num_quantizers <= 0) { + throw std::runtime_error("MOSS audio tokenizer codec requires a positive quantizer count"); + } + impl_->quantizer = std::make_shared( + *impl_->source, + impl_->num_quantizers, + impl_->config.quantizer); +} + +MossAudioTokenizerCodecRuntime::~MossAudioTokenizerCodecRuntime() = default; + +int64_t MossAudioTokenizerCodecRuntime::sampling_rate() const noexcept { + return impl_->config.sampling_rate; +} + +void MossAudioTokenizerCodecRuntime::prepare_encoder() { + (void) impl_->require_encoder(); +} + +void MossAudioTokenizerCodecRuntime::prepare_decoder() { + (void) impl_->require_decoder(); +} + +MossAudioTokenizerCodes MossAudioTokenizerCodecRuntime::encode(const MossAudioTokenizerAudio & audio) { + return impl_->require_encoder().encode(audio); +} + +MossAudioTokenizerAudio MossAudioTokenizerCodecRuntime::decode(const MossAudioTokenizerCodes & codes) { + return impl_->require_decoder().decode(codes); +} + +void MossAudioTokenizerCodecRuntime::release_runtime_graphs() { + if (impl_->encoder != nullptr) { + impl_->encoder->release_runtime_graphs(); + } + if (impl_->decoder != nullptr) { + impl_->decoder->release_runtime_graphs(); + } +} + +} // namespace engine::codecs diff --git a/src/framework/modules/multi_codebook_embedding.cpp b/src/framework/modules/multi_codebook_embedding.cpp new file mode 100644 index 000000000..90371e60e --- /dev/null +++ b/src/framework/modules/multi_codebook_embedding.cpp @@ -0,0 +1,64 @@ +#include "engine/framework/modules/multi_codebook_embedding.h" + +#include + +namespace engine::modules { + +MultiCodebookEmbedding::MultiCodebookEmbedding(const assets::TensorSource & source, MultiCodebookEmbeddingSpec spec) + : hidden_size_(spec.hidden_size), + num_codebooks_(spec.num_codebooks), + pad_token_id_(static_cast(spec.pad_token_id)) { + if (hidden_size_ <= 0 || num_codebooks_ <= 0) { + throw std::runtime_error("multi-codebook embedding requires positive dimensions"); + } + if (!spec.codebook_sizes.empty() && static_cast(spec.codebook_sizes.size()) != num_codebooks_) { + throw std::runtime_error("multi-codebook embedding codebook size count mismatch"); + } + embeddings_.reserve(static_cast(num_codebooks_)); + for (int64_t codebook = 0; codebook < num_codebooks_; ++codebook) { + const int64_t size = spec.codebook_sizes.empty() + ? spec.vocab_size + : spec.codebook_sizes[static_cast(codebook)]; + if (size <= 0) { + throw std::runtime_error("multi-codebook embedding has an invalid codebook size"); + } + embeddings_.push_back(source.require_f32( + spec.tensor_prefix + "." + std::to_string(codebook) + ".weight", {size, hidden_size_})); + } +} + +int64_t MultiCodebookEmbedding::codebook_size(int64_t codebook) const { + if (codebook < 0 || codebook >= num_codebooks_) { + throw std::runtime_error("multi-codebook embedding index is out of range"); + } + return static_cast(embeddings_[static_cast(codebook)].size()) / hidden_size_; +} + +const float * MultiCodebookEmbedding::embedding(int64_t codebook, int32_t code) const { + const int64_t size = codebook_size(codebook); + if (code < 0 || code >= size) { + throw std::runtime_error("multi-codebook embedding code is out of range"); + } + return embeddings_[static_cast(codebook)].data() + static_cast(code) * hidden_size_; +} + +void MultiCodebookEmbedding::add_bias(const int32_t * codes, float * bias) const { + for (int64_t codebook = 0; codebook < num_codebooks_; ++codebook) { + const int32_t code = codes[codebook]; + if (code == pad_token_id_) { + continue; + } + const float * row = embedding(codebook, code); + for (int64_t index = 0; index < hidden_size_; ++index) { + bias[static_cast(index)] += row[index]; + } + } +} + +std::vector MultiCodebookEmbedding::bias_for(const int32_t * codes) const { + std::vector bias(static_cast(hidden_size_), 0.0F); + add_bias(codes, bias.data()); + return bias; +} + +} // namespace engine::modules diff --git a/src/models/moss/moss_tts_local/generator.cpp b/src/models/moss/moss_tts_local/generator.cpp index de04691ed..8af860317 100644 --- a/src/models/moss/moss_tts_local/generator.cpp +++ b/src/models/moss/moss_tts_local/generator.cpp @@ -195,13 +195,13 @@ MossGenerator::MossGenerator( throw std::runtime_error("MOSS-TTS-Local generator only supports the binary local text head"); } const auto & source = *assets_->model_weights; - moss::AudioCodebookSpec codebooks; + engine::modules::MultiCodebookEmbeddingSpec codebooks; codebooks.hidden_size = hidden_size_; codebooks.num_codebooks = num_codebooks_; - codebooks.audio_vocab_size = config.audio_vocab_size; - codebooks.audio_codebook_sizes = config.audio_codebook_sizes; - codebooks.audio_pad_token_id = config.audio_pad_token_id; - audio_codebooks_ = std::make_unique(source, std::move(codebooks)); + codebooks.vocab_size = config.audio_vocab_size; + codebooks.codebook_sizes = config.audio_codebook_sizes; + codebooks.pad_token_id = config.audio_pad_token_id; + audio_codebooks_ = std::make_unique(source, std::move(codebooks)); local_text_head_ = source.require_f32("local_text_lm_head.weight", {2, hidden_size_}); projection_ = std::make_unique( *assets_, diff --git a/src/models/moss/moss_tts_local/session.cpp b/src/models/moss/moss_tts_local/session.cpp index 90ff4b615..4bf468479 100644 --- a/src/models/moss/moss_tts_local/session.cpp +++ b/src/models/moss/moss_tts_local/session.cpp @@ -29,7 +29,6 @@ constexpr size_t kGeneratorProjectionWeightContextBytes = 16ull * 1024 * 1024; constexpr size_t kGeneratorProjectionGraphArenaBytes = 16ull * 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; constexpr int64_t kDefaultTextChunkSize = 2048; @@ -246,12 +245,16 @@ MossTTSLocalSession::MossTTSLocalSession( depth_ = std::make_unique( assets_, execution_context(), kDepthGraphArenaBytes, kDepthWeightContextBytes); processor_ = std::make_unique(assets_); - codec_ = std::make_unique( - *assets_->audio_tokenizer_weights, + codec_ = std::make_unique( + assets_->audio_tokenizer_weights, execution_context(), assets_->config.num_codebooks, - kCodecWeightContextBytes, - kCodecGraphArenaBytes); + engine::codecs::MossAudioTokenizerCodecRuntimeOptions{ + kCodecWeightContextBytes, + kEncoderGraphArenaBytes, + kCodecGraphArenaBytes, + true, + }); generator_ = std::make_unique( assets_, execution_context(), @@ -259,6 +262,7 @@ MossTTSLocalSession::MossTTSLocalSession( kGeneratorProjectionWeightContextBytes, *backbone_, *depth_); + codec_->prepare_decoder(); assets_->model_weights->release_storage(); } @@ -278,24 +282,11 @@ void MossTTSLocalSession::prepare(const runtime::SessionPreparationRequest & req const bool has_reference = request.voice.has_value() && request.voice->speaker.has_value() && request.voice->speaker->audio.has_value(); if (has_reference) { - (void) encoder(); + codec_->prepare_encoder(); } mark_prepared(); } -moss::MossAudioTokenizerEncoder & MossTTSLocalSession::encoder() { - if (encoder_ == nullptr) { - reference_encoder_execution_context_ = std::make_unique(options().backend); - encoder_ = std::make_unique( - *assets_->audio_tokenizer_weights, - *reference_encoder_execution_context_, - assets_->config.num_codebooks, - kEncoderWeightContextBytes, - kEncoderGraphArenaBytes); - } - return *encoder_; -} - bool MossTTSLocalSession::ReferenceAudioCacheKeyEqual::operator()( const ReferenceAudioCacheKey & lhs, const ReferenceAudioCacheKey & rhs) const { @@ -364,7 +355,10 @@ runtime::TaskResult MossTTSLocalSession::run(const runtime::TaskRequest & reques stereo = reference_to_codec_stereo(reference_audio); }); time_once(reference_encode_ms, [&]() { - reference_codes = encoder().encode(stereo); + reference_codes = codec_->encode(engine::codecs::MossAudioTokenizerAudio{ + codec_->sampling_rate(), + std::move(stereo), + }).codebooks; }); ReferenceVoiceCacheEntry entry; entry.codes = reference_codes; @@ -443,7 +437,10 @@ runtime::TaskResult MossTTSLocalSession::run(const runtime::TaskRequest & reques std::vector> channels; time_once(codec_decode_ms, [&]() { - channels = codec_->decode(codes); + channels = codec_->decode(engine::codecs::MossAudioTokenizerCodes{ + static_cast(codes.empty() ? 0 : codes.front().size()), + std::move(codes), + }).channels; }); const int channel_count = static_cast(channels.size()); const size_t samples_per_channel = channels.empty() ? 0 : channels.front().size(); diff --git a/src/models/moss/moss_tts_nano/session.cpp b/src/models/moss/moss_tts_nano/session.cpp index e724d5fa7..437fcdea7 100644 --- a/src/models/moss/moss_tts_nano/session.cpp +++ b/src/models/moss/moss_tts_nano/session.cpp @@ -4,7 +4,7 @@ #include "engine/framework/debug/profiler.h" #include "engine/framework/runtime/options.h" #include "engine/framework/text/chunking.h" -#include "engine/models/moss/shared/audio_tokenizer_config.h" +#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" #include #include @@ -205,19 +205,24 @@ MossTTSNanoSession::MossTTSNanoSession( local_frame_weight_context_bytes_, local_frame_weight_storage_type_), generator_(global_transformer_, local_frame_decoder_), - decoder_( - *assets_->audio_tokenizer_weights, + codec_( + assets_->audio_tokenizer_weights, execution_context(), assets_->config.n_vq, - audio_tokenizer_weight_context_bytes_, - audio_tokenizer_decoder_graph_arena_bytes_, - moss::moss_audio_tokenizer_nano_config()) { + engine::codecs::MossAudioTokenizerCodecRuntimeOptions{ + audio_tokenizer_weight_context_bytes_, + audio_tokenizer_encoder_graph_arena_bytes_, + audio_tokenizer_decoder_graph_arena_bytes_, + true, + }, + engine::codecs::moss_audio_tokenizer_nano_config()) { if (task_.task != runtime::VoiceTaskKind::Tts && task_.task != runtime::VoiceTaskKind::VoiceCloning) { throw std::runtime_error("MOSS-TTS-Nano only supports the Tts and VoiceCloning tasks"); } if (task_.mode != runtime::RunMode::Offline) { throw std::runtime_error("MOSS-TTS-Nano currently supports offline sessions"); } + codec_.prepare_decoder(); for (const auto & [key, value] : options.options) { (void) value; if (key.rfind("moss_tts_nano.", 0) == 0 && @@ -267,29 +272,19 @@ void MossTTSNanoSession::prepare(const runtime::SessionPreparationRequest & requ mark_prepared(); } -moss::MossAudioTokenizerEncoder & MossTTSNanoSession::encoder() { - if (encoder_ == nullptr) { - reference_encoder_execution_context_ = std::make_unique(options().backend); - encoder_ = std::make_unique( - *assets_->audio_tokenizer_weights, - *reference_encoder_execution_context_, - assets_->config.n_vq, - audio_tokenizer_weight_context_bytes_, - audio_tokenizer_encoder_graph_arena_bytes_, - moss::moss_audio_tokenizer_nano_config()); - } - return *encoder_; -} - MossTTSNanoAudioCodes MossTTSNanoSession::encode_reference_audio( const runtime::AudioBuffer & audio, int64_t active_codebooks) { - const auto stereo = reference_to_audio_tokenizer_stereo(audio); - const auto codes = encoder().encode(stereo); + auto stereo = reference_to_audio_tokenizer_stereo(audio); + const auto encoded = codec_.encode(engine::codecs::MossAudioTokenizerAudio{ + codec_.sampling_rate(), + std::move(stereo), + }); + const auto & codes = encoded.codebooks; if (static_cast(codes.size()) < active_codebooks) { throw std::runtime_error("MOSS-TTS-Nano reference encoder returned too few codebooks"); } - const int64_t frames = codes.empty() ? 0 : static_cast(codes.front().size()); + const int64_t frames = encoded.frames; if (frames <= 0) { throw std::runtime_error("MOSS-TTS-Nano reference encoder returned no frames"); } @@ -327,14 +322,18 @@ runtime::AudioBuffer MossTTSNanoSession::decode_generated_audio( codes.token_ids[static_cast(frame * codes.codebooks + codebook)]; } } - const auto channels = decoder_.decode(tokenizer_codes); + const auto decoded = codec_.decode(engine::codecs::MossAudioTokenizerCodes{ + codes.frames, + std::move(tokenizer_codes), + }); + const auto & channels = decoded.channels; const int channel_count = static_cast(channels.size()); const size_t samples_per_channel = channels.empty() ? 0 : channels.front().size(); if (channel_count <= 0 || samples_per_channel == 0) { throw std::runtime_error("MOSS-TTS-Nano audio tokenizer decoder produced no audio"); } runtime::AudioBuffer audio; - audio.sample_rate = static_cast(decoder_.sampling_rate()); + audio.sample_rate = static_cast(decoded.sampling_rate); audio.channels = channel_count; audio.samples.resize(samples_per_channel * static_cast(channel_count)); const int64_t sample_count = static_cast(samples_per_channel); diff --git a/src/models/moss/shared/audio_tokenizer_config.cpp b/src/models/moss/shared/audio_tokenizer_config.cpp deleted file mode 100644 index ec887745a..000000000 --- a/src/models/moss/shared/audio_tokenizer_config.cpp +++ /dev/null @@ -1,106 +0,0 @@ -#include "engine/models/moss/shared/audio_tokenizer_config.h" - -namespace engine::models::moss { - -// MOSS-Audio-Tokenizer v1 - 24 kHz mono, hop 1920, four transformer stages per side. -// Same layer math as v2 with two fewer stages, and one 10 s attention window per stage -// rather than v2's per-stage durations, so the context lengths are just the stage frame -// rate times ten. -AudioTokenizerConfig moss_audio_tokenizer_v1_config() { - AudioTokenizerConfig config; - config.sampling_rate = 24000; - config.samples_per_frame = 1920; - config.channels = 1; - config.quantizer = AudioTokenizerQuantizerConfig{ - 1024, - 8, - 512, - 768, - 32, - }; - config.encoder_stages = { - {240, 384, 768, 12, 12, 3072, 1000, 240}, - {768, 384, 768, 12, 12, 3072, 500, 2}, - {768, 640, 768, 12, 12, 3072, 250, 2}, - {1280, 768, 1280, 20, 32, 5120, 125, 2}, - }; - config.decoder_stages = { - {768, 1280, 1280, 20, 32, 5120, 125, 2}, - {640, 768, 768, 12, 12, 3072, 250, 2}, - {384, 768, 768, 12, 12, 3072, 500, 2}, - {384, 240, 768, 12, 12, 3072, 1000, 240}, - }; - config.encoder_module_start = 1; - config.encoder_module_stride = 2; - config.decoder_module_start = 0; - config.decoder_module_stride = 2; - return config; -} - -AudioTokenizerConfig moss_audio_tokenizer_v2_config() { - AudioTokenizerConfig config; - config.sampling_rate = 48000; - config.samples_per_frame = 3840; - config.quantizer = AudioTokenizerQuantizerConfig{ - 1024, - 8, - 512, - 768, - 12, - }; - config.encoder_stages = { - {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}, - }; - config.decoder_stages = { - {768, 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}, - }; - config.encoder_module_start = 1; - config.encoder_module_stride = 2; - config.decoder_module_start = 0; - config.decoder_module_stride = 2; - return config; -} - -AudioTokenizerConfig moss_audio_tokenizer_nano_config() { - AudioTokenizerConfig config; - config.sampling_rate = 48000; - config.samples_per_frame = 3840; - config.quantizer = AudioTokenizerQuantizerConfig{ - 1024, - 8, - 512, - 768, - 16, - }; - config.encoder_stages = { - {240, 384, 256, 4, 4, 1024, 1600, 240}, - {768, 384, 256, 4, 2, 1024, 1200, 2}, - {768, 384, 256, 4, 2, 1024, 800, 2}, - {768, 192, 256, 4, 4, 1024, 500, 2}, - }; - config.decoder_stages = { - {192, 768, 256, 4, 4, 1024, 1000, 2}, - {384, 768, 256, 4, 2, 1024, 1600, 2}, - {384, 768, 256, 4, 2, 1024, 2400, 2}, - {384, 240, 256, 4, 4, 1024, 3200, 240}, - }; - config.encoder_final_patch = 4; - config.decoder_initial_patch = 4; - config.encoder_module_start = 1; - config.encoder_module_stride = 2; - config.decoder_module_start = 1; - config.decoder_module_stride = 2; - return config; -} - -} // namespace engine::models::moss diff --git a/src/models/moss/shared/audio_tokenizer_decoder.cpp b/src/models/moss/shared/audio_tokenizer_decoder.cpp deleted file mode 100644 index 933333ed8..000000000 --- a/src/models/moss/shared/audio_tokenizer_decoder.cpp +++ /dev/null @@ -1,339 +0,0 @@ -#include "engine/models/moss/shared/audio_tokenizer_decoder.h" - -#include "engine/framework/core/backend.h" -#include "engine/framework/core/module.h" -#include "engine/framework/debug/profiler.h" -#include "engine/models/moss/shared/audio_tokenizer_quantizer.h" -#include "engine/models/moss/shared/audio_tokenizer_transformer.h" - -#include -#include - -#include -#include -#include -#include -#include -#include - -namespace engine::models::moss { -namespace { - -namespace cd = codec_detail; - -constexpr int64_t kAttentionQueryChunk = 1500; - -cd::TransformerSpec to_transformer_spec(const AudioTokenizerTransformerStage & stage) { - return { - stage.input_dimension, - stage.output_dimension, - stage.model_dimension, - stage.num_heads, - stage.num_layers, - stage.feedforward_dimension, - stage.context_window, - stage.patch_size, - }; -} - -// 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 = core::ensure_backend_addressable_layout(ctx, input); - const int64_t length = contiguous.shape.dims[1]; - const int64_t packed = contiguous.shape.dims[2]; - const int64_t channels = packed / patch; - auto reshaped = engine::modules::ReshapeModule({ - core::TensorShape::from_dims({1, length, channels, patch}), - }).build(ctx, contiguous); - auto transposed = engine::modules::TransposeModule({{0, 1, 3, 2}, reshaped.shape.rank}).build(ctx, reshaped); - return engine::modules::ReshapeModule({ - core::TensorShape::from_dims({1, length * patch, channels}), - }).build(ctx, core::ensure_backend_addressable_layout(ctx, transposed)); -} - -} // namespace - -struct MossAudioTokenizerDecoder::Impl { - ggml_backend_t backend = nullptr; - core::BackendType backend_type = core::BackendType::Cpu; - int threads = 1; - int64_t sampling_rate = 48000; - size_t graph_arena_bytes = 0; - AudioTokenizerConfig config; - std::unique_ptr dequantizer; - std::unique_ptr store; - std::vector transformers; -}; - -MossAudioTokenizerDecoder::MossAudioTokenizerDecoder( - const assets::TensorSource & source, - core::ExecutionContext & execution_context, - int64_t num_quantizers, - size_t weight_context_bytes, - size_t graph_arena_bytes, - AudioTokenizerConfig config) - : 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_->threads = execution_context.config().threads; - impl_->config = config; - impl_->sampling_rate = config.sampling_rate; - impl_->graph_arena_bytes = graph_arena_bytes; - impl_->dequantizer = std::make_unique(source, num_quantizers, config.quantizer); - - cd::CodecWeights weights(source); - impl_->store = std::make_unique( - impl_->backend, impl_->backend_type, "moss.audio_tokenizer.decoder", weight_context_bytes); - impl_->transformers.reserve(config.decoder_stages.size()); - for (size_t index = 0; index < config.decoder_stages.size(); ++index) { - const int64_t module_index = - config.decoder_module_start + static_cast(index) * config.decoder_module_stride; - impl_->transformers.push_back(cd::load_transformer( - *impl_->store, weights, to_transformer_spec(config.decoder_stages[index]), "decoder", module_index)); - } - impl_->store->upload(); -} - -MossAudioTokenizerDecoder::~MossAudioTokenizerDecoder() = default; - -int64_t MossAudioTokenizerDecoder::sampling_rate() const noexcept { - return impl_->sampling_rate; -} - -std::vector> MossAudioTokenizerDecoder::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. - double dequant_ms = 0.0; - double latent_pack_ms = 0.0; - double graph_build_ms = 0.0; - double input_upload_ms = 0.0; - double graph_compute_ms = 0.0; - double output_read_ms = 0.0; - double deinterleave_ms = 0.0; - const bool collect_timing = engine::debug::timing_log_enabled(); - std::vector latent; - if (collect_timing) { - dequant_ms = engine::debug::measure_ms([&]() { - latent = impl_->dequantizer->decode(codes); - }); - } else { - latent = impl_->dequantizer->decode(codes); - } - std::vector latent_input; - if (collect_timing) { - latent_pack_ms = engine::debug::measure_ms([&]() { - latent_input.resize(static_cast(frames * cd::kCodeDim)); -#ifdef _OPENMP -#pragma omp parallel for collapse(2) if(frames * cd::kCodeDim >= 4096) -#endif - for (int64_t channel = 0; channel < cd::kCodeDim; ++channel) { - for (int64_t step = 0; step < frames; ++step) { - latent_input[static_cast(step * cd::kCodeDim + channel)] = - latent[static_cast(channel * frames + step)]; - } - } - }); - } else { - latent_input.resize(static_cast(frames * cd::kCodeDim)); -#ifdef _OPENMP -#pragma omp parallel for collapse(2) if(frames * cd::kCodeDim >= 4096) -#endif - for (int64_t channel = 0; channel < cd::kCodeDim; ++channel) { - for (int64_t step = 0; step < frames; ++step) { - latent_input[static_cast(step * cd::kCodeDim + channel)] = - latent[static_cast(channel * frames + step)]; - } - } - } - - const auto graph_build_start = std::chrono::steady_clock::now(); - 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.audio_tokenizer.decode", impl_->backend_type}; - - auto latent_tensor = - core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, frames, cd::kCodeDim})); - ggml_set_input(latent_tensor.tensor); - - struct StageInput { - struct MaskInput { - ggml_tensor * tensor; - std::vector host; - }; - - ggml_tensor * positions; - std::vector position_host; - std::vector masks; - }; - std::vector stage_inputs; - stage_inputs.reserve(impl_->transformers.size()); - - auto hidden = latent_tensor; - int64_t steps = frames; - if (impl_->config.decoder_initial_patch > 1) { - hidden = patch_upsample(ctx, hidden, impl_->config.decoder_initial_patch); - steps *= impl_->config.decoder_initial_patch; - } - 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); - const bool use_windowed_attention = steps > kAttentionQueryChunk; - core::TensorValue mask; - if (!use_windowed_attention) { - 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); - } - std::vector windows; - if (use_windowed_attention) { - for (int64_t query_start = 0; query_start < steps; query_start += kAttentionQueryChunk) { - const int64_t query_steps = std::min(kAttentionQueryChunk, steps - query_start); - const int64_t key_start = std::max(0, query_start - transformer.spec.context + 1); - const int64_t key_steps = query_start + query_steps - key_start; - auto window_mask = core::make_tensor( - ctx, - GGML_TYPE_F32, - core::TensorShape::from_dims({1, 1, query_steps, key_steps})); - ggml_set_input(window_mask.tensor); - stage.masks.push_back(StageInput::MaskInput{ - window_mask.tensor, - cd::causal_context_mask_window( - query_start, - query_steps, - key_start, - key_steps, - transformer.spec.context), - }); - windows.push_back(cd::AttentionWindow{ - query_start, - query_steps, - key_start, - key_steps, - window_mask, - }); - } - } else { - stage.masks.push_back(StageInput::MaskInput{mask.tensor, cd::causal_context_mask(steps, transformer.spec.context)}); - } - stage_inputs.push_back(std::move(stage)); - - hidden = windows.empty() - ? cd::run_transformer(ctx, hidden, transformer, positions, mask, steps) - : cd::run_transformer(ctx, hidden, transformer, positions, windows.front().mask, steps, &windows); - hidden = patch_upsample(ctx, hidden, transformer.spec.patch); - steps *= transformer.spec.patch; - } - - hidden = core::ensure_backend_addressable_layout(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"); - } - if (collect_timing) { - graph_build_ms = engine::debug::elapsed_ms(graph_build_start); - } - - const auto upload_start = std::chrono::steady_clock::now(); - 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)); - for (const auto & mask : stage.masks) { - ggml_backend_tensor_set(mask.tensor, mask.host.data(), 0, mask.host.size() * sizeof(float)); - } - } - if (collect_timing) { - input_upload_ms = engine::debug::elapsed_ms(upload_start); - } - - const auto compute_start = std::chrono::steady_clock::now(); - core::set_backend_threads(impl_->backend, impl_->threads); - const ggml_status status = ggml_backend_graph_compute(impl_->backend, graph); - ggml_backend_synchronize(impl_->backend); - if (collect_timing) { - graph_compute_ms = engine::debug::elapsed_ms(compute_start); - } - 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 * samples_per_frame * channels - std::vector flat(static_cast(interleaved)); - const auto read_start = std::chrono::steady_clock::now(); - ggml_backend_tensor_get(hidden.tensor, flat.data(), 0, flat.size() * sizeof(float)); - ggml_gallocr_free(gallocr); - if (collect_timing) { - output_read_ms = engine::debug::elapsed_ms(read_start); - } - - // 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 * impl_->config.samples_per_frame; - if (impl_->config.channels == 1) { - // Mono codecs (v1) emit the waveform directly; there is nothing to de-interleave. - return {std::move(flat)}; - } - std::vector> stereo(2, std::vector(static_cast(per_channel))); - if (collect_timing) { - deinterleave_ms = engine::debug::measure_ms([&]() { -#ifdef _OPENMP -#pragma omp parallel for if(per_channel >= 4096) -#endif - 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)]; - } - }); - engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.dequant_ms", dequant_ms); - engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.latent_pack_ms", latent_pack_ms); - engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.graph_build_ms", graph_build_ms); - engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.input_upload_ms", input_upload_ms); - engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.graph_compute_ms", graph_compute_ms); - engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.output_read_ms", output_read_ms); - engine::debug::timing_log_scalar("moss.audio_tokenizer.decode.deinterleave_ms", deinterleave_ms); - } else { -#ifdef _OPENMP -#pragma omp parallel for if(per_channel >= 4096) -#endif - 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 diff --git a/src/models/moss/shared/audio_tokenizer_encoder.cpp b/src/models/moss/shared/audio_tokenizer_encoder.cpp deleted file mode 100644 index 52e040a19..000000000 --- a/src/models/moss/shared/audio_tokenizer_encoder.cpp +++ /dev/null @@ -1,230 +0,0 @@ -#include "engine/models/moss/shared/audio_tokenizer_encoder.h" - -#include "engine/framework/core/backend.h" -#include "engine/framework/core/module.h" -#include "engine/models/moss/shared/audio_tokenizer_quantizer.h" -#include "engine/models/moss/shared/audio_tokenizer_transformer.h" - -#include -#include - -#include -#include -#include -#include - -namespace engine::models::moss { -namespace { - -namespace cd = codec_detail; - -cd::TransformerSpec to_transformer_spec(const AudioTokenizerTransformerStage & stage) { - return { - stage.input_dimension, - stage.output_dimension, - stage.model_dimension, - stage.num_heads, - stage.num_layers, - stage.feedforward_dimension, - stage.context_window, - stage.patch_size, - }; -} - -// 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 = core::ensure_backend_addressable_layout(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; - auto reshaped = engine::modules::ReshapeModule({ - core::TensorShape::from_dims({1, length, patch, channels}), - }).build(ctx, contiguous); - auto transposed = engine::modules::TransposeModule({{0, 1, 3, 2}, reshaped.shape.rank}).build(ctx, reshaped); - return engine::modules::ReshapeModule({ - core::TensorShape::from_dims({1, length, channels * patch}), - }).build(ctx, core::ensure_backend_addressable_layout(ctx, transposed)); -} - -} // namespace - -struct MossAudioTokenizerEncoder::Impl { - ggml_backend_t backend = nullptr; - core::BackendType backend_type = core::BackendType::Cpu; - int threads = 1; - int64_t samples_per_frame = 3840; - size_t graph_arena_bytes = 0; - AudioTokenizerConfig config; - std::unique_ptr quantizer; - std::unique_ptr store; - std::vector transformers; -}; - -MossAudioTokenizerEncoder::MossAudioTokenizerEncoder( - const assets::TensorSource & source, - core::ExecutionContext & execution_context, - int64_t num_quantizers, - size_t weight_context_bytes, - size_t graph_arena_bytes, - AudioTokenizerConfig config) - : 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_->threads = execution_context.config().threads; - impl_->config = config; - impl_->samples_per_frame = config.samples_per_frame; - impl_->graph_arena_bytes = graph_arena_bytes; - impl_->quantizer = std::make_unique(source, num_quantizers, config.quantizer); - - cd::CodecWeights weights(source); - impl_->store = std::make_unique( - impl_->backend, impl_->backend_type, "moss.audio_tokenizer.encoder", weight_context_bytes); - impl_->transformers.reserve(config.encoder_stages.size()); - for (size_t index = 0; index < config.encoder_stages.size(); ++index) { - const int64_t module_index = - config.encoder_module_start + static_cast(index) * config.encoder_module_stride; - impl_->transformers.push_back(cd::load_transformer( - *impl_->store, weights, to_transformer_spec(config.encoder_stages[index]), "encoder", module_index)); - } - impl_->store->upload(); -} - -MossAudioTokenizerEncoder::~MossAudioTokenizerEncoder() = default; - -std::vector> MossAudioTokenizerEncoder::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 for the encoder graph, - // but keep the official valid code length as floor(valid_samples / samples_per_frame). - // MossAudioTokenizerPatchedPretransform pads the tensor and propagates input_lengths - // with integer division, then slices audio_codes to audio_codes_lengths. - const int64_t frames = (raw_per_channel + impl_->samples_per_frame - 1) / impl_->samples_per_frame; - const int64_t valid_frames = raw_per_channel / impl_->samples_per_frame; - const int64_t per_channel = frames * impl_->samples_per_frame; - const int64_t interleaved = per_channel * 2; - std::vector waveform(static_cast(interleaved), 0.0F); -#ifdef _OPENMP -#pragma omp parallel for if(raw_per_channel >= 4096) -#endif - 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.audio_tokenizer.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); - } - if (impl_->config.encoder_final_patch > 1) { - hidden = patch_downsample(ctx, hidden, impl_->config.encoder_final_patch); - steps /= impl_->config.encoder_final_patch; - } - - hidden = core::ensure_backend_addressable_layout(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)); - } - - core::set_backend_threads(impl_->backend, impl_->threads); - 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 MossAudioTokenizerQuantizer::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); - - if (valid_frames <= 0) { - throw std::runtime_error("MOSS codec encoder input is shorter than one codec frame"); - } - latent.resize(static_cast(valid_frames * cd::kCodeDim)); - return impl_->quantizer->encode(latent, valid_frames); -} - -} // namespace engine::models::moss diff --git a/src/models/moss/shared/audio_tokenizer_quantizer.cpp b/src/models/moss/shared/audio_tokenizer_quantizer.cpp deleted file mode 100644 index 9fedb85a4..000000000 --- a/src/models/moss/shared/audio_tokenizer_quantizer.cpp +++ /dev/null @@ -1,273 +0,0 @@ -#include "engine/models/moss/shared/audio_tokenizer_quantizer.h" - -#include "engine/framework/assets/tensor_source.h" - -#include -#include -#include -#include -#include -#include - -namespace engine::models::moss { -namespace { - -// 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)); -#ifdef _OPENMP -#pragma omp parallel for if(out_channels * in_channels >= 4096) -#endif - 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 assets::TensorSource & source, - const std::string & prefix, - int64_t out_channels, - int64_t in_channels) { - const auto g = source.require_f32(prefix + ".parametrizations.weight.original0"); - const auto v = source.require_f32(prefix + ".parametrizations.weight.original1"); - return reconstruct_weight_norm(g, v, out_channels, in_channels); -} - -} // namespace - -MossAudioTokenizerQuantizer::MossAudioTokenizerQuantizer( - const assets::TensorSource & source, - int64_t num_quantizers, - AudioTokenizerQuantizerConfig config) - : codebook_size_(config.codebook_size), - codebook_dim_(config.codebook_dim), - rvq_dim_(config.rvq_dim), - code_dim_(config.code_dim), - num_quantizers_(num_quantizers) { - if (num_quantizers_ <= 0) { - throw std::runtime_error("MOSS codec dequantizer requires a positive quantizer count"); - } - - output_weight_ = load_wn_conv_weight(source, "quantizer.output_proj", code_dim_, rvq_dim_); - output_bias_ = source.require_f32("quantizer.output_proj.bias"); - - 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 = source.require_f32(prefix + ".codebook.weight"); - codebook.out_weight = load_wn_conv_weight(source, prefix + ".out_proj", rvq_dim_, codebook_dim_); - codebook.out_bias = source.require_f32(prefix + ".out_proj.bias"); - std::vector combined_bias(static_cast(code_dim_)); - std::vector combined_weight(static_cast(code_dim_ * codebook_dim_)); -#ifdef _OPENMP -#pragma omp parallel for if(code_dim_ >= 256) -#endif - for (int64_t out = 0; out < code_dim_; ++out) { - const float * output_row = &output_weight_[static_cast(out * rvq_dim_)]; - float bias_sum = 0.0F; - for (int64_t rvq = 0; rvq < rvq_dim_; ++rvq) { - bias_sum += output_row[rvq] * codebook.out_bias[static_cast(rvq)]; - } - combined_bias[static_cast(out)] = bias_sum; - for (int64_t k = 0; k < codebook_dim_; ++k) { - float sum = 0.0F; - for (int64_t rvq = 0; rvq < rvq_dim_; ++rvq) { - sum += output_row[rvq] * codebook.out_weight[static_cast(rvq * codebook_dim_ + k)]; - } - combined_weight[static_cast(out * codebook_dim_ + k)] = sum; - } - } - codebook.latent_table.resize(static_cast(codebook_size_ * code_dim_)); -#ifdef _OPENMP -#pragma omp parallel for collapse(2) if(codebook_size_ * code_dim_ >= 4096) -#endif - for (int64_t code = 0; code < codebook_size_; ++code) { - for (int64_t out = 0; out < code_dim_; ++out) { - const float * embedding = &codebook.table[static_cast(code * codebook_dim_)]; - float sum = combined_bias[static_cast(out)]; - const float * row = &combined_weight[static_cast(out * codebook_dim_)]; - for (int64_t k = 0; k < codebook_dim_; ++k) { - sum += row[k] * embedding[k]; - } - codebook.latent_table[static_cast(code * code_dim_ + out)] = sum; - } - } - codebook.in_weight = load_wn_conv_weight(source, prefix + ".in_proj", codebook_dim_, rvq_dim_); - codebook.in_bias = source.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; -#ifdef _OPENMP -#pragma omp parallel for if(codebook_size_ * codebook_dim_ >= 4096) -#endif - 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)); - } - - input_weight_ = load_wn_conv_weight(source, "quantizer.input_proj", rvq_dim_, code_dim_); - input_bias_ = source.require_f32("quantizer.input_proj.bias"); -} - -std::vector MossAudioTokenizerQuantizer::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"); - } - for (int64_t step = 0; step < steps; ++step) { - for (int64_t index = 0; index < num_quantizers_; ++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"); - } - } - } - - std::vector latent(static_cast(code_dim_ * steps)); -#ifdef _OPENMP -#pragma omp parallel for if(steps * code_dim_ >= 4096) -#endif - for (int64_t step = 0; step < steps; ++step) { - for (int64_t out = 0; out < code_dim_; ++out) { - float value = output_bias_[static_cast(out)]; - 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)]; - const float * decoded = &codebook.latent_table[static_cast(code * code_dim_)]; - value += decoded[static_cast(out)]; - } - latent[static_cast(out * steps + step)] = value; - } - } - return latent; -} - -std::vector> MossAudioTokenizerQuantizer::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))); - - const auto encode_frame = [&](int64_t step, std::vector & residual, std::vector & encoding) { - std::fill(residual.begin(), residual.end(), 0.0); - std::fill(encoding.begin(), encoding.end(), 0.0); - - // 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; - } - } - }; - -#ifdef _OPENMP - if (frames >= 8) { -#pragma omp parallel - { - std::vector residual(static_cast(rvq_dim_)); - std::vector encoding(static_cast(codebook_dim_)); -#pragma omp for - for (int64_t step = 0; step < frames; ++step) { - encode_frame(step, residual, encoding); - } - } - } else -#endif - { - std::vector residual(static_cast(rvq_dim_)); - std::vector encoding(static_cast(codebook_dim_)); - for (int64_t step = 0; step < frames; ++step) { - encode_frame(step, residual, encoding); - } - } - return codes; -} - -} // namespace engine::models::moss diff --git a/src/models/moss/shared/token_rows.cpp b/src/models/moss/shared/token_rows.cpp index 698e5edef..0d55e6caa 100644 --- a/src/models/moss/shared/token_rows.cpp +++ b/src/models/moss/shared/token_rows.cpp @@ -1,6 +1,5 @@ #include "engine/models/moss/shared/token_rows.h" -#include #include namespace engine::models::moss { @@ -52,58 +51,4 @@ TokenRows TokenRowBuilder::finish() { return std::move(rows_); } -AudioCodebookEmbeddings::AudioCodebookEmbeddings(const assets::TensorSource & source, AudioCodebookSpec spec) - : hidden_size_(spec.hidden_size), - num_codebooks_(spec.num_codebooks), - audio_pad_token_id_(static_cast(spec.audio_pad_token_id)) { - if (hidden_size_ <= 0 || num_codebooks_ <= 0) { - throw std::runtime_error("MOSS audio codebook embeddings require positive dimensions"); - } - embeddings_.reserve(static_cast(num_codebooks_)); - for (int64_t codebook = 0; codebook < num_codebooks_; ++codebook) { - const int64_t size = spec.audio_codebook_sizes.empty() - ? spec.audio_vocab_size - : spec.audio_codebook_sizes[static_cast(codebook)]; - if (size <= 0) { - throw std::runtime_error("MOSS audio codebook has an invalid size"); - } - embeddings_.push_back(source.require_f32( - spec.tensor_prefix + "." + std::to_string(codebook) + ".weight", {size, hidden_size_})); - } -} - -int64_t AudioCodebookEmbeddings::codebook_size(int64_t codebook) const { - if (codebook < 0 || codebook >= num_codebooks_) { - throw std::runtime_error("MOSS audio codebook index is out of range"); - } - return static_cast(embeddings_[static_cast(codebook)].size()) / hidden_size_; -} - -const float * AudioCodebookEmbeddings::embedding(int64_t codebook, int32_t code) const { - const int64_t size = codebook_size(codebook); - if (code < 0 || code >= size) { - throw std::runtime_error("MOSS audio code is out of range"); - } - return embeddings_[static_cast(codebook)].data() + static_cast(code) * hidden_size_; -} - -void AudioCodebookEmbeddings::add_bias(const int32_t * codes, float * bias) const { - for (int64_t codebook = 0; codebook < num_codebooks_; ++codebook) { - const int32_t code = codes[codebook]; - if (code == audio_pad_token_id_) { - continue; - } - const float * row = embedding(codebook, code); - for (int64_t index = 0; index < hidden_size_; ++index) { - bias[static_cast(index)] += row[index]; - } - } -} - -std::vector AudioCodebookEmbeddings::bias_for(const int32_t * codes) const { - std::vector bias(static_cast(hidden_size_), 0.0F); - add_bias(codes, bias.data()); - return bias; -} - } // namespace engine::models::moss diff --git a/tests/moss_tts_local/codec_decode_parity.cpp b/tests/moss_tts_local/codec_decode_parity.cpp index 93eef2907..97d201534 100644 --- a/tests/moss_tts_local/codec_decode_parity.cpp +++ b/tests/moss_tts_local/codec_decode_parity.cpp @@ -5,7 +5,7 @@ #include "engine/framework/core/backend.h" #include "engine/framework/core/execution_context.h" -#include "engine/models/moss/shared/audio_tokenizer_decoder.h" +#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" #include #include @@ -16,6 +16,7 @@ #include #include #include +#include #include namespace { @@ -101,11 +102,21 @@ int main(int argc, char ** argv) { std::cout << "codec=" << codec_dir.string() << "\n"; std::cout << "loading decoder weights...\n" << std::flush; - engine::models::moss::MossAudioTokenizerDecoder decoder( - codec_dir, execution_context, num_quantizers, kWeightContextBytes, kGraphArenaBytes); + auto codec_weights = engine::assets::open_tensor_source(codec_dir); + engine::codecs::MossAudioTokenizerCodecRuntime decoder( + codec_weights, + execution_context, + num_quantizers, + engine::codecs::MossAudioTokenizerCodecRuntimeOptions{ + kWeightContextBytes, + kGraphArenaBytes, + kGraphArenaBytes, + false, + }); std::cout << "decoding " << frames << " frames...\n" << std::flush; - const auto stereo = decoder.decode(codes); + const auto decoded = decoder.decode(engine::codecs::MossAudioTokenizerCodes{frames, std::move(codes)}); + const auto & stereo = decoded.channels; std::cout << "\n=== RESULT ===\n"; std::cout << "channels=" << stereo.size() << " samples_per_channel=" << stereo.front().size() << "\n"; diff --git a/tests/moss_tts_local/codec_dequant_parity.cpp b/tests/moss_tts_local/codec_dequant_parity.cpp deleted file mode 100644 index 04925bd65..000000000 --- a/tests/moss_tts_local/codec_dequant_parity.cpp +++ /dev/null @@ -1,70 +0,0 @@ -// 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/shared/audio_tokenizer_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::MossAudioTokenizerQuantizer 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; -} diff --git a/tests/moss_tts_local/codec_encode_parity.cpp b/tests/moss_tts_local/codec_encode_parity.cpp index d021bdbf4..6ab1be093 100644 --- a/tests/moss_tts_local/codec_encode_parity.cpp +++ b/tests/moss_tts_local/codec_encode_parity.cpp @@ -8,7 +8,7 @@ #include "engine/framework/core/backend.h" #include "engine/framework/core/execution_context.h" -#include "engine/models/moss/shared/audio_tokenizer_encoder.h" +#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" #include #include @@ -16,6 +16,7 @@ #include #include #include +#include #include namespace { @@ -115,12 +116,25 @@ int main(int argc, char ** argv) { engine::core::ExecutionContext execution_context(backend_config); std::cout << "loading codec encoder weights...\n" << std::flush; - engine::models::moss::MossAudioTokenizerEncoder encoder( - codec_dir, execution_context, num_quantizers, kWeightContextBytes, kGraphArenaBytes); + auto codec_weights = engine::assets::open_tensor_source(codec_dir); + engine::codecs::MossAudioTokenizerCodecRuntime encoder( + codec_weights, + execution_context, + num_quantizers, + engine::codecs::MossAudioTokenizerCodecRuntimeOptions{ + kWeightContextBytes, + kGraphArenaBytes, + kGraphArenaBytes, + false, + }); std::cout << "encoding...\n" << std::flush; - const auto codes = encoder.encode(stereo); - const int64_t frames = codes.empty() ? 0 : static_cast(codes.front().size()); + const auto encoded = encoder.encode(engine::codecs::MossAudioTokenizerAudio{ + 48000, + std::move(stereo), + }); + const auto & codes = encoded.codebooks; + const int64_t frames = encoded.frames; std::cout << "produced codes [" << codes.size() << "," << frames << "]\n"; const auto ref = read_codes_csv(ref_codes_path); diff --git a/tests/moss_voicegen/backbone_parity.cpp b/tests/moss_voicegen/backbone_parity.cpp index 6f8bbe325..1a7838df8 100644 --- a/tests/moss_voicegen/backbone_parity.cpp +++ b/tests/moss_voicegen/backbone_parity.cpp @@ -15,6 +15,7 @@ #include "engine/framework/core/backend.h" #include "engine/framework/core/execution_context.h" #include "engine/framework/io/json.h" +#include "engine/framework/modules/multi_codebook_embedding.h" #include "engine/models/moss/shared/token_rows.h" #include @@ -138,13 +139,13 @@ int main(int argc, char ** argv) { // handed in as a per-position bias. The shared helper skips pad codes, which is // exact for this checkpoint: every emb_ext table's pad row (index audio_vocab_size) // is all zeros, so adding it the way the reference does changes nothing. - engine::models::moss::AudioCodebookSpec codebook_spec; + engine::modules::MultiCodebookEmbeddingSpec codebook_spec; codebook_spec.hidden_size = hidden_size; codebook_spec.num_codebooks = config.num_codebooks; - codebook_spec.audio_vocab_size = config.audio_vocab_size + 1; - codebook_spec.audio_pad_token_id = config.audio_pad_code; + codebook_spec.vocab_size = config.audio_vocab_size + 1; + codebook_spec.pad_token_id = config.audio_pad_code; codebook_spec.tensor_prefix = "emb_ext"; - const engine::models::moss::AudioCodebookEmbeddings codebooks(*assets->model_weights, codebook_spec); + const engine::modules::MultiCodebookEmbedding codebooks(*assets->model_weights, codebook_spec); std::vector audio_bias(static_cast(steps * hidden_size), 0.0F); for (int64_t row = 0; row < steps; ++row) { diff --git a/tests/moss_voicegen/codec_decode_parity.cpp b/tests/moss_voicegen/codec_decode_parity.cpp index 074d9296a..bd4b75640 100644 --- a/tests/moss_voicegen/codec_decode_parity.cpp +++ b/tests/moss_voicegen/codec_decode_parity.cpp @@ -14,7 +14,7 @@ #include "engine/framework/core/backend.h" #include "engine/framework/core/execution_context.h" #include "engine/framework/io/json.h" -#include "engine/models/moss/shared/audio_tokenizer_decoder.h" +#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" #include #include @@ -73,20 +73,24 @@ int main(int argc, char ** argv) { backend_config.threads = std::stoi(arg_value(argc, argv, "--threads", "8")); engine::core::ExecutionContext execution_context(backend_config); - const engine::models::moss::MossAudioTokenizerDecoder codec( - *assets->audio_tokenizer_weights, + engine::codecs::MossAudioTokenizerCodecRuntime codec( + assets->audio_tokenizer_weights, execution_context, kCodebooks, - 4096ull * 1024ull * 1024ull, - 2048ull * 1024ull * 1024ull, - engine::models::moss::moss_audio_tokenizer_v1_config()); - - const auto channels = codec.decode(code_matrix()); - if (channels.size() != 1) { - std::cerr << "FAIL: v1 is mono but the decoder returned " << channels.size() << " channels\n"; + engine::codecs::MossAudioTokenizerCodecRuntimeOptions{ + 4096ull * 1024ull * 1024ull, + 2048ull * 1024ull * 1024ull, + 2048ull * 1024ull * 1024ull, + false, + }, + engine::codecs::moss_audio_tokenizer_v1_config()); + + const auto decoded = codec.decode(engine::codecs::MossAudioTokenizerCodes{kFrames, code_matrix()}); + if (decoded.channels.size() != 1) { + std::cerr << "FAIL: v1 is mono but the decoder returned " << decoded.channels.size() << " channels\n"; return 1; } - const auto & audio = channels.front(); + const auto & audio = decoded.channels.front(); double peak = 0.0; double energy = 0.0; diff --git a/tests/moss_voicegen/generation_parity.cpp b/tests/moss_voicegen/generation_parity.cpp index 6ad67d9bc..31997b199 100644 --- a/tests/moss_voicegen/generation_parity.cpp +++ b/tests/moss_voicegen/generation_parity.cpp @@ -14,6 +14,7 @@ #include "engine/framework/core/backend.h" #include "engine/framework/core/execution_context.h" #include "engine/framework/io/json.h" +#include "engine/framework/modules/multi_codebook_embedding.h" #include "engine/models/moss/shared/token_rows.h" #include @@ -95,13 +96,13 @@ int main(int argc, char ** argv) { const int64_t prompt_rows = static_cast(prompt.text_tokens.size()); const auto steps = static_cast(expected.size()); - engine::models::moss::AudioCodebookSpec codebook_spec; + engine::modules::MultiCodebookEmbeddingSpec codebook_spec; codebook_spec.hidden_size = hidden_size; codebook_spec.num_codebooks = n_vq; - codebook_spec.audio_vocab_size = config.audio_vocab_size + 1; - codebook_spec.audio_pad_token_id = config.audio_pad_code; + codebook_spec.vocab_size = config.audio_vocab_size + 1; + codebook_spec.pad_token_id = config.audio_pad_code; codebook_spec.tensor_prefix = "emb_ext"; - const engine::models::moss::AudioCodebookEmbeddings codebooks(*assets->model_weights, codebook_spec); + const engine::modules::MultiCodebookEmbedding codebooks(*assets->model_weights, codebook_spec); std::vector prompt_bias(static_cast(prompt_rows * hidden_size), 0.0F); for (int64_t row = 0; row < prompt_rows; ++row) { diff --git a/tests/moss_voicegen/voicegen_smoke.cpp b/tests/moss_voicegen/voicegen_smoke.cpp index febd3b1f8..bb725c4f2 100644 --- a/tests/moss_voicegen/voicegen_smoke.cpp +++ b/tests/moss_voicegen/voicegen_smoke.cpp @@ -12,7 +12,8 @@ #include "engine/framework/audio/wav_writer.h" #include "engine/framework/core/backend.h" #include "engine/framework/core/execution_context.h" -#include "engine/models/moss/shared/audio_tokenizer_decoder.h" +#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" +#include "engine/framework/modules/multi_codebook_embedding.h" #include "engine/models/moss/shared/token_rows.h" #include @@ -23,6 +24,7 @@ #include #include #include +#include #include namespace { @@ -86,13 +88,13 @@ int main(int argc, char ** argv) { std::cout << "prompt rows=" << prompt_rows << " frame bounds=[" << min_frames << ", " << max_frames << "] max steps=" << max_steps << "\n"; - engine::models::moss::AudioCodebookSpec codebook_spec; + engine::modules::MultiCodebookEmbeddingSpec codebook_spec; codebook_spec.hidden_size = hidden_size; codebook_spec.num_codebooks = n_vq; - codebook_spec.audio_vocab_size = config.audio_vocab_size + 1; - codebook_spec.audio_pad_token_id = config.audio_pad_code; + codebook_spec.vocab_size = config.audio_vocab_size + 1; + codebook_spec.pad_token_id = config.audio_pad_code; codebook_spec.tensor_prefix = "emb_ext"; - const engine::models::moss::AudioCodebookEmbeddings codebooks(*assets->model_weights, codebook_spec); + const engine::modules::MultiCodebookEmbedding codebooks(*assets->model_weights, codebook_spec); std::vector prompt_bias(static_cast(prompt_rows * hidden_size), 0.0F); for (int64_t row = 0; row < prompt_rows; ++row) { @@ -191,18 +193,25 @@ int main(int argc, char ** argv) { } const auto decode_start = std::chrono::steady_clock::now(); - const engine::models::moss::MossAudioTokenizerDecoder codec( - *assets->audio_tokenizer_weights, + engine::codecs::MossAudioTokenizerCodecRuntime codec( + assets->audio_tokenizer_weights, execution_context, codebooks_out, - 4096ull * 1024ull * 1024ull, - 2048ull * 1024ull * 1024ull, - engine::models::moss::moss_audio_tokenizer_v1_config()); - const auto channels = codec.decode(codes); + engine::codecs::MossAudioTokenizerCodecRuntimeOptions{ + 4096ull * 1024ull * 1024ull, + 2048ull * 1024ull * 1024ull, + 2048ull * 1024ull * 1024ull, + false, + }, + engine::codecs::moss_audio_tokenizer_v1_config()); + const auto decoded = codec.decode(engine::codecs::MossAudioTokenizerCodes{ + codes.empty() ? 0 : static_cast(codes.front().size()), + std::move(codes), + }); const double decode_seconds = std::chrono::duration(std::chrono::steady_clock::now() - decode_start).count(); - const auto & waveform = channels.at(0); + const auto & waveform = decoded.channels.at(0); double peak = 0.0; double energy = 0.0; for (const float sample : waveform) { diff --git a/tools/audiocpp_cli/audiocpp_cli_path_cases.json b/tools/audiocpp_cli/audiocpp_cli_path_cases.json index d52463bf9..20c2b69d3 100644 --- a/tools/audiocpp_cli/audiocpp_cli_path_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_path_cases.json @@ -1914,7 +1914,7 @@ "id": "moss_tts_nano_voice_clone_long", "coverage": "MOSS-TTS-Nano voice clone, text/audio token generation, MOSS-Audio-Tokenizer-Nano decode", "family": "moss_tts_nano", - "model": "models/MOSS-TTS-Nano-100M", + "model": "models/MOSS-TTS-Nano-100M-GGUF/moss-tts-nano-100m-q8_0.gguf", "task": "tts", "mode": "offline", "outputs": [ @@ -1934,7 +1934,7 @@ "id": "moss_tts_nano_sampled_multi_request", "coverage": "MOSS-TTS-Nano seeded sampled voice clone across varied request lengths in one session", "family": "moss_tts_nano", - "model": "models/MOSS-TTS-Nano-100M", + "model": "models/MOSS-TTS-Nano-100M-GGUF/moss-tts-nano-100m-q8_0.gguf", "task": "tts", "mode": "offline", "outputs": [ @@ -1969,7 +1969,7 @@ "id": "moss_tts_nano_text_chunk_mode", "coverage": "MOSS-TTS-Nano framework text chunk mode override through request options", "family": "moss_tts_nano", - "model": "models/MOSS-TTS-Nano-100M", + "model": "models/MOSS-TTS-Nano-100M-GGUF/moss-tts-nano-100m-q8_0.gguf", "task": "tts", "mode": "offline", "outputs": [ @@ -1993,7 +1993,7 @@ "id": "moss_tts_nano_continuation_text_chunk", "coverage": "MOSS-TTS-Nano continuation mode without reference audio, including framework text chunking", "family": "moss_tts_nano", - "model": "models/MOSS-TTS-Nano-100M", + "model": "models/MOSS-TTS-Nano-100M-GGUF/moss-tts-nano-100m-q8_0.gguf", "task": "tts", "mode": "offline", "outputs": [ @@ -2016,7 +2016,7 @@ "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", + "model": "models/MOSS-TTS-Local-v1.5-GGUF/moss-tts-local-v1.5-q8_0.gguf", "task": "tts", "mode": "offline", "outputs": [ @@ -2035,7 +2035,7 @@ "id": "moss_tts_local_voice_clone_sampled", "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", + "model": "models/MOSS-TTS-Local-v1.5-GGUF/moss-tts-local-v1.5-q8_0.gguf", "task": "tts", "mode": "offline", "outputs": [ @@ -2055,7 +2055,7 @@ "id": "moss_tts_local_sampled_language", "coverage": "MOSS-TTS-Local default sampled generation path with an explicit language template slot", "family": "moss_tts_local", - "model": "models/MOSS-TTS-Local-Transformer-v1.5", + "model": "models/MOSS-TTS-Local-v1.5-GGUF/moss-tts-local-v1.5-q8_0.gguf", "task": "tts", "mode": "offline", "outputs": [ @@ -2074,7 +2074,7 @@ "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", + "model": "models/MOSS-TTS-Local-v1.5-GGUF/moss-tts-local-v1.5-q8_0.gguf", "task": "tts", "mode": "offline", "outputs": [ @@ -2125,6 +2125,28 @@ } ] }, + { + "id": "moss_voicegen_radio_voice_design", + "coverage": "MOSS-VoiceGenerator voice design path: Qwen3 delay backbone, multi-codebook embeddings, delay decoder, heads, and MOSS-Audio-Tokenizer-v1 decode", + "family": "moss_voicegen", + "model": "models/MOSS-VoiceGenerator-GGUF/moss_voicegen_bf16_codec_f16_decode.gguf", + "task": "vdes", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "radio_voice_design", + "text": "Good evening, and welcome back to the late show. Tonight we are testing the audio.cpp path that designs a voice from text, generates codec tokens, and decodes them into speech.", + "language": "English", + "instruct": "A warm male radio host in his fifties, calm, confident, and never shrill.", + "seed": 1234, + "max_tokens": 220, + "do_sample": false + } + ] + }, { "id": "qwen3_asr_offline", "coverage": "Qwen3 ASR offline audio encoder and thinker decode", From b0c05e4171a1be8dfb224b1d6b8b6eaedeab0e9e Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Sat, 29 Aug 2026 03:54:45 -0400 Subject: [PATCH 2/3] Use shared HF sampler for MOSS models --- CMakeLists.txt | 1 - .../moss_voicegen/delay_decoder.h | 3 +- .../moss/moss_tts_nano/local_frame_decoder.h | 2 + include/engine/models/moss/shared/sampling.h | 31 ---- .../moss_voicegen/delay_decoder.cpp | 70 ++++----- src/models/moss/moss_tts_local/generator.cpp | 86 ++++++----- .../moss_tts_nano/local_frame_decoder.cpp | 68 ++++---- src/models/moss/shared/sampling.cpp | 146 ------------------ 8 files changed, 123 insertions(+), 284 deletions(-) delete mode 100644 include/engine/models/moss/shared/sampling.h delete mode 100644 src/models/moss/shared/sampling.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 91060e17a..3b416870f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -801,7 +801,6 @@ audiocpp_add_model(moss src/models/moss/moss_tts_nano/prompt_builder.cpp src/models/moss/moss_tts_nano/session.cpp src/models/moss/moss_tts_nano/tokenizer_text.cpp - src/models/moss/shared/sampling.cpp src/models/moss/shared/token_rows.cpp src/models/moss/moss_tts_local/depth_transformer.cpp src/models/moss/moss_tts_local/generator.cpp diff --git a/include/engine/community_models/moss_voicegen/delay_decoder.h b/include/engine/community_models/moss_voicegen/delay_decoder.h index 58dbe4ba3..4bd6346a0 100644 --- a/include/engine/community_models/moss_voicegen/delay_decoder.h +++ b/include/engine/community_models/moss_voicegen/delay_decoder.h @@ -2,6 +2,7 @@ #include "engine/community_models/moss_voicegen/assets.h" #include "engine/community_models/moss_voicegen/heads.h" +#include "engine/framework/sampling/hf_sampler.h" #include #include @@ -68,8 +69,8 @@ class MossVoiceGenDelayDecoder { MossVoiceGenConfig config_; MossVoiceGenSamplingOptions sampling_; MossVoiceGenLengthBounds bounds_; - uint32_t seed_ = 0; std::mt19937 rng_; + engine::sampling::HfSamplerScratch sampler_scratch_; uint64_t sample_call_index_ = 0; int64_t step_index_ = 0; diff --git a/include/engine/models/moss/moss_tts_nano/local_frame_decoder.h b/include/engine/models/moss/moss_tts_nano/local_frame_decoder.h index 53881058c..c77fa814d 100644 --- a/include/engine/models/moss/moss_tts_nano/local_frame_decoder.h +++ b/include/engine/models/moss/moss_tts_nano/local_frame_decoder.h @@ -1,6 +1,7 @@ #pragma once #include "engine/framework/core/execution_context.h" +#include "engine/framework/sampling/hf_sampler.h" #include "engine/framework/sampling/torch_random.h" #include "engine/models/moss/moss_tts_nano/assets.h" #include "engine/models/moss/moss_tts_nano/types.h" @@ -42,6 +43,7 @@ class MossTTSNanoLocalFrameDecoderRuntime { std::shared_ptr assets_; core::ExecutionContext & execution_context_; engine::sampling::TorchCudaSamplingPolicy sampling_policy_; + engine::sampling::HfSamplerScratch sampler_scratch_; std::shared_ptr weights_; size_t graph_arena_bytes_ = 0; std::unique_ptr text_graph_; diff --git a/include/engine/models/moss/shared/sampling.h b/include/engine/models/moss/shared/sampling.h deleted file mode 100644 index 350251d59..000000000 --- a/include/engine/models/moss/shared/sampling.h +++ /dev/null @@ -1,31 +0,0 @@ -#pragma once - -#include "engine/framework/sampling/torch_random.h" - -#include -#include -#include -#include - -namespace engine::models::moss { - -int32_t argmax_index(const std::vector & logits, std::string_view context); - -void apply_repetition_penalty( - std::vector & logits, - const std::vector & previous_token_ids, - float penalty, - std::string_view context); - -int32_t sample_index( - const std::vector & logits, - int top_k, - float top_p, - float temperature, - std::mt19937 & rng, - std::string_view context, - const engine::sampling::TorchCudaSamplingPolicy * sampling_policy = nullptr, - uint64_t seed = 0, - uint64_t call_index = 0); - -} // namespace engine::models::moss diff --git a/src/community_models/moss_voicegen/delay_decoder.cpp b/src/community_models/moss_voicegen/delay_decoder.cpp index 0e2a07a1f..a32977bdd 100644 --- a/src/community_models/moss_voicegen/delay_decoder.cpp +++ b/src/community_models/moss_voicegen/delay_decoder.cpp @@ -1,7 +1,5 @@ #include "engine/community_models/moss_voicegen/delay_decoder.h" -#include "engine/models/moss/shared/sampling.h" - #include #include #include @@ -17,15 +15,6 @@ void forbid(std::vector & logits, int64_t token_id) { } } -void apply_temperature(std::vector & logits, float temperature) { - if (temperature == 1.0F || temperature <= 0.0F) { - return; - } - for (float & value : logits) { - value /= temperature; - } -} - } // namespace MossVoiceGenDelayDecoder::MossVoiceGenDelayDecoder( @@ -36,7 +25,6 @@ MossVoiceGenDelayDecoder::MossVoiceGenDelayDecoder( : config_(std::move(config)), sampling_(sampling), bounds_(bounds), - seed_(seed), rng_(seed) { if (config_.num_codebooks <= 0) { throw std::runtime_error("MOSS-VoiceGenerator delay decoder requires a positive codebook count"); @@ -44,48 +32,58 @@ MossVoiceGenDelayDecoder::MossVoiceGenDelayDecoder( } int32_t MossVoiceGenDelayDecoder::sample_text(std::vector & logits) { - if (!sampling_.do_sample) { - return engine::models::moss::argmax_index(logits, "moss_voicegen.text"); - } - return engine::models::moss::sample_index( + engine::sampling::HfSampler sampler; + engine::sampling::HfSamplingOptions options; + options.do_sample = sampling_.do_sample; + options.temperature = sampling_.text_temperature; + options.top_k = sampling_.text_top_k; + options.top_p = sampling_.text_top_p; + const int32_t token = sampler.sample( logits, - sampling_.text_top_k, - sampling_.text_top_p, - 1.0F, // the temperature is already folded into the logits + {}, + options, + sampler_scratch_, rng_, - "moss_voicegen.text", nullptr, - seed_, - sample_call_index_++); + "moss_voicegen.text"); + if (sampling_.do_sample) { + ++sample_call_index_; + } + return token; } int32_t MossVoiceGenDelayDecoder::sample_code(std::vector & logits, int64_t codebook) { + std::vector previous; if (sampling_.audio_repetition_penalty != 1.0F) { // The reference penalises against every earlier row of this codebook, prompt rows // included. Those are all pad, and pad is masked to -inf just below, so restricting // this to the generated history gives the same result. - std::vector previous; previous.reserve(history_.size()); for (const auto & row : history_) { previous.push_back(row.codes[static_cast(codebook)]); } - engine::models::moss::apply_repetition_penalty( - logits, previous, sampling_.audio_repetition_penalty, "moss_voicegen.audio"); } forbid(logits, config_.audio_pad_code); - if (!sampling_.do_sample) { - return engine::models::moss::argmax_index(logits, "moss_voicegen.audio"); - } - return engine::models::moss::sample_index( + + engine::sampling::HfSampler sampler; + engine::sampling::HfSamplingOptions options; + options.do_sample = sampling_.do_sample; + options.temperature = sampling_.audio_temperature; + options.top_k = sampling_.audio_top_k; + options.top_p = sampling_.audio_top_p; + options.repetition_penalty = sampling_.audio_repetition_penalty; + const int32_t token = sampler.sample( logits, - sampling_.audio_top_k, - sampling_.audio_top_p, - 1.0F, + previous, + options, + sampler_scratch_, rng_, - "moss_voicegen.audio", nullptr, - seed_, - sample_call_index_++); + "moss_voicegen.audio"); + if (sampling_.do_sample) { + ++sample_call_index_; + } + return token; } MossVoiceGenDelayRow MossVoiceGenDelayDecoder::step(MossVoiceGenStepLogits & logits) { @@ -123,7 +121,6 @@ MossVoiceGenDelayRow MossVoiceGenDelayDecoder::step(MossVoiceGenStepLogits & log } if (sample_text_token) { - apply_temperature(logits.text, sampling_.text_temperature); if (in_audio_) { // Mid-utterance the only legal continuations are "another audio frame" and // "start the flush", so everything else is masked out. @@ -179,7 +176,6 @@ MossVoiceGenDelayRow MossVoiceGenDelayDecoder::step(MossVoiceGenStepLogits & log continue; } auto & codebook_logits = logits.audio[static_cast(codebook)]; - apply_temperature(codebook_logits, sampling_.audio_temperature); row.codes[static_cast(codebook)] = sample_code(codebook_logits, codebook); } diff --git a/src/models/moss/moss_tts_local/generator.cpp b/src/models/moss/moss_tts_local/generator.cpp index 8af860317..a5c14a3e3 100644 --- a/src/models/moss/moss_tts_local/generator.cpp +++ b/src/models/moss/moss_tts_local/generator.cpp @@ -7,8 +7,8 @@ #include "engine/framework/core/module.h" #include "engine/framework/modules/linear_module.h" #include "engine/framework/modules/weight_binding.h" +#include "engine/framework/sampling/hf_sampler.h" #include "engine/framework/sampling/torch_random.h" -#include "engine/models/moss/shared/sampling.h" #include #include @@ -248,14 +248,8 @@ std::vector> MossGenerator::generate( std::vector> generated_frames; generated_frames.reserve(static_cast(options.max_new_frames)); std::vector> code_history(static_cast(n_vq)); - const bool use_repetition_penalty = options.audio_repetition_penalty != 1.0F; - std::vector> code_seen(static_cast(n_vq)); - if (use_repetition_penalty) { - for (int64_t codebook = 0; codebook < n_vq; ++codebook) { - const int64_t codebook_size = audio_codebooks_->codebook_size(codebook); - code_seen[static_cast(codebook)].assign(static_cast(codebook_size), 0); - } - } + engine::sampling::HfSampler sampler; + engine::sampling::HfSamplerScratch sampler_scratch; std::mt19937 rng(options.seed); uint64_t sample_call_index = 0; double bias_ms = 0.0; @@ -304,18 +298,27 @@ std::vector> MossGenerator::generate( int32_t gate_index = 0; add_timing(gate_ms, [&]() { gate_logits = project(local_text_head_, local_hidden, 2, hidden); - gate_index = options.do_sample - ? moss::sample_index( - gate_logits, - options.text_top_k, - options.text_top_p, - options.text_temperature, - rng, - "MOSS-TTS-Local sampler", - &sampling_policy_, - options.seed, - sample_call_index++) - : moss::argmax_index(gate_logits, "MOSS-TTS-Local sampler"); + engine::sampling::HfSamplingOptions hf_options; + hf_options.do_sample = options.do_sample; + hf_options.temperature = options.text_temperature; + hf_options.top_k = options.text_top_k; + hf_options.top_p = options.text_top_p; + const engine::sampling::HfTorchSamplingState torch_state{ + &sampling_policy_, + options.seed, + sample_call_index, + }; + gate_index = sampler.sample( + gate_logits, + {}, + hf_options, + sampler_scratch, + rng, + options.do_sample && sampling_policy_.cuda_fast_path ? &torch_state : nullptr, + "MOSS-TTS-Local sampler"); + if (options.do_sample) { + ++sample_call_index; + } }); if (gate_index != 0) { break; @@ -334,31 +337,32 @@ std::vector> MossGenerator::generate( } int32_t code = 0; add_timing(sampling_ms, [&]() { - moss::apply_repetition_penalty( + engine::sampling::HfSamplingOptions hf_options; + hf_options.do_sample = options.do_sample; + hf_options.temperature = options.audio_temperature; + hf_options.top_k = options.audio_top_k; + hf_options.top_p = options.audio_top_p; + hf_options.repetition_penalty = options.audio_repetition_penalty; + const engine::sampling::HfTorchSamplingState torch_state{ + &sampling_policy_, + options.seed, + sample_call_index, + }; + code = sampler.sample( logits, code_history[static_cast(codebook)], - options.audio_repetition_penalty, + hf_options, + sampler_scratch, + rng, + options.do_sample && sampling_policy_.cuda_fast_path ? &torch_state : nullptr, "MOSS-TTS-Local sampler"); - code = options.do_sample - ? moss::sample_index( - logits, - options.audio_top_k, - options.audio_top_p, - options.audio_temperature, - rng, - "MOSS-TTS-Local sampler", - &sampling_policy_, - options.seed, - sample_call_index++) - : moss::argmax_index(logits, "MOSS-TTS-Local sampler"); + if (options.do_sample) { + ++sample_call_index; + } }); frame_codes[static_cast(codebook)] = code; - if (use_repetition_penalty) { - auto & seen = code_seen[static_cast(codebook)]; - if (code >= 0 && static_cast(code) < seen.size() && seen[static_cast(code)] == 0) { - seen[static_cast(code)] = 1; - code_history[static_cast(codebook)].push_back(code); - } + if (code >= 0) { + code_history[static_cast(codebook)].push_back(code); } if (codebook + 1 < n_vq) { diff --git a/src/models/moss/moss_tts_nano/local_frame_decoder.cpp b/src/models/moss/moss_tts_nano/local_frame_decoder.cpp index 41d8f1a2c..d2ead5cf3 100644 --- a/src/models/moss/moss_tts_nano/local_frame_decoder.cpp +++ b/src/models/moss/moss_tts_nano/local_frame_decoder.cpp @@ -10,7 +10,6 @@ #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/shared/sampling.h" #include #include @@ -439,19 +438,29 @@ std::vector MossTTSNanoLocalFrameDecoderRuntime::generate_frame( text_logits.at(static_cast(text_candidates[0])), text_logits.at(static_cast(text_candidates[1])), }; + engine::sampling::HfSampler sampler; + engine::sampling::HfSamplingOptions text_sampling; + text_sampling.do_sample = sampling.do_sample; + text_sampling.temperature = sampling.text_temperature; + text_sampling.top_k = sampling.text_top_k; + text_sampling.top_p = sampling.text_top_p; + const engine::sampling::HfTorchSamplingState text_torch_state{ + &sampling_policy_, + seed, + sample_call_index, + }; const int32_t best_text = text_candidates[static_cast( - sampling.do_sample - ? moss::sample_index( - text_candidate_logits, - sampling.text_top_k, - sampling.text_top_p, - sampling.text_temperature, - rng, - "MOSS-TTS-Nano local decoder", - &sampling_policy_, - seed, - sample_call_index++) - : moss::argmax_index(text_candidate_logits, "MOSS-TTS-Nano local decoder"))]; + sampler.sample( + text_candidate_logits, + {}, + text_sampling, + sampler_scratch_, + rng, + sampling.do_sample && sampling_policy_.cuda_fast_path ? &text_torch_state : nullptr, + "MOSS-TTS-Nano local decoder"))]; + if (sampling.do_sample) { + ++sample_call_index; + } if (best_text == assets_->config.audio_end_token_id) { return {}; } @@ -470,23 +479,28 @@ std::vector MossTTSNanoLocalFrameDecoderRuntime::generate_frame( for (size_t frame = 0; frame * static_cast(assets_->config.n_vq) + static_cast(q) < history.size(); ++frame) { codebook_history.push_back(history[frame * static_cast(assets_->config.n_vq) + static_cast(q)]); } - moss::apply_repetition_penalty( + engine::sampling::HfSamplingOptions audio_sampling; + audio_sampling.do_sample = sampling.do_sample; + audio_sampling.temperature = sampling.audio_temperature; + audio_sampling.top_k = sampling.audio_top_k; + audio_sampling.top_p = sampling.audio_top_p; + audio_sampling.repetition_penalty = sampling.audio_repetition_penalty; + const engine::sampling::HfTorchSamplingState audio_torch_state{ + &sampling_policy_, + seed, + sample_call_index, + }; + const int32_t token = sampler.sample( logits, codebook_history, - sampling.audio_repetition_penalty, + audio_sampling, + sampler_scratch_, + rng, + sampling.do_sample && sampling_policy_.cuda_fast_path ? &audio_torch_state : nullptr, "MOSS-TTS-Nano local decoder"); - const int32_t token = sampling.do_sample - ? moss::sample_index( - logits, - sampling.audio_top_k, - sampling.audio_top_p, - sampling.audio_temperature, - rng, - "MOSS-TTS-Nano local decoder", - &sampling_policy_, - seed, - sample_call_index++) - : moss::argmax_index(logits, "MOSS-TTS-Nano local decoder"); + if (sampling.do_sample) { + ++sample_call_index; + } frame[static_cast(q)] = token; previous.push_back(token); } diff --git a/src/models/moss/shared/sampling.cpp b/src/models/moss/shared/sampling.cpp deleted file mode 100644 index f37a39803..000000000 --- a/src/models/moss/shared/sampling.cpp +++ /dev/null @@ -1,146 +0,0 @@ -#include "engine/models/moss/shared/sampling.h" - -#include -#include -#include -#include -#include -#include - -namespace engine::models::moss { - -int32_t argmax_index(const std::vector & logits, std::string_view context) { - if (logits.empty()) { - throw std::runtime_error(std::string(context) + " 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_token_ids, - float penalty, - std::string_view context) { - if (penalty == 1.0F || previous_token_ids.empty()) { - return; - } - if (penalty <= 0.0F) { - throw std::runtime_error(std::string(context) + " repetition penalty must be positive"); - } - std::unordered_set seen; - seen.reserve(previous_token_ids.size()); - for (const int32_t token : previous_token_ids) { - if (token < 0 || token >= static_cast(logits.size())) { - continue; - } - if (!seen.insert(token).second) { - continue; - } - float & value = logits[static_cast(token)]; - 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, - std::string_view context, - const engine::sampling::TorchCudaSamplingPolicy * sampling_policy, - uint64_t seed, - uint64_t call_index) { - if (temperature <= 0.0F) { - throw std::runtime_error(std::string(context) + " 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(std::string(context) + " sampler has no finite logits"); - } - if (top_k > 0 && static_cast(indices.size()) > top_k) { - std::vector ranked = indices; - const auto keep = ranked.begin() + top_k - 1; - std::nth_element(ranked.begin(), keep, ranked.end(), [&](int32_t lhs, int32_t rhs) { - return logits[static_cast(lhs)] > logits[static_cast(rhs)]; - }); - const float kth_logit = logits[static_cast(*keep)]; - indices.erase( - std::remove_if( - indices.begin(), - indices.end(), - [&](int32_t index) { - return logits[static_cast(index)] < kth_logit; - }), - indices.end()); - } - std::sort(indices.begin(), indices.end(), [&](int32_t lhs, int32_t rhs) { - const float lhs_logit = logits[static_cast(lhs)]; - const float rhs_logit = logits[static_cast(rhs)]; - if (lhs_logit == rhs_logit) { - return lhs < rhs; - } - return lhs_logit > rhs_logit; - }); - - 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); - } - if (sampling_policy != nullptr && sampling_policy->cuda_fast_path) { - double best_rank = -std::numeric_limits::infinity(); - int32_t best_token = -1; - for (size_t index = 0; index < indices.size(); ++index) { - const float exponential = engine::sampling::torch_cuda_tensor_iterator_exponential_element( - seed, - static_cast(logits.size()), - static_cast(indices[index]), - call_index, - sampling_policy->multiprocessor_count, - sampling_policy->max_threads_per_multiprocessor); - const double rank = weights[index] / static_cast(exponential); - if (rank > best_rank) { - best_rank = rank; - best_token = indices[index]; - } - } - if (best_token < 0) { - throw std::runtime_error(std::string(context) + " CUDA sampler failed to select a token"); - } - return best_token; - } - std::discrete_distribution distribution(weights.begin(), weights.end()); - return indices[distribution(rng)]; -} - -} // namespace engine::models::moss From 628bc0ae35d85e70aaf5ae169c27596953c8965a Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Sat, 29 Aug 2026 04:19:48 -0400 Subject: [PATCH 3/3] Move MOSS token row builder into codec runtime --- CMakeLists.txt | 1 - .../community_models/moss_voicegen/session.h | 1 - .../moss_voicegen/tokenizer_text.h | 4 +- .../moss_audio_tokenizer_codec_runtime.h | 21 ++++ .../models/moss/moss_tts_local/generator.h | 1 - .../moss/moss_tts_local/tokenizer_text.h | 4 +- .../engine/models/moss/shared/token_rows.h | 28 ----- .../moss_voicegen/tokenizer_text.cpp | 4 +- .../moss_audio_tokenizer_codec_runtime.cpp | 58 ++++++++++ src/models/moss/moss_tts_local/session.cpp | 4 +- .../moss/moss_tts_local/tokenizer_text.cpp | 14 +-- .../moss/moss_tts_nano/prompt_builder.cpp | 101 +++++++++--------- src/models/moss/shared/token_rows.cpp | 54 ---------- tests/moss_voicegen/backbone_parity.cpp | 23 ++-- tests/moss_voicegen/codec_decode_parity.cpp | 25 +++-- tests/moss_voicegen/generation_parity.cpp | 17 +-- tests/moss_voicegen/voicegen_smoke.cpp | 13 +-- 17 files changed, 186 insertions(+), 187 deletions(-) delete mode 100644 include/engine/models/moss/shared/token_rows.h delete mode 100644 src/models/moss/shared/token_rows.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 3b416870f..a5d2dd1cc 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -801,7 +801,6 @@ audiocpp_add_model(moss src/models/moss/moss_tts_nano/prompt_builder.cpp src/models/moss/moss_tts_nano/session.cpp src/models/moss/moss_tts_nano/tokenizer_text.cpp - src/models/moss/shared/token_rows.cpp src/models/moss/moss_tts_local/depth_transformer.cpp src/models/moss/moss_tts_local/generator.cpp src/models/moss/moss_tts_local/loader.cpp diff --git a/include/engine/community_models/moss_voicegen/session.h b/include/engine/community_models/moss_voicegen/session.h index 8c114a2e9..2d5fb49e1 100644 --- a/include/engine/community_models/moss_voicegen/session.h +++ b/include/engine/community_models/moss_voicegen/session.h @@ -9,7 +9,6 @@ #include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" #include "engine/framework/modules/multi_codebook_embedding.h" #include "engine/framework/runtime/session_base.h" -#include "engine/models/moss/shared/token_rows.h" #include #include diff --git a/include/engine/community_models/moss_voicegen/tokenizer_text.h b/include/engine/community_models/moss_voicegen/tokenizer_text.h index f9a0f7683..90f8a5519 100644 --- a/include/engine/community_models/moss_voicegen/tokenizer_text.h +++ b/include/engine/community_models/moss_voicegen/tokenizer_text.h @@ -1,7 +1,7 @@ #pragma once #include "engine/community_models/moss_voicegen/assets.h" -#include "engine/models/moss/shared/token_rows.h" +#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" #include #include @@ -23,7 +23,7 @@ class MossVoiceGenTextProcessor { // `instruction` describes the speaker to design. `language` must be the full language // name the model was trained on ("English", not "en"); an empty value renders "None". - moss::TokenRows build_generation_prefix( + engine::codecs::MossTokenRows build_generation_prefix( const std::string & text, const std::optional & instruction, const std::optional & language) const; diff --git a/include/engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h b/include/engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h index 1fa685ea0..403b497b1 100644 --- a/include/engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h +++ b/include/engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h @@ -58,6 +58,27 @@ struct MossAudioTokenizerCodes { std::vector> codebooks; }; +struct MossTokenRows { + std::vector text_tokens; + std::vector audio_codes; +}; + +class MossTokenRowBuilder { +public: + MossTokenRowBuilder(int64_t num_codebooks, int32_t audio_pad_token_id); + + void push_text_token(int32_t token_id); + void push_text_tokens(const std::vector & token_ids); + void push_audio_row(int32_t text_slot_token_id, const int32_t * codes, int64_t num_codebooks); + void push_audio_row(int32_t text_slot_token_id, const std::vector> & codes, int64_t frame); + MossTokenRows finish(); + +private: + int64_t num_codebooks_ = 0; + int32_t audio_pad_token_id_ = 0; + MossTokenRows rows_; +}; + struct MossAudioTokenizerCodecRuntimeOptions { size_t weight_context_bytes = 256ull * 1024ull * 1024ull; size_t encoder_graph_arena_bytes = 2048ull * 1024ull * 1024ull; diff --git a/include/engine/models/moss/moss_tts_local/generator.h b/include/engine/models/moss/moss_tts_local/generator.h index 03b120d23..a67d0e759 100644 --- a/include/engine/models/moss/moss_tts_local/generator.h +++ b/include/engine/models/moss/moss_tts_local/generator.h @@ -4,7 +4,6 @@ #include "engine/models/moss/moss_tts_local/backbone.h" #include "engine/models/moss/moss_tts_local/depth_transformer.h" #include "engine/framework/modules/multi_codebook_embedding.h" -#include "engine/models/moss/shared/token_rows.h" #include "engine/framework/sampling/torch_random.h" #include diff --git a/include/engine/models/moss/moss_tts_local/tokenizer_text.h b/include/engine/models/moss/moss_tts_local/tokenizer_text.h index da6d5be0a..28aa1e494 100644 --- a/include/engine/models/moss/moss_tts_local/tokenizer_text.h +++ b/include/engine/models/moss/moss_tts_local/tokenizer_text.h @@ -1,6 +1,6 @@ #pragma once -#include "engine/models/moss/shared/token_rows.h" +#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" #include "engine/models/moss/moss_tts_local/assets.h" #include @@ -14,7 +14,7 @@ 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. -using MossGenerationPrefix = moss::TokenRows; +using MossGenerationPrefix = engine::codecs::MossTokenRows; // Reproduces the direct-generation branch of MossTTSLocalProcessor: it renders the // template, byte-level BPE encodes each piece with the Qwen tokenizer, and diff --git a/include/engine/models/moss/shared/token_rows.h b/include/engine/models/moss/shared/token_rows.h deleted file mode 100644 index 76af2fdb5..000000000 --- a/include/engine/models/moss/shared/token_rows.h +++ /dev/null @@ -1,28 +0,0 @@ -#pragma once - -#include -#include - -namespace engine::models::moss { - -struct TokenRows { - std::vector text_tokens; - std::vector audio_codes; -}; - -class TokenRowBuilder { -public: - TokenRowBuilder(int64_t num_codebooks, int32_t audio_pad_token_id); - - void push_text_token(int32_t token_id); - void push_text_tokens(const std::vector & token_ids); - void push_audio_row(int32_t text_slot_token_id, const std::vector> & codes, int64_t frame); - TokenRows finish(); - -private: - int64_t num_codebooks_ = 0; - int32_t audio_pad_token_id_ = 0; - TokenRows rows_; -}; - -} // namespace engine::models::moss diff --git a/src/community_models/moss_voicegen/tokenizer_text.cpp b/src/community_models/moss_voicegen/tokenizer_text.cpp index 6a4e28051..9295df47e 100644 --- a/src/community_models/moss_voicegen/tokenizer_text.cpp +++ b/src/community_models/moss_voicegen/tokenizer_text.cpp @@ -85,12 +85,12 @@ MossVoiceGenTextProcessor::MossVoiceGenTextProcessor(std::shared_ptr & instruction, const std::optional & language) const { const auto & config = impl_->assets->config; - moss::TokenRowBuilder builder(config.num_codebooks, static_cast(config.audio_pad_code)); + engine::codecs::MossTokenRowBuilder builder(config.num_codebooks, static_cast(config.audio_pad_code)); // Render the whole turn as one string and encode it in a single pass, the way the // reference processor does. Encoding the fragments separately would split merges diff --git a/src/framework/codecs/moss_audio_tokenizer_codec_runtime.cpp b/src/framework/codecs/moss_audio_tokenizer_codec_runtime.cpp index 1feffd1cf..a4efee5eb 100644 --- a/src/framework/codecs/moss_audio_tokenizer_codec_runtime.cpp +++ b/src/framework/codecs/moss_audio_tokenizer_codec_runtime.cpp @@ -353,6 +353,64 @@ inline std::vector causal_context_mask_window( namespace engine::codecs { +MossTokenRowBuilder::MossTokenRowBuilder(int64_t num_codebooks, int32_t audio_pad_token_id) + : num_codebooks_(num_codebooks), + audio_pad_token_id_(audio_pad_token_id) { + if (num_codebooks_ <= 0) { + throw std::runtime_error("MOSS token row builder requires a positive codebook count"); + } +} + +void MossTokenRowBuilder::push_text_token(int32_t token_id) { + rows_.text_tokens.push_back(token_id); + rows_.audio_codes.insert(rows_.audio_codes.end(), static_cast(num_codebooks_), audio_pad_token_id_); +} + +void MossTokenRowBuilder::push_text_tokens(const std::vector & token_ids) { + for (const int32_t token_id : token_ids) { + push_text_token(token_id); + } +} + +void MossTokenRowBuilder::push_audio_row(int32_t text_slot_token_id, const int32_t * codes, int64_t num_codebooks) { + if (num_codebooks != num_codebooks_) { + throw std::runtime_error("MOSS audio row codebook count mismatch"); + } + if (codes == nullptr) { + throw std::runtime_error("MOSS audio row codes are missing"); + } + rows_.text_tokens.push_back(text_slot_token_id); + rows_.audio_codes.insert(rows_.audio_codes.end(), codes, codes + num_codebooks); +} + +void MossTokenRowBuilder::push_audio_row( + int32_t text_slot_token_id, + const std::vector> & codes, + int64_t frame) { + if (static_cast(codes.size()) != num_codebooks_) { + throw std::runtime_error("MOSS audio row codebook count mismatch"); + } + rows_.text_tokens.push_back(text_slot_token_id); + for (int64_t codebook = 0; codebook < num_codebooks_; ++codebook) { + const auto & channel = codes[static_cast(codebook)]; + if (frame < 0 || static_cast(frame) >= channel.size()) { + throw std::runtime_error("MOSS audio row frame index is out of range"); + } + rows_.audio_codes.push_back(channel[static_cast(frame)]); + } +} + +MossTokenRows MossTokenRowBuilder::finish() { + if (rows_.text_tokens.empty()) { + throw std::runtime_error("MOSS token rows must not be empty"); + } + if (static_cast(rows_.audio_codes.size()) != + static_cast(rows_.text_tokens.size()) * num_codebooks_) { + throw std::runtime_error("MOSS token rows audio code shape mismatch"); + } + return std::move(rows_); +} + // 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 diff --git a/src/models/moss/moss_tts_local/session.cpp b/src/models/moss/moss_tts_local/session.cpp index 4bf468479..5e0c1bfaa 100644 --- a/src/models/moss/moss_tts_local/session.cpp +++ b/src/models/moss/moss_tts_local/session.cpp @@ -39,7 +39,7 @@ uint64_t mix_reference_audio_key(uint64_t key, uint64_t value) { return key; } -uint64_t prefix_hash(const moss::TokenRows & prefix, int64_t num_codebooks) { +uint64_t prefix_hash(const engine::codecs::MossTokenRows & prefix, int64_t num_codebooks) { uint64_t key = 1469598103934665603ull; for (size_t row = 0; row < prefix.text_tokens.size(); ++row) { key = mix_reference_audio_key(key, static_cast(prefix.text_tokens[row])); @@ -53,7 +53,7 @@ uint64_t prefix_hash(const moss::TokenRows & prefix, int64_t num_codebooks) { return key; } -int64_t prefix_audio_nonpad_count(const moss::TokenRows & prefix, int32_t audio_pad_token_id) { +int64_t prefix_audio_nonpad_count(const engine::codecs::MossTokenRows & prefix, int32_t audio_pad_token_id) { int64_t count = 0; for (const int32_t code : prefix.audio_codes) { if (code != audio_pad_token_id) { diff --git a/src/models/moss/moss_tts_local/tokenizer_text.cpp b/src/models/moss/moss_tts_local/tokenizer_text.cpp index a7184cb2b..7b2450fac 100644 --- a/src/models/moss/moss_tts_local/tokenizer_text.cpp +++ b/src/models/moss/moss_tts_local/tokenizer_text.cpp @@ -51,13 +51,13 @@ struct MossTextProcessor::Impl { std::shared_ptr assets; std::shared_ptr tokenizer; - moss::TokenRowBuilder make_row_builder() const { - return moss::TokenRowBuilder( + engine::codecs::MossTokenRowBuilder make_row_builder() const { + return engine::codecs::MossTokenRowBuilder( assets->config.num_codebooks, static_cast(assets->config.audio_pad_token_id)); } - void push_text(moss::TokenRowBuilder & builder, const std::string & text) const { + void push_text(engine::codecs::MossTokenRowBuilder & builder, const std::string & text) const { builder.push_text_tokens(tokenizer->encode(text)); } @@ -66,9 +66,9 @@ struct MossTextProcessor::Impl { MossGenerationPrefix build_prefix( const std::string & text, const std::optional & language, - const std::function & emit_reference) const { + const std::function & emit_reference) const { const auto & config = assets->config; - moss::TokenRowBuilder builder = make_row_builder(); + engine::codecs::MossTokenRowBuilder builder = make_row_builder(); builder.push_text_token(static_cast(config.im_start_token_id)); push_text(builder, kUserRolePrefix); push_text(builder, kUserReferencePrefix); @@ -108,7 +108,7 @@ MossTextProcessor::~MossTextProcessor() = default; MossGenerationPrefix MossTextProcessor::build_generation_prefix( const std::string & text, const std::optional & language) const { - return impl_->build_prefix(text, language, [this](moss::TokenRowBuilder & builder) { + return impl_->build_prefix(text, language, [this](engine::codecs::MossTokenRowBuilder & builder) { impl_->push_text(builder, kNoneValue); }); } @@ -132,7 +132,7 @@ MossGenerationPrefix MossTextProcessor::build_clone_prefix( } } - return impl_->build_prefix(text, language, [&](moss::TokenRowBuilder & builder) { + return impl_->build_prefix(text, language, [&](engine::codecs::MossTokenRowBuilder & 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)); diff --git a/src/models/moss/moss_tts_nano/prompt_builder.cpp b/src/models/moss/moss_tts_nano/prompt_builder.cpp index d5219dd20..1fec073ad 100644 --- a/src/models/moss/moss_tts_nano/prompt_builder.cpp +++ b/src/models/moss/moss_tts_nano/prompt_builder.cpp @@ -1,6 +1,9 @@ #include "engine/models/moss/moss_tts_nano/prompt_builder.h" +#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" + #include +#include #include #include @@ -26,41 +29,6 @@ void append_text_tokens(std::vector & out, const MossTTSNanoTextTokeniz out.insert(out.end(), tokens.begin(), tokens.end()); } -void append_text_row(MossTTSNanoPrompt & prompt, int32_t text_token, int64_t audio_pad) { - prompt.input_ids.push_back(text_token); - for (int64_t i = 0; i < prompt.row_width - 1; ++i) { - prompt.input_ids.push_back(static_cast(audio_pad)); - } - prompt.attention_mask.push_back(1); - ++prompt.rows; -} - -void append_text_rows(MossTTSNanoPrompt & prompt, const std::vector & text_tokens, int64_t audio_pad) { - for (const int32_t token : text_tokens) { - append_text_row(prompt, token, audio_pad); - } -} - -void append_audio_rows( - MossTTSNanoPrompt & prompt, - const MossTTSNanoAudioCodes & codes, - int64_t text_slot_token, - int64_t audio_pad) { - if (codes.codebooks != prompt.row_width - 1) { - throw std::runtime_error("MOSS-TTS-Nano prompt audio codebook count does not match prompt row width"); - } - for (int64_t frame = 0; frame < codes.frames; ++frame) { - prompt.input_ids.push_back(static_cast(text_slot_token)); - for (int64_t codebook = 0; codebook < codes.codebooks; ++codebook) { - const auto index = static_cast(frame * codes.codebooks + codebook); - const int32_t code = index < codes.token_ids.size() ? codes.token_ids[index] : static_cast(audio_pad); - prompt.input_ids.push_back(code); - } - prompt.attention_mask.push_back(1); - ++prompt.rows; - } -} - std::vector build_user_prompt_prefix( const MossTTSNanoConfig & config, const MossTTSNanoTextTokenizer & tokenizer) { @@ -82,6 +50,26 @@ std::vector build_assistant_prompt_prefix( return tokens; } +void append_token_rows(MossTTSNanoPrompt & prompt, const engine::codecs::MossTokenRows & rows) { + if (prompt.row_width <= 1) { + throw std::runtime_error("MOSS-TTS-Nano prompt row width is invalid"); + } + const auto codebooks = static_cast(prompt.row_width - 1); + if (rows.audio_codes.size() != rows.text_tokens.size() * codebooks) { + throw std::runtime_error("MOSS-TTS-Nano token rows audio code shape mismatch"); + } + for (size_t row = 0; row < rows.text_tokens.size(); ++row) { + prompt.input_ids.push_back(rows.text_tokens[row]); + const size_t audio_offset = row * codebooks; + prompt.input_ids.insert( + prompt.input_ids.end(), + rows.audio_codes.begin() + static_cast(audio_offset), + rows.audio_codes.begin() + static_cast(audio_offset + codebooks)); + prompt.attention_mask.push_back(1); + ++prompt.rows; + } +} + } // namespace MossTTSNanoPromptBuilder::MossTTSNanoPromptBuilder( @@ -103,6 +91,7 @@ MossTTSNanoPrompt MossTTSNanoPromptBuilder::build( if (request.text.empty()) { throw std::runtime_error("MOSS-TTS-Nano prompt requires target text"); } + engine::codecs::MossTokenRowBuilder builder(config.n_vq, static_cast(config.audio_pad_token_id)); if (prompt_codes == nullptr) { if (!request.prompt_text.empty() || request.has_prompt_audio) { throw std::runtime_error("MOSS-TTS-Nano continuation prompt requires matching reference audio codes"); @@ -115,24 +104,34 @@ MossTTSNanoPrompt MossTTSNanoPromptBuilder::build( const auto assistant = build_assistant_prompt_prefix(config, tokenizer_); tokens.insert(tokens.end(), assistant.begin(), assistant.end()); tokens.push_back(static_cast(config.audio_start_token_id)); - append_text_rows(prompt, tokens, config.audio_pad_token_id); - if (prompt.rows <= 0 || static_cast(prompt.input_ids.size()) != prompt.rows * prompt.row_width) { - throw std::runtime_error("MOSS-TTS-Nano prompt builder produced invalid input shape"); + builder.push_text_tokens(tokens); + } else { + if (prompt_codes->codebooks != config.n_vq) { + throw std::runtime_error("MOSS-TTS-Nano prompt audio codebook count does not match prompt row width"); + } + const auto expected = static_cast(prompt_codes->frames * prompt_codes->codebooks); + if (prompt_codes->token_ids.size() != expected) { + throw std::runtime_error("MOSS-TTS-Nano prompt audio code shape mismatch"); } - return prompt; + auto prefix = build_user_prompt_prefix(config, tokenizer_); + prefix.push_back(static_cast(config.audio_start_token_id)); + builder.push_text_tokens(prefix); + for (int64_t frame = 0; frame < prompt_codes->frames; ++frame) { + builder.push_audio_row( + static_cast(config.audio_user_slot_token_id), + prompt_codes->token_ids.data() + static_cast(frame * prompt_codes->codebooks), + prompt_codes->codebooks); + } + std::vector suffix{static_cast(config.audio_end_token_id)}; + append_text_tokens(suffix, tokenizer_, kUserTemplateAfterReference); + const auto text_tokens = tokenizer_.encode(request.text); + suffix.insert(suffix.end(), text_tokens.begin(), text_tokens.end()); + const auto assistant = build_assistant_prompt_prefix(config, tokenizer_); + suffix.insert(suffix.end(), assistant.begin(), assistant.end()); + suffix.push_back(static_cast(config.audio_start_token_id)); + builder.push_text_tokens(suffix); } - auto prefix = build_user_prompt_prefix(config, tokenizer_); - prefix.push_back(static_cast(config.audio_start_token_id)); - append_text_rows(prompt, prefix, config.audio_pad_token_id); - append_audio_rows(prompt, *prompt_codes, config.audio_user_slot_token_id, config.audio_pad_token_id); - std::vector suffix{static_cast(config.audio_end_token_id)}; - append_text_tokens(suffix, tokenizer_, kUserTemplateAfterReference); - const auto text_tokens = tokenizer_.encode(request.text); - suffix.insert(suffix.end(), text_tokens.begin(), text_tokens.end()); - const auto assistant = build_assistant_prompt_prefix(config, tokenizer_); - suffix.insert(suffix.end(), assistant.begin(), assistant.end()); - suffix.push_back(static_cast(config.audio_start_token_id)); - append_text_rows(prompt, suffix, config.audio_pad_token_id); + append_token_rows(prompt, builder.finish()); if (prompt.rows <= 0 || static_cast(prompt.input_ids.size()) != prompt.rows * prompt.row_width) { throw std::runtime_error("MOSS-TTS-Nano prompt builder produced invalid input shape"); } diff --git a/src/models/moss/shared/token_rows.cpp b/src/models/moss/shared/token_rows.cpp deleted file mode 100644 index 0d55e6caa..000000000 --- a/src/models/moss/shared/token_rows.cpp +++ /dev/null @@ -1,54 +0,0 @@ -#include "engine/models/moss/shared/token_rows.h" - -#include - -namespace engine::models::moss { - -TokenRowBuilder::TokenRowBuilder(int64_t num_codebooks, int32_t audio_pad_token_id) - : num_codebooks_(num_codebooks), - audio_pad_token_id_(audio_pad_token_id) { - if (num_codebooks_ <= 0) { - throw std::runtime_error("MOSS token row builder requires a positive codebook count"); - } -} - -void TokenRowBuilder::push_text_token(int32_t token_id) { - rows_.text_tokens.push_back(token_id); - rows_.audio_codes.insert(rows_.audio_codes.end(), static_cast(num_codebooks_), audio_pad_token_id_); -} - -void TokenRowBuilder::push_text_tokens(const std::vector & token_ids) { - for (const int32_t token_id : token_ids) { - push_text_token(token_id); - } -} - -void TokenRowBuilder::push_audio_row( - int32_t text_slot_token_id, - const std::vector> & codes, - int64_t frame) { - if (static_cast(codes.size()) != num_codebooks_) { - throw std::runtime_error("MOSS audio row codebook count mismatch"); - } - rows_.text_tokens.push_back(text_slot_token_id); - for (int64_t codebook = 0; codebook < num_codebooks_; ++codebook) { - const auto & channel = codes[static_cast(codebook)]; - if (frame < 0 || static_cast(frame) >= channel.size()) { - throw std::runtime_error("MOSS audio row frame index is out of range"); - } - rows_.audio_codes.push_back(channel[static_cast(frame)]); - } -} - -TokenRows TokenRowBuilder::finish() { - if (rows_.text_tokens.empty()) { - throw std::runtime_error("MOSS token rows must not be empty"); - } - if (static_cast(rows_.audio_codes.size()) != - static_cast(rows_.text_tokens.size()) * num_codebooks_) { - throw std::runtime_error("MOSS token rows audio code shape mismatch"); - } - return std::move(rows_); -} - -} // namespace engine::models::moss diff --git a/tests/moss_voicegen/backbone_parity.cpp b/tests/moss_voicegen/backbone_parity.cpp index 1a7838df8..20a725c93 100644 --- a/tests/moss_voicegen/backbone_parity.cpp +++ b/tests/moss_voicegen/backbone_parity.cpp @@ -1,13 +1,15 @@ -// Backbone parity for MOSS-VoiceGenerator. Builds the prompt with the audio.cpp text -// processor, embeds it the way MossTTSDelayModel.get_input_embeddings does (text embedding -// plus one embedding per audio codebook), runs the Qwen3 backbone, and compares the -// resulting hidden states against a dump from the reference PyTorch model. -// -// It also re-runs the last position through the cached single-step path, which catches -// rope/mask/cache-slot mistakes that a prefill-only comparison would miss. -// -// moss_voicegen_backbone_parity --model --prompt \ -// --hidden [--weight-type bf16] [--tolerance 0.02] +/* + * Backbone parity for MOSS-VoiceGenerator. Builds the prompt with the audio.cpp text + * processor, embeds it the way MossTTSDelayModel.get_input_embeddings does (text embedding + * plus one embedding per audio codebook), runs the Qwen3 backbone, and compares the + * resulting hidden states against a dump from the reference PyTorch model. + * + * It also re-runs the last position through the cached single-step path, which catches + * rope/mask/cache-slot mistakes that a prefill-only comparison would miss. + * + * moss_voicegen_backbone_parity --model --prompt \ + * --hidden [--weight-type bf16] [--tolerance 0.02] + */ #include "engine/community_models/moss_voicegen/assets.h" #include "engine/community_models/moss_voicegen/backbone.h" @@ -16,7 +18,6 @@ #include "engine/framework/core/execution_context.h" #include "engine/framework/io/json.h" #include "engine/framework/modules/multi_codebook_embedding.h" -#include "engine/models/moss/shared/token_rows.h" #include #include diff --git a/tests/moss_voicegen/codec_decode_parity.cpp b/tests/moss_voicegen/codec_decode_parity.cpp index bd4b75640..b900c92ea 100644 --- a/tests/moss_voicegen/codec_decode_parity.cpp +++ b/tests/moss_voicegen/codec_decode_parity.cpp @@ -1,13 +1,16 @@ -// Codec parity for MOSS-Audio-Tokenizer v1 (codes -> 24 kHz mono waveform). Decodes the -// same fixed, deterministic code matrix as the reference dumper and compares length, peak, -// RMS and a spread of individual samples against -// tests/moss_voicegen/reference/ref_codec_v1.json. -// -// This is the check that the v1 additions to moss/shared are right: the mono tail, the hop -// taken from the config, the optional stage output projection, and the v1 tensor names. -// -// moss_voicegen_codec_parity --codec \ -// --reference tests/moss_voicegen/reference/ref_codec_v1.json [--out out.wav] +/* + * Codec parity for MOSS-Audio-Tokenizer v1 (codes -> 24 kHz mono waveform). Decodes the + * same fixed, deterministic code matrix as the reference dumper and compares length, peak, + * RMS and a spread of individual samples against + * tests/moss_voicegen/reference/ref_codec_v1.json. + * + * This is the check that the v1 additions to the MOSS audio tokenizer runtime are right: + * the mono tail, the hop taken from the config, the optional stage output projection, and + * the v1 tensor names. + * + * moss_voicegen_codec_parity --codec \ + * --reference tests/moss_voicegen/reference/ref_codec_v1.json [--out out.wav] + */ #include "engine/community_models/moss_voicegen/assets.h" #include "engine/framework/audio/wav_writer.h" @@ -112,7 +115,7 @@ int main(int argc, char ** argv) { std::cout << "samples=" << audio.size() << " (reference " << expected_samples << ")\n"; std::cout << "peak=" << peak << " (reference " << expected_peak << ")\n"; std::cout << "rms=" << rms << " (reference " << expected_rms << ")\n"; - std::cout << "sample rate=" << codec.sampling_rate() << " Hz, channels=" << channels.size() << "\n"; + std::cout << "sample rate=" << codec.sampling_rate() << " Hz, channels=" << decoded.channels.size() << "\n"; bool passed = non_finite == 0 && static_cast(audio.size()) == expected_samples diff --git a/tests/moss_voicegen/generation_parity.cpp b/tests/moss_voicegen/generation_parity.cpp index 31997b199..1b6f9b9f4 100644 --- a/tests/moss_voicegen/generation_parity.cpp +++ b/tests/moss_voicegen/generation_parity.cpp @@ -1,10 +1,12 @@ -// Generation parity for MOSS-VoiceGenerator: runs the backbone, the 1 + n_vq heads and the -// delay-pattern state machine greedily, and compares the emitted rows against a dump from -// the reference PyTorch generate(). Greedy with the repetition penalty off, so a mismatch -// is a real divergence rather than an RNG difference. -// -// moss_voicegen_generation_parity --model --prompt \ -// --generation [--weight-type bf16] [--threads N] +/* + * Generation parity for MOSS-VoiceGenerator: runs the backbone, the 1 + n_vq heads and the + * delay-pattern state machine greedily, and compares the emitted rows against a dump from + * the reference PyTorch generate(). Greedy with the repetition penalty off, so a mismatch + * is a real divergence rather than an RNG difference. + * + * moss_voicegen_generation_parity --model --prompt \ + * --generation [--weight-type bf16] [--threads N] + */ #include "engine/community_models/moss_voicegen/assets.h" #include "engine/community_models/moss_voicegen/backbone.h" @@ -15,7 +17,6 @@ #include "engine/framework/core/execution_context.h" #include "engine/framework/io/json.h" #include "engine/framework/modules/multi_codebook_embedding.h" -#include "engine/models/moss/shared/token_rows.h" #include #include diff --git a/tests/moss_voicegen/voicegen_smoke.cpp b/tests/moss_voicegen/voicegen_smoke.cpp index bb725c4f2..86e41a662 100644 --- a/tests/moss_voicegen/voicegen_smoke.cpp +++ b/tests/moss_voicegen/voicegen_smoke.cpp @@ -1,8 +1,10 @@ -// End-to-end smoke test for MOSS-VoiceGenerator: designs a voice from a written -// instruction, speaks the given text in it, and writes a WAV. -// -// moss_voicegen_smoke --model --instruct "" --text "" \ -// --language English --output out.wav [--weight-type bf16] [--seed 0] [--threads N] +/* + * End-to-end smoke test for MOSS-VoiceGenerator: designs a voice from a written + * instruction, speaks the given text in it, and writes a WAV. + * + * moss_voicegen_smoke --model --instruct "" --text "" \ + * --language English --output out.wav [--weight-type bf16] [--seed 0] [--threads N] + */ #include "engine/community_models/moss_voicegen/assets.h" #include "engine/community_models/moss_voicegen/backbone.h" @@ -14,7 +16,6 @@ #include "engine/framework/core/execution_context.h" #include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h" #include "engine/framework/modules/multi_codebook_embedding.h" -#include "engine/models/moss/shared/token_rows.h" #include #include