diff --git a/CMakeLists.txt b/CMakeLists.txt index 6f6bfee09..e9f545381 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1481,6 +1481,21 @@ audiocpp_add_model(soprano_tts engine::community_models::soprano_tts::make_soprano_tts_loader ) +audiocpp_add_model(mira_tts + SOURCES + src/community_models/mira_tts/assets.cpp + src/community_models/mira_tts/decoder.cpp + src/community_models/mira_tts/generator.cpp + src/community_models/mira_tts/processor.cpp + src/community_models/mira_tts/prompt.cpp + src/community_models/mira_tts/session.cpp + src/community_models/mira_tts/speaker_encoder.cpp + INCLUDES + engine/community_models/mira_tts/session.h + LOADERS + engine::community_models::mira_tts::make_mira_tts_loader +) + audiocpp_add_model(midashenglm_gen SOURCES @@ -1910,6 +1925,7 @@ if (ENGINE_BUILD_WARMBENCH) add_engine_warmbench(irodori_tts_warm_bench tests/irodori_tts/irodori_tts_warm_bench.cpp) add_engine_warmbench(marblenet_vad_warm_bench tests/marblenet_vad/marblenet_vad_warm_bench.cpp) add_engine_warmbench(miocodec_warm_bench tests/miocodec/miocodec_warm_bench.cpp) + add_engine_warmbench(mira_tts_warm_bench tests/mira_tts/mira_tts_warm_bench.cpp) add_engine_warmbench(muscriptor_warm_bench tests/muscriptor/muscriptor_warm_bench.cpp) add_engine_warmbench(moss_tts_nano_warm_bench tests/moss_tts_nano/moss_tts_nano_warm_bench.cpp) add_engine_warmbench(moss_tts_local_warm_bench tests/moss_tts_local/moss_tts_local_warm_bench.cpp) diff --git a/README.md b/README.md index c24e0951a..606795ae3 100644 --- a/README.md +++ b/README.md @@ -152,6 +152,7 @@ Community model ports live under `community_models` to make the ownership bounda | **kroko_asr** | ASR | de, en, es, fr, it, he, nl, pt, sv, tr | Safetensors, GGUF Q8 | Mirek [@mirek190](https://github.com/mirek190) | [Kroko Community ASR](docs/community_models/kroko_asr.md) native offline/streaming Zipformer2/RNN-T transcription with word timestamps | | **minimax_h3** | Video, Music, TTS/Dialogue | auto | GGUF Q4/INT8 | [@0xShug0](https://github.com/0xShug0) | [MiniMax-H3](docs/community_models/minimax_h3.md) text-to-audio/video generation with Q4_K and optional INT8 ConvRot DiT | | **minimax_music3** | Music | auto | GGUF Q4/Q8 | [@0xShug0](https://github.com/0xShug0), [@JoeMattie](https://github.com/JoeMattie) | [MiniMax Music 3](docs/community_models/minimax_music3.md) text-to-music generation with lyrics conditioning | +| **mira_tts** | TTS, Clone | en | Local conversion | Mirek [@mirek190](https://github.com/mirek190) | [MiraTTS](docs/community_models/mira_tts.md) experimental native Qwen2 + ECAPA/Perceiver zero-shot voice cloning with progressive segment streaming (CC-BY-NC-SA-4.0 weights) | | **mms_forced_aligner** | Align | nl (nld), en (eng); pre-romanized Latin | Safetensors, GGUF 16/Q8 | Community | [MMS-300M-1130 Forced Aligner](docs/community_models/mms_forced_aligner.md) word-timestamp alignment from a wav2vec2 CTC checkpoint (safetensors or local GGUF) | | **moss_tts_local** | TTS, Clone, Ctrl | auto, optional language hint | GGUF | [@justinjohn0306](https://github.com/justinjohn0306) | MOSS-TTS-Local Transformer v1.5 support | | **moss_voicegen** | Voice Design | en, zh | GGUF | Joost [@jrohde](https://github.com/jrohde) | [MOSS-VoiceGenerator](docs/community_models/moss_voicegen.md) speech in a voice designed from a written instruction | diff --git a/docs/community_models/mira_tts.md b/docs/community_models/mira_tts.md new file mode 100644 index 000000000..b3dcedcce --- /dev/null +++ b/docs/community_models/mira_tts.md @@ -0,0 +1,78 @@ +# MiraTTS + +MiraTTS is an experimental community port of +[ysharma3501/MiraTTS](https://github.com/ysharma3501/MiraTTS), a zero-shot +voice-cloning text-to-speech model. The native path includes the Qwen2 speech +token generator, ECAPA-TDNN plus Perceiver reference encoder, finite-scalar +speaker tokenizer, conditional acoustic processor, and DAC waveform decoder. + +## Model and license + +The upstream checkpoint is +[YatharthS/MiraTTS](https://huggingface.co/YatharthS/MiraTTS). Its model card +declares `CC-BY-NC-SA-4.0`; this is a non-commercial, attribution, share-alike +license. Review that license before downloading or redistributing converted +weights. audio.cpp does not redistribute the checkpoint. + +No ready-to-run GGUF package is published yet, so MiraTTS intentionally has no +entry in the built-in download catalog. Convert a locally obtained upstream +checkpoint with: + +```bash +python tools/community_models/convert_mira_tts.py /path/to/MiraTTS /path/to/mira-native +audiocpp_gguf \ + --input language_model=/path/to/mira-native/language_model.safetensors \ + --input speaker_encoder=/path/to/mira-native/speaker_encoder.safetensors \ + --input processor=/path/to/mira-native/processor.safetensors \ + --input decoder=/path/to/mira-native/decoder.safetensors \ + --input upsampler=/path/to/mira-native/upsampler.safetensors \ + --output /path/to/mira-native/mira-tts.gguf --type bf16 \ + --family mira_tts --root /path/to/mira-native +``` + +Use BF16 for the first parity-oriented conversion. Converting the natively +BF16 Qwen backbone to F16 can overflow and produce non-finite logits. + +The converter also imports the official FastBiCodec and FlashSR component +checkpoints referenced by the upstream repository. Use `--help` to see its +component path overrides. + +## Run + +MiraTTS requires reference audio. The CLI voice-cloning request accepts the +converted model directory, target text, and a short clean reference WAV through +the normal audio.cpp TTS/clone arguments. Sampling defaults reproduce upstream: +temperature `0.8`, top-k `50`, top-p `0.95`, min-p `0.05`, and repetition +penalty `1.2`. + +The upstream pipeline decodes at 16 kHz and applies its learned FlashSR +upsampler. The native runtime executes both stages and returns 48 kHz audio. + +MiraTTS also exposes a streaming session. It splits long input at natural text +boundaries, reuses one encoded speaker identity for the whole request, and emits +each completed 48 kHz segment immediately. `text_chunk_size` controls the +maximum segment size (160 codepoints by default), while `text_chunk_mode` +selects the framework chunker. This is segment-level progressive synthesis; +the acoustic processor, DAC, and FlashSR still decode each segment as a unit. + +The session caches one encoded reference voice by default, so repeated requests +with the same audio do not rerun the speaker encoder. Increase the bounded cache +with `--session-option reference_cache_slots=`, or set it to `0` to disable +reference reuse. + +## Validation status + +- Official checkpoint conversion: validated. +- Native CUDA build: validated. +- Native CUDA smoke synthesis through all converted model components: validated. +- Segment-level streaming synthesis: validated through the native streaming + session and `/v1/audio/speech/live` route. +- Deterministic upstream comparison with identical speech/context tokens: + validated (48 kHz waveform correlation 0.99996, SNR 41.1 dB). +- End-to-end generation comparison: validated through matching tokenization, + identical 32-token speaker codes, and the first six greedy LM tokens. Later + autoregressive tokens can diverge between LMDeploy, Transformers, and the + native backend because of backend floating-point differences. + +Until end-to-end measurements are published, the family remains experimental +and is not advertised as an installable WebUI package. diff --git a/docs/community_models/models.md b/docs/community_models/models.md index fc6bc07cd..50f41828c 100644 --- a/docs/community_models/models.md +++ b/docs/community_models/models.md @@ -26,6 +26,7 @@ Practical expectations: | **mms_forced_aligner** | Align | nl (nld), en (eng); pre-romanized Latin | Community | [MMS-300M-1130 Forced Aligner](mms_forced_aligner.md) word-timestamp alignment from a wav2vec2 CTC checkpoint (safetensors or local GGUF) | | **minimax_h3** | Video, Music, TTS/Dialogue | auto | [@0xShug0](https://github.com/0xShug0) | [MiniMax-H3](minimax_h3.md) text-to-audio/video generation with Q4_K and optional INT8 ConvRot DiT | | **minimax_music3** | Music | auto | [@0xShug0](https://github.com/0xShug0) | [MiniMax Music 3](minimax_music3.md) text-to-music generation with lyrics conditioning | +| **mira_tts** | TTS, voice cloning | en | Mirek [@mirek190](https://github.com/mirek190) | [MiraTTS](mira_tts.md) experimental native Qwen2 + ECAPA/Perceiver zero-shot cloning; local conversion only (CC-BY-NC-SA-4.0 weights) | | **moss_tts_local** | TTS, voice cloning | auto, optional language hint | [@justinjohn0306](https://github.com/justinjohn0306) | [MOSS-TTS-Local Transformer v1.5](../models/moss_tts.md) support in the core model tree | | **outetts** | TTS, voice cloning | en, ar, zh, nl, fr, de, it, ja, ko, lt, ru, es, pt, be, bn, ka, hu, lv, fa, pl, sw, ta, uk | Mirek [@mirek190](https://github.com/mirek190) | [Llama-OuteTTS-1.0-1B](outetts.md) TTS and voice cloning support | | **voxcpm1** | TTS, voice cloning | zh, en, ja, ko | Community | [VoxCPM1](voxcpm1.md) tokenizer-free 0.5B TTS with 16 kHz output, streaming, and continuation-mode voice cloning | diff --git a/docs/maintainers/model_specs.md b/docs/maintainers/model_specs.md index 87a30f0fd..57c489144 100644 --- a/docs/maintainers/model_specs.md +++ b/docs/maintainers/model_specs.md @@ -156,6 +156,11 @@ packages come from the same repo, put the shared source in `package_defaults.download` and keep package-level `download` only for overrides. +Experimental ports may use an empty `packages` array while conversion and +runtime validation are still local-only. In that case `ui.recommended_package` +is omitted, so model managers do not advertise a download that cannot yet be +loaded. Community and supported families must publish at least one package. + ```json { "package_defaults": { diff --git a/include/engine/community_models/mira_tts/assets.h b/include/engine/community_models/mira_tts/assets.h new file mode 100644 index 000000000..47a0ffddb --- /dev/null +++ b/include/engine/community_models/mira_tts/assets.h @@ -0,0 +1,56 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/assets/tensor_source.h" + +#include +#include +#include + +namespace engine::community_models::mira_tts { + +struct MiraTTSConfig { + int64_t hidden_size = 896; + int64_t intermediate_size = 4864; + int64_t layers = 24; + int64_t attention_heads = 14; + int64_t kv_heads = 2; + int64_t head_dim = 64; + int64_t vocab_size = 166000; + int64_t max_position_embeddings = 32768; + float rms_norm_eps = 1.0e-6F; + float rope_theta = 1.0e6F; + int32_t bos_token_id = 151643; + int32_t eos_token_id = 151645; + int32_t speech_token_start = 155761; + int32_t speech_token_end = 163952; + int32_t prompt_speech_start = 165151; + int32_t sample_rate = 16000; + int32_t output_sample_rate = 48000; +}; + +struct MiraTTSAssets { + assets::ResourceBundle resources; + MiraTTSConfig config; + std::shared_ptr language_model_weights; + std::shared_ptr speaker_encoder_weights; + std::shared_ptr processor_weights; + std::shared_ptr decoder_weights; + std::shared_ptr upsampler_weights; +}; + +struct MiraGenerationOptions { + int64_t max_new_tokens = 1024; + int64_t top_k = 50; + float top_p = 0.95F; + float min_p = 0.05F; + float temperature = 0.8F; + float repetition_penalty = 1.2F; + uint64_t seed = 0; + bool has_seed = false; +}; + +std::shared_ptr load_mira_tts_assets( + const std::filesystem::path & model_path); + +} // namespace engine::community_models::mira_tts diff --git a/include/engine/community_models/mira_tts/decoder.h b/include/engine/community_models/mira_tts/decoder.h new file mode 100644 index 000000000..d9e6bc656 --- /dev/null +++ b/include/engine/community_models/mira_tts/decoder.h @@ -0,0 +1,33 @@ +#pragma once + +#include "engine/community_models/mira_tts/assets.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include + +namespace engine::community_models::mira_tts { + +class MiraDecoder final { +public: + MiraDecoder( + const MiraTTSAssets & assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + size_t graph_context_bytes, + assets::TensorStorageType storage_type); + ~MiraDecoder(); + + runtime::AudioBuffer decode( + const std::vector & latents, + int64_t frames); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::community_models::mira_tts diff --git a/include/engine/community_models/mira_tts/generator.h b/include/engine/community_models/mira_tts/generator.h new file mode 100644 index 000000000..2b253bd3c --- /dev/null +++ b/include/engine/community_models/mira_tts/generator.h @@ -0,0 +1,34 @@ +#pragma once + +#include "engine/community_models/mira_tts/assets.h" +#include "engine/framework/core/execution_context.h" + +#include +#include +#include +#include + +namespace engine::community_models::mira_tts { + +class MiraGenerator final { +public: + MiraGenerator( + const MiraTTSAssets & assets, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + ~MiraGenerator(); + + std::vector generate( + const std::vector & prompt_ids, + const MiraGenerationOptions & options); + void release_runtime_graphs(); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::community_models::mira_tts diff --git a/include/engine/community_models/mira_tts/processor.h b/include/engine/community_models/mira_tts/processor.h new file mode 100644 index 000000000..2caf27a8a --- /dev/null +++ b/include/engine/community_models/mira_tts/processor.h @@ -0,0 +1,33 @@ +#pragma once + +#include "engine/community_models/mira_tts/assets.h" +#include "engine/framework/core/execution_context.h" + +#include +#include +#include +#include + +namespace engine::community_models::mira_tts { + +class MiraAcousticProcessor final { +public: + MiraAcousticProcessor( + const MiraTTSAssets & assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + size_t graph_context_bytes, + assets::TensorStorageType linear_storage_type, + assets::TensorStorageType conv_storage_type); + ~MiraAcousticProcessor(); + + std::vector process( + const std::vector & speech_codes, + const std::vector & context_codes); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::community_models::mira_tts diff --git a/include/engine/community_models/mira_tts/prompt.h b/include/engine/community_models/mira_tts/prompt.h new file mode 100644 index 000000000..188b38f41 --- /dev/null +++ b/include/engine/community_models/mira_tts/prompt.h @@ -0,0 +1,26 @@ +#pragma once + +#include "engine/community_models/mira_tts/assets.h" + +#include +#include +#include +#include + +namespace engine::community_models::mira_tts { + +class MiraPromptBuilder final { +public: + explicit MiraPromptBuilder(std::shared_ptr assets); + ~MiraPromptBuilder(); + + std::vector build( + const std::string & text, + const std::vector & context_codes) const; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::community_models::mira_tts diff --git a/include/engine/community_models/mira_tts/session.h b/include/engine/community_models/mira_tts/session.h new file mode 100644 index 000000000..c5998fd8f --- /dev/null +++ b/include/engine/community_models/mira_tts/session.h @@ -0,0 +1,96 @@ +#pragma once + +#include "engine/community_models/mira_tts/assets.h" +#include "engine/framework/model_spec/metadata.h" +#include "engine/framework/runtime/cache_slots.h" +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/session_base.h" + +#include +#include +#include +#include +#include + +namespace engine::community_models::mira_tts { + +std::shared_ptr make_mira_tts_loader(); + +class MiraPromptBuilder; +class MiraSpeakerEncoder; +class MiraGenerator; +class MiraAcousticProcessor; +class MiraDecoder; + +class MiraTTSOfflineSession final : public runtime::RuntimeSessionBase, + public runtime::IOfflineVoiceTaskSession, + public runtime::IStreamingVoiceTaskSession { +public: + MiraTTSOfflineSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract); + ~MiraTTSOfflineSession() override; + + std::string family() const override; + runtime::VoiceTaskKind task_kind() const override; + runtime::RunMode run_mode() const override; + void prepare(const runtime::SessionPreparationRequest & request) override; + runtime::TaskResult run(const runtime::TaskRequest & request) override; + runtime::StreamingPolicy streaming_policy() const override; + void start_stream(const runtime::TaskRequest & request) override; + std::optional next_stream_event() override; + void set_stream_event_sink(runtime::StreamEventCallback sink) override; + runtime::TaskResult finish_stream() override; + void reset() override; + runtime::StreamEvent process_audio_chunk(const runtime::AudioChunk & chunk) override; + runtime::TaskResult finalize() override; + +private: + struct ReferenceCacheKey { + int sample_rate = 0; + int channels = 0; + uint64_t sample_count = 0; + uint64_t sample_hash = 0; + }; + + struct ReferenceCacheKeyEqual { + bool operator()( + const ReferenceCacheKey & lhs, + const ReferenceCacheKey & rhs) const noexcept; + }; + + static ReferenceCacheKey make_reference_cache_key( + const runtime::AudioBuffer & audio); + const runtime::AudioBuffer & reference_audio( + const runtime::TaskRequest & request) const; + const std::vector & context_codes( + const runtime::AudioBuffer & reference); + runtime::AudioBuffer synthesize_text( + const std::string & text, + const std::vector & context_codes, + const MiraGenerationOptions & options); + + runtime::TaskSpec task_; + std::shared_ptr assets_; + std::shared_ptr contract_; + std::optional prepared_reference_; + std::unique_ptr prompt_; + std::unique_ptr speaker_encoder_; + std::unique_ptr generator_; + std::unique_ptr processor_; + std::unique_ptr decoder_; + runtime::CacheSlots, ReferenceCacheKeyEqual> + reference_cache_; + std::optional> uncached_context_codes_; + std::vector streaming_context_codes_; + std::vector streaming_text_chunks_; + std::vector streaming_audio_chunks_; + std::optional streaming_generation_; + runtime::StreamEventCallback stream_sink_; + size_t streaming_chunk_index_ = 0; + bool streaming_started_ = false; +}; + +} // namespace engine::community_models::mira_tts diff --git a/include/engine/community_models/mira_tts/speaker_encoder.h b/include/engine/community_models/mira_tts/speaker_encoder.h new file mode 100644 index 000000000..fc58b5055 --- /dev/null +++ b/include/engine/community_models/mira_tts/speaker_encoder.h @@ -0,0 +1,34 @@ +#pragma once + +#include "engine/community_models/mira_tts/assets.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/runtime/model.h" + +#include +#include +#include +#include + +namespace engine::community_models::mira_tts { + +class MiraSpeakerEncoder final { +public: + MiraSpeakerEncoder( + const MiraTTSAssets & assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + size_t graph_context_bytes, + assets::TensorStorageType linear_storage_type, + assets::TensorStorageType conv_storage_type); + ~MiraSpeakerEncoder(); + + // Returns the 32 discrete context codes consumed by Mira's prompt and + // acoustic processor. + std::vector encode(const runtime::AudioBuffer & reference_audio); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::community_models::mira_tts diff --git a/include/engine/framework/audio/flashsr.h b/include/engine/framework/audio/flashsr.h index 2b5afa519..dcb6949cc 100644 --- a/include/engine/framework/audio/flashsr.h +++ b/include/engine/framework/audio/flashsr.h @@ -7,6 +7,10 @@ #include #include +namespace engine::assets { +class TensorSource; +} + namespace engine::audio { struct FlashSrWeights; @@ -21,6 +25,9 @@ class FlashSrModel { public: static FlashSrModel load_from_directory(const std::filesystem::path & model_dir); static FlashSrModel load_from_directory(const std::filesystem::path & model_dir, const core::BackendConfig & backend_config); + static FlashSrModel load_from_tensor_source( + std::shared_ptr source, + const core::BackendConfig & backend_config); FlashSrModel(); ~FlashSrModel(); diff --git a/include/engine/framework/sampling/hf_sampler.h b/include/engine/framework/sampling/hf_sampler.h index d4c7ebb5c..21784c184 100644 --- a/include/engine/framework/sampling/hf_sampler.h +++ b/include/engine/framework/sampling/hf_sampler.h @@ -15,6 +15,7 @@ struct HfSamplingOptions { float temperature = 1.0F; int64_t top_k = 0; float top_p = 1.0F; + float min_p = 0.0F; int64_t min_tokens_to_keep = 1; float repetition_penalty = 1.0F; }; @@ -75,6 +76,12 @@ class HfLogitsProcessor { int64_t min_tokens_to_keep, HfSamplerScratch & scratch); + static void apply_min_p( + std::vector & scores, + float min_p, + int64_t min_tokens_to_keep, + HfSamplerScratch & scratch); + static void apply_temperature(std::vector & scores, float temperature); static void build_candidates( diff --git a/model_specs/mira_tts.json b/model_specs/mira_tts.json new file mode 100644 index 000000000..9ad42c8d2 --- /dev/null +++ b/model_specs/mira_tts.json @@ -0,0 +1,215 @@ +{ + "schema_version": 1, + "family": "mira_tts", + "display_name": "MiraTTS", + "description": "MiraTTS is a community voice-cloning TTS model with a Qwen2 autoregressive speech-token generator, ECAPA/Perceiver speaker tokenizer, conditional acoustic processor, and DAC decoder.", + "category": "tts", + "status": "experimental", + "tasks": [ + "tts", + "clone" + ], + "modes": [ + "offline", + "streaming" + ], + "languages": [ + "en" + ], + "runtime": { + "tags": [ + "gguf" + ] + }, + "capabilities": { + "tts": [ + "speaker_reference" + ], + "clone": [ + "speaker_reference" + ] + }, + "options": { + "request": [ + { + "name": "max_tokens", + "type": "int", + "description": "Maximum autoregressive speech tokens; default 1024.", + "required": false, + "min": 1, + "default": 1024 + }, + { + "name": "temperature", + "type": "float", + "description": "Sampling temperature; default 0.8.", + "required": false, + "min": 0.000001, + "default": 0.8 + }, + { + "name": "top_k", + "type": "int", + "description": "Top-k sampling cutoff; default 50.", + "required": false, + "min": 1, + "default": 50 + }, + { + "name": "top_p", + "type": "float", + "description": "Nucleus sampling probability; default 0.95.", + "required": false, + "min": 0.000001, + "max": 1.0, + "default": 0.95 + }, + { + "name": "min_p", + "type": "float", + "description": "Minimum probability relative to the most likely token; default 0.05.", + "required": false, + "min": 0.0, + "max": 1.0, + "default": 0.05 + }, + { + "name": "repetition_penalty", + "type": "float", + "description": "Autoregressive repetition penalty; default 1.2.", + "required": false, + "min": 1.0, + "default": 1.2 + }, + { + "name": "text_chunk_size", + "type": "int", + "description": "Maximum codepoints per progressively emitted streaming segment; default 160.", + "required": false, + "min": 1, + "default": 160 + }, + { + "name": "text_chunk_mode", + "type": "enum", + "description": "Framework text chunking mode used by streaming synthesis.", + "values": [ + "default", + "tag_aware", + "japanese", + "endline" + ], + "required": false, + "default": "default" + }, + { + "name": "seed", + "type": "int", + "description": "Sampling seed; omitted requests choose a random seed.", + "required": false, + "min": 0 + } + ], + "session": [ + { + "name": "reference_cache_slots", + "type": "int", + "description": "Encoded reference-voice cache slots; default 1. Set to 0 to disable reuse.", + "required": false, + "min": 0, + "default": 1 + } + ], + "load": [ + { + "name": "backbone_weight_type", + "type": "enum", + "preset": "weight_type_full", + "required": false, + "default": "native", + "description": "Storage type for the Qwen2 language-model weights." + }, + { + "name": "linear_weight_type", + "type": "enum", + "preset": "weight_type_full", + "required": false, + "default": "native", + "description": "Storage type for non-convolutional speaker, processor, and decoder weights." + }, + { + "name": "conv_weight_type", + "type": "enum", + "preset": "weight_type_conv", + "required": false, + "default": "f32", + "description": "Storage type for speaker and processor convolution weights. The DAC transposed-convolution decoder remains F32 for CUDA compatibility." + } + ] + }, + "packages": [], + "dependencies": [], + "ui": { + "tags": [ + "TTS", + "Clone" + ], + "docs": [ + "docs/community_models/mira_tts.md" + ] + }, + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "files": { + "config": "model:config.json", + "tokenizer_config": "model:tokenizer_config.json", + "tokenizer_json": "model:tokenizer.json" + }, + "tensors": { + "language_model": { + "source": "weights:", + "prefix": "language_model" + }, + "speaker_encoder": { + "source": "weights:", + "prefix": "speaker_encoder" + }, + "processor": { + "source": "weights:", + "prefix": "processor" + }, + "decoder": { + "source": "weights:", + "prefix": "decoder" + }, + "upsampler": { + "source": "weights:", + "prefix": "upsampler" + } + } + }, + { + "format": "safetensors", + "roots": { + "model": "." + }, + "files": { + "config": "model:config.json", + "tokenizer_config": "model:tokenizer_config.json", + "tokenizer_json": "model:tokenizer.json" + }, + "tensors": { + "language_model": "model:language_model.safetensors", + "speaker_encoder": "model:speaker_encoder.safetensors", + "processor": "model:processor.safetensors", + "decoder": "model:decoder.safetensors", + "upsampler": "model:upsampler.safetensors" + } + } + ] +} diff --git a/src/community_models/mira_tts/assets.cpp b/src/community_models/mira_tts/assets.cpp new file mode 100644 index 000000000..8f7f69c92 --- /dev/null +++ b/src/community_models/mira_tts/assets.cpp @@ -0,0 +1,54 @@ +#include "engine/community_models/mira_tts/assets.h" + +#include "engine/framework/io/json.h" +#include "engine/framework/model_spec/package.h" + +#include + +namespace engine::community_models::mira_tts { +namespace { + +namespace json = engine::io::json; +constexpr const char * kFamily = "mira_tts"; + +MiraTTSConfig parse_config(const assets::ResourceBundle & resources) { + const auto root = resources.parse_json("config"); + if (json::require_string(root, "model_type") != "qwen2") { + throw std::runtime_error("MiraTTS config must use model_type qwen2"); + } + MiraTTSConfig out; + out.hidden_size = json::require_i64(root, "hidden_size"); + out.intermediate_size = json::require_i64(root, "intermediate_size"); + out.layers = json::require_i64(root, "num_hidden_layers"); + out.attention_heads = json::require_i64(root, "num_attention_heads"); + out.kv_heads = json::require_i64(root, "num_key_value_heads"); + out.head_dim = out.hidden_size / out.attention_heads; + out.vocab_size = json::require_i64(root, "vocab_size"); + out.max_position_embeddings = json::optional_i64( + root, "max_position_embeddings", out.max_position_embeddings); + out.rms_norm_eps = json::optional_f32(root, "rms_norm_eps", out.rms_norm_eps); + out.rope_theta = json::optional_f32(root, "rope_theta", out.rope_theta); + out.bos_token_id = static_cast(json::optional_i64( + root, "bos_token_id", out.bos_token_id)); + out.eos_token_id = static_cast(json::optional_i64( + root, "eos_token_id", out.eos_token_id)); + return out; +} + +} // namespace + +std::shared_ptr load_mira_tts_assets( + const std::filesystem::path & model_path) { + auto out = std::make_shared(); + out->resources = engine::model_spec::load_resource_bundle( + model_path, engine::model_spec::default_spec_path(kFamily)); + out->config = parse_config(out->resources); + out->language_model_weights = out->resources.open_tensor_source("language_model"); + out->speaker_encoder_weights = out->resources.open_tensor_source("speaker_encoder"); + out->processor_weights = out->resources.open_tensor_source("processor"); + out->decoder_weights = out->resources.open_tensor_source("decoder"); + out->upsampler_weights = out->resources.open_tensor_source("upsampler"); + return out; +} + +} // namespace engine::community_models::mira_tts diff --git a/src/community_models/mira_tts/decoder.cpp b/src/community_models/mira_tts/decoder.cpp new file mode 100644 index 000000000..b33ee2587 --- /dev/null +++ b/src/community_models/mira_tts/decoder.cpp @@ -0,0 +1,323 @@ +#include "engine/community_models/mira_tts/decoder.h" + +#include "engine/framework/audio/flashsr.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::community_models::mira_tts { +namespace { + +using Clock = std::chrono::steady_clock; + +struct ContextDeleter { + void operator()(ggml_context * context) const noexcept { + if (context != nullptr) ggml_free(context); + } +}; + +struct SnakeWeights { core::TensorValue alpha; }; +struct ConvWeights { + modules::Conv1dWeights value; + int64_t in_channels = 0; + int64_t out_channels = 0; + int64_t kernel = 0; +}; +struct UpWeights { + modules::ConvTranspose1dWeights value; + int64_t in_channels = 0; + int64_t out_channels = 0; + int64_t kernel = 0; +}; +struct ResidualWeights { + SnakeWeights snake1; + ConvWeights conv1; + SnakeWeights snake2; + ConvWeights conv2; +}; +struct BlockWeights { + SnakeWeights snake; + UpWeights up; + std::vector residuals; + int stride = 1; +}; +struct Weights { + std::shared_ptr store; + ConvWeights first; + std::vector blocks; + SnakeWeights final_snake; + ConvWeights final_conv; +}; + +ConvWeights load_conv( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t out_channels, + int64_t in_channels, + int64_t kernel, + assets::TensorStorageType storage_type) { + ConvWeights out; + out.in_channels = in_channels; + out.out_channels = out_channels; + out.kernel = kernel; + out.value.weight = store.load_tensor( + source, prefix + ".weight", storage_type, + {out_channels, in_channels, kernel}); + out.value.bias = store.load_f32_tensor( + source, prefix + ".bias", {out_channels}); + return out; +} + +UpWeights load_up( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t in_channels, + int64_t out_channels, + int64_t kernel, + assets::TensorStorageType storage_type) { + UpWeights out; + out.in_channels = in_channels; + out.out_channels = out_channels; + out.kernel = kernel; + out.value.weight = store.load_tensor( + source, prefix + ".weight", storage_type, + {in_channels, out_channels, kernel}); + out.value.bias = store.load_f32_tensor( + source, prefix + ".bias", {out_channels}); + return out; +} + +SnakeWeights load_snake( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + int64_t channels) { + return {store.make_from_f32( + core::TensorShape::from_dims({channels}), + assets::TensorStorageType::F32, + source.require_f32(name, {1, channels, 1}))}; +} + +ResidualWeights load_residual( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t channels, + assets::TensorStorageType storage_type) { + ResidualWeights out; + out.snake1 = load_snake(store, source, prefix + ".block.0.alpha", channels); + out.conv1 = load_conv(store, source, prefix + ".block.1", channels, channels, 7, storage_type); + out.snake2 = load_snake(store, source, prefix + ".block.2.alpha", channels); + out.conv2 = load_conv(store, source, prefix + ".block.3", channels, channels, 1, storage_type); + return out; +} + +Weights load_weights( + const MiraTTSAssets & assets, + core::ExecutionContext & execution, + size_t context_bytes, + assets::TensorStorageType storage_type) { + Weights out; + out.store = std::make_shared( + execution.backend(), execution.backend_type(), + "mira_tts.decoder.weights", context_bytes); + const auto & source = *assets.decoder_weights; + out.first = load_conv(*out.store, source, "model.0", 1536, 1024, 7, storage_type); + const int strides[] = {8, 5, 4, 2}; + const int kernels[] = {16, 11, 8, 4}; + int64_t channels = 1536; + for (int stage = 0; stage < 4; ++stage) { + const int64_t out_channels = channels / 2; + const std::string prefix = "model." + std::to_string(stage + 1) + ".block"; + BlockWeights block; + block.stride = strides[stage]; + block.snake = load_snake(*out.store, source, prefix + ".0.alpha", channels); + block.up = load_up( + *out.store, source, prefix + ".1", channels, out_channels, + kernels[stage], storage_type); + for (int residual = 0; residual < 3; ++residual) { + block.residuals.push_back(load_residual( + *out.store, source, + prefix + "." + std::to_string(residual + 2), + out_channels, storage_type)); + } + out.blocks.push_back(std::move(block)); + channels = out_channels; + } + out.final_snake = load_snake(*out.store, source, "model.5.alpha", 96); + out.final_conv = load_conv(*out.store, source, "model.6", 1, 96, 7, storage_type); + out.store->upload(); + return out; +} + +core::TensorValue conv( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const ConvWeights & weights, + int padding, + int dilation = 1) { + return modules::Conv1dModule({ + weights.in_channels, weights.out_channels, weights.kernel, + 1, padding, dilation, true}).build(ctx, input, weights.value); +} + +core::TensorValue snake( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const SnakeWeights & weights) { + return modules::Snake1dModule({input.shape.dims[1]}).build( + ctx, input, {weights.alpha}); +} + +core::TensorValue residual( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const ResidualWeights & weights, + int dilation) { + auto x = snake(ctx, input, weights.snake1); + x = conv(ctx, x, weights.conv1, 3 * dilation, dilation); + x = snake(ctx, x, weights.snake2); + x = conv(ctx, x, weights.conv2, 0); + return modules::AddModule{}.build(ctx, input, x); +} + +} // namespace + +struct MiraDecoder::Impl { + Impl( + const MiraTTSAssets & assets, + core::ExecutionContext & execution_in, + size_t weight_bytes, + size_t graph_bytes, + assets::TensorStorageType storage_type) + : execution(execution_in), + graph_context_bytes(graph_bytes), + weights(load_weights(assets, execution_in, weight_bytes, storage_type)), + upsampler(audio::FlashSrModel::load_from_tensor_source( + assets.upsampler_weights, execution_in.config())) {} + + runtime::AudioBuffer decode(const std::vector & latents, int64_t frames) { + if (frames <= 0 || latents.size() != static_cast(frames * 1024)) { + throw std::runtime_error("MiraTTS decoder expects [1024, frames] latents"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_context_bytes, nullptr, true}; + std::unique_ptr context(ggml_init(params)); + if (!context) throw std::runtime_error("failed to create MiraTTS decoder graph context"); + core::ModuleBuildContext build{ + context.get(), "mira_tts.decoder", execution.backend_type()}; + auto * input = ggml_new_tensor_3d(context.get(), GGML_TYPE_F32, frames, 1024, 1); + ggml_set_input(input); + auto x = core::wrap_tensor( + input, core::TensorShape::from_dims({1, 1024, frames}), GGML_TYPE_F32); + x = conv(build, x, weights.first, 3); + for (const auto & block : weights.blocks) { + x = snake(build, x, block.snake); + auto upsampled = modules::ConvTranspose1dModule({ + block.up.in_channels, block.up.out_channels, block.up.kernel, + block.stride, 0, 1, true}).build(build, x, block.up.value); + const int padding = static_cast(std::ceil(block.stride / 2.0)); + x = modules::SliceModule({2, padding, upsampled.shape.dims[2] - 2 * padding}) + .build(build, upsampled); + x = residual(build, x, block.residuals[0], 1); + x = residual(build, x, block.residuals[1], 3); + x = residual(build, x, block.residuals[2], 9); + } + x = snake(build, x, weights.final_snake); + x = conv(build, x, weights.final_conv, 3); + x = modules::TanhModule{}.build(build, x); + x = core::ensure_backend_addressable_layout(build, x); + ggml_set_output(x.tensor); + auto * graph = ggml_new_graph_custom(context.get(), 65536, false); + ggml_build_forward_expand(graph, x.tensor); + ggml_gallocr_t allocator = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(execution.backend())); + if (allocator == nullptr || + !ggml_gallocr_reserve(allocator, graph) || + !ggml_gallocr_alloc_graph(allocator, graph)) { + if (allocator != nullptr) ggml_gallocr_free(allocator); + throw std::runtime_error("failed to allocate MiraTTS decoder graph"); + } + engine::debug::timing_log_scalar( + "mira_tts.decoder.dac.build_allocate_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + const auto upload_start = Clock::now(); + ggml_backend_tensor_set( + input, latents.data(), 0, latents.size() * sizeof(float)); + engine::debug::timing_log_scalar( + "mira_tts.decoder.dac.upload_ms", + engine::debug::elapsed_ms(upload_start, Clock::now())); + core::set_backend_threads(execution.backend(), std::max(1, execution.config().threads)); + const auto compute_start = Clock::now(); + const auto status = core::compute_backend_graph(execution.backend(), graph); + ggml_backend_synchronize(execution.backend()); + engine::debug::timing_log_scalar( + "mira_tts.decoder.dac.compute_ms", + engine::debug::elapsed_ms(compute_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + core::release_backend_graph_resources(execution.backend(), graph); + ggml_gallocr_free(allocator); + throw std::runtime_error("MiraTTS decoder graph compute failed"); + } + const auto readback_start = Clock::now(); + std::vector decoded(static_cast(x.shape.dims[2])); + ggml_backend_tensor_get(x.tensor, decoded.data(), 0, decoded.size() * sizeof(float)); + engine::debug::timing_log_scalar( + "mira_tts.decoder.dac.readback_ms", + engine::debug::elapsed_ms(readback_start, Clock::now())); + core::release_backend_graph_resources(execution.backend(), graph); + ggml_gallocr_free(allocator); + const auto upsample_start = Clock::now(); + const auto enhanced = upsampler.super_resolve_mono_16k(decoded); + engine::debug::timing_log_scalar( + "mira_tts.decoder.flashsr_ms", + engine::debug::elapsed_ms(upsample_start, Clock::now())); + runtime::AudioBuffer audio; + audio.samples = enhanced.samples; + audio.sample_rate = enhanced.sample_rate; + audio.channels = 1; + return audio; + } + + core::ExecutionContext & execution; + size_t graph_context_bytes; + Weights weights; + audio::FlashSrModel upsampler; +}; + +MiraDecoder::MiraDecoder( + const MiraTTSAssets & assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + size_t graph_context_bytes, + assets::TensorStorageType storage_type) + : impl_(std::make_unique( + assets, execution, weight_context_bytes, graph_context_bytes, storage_type)) {} + +MiraDecoder::~MiraDecoder() = default; + +runtime::AudioBuffer MiraDecoder::decode( + const std::vector & latents, + int64_t frames) { + return impl_->decode(latents, frames); +} + +} // namespace engine::community_models::mira_tts diff --git a/src/community_models/mira_tts/generator.cpp b/src/community_models/mira_tts/generator.cpp new file mode 100644 index 000000000..a5eb78b0a --- /dev/null +++ b/src/community_models/mira_tts/generator.cpp @@ -0,0 +1,319 @@ +#include "engine/community_models/mira_tts/generator.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/modules/transformers/qwen_causal_decode_runtime.h" +#include "engine/framework/modules/weight_binding.h" +#include "engine/framework/sampling/hf_sampler.h" + +#include +#include +#include +#include +#include +#include + +namespace engine::community_models::mira_tts { +namespace { + +namespace binding = engine::modules::binding; + +struct MiraQwenWeights { + std::shared_ptr store; + core::TensorValue token_embedding; + modules::QwenDecoderStackWeights stack; + modules::NormWeights final_norm; +}; + +modules::QwenDecoderLayerWeights load_layer( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const MiraTTSConfig & config, + assets::TensorStorageType storage_type, + int64_t layer) { + const std::string prefix = "model.layers." + std::to_string(layer); + modules::QwenDecoderLayerWeights out; + out.input_norm = binding::norm_weight_from_source( + store, source, prefix + ".input_layernorm", config.hidden_size); + + const int64_t q_out = config.attention_heads * config.head_dim; + const int64_t kv_out = config.kv_heads * config.head_dim; + std::vector qkv = source.require_f32( + prefix + ".self_attn.q_proj.weight", {q_out, config.hidden_size}); + const auto k = source.require_f32( + prefix + ".self_attn.k_proj.weight", {kv_out, config.hidden_size}); + const auto v = source.require_f32( + prefix + ".self_attn.v_proj.weight", {kv_out, config.hidden_size}); + qkv.insert(qkv.end(), k.begin(), k.end()); + qkv.insert(qkv.end(), v.begin(), v.end()); + out.self_attention.qkv_weight = store.make_from_f32( + core::TensorShape::from_dims({q_out + 2 * kv_out, config.hidden_size}), + storage_type, + std::move(qkv)); + std::vector qkv_bias = source.require_f32( + prefix + ".self_attn.q_proj.bias", {q_out}); + const auto k_bias = source.require_f32( + prefix + ".self_attn.k_proj.bias", {kv_out}); + const auto v_bias = source.require_f32( + prefix + ".self_attn.v_proj.bias", {kv_out}); + qkv_bias.insert(qkv_bias.end(), k_bias.begin(), k_bias.end()); + qkv_bias.insert(qkv_bias.end(), v_bias.begin(), v_bias.end()); + out.self_attention.qkv_bias = store.make_f32( + core::TensorShape::from_dims({q_out + 2 * kv_out}), + qkv_bias); + out.self_attention.out_weight = store.load_tensor( + source, + prefix + ".self_attn.o_proj.weight", + storage_type, + {config.hidden_size, q_out}); + out.post_norm = binding::norm_weight_from_source( + store, source, prefix + ".post_attention_layernorm", config.hidden_size); + + std::vector gate_up = source.require_f32( + prefix + ".mlp.gate_proj.weight", + {config.intermediate_size, config.hidden_size}); + const auto up = source.require_f32( + prefix + ".mlp.up_proj.weight", + {config.intermediate_size, config.hidden_size}); + gate_up.insert(gate_up.end(), up.begin(), up.end()); + out.mlp.gate_up_proj = modules::LinearWeights{ + store.make_from_f32( + core::TensorShape::from_dims( + {2 * config.intermediate_size, config.hidden_size}), + storage_type, + std::move(gate_up)), + std::nullopt}; + out.mlp.down_proj = binding::linear_from_source( + store, + source, + prefix + ".mlp.down_proj", + storage_type, + config.hidden_size, + config.intermediate_size, + false); + return out; +} + +modules::QwenCausalDecoderConfig decoder_config( + const MiraTTSConfig & config, + core::BackendType backend_type) { + modules::QwenCausalDecoderConfig out; + out.stack.hidden_size = config.hidden_size; + out.stack.num_attention_heads = config.attention_heads; + out.stack.num_key_value_heads = config.kv_heads; + out.stack.head_dim = config.head_dim; + out.stack.intermediate_size = config.intermediate_size; + out.stack.layers = config.layers; + out.stack.rms_norm_eps = config.rms_norm_eps; + out.stack.rope_theta = config.rope_theta; + out.stack.rope_type = GGML_ROPE_TYPE_NEOX; + out.stack.use_qk_norm = false; + out.stack.qkv_layout = modules::QwenDecoderQKVLayout::PackedQKV; + out.stack.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; + out.stack.runtime.attention.prefill_mode = + modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.stack.runtime.attention.static_mode = + modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.stack.runtime.static_cache.update_mode = + modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.logits_size = config.vocab_size; + out.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; + out.use_lm_head_bias = false; + if (backend_type == core::BackendType::Vulkan || + backend_type == core::BackendType::Metal) { + out.lm_head_input_type = GGML_TYPE_F16; + } else if (backend_type != core::BackendType::Cpu) { + out.lm_head_input_type = GGML_TYPE_BF16; + } + return out; +} + +std::vector generation_token_ids(const MiraTTSConfig & config) { + std::vector out; + out.reserve(static_cast( + config.speech_token_end - config.speech_token_start + 2)); + for (int32_t token = config.speech_token_start; + token <= config.speech_token_end; + ++token) { + out.push_back(token); + } + out.push_back(config.eos_token_id); + return out; +} + +std::shared_ptr load_weights( + const MiraTTSAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t context_bytes, + assets::TensorStorageType storage_type) { + auto out = std::make_shared(); + out->store = std::make_shared( + backend, backend_type, "mira_tts.lm.weights", context_bytes); + const auto & config = assets.config; + const auto & source = *assets.language_model_weights; + out->token_embedding = out->store->load_tensor( + source, + "model.embed_tokens.weight", + storage_type, + {config.vocab_size, config.hidden_size}); + out->stack.layers.reserve(static_cast(config.layers)); + for (int64_t layer = 0; layer < config.layers; ++layer) { + out->stack.layers.push_back(load_layer( + *out->store, source, config, storage_type, layer)); + } + out->final_norm = binding::norm_weight_from_source( + *out->store, source, "model.norm", config.hidden_size); + out->store->upload(); + return out; +} + +modules::QwenCausalDecodeRuntimeConfig runtime_config( + const MiraTTSConfig & config, + core::BackendType backend_type, + size_t prefill_bytes, + size_t decode_bytes) { + modules::QwenCausalDecodeRuntimeConfig out; + out.trace_name = "mira_tts.lm"; + out.decoder = decoder_config(config, backend_type); + // MiraTTS generation emits one of the 8192 speech-code tokens or EOS. + // Keep the tied full-vocabulary projection for exact model semantics, but + // read back and sample only its generative alphabet. This avoids copying + // and sorting 166k logits on the host for every autoregressive step. + out.decoder.logits_size = config.vocab_size; + out.decoder.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; + out.decoder.use_lm_head_bias = false; + out.logits_readback_token_ids = generation_token_ids(config); + out.prefill_graph_arena_bytes = prefill_bytes; + out.decode_graph_arena_bytes = decode_bytes; + return out; +} + +modules::QwenCausalDecodeRuntimeWeights runtime_weights( + const MiraQwenWeights & weights) { + modules::QwenCausalDecodeRuntimeWeights out; + out.token_embedding = weights.token_embedding; + out.stack = weights.stack; + out.final_norm = weights.final_norm; + out.lm_head = modules::LinearWeights{weights.token_embedding, std::nullopt}; + return out; +} + +} // namespace + +struct MiraGenerator::Impl { + Impl( + const MiraTTSAssets & assets, + core::ExecutionContext & execution, + size_t prefill_bytes, + size_t decode_bytes, + size_t weight_bytes, + assets::TensorStorageType storage_type) + : config(assets.config), + weights(load_weights( + assets, + execution.backend(), + execution.backend_type(), + weight_bytes, + storage_type)), + runtime(std::make_unique( + execution, + runtime_config(config, execution.backend_type(), prefill_bytes, decode_bytes), + runtime_weights(*weights))) {} + + std::vector generate( + const std::vector & prompt, + const MiraGenerationOptions & options) { + if (prompt.empty()) { + throw std::runtime_error("MiraTTS LM prompt is empty"); + } + const int64_t room = config.max_position_embeddings - + static_cast(prompt.size()); + const int64_t max_tokens = std::min(options.max_new_tokens, room); + if (max_tokens <= 0) { + throw std::runtime_error("MiraTTS LM prompt exceeds its context window"); + } + auto prefill = runtime->prefill_tokens(prompt); + runtime->start_decode_tokens( + prefill.state, static_cast(prompt.size()) + max_tokens); + + sampling::HfSamplingOptions sampling_options; + sampling_options.do_sample = true; + sampling_options.temperature = options.temperature; + sampling_options.top_k = options.top_k; + sampling_options.top_p = options.top_p; + sampling_options.min_p = options.min_p; + sampling_options.repetition_penalty = options.repetition_penalty; + sampling_options.min_tokens_to_keep = 1; + sampling::HfSampler sampler; + sampling::HfSamplerScratch scratch; + scratch.reserve_vocab(static_cast(config.vocab_size)); + std::mt19937 rng(static_cast(options.seed)); + // Logits are compacted to [speech codes..., EOS]. Prompt tokens do not + // overlap this alphabet, so only generated compact ids participate in + // repetition penalty bookkeeping. + std::vector history; + std::vector codes; + auto logits = std::move(prefill.logits); + for (int64_t step = 0; step < max_tokens; ++step) { + const int32_t compact_token = sampler.sample( + logits, + history, + sampling_options, + scratch, + rng, + nullptr, + "MiraTTS LM"); + const int32_t token = compact_token == + static_cast(config.speech_token_end - + config.speech_token_start + 1) + ? config.eos_token_id + : config.speech_token_start + compact_token; + if (token == config.eos_token_id) { + break; + } + history.push_back(compact_token); + if (token >= config.speech_token_start && token <= config.speech_token_end) { + codes.push_back(token - config.speech_token_start); + } + logits = runtime->decode_token(token).logits; + } + if (codes.empty()) { + throw std::runtime_error("MiraTTS LM produced no speech tokens"); + } + return codes; + } + + MiraTTSConfig config; + std::shared_ptr weights; + std::unique_ptr runtime; +}; + +MiraGenerator::MiraGenerator( + const MiraTTSAssets & assets, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : impl_(std::make_unique( + assets, + execution, + prefill_graph_arena_bytes, + decode_graph_arena_bytes, + weight_context_bytes, + weight_storage_type)) {} + +MiraGenerator::~MiraGenerator() = default; + +std::vector MiraGenerator::generate( + const std::vector & prompt_ids, + const MiraGenerationOptions & options) { + return impl_->generate(prompt_ids, options); +} + +void MiraGenerator::release_runtime_graphs() { + impl_->runtime->release_runtime_graphs(); +} + +} // namespace engine::community_models::mira_tts diff --git a/src/community_models/mira_tts/processor.cpp b/src/community_models/mira_tts/processor.cpp new file mode 100644 index 000000000..046b10f40 --- /dev/null +++ b/src/community_models/mira_tts/processor.cpp @@ -0,0 +1,429 @@ +#include "engine/community_models/mira_tts/processor.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/streaming_conv_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace engine::community_models::mira_tts { +namespace { + +namespace binding = modules::binding; + +struct ContextDeleter { + void operator()(ggml_context * context) const noexcept { + if (context != nullptr) ggml_free(context); + } +}; + +struct ConvNeXtWeights { + modules::Conv1dWeights depthwise; + modules::NormWeights norm; + modules::LinearWeights first; + modules::LinearWeights second; + core::TensorValue gamma; +}; + +struct PlainStageWeights { + modules::Conv1dWeights embed; + modules::NormWeights norm; + std::vector blocks; + modules::NormWeights final_norm; +}; + +struct ConditionalNormWeights { + modules::LinearWeights scale; + modules::LinearWeights shift; +}; + +struct ConditionalBlockWeights { + modules::Conv1dWeights depthwise; + ConditionalNormWeights norm; + modules::LinearWeights first; + modules::LinearWeights second; + core::TensorValue gamma; +}; + +struct ProcessorWeights { + std::shared_ptr store; + core::TensorValue speech_codebook; + modules::Conv1dWeights speech_projection; + modules::LinearWeights speech_linear; + core::TensorValue context_codebook; + modules::LinearWeights context_project_out; + modules::LinearWeights speaker_project; + std::vector downsample; + modules::Conv1dWeights backbone_embed; + ConditionalNormWeights backbone_norm; + std::vector backbone_blocks; + modules::NormWeights final_norm; + modules::LinearWeights output_linear; +}; + +modules::Conv1dWeights load_conv( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t out_channels, + int64_t in_channels, + int64_t kernel, + int64_t groups = 1) { + modules::Conv1dWeights out; + out.weight = store.load_tensor( + source, prefix + ".weight", storage_type, + {out_channels, in_channels / groups, kernel}); + out.bias = store.load_f32_tensor(source, prefix + ".bias", {out_channels}); + return out; +} + +modules::LinearWeights load_linear( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t out_features, + int64_t in_features, + bool bias = true) { + return binding::linear_from_source( + store, source, prefix, storage_type, out_features, in_features, bias); +} + +modules::NormWeights load_norm( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t hidden) { + return { + store.load_f32_tensor(source, prefix + ".weight", {hidden}), + store.load_f32_tensor(source, prefix + ".bias", {hidden})}; +} + +ConvNeXtWeights load_plain_block( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType linear_type, + assets::TensorStorageType conv_type) { + ConvNeXtWeights out; + out.depthwise = load_conv(store, source, prefix + ".dwconv", conv_type, 384, 384, 7, 384); + out.norm = load_norm(store, source, prefix + ".norm", 384); + out.first = load_linear(store, source, prefix + ".pwconv1", linear_type, 2048, 384); + out.second = load_linear(store, source, prefix + ".pwconv2", linear_type, 384, 2048); + out.gamma = store.load_f32_tensor(source, prefix + ".gamma", {384}); + return out; +} + +ConditionalNormWeights load_cond_norm( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type) { + return { + load_linear(store, source, prefix + ".scale", storage_type, 384, 1024), + load_linear(store, source, prefix + ".shift", storage_type, 384, 1024)}; +} + +ProcessorWeights load_weights( + const MiraTTSAssets & assets, + core::ExecutionContext & execution, + size_t context_bytes, + assets::TensorStorageType linear_type, + assets::TensorStorageType conv_type) { + ProcessorWeights out; + out.store = std::make_shared( + execution.backend(), execution.backend_type(), + "mira_tts.processor.weights", context_bytes); + const auto & source = *assets.processor_weights; + out.speech_codebook = out.store->load_tensor( + source, "quantizer.codebook.weight", assets::TensorStorageType::F32, {8192, 8}); + out.speech_projection = load_conv( + *out.store, source, "quantizer.out_project", conv_type, 1024, 8, 1); + out.speech_linear = load_linear( + *out.store, source, "prenet.linear_pre", linear_type, 384, 1024); + out.context_codebook = out.store->load_tensor( + source, "speaker_encoder.context_codebook", assets::TensorStorageType::F32, + {4096, 6}); + out.context_project_out = load_linear( + *out.store, source, "speaker_encoder.quantizer.project_out", + linear_type, 128, 6); + out.speaker_project = load_linear( + *out.store, source, "speaker_encoder.project", linear_type, 1024, 4096); + for (int stage = 0; stage < 2; ++stage) { + const std::string prefix = "prenet.downsample." + std::to_string(stage) + ".1"; + PlainStageWeights item; + item.embed = load_conv(*out.store, source, prefix + ".embed", conv_type, 384, 384, 7); + item.norm = load_norm(*out.store, source, prefix + ".norm", 384); + for (int block = 0; block < 2; ++block) { + item.blocks.push_back(load_plain_block( + *out.store, source, + prefix + ".convnext." + std::to_string(block), + linear_type, conv_type)); + } + item.final_norm = load_norm(*out.store, source, prefix + ".final_layer_norm", 384); + out.downsample.push_back(std::move(item)); + } + out.backbone_embed = load_conv( + *out.store, source, "prenet.vocos_backbone.embed", conv_type, 384, 384, 7); + out.backbone_norm = load_cond_norm( + *out.store, source, "prenet.vocos_backbone.norm", linear_type); + for (int block = 0; block < 12; ++block) { + const std::string prefix = + "prenet.vocos_backbone.convnext." + std::to_string(block); + ConditionalBlockWeights item; + item.depthwise = load_conv( + *out.store, source, prefix + ".dwconv", conv_type, 384, 384, 7, 384); + item.norm = load_cond_norm(*out.store, source, prefix + ".norm", linear_type); + item.first = load_linear(*out.store, source, prefix + ".pwconv1", linear_type, 2048, 384); + item.second = load_linear(*out.store, source, prefix + ".pwconv2", linear_type, 384, 2048); + item.gamma = out.store->load_f32_tensor(source, prefix + ".gamma", {384}); + out.backbone_blocks.push_back(std::move(item)); + } + out.final_norm = load_norm( + *out.store, source, "prenet.vocos_backbone.final_layer_norm", 384); + out.output_linear = load_linear( + *out.store, source, "prenet.linear", linear_type, 1024, 384); + out.store->upload(); + return out; +} + +core::TensorValue linear( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::LinearWeights & weights, + int64_t in_features, + int64_t out_features) { + return modules::LinearModule({in_features, out_features, weights.bias.has_value()}) + .build(ctx, input, weights); +} + +core::TensorValue conv( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::Conv1dWeights & weights, + int64_t channels, + int64_t kernel, + int64_t groups = 1) { + if (groups == channels) { + return modules::DepthwiseConv1dModule({ + channels, kernel, 1, static_cast(kernel / 2), 1, true}) + .build(ctx, input, {weights.weight, weights.bias}); + } + if (groups != 1) { + throw std::runtime_error("MiraTTS processor only supports regular or depthwise convolution"); + } + return modules::Conv1dModule({ + channels, channels, kernel, 1, static_cast(kernel / 2), 1, true}) + .build(ctx, input, weights); +} + +core::TensorValue layer_norm( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::NormWeights & weights) { + return modules::LayerNormModule({384, 1.0e-5F, true, true}) + .build(ctx, input, weights); +} + +core::TensorValue scale_last( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & scale) { + auto shaped = core::reshape_tensor( + ctx, scale, core::TensorShape::from_dims({1, 1, 384})); + auto repeated = core::wrap_tensor( + ggml_repeat(ctx.ggml, shaped.tensor, input.tensor), input.shape, GGML_TYPE_F32); + return modules::MulModule{}.build(ctx, input, repeated); +} + +core::TensorValue plain_block( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const ConvNeXtWeights & weights) { + auto x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, input); + x = conv(ctx, x, weights.depthwise, 384, 7, 384); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); + x = layer_norm(ctx, x, weights.norm); + x = linear(ctx, x, weights.first, 384, 2048); + x = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, x); + x = linear(ctx, x, weights.second, 2048, 384); + x = scale_last(ctx, x, weights.gamma); + return modules::AddModule{}.build(ctx, input, x); +} + +core::TensorValue conditional_norm( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & condition, + const ConditionalNormWeights & weights) { + auto normalized = modules::LayerNormModule({384, 1.0e-5F, false, false}) + .build(ctx, input, {}); + auto scale = linear(ctx, condition, weights.scale, 1024, 384); + auto shift = linear(ctx, condition, weights.shift, 1024, 384); + scale = core::reshape_tensor(ctx, scale, core::TensorShape::from_dims({1, 1, 384})); + shift = core::reshape_tensor(ctx, shift, core::TensorShape::from_dims({1, 1, 384})); + auto scale_rep = core::wrap_tensor( + ggml_repeat(ctx.ggml, scale.tensor, normalized.tensor), normalized.shape, GGML_TYPE_F32); + auto shift_rep = core::wrap_tensor( + ggml_repeat(ctx.ggml, shift.tensor, normalized.tensor), normalized.shape, GGML_TYPE_F32); + auto x = modules::MulModule{}.build(ctx, normalized, scale_rep); + return modules::AddModule{}.build(ctx, x, shift_rep); +} + +} // namespace + +struct MiraAcousticProcessor::Impl { + Impl( + const MiraTTSAssets & assets, + core::ExecutionContext & execution_in, + size_t weight_bytes, + size_t graph_bytes, + assets::TensorStorageType linear_type, + assets::TensorStorageType conv_type) + : execution(execution_in), + graph_context_bytes(graph_bytes), + weights(load_weights( + assets, execution_in, weight_bytes, linear_type, conv_type)) {} + + std::vector process( + const std::vector & speech_codes, + const std::vector & context_codes) { + if (speech_codes.empty()) throw std::runtime_error("MiraTTS processor requires speech codes"); + if (context_codes.size() != 32) throw std::runtime_error("MiraTTS processor requires 32 context codes"); + const int64_t frames = static_cast(speech_codes.size()); + ggml_init_params params{graph_context_bytes, nullptr, true}; + std::unique_ptr context(ggml_init(params)); + if (!context) throw std::runtime_error("failed to create MiraTTS processor graph context"); + auto * speech = ggml_new_tensor_2d(context.get(), GGML_TYPE_I32, frames, 1); + auto * speaker = ggml_new_tensor_2d(context.get(), GGML_TYPE_I32, 32, 1); + ggml_set_input(speech); + ggml_set_input(speaker); + core::ModuleBuildContext build{context.get(), "mira_tts.processor"}; + + auto speech_ids = core::wrap_tensor( + speech, core::TensorShape::from_dims({1, frames}), GGML_TYPE_I32); + auto x = modules::EmbeddingModule({8192, 8}).build( + build, speech_ids, weights.speech_codebook); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build, x); + x = modules::Conv1dModule({8, 1024, 1, 1, 0, 1, true}) + .build(build, x, weights.speech_projection); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build, x); + x = linear(build, x, weights.speech_linear, 1024, 384); + + auto context_ids = core::wrap_tensor( + speaker, core::TensorShape::from_dims({1, 32}), GGML_TYPE_I32); + auto condition = modules::EmbeddingModule({4096, 6}).build( + build, context_ids, weights.context_codebook); + condition = linear(build, condition, weights.context_project_out, 6, 128); + // The exported processor flattens [B, 128, 32], not [B, 32, 128]. + // Preserve that channel-major speaker-conditioning order. + condition = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build, condition); + condition = core::reshape_tensor( + build, + core::ensure_backend_addressable_layout(build, condition), + core::TensorShape::from_dims({1, 4096})); + condition = linear(build, condition, weights.speaker_project, 4096, 1024); + + for (const auto & stage : weights.downsample) { + x = core::wrap_tensor(ggml_scale(build.ggml, x.tensor, 3.0F), x.shape, x.type); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build, x); + x = conv(build, x, stage.embed, 384, 7); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build, x); + x = layer_norm(build, x, stage.norm); + for (const auto & block : stage.blocks) x = plain_block(build, x, block); + x = layer_norm(build, x, stage.final_norm); + } + + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build, x); + x = conv(build, x, weights.backbone_embed, 384, 7); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build, x); + x = conditional_norm(build, x, condition, weights.backbone_norm); + for (const auto & block : weights.backbone_blocks) { + auto residual = x; + auto hidden = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build, x); + hidden = conv(build, hidden, block.depthwise, 384, 7, 384); + hidden = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build, hidden); + hidden = conditional_norm(build, hidden, condition, block.norm); + hidden = linear(build, hidden, block.first, 384, 2048); + hidden = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(build, hidden); + hidden = linear(build, hidden, block.second, 2048, 384); + hidden = scale_last(build, hidden, block.gamma); + x = modules::AddModule{}.build(build, residual, hidden); + } + x = layer_norm(build, x, weights.final_norm); + x = linear(build, x, weights.output_linear, 384, 1024); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build, x); + auto cond_bct = core::reshape_tensor( + build, condition, core::TensorShape::from_dims({1, 1024, 1})); + cond_bct = core::wrap_tensor( + ggml_repeat(build.ggml, cond_bct.tensor, x.tensor), x.shape, GGML_TYPE_F32); + x = modules::AddModule{}.build(build, x, cond_bct); + x = core::ensure_backend_addressable_layout(build, x); + ggml_set_output(x.tensor); + auto * graph = ggml_new_graph_custom(context.get(), 65536, false); + ggml_build_forward_expand(graph, x.tensor); + auto allocator = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(execution.backend())); + if (allocator == nullptr || !ggml_gallocr_alloc_graph(allocator, graph)) { + if (allocator != nullptr) ggml_gallocr_free(allocator); + throw std::runtime_error("failed to allocate MiraTTS processor graph"); + } + ggml_backend_tensor_set(speech, speech_codes.data(), 0, speech_codes.size() * sizeof(int32_t)); + ggml_backend_tensor_set(speaker, context_codes.data(), 0, context_codes.size() * sizeof(int32_t)); + core::set_backend_threads(execution.backend(), std::max(1, execution.config().threads)); + const auto status = core::compute_backend_graph(execution.backend(), graph); + ggml_backend_synchronize(execution.backend()); + if (status != GGML_STATUS_SUCCESS) { + core::release_backend_graph_resources(execution.backend(), graph); + ggml_gallocr_free(allocator); + throw std::runtime_error("MiraTTS processor graph compute failed"); + } + std::vector output(static_cast(1024 * frames)); + ggml_backend_tensor_get(x.tensor, output.data(), 0, output.size() * sizeof(float)); + core::release_backend_graph_resources(execution.backend(), graph); + ggml_gallocr_free(allocator); + return output; + } + + core::ExecutionContext & execution; + size_t graph_context_bytes; + ProcessorWeights weights; +}; + +MiraAcousticProcessor::MiraAcousticProcessor( + const MiraTTSAssets & assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + size_t graph_context_bytes, + assets::TensorStorageType linear_storage_type, + assets::TensorStorageType conv_storage_type) + : impl_(std::make_unique( + assets, execution, weight_context_bytes, graph_context_bytes, + linear_storage_type, conv_storage_type)) {} + +MiraAcousticProcessor::~MiraAcousticProcessor() = default; + +std::vector MiraAcousticProcessor::process( + const std::vector & speech_codes, + const std::vector & context_codes) { + return impl_->process(speech_codes, context_codes); +} + +} // namespace engine::community_models::mira_tts diff --git a/src/community_models/mira_tts/prompt.cpp b/src/community_models/mira_tts/prompt.cpp new file mode 100644 index 000000000..1cfec568b --- /dev/null +++ b/src/community_models/mira_tts/prompt.cpp @@ -0,0 +1,87 @@ +#include "engine/community_models/mira_tts/prompt.h" + +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include + +namespace engine::community_models::mira_tts { +namespace { + +int32_t require_token_id( + const tokenizers::LlamaBpeTokenizer & tokenizer, + const std::string & token) { + const auto id = tokenizer.find_token_id(token); + if (!id.has_value()) { + throw std::runtime_error("MiraTTS tokenizer is missing token " + token); + } + return *id; +} + +} // namespace + +struct MiraPromptBuilder::Impl { + explicit Impl(std::shared_ptr input_assets) + : assets(std::move(input_assets)) { + if (assets == nullptr) { + throw std::runtime_error("MiraTTS prompt builder requires assets"); + } + tokenizers::LlamaBpeTokenizerSpec spec; + spec.tokenizer_json_path = assets->resources.require_file("tokenizer_json"); + spec.tokenizer_config_path = assets->resources.require_file("tokenizer_config"); + spec.pre_type = tokenizers::LlamaBpePreTokenizer::Qwen2; + tokenizer = tokenizers::load_llama_bpe_tokenizer(spec); + task_tts = require_token_id(*tokenizer, "<|task_tts|>"); + start_text = require_token_id(*tokenizer, "<|start_text|>"); + end_text = require_token_id(*tokenizer, "<|end_text|>"); + context_start = require_token_id(*tokenizer, "<|context_audio_start|>"); + context_end = require_token_id(*tokenizer, "<|context_audio_end|>"); + speech_start = require_token_id(*tokenizer, "<|prompt_speech_start|>"); + context_token_start = require_token_id(*tokenizer, "<|context_token_0|>"); + } + + std::shared_ptr assets; + std::shared_ptr tokenizer; + int32_t task_tts = 0; + int32_t start_text = 0; + int32_t end_text = 0; + int32_t context_start = 0; + int32_t context_end = 0; + int32_t speech_start = 0; + int32_t context_token_start = 0; +}; + +MiraPromptBuilder::MiraPromptBuilder(std::shared_ptr assets) + : impl_(std::make_unique(std::move(assets))) {} + +MiraPromptBuilder::~MiraPromptBuilder() = default; + +std::vector MiraPromptBuilder::build( + const std::string & text, + const std::vector & context_codes) const { + if (text.empty()) { + throw std::runtime_error("MiraTTS requires non-empty text"); + } + if (context_codes.size() != 32) { + throw std::runtime_error("MiraTTS speaker encoder must produce 32 context codes"); + } + auto text_ids = impl_->tokenizer->encode(text, false); + std::vector out; + out.reserve(text_ids.size() + context_codes.size() + 6); + out.push_back(impl_->task_tts); + out.push_back(impl_->start_text); + out.insert(out.end(), text_ids.begin(), text_ids.end()); + out.push_back(impl_->end_text); + out.push_back(impl_->context_start); + for (const int32_t code : context_codes) { + if (code < 0 || code >= 4096) { + throw std::runtime_error("MiraTTS context code is outside [0, 4096)"); + } + out.push_back(impl_->context_token_start + code); + } + out.push_back(impl_->context_end); + out.push_back(impl_->speech_start); + return out; +} + +} // namespace engine::community_models::mira_tts diff --git a/src/community_models/mira_tts/session.cpp b/src/community_models/mira_tts/session.cpp new file mode 100644 index 000000000..60eecdacb --- /dev/null +++ b/src/community_models/mira_tts/session.cpp @@ -0,0 +1,393 @@ +#include "engine/community_models/mira_tts/session.h" + +#include "engine/community_models/mira_tts/decoder.h" +#include "engine/community_models/mira_tts/generator.h" +#include "engine/community_models/mira_tts/processor.h" +#include "engine/community_models/mira_tts/prompt.h" +#include "engine/community_models/mira_tts/speaker_encoder.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/runtime/spec_backed_model.h" +#include "engine/framework/text/chunking.h" + +#include +#include +#include +#include +#include +#include + +namespace engine::community_models::mira_tts { +namespace { + +constexpr const char * kFamily = "mira_tts"; +constexpr size_t kGraphBytes = 512ull * 1024ull * 1024ull; +constexpr size_t kWeightBytes = 256ull * 1024ull * 1024ull; +constexpr size_t kDefaultReferenceCacheSlots = 1; + +using Clock = std::chrono::steady_clock; + +std::shared_ptr require_assets( + std::shared_ptr value) { + if (value == nullptr) throw std::runtime_error("MiraTTS session requires assets"); + return value; +} + +std::shared_ptr require_contract( + std::shared_ptr value) { + if (value == nullptr) throw std::runtime_error("MiraTTS session requires a model contract"); + return value; +} + +MiraGenerationOptions generation_options(const runtime::TaskRequest & request) { + MiraGenerationOptions out; + if (const auto value = runtime::parse_i64_option(request.options, {"max_tokens"})) { + out.max_new_tokens = *value; + } + if (const auto value = runtime::parse_i64_option(request.options, {"top_k"})) { + out.top_k = *value; + } + if (const auto value = runtime::parse_finite_float_option(request.options, {"top_p"})) { + out.top_p = *value; + } + if (const auto value = runtime::parse_finite_float_option(request.options, {"min_p"})) { + out.min_p = *value; + } + if (const auto value = runtime::parse_finite_float_option(request.options, {"temperature"})) { + out.temperature = *value; + } + if (const auto value = runtime::parse_finite_float_option( + request.options, {"repetition_penalty"})) { + out.repetition_penalty = *value; + } + if (const auto value = runtime::parse_u64_option(request.options, {"seed"})) { + out.seed = *value; + out.has_seed = true; + } + if (!out.has_seed) out.seed = runtime::random_u64_seed(); + if (out.max_new_tokens < 1 || out.top_k < 1 || out.temperature <= 0.0F || + out.top_p <= 0.0F || out.top_p > 1.0F || out.min_p < 0.0F || + out.min_p > 1.0F || out.repetition_penalty < 1.0F) { + throw std::runtime_error("MiraTTS generation options are outside their valid ranges"); + } + return out; +} + +size_t reference_cache_slots(const runtime::SessionOptions & options) { + const int64_t slots = runtime::parse_i64_option( + options.options, + {"mira_tts.reference_cache_slots", "reference_cache_slots"}) + .value_or(static_cast(kDefaultReferenceCacheSlots)); + if (slots < 0) { + throw std::runtime_error( + "mira_tts.reference_cache_slots must be non-negative"); + } + if (static_cast(slots) > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "mira_tts.reference_cache_slots is too large"); + } + return static_cast(slots); +} + +uint64_t mix_reference_hash(uint64_t hash, uint64_t value) { + hash ^= value; + hash *= 1099511628211ull; + return hash; +} + +std::unique_ptr create_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options, + std::shared_ptr assets, + std::shared_ptr contract) { + return std::make_unique( + task, options, std::move(assets), std::move(contract)); +} + +} // namespace + +MiraTTSOfflineSession::MiraTTSOfflineSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract) + : RuntimeSessionBase(options), + task_(task), + assets_(require_assets(std::move(assets))), + contract_(require_contract(std::move(contract))), + reference_cache_(reference_cache_slots(this->options())) { + runtime::validate_spec_backed_session_options( + options, *contract_, kFamily, "MiraTTS"); + if ((task.mode != runtime::RunMode::Offline && + task.mode != runtime::RunMode::Streaming) || + (task.task != runtime::VoiceTaskKind::Tts && + task.task != runtime::VoiceTaskKind::VoiceCloning)) { + throw std::runtime_error( + "MiraTTS supports offline and streaming TTS/voice cloning only"); + } + const auto lm_type = runtime::parse_tensor_storage_option( + options.options, "backbone_weight_type", assets::TensorStorageType::Native, + {assets::TensorStorageType::Native, assets::TensorStorageType::F32, + assets::TensorStorageType::F16, assets::TensorStorageType::BF16, + assets::TensorStorageType::Q8_0}); + const auto linear_type = runtime::parse_tensor_storage_option( + options.options, "linear_weight_type", assets::TensorStorageType::Native, + {assets::TensorStorageType::Native, assets::TensorStorageType::F32, + assets::TensorStorageType::F16, assets::TensorStorageType::BF16, + assets::TensorStorageType::Q8_0}); + const auto conv_type = runtime::parse_tensor_storage_option( + options.options, "conv_weight_type", assets::TensorStorageType::F32, + {assets::TensorStorageType::Native, assets::TensorStorageType::F32, + assets::TensorStorageType::F16, assets::TensorStorageType::BF16}); + auto & execution = execution_context(); + prompt_ = std::make_unique(assets_); + speaker_encoder_ = std::make_unique( + *assets_, execution, kWeightBytes, kGraphBytes, linear_type, conv_type); + generator_ = std::make_unique( + *assets_, execution, kGraphBytes, kGraphBytes, kWeightBytes, lm_type); + processor_ = std::make_unique( + *assets_, execution, kWeightBytes, kGraphBytes, linear_type, conv_type); + // ggml's current CUDA ConvTranspose1d kernel requires F32 weights. + decoder_ = std::make_unique( + *assets_, execution, kWeightBytes, kGraphBytes, + assets::TensorStorageType::F32); +} + +MiraTTSOfflineSession::~MiraTTSOfflineSession() = default; + +std::string MiraTTSOfflineSession::family() const { return kFamily; } + +runtime::VoiceTaskKind MiraTTSOfflineSession::task_kind() const { + return task_.task; +} + +runtime::RunMode MiraTTSOfflineSession::run_mode() const { return task_.mode; } + +void MiraTTSOfflineSession::prepare( + const runtime::SessionPreparationRequest & request) { + runtime::validate_spec_backed_request_options( + request.options, *contract_, "MiraTTS"); + prepared_reference_.reset(); + if (request.voice.has_value() && request.voice->speaker.has_value() && + request.voice->speaker->audio.has_value()) { + prepared_reference_ = *request.voice->speaker->audio; + (void)context_codes(*prepared_reference_); + } + mark_prepared(); +} + +bool MiraTTSOfflineSession::ReferenceCacheKeyEqual::operator()( + const ReferenceCacheKey & lhs, + const ReferenceCacheKey & rhs) const noexcept { + return lhs.sample_rate == rhs.sample_rate && + lhs.channels == rhs.channels && + lhs.sample_count == rhs.sample_count && + lhs.sample_hash == rhs.sample_hash; +} + +MiraTTSOfflineSession::ReferenceCacheKey +MiraTTSOfflineSession::make_reference_cache_key( + const runtime::AudioBuffer & audio) { + uint64_t hash = 1469598103934665603ull; + for (const float sample : audio.samples) { + uint32_t bits = 0; + std::memcpy(&bits, &sample, sizeof(bits)); + hash = mix_reference_hash(hash, static_cast(bits)); + } + return ReferenceCacheKey{ + audio.sample_rate, + audio.channels, + static_cast(audio.samples.size()), + hash, + }; +} + +const runtime::AudioBuffer & MiraTTSOfflineSession::reference_audio( + const runtime::TaskRequest & request) const { + if (request.voice.has_value() && request.voice->speaker.has_value() && + request.voice->speaker->audio.has_value()) { + return *request.voice->speaker->audio; + } + if (request.audio_input.has_value()) return *request.audio_input; + if (prepared_reference_.has_value()) return *prepared_reference_; + throw std::runtime_error( + "MiraTTS requires a reference voice in voice.speaker.audio or audio_input"); +} + +const std::vector & MiraTTSOfflineSession::context_codes( + const runtime::AudioBuffer & reference) { + const auto key_start = Clock::now(); + auto key = make_reference_cache_key(reference); + engine::debug::timing_log_scalar( + "mira_tts.reference.hash_ms", engine::debug::elapsed_ms(key_start)); + if (const auto * cached = reference_cache_.find(key)) { + engine::debug::trace_log_scalar("mira_tts.reference.cache_hit", true); + return *cached; + } + + engine::debug::trace_log_scalar("mira_tts.reference.cache_hit", false); + const auto encode_start = Clock::now(); + auto encoded = speaker_encoder_->encode(reference); + engine::debug::timing_log_scalar( + "mira_tts.reference.encode_ms", engine::debug::elapsed_ms(encode_start)); + if (reference_cache_.capacity() == 0) { + uncached_context_codes_ = std::move(encoded); + return *uncached_context_codes_; + } + reference_cache_.put(key, std::move(encoded)); + return *reference_cache_.find(key); +} + +runtime::TaskResult MiraTTSOfflineSession::run( + const runtime::TaskRequest & request) { + require_prepared("MiraTTS run"); + runtime::validate_spec_backed_request_options( + request.options, *contract_, "MiraTTS"); + if (!request.text_input.has_value() || request.text_input->text.empty()) { + throw std::runtime_error("MiraTTS requires non-empty text input"); + } + if (task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("MiraTTS run requires an offline session"); + } + const auto & codes = context_codes(reference_audio(request)); + runtime::TaskResult result; + result.audio_output = synthesize_text( + request.text_input->text, codes, generation_options(request)); + return result; +} + +runtime::AudioBuffer MiraTTSOfflineSession::synthesize_text( + const std::string & text, + const std::vector & context_codes, + const MiraGenerationOptions & options) { + const auto prompt_start = Clock::now(); + const auto prompt_ids = prompt_->build(text, context_codes); + engine::debug::timing_log_scalar( + "mira_tts.prompt_ms", engine::debug::elapsed_ms(prompt_start)); + const auto generator_start = Clock::now(); + const auto speech_codes = generator_->generate( + prompt_ids, options); + engine::debug::timing_log_scalar( + "mira_tts.generator_ms", engine::debug::elapsed_ms(generator_start)); + if (speech_codes.empty()) { + throw std::runtime_error("MiraTTS generated no speech tokens"); + } + const auto processor_start = Clock::now(); + const auto latents = processor_->process(speech_codes, context_codes); + engine::debug::timing_log_scalar( + "mira_tts.processor_ms", engine::debug::elapsed_ms(processor_start)); + const auto decoder_start = Clock::now(); + auto audio = decoder_->decode( + latents, static_cast(speech_codes.size())); + engine::debug::timing_log_scalar( + "mira_tts.decoder_ms", engine::debug::elapsed_ms(decoder_start)); + return audio; +} + +runtime::StreamingPolicy MiraTTSOfflineSession::streaming_policy() const { + runtime::StreamingPolicy policy; + policy.input = runtime::StreamingInputKind::None; + policy.output = runtime::StreamingOutputKind::PullEvents; + return policy; +} + +void MiraTTSOfflineSession::start_stream( + const runtime::TaskRequest & request) { + require_prepared("MiraTTS streaming"); + runtime::validate_spec_backed_request_options( + request.options, *contract_, "MiraTTS"); + if (task_.mode != runtime::RunMode::Streaming) { + throw std::runtime_error("MiraTTS start_stream requires a streaming session"); + } + if (!request.text_input.has_value() || request.text_input->text.empty()) { + throw std::runtime_error("MiraTTS streaming requires non-empty text input"); + } + reset(); + const int64_t chunk_size = engine::text::parse_text_chunk_size_override( + request.options).value_or(160); + const auto chunk_mode = engine::text::parse_text_chunk_mode_override( + request.options).value_or(engine::text::TextChunkMode::Default); + streaming_text_chunks_ = engine::text::split_text_chunks( + request.text_input->text, chunk_size, chunk_mode); + if (streaming_text_chunks_.empty()) { + throw std::runtime_error("MiraTTS streaming text chunking produced no segments"); + } + streaming_context_codes_ = context_codes(reference_audio(request)); + streaming_generation_ = generation_options(request); + streaming_started_ = true; +} + +std::optional +MiraTTSOfflineSession::next_stream_event() { + if (!streaming_started_ || !streaming_generation_.has_value()) { + throw std::runtime_error("MiraTTS streaming has not been started"); + } + if (streaming_chunk_index_ >= streaming_text_chunks_.size()) { + return std::nullopt; + } + const size_t index = streaming_chunk_index_++; + auto options = *streaming_generation_; + options.seed += index; + auto audio = synthesize_text( + streaming_text_chunks_[index], streaming_context_codes_, options); + streaming_audio_chunks_.push_back(audio); + runtime::StreamEvent event; + event.audio_output = std::move(audio); + if (stream_sink_) { + stream_sink_(event); + } + return event; +} + +void MiraTTSOfflineSession::set_stream_event_sink( + runtime::StreamEventCallback sink) { + stream_sink_ = std::move(sink); +} + +runtime::TaskResult MiraTTSOfflineSession::finish_stream() { + if (!streaming_started_) { + throw std::runtime_error("MiraTTS streaming has not been started"); + } + runtime::TaskResult result; + runtime::AudioBuffer merged; + for (const auto & chunk : streaming_audio_chunks_) { + runtime::append_audio_buffer(merged, chunk); + } + if (merged.sample_rate == 0) { + throw std::runtime_error("MiraTTS streaming produced no audio chunks"); + } + result.audio_output = std::move(merged); + reset(); + return result; +} + +void MiraTTSOfflineSession::reset() { + streaming_context_codes_.clear(); + streaming_text_chunks_.clear(); + streaming_audio_chunks_.clear(); + streaming_generation_.reset(); + streaming_chunk_index_ = 0; + streaming_started_ = false; +} + +runtime::StreamEvent MiraTTSOfflineSession::process_audio_chunk( + const runtime::AudioChunk & chunk) { + (void)chunk; + throw std::runtime_error("MiraTTS streaming does not consume audio chunks"); +} + +runtime::TaskResult MiraTTSOfflineSession::finalize() { + return finish_stream(); +} + +std::shared_ptr make_mira_tts_loader() { + runtime::SpecBackedVoiceModelConfig config; + config.family = kFamily; + config.aliases = {"mira", "MiraTTS"}; + config.load_assets = load_mira_tts_assets; + config.create_session = create_session; + return runtime::make_spec_backed_voice_loader(std::move(config)); +} + +} // namespace engine::community_models::mira_tts diff --git a/src/community_models/mira_tts/speaker_encoder.cpp b/src/community_models/mira_tts/speaker_encoder.cpp new file mode 100644 index 000000000..4fe40e03b --- /dev/null +++ b/src/community_models/mira_tts/speaker_encoder.cpp @@ -0,0 +1,584 @@ +#include "engine/community_models/mira_tts/speaker_encoder.h" + +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/dsp.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/conditioning_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/optimizations/fast_conv_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 + +namespace engine::community_models::mira_tts { +namespace { + +namespace binding = modules::binding; + +constexpr int kSampleRate = 16000; +constexpr int64_t kReferenceSamples = 96000; +constexpr int64_t kMelBins = 128; +constexpr int64_t kEcapaChannels = 512; +constexpr int64_t kRes2Width = 64; +constexpr int64_t kPerceiverDim = 128; +constexpr int64_t kPerceiverLatents = 32; +constexpr int64_t kPerceiverHeads = 8; +constexpr int64_t kPerceiverInner = 512; +constexpr float kBatchNormEps = 1.0e-5F; + +struct ContextDeleter { + void operator()(ggml_context * context) const noexcept { + if (context != nullptr) ggml_free(context); + } +}; + +struct ConvWeights { + modules::Conv1dWeights value; + int64_t in_channels = 0; + int64_t out_channels = 0; + int64_t kernel = 1; + int64_t padding = 0; + int64_t dilation = 1; +}; + +struct TdnnWeights { + ConvWeights conv; + modules::BatchNorm1dEvalWeights norm; +}; + +struct SeRes2Weights { + TdnnWeights first; + std::vector res2; + TdnnWeights second; + modules::LinearWeights se_first; + modules::LinearWeights se_second; +}; + +struct PerceiverLayerWeights { + modules::LinearWeights q; + modules::LinearWeights kv; + modules::LinearWeights out; + modules::LinearWeights ff_in; + modules::LinearWeights ff_out; +}; + +struct SpeakerWeights { + std::shared_ptr store; + TdnnWeights input; + std::vector blocks; + ConvWeights mfa; + modules::LinearWeights project_context; + core::TensorValue latents; + std::vector perceiver; + core::TensorValue norm_gamma; + modules::LinearWeights quant_project; +}; + +std::vector require_values( + const assets::TensorSource & source, + const std::string & name, + int64_t size) { + auto tensor = source.require_f32_tensor(name); + if (static_cast(tensor.values.size()) != size) { + throw std::runtime_error("MiraTTS tensor size mismatch: " + name); + } + return std::move(tensor.values); +} + +modules::BatchNorm1dEvalWeights load_batch_norm( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t channels) { + const auto gamma = require_values(source, prefix + ".weight", channels); + const auto beta = require_values(source, prefix + ".bias", channels); + const auto mean = require_values(source, prefix + ".running_mean", channels); + const auto variance = require_values(source, prefix + ".running_var", channels); + std::vector scale(static_cast(channels)); + std::vector bias(static_cast(channels)); + for (int64_t i = 0; i < channels; ++i) { + scale[static_cast(i)] = gamma[static_cast(i)] / + std::sqrt(variance[static_cast(i)] + kBatchNormEps); + bias[static_cast(i)] = beta[static_cast(i)] - + mean[static_cast(i)] * scale[static_cast(i)]; + } + return { + store.make_f32(core::TensorShape::from_dims({channels}), scale), + store.make_f32(core::TensorShape::from_dims({channels}), bias)}; +} + +ConvWeights load_conv( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t out_channels, + int64_t in_channels, + int64_t kernel, + int64_t padding, + int64_t dilation) { + ConvWeights out; + out.in_channels = in_channels; + out.out_channels = out_channels; + out.kernel = kernel; + out.padding = padding; + out.dilation = dilation; + out.value.weight = store.load_tensor( + source, prefix + ".weight", storage_type, + {out_channels, in_channels, kernel}); + out.value.bias = store.load_f32_tensor(source, prefix + ".bias", {out_channels}); + return out; +} + +TdnnWeights load_tdnn( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t out_channels, + int64_t in_channels, + int64_t kernel, + int64_t padding, + int64_t dilation) { + return { + load_conv(store, source, prefix + ".conv", storage_type, + out_channels, in_channels, kernel, padding, dilation), + load_batch_norm(store, source, prefix + ".bn", out_channels)}; +} + +modules::LinearWeights load_linear( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t out_features, + int64_t in_features, + bool bias) { + return binding::linear_from_source( + store, source, prefix, storage_type, out_features, in_features, bias); +} + +SpeakerWeights load_weights( + const MiraTTSAssets & assets, + core::ExecutionContext & execution, + size_t context_bytes, + assets::TensorStorageType linear_type, + assets::TensorStorageType conv_type) { + SpeakerWeights out; + out.store = std::make_shared( + execution.backend(), execution.backend_type(), + "mira_tts.speaker_encoder.weights", context_bytes); + const auto & source = *assets.speaker_encoder_weights; + const std::string root = "speaker_encoder.speaker_encoder."; + out.input = load_tdnn( + *out.store, source, root + "layer1", conv_type, + kEcapaChannels, kMelBins, 5, 2, 1); + for (int layer = 2; layer <= 4; ++layer) { + const std::string prefix = root + "layer" + std::to_string(layer) + ".se_res2block."; + SeRes2Weights block; + block.first = load_tdnn( + *out.store, source, prefix + "0", conv_type, + kEcapaChannels, kEcapaChannels, 1, 0, 1); + const int64_t dilation = layer; + for (int branch = 0; branch < 7; ++branch) { + TdnnWeights branch_weights; + branch_weights.conv = load_conv( + *out.store, source, + prefix + "1.convs." + std::to_string(branch), conv_type, + kRes2Width, kRes2Width, 3, dilation, dilation); + branch_weights.norm = load_batch_norm( + *out.store, source, + prefix + "1.bns." + std::to_string(branch), kRes2Width); + block.res2.push_back(std::move(branch_weights)); + } + block.second = load_tdnn( + *out.store, source, prefix + "2", conv_type, + kEcapaChannels, kEcapaChannels, 1, 0, 1); + block.se_first = load_linear( + *out.store, source, prefix + "3.linear1", linear_type, 128, 512, true); + block.se_second = load_linear( + *out.store, source, prefix + "3.linear2", linear_type, 512, 128, true); + out.blocks.push_back(std::move(block)); + } + out.mfa = load_conv( + *out.store, source, root + "conv", conv_type, + 1536, 1536, 1, 0, 1); + out.project_context.weight = out.store->load_tensor( + source, "perceiver.proj_context.weight", linear_type, {128, 1536}); + out.project_context.bias = out.store->load_f32_tensor( + source, "speaker_encoder.perceiver_sampler.proj_context.bias", {128}); + out.latents = out.store->load_f32_tensor( + source, "perceiver.latents", {1, kPerceiverLatents, kPerceiverDim}); + for (int layer = 0; layer < 2; ++layer) { + const std::string prefix = "perceiver.layers." + std::to_string(layer); + const std::string bias_prefix = "speaker_encoder.perceiver_sampler.layers." + + std::to_string(layer) + ".1."; + PerceiverLayerWeights item; + item.q = load_linear(*out.store, source, prefix + ".attn.q", linear_type, 512, 128, false); + item.kv = load_linear(*out.store, source, prefix + ".attn.kv", linear_type, 1024, 128, false); + item.out = load_linear(*out.store, source, prefix + ".attn.out", linear_type, 128, 512, false); + item.ff_in.weight = out.store->load_tensor( + source, prefix + ".ff.in.weight", linear_type, {682, 128}); + item.ff_in.bias = out.store->load_f32_tensor(source, bias_prefix + "0.bias", {682}); + item.ff_out.weight = out.store->load_tensor( + source, prefix + ".ff.out.weight", linear_type, {128, 341}); + item.ff_out.bias = out.store->load_f32_tensor(source, bias_prefix + "2.bias", {128}); + out.perceiver.push_back(std::move(item)); + } + out.norm_gamma = out.store->load_f32_tensor( + source, "speaker_encoder.perceiver_sampler.norm.gamma", {128}); + out.quant_project.weight = out.store->load_tensor( + source, "quantizer.project_in.weight", linear_type, {6, 128}); + out.quant_project.bias = out.store->load_f32_tensor( + source, "speaker_encoder.quantizer.project_in.bias", {6}); + out.store->upload(); + return out; +} + +core::TensorValue conv1d( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const ConvWeights & weights) { + return modules::FastConv1dModule({ + weights.in_channels, weights.out_channels, weights.kernel, 1, + static_cast(weights.padding), static_cast(weights.dilation), true}, + modules::FastConv1dKind::MinittsFast1dIm2col) + .build(ctx, input, weights.value); +} + +core::TensorValue tdnn( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const TdnnWeights & weights) { + auto x = conv1d(ctx, input, weights.conv); + x = modules::ReluModule{}.build(ctx, x); + return modules::BatchNorm1dEvalModule({weights.conv.out_channels}) + .build(ctx, x, weights.norm); +} + +core::TensorValue se_res2( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const SeRes2Weights & weights) { + auto x = tdnn(ctx, input, weights.first); + core::TensorValue merged; + core::TensorValue previous; + for (int branch = 0; branch < 8; ++branch) { + auto chunk = modules::SliceModule({1, branch * kRes2Width, kRes2Width}) + .build(ctx, x); + core::TensorValue current; + if (branch == 7) { + current = chunk; + } else { + if (branch > 0) { + chunk = modules::AddModule{}.build(ctx, chunk, previous); + } + current = tdnn(ctx, chunk, weights.res2[static_cast(branch)]); + previous = current; + } + merged = merged.valid() + ? modules::ConcatModule({1}).build(ctx, merged, current) + : current; + } + x = tdnn(ctx, merged, weights.second); + auto pooled = modules::ReduceMeanModule({2}).build(ctx, x); + pooled = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, pooled); + auto gate = modules::LinearModule({512, 128, true}).build(ctx, pooled, weights.se_first); + gate = modules::ReluModule{}.build(ctx, gate); + gate = modules::LinearModule({128, 512, true}).build(ctx, gate, weights.se_second); + gate = modules::SigmoidModule{}.build(ctx, gate); + gate = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, gate); + gate = modules::RepeatModule({x.shape}).build(ctx, gate); + return modules::AddModule{}.build( + ctx, input, modules::MulModule{}.build(ctx, x, gate)); +} + +core::TensorValue reshape_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & input) { + auto x = core::ensure_backend_addressable_layout(ctx, input); + x = core::reshape_tensor(ctx, x, core::TensorShape::from_dims( + {input.shape.dims[0], input.shape.dims[1], kPerceiverHeads, + kPerceiverInner / kPerceiverHeads})); + return modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, x); +} + +core::TensorValue scale( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + float value) { + return core::wrap_tensor( + ggml_scale(ctx.ggml, input.tensor, value), input.shape, GGML_TYPE_F32); +} + +core::TensorValue perceiver_attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & latents, + const core::TensorValue & context, + const PerceiverLayerWeights & weights) { + const auto full_context = modules::ConcatModule({1}).build(ctx, latents, context); + auto q = modules::LinearModule({128, 512, false}).build(ctx, latents, weights.q); + auto kv = modules::LinearModule({128, 1024, false}).build(ctx, full_context, weights.kv); + auto k = modules::SliceModule({2, 0, 512}).build(ctx, kv); + auto v = modules::SliceModule({2, 512, 512}).build(ctx, kv); + q = reshape_heads(ctx, q); + k = reshape_heads(ctx, k); + v = reshape_heads(ctx, v); + auto scores = modules::MatMulModule{}.build( + ctx, q, modules::TransposeModule({{0, 1, 3, 2}, 4}).build(ctx, k)); + scores = scale(ctx, scores, 1.0F / std::sqrt(64.0F)); + auto attention = core::wrap_tensor( + ggml_soft_max(ctx.ggml, + core::ensure_backend_addressable_layout(ctx, scores).tensor), + scores.shape, GGML_TYPE_F32); + auto x = modules::MatMulModule{}.build(ctx, attention, v); + x = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, x); + x = core::ensure_backend_addressable_layout(ctx, x); + x = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, 32, 512})); + return modules::LinearModule({512, 128, false}).build(ctx, x, weights.out); +} + +core::TensorValue perceiver_ff( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const PerceiverLayerWeights & weights) { + auto x = modules::LinearModule({128, 682, true}).build(ctx, input, weights.ff_in); + auto value = modules::SliceModule({2, 0, 341}).build(ctx, x); + auto gate = modules::SliceModule({2, 341, 341}).build(ctx, x); + gate = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, gate); + x = modules::MulModule{}.build(ctx, value, gate); + return modules::LinearModule({341, 128, true}).build(ctx, x, weights.ff_out); +} + +core::TensorValue rms_norm( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & gamma) { + auto x = modules::RMSNormModule({128, 1.0e-5F, true, false}) + .build(ctx, input, {gamma, std::nullopt}); + return x; +} + +core::TensorValue build_graph( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const SpeakerWeights & weights) { + auto x = tdnn(ctx, input, weights.input); + std::vector outputs; + for (const auto & block : weights.blocks) { + x = se_res2(ctx, x, block); + outputs.push_back(x); + } + x = modules::ConcatModule({1}).build(ctx, outputs[0], outputs[1]); + x = modules::ConcatModule({1}).build(ctx, x, outputs[2]); + x = modules::ReluModule{}.build(ctx, conv1d(ctx, x, weights.mfa)); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); + x = modules::LinearModule({1536, 128, true}) + .build(ctx, x, weights.project_context); + auto latents = modules::RepeatModule({core::TensorShape::from_dims({1, 32, 128})}) + .build(ctx, weights.latents); + for (const auto & layer : weights.perceiver) { + latents = modules::AddModule{}.build( + ctx, latents, perceiver_attention(ctx, latents, x, layer)); + latents = modules::AddModule{}.build( + ctx, latents, perceiver_ff(ctx, latents, layer)); + } + latents = rms_norm(ctx, latents, weights.norm_gamma); + return modules::LinearModule({128, 6, true}) + .build(ctx, latents, weights.quant_project); +} + +std::vector prepare_reference(const runtime::AudioBuffer & audio) { + if (audio.sample_rate <= 0 || audio.channels <= 0 || audio.samples.empty()) { + throw std::runtime_error("MiraTTS requires non-empty reference audio"); + } + auto mono = engine::audio::convert_interleaved_audio_to_mono_linear_resampled( + audio.samples, audio.sample_rate, audio.channels, kSampleRate); + if (mono.empty()) { + throw std::runtime_error("MiraTTS reference audio contains no samples"); + } + // librosa.load(..., duration=8, sr=16000) in the upstream encoder limits + // the signal before volume normalization and six-second tiling/truncation. + mono.resize(std::min(mono.size(), 8 * kSampleRate)); + std::vector magnitudes; + magnitudes.reserve(mono.size()); + for (float sample : mono) magnitudes.push_back(std::abs(sample)); + std::sort(magnitudes.begin(), magnitudes.end()); + if (magnitudes.back() < 0.1F) { + const float divisor = std::max(magnitudes.back(), 1.0e-3F); + for (float & sample : mono) sample = sample / divisor * 0.1F; + for (float & magnitude : magnitudes) magnitude = magnitude / divisor * 0.1F; + } + const auto first_significant = std::upper_bound( + magnitudes.begin(), magnitudes.end(), 0.01F); + const size_t significant = static_cast(magnitudes.end() - first_significant); + if (significant > 10) { + const size_t begin = static_cast(0.90 * significant); + const size_t end = static_cast(0.99 * significant); + float sum = 0.0F; + for (size_t i = begin; i < end; ++i) { + sum += *(first_significant + static_cast(i)); + } + const float volume = sum / static_cast(std::max(1, end - begin)); + const float gain = std::clamp(0.2F / volume, 0.1F, 10.0F); + for (float & sample : mono) sample *= gain; + } + float peak = 0.0F; + for (float sample : mono) peak = std::max(peak, std::abs(sample)); + if (peak > 1.0F) { + for (float & sample : mono) sample /= peak; + } + std::vector fixed(static_cast(kReferenceSamples)); + for (int64_t i = 0; i < kReferenceSamples; ++i) { + fixed[static_cast(i)] = mono[static_cast(i) % mono.size()]; + } + return fixed; +} + +std::vector extract_mel(const runtime::AudioBuffer & audio, size_t threads) { + const auto waveform = prepare_reference(audio); + const engine::audio::STFTConfig stft{ + 1024, 320, 640, true, + engine::audio::STFTPadMode::Reflect, + // torch.hann_window defaults to periodic=true in upstream MiraTTS. + engine::audio::STFTFamily::Kokoro}; + const auto & window = engine::audio::get_cached_stft_window(stft); + const auto magnitude = engine::audio::STFT().compute_magnitude( + waveform, window, 1, kReferenceSamples, stft, threads); + const int64_t frames = magnitude.shape.at(2); + auto mel = engine::audio::MelFilterbank().compute( + magnitude.values, 1, 513, frames, + engine::audio::MelFilterbankConfig{ + 16000, 1024, 128, 10.0F, 8000.0F, true}); + // AudioTensor is [B, mel, frames], which is already the graph's BCT layout. + return std::move(mel.values); +} + +} // namespace + +struct MiraSpeakerEncoder::Impl { + Impl( + const MiraTTSAssets & assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + size_t graph_context_bytes, + assets::TensorStorageType linear_type, + assets::TensorStorageType conv_type) + : execution(execution), + weights(load_weights( + assets, execution, weight_context_bytes, linear_type, conv_type)), + graph_context_bytes(graph_context_bytes) {} + + ~Impl() { + if (gallocr != nullptr) ggml_gallocr_free(gallocr); + } + + void ensure_graph(int64_t frames) { + if (ctx != nullptr && frames == graph_frames) return; + if (gallocr != nullptr) { + ggml_gallocr_free(gallocr); + gallocr = nullptr; + } + ctx.reset(); + ggml_init_params params{graph_context_bytes, nullptr, true}; + ctx.reset(ggml_init(params)); + if (!ctx) throw std::runtime_error("MiraTTS failed to initialize speaker graph"); + core::ModuleBuildContext build_ctx{ + ctx.get(), "mira_tts.speaker_encoder", execution.backend_type()}; + auto input = core::make_tensor( + build_ctx, GGML_TYPE_F32, + core::TensorShape::from_dims({1, 128, frames})); + input_tensor = input.tensor; + output_tensor = build_graph(build_ctx, input, weights).tensor; + ggml_set_output(output_tensor); + graph = ggml_new_graph_custom(ctx.get(), 65536, false); + ggml_build_forward_expand(graph, output_tensor); + gallocr = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(execution.backend())); + if (gallocr == nullptr || !ggml_gallocr_reserve(gallocr, graph) || + !ggml_gallocr_alloc_graph(gallocr, graph)) { + throw std::runtime_error("MiraTTS failed to allocate speaker graph"); + } + graph_frames = frames; + } + + std::vector encode(const runtime::AudioBuffer & audio) { + auto mel = extract_mel(audio, static_cast( + std::max(1, execution.config().threads))); + if (mel.size() % 128 != 0) { + throw std::runtime_error("MiraTTS mel feature shape mismatch"); + } + const int64_t frames = static_cast(mel.size() / 128); + ensure_graph(frames); + ggml_backend_tensor_set( + input_tensor, mel.data(), 0, mel.size() * sizeof(float)); + if (ggml_backend_graph_compute(execution.backend(), graph) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("MiraTTS speaker graph execution failed"); + } + std::vector projected(32 * 6); + ggml_backend_tensor_get( + output_tensor, projected.data(), 0, + projected.size() * sizeof(float)); + std::vector codes(32, 0); + for (int row = 0; row < 32; ++row) { + int32_t code = 0; + int32_t radix = 1; + for (int dim = 0; dim < 6; ++dim) { + const float value = projected[static_cast(row * 6 + dim)]; + const float bounded = std::tanh(value + 0.3461989760398865F) * + 1.501500129699707F - 0.5F; + const int32_t digit = static_cast(std::nearbyint(bounded)) + 2; + code += std::clamp(digit, 0, 3) * radix; + radix *= 4; + } + codes[static_cast(row)] = code; + } + return codes; + } + + core::ExecutionContext & execution; + SpeakerWeights weights; + size_t graph_context_bytes = 0; + int64_t graph_frames = 0; + std::unique_ptr ctx; + ggml_cgraph * graph = nullptr; + ggml_gallocr_t gallocr = nullptr; + ggml_tensor * input_tensor = nullptr; + ggml_tensor * output_tensor = nullptr; +}; + +MiraSpeakerEncoder::MiraSpeakerEncoder( + const MiraTTSAssets & assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + size_t graph_context_bytes, + assets::TensorStorageType linear_storage_type, + assets::TensorStorageType conv_storage_type) + : impl_(std::make_unique( + assets, execution, weight_context_bytes, graph_context_bytes, + linear_storage_type, conv_storage_type)) {} + +MiraSpeakerEncoder::~MiraSpeakerEncoder() = default; + +std::vector MiraSpeakerEncoder::encode( + const runtime::AudioBuffer & reference_audio) { + return impl_->encode(reference_audio); +} + +} // namespace engine::community_models::mira_tts diff --git a/src/framework/audio/flashsr.cpp b/src/framework/audio/flashsr.cpp index 7c959691d..27babaae1 100644 --- a/src/framework/audio/flashsr.cpp +++ b/src/framework/audio/flashsr.cpp @@ -420,7 +420,17 @@ FlashSrModel FlashSrModel::load_from_directory(const std::filesystem::path & mod FlashSrModel FlashSrModel::load_from_directory( const std::filesystem::path & model_dir, const core::BackendConfig & backend_config) { - auto source = engine::assets::open_tensor_source(model_dir / "flashsr.safetensors"); + return load_from_tensor_source( + engine::assets::open_tensor_source(model_dir / "flashsr.safetensors"), + backend_config); +} + +FlashSrModel FlashSrModel::load_from_tensor_source( + std::shared_ptr source, + const core::BackendConfig & backend_config) { + if (!source) { + throw std::runtime_error("FlashSR tensor source is missing"); + } auto weights = std::make_shared(); weights->backend.reset(core::init_backend(backend_config)); weights->backend_type = core::backend_type(weights->backend.get()); diff --git a/src/framework/model_spec/schema.cpp b/src/framework/model_spec/schema.cpp index fbc49aa9e..1fd22d5ca 100644 --- a/src/framework/model_spec/schema.cpp +++ b/src/framework/model_spec/schema.cpp @@ -546,10 +546,14 @@ ValidatedPackage validate_package(const json::Value & value, std::string_view pa std::unordered_set validate_packages( const json::Value & value, std::string_view path, - bool has_default_download) { + bool has_default_download, + bool allow_empty) { const auto & packages = require_spec_array(value, path); if (packages.empty()) { - fail(path, "packages must not be empty"); + if (allow_empty) { + return {}; + } + fail(path, "packages must not be empty unless status is experimental"); } bool has_default = false; std::unordered_set package_ids; @@ -651,10 +655,14 @@ void validate_dependencies( void validate_ui(const json::Value & value, const std::unordered_set & package_ids, std::string_view path) { require_spec_object(value, path); - const auto recommended = require_spec_string(require_spec_field(value, "recommended_package", path), - std::string(path) + ".recommended_package"); - if (package_ids.find(recommended) == package_ids.end()) { - fail(std::string(path) + ".recommended_package", "unknown package '" + recommended + "'"); + if (const auto * recommended_value = value.find("recommended_package")) { + const auto recommended = require_spec_string( + *recommended_value, std::string(path) + ".recommended_package"); + if (package_ids.find(recommended) == package_ids.end()) { + fail(std::string(path) + ".recommended_package", "unknown package '" + recommended + "'"); + } + } else if (!package_ids.empty()) { + fail(std::string(path) + ".recommended_package", "missing required field"); } if (const auto * min_vram = value.find("min_vram_gb")) { require_spec_number(*min_vram, std::string(path) + ".min_vram_gb"); @@ -675,8 +683,9 @@ void validate_v1(const json::Value & spec, std::string_view source_name) { (void) require_spec_string(require_spec_field(spec, "display_name", source_name), std::string(source_name) + ".display_name"); validate_enum(require_spec_string(require_spec_field(spec, "category", source_name), std::string(source_name) + ".category"), categories(), std::string(source_name) + ".category", "category"); - validate_enum(require_spec_string(require_spec_field(spec, "status", source_name), std::string(source_name) + ".status"), - statuses(), std::string(source_name) + ".status", "status"); + const auto status = require_spec_string( + require_spec_field(spec, "status", source_name), std::string(source_name) + ".status"); + validate_enum(status, statuses(), std::string(source_name) + ".status", "status"); const auto task_ids = validate_nonempty_string_set( require_spec_field(spec, "tasks", source_name), &tasks(), std::string(source_name) + ".tasks", "task"); validate_nonempty_string_set( @@ -698,7 +707,8 @@ void validate_v1(const json::Value & spec, std::string_view source_name) { const auto packages_path = std::string(source_name) + ".packages"; const auto & packages_field = require_spec_field(spec, "packages", source_name); - const auto package_ids = validate_packages(packages_field, packages_path, has_default_download); + const auto package_ids = validate_packages( + packages_field, packages_path, has_default_download, status == "experimental"); validate_dependencies( require_spec_field(spec, "dependencies", source_name), family, diff --git a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp index 8d0e65ad9..6bd4c0c54 100644 --- a/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp +++ b/src/framework/modules/transformers/qwen_causal_decode_runtime.cpp @@ -291,10 +291,12 @@ ggml_tensor * make_logits_readback_token_ids( if (config.logits_readback_token_ids.empty()) { return nullptr; } - return ggml_new_tensor_1d( + auto * tensor = ggml_new_tensor_1d( ctx, GGML_TYPE_I32, static_cast(config.logits_readback_token_ids.size())); + ggml_set_input(tensor); + return tensor; } core::TensorValue wrap_logits_readback_token_ids( @@ -638,7 +640,10 @@ class QwenCausalDecodeRuntime::Impl { }; void ensure_prefill_token_graph(int64_t steps) { - if (prefill_graph_ != nullptr && prefill_input_kind_ == InputKind::Token && prefill_steps_ == steps) { + if (prefill_graph_ != nullptr && + prefill_input_kind_ == InputKind::Token && + prefill_steps_ == steps && + config_.logits_readback_token_ids.empty()) { debug::timing_log_scalar(config_.trace_name + ".prefill.graph.build_ms", 0.0); debug::trace_log_scalar(config_.trace_name + ".prefill.steps", steps); return; @@ -648,7 +653,10 @@ class QwenCausalDecodeRuntime::Impl { } void ensure_prefill_embedding_graph(int64_t steps) { - if (prefill_graph_ != nullptr && prefill_input_kind_ == InputKind::Embedding && prefill_steps_ == steps) { + if (prefill_graph_ != nullptr && + prefill_input_kind_ == InputKind::Embedding && + prefill_steps_ == steps && + config_.logits_readback_token_ids.empty()) { debug::timing_log_scalar(config_.trace_name + ".prefill.graph.build_ms", 0.0); debug::trace_log_scalar(config_.trace_name + ".prefill.steps", steps); return; @@ -819,7 +827,8 @@ class QwenCausalDecodeRuntime::Impl { void ensure_batched_prefill_token_graph(int64_t batch_size, int64_t steps) { if (batched_prefill_graph_ != nullptr && batched_prefill_input_kind_ == InputKind::Token && - batched_prefill_batch_size_ == batch_size && batched_prefill_steps_ == steps) { + batched_prefill_batch_size_ == batch_size && batched_prefill_steps_ == steps && + config_.logits_readback_token_ids.empty()) { debug::timing_log_scalar(config_.trace_name + ".batched_prefill.graph.build_ms", 0.0); return; } @@ -829,7 +838,8 @@ class QwenCausalDecodeRuntime::Impl { void ensure_batched_prefill_embedding_graph(int64_t batch_size, int64_t steps) { if (batched_prefill_graph_ != nullptr && batched_prefill_input_kind_ == InputKind::Embedding && - batched_prefill_batch_size_ == batch_size && batched_prefill_steps_ == steps) { + batched_prefill_batch_size_ == batch_size && batched_prefill_steps_ == steps && + config_.logits_readback_token_ids.empty()) { debug::timing_log_scalar(config_.trace_name + ".batched_prefill.graph.build_ms", 0.0); return; } diff --git a/src/framework/sampling/hf_sampler.cpp b/src/framework/sampling/hf_sampler.cpp index f6cb8cf65..fbbebb6cd 100644 --- a/src/framework/sampling/hf_sampler.cpp +++ b/src/framework/sampling/hf_sampler.cpp @@ -327,6 +327,57 @@ void HfLogitsProcessor::apply_top_p( scratch.probabilities_scores_size_ = scores.size(); } +void HfLogitsProcessor::apply_min_p( + std::vector & scores, + float min_p, + int64_t min_tokens_to_keep, + HfSamplerScratch & scratch) { + require_min_tokens(min_tokens_to_keep); + if (min_p < 0.0F || min_p > 1.0F || !std::isfinite(min_p)) { + throw std::runtime_error("HF sampler min_p must be finite and in [0, 1]"); + } + if (min_p == 0.0F) { + return; + } + if (scores.empty()) { + return; + } + auto & order = scratch.candidates_; + order.clear(); + order.reserve(scores.size()); + for (size_t index = 0; index < scores.size(); ++index) { + if (std::isfinite(scores[index])) { + order.push_back(static_cast(index)); + } + } + if (order.empty()) { + return; + } + const size_t min_keep = std::min( + static_cast(min_tokens_to_keep), order.size()); + std::partial_sort( + order.begin(), + order.begin() + static_cast(min_keep), + order.end(), + [&](int32_t lhs, int32_t rhs) { + const float lhs_score = scores[static_cast(lhs)]; + const float rhs_score = scores[static_cast(rhs)]; + return lhs_score == rhs_score ? lhs < rhs : lhs_score > rhs_score; + }); + const float max_score = scores[static_cast(order.front())]; + const float threshold = max_score + std::log(min_p); + const float protected_threshold = + scores[static_cast(order[min_keep - 1])]; + for (float & score : scores) { + if (score < threshold && score < protected_threshold) { + score = -std::numeric_limits::infinity(); + } + } + scratch.probabilities_ready_ = false; + scratch.probabilities_scores_data_ = nullptr; + scratch.probabilities_scores_size_ = 0; +} + void HfLogitsProcessor::apply_temperature(std::vector & scores, float temperature) { if (!(temperature > 0.0F) || !std::isfinite(temperature)) { throw std::runtime_error("HF sampler temperature must be finite and positive"); @@ -434,7 +485,8 @@ int32_t HfSampler::sample( const bool needs_repetition_penalty = options.repetition_penalty != 1.0F && !history.empty(); const bool needs_sampling_processors = options.do_sample && - (options.temperature != 1.0F || options.top_k > 0 || options.top_p < 1.0F); + (options.temperature != 1.0F || options.top_k > 0 || + options.top_p < 1.0F || options.min_p > 0.0F); if (!needs_repetition_penalty && !needs_sampling_processors) { if (!options.do_sample) { return HfLogitsProcessor::argmax(logits.data(), logits.size(), context); @@ -452,6 +504,7 @@ int32_t HfSampler::sample( HfLogitsProcessor::apply_temperature(scores, options.temperature); HfLogitsProcessor::apply_top_k(scores, options.top_k, options.min_tokens_to_keep, scratch); HfLogitsProcessor::apply_top_p(scores, options.top_p, options.min_tokens_to_keep, scratch); + HfLogitsProcessor::apply_min_p(scores, options.min_p, options.min_tokens_to_keep, scratch); return HfTokenSampler::sample_from_processed_scores(scores, scratch, fallback_rng, torch_state, context, true); } diff --git a/tests/mira_tts/README.md b/tests/mira_tts/README.md new file mode 100644 index 000000000..3d1b33b54 --- /dev/null +++ b/tests/mira_tts/README.md @@ -0,0 +1,113 @@ +# MiraTTS validation + +This directory follows the long-lived validation pattern used for the OuteTTS +port in PR #63. The model is loaded once and exercised by a sequence containing +cold, repeat, changed-prompt, long-form, and post-long-form repeat requests. +Streaming has a separate long-lived sequence so that offline and streaming +sessions do not duplicate model weights in VRAM. + +The benchmark reports one JSON object per request with wall time, generated +duration, real-time factor, sample/frame counts, a deterministic audio hash, +and output path. Streaming runs additionally report first-event latency and the +number of progressively emitted audio events. + +## Build + +```powershell +cmake -S . -B build/windows-cuda-release ` + -DGGML_CUDA=ON -DENGINE_BUILD_WARMBENCH=ON ` + -DCMAKE_BUILD_TYPE=Release +cmake --build build/windows-cuda-release --config Release ` + --target mira_tts_warm_bench -j 8 +``` + +## Native long-lived runs + +Set paths once: + +```powershell +$bench = "build/windows-cuda-release/bin/mira_tts_warm_bench.exe" +$model = "../models_v3_test/MiraTTS-GGUF/mira-tts-bf16.gguf" +$voice = "../models_v3_test/MiraTTS-comparison/00-reference-voice.wav" +$spec = "model_specs/mira_tts.json" +$out = "build/logs/warmbench/mira_tts" +``` + +Run the five-case offline sequence: + +```powershell +& $bench --model $model --model-spec-override $spec --backend cuda ` + --voice-ref $voice --request-file tests/mira_tts/offline_requests.json ` + --audio-out-dir "$out/native-offline" ` + --summary-file "$out/native-offline.json" ` + --log-file "$out/native-offline.log" +``` + +Run the repeated streaming sequence: + +```powershell +& $bench --model $model --model-spec-override $spec --backend cuda ` + --run-mode streaming --voice-ref $voice ` + --request-file tests/mira_tts/streaming_requests.json ` + --audio-out-dir "$out/native-streaming" ` + --summary-file "$out/native-streaming.json" ` + --log-file "$out/native-streaming.log" +``` + +Validate deterministic repeats, streaming events, durations, hashes, and WAV +readability (with FFprobe when available): + +```powershell +python tests/mira_tts/validate_bench.py ` + --offline-summary "$out/native-offline.json" ` + --streaming-summary "$out/native-streaming.json" +``` + +For a memory measurement, start the benchmark with `--hold-seconds 30` and +sample the process plus the selected GPU while it is holding the loaded +session. Record both peak host RSS and peak device memory; do not report only +the final idle value. + +## Trusted Python sequence + +Run the same offline cases through a local, revision-pinned upstream snapshot: + +```powershell +python tools/community_models/mira_tts_reference_bench.py ` + --model-dir --reference $voice ` + --request-file tests/mira_tts/offline_requests.json ` + --output-dir "$out/python-offline" ` + --summary-file "$out/python-offline.json" +``` + +Compare every matched request. Exact frame count is required for the strongest +deterministic parity run: + +```powershell +python tools/community_models/compare_mira_tts_outputs.py ` + --cpp-summary "$out/native-offline.json" ` + --python-summary "$out/python-offline.json" ` + --output "$out/parity.json" --require-exact-frames +``` + +The default gates are waveform cosine >= 0.95 and log-mel cosine >= 0.95. +Threshold changes must be justified by a saved artifact and must not hide a +frame-count, token-boundary, or sampling mismatch. + +## Acceptance matrix + +| Check | Required evidence | +|---|---| +| Cold/repeat determinism | `clone_cold`, `clone_repeat`, and `clone_repeat_after_longform` have identical hashes | +| Long-lived lifecycle | All offline requests finish in one process without reloading the model | +| Long form | `longform` produces non-empty audio within its token budget | +| Streaming | More than one event, finite first-event latency, merged non-empty WAV | +| Streaming determinism | `stream_cold` and `stream_repeat` have identical hashes and event counts | +| Python/native parity | Per-case WAV cosine, log-mel cosine, and frame counts are recorded | +| Audio validity | Every WAV is readable by FFmpeg and has the reported sample rate/channels | +| Resource use | Peak RSS, peak GPU memory, wall time, duration, and RTF are recorded | +| Backends | CUDA is required; CPU/Vulkan results or explicit limitations are documented | + +Run deterministic checks before experimenting with unseeded sampling. If a +repeat hash changes, inspect generated speech-token boundaries, prompt/context +tokens, and component traces before comparing subjective audio quality. diff --git a/tests/mira_tts/VALIDATION.md b/tests/mira_tts/VALIDATION.md new file mode 100644 index 000000000..437e30120 --- /dev/null +++ b/tests/mira_tts/VALIDATION.md @@ -0,0 +1,123 @@ +# MiraTTS local validation record + +Date: 2026-09-02 + +Model: `mira-tts-q8.gguf` + +Reference: `00-reference-voice.wav` + +Host: Windows, NVIDIA GeForce RTX 3090 24 GB + +This record is intentionally local and is not a claim that untested devices or +backends have parity. Generated WAVs and machine-readable summaries are under +`build/logs/warmbench/mira_tts/` and are not source-controlled. + +## CUDA offline — long-lived session + +One model and one session handled all five requests. + +| Request | Wall ms | Audio s | RTF | FNV-1a audio hash | +|---|---:|---:|---:|---| +| clone_cold | 1707.41 | 5.12 | 0.3335 | `426924e4695337c6` | +| clone_repeat | 1560.66 | 5.12 | 0.3048 | `426924e4695337c6` | +| short_second_prompt | 1566.03 | 5.12 | 0.3059 | `8255f4342afeb014` | +| longform | 4771.24 | 15.36 | 0.3106 | `5708fcccf7c25080` | +| clone_repeat_after_longform | 1583.74 | 5.12 | 0.3093 | `426924e4695337c6` | + +The three identical requests remain bit-deterministic before and after the +different-prompt and long-form requests. Internal traces show the decode graph +reused after its first build for matching capacity; larger/different prompt +capacities can build a separate prefill graph. + +## Upstream Python performance comparison + +The revision-pinned upstream MiraTTS implementation was run on the same RTX +3090 with the same prompts, reference voice, sampling parameters, token budgets, +and request order. Model initialization is excluded from both request loops. + +| Test | Original MiraTTS | audio.cpp Q8 | Difference | +|---|---:|---:|---| +| Cold request, 5.12 s audio | 2.227 s | 1.707 s | audio.cpp 1.30× faster* | +| Repeated request, 5.12 s | 0.911 s | 1.561 s | Original 1.71× faster | +| Second short prompt, 5.12 s | 0.895 s | 1.566 s | Original 1.75× faster | +| Long-form, 15.36 s | 2.976 s | 4.771 s | Original 1.60× faster | +| Repeat after long-form, 5.12 s | 0.918 s | 1.584 s | Original 1.73× faster | + +\* The original cold request paid a one-time ONNX Runtime CUDA fallback cost; +the warm measurements are the representative steady-state comparison. + +The three warm short requests average 908.02 ms upstream and 1570.14 ms in +native Q8, making upstream about 1.73x faster in this measurement. For the +15.36-second long-form case, upstream achieved RTF 0.1938 (5.16x real time) +and native Q8 achieved RTF 0.3106 (3.22x real time). + +This is not a precision-matched comparison: upstream uses BF16 while the native +run uses GGUF Q8. Upstream also encodes and retains the reference context before +the timed request sequence, whereas the native request currently processes its +reference input on each run. + +## CUDA streaming + +| Request | Wall ms | First event ms | Events | Audio s | RTF | FNV-1a audio hash | +|---|---:|---:|---:|---:|---:|---| +| stream_cold | 6639.50 | 2489.93 | 4 | 21.14 | 0.3141 | `26eb3542a073a8db` | +| stream_repeat | 6498.52 | 2368.89 | 4 | 21.14 | 0.3074 | `26eb3542a073a8db` | + +The stream emits multiple independently consumable audio events and the merged +repeat is bit-deterministic. + +## Resident memory — Q8 CUDA + +With the Q8 model held after the five-request sequence: + +| Measurement | Resident usage | +|---|---:| +| Process working set | 1429.54 MiB | +| Process private bytes | 5958.28 MiB | +| Windows GPU Process Memory dedicated usage | 2551.55 MiB | +| `nvidia-smi` total device memory in use | 2552 MiB | + +These are held-resident samples, not an instrumented peak across model loading. +The README therefore still requires peak sampling for a formal performance +submission. + +## Python/native decoder parity + +The previously saved exact-token comparison uses identical upstream speech and +context tokens: + +| Metric | Python | Native | Result | +|---|---:|---:|---:| +| Sample rate | 48000 Hz | 48000 Hz | Match | +| Frame count | 122880 | 122880 | Exact match | +| Waveform cosine | — | — | 0.99996735 | +| Log-mel cosine | — | — | 0.99943239 | + +This isolates the native processor/decoder path from autoregressive sampling. + +## Backend coverage + +| Backend | Build | Runtime result | Status | +|---|---|---|---| +| CUDA, RTX 3090 | Passed | Offline, streaming, determinism, and WAV validation passed | Validated | +| Vulkan, RTX 3090 | Passed | Deterministic, but under-generated (3.44 s vs 5.12 s) and diverged on other prompts | Not parity-clean | +| Vulkan, AMD integrated GPU | Passed | Deterministic, but generated only 0.10 s for the 5.12 s CUDA case | Not parity-clean | +| CPU | Not run | No runtime measurement recorded | Untested | + +`mira_tts_warm_bench` builds successfully with the Vulkan configuration. It is +not parity-clean in this test: + +- AMD integrated Vulkan device: deterministic repeats, but only 0.10 seconds + for the 5.12-second CUDA case. +- RTX 3090 Vulkan device: deterministic repeats, but 3.44 seconds for the same + case and severe under-generation on the other prompts. + +Vulkan is therefore recorded as compile-tested but unsupported for MiraTTS +output parity until the backend divergence is diagnosed. CUDA is the validated +runtime for this model. + +## Automated result + +`tests/mira_tts/validate_bench.py` passed all offline/streaming lifecycle, +determinism, duration, event-count, and WAV-readability checks for the CUDA +summaries. diff --git a/tests/mira_tts/mira_tts_warm_bench.cpp b/tests/mira_tts/mira_tts_warm_bench.cpp new file mode 100644 index 000000000..195b7285e --- /dev/null +++ b/tests/mira_tts/mira_tts_warm_bench.cpp @@ -0,0 +1,372 @@ +#include "engine/framework/audio/wav_reader.h" +#include "engine/framework/audio/wav_writer.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/io/json.h" +#include "engine/framework/runtime/registry.h" +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +struct RequestCase { + std::string name; + std::string text; + std::string language; + std::filesystem::path voice_ref; + std::unordered_map options; +}; + +std::string arg_value(int argc, char **argv, const std::string &name, + const std::string &fallback) { + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) + return argv[i + 1]; + } + return fallback; +} + +std::vector arg_values(int argc, char **argv, + const std::string &name) { + std::vector out; + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) + out.emplace_back(argv[i + 1]); + } + return out; +} + +int int_arg(int argc, char **argv, const std::string &name, int fallback) { + return std::stoi(arg_value(argc, argv, name, std::to_string(fallback))); +} + +engine::core::BackendType parse_backend(const std::string &value) { + if (value == "cpu") + return engine::core::BackendType::Cpu; + if (value == "cuda") + return engine::core::BackendType::Cuda; + if (value == "vulkan") + return engine::core::BackendType::Vulkan; + if (value == "best") + return engine::core::BackendType::BestAvailable; + throw std::runtime_error("unsupported backend: " + value); +} + +std::string scalar_option(const engine::io::json::Value &value) { + if (value.is_string()) + return value.as_string(); + if (value.is_bool()) + return value.as_bool() ? "true" : "false"; + if (value.is_number()) + return engine::io::json::stringify_number(value.as_number()); + throw std::runtime_error( + "MiraTTS warm-bench options must be scalar values"); +} + +void copy_option_if_present( + std::unordered_map &options, + const engine::io::json::Value &item, const std::string &name) { + if (const auto *value = item.find(name); + value != nullptr && !value->is_null()) { + options[name] = scalar_option(*value); + } +} + +std::vector load_requests( + const std::filesystem::path &path, + const std::filesystem::path &default_voice_ref, + const std::unordered_map &defaults) { + const auto root = engine::io::json::parse_file(path); + const auto &items = root.require("requests").as_array(); + if (items.empty()) + throw std::runtime_error("MiraTTS request file has no requests"); + + std::vector out; + out.reserve(items.size()); + for (size_t index = 0; index < items.size(); ++index) { + const auto &item = items[index]; + RequestCase request; + request.name = engine::io::json::optional_string( + item, "name", "request_" + std::to_string(index)); + request.text = engine::io::json::require_string(item, "text"); + request.language = + engine::io::json::optional_string(item, "language", "en"); + request.voice_ref = engine::io::json::optional_string( + item, "voice_ref", default_voice_ref.string()); + request.options = defaults; + for (const char *name : + {"max_tokens", "seed", "temperature", "top_k", "top_p", + "min_p", "repetition_penalty", "text_chunk_size", + "text_chunk_mode"}) { + copy_option_if_present(request.options, item, name); + } + if (request.voice_ref.empty()) + throw std::runtime_error("MiraTTS request '" + request.name + + "' has no reference voice"); + out.push_back(std::move(request)); + } + return out; +} + +engine::runtime::AudioBuffer read_audio(const std::filesystem::path &path) { + const auto wav = engine::audio::read_wav_f32(path); + return {wav.sample_rate, wav.channels, wav.samples}; +} + +double audio_seconds(const engine::runtime::AudioBuffer &audio) { + if (audio.sample_rate <= 0 || audio.channels <= 0) + return 0.0; + return static_cast(audio.samples.size()) / + static_cast(audio.sample_rate * audio.channels); +} + +std::string fnv1a64_hex(const engine::runtime::AudioBuffer &audio) { + uint64_t hash = 1469598103934665603ULL; + const auto mix = [&hash](const void *data, size_t size) { + const auto *bytes = static_cast(data); + for (size_t i = 0; i < size; ++i) { + hash ^= static_cast(bytes[i]); + hash *= 1099511628211ULL; + } + }; + mix(&audio.sample_rate, sizeof(audio.sample_rate)); + mix(&audio.channels, sizeof(audio.channels)); + for (float sample : audio.samples) { + uint32_t bits = 0; + std::memcpy(&bits, &sample, sizeof(bits)); + mix(&bits, sizeof(bits)); + } + std::ostringstream out; + out << std::hex << std::setfill('0') << std::setw(16) << hash; + return out.str(); +} + +engine::runtime::TaskRequest make_request( + const RequestCase &request, + std::unordered_map + &audio_cache) { + engine::runtime::TaskRequest out; + out.text_input = + engine::runtime::Transcript{request.text, request.language}; + out.options = request.options; + const std::string key = request.voice_ref.lexically_normal().string(); + auto found = audio_cache.find(key); + if (found == audio_cache.end()) + found = audio_cache.emplace(key, read_audio(request.voice_ref)).first; + out.voice = engine::runtime::VoiceCondition{}; + out.voice->speaker = engine::runtime::VoiceReference{}; + out.voice->speaker->audio = found->second; + return out; +} + +engine::io::json::Value result_json( + const RequestCase &request, int iteration, const std::string &mode, + const engine::runtime::AudioBuffer &audio, double wall_ms, + double first_event_ms, int event_count, + const std::filesystem::path &output_path) { + const double seconds = audio_seconds(audio); + return engine::io::json::Value::make_object({ + {"name", engine::io::json::Value::make_string(request.name)}, + {"iteration", engine::io::json::Value::make_number(iteration)}, + {"mode", engine::io::json::Value::make_string(mode)}, + {"wall_ms", engine::io::json::Value::make_number(wall_ms)}, + {"audio_seconds", engine::io::json::Value::make_number(seconds)}, + {"rtf", engine::io::json::Value::make_number( + seconds > 0.0 ? wall_ms / 1000.0 / seconds : 0.0)}, + {"first_event_ms", + first_event_ms >= 0.0 + ? engine::io::json::Value::make_number(first_event_ms) + : engine::io::json::Value::make_null()}, + {"event_count", engine::io::json::Value::make_number(event_count)}, + {"sample_rate", + engine::io::json::Value::make_number(audio.sample_rate)}, + {"channels", engine::io::json::Value::make_number(audio.channels)}, + {"samples", + engine::io::json::Value::make_number( + static_cast(audio.samples.size()))}, + {"audio_hash", engine::io::json::Value::make_string(fnv1a64_hex(audio))}, + {"audio_out", + engine::io::json::Value::make_string(output_path.string())}, + }); +} + +} // namespace + +int main(int argc, char **argv) try { + const std::filesystem::path model_path = + arg_value(argc, argv, "--model", "models/MiraTTS/model.gguf"); + const std::filesystem::path request_file = + arg_value(argc, argv, "--request-file", ""); + if (request_file.empty()) + throw std::runtime_error("MiraTTS warm bench requires --request-file"); + const std::filesystem::path default_voice_ref = + arg_value(argc, argv, "--voice-ref", ""); + const std::filesystem::path output_dir = + arg_value(argc, argv, "--audio-out-dir", + "build/logs/warmbench/mira_tts_audio"); + const std::filesystem::path log_file = + arg_value(argc, argv, "--log-file", + "build/logs/warmbench/mira_tts.log"); + const std::filesystem::path summary_file = + arg_value(argc, argv, "--summary-file", + "build/logs/warmbench/mira_tts_summary.json"); + const std::filesystem::path spec_override = + arg_value(argc, argv, "--model-spec-override", ""); + const std::string backend_name = arg_value(argc, argv, "--backend", "cuda"); + const std::string mode = arg_value(argc, argv, "--run-mode", "offline"); + const int device = int_arg(argc, argv, "--device", 0); + const int threads = int_arg(argc, argv, "--threads", 8); + const int iterations = int_arg(argc, argv, "--iterations", 1); + const int hold_seconds = int_arg(argc, argv, "--hold-seconds", 0); + if (mode != "offline" && mode != "streaming") + throw std::runtime_error("--run-mode must be offline or streaming"); + if (iterations <= 0) + throw std::runtime_error("--iterations must be positive"); + if (hold_seconds < 0) + throw std::runtime_error("--hold-seconds must be non-negative"); + + std::unordered_map defaults; + for (const auto &option : arg_values(argc, argv, "--request-option")) { + const size_t equals = option.find('='); + if (equals == std::string::npos || equals == 0) + throw std::runtime_error("invalid --request-option: " + option); + defaults[option.substr(0, equals)] = option.substr(equals + 1); + } + const auto requests = + load_requests(request_file, default_voice_ref, defaults); + + std::filesystem::create_directories(output_dir); + if (!log_file.parent_path().empty()) + std::filesystem::create_directories(log_file.parent_path()); + if (!summary_file.parent_path().empty()) + std::filesystem::create_directories(summary_file.parent_path()); + engine::debug::configure_logging( + engine::debug::LoggingConfig{true, log_file.string()}); + + auto registry = engine::runtime::make_default_registry(); + engine::runtime::ModelLoadRequest load_request; + load_request.model_path = model_path; + load_request.family_hint = "mira_tts"; + if (!spec_override.empty()) + load_request.model_spec_override = spec_override; + auto model = registry.load(load_request); + + engine::runtime::SessionOptions session_options; + session_options.backend.type = parse_backend(backend_name); + session_options.backend.device = device; + session_options.backend.threads = threads; + for (const auto &option : arg_values(argc, argv, "--session-option")) { + const size_t equals = option.find('='); + if (equals == std::string::npos || equals == 0) + throw std::runtime_error("invalid --session-option: " + option); + session_options.options[option.substr(0, equals)] = + option.substr(equals + 1); + } + + const auto run_mode = mode == "streaming" + ? engine::runtime::RunMode::Streaming + : engine::runtime::RunMode::Offline; + auto session_base = model->create_task_session( + {engine::runtime::VoiceTaskKind::VoiceCloning, run_mode}, + session_options); + auto *offline = dynamic_cast( + session_base.get()); + auto *streaming = dynamic_cast( + session_base.get()); + if (mode == "offline" && offline == nullptr) + throw std::runtime_error("MiraTTS did not create an offline session"); + if (mode == "streaming" && streaming == nullptr) + throw std::runtime_error("MiraTTS did not create a streaming session"); + + std::unordered_map audio_cache; + auto first_request = make_request(requests.front(), audio_cache); + session_base->prepare( + engine::runtime::build_preparation_request(first_request)); + + engine::io::json::Value::Array results; + for (const auto &request_case : requests) { + for (int iteration = 1; iteration <= iterations; ++iteration) { + auto request = make_request(request_case, audio_cache); + const auto started = std::chrono::steady_clock::now(); + engine::runtime::TaskResult result; + double first_event_ms = -1.0; + int event_count = 0; + if (mode == "streaming") { + streaming->start_stream(request); + while (auto event = streaming->next_stream_event()) { + ++event_count; + if (first_event_ms < 0.0) { + first_event_ms = + std::chrono::duration( + std::chrono::steady_clock::now() - started) + .count(); + } + } + result = streaming->finish_stream(); + } else { + result = offline->run(request); + } + const double wall_ms = + std::chrono::duration( + std::chrono::steady_clock::now() - started) + .count(); + if (!result.audio_output.has_value()) + throw std::runtime_error("MiraTTS produced no audio"); + const auto output_path = + output_dir / + (request_case.name + "_" + mode + "_" + + std::to_string(iteration) + ".wav"); + engine::audio::write_pcm16_wav( + output_path, result.audio_output->sample_rate, + result.audio_output->channels, result.audio_output->samples); + auto item = result_json(request_case, iteration, mode, + *result.audio_output, wall_ms, first_event_ms, + event_count, output_path); + std::cout << "result_json=" << engine::io::json::stringify(item) + << "\n"; + results.push_back(std::move(item)); + } + } + + auto summary = engine::io::json::Value::make_object({ + {"family", engine::io::json::Value::make_string("mira_tts")}, + {"backend", engine::io::json::Value::make_string(backend_name)}, + {"mode", engine::io::json::Value::make_string(mode)}, + {"model", engine::io::json::Value::make_string(model_path.string())}, + {"results", + engine::io::json::Value::make_array(std::move(results))}, + }); + { + std::ofstream output(summary_file, std::ios::binary | std::ios::trunc); + if (!output) + throw std::runtime_error("failed to open summary file: " + + summary_file.string()); + output << engine::io::json::stringify(summary) << "\n"; + } + std::cout << "summary_json=" << summary_file.string() << "\n"; + std::cout << "log_out=" << log_file.string() << "\n"; + if (hold_seconds > 0) { + std::cout << "holding_session_seconds=" << hold_seconds << "\n"; + std::cout.flush(); + std::this_thread::sleep_for(std::chrono::seconds(hold_seconds)); + } + engine::debug::reset_logging(); + return 0; +} catch (const std::exception &error) { + std::cerr << "mira_tts_warm_bench failed: " << error.what() << "\n"; + return 1; +} diff --git a/tests/mira_tts/offline_requests.json b/tests/mira_tts/offline_requests.json new file mode 100644 index 000000000..1c4ccce94 --- /dev/null +++ b/tests/mira_tts/offline_requests.json @@ -0,0 +1,34 @@ +{ + "requests": [ + { + "name": "clone_cold", + "text": "Good morning. This deterministic request checks the first synthesis in a newly prepared session.", + "seed": 1234, + "max_tokens": 256 + }, + { + "name": "clone_repeat", + "text": "Good morning. This deterministic request checks the first synthesis in a newly prepared session.", + "seed": 1234, + "max_tokens": 256 + }, + { + "name": "short_second_prompt", + "text": "The same loaded model now speaks a different sentence without restarting the process.", + "seed": 4321, + "max_tokens": 256 + }, + { + "name": "longform", + "text": "MiraTTS remains loaded while this longer paragraph exercises its autoregressive generator, acoustic processor, low resolution decoder, and audio upsampler. The benchmark records wall time, generated duration, real time factor, and a deterministic audio hash. A second sentence ensures that punctuation and a longer token budget are covered by the same warm session.", + "seed": 2468, + "max_tokens": 768 + }, + { + "name": "clone_repeat_after_longform", + "text": "Good morning. This deterministic request checks the first synthesis in a newly prepared session.", + "seed": 1234, + "max_tokens": 256 + } + ] +} diff --git a/tests/mira_tts/streaming_requests.json b/tests/mira_tts/streaming_requests.json new file mode 100644 index 000000000..8a87214f6 --- /dev/null +++ b/tests/mira_tts/streaming_requests.json @@ -0,0 +1,20 @@ +{ + "requests": [ + { + "name": "stream_cold", + "text": "The first streaming request emits this sentence in several progressive audio segments. The listener can begin playback before the complete paragraph has finished synthesizing.", + "seed": 1234, + "max_tokens": 384, + "text_chunk_size": 72, + "text_chunk_mode": "default" + }, + { + "name": "stream_repeat", + "text": "The first streaming request emits this sentence in several progressive audio segments. The listener can begin playback before the complete paragraph has finished synthesizing.", + "seed": 1234, + "max_tokens": 384, + "text_chunk_size": 72, + "text_chunk_mode": "default" + } + ] +} diff --git a/tests/mira_tts/validate_bench.py b/tests/mira_tts/validate_bench.py new file mode 100644 index 000000000..353eb5d69 --- /dev/null +++ b/tests/mira_tts/validate_bench.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python3 +"""Validate deterministic MiraTTS warm-benchmark artifacts.""" + +from __future__ import annotations + +import argparse +import json +import shutil +import subprocess +from pathlib import Path + + +def load_results(path: Path) -> dict[str, dict[str, object]]: + payload = json.loads(path.read_text(encoding="utf-8")) + return {item["name"]: item for item in payload["results"]} + + +def require(condition: bool, message: str, failures: list[str]) -> None: + if not condition: + failures.append(message) + + +def validate_audio(item: dict[str, object], failures: list[str]) -> None: + path = Path(str(item["audio_out"])) + require(path.is_file(), f"missing WAV: {path}", failures) + require(int(item["sample_rate"]) == 48000, f"unexpected sample rate: {path}", failures) + require(int(item["channels"]) == 1, f"unexpected channel count: {path}", failures) + require(int(item["samples"]) > 0, f"empty audio: {path}", failures) + require(float(item["audio_seconds"]) > 0.0, f"zero duration: {path}", failures) + require(float(item["rtf"]) > 0.0, f"invalid RTF: {path}", failures) + if path.is_file() and shutil.which("ffprobe"): + process = subprocess.run( + [ + "ffprobe", + "-v", + "error", + "-select_streams", + "a:0", + "-show_entries", + "stream=sample_rate,channels", + "-of", + "json", + str(path), + ], + capture_output=True, + text=True, + check=False, + ) + require(process.returncode == 0, f"ffprobe rejected WAV: {path}", failures) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--offline-summary", type=Path, required=True) + parser.add_argument("--streaming-summary", type=Path, required=True) + args = parser.parse_args() + + offline = load_results(args.offline_summary) + streaming = load_results(args.streaming_summary) + failures: list[str] = [] + + expected_offline = { + "clone_cold", + "clone_repeat", + "short_second_prompt", + "longform", + "clone_repeat_after_longform", + } + expected_streaming = {"stream_cold", "stream_repeat"} + require(expected_offline <= set(offline), "offline cases are incomplete", failures) + require(expected_streaming <= set(streaming), "streaming cases are incomplete", failures) + + for item in [*offline.values(), *streaming.values()]: + validate_audio(item, failures) + + if expected_offline <= set(offline): + repeat_hashes = { + str(offline[name]["audio_hash"]) + for name in ( + "clone_cold", + "clone_repeat", + "clone_repeat_after_longform", + ) + } + require(len(repeat_hashes) == 1, "offline deterministic hashes differ", failures) + require( + offline["short_second_prompt"]["audio_hash"] + != offline["clone_cold"]["audio_hash"], + "different prompts unexpectedly produced the same audio hash", + failures, + ) + require( + int(offline["longform"]["samples"]) + > int(offline["clone_cold"]["samples"]), + "long-form case is not longer than the short case", + failures, + ) + + if expected_streaming <= set(streaming): + require( + streaming["stream_cold"]["audio_hash"] + == streaming["stream_repeat"]["audio_hash"], + "streaming deterministic hashes differ", + failures, + ) + require( + int(streaming["stream_cold"]["event_count"]) > 1, + "streaming cold case emitted fewer than two events", + failures, + ) + require( + streaming["stream_cold"]["event_count"] + == streaming["stream_repeat"]["event_count"], + "streaming repeat event counts differ", + failures, + ) + for name in expected_streaming: + require( + float(streaming[name]["first_event_ms"]) > 0.0, + f"{name} has no valid first-event latency", + failures, + ) + + report = { + "passed": not failures, + "failures": failures, + "offline_cases": sorted(offline), + "streaming_cases": sorted(streaming), + } + print(json.dumps(report, indent=2)) + if failures: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tests/unittests/test_hf_sampler.cpp b/tests/unittests/test_hf_sampler.cpp index 421ea47e1..5aa3b783e 100644 --- a/tests/unittests/test_hf_sampler.cpp +++ b/tests/unittests/test_hf_sampler.cpp @@ -283,6 +283,28 @@ void test_matches_python_hf_processor_reference_values() { } } +void test_min_p_masks_relative_probability_tail() { + std::vector scores{4.0F, 2.0F, 1.0F, -1.0F}; + HfSamplerScratch scratch; + engine::sampling::HfLogitsProcessor::apply_min_p(scores, 0.1F, 1, scratch); + + engine::test::require(std::isfinite(scores[0]), "min-p keeps maximum token"); + engine::test::require(std::isfinite(scores[1]), "min-p keeps token above relative threshold"); + engine::test::require( + std::isinf(scores[2]) && scores[2] < 0.0F, + "min-p masks token below relative threshold"); + engine::test::require( + std::isinf(scores[3]) && scores[3] < 0.0F, + "min-p masks probability tail"); + + std::vector protected_scores{4.0F, 0.0F, -1.0F}; + engine::sampling::HfLogitsProcessor::apply_min_p( + protected_scores, 0.9F, 2, scratch); + engine::test::require( + std::isfinite(protected_scores[0]) && std::isfinite(protected_scores[1]), + "min-p honors min_tokens_to_keep"); +} + void test_matches_python_hf_cuda_multinomial_reference_sequence() { TorchCudaSamplingPolicy policy; policy.cuda_fast_path = true; @@ -506,6 +528,7 @@ int main() { try { test_greedy_fast_path_matches_reference(); test_matches_python_hf_processor_reference_values(); + test_min_p_masks_relative_probability_tail(); test_matches_python_hf_cuda_multinomial_reference_sequence(); test_matches_python_hf_cpu_multinomial_reference_sequence(); test_no_processor_sampling_fast_path_matches_reference(); diff --git a/tests/unittests/test_model_spec_system.cpp b/tests/unittests/test_model_spec_system.cpp index 6c9230e23..ace7fd19e 100644 --- a/tests/unittests/test_model_spec_system.cpp +++ b/tests/unittests/test_model_spec_system.cpp @@ -1223,6 +1223,43 @@ void test_legacy_spec_contract_behavior_unchanged() { std::filesystem::remove_all(root); } +void test_experimental_spec_without_installable_package() { + const std::string experimental = R"JSON({ + "schema_version": 1, + "family": "local_only_model", + "display_name": "Local Only Model", + "description": "Requires a local checkpoint conversion.", + "category": "tts", + "status": "experimental", + "tasks": ["tts"], + "modes": ["offline"], + "languages": ["en"], + "runtime": {"tags": ["gguf"]}, + "capabilities": {}, + "options": {"request": [], "session": [], "load": []}, + "packages": [], + "dependencies": [], + "ui": {"tags": ["TTS"], "docs": ["docs/local.md"]}, + "sources": [{ + "format": "safetensors", + "roots": {"model": "."}, + "files": {"config": "model:config.json"}, + "tensors": {"weights": "model:model.safetensors"} + }] + })JSON"; + engine::model_spec::validate_spec( + json::parse(experimental), "experimental_local_only"); + + auto community = experimental; + const auto status = community.find("\"status\": \"experimental\""); + engine::test::require(status != std::string::npos, "experimental status fixture"); + community.replace(status, std::string("\"status\": \"experimental\"").size(), + "\"status\": \"community\""); + expect_rejects( + "community_requires_package", community, + "packages must not be empty unless status is experimental"); +} + void test_contract_spec_prefers_workspace_over_package_local_spec() { const auto root = make_temp_root(); const auto workspace = root / "workspace"; @@ -1291,6 +1328,7 @@ int main() { test_schema_v1_metadata_projection(); test_contract_projection_ignores_package_metadata_validation(); test_legacy_spec_contract_behavior_unchanged(); + test_experimental_spec_without_installable_package(); test_contract_spec_prefers_workspace_over_package_local_spec(); test_loading_and_resource_bundle(); } catch (const std::exception & error) { diff --git a/tools/community_models/compare_mira_tts_outputs.py b/tools/community_models/compare_mira_tts_outputs.py new file mode 100644 index 000000000..866d8115c --- /dev/null +++ b/tools/community_models/compare_mira_tts_outputs.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python3 +"""Compare matched upstream/native MiraTTS WAV outputs.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import librosa +import numpy as np +import soundfile as sf + + +def mono(path: Path) -> tuple[np.ndarray, int]: + audio, sample_rate = sf.read(path, dtype="float32", always_2d=True) + return np.mean(audio, axis=1, dtype=np.float32), sample_rate + + +def cosine(left: np.ndarray, right: np.ndarray) -> float: + denominator = float(np.linalg.norm(left) * np.linalg.norm(right)) + return 1.0 if denominator == 0.0 else float(np.dot(left, right) / denominator) + + +def compare(cpp_path: Path, python_path: Path) -> dict[str, object]: + cpp, cpp_rate = mono(cpp_path) + python, python_rate = mono(python_path) + if cpp_rate != python_rate: + raise RuntimeError( + f"sample-rate mismatch for {cpp_path.name}: {cpp_rate} != {python_rate}" + ) + common = min(cpp.size, python.size) + cpp_common = cpp[:common].astype(np.float64, copy=False) + python_common = python[:common].astype(np.float64, copy=False) + wav_cosine = cosine(cpp_common, python_common) + mel_kwargs = { + "sr": cpp_rate, + "n_fft": 2048, + "hop_length": 512, + "win_length": 2048, + "n_mels": 128, + "power": 2.0, + } + cpp_mel = np.log( + np.maximum(librosa.feature.melspectrogram(y=cpp, **mel_kwargs), 1.0e-10) + ) + python_mel = np.log( + np.maximum(librosa.feature.melspectrogram(y=python, **mel_kwargs), 1.0e-10) + ) + mel_frames = min(cpp_mel.shape[1], python_mel.shape[1]) + log_mel_cosine = cosine( + cpp_mel[:, :mel_frames].reshape(-1).astype(np.float64, copy=False), + python_mel[:, :mel_frames].reshape(-1).astype(np.float64, copy=False), + ) + return { + "cpp_audio": str(cpp_path), + "python_audio": str(python_path), + "sample_rate": cpp_rate, + "cpp_frames": int(cpp.size), + "python_frames": int(python.size), + "exact_frame_count": cpp.size == python.size, + "common_frames": int(common), + "waveform_cosine": wav_cosine, + "log_mel_cosine": log_mel_cosine, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--cpp-summary", type=Path, required=True) + parser.add_argument("--python-summary", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--wav-cosine-min", type=float, default=0.95) + parser.add_argument("--log-mel-cosine-min", type=float, default=0.95) + parser.add_argument("--require-exact-frames", action="store_true") + args = parser.parse_args() + + cpp = json.loads(args.cpp_summary.read_text(encoding="utf-8")) + python = json.loads(args.python_summary.read_text(encoding="utf-8")) + cpp_by_name = {item["name"]: item for item in cpp["results"]} + python_by_name = {item["name"]: item for item in python["results"]} + names = sorted(set(cpp_by_name) & set(python_by_name)) + if not names: + raise RuntimeError("the summaries contain no matching request names") + + comparisons = [] + failed = [] + for name in names: + item = compare( + Path(cpp_by_name[name]["audio_out"]), + Path(python_by_name[name]["audio_out"]), + ) + item["name"] = name + item["passed"] = ( + item["waveform_cosine"] >= args.wav_cosine_min + and item["log_mel_cosine"] >= args.log_mel_cosine_min + and (not args.require_exact_frames or item["exact_frame_count"]) + ) + comparisons.append(item) + if not item["passed"]: + failed.append(name) + print(json.dumps(item)) + + report = { + "thresholds": { + "waveform_cosine": args.wav_cosine_min, + "log_mel_cosine": args.log_mel_cosine_min, + "require_exact_frames": args.require_exact_frames, + }, + "passed": not failed, + "failed": failed, + "comparisons": comparisons, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2), encoding="utf-8") + if failed: + raise SystemExit(f"MiraTTS parity failed: {', '.join(failed)}") + + +if __name__ == "__main__": + main() diff --git a/tools/community_models/convert_mira_tts.py b/tools/community_models/convert_mira_tts.py new file mode 100644 index 000000000..5368b78a8 --- /dev/null +++ b/tools/community_models/convert_mira_tts.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python3 +"""Convert the official MiraTTS checkpoint into audio.cpp tensor assets. + +The upstream package combines a Qwen2 safetensors checkpoint, two ONNX +graphs, a DAC decoder safetensors checkpoint, and a small PyTorch 48 kHz +upsampler. audio.cpp does not execute ONNX or pickle at runtime: this tool +extracts and gives stable names to all learned tensors, fuses upsampler weight +normalization, and writes one ordinary safetensors file per runtime namespace. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +from pathlib import Path + +import numpy as np +import onnx +from onnx import numpy_helper +import torch +from safetensors.torch import load_file, save_file + + +SPEAKER_ANONYMOUS = { + "onnx::MatMul_874": "perceiver.proj_context.weight", + "onnx::Expand_875": "perceiver.latents", + "onnx::MatMul_878": "perceiver.layers.0.attn.q.weight", + "onnx::MatMul_879": "perceiver.layers.0.attn.kv.weight", + "onnx::MatMul_883": "perceiver.layers.0.attn.out.weight", + "onnx::MatMul_884": "perceiver.layers.0.ff.in.weight", + "onnx::MatMul_885": "perceiver.layers.0.ff.out.weight", + "onnx::MatMul_886": "perceiver.layers.1.attn.q.weight", + "onnx::MatMul_887": "perceiver.layers.1.attn.kv.weight", + "onnx::MatMul_891": "perceiver.layers.1.attn.out.weight", + "onnx::MatMul_892": "perceiver.layers.1.ff.in.weight", + "onnx::MatMul_893": "perceiver.layers.1.ff.out.weight", + "onnx::MatMul_896": "quantizer.project_in.weight", +} + + +def processor_anonymous() -> dict[str, str]: + names = { + "onnx::MatMul_978": "speaker_encoder.quantizer.project_out.weight", + "onnx::MatMul_980": "prenet.linear_pre.weight", + "onnx::MatMul_1013": "prenet.linear.weight", + } + index = 981 + for downsample in range(2): + for block in range(2): + base = f"prenet.downsample.{downsample}.1.convnext.{block}" + names[f"onnx::MatMul_{index}"] = f"{base}.pwconv1.weight" + names[f"onnx::MatMul_{index + 1}"] = f"{base}.pwconv2.weight" + index += 2 + for block in range(12): + base = f"prenet.vocos_backbone.convnext.{block}" + names[f"onnx::MatMul_{index}"] = f"{base}.pwconv1.weight" + names[f"onnx::MatMul_{index + 1}"] = f"{base}.pwconv2.weight" + index += 2 + if index != 1013: + raise AssertionError(index) + return names + + +def onnx_tensors(path: Path, namespace: str, anonymous: dict[str, str]) -> dict[str, torch.Tensor]: + model = onnx.load(path, load_external_data=True) + output: dict[str, torch.Tensor] = {} + found_anonymous: set[str] = set() + for tensor in model.graph.initializer: + name = anonymous.get(tensor.name, tensor.name) + if tensor.name.startswith("onnx::MatMul_"): + if tensor.name not in anonymous: + raise SystemExit(f"unmapped learned ONNX tensor: {tensor.name}") + if tensor.name in anonymous: + found_anonymous.add(tensor.name) + value = np.asarray(numpy_helper.to_array(tensor)) + if value.ndim == 0: + value = value.reshape(1) + # ONNX MatMul parameters are [in, out]; audio.cpp Linear parameters + # follow PyTorch's [out, in] convention. + if tensor.name.startswith("onnx::MatMul_"): + value = value.T + output[f"{namespace}.{name}"] = torch.from_numpy(np.array(value, copy=True, order="C")) + missing = set(anonymous) - found_anonymous + if missing: + raise SystemExit(f"ONNX graph is missing expected tensors: {sorted(missing)}") + return output + + +def processor_context_codebook(path: Path) -> np.ndarray: + """Recover the exported FSQ 4096x6 implicit codebook constant.""" + model = onnx.load(path, load_external_data=True) + candidates: list[np.ndarray] = [] + for node in model.graph.node: + if node.op_type != "Constant": + continue + for attribute in node.attribute: + if attribute.type != onnx.AttributeProto.TENSOR: + continue + value = np.asarray(numpy_helper.to_array(attribute.t)) + if value.shape == (4096, 6): + candidates.append(value) + if len(candidates) != 1: + raise SystemExit( + f"expected one processor FSQ [4096, 6] codebook constant, found {len(candidates)}" + ) + return np.ascontiguousarray(candidates[0].astype(np.float32)) + + +def safetensor_tensors(path: Path, namespace: str) -> dict[str, torch.Tensor]: + return { + f"{namespace}.{name}": value.contiguous() + for name, value in load_file(str(path), device="cpu").items() + } + + +def upsampler_tensors(path: Path) -> dict[str, torch.Tensor]: + checkpoint = torch.load(path, map_location="cpu", weights_only=True) + state = checkpoint.get("model", checkpoint) + tensors = { + name: value.detach().cpu().float().numpy() + for name, value in state.items() + if hasattr(value, "detach") and not name.startswith("optimizer.") + } + fused_weights: dict[str, np.ndarray] = {} + for name, value in tensors.items(): + if name.endswith(".weight_g"): + base = name[:-len("_g")] + direction = tensors[base + "_v"].astype(np.float32) + axes = tuple(range(1, direction.ndim)) + norm = np.sqrt(np.sum(direction * direction, axis=axes, keepdims=True)) + fused_weights[base] = value.astype(np.float32) * direction / np.maximum(norm, 1.0e-12) + + # Use the framework FlashSR tensor contract. MiraTTS runs only residual + # blocks 2 and 0, matching FastAudioSR.Generator.forward(). + output: dict[str, torch.Tensor] = { + "upsampler.conv_pre.weight": torch.from_numpy(np.ascontiguousarray(fused_weights["dec.conv_pre.weight"])), + "upsampler.conv_pre.bias": torch.from_numpy(np.ascontiguousarray(tensors["dec.conv_pre.bias"])), + "upsampler.conv_post.weight": torch.from_numpy(np.ascontiguousarray(tensors["dec.conv_post.weight"])), + "upsampler.activation_filter": torch.from_numpy(np.ascontiguousarray( + tensors["dec.activation_post.upsample.filter"])), + } + for block in ("0", "2"): + for group in (1, 2): + for index in range(3): + source = f"dec.resblocks.{block}.convs{group}.{index}" + target = f"upsampler.resblocks.{block}.convs{group}.{index}" + output[target + ".weight"] = torch.from_numpy( + np.ascontiguousarray(fused_weights[source + ".weight"])) + output[target + ".bias"] = torch.from_numpy( + np.ascontiguousarray(tensors[source + ".bias"])) + for index in range(6): + source = f"dec.resblocks.{block}.activations.{index}.act" + target = f"upsampler.resblocks.{block}.activations.{index}" + output[target + ".alpha"] = torch.from_numpy( + np.ascontiguousarray(np.exp(tensors[source + ".alpha"])).reshape(1, 32, 1)) + output[target + ".inv_beta"] = torch.from_numpy( + np.ascontiguousarray(1.0 / (np.exp(tensors[source + ".beta"]) + 1.0e-9)).reshape(1, 32, 1)) + output["upsampler.activation_post.alpha"] = torch.from_numpy( + np.ascontiguousarray(np.exp(tensors["dec.activation_post.act.alpha"])).reshape(1, 32, 1)) + output["upsampler.activation_post.inv_beta"] = torch.from_numpy( + np.ascontiguousarray(1.0 / (np.exp(tensors["dec.activation_post.act.beta"]) + 1.0e-9)).reshape(1, 32, 1)) + return output + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("input", type=Path, help="downloaded YatharthS/MiraTTS directory") + parser.add_argument("output", type=Path, help="output audio.cpp model directory") + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + root = args.input.resolve() + output = args.output.resolve() + namespaces = ("language_model", "speaker_encoder", "processor", "decoder", "upsampler") + destinations = [output / f"{namespace}.safetensors" for namespace in namespaces] + existing = [path for path in destinations if path.exists()] + if existing and not args.overwrite: + raise SystemExit(f"output exists (pass --overwrite): {existing[0]}") + + required = { + "lm": root / "model.safetensors", + "speaker": root / "decoders" / "s_encoder.onnx", + "processor": root / "decoders" / "processer.onnx", + "decoder": root / "decoders" / "detokenizer.safetensors", + "upsampler": root / "decoders" / "upsampler.pth", + "config": root / "config.json", + "tokenizer": root / "tokenizer.json", + "tokenizer_config": root / "tokenizer_config.json", + } + missing = [str(path) for path in required.values() if not path.is_file()] + if missing: + raise SystemExit("missing MiraTTS files: " + ", ".join(missing)) + + tensors: dict[str, torch.Tensor] = {} + tensors.update(safetensor_tensors(required["lm"], "language_model")) + tensors.update(onnx_tensors(required["speaker"], "speaker_encoder", SPEAKER_ANONYMOUS)) + tensors.update(onnx_tensors(required["processor"], "processor", processor_anonymous())) + tensors["processor.speaker_encoder.context_codebook"] = torch.from_numpy( + processor_context_codebook(required["processor"]) + ) + tensors.update(safetensor_tensors(required["decoder"], "decoder")) + tensors.update(upsampler_tensors(required["upsampler"])) + + output.mkdir(parents=True, exist_ok=True) + total = 0 + for namespace in namespaces: + prefix = namespace + "." + scoped = { + name[len(prefix):]: value + for name, value in tensors.items() + if name.startswith(prefix) + } + if not scoped: + raise SystemExit(f"no tensors collected for namespace {namespace}") + save_file( + scoped, + str(output / f"{namespace}.safetensors"), + metadata={ + "format": "pt", + "source": "YatharthS/MiraTTS", + "audiocpp_family": "mira_tts", + "audiocpp_namespace": namespace, + }, + ) + total += len(scoped) + for name in ("config.json", "tokenizer.json", "tokenizer_config.json"): + shutil.copy2(root / name, output / name) + print(f"wrote {len(namespaces)} tensor namespaces ({total} tensors) to {output}") + + +if __name__ == "__main__": + main() diff --git a/tools/community_models/mira_tts_reference_bench.py b/tools/community_models/mira_tts_reference_bench.py new file mode 100644 index 000000000..125ea4297 --- /dev/null +++ b/tools/community_models/mira_tts_reference_bench.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +"""Run a long-lived upstream MiraTTS request sequence for parity testing.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import re +import time +from pathlib import Path + +import huggingface_hub +import numpy as np +import scipy.io.wavfile +import torch + + +def audio_sha256(waveform: np.ndarray) -> str: + pcm = np.clip(waveform, -1.0, 1.0).astype(np.float32, copy=False) + return hashlib.sha256(pcm.tobytes()).hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--model-dir", type=Path, required=True) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--request-file", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--summary-file", type=Path, required=True) + args = parser.parse_args() + + model_dir = args.model_dir.resolve() + huggingface_hub.snapshot_download = ( + lambda *unused_args, **unused_kwargs: str(model_dir) + ) + + # The official checkpoint contains a revision-pinned legacy PyTorch file. + import transformers.modeling_utils + + transformers.modeling_utils.check_torch_load_is_safe = lambda: None + + from lmdeploy import GenerationConfig, TurbomindEngineConfig, pipeline + from ncodec.codec import TTSCodec + + requests = json.loads(args.request_file.read_text(encoding="utf-8"))["requests"] + if not requests: + raise RuntimeError("MiraTTS reference request file is empty") + args.output_dir.mkdir(parents=True, exist_ok=True) + args.summary_file.parent.mkdir(parents=True, exist_ok=True) + + codec = TTSCodec() + context_tokens = codec.encode(str(args.reference.resolve()), encode_semantic=False) + context_ids = np.asarray( + [[[int(token) for token in re.findall(r"context_token_(\d+)", context_tokens)]]], + dtype=np.int32, + ) + backend = TurbomindEngineConfig( + cache_max_entry_count=0.2, + tp=1, + dtype="bfloat16", + enable_prefix_caching=False, + ) + pipe = pipeline(str(model_dir), backend_config=backend) + decoder = codec.audio_decoder + + results: list[dict[str, object]] = [] + for index, request in enumerate(requests): + name = request.get("name", f"request_{index}") + prompt = codec.format_prompt(request["text"], context_tokens, None) + generation = GenerationConfig( + top_p=float(request.get("top_p", 0.95)), + top_k=int(request.get("top_k", 50)), + temperature=float(request.get("temperature", 0.8)), + max_new_tokens=int(request.get("max_tokens", 1024)), + repetition_penalty=float(request.get("repetition_penalty", 1.2)), + min_p=float(request.get("min_p", 0.05)), + do_sample=True, + random_seed=int(request.get("seed", 1234)), + ) + started = time.perf_counter() + response = pipe([prompt], gen_config=generation, do_preprocess=False)[0] + speech_tokens = response.text + speech_ids = np.asarray( + [[int(token) for token in re.findall(r"speech_token_(\d+)", speech_tokens)]], + dtype=np.int64, + ) + latent = decoder.processor_detokenizer.run( + ["preprocessed_output"], + {"context_tokens": context_ids, "speech_tokens": speech_ids}, + )[0] + audio = codec.decode(speech_tokens, context_tokens) + waveform = audio.detach().float().cpu().numpy().reshape(-1) + wall_ms = (time.perf_counter() - started) * 1000.0 + output = args.output_dir / f"{name}.wav" + scipy.io.wavfile.write(output, 48000, np.clip(waveform, -1.0, 1.0)) + results.append( + { + "name": name, + "wall_ms": wall_ms, + "audio_seconds": waveform.size / 48000.0, + "rtf": wall_ms / 1000.0 / (waveform.size / 48000.0), + "sample_rate": 48000, + "samples": int(waveform.size), + "speech_token_count": int(speech_ids.size), + "audio_sha256_f32": audio_sha256(waveform), + "audio_out": str(output), + "processor_latent": { + "shape": list(latent.shape), + "sum": float(np.sum(latent, dtype=np.float64)), + "sum_sq": float(np.sum(np.square(latent), dtype=np.float64)), + }, + } + ) + print(json.dumps(results[-1])) + + summary = { + "family": "mira_tts", + "implementation": "upstream_python", + "model_dir": str(model_dir), + "reference": str(args.reference.resolve()), + "results": results, + } + args.summary_file.write_text(json.dumps(summary, indent=2), encoding="utf-8") + print(args.summary_file) + + +if __name__ == "__main__": + main() diff --git a/tools/community_models/mira_tts_reference_smoke.py b/tools/community_models/mira_tts_reference_smoke.py new file mode 100644 index 000000000..8263a98fd --- /dev/null +++ b/tools/community_models/mira_tts_reference_smoke.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +"""Run the trusted upstream MiraTTS implementation for native parity checks.""" + +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path + +import huggingface_hub +import numpy as np +import scipy.io.wavfile +import torch + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--model-dir", type=Path, required=True) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--text", required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--max-new-tokens", type=int, default=1024) + parser.add_argument("--top-k", type=int, default=50) + parser.add_argument("--top-p", type=float, default=0.95) + parser.add_argument("--min-p", type=float, default=0.05) + parser.add_argument("--temperature", type=float, default=0.8) + parser.add_argument("--repetition-penalty", type=float, default=1.2) + parser.add_argument("--seed", type=int, default=1234) + args = parser.parse_args() + + model_dir = args.model_dir.resolve() + # TTSCodec hardcodes snapshot_download; keep the upstream implementation intact + # while making it use the already verified local snapshot. + huggingface_hub.snapshot_download = lambda *unused_args, **unused_kwargs: str(model_dir) + + # The official checkpoint uses a legacy PyTorch file for FlashSR. The pinned + # Windows Torch in this isolated environment predates Transformers' new guard. + # Only the official, revision-pinned checkpoint is admitted here. + import transformers.modeling_utils + + transformers.modeling_utils.check_torch_load_is_safe = lambda: None + + from lmdeploy import GenerationConfig, TurbomindEngineConfig, pipeline + from ncodec.codec import TTSCodec + + codec = TTSCodec() + context_tokens = codec.encode(str(args.reference), encode_semantic=False) + prompt = codec.format_prompt(args.text, context_tokens, None) + + backend = TurbomindEngineConfig( + cache_max_entry_count=0.2, + tp=1, + dtype="bfloat16", + enable_prefix_caching=False, + ) + pipe = pipeline(str(model_dir), backend_config=backend) + generation = GenerationConfig( + top_p=args.top_p, + top_k=args.top_k, + temperature=args.temperature, + max_new_tokens=args.max_new_tokens, + repetition_penalty=args.repetition_penalty, + min_p=args.min_p, + do_sample=True, + random_seed=args.seed, + ) + response = pipe([prompt], gen_config=generation, do_preprocess=False)[0] + speech_tokens = response.text + speech_ids = np.asarray( + [[int(token) for token in re.findall(r"speech_token_(\d+)", speech_tokens)]], + dtype=np.int64, + ) + context_ids = np.asarray( + [[[int(token) for token in re.findall(r"context_token_(\d+)", context_tokens)]]], + dtype=np.int32, + ) + decoder = codec.audio_decoder + latent = decoder.processor_detokenizer.run( + ["preprocessed_output"], + {"context_tokens": context_ids, "speech_tokens": speech_ids}, + )[0] + lowres = decoder.audio_detokenizer.decode( + torch.from_numpy(latent).to("cuda:0") + ).squeeze().detach().float().cpu().numpy() + audio = codec.decode(speech_tokens, context_tokens) + waveform = audio.detach().float().cpu().numpy().reshape(-1) + + args.output.parent.mkdir(parents=True, exist_ok=True) + scipy.io.wavfile.write(args.output, 48000, np.clip(waveform, -1.0, 1.0)) + scipy.io.wavfile.write( + args.output.with_name(args.output.stem + "-lowres.wav"), + 16000, + np.clip(lowres, -1.0, 1.0), + ) + args.output.with_suffix(".json").write_text( + json.dumps( + { + "model_dir": str(model_dir), + "reference": str(args.reference.resolve()), + "text": args.text, + "context_tokens": context_tokens, + "speech_tokens": speech_tokens, + "sample_rate": 48000, + "samples": int(waveform.size), + "rms": float(np.sqrt(np.mean(np.square(waveform)))), + "peak": float(np.max(np.abs(waveform))), + "processor_latent": { + "shape": list(latent.shape), + "sum": float(np.sum(latent, dtype=np.float64)), + "sum_sq": float(np.sum(np.square(latent), dtype=np.float64)), + "first": latent.reshape(-1)[:10].astype(float).tolist(), + }, + }, + indent=2, + ), + encoding="utf-8", + ) + print(args.output) + + +if __name__ == "__main__": + main()