diff --git a/CMakeLists.txt b/CMakeLists.txt index b994f87c..51b5218d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -541,6 +541,23 @@ audiocpp_add_model(inflect_v2 engine::models::inflect_v2::make_inflect_v2_loader ) +audiocpp_add_model(f5_tts + SOURCES + src/community_models/f5_tts/session.cpp + src/community_models/f5_tts/runtime.cpp + src/community_models/f5_tts/weights.cpp + src/community_models/f5_tts/dit_modules.cpp + src/community_models/f5_tts/synthesize.cpp + INCLUDES + engine/community_models/f5_tts/session.h + engine/community_models/f5_tts/runtime.h + LOADERS + engine::models::f5_tts::make_f5_tts_loader + ALIASES + habibi + habibi_tts +) + audiocpp_add_model(minimax_h3 SOURCES src/community_models/minimax_h3/assets.cpp @@ -1502,6 +1519,24 @@ if (vibevoice IN_LIST AUDIOCPP_LINKED_MODELS) endif() endif() +# F5/Habibi tests: parity harnesses + e2e sample generator. Only when the +# f5_tts model is linked (they call model-internal symbols) AND test hooks +# are explicitly enabled — F5_MEL_TEST must never leak into production +# model libraries (PR #275 review feedback). +option(ENGINE_F5_TEST_HOOKS "Build F5-TTS test hooks (f5_test_* symbols) and parity binaries" OFF) +if (f5_tts IN_LIST AUDIOCPP_LINKED_MODELS AND ENGINE_F5_TEST_HOOKS) + target_compile_definitions(engine_model_f5_tts PRIVATE F5_MEL_TEST=1) + foreach(f5_test IN ITEMS f5_e2e f5_parity f5_cfg_parity f5_tokenizer) + add_executable(${f5_test} tests/${f5_test}_main.cpp) + target_compile_definitions(${f5_test} PRIVATE F5_MEL_TEST=1) + target_link_libraries(${f5_test} PRIVATE engine_runtime ggml) + target_include_directories(${f5_test} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) + if (ENGINE_ENABLE_OPENMP) + target_link_libraries(${f5_test} PRIVATE OpenMP::OpenMP_CXX) + endif() + endforeach() +endif() + if (ENGINE_BUILD_TESTS) add_executable(moss_tts_local_smoke tests/moss_tts_local/moss_tts_local_smoke.cpp diff --git a/docs/community_models/f5_tts.md b/docs/community_models/f5_tts.md new file mode 100644 index 00000000..b49fd418 --- /dev/null +++ b/docs/community_models/f5_tts.md @@ -0,0 +1,105 @@ +# F5-TTS (community model) + +[F5-TTS](https://github.com/SWivid/F5-TTS) is an open-source zero-shot voice-cloning TTS built on a +flow-matching diffusion transformer (DiT) with a ConvNeXt text conditioner and a Vocos vocoder. +This community port also targets the [Habibi-TTS](https://github.com/SWivid/Habibi-TTS) finetune — +a multi-dialect Arabic checkpoint suite (MSA, SAU, UAE, ALG, IRQ, EGY, MAR, OMN, TUN, LEV, SDN, LBY) +from the same authors — which uses the identical architecture, giving Arabic support through the +same family (`habibi` / `habibi_tts` are registered as aliases). + +**Status: M4 — inference wired end to end.** The session (`src/community_models/f5_tts/session.cpp`) +runs full synthesis via `f5_synthesize` (text pipeline with Habibi dialect tokens → batched-CFG DiT +Euler sampler → Vocos vocoder) and is served by `audiocpp_server`. Parity vs the reference PyTorch +implementation is covered by golden harnesses (DiT stage taps, batched CFG incl. the null branch, +tokenizer ids) plus a whisper.cpp ASR pronunciation check; see `/mnt/ai/f5-parity/run_all.sh`. + +## Milestones + +Each milestone is gated on parity against the reference PyTorch implementation (cosine similarity +≥ 0.999 on fixed inputs) plus a listening check, matching the evidence bar described in #54 and +PR #180. + +| Milestone | Scope | Status | +|---|---|---| +| M0 | Family registration, model spec, stub session, this doc | done | +| M1 | Weight loading + mel-Vocos decode path (ConvNeXt + iSTFT) | done (mel-corr 0.9963) | +| M2 | DiT forward (RoPE, adaLN) + ConvNeXt text conditioner | done (all stages cosine 1.000000) | +| M3 | CFM sampler (Euler, sway + EPSS, CFG null-branch parity), inference wiring, En/Ar samples | done | +| M4 | Long-form chunking, RTF/VRAM evidence, server wiring | done (0.51x RTF on RTX 3090; GGUF packages hosted at trklou/audio.cpp) | + +## Server usage + +`audiocpp_server.json` entry: family `f5_tts`, model path = checkpoint directory (a `*.gguf` +package, or a safetensors checkpoint + `vocab.txt` for development). Session options: +`f5_tts.vocos_path` (only needed with safetensors checkpoints lacking a bundled vocoder), +`f5_tts.dialect` (default UNK), `f5_tts.frame_budget` (mel frames per CFM pass, 0 = 2048). +Requests take `reference_text` (required), `dialect`, `speed`, `seed`, `num_inference_steps`, +`cfg_strength` (alias `guidance_scale`), `sway_sampling_coef`, `strip_diacritics`. + +**Diacritics (harakat):** Habibi was trained on ASR transcripts, which are undiacritized — +the harakat/tanwin/shadda tokens exist in the vocab but are severely undertrained, and raw +diacritized input degrades to garbled speech with character repetitions (identical in the +Python reference; not a port issue). By default the frontend strips combining marks +(U+0640, U+064B–U+065F, U+0670) before synthesis, so `أَيْنَ اللَّوْنُ الأَحْمَر؟` reads as +`أين اللون الأحمر؟`. Disable per request with `strip_diacritics=false`. + +## Quickstart (from a fresh clone) + +Everything needed is installable from this repository — no Python inference stack required: + +```bash +# 1. Build (f5_tts is included in AUDIOCPP_MODEL_SET=full, or select it explicitly) +cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DAUDIOCPP_MODEL_SET=custom \ + -DAUDIOCPP_MODELS=f5_tts -DENGINE_ENABLE_CUDA=ON +cmake --build build --parallel --target audiocpp_cli audiocpp_server + +# 2. Download the model package (GGUF: DiT + Vocos vocoder in one file + vocab.txt) +python3 tools/model_manager_v2.py install habibi_unified +# Per-dialect specialized checkpoints (stronger accent): habibi_alg, habibi_egy, +# habibi_irq, habibi_mar, habibi_msa, habibi_sau, habibi_uae +# Standalone vocoder for the original safetensors checkpoints: vocos_mel_24khz + +# 3. Synthesize (zero-shot: any short reference WAV + its transcript) +build/bin/audiocpp_cli --task tts --family habibi \ + --model /Habibi-TTS/Unified \ + --voice-ref /path/to/reference.wav \ + --reference-text 'transcript of the reference audio' \ + --text 'أهلا، هذا نص عربي تجريبي.' \ + --request-option dialect=UNK --out out.wav +``` + +The GGUF package is self-contained: the DiT lives under the `transformer` namespace and the +Vocos vocoder under `vocos`, so no separate vocoder download or `f5_tts.vocos_path` option is +needed (the option and safetensors fallbacks still work for the original HF checkpoints). +Dialects: `UNK MSA SAU UAE ALG IRQ EGY MAR OMN TUN LEV SDN LBY`. +Reference audio longer than ~10.9s is truncated (with a warning) — keep refs under that +and make sure the transcript matches, otherwise the transcript tail leaks into the output. + +### Converting checkpoints to GGUF + +`tools/convert_f5_tts.py` wraps `audiocpp_gguf` and produces one self-contained GGUF per +checkpoint (`transformer` + `vocos` namespaces) plus the standalone vocoder package: + +```bash +python3 tools/convert_f5_tts.py --checkpoint-root /models/Habibi-TTS \ + --vocos /models/vocos-mel-24khz/vocos.safetensors \ + --converter build/bin/audiocpp_gguf --output-dir gguf-packages +``` + +The safetensors source layout stays supported for development; the packaged/default format is GGUF. + +## Relevant building blocks already in-tree + +- Vocos vocoder: `src/models/vevo2/components.cpp`, `src/models/index_tts2/` +- iSTFT: `src/models/miocodec/`, `src/models/seed_vc/` +- Flow matching: `src/models/vevo2/fm.cpp` +- RoPE DiT / adaLN: `src/models/stable_audio/foundation/rf_dit.cpp` + +## Checkpoints + +| Model | Source | License | +|---|---|---| +| F5-TTS Base (en/zh) | `SWivid/F5-TTS` | cc-by-nc-4.0 | +| Habibi Unified (ar) | `SWivid/Habibi-TTS` | cc-by-nc-sa-4.0 | + +Both are non-commercial licenses; keep that in mind before shipping anything built on them. diff --git a/docs/community_models/models.md b/docs/community_models/models.md index 94058c63..97f419ea 100644 --- a/docs/community_models/models.md +++ b/docs/community_models/models.md @@ -16,6 +16,7 @@ Practical expectations: | Family | Task | Supported language(s) | Contributor | What They Added | |---|---|---|---|---| +| **f5_tts** | TTS, voice cloning | en, ar (Habibi) | Community | [F5-TTS](f5_tts.md) flow-matching DiT — M0 scaffolding, aliases `habibi`/`habibi_tts` | | **glm_tts** | TTS, voice cloning | zh, en | Mirek [@mirek190](https://github.com/mirek190) | [GLM-TTS](glm_tts.md) zero-shot synthesis and voice cloning support | | **inflect_v2** | TTS | en | Community | [Inflect Micro v2 and Nano v2](inflect_v2.md) native FP32 offline synthesis | | **kroko_asr** | ASR | de, en, es, fr, it, he, nl, pt, sv, tr | Mirek [@mirek190](https://github.com/mirek190) | [Kroko Community ASR](kroko_asr.md) native offline/streaming Zipformer2/RNN-T transcription with word timestamps | diff --git a/include/engine/community_models/f5_tts/dit_modules.h b/include/engine/community_models/f5_tts/dit_modules.h new file mode 100644 index 00000000..60cbae1f --- /dev/null +++ b/include/engine/community_models/f5_tts/dit_modules.h @@ -0,0 +1,53 @@ +#pragma once + +#include "engine/community_models/f5_tts/runtime.h" +#include "engine/community_models/f5_tts/weights.h" + +#include "engine/framework/core/module.h" + +namespace engine::models::f5_tts { + +struct F5DiTGraphBuild { + core::TensorValue x; // leaf [1, N, 100] + core::TensorValue cond; // leaf [1, N, 100] + core::TensorValue text_ids; // leaf [NT] (i32) + core::TensorValue time_input; // leaf [1, 256] + core::TensorValue output; // [1, N, 100] +}; + +// Builds the DiT velocity graph from framework modules (see dit_modules.cpp). +F5DiTGraphBuild build_dit_modules_graph( + ggml_context * ggml, + const F5DiTWeights & w, + const F5Architecture & arch, + int frames, + int text_len, + core::BackendType backend_type); + +} // namespace engine::models::f5_tts + +namespace engine::models::f5_tts { +// CUDA build-time constant staging (internal; used by runtime.cpp) +struct ConstStage; +std::vector * const_stage_begin(); +void const_stage_bind(std::vector * stage, ggml_backend_t backend); +void const_stage_upload(std::vector * stage, ggml_backend_t backend); +void const_stage_end(std::vector * stage); +} // namespace engine::models::f5_tts + +namespace engine::models::f5_tts { + +// Batched-CFG (B=2) variant of the module-composed DiT graph. +F5DiTGraphBuild build_dit_cfg_modules_graph( + ggml_context * ggml, + const F5DiTWeights & w, + const F5Architecture & arch, + int frames, + int text_len, + core::BackendType backend_type); + +// debug: registered stage taps from the last CFG graph build (F5_DUMP_STAGES=1) +std::vector> & stage_taps(); + +} // namespace engine::models::f5_tts + diff --git a/include/engine/community_models/f5_tts/runtime.h b/include/engine/community_models/f5_tts/runtime.h new file mode 100644 index 00000000..333ddd6a --- /dev/null +++ b/include/engine/community_models/f5_tts/runtime.h @@ -0,0 +1,94 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include +#include + +namespace engine::models::f5_tts { + +// F5TTS_v1_Base / Habibi architecture (validated against the checkpoint: +// 22 blocks, dim 1024, 16 heads x 64, ff_mult 2, text_dim 512, conv_layers 4, +// 100 mel channels, embedding rows 2731). +struct F5Architecture { + int dim = 1024; + int depth = 22; + int heads = 16; + int head_dim = 64; + int ff_mult = 2; + int text_dim = 512; + int conv_layers = 4; + int mel_dim = 100; + int vocab_rows = 2731; + int sample_rate = 24000; + int hop_length = 256; + int n_fft = 1024; +}; + +struct F5SampleOptions { + float cfg_strength = 2.0F; + float sway_sampling_coef = -1.0F; + float speed = 1.0F; + int steps = 32; + uint32_t seed = 0; +}; + +// Compute device for the DiT forward: CUDA device index or CPU threads. +struct F5ComputeDevice { + bool use_cuda = false; + int device = 0; // CUDA device index + int threads = 0; // CPU threads; 0 = hardware concurrency + // FP16 linear weights: GEMMs get ~3x faster on tensor cores but each + // mul_mat converts the F32 activations to F16 first; at F5's GEMM sizes + // (K=1024/2048, N~1022) the conversion overhead outweighs the gain on an + // RTX 3090 (measured 4.0s -> 5.0s per clip). Off by default. + bool fp16_weights = false; +}; + +// Debug taps for parity testing: when non-null, intermediate stage outputs are +// appended (column layout, [features, T] flattened feature-major). +struct F5DebugTaps { + std::vector * text_embed = nullptr; // after lookup + pos (01) + std::vector * text_convnext = nullptr; // after 4 ConvNeXt (02) + std::vector * text_padded = nullptr; // after pad/curtail (03) + std::vector * input_embed = nullptr; // after proj + cpe (04) + std::vector * time_embed = nullptr; // (05) + std::vector * block0 = nullptr; // (07_block0) + std::vector * block21 = nullptr; // (07_block21) +}; + +// Full DiT velocity forward for one Euler step inputs. +// x/cond: [seq_len * mel_dim] row-major (seq-major), text: ids (0 = filler), +// returns [seq_len * mel_dim] column layout (out[m * N + n]). +std::vector f5_dit_forward( + const std::string & weights_path, + const std::vector & x, + const std::vector & cond, + const std::vector & text, + float time_value, + int seq_len, + const F5Architecture & arch, + bool drop_audio_cond, + bool drop_text, + const F5DebugTaps * taps = nullptr, + const F5ComputeDevice * device = nullptr); + +// Batched CFG: one ne3=2 graph compute returning {conditioned, unconditioned} +// velocities. Matches python cfg_infer: the uncond half runs with +// drop_audio_cond (zeroed cond) and drop_text (filler text id 0); the host +// upload in runtime.cpp prepares both halves accordingly. +std::pair, std::vector> f5_dit_forward_cfg( + const std::string & weights_path, + const std::vector & x, + const std::vector & cond, + const std::vector & text, + float time_value, + int seq_len, + const F5Architecture & arch, + const F5ComputeDevice * device = nullptr); + +} // namespace engine::models::f5_tts diff --git a/include/engine/community_models/f5_tts/session.h b/include/engine/community_models/f5_tts/session.h new file mode 100644 index 00000000..6cbfe642 --- /dev/null +++ b/include/engine/community_models/f5_tts/session.h @@ -0,0 +1,54 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/session.h" +#include "engine/framework/runtime/spec_backed_model.h" + +#include +#include +#include + +namespace engine::models::f5_tts { + +// F5-TTS community model assets: resource bundle + resolved checkpoint path. +// Weights are loaded lazily by the runtime on first synthesis (graph cache). +struct F5TTSAssets { + assets::ResourceBundle resources; + std::filesystem::path checkpoint; // DiT *.safetensors (ema weights) +}; + +class F5TTSSession final : public runtime::IOfflineVoiceTaskSession { +public: + F5TTSSession( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options, + std::shared_ptr assets, + std::shared_ptr contract); + + std::string family() const noexcept override; + runtime::VoiceTaskKind task_kind() const noexcept override; + runtime::RunMode run_mode() const noexcept override; + void prepare(const runtime::SessionPreparationRequest & request) override; + + runtime::TaskResult run(const runtime::TaskRequest & request) override; + +private: + runtime::VoiceTaskKind task_kind_; + runtime::RunMode run_mode_; + std::shared_ptr assets_; + std::shared_ptr contract_; + std::string vocos_path_; + std::string dialect_ = "UNK"; + int frame_budget_ = 0; // 0 = default 2048 + bool use_cuda_ = false; + int cuda_device_ = 0; + int threads_ = 0; +}; + +std::shared_ptr load_f5_tts_assets( + const std::filesystem::path & model_path); + +std::shared_ptr make_f5_tts_loader(); + +} // namespace engine::models::f5_tts diff --git a/include/engine/community_models/f5_tts/synthesize.h b/include/engine/community_models/f5_tts/synthesize.h new file mode 100644 index 00000000..bfcf3f41 --- /dev/null +++ b/include/engine/community_models/f5_tts/synthesize.h @@ -0,0 +1,58 @@ +#pragma once + +#include "engine/community_models/f5_tts/runtime.h" +#include +#include +#include + +namespace engine::models::f5_tts { + +// ---- inference configuration ---- +struct F5SynthesisRequest { + std::string text; // text to speak (dialect token prepended internally) + std::string dialect = "UNK"; // UNK MSA SAU UAE ALG IRQ EGY MAR OMN TUN LEV SDN LBY + std::vector ref_audio; // mono float samples, ANY sample rate (resampled to 24k) + int ref_sample_rate = 24000; + std::string ref_text; // transcript of ref_audio + float speed = 1.0F; + int steps = 32; + float cfg_strength = 2.0F; + float sway_sampling_coef = -1.0F; + uint32_t seed = 0; + bool fixed_seed = false; + int frame_budget = 0; // total mel frames per CFM pass; 0 = default 2048 + // Strip harakat/tanwin/shadda/tatweel from the input (Habibi is trained + // on undiacritized ASR transcripts; diacritized input garbles). + bool strip_diacritics = true; + int threads = 0; // 0 = hardware concurrency + bool use_cuda = false; + int cuda_device = 0; +}; + +struct F5SynthesisResult { + std::vector audio; // 24 kHz mono + int64_t sample_rate = 24000; + double generation_seconds = 0.0; // wall time +}; + +// Full Habibi/F5 inference: text + ref audio -> waveform. +F5SynthesisResult f5_synthesize( + const std::string & model_path, + const std::string & vocos_path, + const F5SynthesisRequest & request); + +#ifdef F5_MEL_TEST +// test hook: log-mel frontend for parity tests +std::vector f5_test_mel(const std::vector & wav); +std::vector f5_test_vocos(const std::string & vocos_path, const std::vector & mel); +std::vector f5_test_vocos_gpu(const std::string & vocos_path, const std::vector & mel, const F5ComputeDevice & dev); +// test hook: full text pipeline (dialect wrap + ref trailing-space rule + +// UTF-8 char tokenization) -> vocab ids, for parity vs python list_str_to_idx +std::vector f5_test_token_ids( + const std::string & model_path, + const std::string & dialect, + const std::string & ref_text, + const std::string & gen_text); +#endif + +} // namespace engine::models::f5_tts diff --git a/include/engine/community_models/f5_tts/weights.h b/include/engine/community_models/f5_tts/weights.h new file mode 100644 index 00000000..36390b80 --- /dev/null +++ b/include/engine/community_models/f5_tts/weights.h @@ -0,0 +1,71 @@ +#pragma once + +#include "engine/community_models/f5_tts/runtime.h" +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/module.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/attention/scaled_dot_product_attention.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/positional_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 +#include +#include +#include + +namespace engine::models::f5_tts { + +// --------------------------------------------------------------------------- +// Weights, expressed with framework module weight types (dev-branch pattern: +// the weight store loads torch-shaped tensors once; modules consume them). +// --------------------------------------------------------------------------- + +struct F5TextConvNextWeights { + modules::DepthwiseConv1dWeights dwconv; // [C, 1, 7] + [C] + modules::NormWeights norm; // [C] + [C] + modules::LinearWeights pw1; // [1024, 512] + modules::LinearWeights pw2; // [512, 1024] + std::vector grn_gamma; // [1024] (host: GRN has no module) + std::vector grn_beta; // [1024] +}; + +struct F5BlockWeights { + modules::LinearWeights attn_norm; // adaLN: [6*D, D] + modules::LinearWeights to_q; // [D, D] + modules::LinearWeights to_k; + modules::LinearWeights to_v; + modules::LinearWeights to_out; + modules::LinearWeights ff0; // [2*D, D] + modules::LinearWeights ff2; // [D, 2*D] +}; + +struct F5DiTWeights { + int64_t vocab_size = 2731; // from the checkpoint's text embedding + std::shared_ptr store; + core::TensorValue text_embedding; // [vocab, 512] + modules::LinearWeights input_proj; // [1024, 712] + modules::Conv1dWeights cpe0; // [1024, 64, 31] (grouped g=16) + modules::Conv1dWeights cpe2; + modules::LinearWeights time0; // [256, 256] + modules::LinearWeights time2; // [1024, 256] + std::vector text_blocks; // x4 + std::vector blocks; // x22 + modules::LinearWeights norm_out; // [2*D, D] + modules::LinearWeights proj_out; // [100, D] +}; + +// Loads the EMA weights (prefix stripped) into module weight types. +F5DiTWeights load_dit_weights( + const assets::TensorSource & source, + ggml_backend_t backend, + core::BackendType backend_type); + +} // namespace engine::models::f5_tts diff --git a/include/engine/framework/runtime/model.h b/include/engine/framework/runtime/model.h index 995e05bc..8a089044 100644 --- a/include/engine/framework/runtime/model.h +++ b/include/engine/framework/runtime/model.h @@ -124,6 +124,8 @@ class IVoiceModelLoader { virtual ~IVoiceModelLoader() = default; virtual std::string family() const = 0; + /** Alternate family names accepted for family hints (e.g. "habibi" for "f5_tts"). */ + virtual std::vector family_aliases() const { return {}; } virtual bool can_load(const ModelLoadRequest & request) const = 0; virtual ModelInspection inspect(const ModelLoadRequest & request) const = 0; virtual std::unique_ptr load(const ModelLoadRequest & request) const = 0; diff --git a/include/engine/framework/runtime/spec_backed_model.h b/include/engine/framework/runtime/spec_backed_model.h index 59921d9c..12185f3f 100644 --- a/include/engine/framework/runtime/spec_backed_model.h +++ b/include/engine/framework/runtime/spec_backed_model.h @@ -4,6 +4,7 @@ #include "engine/framework/model_spec/package.h" #include "engine/framework/runtime/model.h" +#include #include #include #include @@ -24,6 +25,7 @@ struct OptionV1CompatibilityAlias { template struct SpecBackedVoiceModelConfig { std::string family; + std::vector aliases; // accepted family hints (spec resolution still uses family) std::function(const std::filesystem::path &)> load_assets; std::function( const TaskSpec &, @@ -159,12 +161,21 @@ class SpecBackedVoiceModelLoader final : public IVoiceModelLoader { return config_.family; } + std::vector family_aliases() const override { + return config_.aliases; + } + + bool family_matches(const std::string & hint) const { + if (hint == config_.family) return true; + return std::find(config_.aliases.begin(), config_.aliases.end(), hint) != config_.aliases.end(); + } + CapabilitySet advertised_capabilities() const override { return require_model_contract(config_.family)->capabilities; } bool can_load(const ModelLoadRequest & request) const override { - if (request.family_hint.has_value() && *request.family_hint != config_.family) { + if (request.family_hint.has_value() && !family_matches(*request.family_hint)) { return false; } try { diff --git a/model_specs/f5_tts.json b/model_specs/f5_tts.json new file mode 100644 index 00000000..c26a737f --- /dev/null +++ b/model_specs/f5_tts.json @@ -0,0 +1,314 @@ +{ + "schema_version": 1, + "family": "f5_tts", + "display_name": "F5-TTS", + "description": "Community F5-TTS flow-matching diffusion transformer for zero-shot voice cloning, including Arabic via the Habibi-TTS finetune (SWivid). Text conv conditioner, RoPE DiT, CFM sampler, Vocos vocoder.", + "category": "tts", + "status": "community", + "tasks": [ + "tts", + "clone" + ], + "modes": [ + "offline" + ], + "languages": [ + "en", + "ar" + ], + "runtime": { + "tags": [ + "server" + ] + }, + "capabilities": { + "clone": [ + "speaker_reference" + ] + }, + "options": { + "request": [ + { + "name": "reference_text", + "type": "string", + "description": "Transcript matching the reference voice audio; required by F5-TTS zero-shot cloning.", + "required": true + }, + { + "name": "dialect", + "type": "string", + "description": "Habibi dialect token: UNK MSA SAU UAE ALG IRQ EGY MAR OMN TUN LEV SDN LBY; default UNK.", + "required": false, + "default": "UNK" + }, + { + "name": "cfg_strength", + "type": "float", + "description": "Classifier-free guidance strength; default 2.0.", + "required": false, + "min": 0.0, + "max": 10.0, + "default": 2.0 + }, + { + "name": "sway_sampling_coef", + "type": "float", + "description": "Sway sampling coefficient for inference timesteps; default -1.0.", + "required": false, + "min": -5.0, + "max": 5.0, + "default": -1.0 + }, + { + "name": "speed", + "type": "float", + "description": "Speech speed multiplier applied via the F5 fix_duration parameter; default 1.0.", + "required": false, + "min": 0.5, + "max": 2.0, + "default": 1.0 + }, + { + "name": "seed", + "type": "int", + "description": "Non-negative noise seed; default 0.", + "required": false, + "min": 0, + "default": 0 + }, + { + "name": "num_inference_steps", + "type": "int", + "description": "CFM Euler steps (NFE); default 32.", + "required": false, + "min": 1, + "default": 32 + }, + { + "name": "guidance_scale", + "type": "float", + "description": "Classifier-free guidance strength; default 2.0.", + "required": false, + "min": 0.0, + "max": 10.0, + "default": 2.0 + }, + { + "name": "strip_diacritics", + "type": "bool", + "description": "Strip Arabic combining marks (harakat/tanwin/shadda/tatweel) before synthesis; Habibi is trained on undiacritized ASR transcripts and garbles diacritized input. Default true.", + "required": false, + "default": true + } + ], + "session": [ + { + "name": "vocos_path", + "type": "string", + "description": "Path to the Vocos vocoder checkpoint (vocos.safetensors); required unless placed next to the DiT checkpoint.", + "required": false + }, + { + "name": "dialect", + "type": "string", + "description": "Default Habibi dialect token for requests that do not set it; default UNK.", + "required": false, + "default": "UNK" + }, + { + "name": "frame_budget", + "type": "int", + "description": "Total mel frames per CFM pass (ref + generated); larger budgets allow sentence-scale chunks at higher VRAM. 0 = default 2048.", + "required": false, + "min": 0, + "default": 0 + } + ], + "load": [] + }, + "package_defaults": { + "download": { + "kind": "huggingface_snapshot", + "repo": "trklou/audio.cpp", + "revision": "main", + "gated": false + } + }, + "packages": [ + { + "id": "habibi_unified", + "display_name": "Habibi-TTS Unified (Arabic, multi-dialect)", + "description": "Unified multi-dialect Arabic checkpoint from SWivid/Habibi-TTS, converted to GGUF (DiT transformer + Vocos vocoder namespaces, f32).", + "default": true, + "format": "gguf", + "precision": "orig", + "target_directory": "Habibi-TTS/Unified", + "files": [ + "habibi-unified/habibi-unified-orig.gguf", + "habibi-unified/vocab.txt" + ], + "strip_prefix": "habibi-unified" + }, + { + "id": "vocos_mel_24khz", + "display_name": "Vocos mel 24kHz vocoder (GGUF)", + "description": "Standalone Vocos vocoder GGUF for use with the original safetensors checkpoints.", + "default": false, + "format": "gguf", + "precision": "orig", + "target_directory": "vocos-mel-24khz", + "files": [ + "vocos-mel-24khz/vocos-mel-24khz-orig.gguf" + ], + "strip_prefix": "vocos-mel-24khz" + }, + { + "id": "habibi_alg", + "display_name": "Habibi-TTS ALG specialized checkpoint (GGUF)", + "description": "Single-dialect ALG checkpoint from SWivid/Habibi-TTS, converted to GGUF (DiT transformer + Vocos vocoder namespaces, f32). Stronger ALG accent than the unified model.", + "default": false, + "format": "gguf", + "precision": "orig", + "target_directory": "Habibi-TTS/Specialized/ALG", + "files": [ + "habibi-alg/habibi-alg-orig.gguf", + "habibi-alg/vocab.txt" + ], + "strip_prefix": "habibi-alg" + }, + { + "id": "habibi_egy", + "display_name": "Habibi-TTS EGY specialized checkpoint (GGUF)", + "description": "Single-dialect EGY checkpoint from SWivid/Habibi-TTS, converted to GGUF (DiT transformer + Vocos vocoder namespaces, f32). Stronger EGY accent than the unified model.", + "default": false, + "format": "gguf", + "precision": "orig", + "target_directory": "Habibi-TTS/Specialized/EGY", + "files": [ + "habibi-egy/habibi-egy-orig.gguf", + "habibi-egy/vocab.txt" + ], + "strip_prefix": "habibi-egy" + }, + { + "id": "habibi_irq", + "display_name": "Habibi-TTS IRQ specialized checkpoint (GGUF)", + "description": "Single-dialect IRQ checkpoint from SWivid/Habibi-TTS, converted to GGUF (DiT transformer + Vocos vocoder namespaces, f32). Stronger IRQ accent than the unified model.", + "default": false, + "format": "gguf", + "precision": "orig", + "target_directory": "Habibi-TTS/Specialized/IRQ", + "files": [ + "habibi-irq/habibi-irq-orig.gguf", + "habibi-irq/vocab.txt" + ], + "strip_prefix": "habibi-irq" + }, + { + "id": "habibi_mar", + "display_name": "Habibi-TTS MAR specialized checkpoint (GGUF)", + "description": "Single-dialect MAR checkpoint from SWivid/Habibi-TTS, converted to GGUF (DiT transformer + Vocos vocoder namespaces, f32). Stronger MAR accent than the unified model.", + "default": false, + "format": "gguf", + "precision": "orig", + "target_directory": "Habibi-TTS/Specialized/MAR", + "files": [ + "habibi-mar/habibi-mar-orig.gguf", + "habibi-mar/vocab.txt" + ], + "strip_prefix": "habibi-mar" + }, + { + "id": "habibi_msa", + "display_name": "Habibi-TTS MSA specialized checkpoint (GGUF)", + "description": "Single-dialect MSA checkpoint from SWivid/Habibi-TTS, converted to GGUF (DiT transformer + Vocos vocoder namespaces, f32). Stronger MSA accent than the unified model.", + "default": false, + "format": "gguf", + "precision": "orig", + "target_directory": "Habibi-TTS/Specialized/MSA", + "files": [ + "habibi-msa/habibi-msa-orig.gguf", + "habibi-msa/vocab.txt" + ], + "strip_prefix": "habibi-msa" + }, + { + "id": "habibi_sau", + "display_name": "Habibi-TTS SAU specialized checkpoint (GGUF)", + "description": "Single-dialect SAU checkpoint from SWivid/Habibi-TTS, converted to GGUF (DiT transformer + Vocos vocoder namespaces, f32). Stronger SAU accent than the unified model.", + "default": false, + "format": "gguf", + "precision": "orig", + "target_directory": "Habibi-TTS/Specialized/SAU", + "files": [ + "habibi-sau/habibi-sau-orig.gguf", + "habibi-sau/vocab.txt" + ], + "strip_prefix": "habibi-sau" + }, + { + "id": "habibi_uae", + "display_name": "Habibi-TTS UAE specialized checkpoint (GGUF)", + "description": "Single-dialect UAE checkpoint from SWivid/Habibi-TTS, converted to GGUF (DiT transformer + Vocos vocoder namespaces, f32). Stronger UAE accent than the unified model.", + "default": false, + "format": "gguf", + "precision": "orig", + "target_directory": "Habibi-TTS/Specialized/UAE", + "files": [ + "habibi-uae/habibi-uae-orig.gguf", + "habibi-uae/vocab.txt" + ], + "strip_prefix": "habibi-uae" + } + ], + "dependencies": [], + "ui": { + "recommended_package": "habibi_unified", + "tags": [ + "TTS", + "Clone" + ], + "docs": [ + "docs/community_models/f5_tts.md" + ] + }, + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "files": { + "vocab": "model:vocab.txt" + }, + "tensors": { + "transformer": { + "source": "weights:", + "prefix": "transformer" + } + }, + "optional_tensors": { + "vocos_vocoder": { + "source": "weights:", + "prefix": "vocos" + } + } + }, + { + "format": "safetensors", + "roots": { + "model": "." + }, + "files": { + "vocab": "model:vocab.txt" + }, + "tensors": { + "transformer": { + "source": "model:model_200000.safetensors", + "prefix": "ema_model.transformer" + } + } + } + ] +} diff --git a/src/community_models/f5_tts/cpu_graph_compute.h b/src/community_models/f5_tts/cpu_graph_compute.h new file mode 100644 index 00000000..a4ceaec1 --- /dev/null +++ b/src/community_models/f5_tts/cpu_graph_compute.h @@ -0,0 +1,72 @@ +#pragma once + +// ggml_graph_compute_with_ctx lives in the CPU backend (ggml-cpu). With +// GGML_BACKEND_DL builds (ENGINE_ENABLE_CPU_ALL_VARIANTS, e.g. the docker +// image) that backend is a dlopen'd MODULE library, so the symbol is not +// linkable from engine_runtime. Resolve it at runtime instead: +// - static-link builds bind the weak reference directly; +// - dlopen'd builds find the already-loaded libggml-cpu via the dynamic +// linker's loaded-object list (it is RTLD_LOCAL, so plain +// dlsym(RTLD_DEFAULT) would miss it). + +#include "ggml.h" + +#include +#include + +#if defined(__linux__) +#include +#include +#endif + +namespace engine::models::f5_tts { + +#if defined(__GNUC__) +extern "C" ggml_status ggml_graph_compute_with_ctx( + ggml_context * ctx, ggml_cgraph * cgraph, int n_threads) __attribute__((weak)); +#endif + +using F5CpuGraphComputeFn = ggml_status (*)(ggml_context *, ggml_cgraph *, int); + +inline F5CpuGraphComputeFn f5_cpu_graph_compute_fn() { +#if defined(__GNUC__) + if (ggml_graph_compute_with_ctx != nullptr) { + return ggml_graph_compute_with_ctx; // statically linked build + } +#endif +#if defined(__linux__) + static F5CpuGraphComputeFn fn = [] { + void * sym = nullptr; + dl_iterate_phdr( + [](struct dl_phdr_info * info, size_t, void * data) -> int { + if (info->dlpi_name == nullptr || + std::strstr(info->dlpi_name, "libggml-cpu") == nullptr) { + return 0; + } + void * h = dlopen(info->dlpi_name, RTLD_NOW | RTLD_NOLOAD); + if (h == nullptr) { + return 0; + } + void * s = dlsym(h, "ggml_graph_compute_with_ctx"); + if (s != nullptr) { + *reinterpret_cast(data) = s; + return 1; + } + return 0; + }, + &sym); + return reinterpret_cast(sym); + }(); + if (fn != nullptr) { + return fn; + } +#endif + throw std::runtime_error( + "F5-TTS: ggml_graph_compute_with_ctx unavailable (CPU backend not loaded)"); +} + +inline ggml_status f5_cpu_graph_compute(ggml_context * ctx, ggml_cgraph * graph, int threads) { + return f5_cpu_graph_compute_fn()(ctx, graph, threads); +} + +} // namespace engine::models::f5_tts diff --git a/src/community_models/f5_tts/dit_modules.cpp b/src/community_models/f5_tts/dit_modules.cpp new file mode 100644 index 00000000..b579a798 --- /dev/null +++ b/src/community_models/f5_tts/dit_modules.cpp @@ -0,0 +1,742 @@ +// F5 DiT forward composed from framework modules (dev-branch pattern). +// +// All activations use the framework's logical [batch, frames, features] +// layout; modules handle the ggml mapping. Verified stage-by-stage against +// the same goldens as the original raw-ggml implementation (cosine 1.0). +#include "engine/community_models/f5_tts/weights.h" + +#include "engine/framework/core/module.h" + +#include "ggml-backend.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/attention/scaled_dot_product_attention.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/positional_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" + +#include +#include +#include +#include + +namespace engine::models::f5_tts { + +// Build-time constant staging: on CPU (inline ctx) values are written +// directly; on CUDA (no_alloc ctx) they are staged for the caller to upload +// via ggml_backend_tensor_set after allocation. +struct ConstStage { + ggml_tensor * tensor; + std::vector bytes; + ggml_backend_buffer_t owned_buffer = nullptr; +}; +thread_local std::vector * t_const_stage = nullptr; + +namespace { + +namespace mod = engine::modules; + +core::ModuleBuildContext make_ctx( + ggml_context * ggml, const char * name, core::BackendType type) { + return core::ModuleBuildContext{ggml, name, type}; +} + +// ---- grouped conv1d (groups=16, k=31) via per-group im2col + matmul ---- +// The framework has no grouped Conv1d module yet; lower it with the same +// primitives Conv1dModule uses (im2col + mul_mat), expressed as modules. +core::TensorValue grouped_conv1d( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, // [B, C_in, T] + const core::TensorValue & weight, // [C_out, C_in/g, k] + const core::TensorValue & bias, // [C_out] + int64_t groups) { + const int64_t frames = input.shape.dims[2]; + const int64_t c_in = input.shape.dims[1]; + const int64_t c_out = weight.shape.dims[0]; + const int64_t kernel = weight.shape.dims[2]; + const int64_t cg_in = c_in / groups; + const int64_t cg_out = c_out / groups; + + std::vector group_outputs; + group_outputs.reserve(static_cast(groups)); + for (int64_t g = 0; g < groups; ++g) { + // input slice [B, cg_in, T] on axis 1; im2col needs contiguous input + auto in_sliced = mod::SliceModule({1, g * cg_in, cg_in}).build(ctx, input); + auto in_g = core::ensure_backend_addressable_layout(ctx, in_sliced); + auto * cols = ggml_im2col( + ctx.ggml, weight.tensor, in_g.tensor, + 1, 1, kernel / 2, 0, 1, 1, false, GGML_TYPE_F32); + // ggml im2col output ne [cg_in*k, T]; mul_mat with the per-group + // weight slice (ggml ne [k, cg_in, cg_out] -> view as [K, cg_out]) + auto * w_g = ggml_view_3d( + ctx.ggml, weight.tensor, + kernel, cg_in, cg_out, + weight.tensor->nb[1], weight.tensor->nb[2], + g * cg_out * weight.tensor->nb[2]); + auto * w2 = ggml_reshape_2d(ctx.ggml, w_g, cg_in * kernel, cg_out); + auto * y_raw = ggml_mul_mat(ctx.ggml, w2, cols); // ggml [cg_out, T] + // logical [T, cg_out] shares the same memory; wrap for module use + auto y = core::wrap_tensor( + y_raw, core::TensorShape::from_dims({frames, cg_out}), GGML_TYPE_F32); + auto b_g = mod::SliceModule({0, g * cg_out, cg_out}).build(ctx, bias); + auto b_row = core::reshape_tensor( + ctx, b_g, core::TensorShape::from_dims({1, cg_out})); + y = mod::AddModule().build( + ctx, y, mod::RepeatModule({y.shape}).build(ctx, b_row)); + group_outputs.push_back(y); + } + auto out = group_outputs[0]; + for (size_t i = 1; i < group_outputs.size(); ++i) { + out = mod::ConcatModule({1}).build(ctx, out, group_outputs[i]); // [T, c_out] + } + return out; // [T, c_out] logical (batch folded into frames column) +} + +core::TensorValue ctx_store_f32( + core::ModuleBuildContext & ctx, const core::TensorShape & shape, + const std::vector & values) { + auto t = core::make_tensor(ctx, GGML_TYPE_F32, shape); + ggml_set_input(t.tensor); // literal data: never scratch for the allocator + if (t.tensor->data != nullptr) { + std::memcpy(t.tensor->data, values.data(), values.size() * sizeof(float)); + } else if (t_const_stage != nullptr) { + const auto * b = reinterpret_cast(values.data()); + t_const_stage->push_back({t.tensor, std::vector(b, b + values.size() * sizeof(float))}); + } + return t; +} + +// softplus(x): ggml's numerically stable primitive (same op the raw path used) +core::TensorValue exp_log_softplus(core::ModuleBuildContext & ctx, const core::TensorValue & x) { + return core::wrap_tensor(ggml_softplus(ctx.ggml, x.tensor), x.shape, GGML_TYPE_F32); +} + +// lift a [1, C] row to [1, 1, C] so Repeat can broadcast over frames +core::TensorValue lift_row( + core::ModuleBuildContext & ctx, const core::TensorValue & row) { + if (row.shape.rank == 3) { + return row; + } + return core::reshape_tensor( + ctx, core::ensure_backend_addressable_layout(ctx, row), + core::TensorShape::from_dims({1, 1, row.shape.dims[row.shape.rank - 1]})); +} + +// adaLN modulate: x * (1 + scale) + shift, scale/shift [1, C] +core::TensorValue modulate( + core::ModuleBuildContext & ctx, + const core::TensorValue & x, + const core::TensorValue & scale, + const core::TensorValue & shift) { + // constants sized [1, 1, C] (NOT [B, T, C]): one shared ones-row, and the + // repeat broadcasts handle the expansion in the compute buffer. + const auto row_shape = core::TensorShape::from_dims( + {1, 1, x.shape.dims[x.shape.rank - 1]}); + auto ones = ctx_store_f32( + ctx, row_shape, std::vector(static_cast(x.shape.dims[x.shape.rank - 1]), 1.0F)); + auto s_row = lift_row(ctx, scale); + auto sh_row = lift_row(ctx, shift); + // 1 + scale, then broadcast once + auto one_plus_s = mod::AddModule().build(ctx, ones, s_row); // [1,1,C] + auto scale_b = mod::RepeatModule({x.shape}).build(ctx, one_plus_s); + auto shift_b = mod::RepeatModule({x.shape}).build(ctx, sh_row); + return mod::AddModule().build( + ctx, mod::MulModule().build(ctx, x, scale_b), shift_b); +} + +// ---- GRN (global response norm): no framework module; expressed with +// primitives on [B, T, C] ---- +core::TensorValue grn( + core::ModuleBuildContext & ctx, + const core::TensorValue & h, // [B, T, C] + const std::vector & gamma, + const std::vector & beta, + core::TensorValue & gamma_out, + core::TensorValue & beta_out) { + // per-channel L2 norm over T, then normalize by mean over channels: + // gx[c] = ||h[:, :, c]||_2 ; nx[c] = gx[c] / (mean(gx) + 1e-6) + const int64_t channels = h.shape.dims[2]; + // squares -> sum over T -> sqrt (reduce axis=1) + auto sq = mod::MulModule().build(ctx, h, h); // elementwise + auto sums = mod::ReduceSumModule({1}).build(ctx, sq); // [B, 1, C] + auto gx = mod::SqrtModule().build(ctx, sums); + auto mean = mod::ReduceMeanModule({2}).build(ctx, gx); // [B, 1, 1] + // nx = gx / (mean + eps): eps via a constant scalar broadcast + // (ReduceMean + add-eps fused: use MeanModule output + eps tensor) + // eps as a [1,1,1] constant: + auto eps = ctx_store_f32(ctx, mean.shape, std::vector{1e-6F}); + auto denom = mod::AddModule().build(ctx, mean, eps); + auto nx = core::wrap_tensor( + ggml_div(ctx.ggml, gx.tensor, denom.tensor), gx.shape, GGML_TYPE_F32); + // scale h by nx per channel + gamma*x + beta + residual + auto nx_b = mod::RepeatModule({h.shape}).build(ctx, nx); // [B,1,C]->[B,T,C] via repeat + auto scaled = mod::MulModule().build(ctx, h, nx_b); + // gamma/beta as [1, 1, C] constants + gamma_out = ctx_store_f32(ctx, core::TensorShape::from_dims({1, 1, channels}), gamma); + beta_out = ctx_store_f32(ctx, core::TensorShape::from_dims({1, 1, channels}), beta); + auto g_b = mod::RepeatModule({h.shape}).build(ctx, gamma_out); + auto b_b = mod::RepeatModule({h.shape}).build(ctx, beta_out); + auto out = mod::AddModule().build(ctx, mod::MulModule().build(ctx, scaled, g_b), b_b); + return mod::AddModule().build(ctx, out, h); // + residual +} + + +} // namespace + +// Builds the full DiT velocity graph. Leaves: x/cond [B=1, T, MEL], text ids +// [NT], time-embedding input [1, 256]. Returns the output TensorValue. +struct F5DiTGraphBuild { + core::TensorValue x; + core::TensorValue cond; + core::TensorValue text_ids; + core::TensorValue time_input; + core::TensorValue output; +}; + + +std::vector * const_stage_begin() { + t_const_stage = new std::vector(); + return t_const_stage; +} +// Bind staged constants to private backend buffers BEFORE the gallocr +// reserves the compute arena: a tensor with data already set is treated as +// externally owned and never aliased by scratch reuse. + +// --- cross-val stage dumps (debug only; compiled out of production) --- +std::vector> g_stage_taps; +static void tap_stage(const char * name, const core::TensorValue & t) { +#ifdef F5_MEL_TEST + if (std::getenv("F5_DUMP_STAGES") == nullptr) return; + ggml_set_output(t.tensor); // protect from arena reuse so taps are readable post-compute + g_stage_taps.emplace_back(name, t.tensor); +#else + (void) name; (void) t; +#endif +} + +static void tap_stage_cond(bool cond, const char * name, const core::TensorValue & t) { + if (cond) tap_stage(name, t); +} + +std::vector> & stage_taps() { + return g_stage_taps; +} + +void const_stage_bind(std::vector * stage, ggml_backend_t backend) { + if (stage == nullptr || backend == nullptr) return; + for (auto & c : *stage) { + if (c.tensor->data != nullptr) continue; // inline ctx already + const size_t nbytes = ggml_nbytes(c.tensor); + ggml_backend_buffer_t buf = ggml_backend_alloc_buffer(backend, nbytes); + if (buf == nullptr) continue; + c.tensor->buffer = buf; + c.tensor->data = ggml_backend_buffer_get_base(buf); + c.owned_buffer = buf; // leaked with the graph (driver shutdown) + } +} + +void const_stage_upload(std::vector * stage, ggml_backend_t backend) { + (void)backend; + // Give each constant its own backend buffer OUTSIDE the compute arena: + // the gallocr may reuse the arena slot of an early-consumed input for + // later intermediates, silently corrupting constants between computes. + // (Observed: pe table and per-block ones rows drifted after one compute.) + for (auto & c : *stage) { + ggml_backend_tensor_set(c.tensor, c.bytes.data(), 0, c.bytes.size()); + } +} +void const_stage_end(std::vector * stage) { + // detach only; the vector is deliberately leaked (a few KB per graph + // build) to avoid ownership hazards across build paths. + t_const_stage = nullptr; + (void)stage; +} + + +F5DiTGraphBuild build_dit_modules_graph( + ggml_context * ggml, + const F5DiTWeights & w, + const F5Architecture & arch, + int frames, + int text_len, + core::BackendType backend_type) { + (void)arch; // validated at weight-load time; dimensions are architecture constants + auto ctx = make_ctx(ggml, "f5.dit", backend_type); + constexpr int64_t kMel = 100, kTextDim = 512, kDim = 1024; + constexpr int64_t kHeads = 16, kHeadDim = 64; + const int64_t N = frames; + const int64_t NT = text_len; + + F5DiTGraphBuild io; + io.x = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, N, kMel})); + io.cond = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, N, kMel})); + io.text_ids = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({NT})); + io.time_input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 256})); + + // ---- text embed: lookup + sinusoidal pe (both halves share) ---- + auto te = mod::EmbeddingModule({w.vocab_size, kTextDim}).build(ctx, io.text_ids, w.text_embedding); + // batch the text stream: [NT, C] -> [1, NT, C] + te = core::reshape_tensor( + ctx, core::ensure_backend_addressable_layout(ctx, te), + core::TensorShape::from_dims({1, NT, kTextDim})); + // pe table [1, NT, kTextDim] constant (verified layout: cos|sin halves) + { + std::vector pe(static_cast(NT) * kTextDim); + const int64_t half = kTextDim / 2; + for (int64_t pos = 0; pos < NT; ++pos) { + for (int64_t i = 0; i < half; ++i) { + const float inv = std::pow(10000.0F, -2.0F * i / static_cast(kTextDim)); + const float f = static_cast(pos) * inv; + pe[static_cast(pos) * kTextDim + i] = std::cos(f); + pe[static_cast(pos) * kTextDim + half + i] = std::sin(f); + } + } + auto pe_t = ctx_store_f32( + ctx, core::TensorShape::from_dims({1, NT, kTextDim}), pe); + te = mod::AddModule().build(ctx, te, pe_t); + } + + // ---- Python semantics: text encoder over padded length with filler + // positions masked after pe and after every block ---- + core::TensorValue te_mask_on; + if (NT < N) { + const auto pad_shape = core::TensorShape::from_dims({1, N - NT, kTextDim}); + std::vector zv(static_cast(pad_shape.num_elements()), 0.0F); + te = mod::ConcatModule({1}).build(ctx, te, ctx_store_f32(ctx, pad_shape, zv)); + std::vector ones(static_cast(NT), 1.0F); + std::vector zeros(static_cast(N - NT), 0.0F); + auto m_on = ctx_store_f32(ctx, core::TensorShape::from_dims({1, NT, 1}), ones); + auto m_off = ctx_store_f32(ctx, core::TensorShape::from_dims({1, N - NT, 1}), zeros); + te_mask_on = mod::ConcatModule({1}).build(ctx, m_on, m_off); + te_mask_on = mod::RepeatModule( + {core::TensorShape::from_dims({1, N, kTextDim})}).build(ctx, te_mask_on); + } else { + te_mask_on = core::TensorValue{}; + } + if (te_mask_on.tensor != nullptr) { + te = mod::MulModule().build(ctx, te, te_mask_on); + } + + // ---- 4x ConvNeXt text blocks (dwconv k7, LN, pw1+GELU, GRN, pw2, residual) ---- + for (int bi = 0; bi < 4; ++bi) { + const auto & B = w.text_blocks[static_cast(bi)]; + // ConvNeXt operates channel-major: [B, C, T]; our te is [B, T, C] + auto te_c = mod::TransposeModule({{0, 2, 1}, 3}).build(ctx, te); + auto dw = mod::DepthwiseConv1dModule({kTextDim, 7, 1, 3, 1, true}).build(ctx, te_c, B.dwconv); + auto dw_t = mod::TransposeModule({{0, 2, 1}, 3}).build(ctx, dw); // [B, T, C] + auto nrm = mod::LayerNormModule({kTextDim, 1e-6F, true, true}).build(ctx, dw_t, B.norm); + auto h1 = mod::LinearModule({kTextDim, kDim, true}).build(ctx, nrm, B.pw1); + h1 = mod::GeluModule({mod::GeluApproximation::ExactErf}).build(ctx, h1); + core::TensorValue g_dummy, b_dummy; + auto grn_out = grn(ctx, h1, B.grn_gamma, B.grn_beta, g_dummy, b_dummy); + auto h2 = mod::LinearModule({kDim, kTextDim, true}).build(ctx, grn_out, B.pw2); + te = mod::AddModule().build(ctx, te, h2); + if (te_mask_on.tensor != nullptr) { + te = mod::MulModule().build(ctx, te, te_mask_on); // re-zero pads + } + tap_stage_cond(bi == 0, "txt_h1", h1); + tap_stage_cond(bi == 0, "txt_grn", grn_out); + tap_stage_cond(bi == 0, "txt_after_block", te); + } + + // text is already exactly N frames (padded + masked before the blocks) + const auto & te_pad = te; + + // ---- input embed: concat features [x | cond | text] -> proj -> CPE ---- + auto cat0 = mod::ConcatModule({2}).build(ctx, io.x, io.cond); + auto cat1 = mod::ConcatModule({2}).build(ctx, cat0, te_pad); // [1, N, 712] + auto inp = mod::LinearModule({712LL, kDim, true}).build(ctx, cat1, w.input_proj); + + // conv pos embed (grouped k31 g16, Mish x2): + // inp += mish(conv1(mish(conv0(inp)))) + { + auto conv_mish = [&](const core::TensorValue & x_bnd, + const core::TensorValue & cweight, + const core::TensorValue & cbias) -> core::TensorValue { + // im2col requires a contiguous, time-fastest input; ggml_transpose + // yields a strided view, so materialize it first. + auto x_t = mod::TransposeModule({{0, 2, 1}, 3}).build(ctx, x_bnd); // logical [1, D, N] + auto x_c = core::wrap_tensor( + ggml_cont(ctx.ggml, x_t.tensor), x_t.shape, GGML_TYPE_F32); + auto r = grouped_conv1d(ctx, x_c, cweight, cbias, 16); // [N, D] + auto sp = exp_log_softplus(ctx, r); + auto mish = mod::MulModule().build(ctx, r, mod::TanhModule().build(ctx, sp)); + return core::reshape_tensor( + ctx, core::ensure_backend_addressable_layout(ctx, mish), + core::TensorShape::from_dims({1, N, kDim})); + }; + auto r0 = conv_mish(inp, w.cpe0.weight, *w.cpe0.bias); + auto r1 = conv_mish(r0, w.cpe2.weight, *w.cpe2.bias); + inp = mod::AddModule().build(ctx, inp, r1); + } + + // ---- time embedding: shared MLP over the per-call leaf ---- + auto t0 = mod::LinearModule({256, kDim, true}).build(ctx, io.time_input, w.time0); + t0 = mod::SiluModule().build(ctx, t0); + auto t_emb = mod::LinearModule({kDim, kDim, true}).build(ctx, t0, w.time2); // [1, 1024] + + // ---- RoPE positions [N] constant ---- + core::TensorValue positions; + { + std::vector pos(static_cast(N)); + for (int64_t i = 0; i < N; ++i) { + pos[static_cast(i)] = static_cast(i); + } + auto p = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({N})); + ggml_set_input(p.tensor); + if (p.tensor->data != nullptr) { + std::memcpy(p.tensor->data, pos.data(), pos.size() * sizeof(int32_t)); + } else if (t_const_stage != nullptr) { + const auto * b = reinterpret_cast(pos.data()); + t_const_stage->push_back({p.tensor, std::vector(b, b + pos.size() * sizeof(int32_t))}); + } + positions = p; + } + + // ---- 22 DiT blocks ---- + auto h = inp; + for (int bi = 0; bi < 22; ++bi) { + const auto & B = w.blocks[static_cast(bi)]; + // adaLN modulation: 6 chunks from silu(t_emb) @ attn_norm + auto emb = mod::LinearModule({kDim, 6 * kDim, true}).build( + ctx, mod::SiluModule().build(ctx, t_emb), B.attn_norm); // [1, 6144] + // F5 chunk order: shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp + auto shift_msa = mod::SliceModule({1, 0 * kDim, kDim}).build(ctx, emb); + auto scale_msa = mod::SliceModule({1, 1 * kDim, kDim}).build(ctx, emb); + auto gate_msa = mod::SliceModule({1, 2 * kDim, kDim}).build(ctx, emb); + auto shift_mlp = mod::SliceModule({1, 3 * kDim, kDim}).build(ctx, emb); + auto scale_mlp = mod::SliceModule({1, 4 * kDim, kDim}).build(ctx, emb); + auto gate_mlp = mod::SliceModule({1, 5 * kDim, kDim}).build(ctx, emb); + (void)shift_msa; + + // modulate: x * (1 + scale) + shift + auto norm = modulate(ctx, mod::LayerNormModule({kDim, 1e-6F, false, false}).build(ctx, h, mod::NormWeights{}), scale_msa, shift_msa); + auto q = mod::LinearModule({kDim, kDim, true}).build(ctx, norm, B.to_q); + auto k = mod::LinearModule({kDim, kDim, true}).build(ctx, norm, B.to_k); + auto v = mod::LinearModule({kDim, kDim, true}).build(ctx, norm, B.to_v); + // heads: [1, N, H, DH] (roformer reshape_heads pattern) + auto to_heads = [&](core::TensorValue t) { + return core::reshape_tensor( + ctx, core::ensure_backend_addressable_layout(ctx, t), + core::TensorShape::from_dims({t.shape.dims[0], t.shape.dims[1], kHeads, kHeadDim})); + }; + q = to_heads(q); + k = to_heads(k); + v = to_heads(v); + // ggml_rope_ext on a strided view corrupts arena neighbors (root + // cause of the garbled-output regression): materialize q/k first. + q = core::wrap_tensor(ggml_cont(ctx.ggml, q.tensor), q.shape, GGML_TYPE_F32); + k = core::wrap_tensor(ggml_cont(ctx.ggml, k.tensor), k.shape, GGML_TYPE_F32); + q = mod::RoPEModule({kHeadDim, GGML_ROPE_TYPE_NORMAL, 10000.0F}).build(ctx, q, positions); + k = mod::RoPEModule({kHeadDim, GGML_ROPE_TYPE_NORMAL, 10000.0F}).build(ctx, k, positions); + auto q_heads = mod::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, q); + auto k_heads = mod::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, k); + auto v_heads = mod::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, v); + // flash-attn requires dense, contiguous q/k/v: materialize the + // strided transposes explicitly (a strided input corrupts arena + // neighbors on replay — root cause of the noise regression) + auto dense4 = [&](const core::TensorValue & t) { + return core::wrap_tensor( + ggml_cont(ctx.ggml, t.tensor), t.shape, GGML_TYPE_F32); + }; + q_heads = dense4(q_heads); + k_heads = dense4(k_heads); + v_heads = dense4(v_heads); + auto attn = mod::ScaledDotProductAttentionModule({ + kHeadDim, + mod::ScaledDotProductAttentionLowering::Flash, + GGML_PREC_F32, + mod::AttentionCausality::NonCausal, + }).build(ctx, q_heads, k_heads, v_heads); // [1, N, H, DH] + auto attn_flat = core::reshape_tensor( + ctx, core::ensure_backend_addressable_layout(ctx, attn), + core::TensorShape::from_dims({attn.shape.dims[0], attn.shape.dims[1], kDim})); + auto proj = mod::LinearModule({kDim, kDim, true}).build(ctx, attn_flat, B.to_out); + // gated residual: h + proj * gate + h = mod::AddModule().build( + ctx, h, mod::MulModule().build( + ctx, proj, mod::RepeatModule({proj.shape}).build(ctx, lift_row(ctx, gate_msa)))); + + auto norm2 = modulate(ctx, mod::LayerNormModule({kDim, 1e-6F, false, false}).build(ctx, h, mod::NormWeights{}), scale_mlp, shift_mlp); + auto f1 = mod::LinearModule({kDim, 2048, true}).build(ctx, norm2, B.ff0); + f1 = mod::GeluModule({mod::GeluApproximation::Tanh}).build(ctx, f1); + auto f2 = mod::LinearModule({2048, kDim, true}).build(ctx, f1, B.ff2); + h = mod::AddModule().build( + ctx, h, mod::MulModule().build( + ctx, f2, mod::RepeatModule({f2.shape}).build(ctx, lift_row(ctx, gate_mlp)))); + } + + // ---- final adaLN + projection to mel ---- + { + auto emb = mod::LinearModule({kDim, 2 * kDim, true}).build( + ctx, mod::SiluModule().build(ctx, t_emb), w.norm_out); + auto scale = mod::SliceModule({1, 0, kDim}).build(ctx, emb); + auto shift = mod::SliceModule({1, kDim, kDim}).build(ctx, emb); + auto norm = modulate(ctx, mod::LayerNormModule({kDim, 1e-6F, false, false}).build(ctx, h, mod::NormWeights{}), scale, shift); + io.output = mod::LinearModule({kDim, kMel, true}).build(ctx, norm, w.proj_out); // [1, N, 100] + } + return io; +} + + + +// Batched-CFG variant: B=2 halves share every weight; halves differ only in +// text ids (cond half: ids+1 offset embeds 〈ref+text〉, uncond half: filler +// id 0 + zeroed cond, per python cfg_infer drop_audio_cond/drop_text). +F5DiTGraphBuild build_dit_cfg_modules_graph( + ggml_context * ggml, + const F5DiTWeights & w, + const F5Architecture & arch, + int frames, + int text_len, + core::BackendType backend_type) { + (void)arch; // validated at weight-load time; dimensions are architecture constants + auto ctx = make_ctx(ggml, "f5.dit.cfg", backend_type); + constexpr int64_t kMel = 100, kTextDim = 512, kDim = 1024; + constexpr int64_t kHeads = 16, kHeadDim = 64; + const int64_t N = frames; + const int64_t NT = text_len; + + F5DiTGraphBuild io; + // leaves: x/cond [2, N, 100] (both halves identical values), ids [2*NT] + io.x = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({2, N, kMel})); + io.cond = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({2, N, kMel})); + io.text_ids = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({2 * NT})); + io.time_input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 256})); + + // per-half embeddings then concat on the batch axis + auto ids_c = mod::SliceModule({0, 0, NT}).build(ctx, io.text_ids); + auto ids_u = mod::SliceModule({0, NT, NT}).build(ctx, io.text_ids); + auto emb = [&](const core::TensorValue & ids) { + auto e = mod::EmbeddingModule({w.vocab_size, kTextDim}).build(ctx, ids, w.text_embedding); + return core::reshape_tensor( + ctx, core::ensure_backend_addressable_layout(ctx, e), + core::TensorShape::from_dims({1, NT, kTextDim})); + }; + auto te_c = emb(ids_c); + auto te_u = emb(ids_u); + // shared pe over positions + std::vector pe(static_cast(NT) * kTextDim); + { + const int64_t half = kTextDim / 2; + for (int64_t pos = 0; pos < NT; ++pos) { + for (int64_t i = 0; i < half; ++i) { + const float inv = std::pow(10000.0F, -2.0F * i / static_cast(kTextDim)); + const float f = static_cast(pos) * inv; + pe[static_cast(pos) * kTextDim + i] = std::cos(f); + pe[static_cast(pos) * kTextDim + half + i] = std::sin(f); + } + } + } + auto pe_t = ctx_store_f32(ctx, core::TensorShape::from_dims({1, NT, kTextDim}), pe); + { + core::TensorValue wt; + wt.tensor = w.text_embedding.tensor; + wt.shape = w.text_embedding.shape; + tap_stage("emb_table", wt); + tap_stage("txt_ids", io.text_ids); + } + tap_stage("txt_emb_raw_c", te_c); + tap_stage("txt_pe", pe_t); + te_c = mod::AddModule().build(ctx, te_c, pe_t); + te_u = mod::AddModule().build(ctx, te_u, pe_t); + auto te = mod::ConcatModule({0}).build(ctx, te_c, te_u); // [2, NT, 512] + + // Python runs the text encoder over the FULL padded length with filler + // positions masked to zero after the pe add and after every block (the + // GRN norm and depthwise context depend on it). Pad to N frames up front. + core::TensorValue te_mask_on; // [1, N, 1] 1.0 on real cols, 0 on pads + if (NT < N) { + const auto pad_shape = core::TensorShape::from_dims({2, N - NT, kTextDim}); + std::vector zv(static_cast(pad_shape.num_elements()), 0.0F); + te = mod::ConcatModule({1}).build( + ctx, te, ctx_store_f32(ctx, pad_shape, zv)); + // mask: ones [2, NT, 1] concat zeros [2, N-NT, 1] + std::vector ones(static_cast(2 * NT), 1.0F); + std::vector zeros(static_cast(2 * (N - NT)), 0.0F); + auto m_on = ctx_store_f32(ctx, core::TensorShape::from_dims({2, NT, 1}), ones); + auto m_off = ctx_store_f32(ctx, core::TensorShape::from_dims({2, N - NT, 1}), zeros); + te_mask_on = mod::ConcatModule({1}).build(ctx, m_on, m_off); + te_mask_on = mod::RepeatModule( + {core::TensorShape::from_dims({2, N, kTextDim})}).build(ctx, te_mask_on); + } else { + // no padding: identity mask (skip the mul entirely below) + te_mask_on = core::TensorValue{}; + } + // apply the mask right after the pe add (zero pad columns) + if (te_mask_on.tensor != nullptr) { + te = mod::MulModule().build(ctx, te, te_mask_on); + } + + // text ConvNeXt x4 (batch-aware: dwconv input [B, C, T]) + tap_stage("txt_in", te); + for (int bi = 0; bi < 4; ++bi) { + const auto & B = w.text_blocks[static_cast(bi)]; + auto te_c2 = mod::TransposeModule({{0, 2, 1}, 3}).build(ctx, te); // [B, C, T] + // depthwise over batch: module requires rank-3 [B, C, T] — supported + auto dw = mod::DepthwiseConv1dModule({kTextDim, 7, 1, 3, 1, true}).build(ctx, te_c2, B.dwconv); + auto dw_t = mod::TransposeModule({{0, 2, 1}, 3}).build(ctx, dw); // [B, T, C] + auto nrm = mod::LayerNormModule({kTextDim, 1e-6F, true, true}).build(ctx, dw_t, B.norm); + auto h1 = mod::LinearModule({kTextDim, kDim, true}).build(ctx, nrm, B.pw1); + h1 = mod::GeluModule({mod::GeluApproximation::ExactErf}).build(ctx, h1); + core::TensorValue g_dummy, b_dummy; + auto grn_out = grn(ctx, h1, B.grn_gamma, B.grn_beta, g_dummy, b_dummy); + auto h2 = mod::LinearModule({kDim, kTextDim, true}).build(ctx, grn_out, B.pw2); + te = mod::AddModule().build(ctx, te, h2); + if (te_mask_on.tensor != nullptr) { + te = mod::MulModule().build(ctx, te, te_mask_on); // re-zero pads + } + tap_stage_cond(bi == 0, "txt_h1", h1); + tap_stage_cond(bi == 0, "txt_grn", grn_out); + tap_stage_cond(bi == 0, "txt_after_block", te); + } + + // text is already exactly N frames (padded + masked before the blocks) + const auto & te_pad = te; + + // input embed + auto cat0 = mod::ConcatModule({2}).build(ctx, io.x, io.cond); + auto cat1 = mod::ConcatModule({2}).build(ctx, cat0, te_pad); // [2, N, 712] + tap_stage("te", te); // [2, NT, 512] before pad (post-convnext) + tap_stage("te_pad", te_pad); + auto inp = mod::LinearModule({712LL, kDim, true}).build(ctx, cat1, w.input_proj); + tap_stage("inp_proj", inp); + + // CPE: grouped conv per half (B=1) — folding the batch into the conv's + // time axis would bleed the zero-padding across the batch seam; run each + // half separately and concat back on the batch axis. + { + auto conv_mish_half = [&](const core::TensorValue & x_nd, // [N, D] rows + const core::TensorValue & cweight, + const core::TensorValue & cbias) -> core::TensorValue { + auto b1 = core::reshape_tensor( + ctx, core::ensure_backend_addressable_layout(ctx, x_nd), + core::TensorShape::from_dims({1, x_nd.shape.dims[0], x_nd.shape.dims[1]})); + auto x_c = mod::TransposeModule({{0, 2, 1}, 3}).build(ctx, b1); // [1, D, N] + auto x_cc = core::wrap_tensor(ggml_cont(ctx.ggml, x_c.tensor), x_c.shape, GGML_TYPE_F32); + auto r = grouped_conv1d(ctx, x_cc, cweight, cbias, 16); // [N, D] + auto sp = exp_log_softplus(ctx, r); + return mod::MulModule().build(ctx, r, mod::TanhModule().build(ctx, sp)); + }; + // split halves, conv each, concat + auto half = [&](int64_t b) { + auto rows = core::reshape_tensor( + ctx, core::ensure_backend_addressable_layout(ctx, inp), + core::TensorShape::from_dims({2 * N, kDim})); + return mod::SliceModule({0, b * N, N}).build(ctx, rows); + }; + auto r0_c = conv_mish_half(half(0), w.cpe0.weight, *w.cpe0.bias); + auto r0_u = conv_mish_half(half(1), w.cpe0.weight, *w.cpe0.bias); + auto r0 = mod::ConcatModule({0}).build(ctx, r0_c, r0_u); // [2N, D] + auto slice_of = [&](const core::TensorValue & rows, int64_t b) { + return mod::SliceModule({0, b * N, N}).build(ctx, rows); + }; + auto r1_c = conv_mish_half(slice_of(r0, 0), w.cpe2.weight, *w.cpe2.bias); + auto r1_u = conv_mish_half(slice_of(r0, 1), w.cpe2.weight, *w.cpe2.bias); + auto r1 = mod::ConcatModule({0}).build(ctx, r1_c, r1_u); // [2N, D] + auto r1_b = core::reshape_tensor( + ctx, core::ensure_backend_addressable_layout(ctx, r1), + core::TensorShape::from_dims({2, N, kDim})); + inp = mod::AddModule().build(ctx, inp, r1_b); + tap_stage("inp", inp); + } + + // time embedding (shared across halves) + auto t0 = mod::LinearModule({256, kDim, true}).build(ctx, io.time_input, w.time0); + t0 = mod::SiluModule().build(ctx, t0); + auto t_emb = mod::LinearModule({kDim, kDim, true}).build(ctx, t0, w.time2); // [1, 1024] + + core::TensorValue positions; + { + std::vector pos(static_cast(N)); + for (int64_t i = 0; i < N; ++i) pos[static_cast(i)] = static_cast(i); + auto p = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({N})); + ggml_set_input(p.tensor); + if (p.tensor->data != nullptr) { + std::memcpy(p.tensor->data, pos.data(), pos.size() * sizeof(int32_t)); + } else if (t_const_stage != nullptr) { + const auto * b = reinterpret_cast(pos.data()); + t_const_stage->push_back({p.tensor, std::vector(b, b + pos.size() * sizeof(int32_t))}); + } + positions = p; + } + + // 22 DiT blocks: identical to the B=1 graph; all modules are batch-aware + auto h = inp; + for (int bi = 0; bi < 22; ++bi) { + const auto & B = w.blocks[static_cast(bi)]; + auto emb6 = mod::LinearModule({kDim, 6 * kDim, true}).build( + ctx, mod::SiluModule().build(ctx, t_emb), B.attn_norm); // [1, 6144] + auto shift_msa = mod::SliceModule({1, 0 * kDim, kDim}).build(ctx, emb6); + auto scale_msa = mod::SliceModule({1, 1 * kDim, kDim}).build(ctx, emb6); + auto gate_msa = mod::SliceModule({1, 2 * kDim, kDim}).build(ctx, emb6); + auto shift_mlp = mod::SliceModule({1, 3 * kDim, kDim}).build(ctx, emb6); + auto scale_mlp = mod::SliceModule({1, 4 * kDim, kDim}).build(ctx, emb6); + auto gate_mlp = mod::SliceModule({1, 5 * kDim, kDim}).build(ctx, emb6); + (void)shift_msa; + + auto norm = modulate(ctx, mod::LayerNormModule({kDim, 1e-6F, false, false}).build(ctx, h, mod::NormWeights{}), scale_msa, shift_msa); + auto q = mod::LinearModule({kDim, kDim, true}).build(ctx, norm, B.to_q); + auto k = mod::LinearModule({kDim, kDim, true}).build(ctx, norm, B.to_k); + auto v = mod::LinearModule({kDim, kDim, true}).build(ctx, norm, B.to_v); + auto to_heads = [&](core::TensorValue t) { + return core::reshape_tensor( + ctx, core::ensure_backend_addressable_layout(ctx, t), + core::TensorShape::from_dims({t.shape.dims[0], t.shape.dims[1], kHeads, kHeadDim})); + }; + q = to_heads(q); k = to_heads(k); v = to_heads(v); + q = mod::RoPEModule({kHeadDim, GGML_ROPE_TYPE_NORMAL, 10000.0F}).build(ctx, q, positions); + k = mod::RoPEModule({kHeadDim, GGML_ROPE_TYPE_NORMAL, 10000.0F}).build(ctx, k, positions); + auto q_heads = mod::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, q); + auto k_heads = mod::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, k); + auto v_heads = mod::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, v); + // flash-attn requires dense, contiguous q/k/v: materialize the + // strided transposes explicitly (a strided input corrupts arena + // neighbors on replay — root cause of the noise regression) + auto dense4 = [&](const core::TensorValue & t) { + return core::wrap_tensor( + ggml_cont(ctx.ggml, t.tensor), t.shape, GGML_TYPE_F32); + }; + q_heads = dense4(q_heads); + k_heads = dense4(k_heads); + v_heads = dense4(v_heads); + auto attn = mod::ScaledDotProductAttentionModule({ + kHeadDim, + mod::ScaledDotProductAttentionLowering::Flash, + GGML_PREC_F32, + mod::AttentionCausality::NonCausal, + }).build(ctx, q_heads, k_heads, v_heads); + tap_stage_cond(bi == 0, "block0.attn", attn); + tap_stage_cond(bi == 21, "block21.attn", attn); // [2, N, H, DH] + auto attn_flat = core::reshape_tensor( + ctx, core::ensure_backend_addressable_layout(ctx, attn), + core::TensorShape::from_dims({attn.shape.dims[0], attn.shape.dims[1], kDim})); + auto proj = mod::LinearModule({kDim, kDim, true}).build(ctx, attn_flat, B.to_out); + h = mod::AddModule().build( + ctx, h, mod::MulModule().build( + ctx, proj, mod::RepeatModule({proj.shape}).build(ctx, lift_row(ctx, gate_msa)))); + + auto norm2 = modulate(ctx, mod::LayerNormModule({kDim, 1e-6F, false, false}).build(ctx, h, mod::NormWeights{}), scale_mlp, shift_mlp); + auto f1 = mod::LinearModule({kDim, 2048, true}).build(ctx, norm2, B.ff0); + f1 = mod::GeluModule({mod::GeluApproximation::Tanh}).build(ctx, f1); + auto f2 = mod::LinearModule({2048, kDim, true}).build(ctx, f1, B.ff2); + h = mod::AddModule().build( + ctx, h, mod::MulModule().build( + ctx, f2, mod::RepeatModule({f2.shape}).build(ctx, lift_row(ctx, gate_mlp)))); + } + + { + auto emb2 = mod::LinearModule({kDim, 2 * kDim, true}).build( + ctx, mod::SiluModule().build(ctx, t_emb), w.norm_out); + auto scale = mod::SliceModule({1, 0, kDim}).build(ctx, emb2); + auto shift = mod::SliceModule({1, kDim, kDim}).build(ctx, emb2); + auto norm = modulate(ctx, mod::LayerNormModule({kDim, 1e-6F, false, false}).build(ctx, h, mod::NormWeights{}), scale, shift); + io.output = mod::LinearModule({kDim, kMel, true}).build(ctx, norm, w.proj_out); // [2, N, 100] + } + return io; +} + +} // namespace engine::models::f5_tts diff --git a/src/community_models/f5_tts/runtime.cpp b/src/community_models/f5_tts/runtime.cpp new file mode 100644 index 00000000..ee844c9d --- /dev/null +++ b/src/community_models/f5_tts/runtime.cpp @@ -0,0 +1,709 @@ +#include "engine/community_models/f5_tts/runtime.h" + +#include "cpu_graph_compute.h" + +#include "engine/community_models/f5_tts/dit_modules.h" +#include "engine/community_models/f5_tts/weights.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/module.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/weight_binding.h" + +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::f5_tts { +namespace { + +namespace core = engine::core; +namespace modules = engine::modules; + +// Column convention throughout: tensors are [features, seq] (ggml ne0 = +// features), so ggml_norm == LayerNorm over features and ggml_mul_mat(W, X) +// matches torch X @ W^T with torch weights stored [out, in]. + +struct F5Linear { + core::TensorValue weight; + core::TensorValue bias; +}; + +struct F5ConvNeXt { + modules::DepthwiseConv1dWeights dwconv; + core::TensorValue norm_w, norm_b; + F5Linear pw1, pw2; + std::vector grn_gamma, grn_beta; +}; + +struct F5Block { + F5Linear attn_norm; // -> 6*dim + F5Linear to_q, to_k, to_v, to_out; + F5Linear ff0, ff2; +}; + +struct F5Weights { + std::shared_ptr store; + core::TensorValue text_embedding; // [512, 2731] + F5Linear input_proj; + modules::Conv1dWeights cpe0, cpe2; // grouped k31 g16 + F5Linear time0, time2; + std::vector inv_freq; + std::vector text_blocks; + std::vector blocks; + F5Linear norm_out; + F5Linear proj_out; +}; + +struct LoadedModel { + F5Weights w; + F5Architecture arch; + ggml_backend_t backend = nullptr; // backend owning the weight store + core::BackendType backend_type = core::BackendType::Cpu; +}; + +struct BackendOwner { + ggml_backend_t value = nullptr; + // Intentionally never freed: CUDA backends must be released before the + // driver shuts down at static destruction, which we cannot order reliably. + // Leaking a few backends at process exit is harmless. + ~BackendOwner() = default; +}; + +F5Weights load_weights( + const engine::assets::TensorSource & source, + ggml_backend_t backend, + core::BackendType backend_type, + bool fp16_linears) { + F5Weights w; + w.store = std::make_shared( + backend, backend_type, "f5_tts.weights", 2ULL * 1024ULL * 1024ULL * 1024ULL); + const auto f32 = [&](const std::string & n) { + return w.store->load_f32_tensor( + source, n, source.require_metadata(n).shape); + }; + // Linear weights consumed only by ggml_mul_mat: FP16 storage halves GEMM + // time on tensor cores (mul_mat returns F32 regardless of weight dtype). + // Everything else (embeddings consumed by get_rows, dwconv kernels read + // on host, biases/norms broadcast) stays F32. + const auto gemm_w = [&](const std::string & n) { + if (fp16_linears) { + return w.store->load_tensor( + source, n, engine::assets::TensorStorageType::F16, + source.require_metadata(n).shape); + } + return f32(n); + }; + const auto lin_f32bias = [&](const std::string & n) { + return F5Linear{gemm_w(n + ".weight"), f32(n + ".bias")}; + }; + + w.text_embedding = f32("text_embed.text_embed.weight"); // get_rows consumer: keep F32 + w.input_proj = lin_f32bias("input_embed.proj"); + w.cpe0.weight = gemm_w("input_embed.conv_pos_embed.conv1d.0.weight"); + w.cpe0.bias = f32("input_embed.conv_pos_embed.conv1d.0.bias"); + w.cpe2.weight = gemm_w("input_embed.conv_pos_embed.conv1d.2.weight"); + w.cpe2.bias = f32("input_embed.conv_pos_embed.conv1d.2.bias"); + w.time0 = lin_f32bias("time_embed.time_mlp.0"); + w.time2 = lin_f32bias("time_embed.time_mlp.2"); + w.inv_freq = source.require_f32("rotary_embed.inv_freq"); + + w.text_blocks.reserve(4); + for (int i = 0; i < 4; ++i) { + const std::string p = "text_embed.text_blocks." + std::to_string(i); + F5ConvNeXt b; + b.dwconv.weight = f32(p + ".dwconv.weight"); // host-side kernel read + b.dwconv.bias = f32(p + ".dwconv.bias"); + b.norm_w = f32(p + ".norm.weight"); + b.norm_b = f32(p + ".norm.bias"); + b.pw1 = lin_f32bias(p + ".pwconv1"); + b.pw2 = lin_f32bias(p + ".pwconv2"); + b.grn_gamma = source.require_f32(p + ".grn.gamma"); + b.grn_beta = source.require_f32(p + ".grn.beta"); + w.text_blocks.push_back(std::move(b)); + } + w.blocks.reserve(22); + for (int i = 0; i < 22; ++i) { + const std::string p = "transformer_blocks." + std::to_string(i); + F5Block b; + b.attn_norm = lin_f32bias(p + ".attn_norm.linear"); + b.to_q = lin_f32bias(p + ".attn.to_q"); + b.to_k = lin_f32bias(p + ".attn.to_k"); + b.to_v = lin_f32bias(p + ".attn.to_v"); + b.to_out = lin_f32bias(p + ".attn.to_out.0"); + b.ff0 = lin_f32bias(p + ".ff.ff.0.0"); + b.ff2 = lin_f32bias(p + ".ff.ff.2"); + w.blocks.push_back(std::move(b)); + } + w.norm_out = lin_f32bias("norm_out.linear"); + w.proj_out = lin_f32bias("proj_out"); + w.store->upload(); + source.release_storage(); + return w; +} + +// write a small view that strips that dotted prefix. +class StrippedView final : public engine::assets::TensorSource { +public: + static constexpr std::string_view kPrefix = "ema_model.transformer."; + + explicit StrippedView(std::shared_ptr inner) + : inner_(std::move(inner)) { + for (const auto & t : inner_->tensors()) { + if (t.name.rfind(kPrefix, 0) == 0) { + routes_.emplace(t.name.substr(kPrefix.size()), t.name); + } + } + if (routes_.empty()) { + throw std::runtime_error("F5 checkpoint has no ema_model.transformer.* tensors"); + } + } + const std::filesystem::path & source_path() const noexcept override { + return inner_->source_path(); + } + bool has_tensor(std::string_view name) const noexcept override { + return routes_.find(std::string(name)) != routes_.end(); + } + engine::assets::TensorMetadata require_metadata(std::string_view name) const override { + auto m = inner_->require_metadata(require(name)); + m.name = std::string(name); + return m; + } + std::vector tensors() const override { + std::vector out; + out.reserve(routes_.size()); + for (const auto & [n, _] : routes_) { + out.push_back(require_metadata(n)); + } + return out; + } + void release_storage() const override { inner_->release_storage(); } + engine::assets::RawTensorData require_tensor_data(std::string_view name) const override { + auto d = inner_->require_tensor_data(require(name)); + d.metadata.name = std::string(name); + return d; + } + std::vector require_f32( + std::string_view name, + const std::optional> & expected_shape) const override { + return inner_->require_f32(require(name), expected_shape); + } + std::optional> optional_f32( + std::string_view name, + const std::optional> & expected_shape) const override { + const auto found = routes_.find(std::string(name)); + if (found == routes_.end()) { + return std::nullopt; + } + return inner_->optional_f32(found->second, expected_shape); + } + int64_t require_i64_scalar(std::string_view name) const override { + return inner_->require_i64_scalar(require(name)); + } + +private: + const std::string & require(std::string_view name) const { + const auto found = routes_.find(std::string(name)); + if (found == routes_.end()) { + throw std::runtime_error("F5 missing tensor: " + std::string(name)); + } + return found->second; + } + std::shared_ptr inner_; + std::unordered_map routes_; +}; + +const LoadedModel & load_model_once(const std::string & path, const F5ComputeDevice & dev); + +// Open the DiT checkpoint as a stripped-name tensor source. Safetensors +// checkpoints carry raw torch names ("ema_model.transformer.*"); GGUF +// packages store the same tensors under the "transformer" namespace, so the +// namespace view is stripped first and both formats converge on the same +// logical names. +std::shared_ptr open_dit_source(const std::string & path) { + auto source = engine::assets::open_tensor_source(path); + std::shared_ptr base = std::move(source); + if (std::filesystem::path(path).extension() == ".gguf") { + base = engine::assets::make_prefixed_tensor_source(base, "transformer"); + } + return std::make_shared(std::move(base)); +} + +// Module-typed DiT weights cache (per path + device). Leaked at exit like +// the raw-weight cache: CUDA buffers cannot be freed after driver shutdown. +const F5DiTWeights & load_dit_weights_once( + const std::string & path, const F5ComputeDevice & dev) { + struct Entry { F5DiTWeights w; }; + static auto * cache = new std::unordered_map(); + const std::string key = + (dev.use_cuda ? "cuda" + std::to_string(dev.device) : "cpu") + ":" + path; + const auto found = cache->find(key); + if (found != cache->end()) { + return found->second.w; + } + const auto & model = load_model_once(path, dev); + auto stripped = open_dit_source(path); + Entry entry; + entry.w = load_dit_weights(*stripped, model.backend, model.backend_type); + return cache->emplace(key, std::move(entry)).first->second.w; +} + +const LoadedModel & load_model_once(const std::string & path, const F5ComputeDevice & dev) { + // Deliberately leaked at exit (never destroyed): CUDA weight buffers must + // be freed before driver shutdown; static destruction order cannot + // guarantee that. Model caches are process-lifetime anyway. + static auto * cache = new std::unordered_map(); + // The backend (and its weight buffer) must outlive the cache entry, so it + // is owned by a static owner freed after the cache at exit. + static std::vector> owners; + const std::string key = (dev.use_cuda ? "cuda" + std::to_string(dev.device) : "cpu") + ":" + path; + if (const auto found = cache->find(key); found != cache->end()) { + return found->second; + } + auto stripped = open_dit_source(path); + auto owner = std::make_unique(); + const core::BackendType type = dev.use_cuda ? core::BackendType::Cuda : core::BackendType::Cpu; + core::BackendConfig cfg{type, dev.use_cuda ? dev.device : 0, dev.use_cuda ? 1 : std::max(1, dev.threads)}; + owner->value = core::init_backend(cfg); + if (!dev.use_cuda) { + core::set_backend_threads(owner->value, std::max(1, dev.threads)); + } + LoadedModel model; + model.arch = F5Architecture{}; + model.backend = owner->value; + model.backend_type = type; + // FP16 linear weights on CUDA only (CPU mul_mat with F16 weights is slow + // via the fallback path; CUDA hits tensor cores). + model.w = load_weights(*stripped, owner->value, type, dev.use_cuda && dev.fp16_weights); + owners.push_back(std::move(owner)); + return cache->emplace(key, std::move(model)).first->second; +} + +} // namespace + +// Batched CFG forward: one graph, ne3=2 batch (half 0 = conditioned with +// text_ids, half 1 = uncond: filler text ids + zeroed cond, uploaded from +// the host — matches python cfg_infer drop_audio_cond=True/drop_text=True). +// Same per-half math as two f5_dit_forward calls; halves share +// weights/time-embed/positions. +// Returns {cond, null} mel-major [MEL*N] each. +std::pair, std::vector> f5_dit_forward_cfg( + const std::string & weights_path, + const std::vector & x_in, + const std::vector & cond_in, + const std::vector & text_in, + float time_value, + int seq_len, + const F5Architecture & arch, + const F5ComputeDevice * device) { + static const F5ComputeDevice kDefaultDevice{}; + const F5ComputeDevice & dev = device != nullptr ? *device : kDefaultDevice; + const auto & model = load_model_once(weights_path, dev); + const int N = seq_len; + const int MEL = arch.mel_dim; + const int NT = static_cast(text_in.size()); + const bool is_cuda = model.backend_type == core::BackendType::Cuda; + + struct CfgGraph { + ggml_context * ctx = nullptr; + ggml_cgraph * graph = nullptr; + ggml_gallocr_t gallocr = nullptr; + ggml_backend_buffer_t io_buffer = nullptr; + ggml_tensor * x = nullptr; // [MEL, N, 1, 2] + ggml_tensor * cond = nullptr; // [MEL, N, 1, 2] + ggml_tensor * text_ids = nullptr; // [NT, 1, 2, 1] (per-half ids) + ggml_tensor * th_t = nullptr; // [256, 1] + ggml_tensor * output = nullptr; // [MEL, N] (half 0, conditioned) + ggml_tensor * out_u = nullptr; // [MEL, N] (half 1, uncond) + }; + static auto * cache = new std::map, std::unique_ptr>(); + const auto ckey = std::make_tuple(&model, N, NT); + auto it = cache->find(ckey); + if (it == cache->end()) { + auto gnew = std::make_unique(); + const size_t ctx_bytes = std::min( + std::max(1536ULL << 20, static_cast(N) * (8ULL << 20)), + 12288ULL << 20); + gnew->ctx = ggml_init({ctx_bytes, nullptr, is_cuda}); + ggml_context * ctx = gnew->ctx; + // ---- module-composed batched-CFG graph (B=2) ---- + // Leaves: x/cond [2, N, 100] (same values in both halves), ids + // [2*NT] (cond half, then uncond half), time input [1, 256]. + // Memory layout of [2, N, 100] rows == the old 4d (MEL, N, 1, 2) + // interleaved halves? No: halves are now batch-major. The per-call + // upload path below fills each half explicitly. + const F5DiTWeights & dit_w = load_dit_weights_once(weights_path, dev); + std::vector * cfg_staged = nullptr; + auto * stage = const_stage_begin(); + auto io = build_dit_cfg_modules_graph( + ctx, dit_w, arch, N, NT, model.backend_type); + const_stage_end(stage); + cfg_staged = stage; + // Mark I/O explicitly: without ggml_set_input/output the gallocr + // treats leaves as scratch and may clobber them across replays. + ggml_set_input(io.x.tensor); + ggml_set_input(io.cond.tensor); + ggml_set_input(io.text_ids.tensor); + ggml_set_input(io.time_input.tensor); + ggml_set_output(io.output.tensor); + ggml_tensor * output = io.output.tensor; // [2, N, 100] + ggml_tensor * th_t = io.time_input.tensor; + gnew->output = output; // cond half via view at read time + gnew->out_u = output; // uncond half: same tensor, offset read + gnew->x = io.x.tensor; + gnew->cond = io.cond.tensor; + gnew->text_ids = io.text_ids.tensor; + gnew->th_t = th_t; + gnew->graph = ggml_new_graph_custom(ctx, 262144, false); + ggml_build_forward_expand(gnew->graph, output); + core::validate_backend_graph_supported(model.backend, gnew->graph, "f5_dit_cfg"); + if (!is_cuda) { + const int threads = dev.threads > 0 ? dev.threads : static_cast(std::thread::hardware_concurrency()); + core::set_backend_threads(model.backend, threads); + } + if (is_cuda) { + // gallocr-only flow; constants first get PRIVATE buffers so the + // arena never aliases them (root cause of the noise regression) + const_stage_bind(cfg_staged, model.backend); + gnew->gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); + if (gnew->gallocr == nullptr || !ggml_gallocr_reserve(gnew->gallocr, gnew->graph) || + !ggml_gallocr_alloc_graph(gnew->gallocr, gnew->graph)) { + throw std::runtime_error("F5 DiT CUDA graph alloc failed"); + } + if (cfg_staged != nullptr) { + const_stage_upload(cfg_staged, is_cuda ? model.backend : nullptr); + const_stage_end(cfg_staged); + } + (void)0; + } + it = cache->emplace(ckey, std::move(gnew)).first; + // bound the graph cache: each entry holds a CUDA io buffer + gallocr + // arena (~1-5 GiB at N~1000). Keep at most 2; evict the oldest + // (map order == insertion order), never the entry just added. + while (cache->size() > 2) { + bool evicted = false; + for (auto cit = cache->begin(); cit != cache->end(); ++cit) { + if (cit->first != ckey) { + cache->erase(cit); + evicted = true; + break; + } + } + if (!evicted) break; + } + } + CfgGraph & g = *it->second; + + // ---- per-call uploads: x/cond duplicated to both halves, ids per half ---- + { + const size_t half_bytes = static_cast(N) * MEL * sizeof(float); + std::vector xb(x_in.size() * 2); + std::memcpy(xb.data(), x_in.data(), half_bytes); + std::memcpy(xb.data() + x_in.size(), x_in.data(), half_bytes); + // Python cfg_infer: the uncond half runs with drop_audio_cond=True + // (cond = zeros) and drop_text=True (text zeroed AFTER the +1 offset, + // so the embedding sees the filler row 0, not space row 1). + std::vector cb(cond_in.size() * 2, 0.0F); + std::memcpy(cb.data(), cond_in.data(), half_bytes); // cond half only + std::vector ids(NT * 2); + for (int i = 0; i < NT; ++i) { + ids[i] = text_in[i] + 1; // cond half + ids[NT + i] = 0; // uncond: drop_text zeros -> filler token 0 + } + std::vector th(256); + { + const float log_base = std::log(10000.0F) / 127.0F; + for (int i = 0; i < 128; ++i) { + const float f = 1000.0F * time_value * std::exp(-log_base * i); + th[i] = std::sin(f); + th[128 + i] = std::cos(f); + } + } + if (is_cuda) { + ggml_backend_tensor_set(g.x, xb.data(), 0, xb.size() * sizeof(float)); + ggml_backend_tensor_set(g.cond, cb.data(), 0, cb.size() * sizeof(float)); + ggml_backend_tensor_set(g.text_ids, ids.data(), 0, ids.size() * sizeof(int32_t)); + ggml_backend_tensor_set(g.th_t, th.data(), 0, th.size() * sizeof(float)); + } else { + std::memcpy(g.x->data, xb.data(), xb.size() * sizeof(float)); + std::memcpy(g.cond->data, cb.data(), cb.size() * sizeof(float)); + std::memcpy(g.text_ids->data, ids.data(), ids.size() * sizeof(int32_t)); + std::memcpy(g.th_t->data, th.data(), th.size() * sizeof(float)); + } + } + + // ---- compute + read both halves ---- + + + + const auto status = is_cuda + ? core::compute_backend_graph(model.backend, g.graph, nullptr, "f5_dit_cfg") + : f5_cpu_graph_compute(g.ctx, g.graph, + dev.threads > 0 ? dev.threads : static_cast(std::thread::hardware_concurrency())); + if (is_cuda) ggml_backend_synchronize(model.backend); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("F5 DiT CFG graph compute failed"); + } + std::pair, std::vector> out; + // single [2, N, 100] output tensor, batch-major halves + const size_t half_floats = ggml_nelements(g.output) / 2; + out.first.resize(half_floats); + out.second.resize(half_floats); + if (is_cuda) { + ggml_backend_tensor_get(g.output, out.first.data(), 0, half_floats * sizeof(float)); + ggml_backend_tensor_get(g.output, out.second.data(), half_floats * sizeof(float), half_floats * sizeof(float)); + } else { + const float * all = reinterpret_cast(ggml_get_data(g.output)); + std::memcpy(out.first.data(), all, half_floats * sizeof(float)); + std::memcpy(out.second.data(), all + half_floats, half_floats * sizeof(float)); + } + return out; +} + +} // namespace engine::models::f5_tts + + +namespace engine::models::f5_tts { + + +std::vector f5_dit_forward( + const std::string & weights_path, + const std::vector & x_in, + const std::vector & cond_in, + const std::vector & text_in, + float time_value, + int seq_len, + const F5Architecture & arch, + bool drop_audio_cond, + bool drop_text, + const F5DebugTaps * taps, + const F5ComputeDevice * device) { + static const F5ComputeDevice kDefaultDevice{}; + const F5ComputeDevice & dev = device != nullptr ? *device : kDefaultDevice; + const auto & model = load_model_once(weights_path, dev); + const int N = seq_len; + const int MEL = arch.mel_dim; + const int NT = static_cast(text_in.size()); + const bool is_cuda = model.backend_type == core::BackendType::Cuda; + + // ---- cached graph per (model, N, NT, with/without taps) ---- + // Taps change the graph (extra roots), so key on their presence. The + // per-call leaves (x, cond, text_ids, th) are uploaded each invocation; + // everything else (pos ids, pe, ones, constants) is uploaded once. + struct DiTGraph { + ggml_context * ctx = nullptr; + ggml_cgraph * graph = nullptr; + ggml_gallocr_t gallocr = nullptr; + ggml_backend_buffer_t io_buffer = nullptr; // CUDA leaves buffer + core::HostGraphPlan host_plan; // CPU plan reuse + ggml_tensor * x = nullptr; + ggml_tensor * cond = nullptr; + ggml_tensor * text_ids = nullptr; + ggml_tensor * th_t = nullptr; + ggml_tensor * output = nullptr; + ggml_tensor * tap_text_embed = nullptr; + ggml_tensor * tap_text_convnext = nullptr; + ggml_tensor * tap_text_padded = nullptr; + ggml_tensor * tap_input_embed = nullptr; + ggml_tensor * tap_time_embed = nullptr; + ggml_tensor * tap_block0 = nullptr; + ggml_tensor * tap_block21 = nullptr; + ~DiTGraph() { + // Leaked by design when cached (freed only on program-graph reset); + // destroyed here only when construction throws mid-build. + if (gallocr != nullptr) ggml_gallocr_free(gallocr); + if (io_buffer != nullptr) ggml_backend_buffer_free(io_buffer); + if (ctx != nullptr) ggml_free(ctx); + } + }; + + // Graph cache is process-lifetime (CUDA buffers cannot be safely freed + // after driver shutdown in static destruction). + static auto * graph_cache = + new std::map, std::unique_ptr>(); + const bool want_taps = taps != nullptr; + const auto cache_key = std::make_tuple(&model, N, NT, want_taps); + auto it = graph_cache->find(cache_key); + if (it == graph_cache->end()) { + auto gnew = std::make_unique(); + const size_t ctx_bytes = std::min( + std::max(1536ULL << 20, static_cast(N) * (6ULL << 20)), + 6144ULL << 20); + gnew->ctx = ggml_init({ctx_bytes, nullptr, is_cuda}); + ggml_context * ctx = gnew->ctx; + // On CUDA the ctx is no_alloc: leaf tensors get device storage after + // ggml_backend_alloc_ctx_tensors, values uploaded from staging vectors. + (void)MEL; + ggml_tensor * output = nullptr; + // NOTE: stage taps are not wired in the module-composed graph; the parity + // harness compares the final output (and uses the raw path where needed). + ggml_tensor * tap_text_embed = nullptr; + ggml_tensor * tap_text_convnext = nullptr; + ggml_tensor * tap_text_padded = nullptr; + ggml_tensor * tap_input_embed = nullptr; + ggml_tensor * tap_time_embed = nullptr; + ggml_tensor * tap_block0 = nullptr; + ggml_tensor * tap_block21 = nullptr; + // ---- module-composed graph build (dev-branch reviewer pattern) ---- + // All ops expressed via framework modules; leaves are [1, N, 100] + // row-major (same memory as the old mel-major columns) so the upload + // and readback paths are unchanged. + const F5DiTWeights & dit_w = load_dit_weights_once(weights_path, dev); + std::vector * staged_module_consts = nullptr; + auto * stage = const_stage_begin(); + auto io = build_dit_modules_graph( + ctx, dit_w, arch, N, NT, model.backend_type); + const_stage_end(stage); + // staged constants upload after ggml_backend_alloc_ctx_tensors below + staged_module_consts = stage; + ggml_set_input(io.x.tensor); + ggml_set_input(io.cond.tensor); + ggml_set_input(io.text_ids.tensor); + ggml_set_input(io.time_input.tensor); + ggml_set_output(io.output.tensor); + output = io.output.tensor; + gnew->x = io.x.tensor; + gnew->cond = io.cond.tensor; + gnew->text_ids = io.text_ids.tensor; + gnew->th_t = io.time_input.tensor; + gnew->output = output; + gnew->graph = ggml_new_graph_custom(ctx, 262144, false); + ggml_build_forward_expand(gnew->graph, output); + for (ggml_tensor * tap : + {tap_text_embed, tap_text_convnext, tap_text_padded, tap_input_embed, + tap_time_embed, tap_block0, tap_block21}) { + if (tap != nullptr) { + ggml_build_forward_expand(gnew->graph, tap); + } + } + core::validate_backend_graph_supported(model.backend, gnew->graph, "f5_dit"); + if (!is_cuda) { + const int threads = dev.threads > 0 + ? dev.threads + : static_cast(std::thread::hardware_concurrency()); + core::set_backend_threads(model.backend, threads); + } + if (is_cuda) { + const_stage_bind(staged_module_consts, model.backend); + gnew->gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); + if (gnew->gallocr == nullptr || !ggml_gallocr_reserve(gnew->gallocr, gnew->graph) || + !ggml_gallocr_alloc_graph(gnew->gallocr, gnew->graph)) { + throw std::runtime_error("F5 DiT CUDA graph alloc failed"); + } + if (staged_module_consts != nullptr) { + const_stage_upload(staged_module_consts, is_cuda ? model.backend : nullptr); + const_stage_end(staged_module_consts); + staged_module_consts = nullptr; + } + } else { + if (staged_module_consts != nullptr) { + const_stage_end(staged_module_consts); + staged_module_consts = nullptr; + } + // CPU: graph data tensors live in the ctx pool (inline alloc); + // ggml_graph_compute_with_ctx replays without realloc. + } + it = graph_cache->emplace(cache_key, std::move(gnew)).first; + } // if build + DiTGraph & g = *it->second; + + // ---- per-call leaf uploads ---- + { + // x/cond arrive mel-major [N][mel]; ggml [MEL, N] wants (m, n) at + // n * MEL + m -- identical layout, direct copy. drop_audio_cond + // zeroes cond on the host before upload. + std::vector cond_col; + const float * cond_src = cond_in.data(); + if (drop_audio_cond) { + cond_col.assign(cond_in.size(), 0.0F); + cond_src = cond_col.data(); + } + if (is_cuda) { + ggml_backend_tensor_set(g.x, x_in.data(), 0, x_in.size() * sizeof(float)); + ggml_backend_tensor_set(g.cond, cond_src, 0, cond_in.size() * sizeof(float)); + } else { + std::memcpy(g.x->data, x_in.data(), x_in.size() * sizeof(float)); + std::memcpy(g.cond->data, cond_src, cond_in.size() * sizeof(float)); + } + std::vector ids(NT); + for (int i = 0; i < NT; ++i) { + ids[i] = drop_text ? 0 : (text_in[i] + 1); + } + if (is_cuda) { + ggml_backend_tensor_set(g.text_ids, ids.data(), 0, ids.size() * sizeof(int32_t)); + } else { + std::memcpy(g.text_ids->data, ids.data(), ids.size() * sizeof(int32_t)); + } + std::vector th(256); + { + const float log_base = std::log(10000.0F) / 127.0F; + for (int i = 0; i < 128; ++i) { + const float f = 1000.0F * time_value * std::exp(-log_base * i); + th[i] = std::sin(f); + th[128 + i] = std::cos(f); + } + } + if (is_cuda) { + ggml_backend_tensor_set(g.th_t, th.data(), 0, th.size() * sizeof(float)); + } else { + std::memcpy(g.th_t->data, th.data(), th.size() * sizeof(float)); + } + } + + // ---- compute ---- + std::vector out; + const auto status = is_cuda + ? core::compute_backend_graph(model.backend, g.graph, nullptr, "f5_dit") + : f5_cpu_graph_compute(g.ctx, g.graph, + dev.threads > 0 ? dev.threads + : static_cast(std::thread::hardware_concurrency())); + if (is_cuda) { + ggml_backend_synchronize(model.backend); + } + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("F5 DiT graph compute failed"); + } + out.resize(ggml_nelements(g.output)); + if (is_cuda) { + ggml_backend_tensor_get(g.output, out.data(), 0, out.size() * sizeof(float)); + } else { + std::memcpy(out.data(), ggml_get_data(g.output), out.size() * sizeof(float)); + } + if (taps != nullptr) { + const auto read_tap = [&](ggml_tensor * t, std::vector * dst) { + if (t != nullptr && dst != nullptr) { + dst->resize(ggml_nelements(t)); + if (is_cuda) { + ggml_backend_tensor_get(t, dst->data(), 0, dst->size() * sizeof(float)); + } else { + std::memcpy(dst->data(), ggml_get_data(t), dst->size() * sizeof(float)); + } + } + }; + read_tap(g.tap_text_embed, taps->text_embed); + read_tap(g.tap_text_convnext, taps->text_convnext); + read_tap(g.tap_text_padded, taps->text_padded); + read_tap(g.tap_input_embed, taps->input_embed); + read_tap(g.tap_time_embed, taps->time_embed); + read_tap(g.tap_block0, taps->block0); + read_tap(g.tap_block21, taps->block21); + } + return out; // [MEL * N] mel-major columns: out[m * N + n] +} + + +} // namespace engine::models::f5_tts + diff --git a/src/community_models/f5_tts/session.cpp b/src/community_models/f5_tts/session.cpp new file mode 100644 index 00000000..c847aec1 --- /dev/null +++ b/src/community_models/f5_tts/session.cpp @@ -0,0 +1,241 @@ +#include "engine/community_models/f5_tts/session.h" + +#include "engine/community_models/f5_tts/synthesize.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/runtime/spec_backed_model.h" + +#include +#include +#include +#include +#include +#include + +namespace engine::models::f5_tts { +namespace { + +constexpr const char * kFamily = "f5_tts"; + +const runtime::AudioBuffer * reference_audio(const runtime::TaskRequest & request) { + if (request.voice.has_value() && + request.voice->speaker.has_value() && + request.voice->speaker->audio.has_value()) { + return &*request.voice->speaker->audio; + } + return request.audio_input.has_value() + ? &*request.audio_input + : nullptr; +} + +// Locate the DiT checkpoint inside the model directory: exactly one +// *.safetensors / *.gguf is expected (Habibi Unified/Specialized layout). +// GGUF is preferred when both are present (it is the packaged format). +std::filesystem::path find_checkpoint(const std::filesystem::path & model_path) { + namespace fs = std::filesystem; + if (fs::is_regular_file(model_path)) { + return model_path; // direct path to the checkpoint file + } + std::vector ggufs, safetensors; + for (const auto & entry : fs::directory_iterator(model_path)) { + if (entry.path().extension() == ".gguf") ggufs.push_back(entry.path()); + else if (entry.path().extension() == ".safetensors") safetensors.push_back(entry.path()); + } + std::sort(ggufs.begin(), ggufs.end()); + std::sort(safetensors.begin(), safetensors.end()); + if (!ggufs.empty()) { + return ggufs.back(); // prefer gguf; highest sort key if several + } + if (!safetensors.empty()) { + return safetensors.back(); // highest-numbered model_*.safetensors + } + throw std::runtime_error( + "F5-TTS: no .gguf/.safetensors checkpoint found in " + model_path.string()); +} + +// First tensor-source file (safetensors preferred, then gguf) in a directory. +std::optional find_tensor_file(const std::filesystem::path & dir) { + namespace fs = std::filesystem; + if (!fs::is_directory(dir)) return std::nullopt; + for (const char * ext : {".safetensors", ".gguf"}) { + for (const auto & entry : fs::directory_iterator(dir)) { + if (entry.path().extension() == ext) return entry.path(); + } + } + return std::nullopt; +} + +} // namespace + +std::shared_ptr load_f5_tts_assets( + const std::filesystem::path & model_path) { + auto assets = std::make_shared(); + assets->resources = assets::ResourceBundle(model_path); + assets->checkpoint = find_checkpoint(model_path); + return assets; +} + +F5TTSSession::F5TTSSession( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options, + std::shared_ptr assets, + std::shared_ptr contract) + : task_kind_(task.task), + run_mode_(task.mode), + assets_(std::move(assets)), + contract_(std::move(contract)) { + if (assets_ == nullptr) { + throw std::runtime_error("F5-TTS session requires assets"); + } + if (contract_ == nullptr) { + throw std::runtime_error("F5-TTS session requires a model contract"); + } + // Vocos vocoder resolution order: + // 1. f5_tts.vocos_path session option + // 2. bundled "vocos" namespace inside a GGUF checkpoint + // 3. vocos.safetensors next to the checkpoint + // 4. the vocos-mel-24khz package installed next to the model directory + const auto vocos_opt = runtime::find_option( + options.options, {"f5_tts.vocos_path", "vocos_path"}); + namespace fs = std::filesystem; + if (vocos_opt.has_value()) { + vocos_path_ = *vocos_opt; + } else { + const fs::path ckpt_dir = assets_->checkpoint.parent_path(); + const fs::path models_root = ckpt_dir.parent_path().parent_path(); + if (assets_->checkpoint.extension() == ".gguf") { + const auto probe = assets::open_tensor_source(assets_->checkpoint); + // packed GGUF namespaces are slash-separated (vocos/backbone...) + if (probe->has_tensor("vocos/backbone.embed.weight") || + probe->has_tensor("vocos.backbone.embed.weight")) { + vocos_path_ = assets_->checkpoint.string(); + } + } + if (vocos_path_.empty()) { + const fs::path sibling = ckpt_dir / "vocos.safetensors"; + if (fs::exists(sibling)) { + vocos_path_ = sibling.string(); + } else if (const auto pkg = find_tensor_file(models_root / "vocos-mel-24khz")) { + vocos_path_ = pkg->string(); + } + } + if (vocos_path_.empty()) { + throw std::runtime_error( + "F5-TTS: no vocos vocoder found; install the vocos_mel_24khz " + "package or set session option f5_tts.vocos_path"); + } + } + if (const auto d = runtime::find_option(options.options, {"f5_tts.dialect", "dialect"})) { + dialect_ = *d; + } + if (const auto fb = runtime::find_option(options.options, {"f5_tts.frame_budget", "frame_budget"})) { + frame_budget_ = std::stoi(*fb); + if (frame_budget_ < 256 || frame_budget_ > 8192) { + throw std::runtime_error( + "f5_tts.frame_budget must be within [256, 8192] mel frames"); + } + } + use_cuda_ = options.backend.type == core::BackendType::Cuda; + cuda_device_ = options.backend.device; + threads_ = options.backend.threads; +} + +std::string F5TTSSession::family() const noexcept { + return kFamily; +} + +runtime::VoiceTaskKind F5TTSSession::task_kind() const noexcept { + return task_kind_; +} + +runtime::RunMode F5TTSSession::run_mode() const noexcept { + return run_mode_; +} + +void F5TTSSession::prepare(const runtime::SessionPreparationRequest & request) { + (void) request; + // Graphs are built lazily on first synthesis (bucketed by duration). +} + +runtime::TaskResult F5TTSSession::run(const runtime::TaskRequest & request) { + if (!request.text_input.has_value() || request.text_input->text.empty()) { + throw std::runtime_error("F5-TTS requires input text"); + } + const runtime::AudioBuffer * ref = reference_audio(request); + if (ref == nullptr || ref->samples.empty()) { + throw std::runtime_error( + "F5-TTS requires reference voice audio (voice preset or voice_ref)"); + } + const auto ref_text_it = request.options.find("reference_text"); + if (ref_text_it == request.options.end() || ref_text_it->second.empty()) { + throw std::runtime_error( + "F5-TTS requires reference_text (transcript of the reference audio)"); + } + + F5SynthesisRequest req; + req.text = request.text_input->text; + req.ref_audio = ref->samples; + req.ref_sample_rate = ref->sample_rate; + req.ref_text = ref_text_it->second; + if (const auto v = runtime::find_option(request.options, {"dialect"})) { + req.dialect = *v; + } else { + req.dialect = dialect_; + } + if (const auto v = runtime::find_option(request.options, {"speed"})) { + req.speed = std::stof(*v); + } + if (const auto v = runtime::find_option(request.options, {"num_inference_steps"})) { + req.steps = std::stoi(*v); + } + if (const auto v = runtime::find_option(request.options, {"cfg_strength", "guidance_scale"})) { + req.cfg_strength = std::stof(*v); + } + if (const auto v = runtime::find_option(request.options, {"sway_sampling_coef"})) { + req.sway_sampling_coef = std::stof(*v); + } + if (const auto v = runtime::find_option(request.options, {"seed"})) { + req.seed = static_cast(std::stoul(*v)); + req.fixed_seed = true; + } + if (const auto v = runtime::find_option(request.options, {"strip_diacritics"})) { + req.strip_diacritics = runtime::parse_bool_option(*v, "strip_diacritics"); + } + req.use_cuda = use_cuda_; + req.frame_budget = frame_budget_; + req.cuda_device = cuda_device_; + req.threads = threads_; + + const auto out = f5_synthesize( + assets_->checkpoint.string(), vocos_path_, req); + + runtime::TaskResult result; + runtime::AudioBuffer audio; + audio.sample_rate = static_cast(out.sample_rate); + audio.channels = 1; + audio.samples = std::move(out.audio); + result.audio_output = std::move(audio); + return result; +} + +std::shared_ptr make_f5_tts_loader() { + runtime::SpecBackedVoiceModelConfig config; + config.family = std::string(kFamily); + config.aliases = {"habibi", "habibi_tts"}; + config.load_assets = load_f5_tts_assets; + config.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)); + }; + return runtime::make_spec_backed_voice_loader(std::move(config)); +} + +} // namespace engine::models::f5_tts diff --git a/src/community_models/f5_tts/synthesize.cpp b/src/community_models/f5_tts/synthesize.cpp new file mode 100644 index 00000000..8b1e981b --- /dev/null +++ b/src/community_models/f5_tts/synthesize.cpp @@ -0,0 +1,1290 @@ +#include "engine/community_models/f5_tts/synthesize.h" + +#include "cpu_graph_compute.h" + +#include "engine/community_models/f5_tts/runtime.h" + +#include "cpu_graph_compute.h" + +#include "engine/framework/core/backend.h" + +#include "ggml.h" +#include "ggml-cpu.h" +#include "engine/framework/assets/tensor_source.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::f5_tts { +namespace { + +// ---- mel filterbank (librosa-compatible, htk-free slaney, 24 kHz) ---------- +// MSVC does not define kPi; use our own constant everywhere. +constexpr float kPi = 3.14159265358979323846F; +constexpr int kSampleRate = 24000; +constexpr int kNfft = 1024; +constexpr int kHop = 256; +constexpr int kNMel = 100; +constexpr float kFMin = 0.0F; +constexpr float kFMax = 12000.0F; + +// torchaudio defaults used by F5: htk mel scale, norm=None (no slaney +// normalization), power=1 (magnitude), f_max=sample_rate/2. +float hz_to_mel_htk(float hz) { return 2595.0F * std::log10(1.0F + hz / 700.0F); } +float mel_to_hz_htk(float mel) { return 700.0F * (std::pow(10.0F, mel / 2595.0F) - 1.0F); } + +const std::vector & mel_filterbank() { + static std::vector fb; // [n_freqs, kNMel] like torch fb (freq-major) + static std::once_flag once; + std::call_once(once, [] { + const int n_freqs = kNfft / 2 + 1; + const float m_min = hz_to_mel_htk(kFMin); + const float m_max = hz_to_mel_htk(kFMax); + std::vector mels(kNMel + 2); + for (int i = 0; i < kNMel + 2; ++i) { + mels[i] = mel_to_hz_htk(m_min + (m_max - m_min) * i / (kNMel + 1)); + } + fb.assign(static_cast(n_freqs) * kNMel, 0.0F); + for (int m = 0; m < kNMel; ++m) { + const float lo = mels[m]; + const float mid = mels[m + 1]; + const float hi = mels[m + 2]; + for (int f = 0; f < n_freqs; ++f) { + const float freq = static_cast(f) * kSampleRate / kNfft; + if (freq <= lo || freq >= hi) continue; + const float w = freq <= mid + ? (freq - lo) / (mid - lo) + : (hi - freq) / (hi - mid); + fb[static_cast(f) * kNMel + m] = w; // freq-major like torch + } + } + }); + return fb; +} + +void fft_inplace(std::vector & re, std::vector & im, bool inverse) { + const size_t n = re.size(); + for (size_t i = 1, j = 0; i < n; ++i) { + size_t bit = n >> 1; + for (; j & bit; bit >>= 1) j ^= bit; + j ^= bit; + if (i < j) { + std::swap(re[i], re[j]); + std::swap(im[i], im[j]); + } + } + for (size_t len = 2; len <= n; len <<= 1) { + const float ang = static_cast(2.0 * kPi / static_cast(len)) * (inverse ? 1 : -1); + for (size_t i = 0; i < n; i += len) { + for (size_t k = 0; k < len / 2; ++k) { + const float wr = std::cos(ang * static_cast(k)); + const float wi = std::sin(ang * static_cast(k)); + const size_t a = i + k; + const size_t b = i + k + len / 2; + const float vr = re[b] * wr - im[b] * wi; + const float vi = re[b] * wi + im[b] * wr; + const float ur = re[a]; + const float ui = im[a]; + re[a] = ur + vr; + im[a] = ui + vi; + re[b] = ur - vr; + im[b] = ui - vi; + } + } + } + if (inverse) { + for (size_t i = 0; i < n; ++i) { + re[i] /= static_cast(n); + im[i] /= static_cast(n); + } + } +} + +// log-mel with reflect-padded center STFT (librosa semantics), log clamp 1e-5. +std::vector compute_mel(const std::vector & wav) { + const int n_freqs = kNfft / 2 + 1; + // torchaudio center=True: n_frames = 1 + floor(len / hop) + const int frames = std::max(1, 1 + static_cast(wav.size() / kHop)); + std::vector hann(kNfft); + for (int i = 0; i < kNfft; ++i) { + hann[i] = 0.5F * (1.0F - std::cos(2.0F * static_cast(kPi) * i / kNfft)); + } + const auto & fb = mel_filterbank(); + std::vector mel(static_cast(kNMel) * frames); + std::vector re(kNfft), im(kNfft); + std::vector spec_pow(n_freqs); + for (int t = 0; t < frames; ++t) { + std::fill(re.begin(), re.end(), 0.0F); + std::fill(im.begin(), im.end(), 0.0F); + // torchaudio center=True alignment: frame t center = t*hop, achieved by + // starting one hop earlier than librosa's default convention. + const int start = (t - 2) * kHop; + for (int i = 0; i < kNfft; ++i) { + int r = start + i; + // torch reflect pad: wav[-k] = wav[k], wav[L+k] = wav[L-2-k] + if (r < 0) r = -r; + if (r >= static_cast(wav.size())) r = 2 * static_cast(wav.size()) - 2 - r; + r = std::clamp(r, 0, static_cast(wav.size()) - 1); + re[i] = wav[r] * hann[i]; + } + fft_inplace(re, im, false); + for (int f = 0; f < n_freqs; ++f) { + spec_pow[f] = std::sqrt(re[f] * re[f] + im[f] * im[f]); // power=1 magnitude + } + for (int m = 0; m < kNMel; ++m) { + float acc = 0.0F; + for (int f = 0; f < n_freqs; ++f) { + acc += fb[static_cast(f) * kNMel + m] * spec_pow[f]; + } + mel[static_cast(m) * frames + t] = std::log(std::max(acc, 1e-5F)); + } + } + return mel; // [mel, frames] feature-fastest memory +} + +std::vector resample(const std::vector & in, int sr_in, int sr_out) { + if (sr_in == sr_out || in.empty()) return in; + const double ratio = static_cast(sr_out) / sr_in; + const size_t out_n = static_cast(static_cast(in.size()) * ratio); + std::vector out(out_n); + for (size_t i = 0; i < out_n; ++i) { + const double pos = static_cast(i) / ratio; + const size_t i0 = static_cast(pos); + const size_t i1 = std::min(i0 + 1, in.size() - 1); + const double frac = pos - i0; + out[i] = static_cast(in[i0] * (1 - frac) + in[i1] * frac); + } + return out; +} + +std::unordered_map load_vocab(const std::string & dir) { + std::unordered_map map; + std::ifstream f(dir + "/vocab.txt", std::ios::binary); + if (!f) throw std::runtime_error("cannot open vocab.txt in " + dir); + std::string content((std::istreambuf_iterator(f)), std::istreambuf_iterator()); + std::vector lines; + std::string cur; + for (size_t i = 0; i < content.size();) { + if (content[i] == '\n') { + lines.push_back(cur); + cur.clear(); + ++i; + continue; + } + size_t len = 1; + const auto c = static_cast(content[i]); + if (c >= 0xF0) len = 4; + else if (c >= 0xE0) len = 3; + else if (c >= 0xC0) len = 2; + cur.append(content, i, len); + i += len; + } + if (!cur.empty()) lines.push_back(cur); + for (size_t i = 0; i < lines.size(); ++i) { + if (!lines[i].empty()) map.emplace(lines[i], static_cast(i)); + } + return map; +} + +const char * dialect_token(const std::string & dialect) { + if (dialect == "MSA") return "\xE2\x91\xA0"; + if (dialect == "SAU") return "\xE2\x91\xA1"; + if (dialect == "UAE") return "\xE2\x91\xA2"; + if (dialect == "ALG") return "\xE2\x91\xA3"; + if (dialect == "IRQ") return "\xE2\x91\xA4"; + if (dialect == "EGY") return "\xE2\x91\xA5"; + if (dialect == "MAR") return "\xE2\x91\xA6"; + if (dialect == "OMN") return "\xE2\x91\xA7"; + if (dialect == "TUN") return "\xE2\x91\xA8"; + if (dialect == "LEV") return "\xE2\x91\xA9"; + if (dialect == "SDN") return "\xE2\x91\xAA"; + if (dialect == "LBY") return "\xE2\x91\xAB"; + return "\xE2\x93\xAA"; // ⓪ UNK +} + +std::vector utf8_chars(const std::string & s) { + std::vector out; + for (size_t i = 0; i < s.size();) { + size_t len = 1; + const auto c = static_cast(s[i]); + if (c >= 0xF0) len = 4; + else if (c >= 0xE0) len = 3; + else if (c >= 0xC0) len = 2; + out.emplace_back(s, i, len); + i += len; + } + return out; +} + +// Strip Arabic combining marks the model cannot read: Habibi was trained on +// ASR transcripts, which are undiacritized, so harakat/tanwin/shadda tokens +// are severely undertrained and diacritized input degrades to garbled speech +// with character repetitions (verified identical in the Python reference). +// Removes tatweel U+0640, harakat U+064B..U+065F, dagger alif U+0670 — +// all encoded as 0xD9-prefixed two-byte sequences. +std::string strip_arabic_diacritics(const std::string & s) { + std::string out; + out.reserve(s.size()); + for (size_t i = 0; i < s.size();) { + const auto c = static_cast(s[i]); + if (c == 0xD9 && i + 1 < s.size()) { + const auto c2 = static_cast(s[i + 1]); + if (c2 == 0x80 || (c2 >= 0x8B && c2 <= 0x9F) || c2 == 0xB0) { + i += 2; // drop the combining mark + continue; + } + } + size_t len = 1; + if (c >= 0xF0) len = 4; + else if (c >= 0xE0) len = 3; + else if (c >= 0xC0) len = 2; + out.append(s, i, len); + i += len; + } + return out; +} + +// Python: if ref_text ends with a single-byte char (ASCII), a space is +// appended so ref and gen text do not fuse into one token stream. +std::string apply_ref_trailing_space(std::string ref_text) { + if (!ref_text.empty() && (static_cast(ref_text.back()) & 0x80) == 0) { + ref_text += " "; + } + return ref_text; +} + +// Habibi prompt assembly: dialect token + 〈ref_text + gen_chunk〉. +std::string assemble_chunk_text( + const std::string & dialect, const std::string & ref_text, const std::string & chunk) { + return std::string(dialect_token(dialect)) + "\xE3\x80\x88" + ref_text + chunk + "\xE3\x80\x89"; +} + +std::vector tokenize_text( + const std::unordered_map & vocab, const std::string & s) { + std::vector ids; + for (const auto & ch : utf8_chars(s)) { + const auto it = vocab.find(ch); + ids.push_back(it != vocab.end() ? it->second : 0); + } + return ids; +} + +struct Rng { + uint64_t state; + explicit Rng(uint64_t seed) : state(seed ? seed : 0x9E3779B97F4A7C15ULL) {} + uint64_t next_u64() { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + return state; + } + float next_f32() { + return static_cast(static_cast(next_u64() >> 11) / 9007199254740992.0); + } + float normal() { + // Box-Muller + float u1 = std::max(next_f32(), 1e-7F); + float u2 = next_f32(); + return std::sqrt(-2.0F * std::log(u1)) * std::cos(2.0F * static_cast(kPi) * u2); + } +}; + +// Python F5 uses Empirically Pruned Step Sampling (EPSS) for low NFE: +// non-uniform grids on a 1/32 quantum (get_epss_timesteps); uniform +// linspace otherwise. The sway transform is applied on top either way. +std::vector sway_timesteps(int steps, float coef) { + static const std::map> kEpss = { + {5, {0, 2, 4, 8, 16, 32}}, + {6, {0, 2, 4, 6, 8, 16, 32}}, + {7, {0, 2, 4, 6, 8, 16, 24, 32}}, + {10, {0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32}}, + {12, {0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32}}, + {16, {0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32}}, + }; + std::vector t; + const auto it = kEpss.find(steps); + if (steps == 32 || it == kEpss.end()) { + t.resize(static_cast(steps) + 1); + for (int i = 0; i <= steps; ++i) { + t[static_cast(i)] = static_cast(i) / steps; + } + } else { + for (const int v : it->second) t.push_back(static_cast(v) / 32.0F); + } + for (auto & v : t) { + v = v + coef * (std::cos(static_cast(kPi) / 2 * v) - 1 + v); + } + return t; +} + +// ---- vocos vocoder ---------------------------------------------------------- +// Loads vocos.safetensors and decodes [frames][100] log-mel -> waveform. +struct VocosWeights { + std::shared_ptr source; + std::vector embed_w; // [512, 100, 7] torch + std::vector embed_b; // [512] + std::vector input_nw, input_nb; // backbone.norm (pre-blocks) + struct Block { + std::vector dw_w, dw_b, n_w, n_b, p1w, p1b, p2w, p2b, gamma; + }; + std::vector blocks; // 8 + std::vector final_nw, final_nb; // [512] + std::vector head_w, head_b; // [1026, 512], [1026] +}; + +const VocosWeights & load_vocos_once(const std::string & path) { + static std::unordered_map cache; + static std::mutex mutex; + std::lock_guard lock(mutex); + if (const auto it = cache.find(path); it != cache.end()) return it->second; + VocosWeights v; + v.source = assets::open_tensor_source(path); + if (std::filesystem::path(path).extension() == ".gguf") { + // GGUF packages store the vocoder under the "vocos" namespace + v.source = assets::make_prefixed_tensor_source(v.source, "vocos"); + } + const auto f32 = [&](const char * n) { return v.source->require_f32(n); }; + v.embed_w = f32("backbone.embed.weight"); + v.embed_b = f32("backbone.embed.bias"); + v.input_nw = f32("backbone.norm.weight"); + v.input_nb = f32("backbone.norm.bias"); + v.blocks.resize(8); + for (int i = 0; i < 8; ++i) { + const std::string p = "backbone.convnext." + std::to_string(i) + "."; + auto & b = v.blocks[i]; + b.dw_w = f32((p + "dwconv.weight").c_str()); + b.dw_b = f32((p + "dwconv.bias").c_str()); + b.n_w = f32((p + "norm.weight").c_str()); + b.n_b = f32((p + "norm.bias").c_str()); + b.p1w = f32((p + "pwconv1.weight").c_str()); + b.p1b = f32((p + "pwconv1.bias").c_str()); + b.p2w = f32((p + "pwconv2.weight").c_str()); + b.p2b = f32((p + "pwconv2.bias").c_str()); + b.gamma = f32((p + "gamma").c_str()); + } + v.final_nw = f32("backbone.final_layer_norm.weight"); + v.final_nb = f32("backbone.final_layer_norm.bias"); + v.head_w = f32("head.out.weight"); + v.head_b = f32("head.out.bias"); + return cache.emplace(path, std::move(v)).first->second; +} + +// CPU decode (pure host math, no graph): frames x 100 -> samples +std::vector vocos_decode(const std::string & vocos_path, const std::vector & mel_rows) { + const auto & v = load_vocos_once(vocos_path); + const int T = static_cast(mel_rows.size()) / kNMel; + const int D = 512; + const int IM = 1536; + + // embed conv k7 pad 3 groups 1: out[t, o] = sum_i sum_k w[o, i, k] * x[t+k-3, i] + b[o] + std::vector h(static_cast(T) * D); + for (int t = 0; t < T; ++t) { + for (int o = 0; o < D; ++o) { + float acc = v.embed_b[o]; + const float * w = v.embed_w.data() + static_cast(o) * kNMel * 7; + for (int k = 0; k < 7; ++k) { + const int tt = t + k - 3; + if (tt < 0 || tt >= T) continue; + const float * x = mel_rows.data() + static_cast(tt) * kNMel; + for (int i = 0; i < kNMel; ++i) { + acc += w[static_cast(i) * 7 + k] * x[i]; + } + } + h[static_cast(t) * D + o] = acc; + } + } + + // input layernorm (backbone.norm) after embed — VocosBackbone.forward applies + // norm BEFORE the convnext stack (easy to miss; verified against torch hook). + for (int t = 0; t < T; ++t) { + float * x = h.data() + static_cast(t) * D; + float mu = 0, var = 0; + for (int c = 0; c < D; ++c) mu += x[c]; + mu /= D; + for (int c = 0; c < D; ++c) { + const float d = x[c] - mu; + var += d * d; + } + var /= D; + const float inv = 1.0F / std::sqrt(var + 1e-6F); + for (int c = 0; c < D; ++c) { + x[c] = (x[c] - mu) * inv * v.input_nw[c] + v.input_nb[c]; + } + } + + // 8 convnext blocks +#ifdef F5_MEL_TEST + { std::ofstream f("/tmp/cpp_vocos_embed.bin", std::ios::binary); + f.write(reinterpret_cast(h.data()), h.size() * 4); } +#endif + for (int bi = 0; bi < 8; ++bi) { + const auto & B = v.blocks[bi]; + std::vector dw(static_cast(T) * D); + for (int t = 0; t < T; ++t) { + for (int c = 0; c < D; ++c) { + float acc = B.dw_b[c]; + const float * w = B.dw_w.data() + static_cast(c) * 7; + for (int k = 0; k < 7; ++k) { + const int tt = t + k - 3; + if (tt < 0 || tt >= T) continue; + acc += w[k] * h[static_cast(tt) * D + c]; + } + dw[static_cast(t) * D + c] = acc; + } + } + // ln -> gelu -> pw1 -> pw2 -> gamma -> +residual + std::vector ln_buf(static_cast(T) * D); + for (int t = 0; t < T; ++t) { + const float * x = dw.data() + static_cast(t) * D; + float mu = 0, var = 0; + for (int c = 0; c < D; ++c) mu += x[c]; + mu /= D; + for (int c = 0; c < D; ++c) { + const float d = x[c] - mu; + var += d * d; + } + var /= D; + const float inv = 1.0F / std::sqrt(var + 1e-6F); + for (int c = 0; c < D; ++c) { + ln_buf[static_cast(t) * D + c] = (x[c] - mu) * inv * B.n_w[c] + B.n_b[c]; + } + } + std::vector mid(static_cast(T) * IM); + for (int t = 0; t < T; ++t) { + for (int o = 0; o < IM; ++o) { + const float * w = B.p1w.data() + static_cast(o) * D; + const float * x = ln_buf.data() + static_cast(t) * D; + float acc = B.p1b[o]; + for (int c = 0; c < D; ++c) acc += w[c] * x[c]; + // gelu exact + const float xg = acc; + const float k0 = 0.7978845608028654F; + const float inner = k0 * xg * (1.0F + 0.044715F * xg * xg); + acc = 0.5F * xg * (1.0F + std::tanh(inner)); + mid[static_cast(t) * IM + o] = acc; + } + } + for (int t = 0; t < T; ++t) { + for (int o = 0; o < D; ++o) { + const float * w = B.p2w.data() + static_cast(o) * IM; + const float * x = mid.data() + static_cast(t) * IM; + float acc = B.p2b[o]; + for (int c = 0; c < IM; ++c) acc += w[c] * x[c]; + h[static_cast(t) * D + o] += B.gamma[o] * acc; + } + } +#ifdef F5_MEL_TEST + if (bi == 0) { std::ofstream f("/tmp/cpp_vocos_blk0.bin", std::ios::binary); + f.write(reinterpret_cast(h.data()), h.size() * 4); } +#endif + } + + // final layernorm +#ifdef F5_MEL_TEST + { + std::ofstream f("/tmp/cpp_vocos_h.bin", std::ios::binary); + f.write(reinterpret_cast(h.data()), h.size() * 4); + } +#endif + for (int t = 0; t < T; ++t) { + float * x = h.data() + static_cast(t) * D; + float mu = 0, var = 0; + for (int c = 0; c < D; ++c) mu += x[c]; + mu /= D; + for (int c = 0; c < D; ++c) { + const float d = x[c] - mu; + var += d * d; + } + var /= D; + const float inv = 1.0F / std::sqrt(var + 1e-6F); + for (int c = 0; c < D; ++c) { + x[c] = (x[c] - mu) * inv * v.final_nw[c] + v.final_nb[c]; + } + } +#ifdef F5_MEL_TEST + { std::ofstream f("/tmp/cpp_vocos_fln.bin", std::ios::binary); + f.write(reinterpret_cast(h.data()), h.size() * 4); } +#endif + + // head: [T, 1026] -> (mag 513, phase 513) + const int n_freqs = kNfft / 2 + 1; + std::vector spec(static_cast(T) * (2 * n_freqs)); + for (int t = 0; t < T; ++t) { + const float * x = h.data() + static_cast(t) * D; + for (int o = 0; o < 2 * n_freqs; ++o) { + const float * w = v.head_w.data() + static_cast(o) * D; + float acc = v.head_b[o]; + for (int c = 0; c < D; ++c) acc += w[c] * x[c]; + spec[static_cast(t) * (2 * n_freqs) + o] = acc; + } + } + +#ifdef F5_MEL_TEST + { std::ofstream f("/tmp/cpp_head_spec.bin", std::ios::binary); + f.write(reinterpret_cast(spec.data()), spec.size() * 4); } +#endif + // ISTFT (center, hann) — build full spectrum frames then overlap-add + const int out_len = (T - 1) * kHop + kNfft; + std::vector out(static_cast(out_len), 0.0F); + std::vector wsum(static_cast(out_len), 0.0F); + std::vector hann(kNfft); + for (int i = 0; i < kNfft; ++i) { + hann[i] = 0.5F * (1.0F - std::cos(2.0F * static_cast(kPi) * i / kNfft)); + } + std::vector re(kNfft), im(kNfft); + for (int t = 0; t < T; ++t) { + const float * row = spec.data() + static_cast(t) * (2 * n_freqs); + std::fill(re.begin(), re.end(), 0.0F); + std::fill(im.begin(), im.end(), 0.0F); + for (int f = 0; f < n_freqs; ++f) { + const float mag = std::min(std::exp(row[f]), 100.0F); + const float ph = row[n_freqs + f]; + re[f] = mag * std::cos(ph); + im[f] = mag * std::sin(ph); + if (f > 0 && f < n_freqs - 1) { + re[kNfft - f] = re[f]; + im[kNfft - f] = -im[f]; + } + } + re[n_freqs - 1] = im[n_freqs - 1] = 0.0F; // nyquist bin real + fft_inplace(re, im, true); + const int start = t * kHop; + for (int i = 0; i < kNfft; ++i) { + out[static_cast(start + i)] += re[i] * hann[i]; + wsum[static_cast(start + i)] += hann[i] * hann[i]; + } + } +#ifdef F5_MEL_TEST + { std::ofstream f("/tmp/cpp_istft_raw.bin", std::ios::binary); + f.write(reinterpret_cast(out.data()), out.size() * 4); } +#endif + std::vector audio(static_cast(T - 1) * kHop + 1); + const int keep = static_cast(audio.size()); + for (int i = 0; i < keep; ++i) { + const size_t idx = i + kNfft / 2; // center: drop first half-frame + audio[static_cast(i)] = + idx < out.size() && wsum[idx] > 1e-8F ? out[idx] / wsum[idx] : 0.0F; + } + return audio; +} + +} // namespace + +// ---- GPU vocoder (ggml graph; CUDA or CPU via F5ComputeDevice) ---- +// Same math as vocos_decode above (verified vs torch at mel-corr 0.9963): +// embed conv k7 -> input LN -> 8x ConvNeXt (dwconv7, LN, pw1+GELU, pw2, +// gamma, residual) -> final LN -> head linear [T,1026]; the ISTFT tail runs +// on the host (O(n) overlap-add, ~minor vs the backbone matmuls). +namespace { + +ggml_tensor * vlin( + ggml_context * ctx, + ggml_tensor * w, // ne [in, out] (torch [out, in] order) + ggml_tensor * b, // ne [out] + ggml_tensor * x) { // ne [in, T] + auto * out = ggml_mul_mat(ctx, w, x); + auto * b2 = ggml_reshape_2d(ctx, b, ggml_nelements(b), 1); + auto * r = ggml_repeat(ctx, b2, out); + return ggml_add(ctx, out, r); +} + +ggml_tensor * vln( + ggml_context * ctx, + ggml_tensor * x, // ne [D, T] + ggml_tensor * w, ggml_tensor * b) { // ne [D] + auto * n = ggml_norm(ctx, x, 1e-6F); + auto * w2 = ggml_reshape_2d(ctx, w, ggml_nelements(w), 1); + auto * b2 = ggml_reshape_2d(ctx, b, ggml_nelements(b), 1); + return ggml_add( + ctx, ggml_mul(ctx, n, ggml_repeat(ctx, w2, n)), ggml_repeat(ctx, b2, n)); +} + +// depthwise k7 pad3 over ne [C, T] (c fastest, time = ne1); wk host data +// (c*7+k), bias leaf [C] +ggml_tensor * leaf_zero_ret( + ggml_context * ctx, + int64_t c, + int64_t t, + const std::function & leaf_zero) { + auto * z = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, c, t); + leaf_zero(z, static_cast(c) * t * sizeof(float)); + return z; +} + +ggml_tensor * vdw7( + ggml_context * ctx, + ggml_tensor * x, + const float * wk_host, + ggml_tensor * b, + const std::function & leaf_write, + const std::function & leaf_zero) { + const int64_t C = x->ne[0]; + const int64_t T = x->ne[1]; + auto * xp = ggml_concat( + ctx, + ggml_concat(ctx, leaf_zero_ret(ctx, C, 3, leaf_zero), x, 1), + leaf_zero_ret(ctx, C, 3, leaf_zero), + 1); // [C, T+6] + ggml_tensor * out = nullptr; + for (int k = 0; k < 7; ++k) { + std::vector wk(C); + for (int c = 0; c < C; ++c) { + wk[c] = wk_host[static_cast(c) * 7 + k]; + } + auto * wk_t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, C, 1); + leaf_write(wk_t, wk.data(), wk.size() * sizeof(float)); + auto * v = ggml_view_2d( + ctx, xp, C, T, xp->nb[1], static_cast(k) * xp->nb[1]); + auto * term = ggml_mul(ctx, v, wk_t); // [C,T] * [C,1] broadcast + out = out == nullptr ? term : ggml_add(ctx, out, term); + } + auto * b2 = ggml_reshape_2d(ctx, b, ggml_nelements(b), 1); + return ggml_add(ctx, out, ggml_repeat(ctx, b2, out)); +} + +struct VocosGraph { + ggml_context * ctx = nullptr; + ggml_cgraph * graph = nullptr; + ggml_gallocr_t gallocr = nullptr; + ggml_backend_buffer_t io_buffer = nullptr; + ggml_tensor * mel = nullptr; // ne [100, T] (m fastest) + ggml_tensor * spec = nullptr; // ne [1026, T] (o fastest; row t at t*1026) +}; + +std::vector vocos_decode_gpu( + const std::string & vocos_path, + const std::vector & mel_rows, + const F5ComputeDevice & dev) { + const auto & v = load_vocos_once(vocos_path); + const int T = static_cast(mel_rows.size()) / kNMel; + const int D = 512; + const int IM = 1536; + const int n_freqs = kNfft / 2 + 1; + + // one backend per cuda device (or cpu); leaked at exit (CUDA driver + // shutdown cannot be ordered before buffer frees in static destruction) + struct BackendOwnerV { + ggml_backend_t value = nullptr; + }; + static auto * owners = new std::map(); + const int key = dev.use_cuda ? dev.device : -1; + ggml_backend_t backend = nullptr; + if (dev.use_cuda) { + auto ob = owners->find(key); + if (ob == owners->end()) { + core::BackendConfig cfg{core::BackendType::Cuda, dev.device, 1}; + ob = owners->emplace(key, BackendOwnerV{core::init_backend(cfg)}).first; + } + backend = ob->second.value; + } else { + auto ob = owners->find(key); + if (ob == owners->end()) { + core::BackendConfig cfg{core::BackendType::Cpu, 0, std::max(1, dev.threads)}; + ob = owners->emplace(key, BackendOwnerV{core::init_backend(cfg)}).first; + } + backend = ob->second.value; + } + const bool is_cuda = dev.use_cuda; + + // graph cache per (T, device) + static auto * cache = new std::map, std::unique_ptr>(); + const auto ckey = std::make_pair(T, key); + auto it = cache->find(ckey); + if (it == cache->end()) { + auto g = std::make_unique(); + const size_t ctx_bytes = 256ULL << 20; + g->ctx = ggml_init({ctx_bytes, nullptr, is_cuda}); + ggml_context * ctx = g->ctx; + std::vector>> pending; + const auto leaf_write = [&](ggml_tensor * t, const void * src, size_t bytes) { + if (!is_cuda) { + std::memcpy(t->data, src, bytes); + } else { + const auto * b = static_cast(src); + pending.emplace_back(t, std::vector(b, b + bytes)); + } + }; + const auto leaf_zero = [&](ggml_tensor * t, size_t bytes) { + if (!is_cuda) { + std::memset(t->data, 0, bytes); + } else { + pending.emplace_back(t, std::vector(bytes, 0)); + } + }; + auto leaf_f32v = [&](const std::vector & src) { + auto * t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, static_cast(src.size())); + leaf_write(t, src.data(), src.size() * sizeof(float)); + return t; + }; + + // mel leaf: ne [100, T], m fastest — mel_rows buffer is exactly that + auto * mel = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kNMel, T); + // weight leaves (torch order == needed ggml order, see notes) + auto * embed_w = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 7, kNMel, D); + leaf_write(embed_w, v.embed_w.data(), v.embed_w.size() * sizeof(float)); + auto * embed_b = leaf_f32v(v.embed_b); + ggml_tensor * var_h = nullptr; + auto * input_nw = leaf_f32v(v.input_nw); + auto * input_nb = leaf_f32v(v.input_nb); + auto * final_nw = leaf_f32v(v.final_nw); + auto * final_nb = leaf_f32v(v.final_nb); + auto * head_w = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, D, 2 * n_freqs); + leaf_write(head_w, v.head_w.data(), v.head_w.size() * sizeof(float)); + auto * head_b = leaf_f32v(v.head_b); + + // embed conv (dense k7, pad 3) via im2col + matmul on time-fastest rows + auto * t_fast = ggml_cont(ctx, ggml_transpose(ctx, mel)); // ne [T, 100] + { + // im2col: weight ne [k, cin, cout]; cols ne [cin*k, T] + auto * in3 = ggml_reshape_3d(ctx, t_fast, T, kNMel, 1); + auto * cols = ggml_im2col( + ctx, embed_w, in3, 1, 1, 3, 0, 1, 1, false, GGML_TYPE_F32); + auto * w2 = ggml_reshape_2d(ctx, embed_w, kNMel * 7, D); + auto * y = ggml_mul_mat(ctx, w2, cols); // [D, T] feature-fastest + auto * b2_ = ggml_reshape_2d(ctx, embed_b, D, 1); + auto * yb = ggml_add(ctx, y, ggml_repeat(ctx, b2_, y)); + var_h = ggml_cont(ctx, yb); + } + auto * h = var_h; + + h = vln(ctx, h, input_nw, input_nb); + ggml_set_name(h, "v_after_input_ln"); + + for (int bi = 0; bi < 8; ++bi) { + const auto & B = v.blocks[bi]; + auto * dwb = leaf_f32v(B.dw_b); + auto * nw = leaf_f32v(B.n_w); + auto * nb = leaf_f32v(B.n_b); + auto * p1w = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, D, IM); + leaf_write(p1w, B.p1w.data(), B.p1w.size() * sizeof(float)); + auto * p1b = leaf_f32v(B.p1b); + auto * p2w = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, IM, D); + leaf_write(p2w, B.p2w.data(), B.p2w.size() * sizeof(float)); + auto * p2b = leaf_f32v(B.p2b); + auto * gam = leaf_f32v(B.gamma); + + auto * dw = vdw7(ctx, h, B.dw_w.data(), dwb, leaf_write, leaf_zero); + ggml_set_name(dw, "v_dw"); + auto * ln = vln(ctx, dw, nw, nb); + auto * mid = vlin(ctx, p1w, p1b, ln); + mid = ggml_gelu(ctx, mid); // exact erf + auto * up = vlin(ctx, p2w, p2b, mid); // [512, T] + auto * g2 = ggml_reshape_2d(ctx, gam, D, 1); + up = ggml_mul(ctx, up, ggml_repeat(ctx, g2, up)); + h = ggml_add(ctx, h, up); + } + + h = vln(ctx, h, final_nw, final_nb); + auto * spec = vlin(ctx, head_w, head_b, h); // ne [1026, T] + + g->mel = mel; + g->spec = spec; + g->graph = ggml_new_graph_custom(ctx, 8192, false); + ggml_build_forward_expand(g->graph, spec); + core::validate_backend_graph_supported(backend, g->graph, "f5_vocos"); + if (is_cuda) { + g->io_buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + for (auto & leaf : pending) { + ggml_backend_tensor_set(leaf.first, leaf.second.data(), 0, leaf.second.size()); + } + g->gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (g->gallocr == nullptr || !ggml_gallocr_reserve(g->gallocr, g->graph) || + !ggml_gallocr_alloc_graph(g->gallocr, g->graph)) { + throw std::runtime_error("vocos CUDA graph alloc failed"); + } + } + it = cache->emplace(ckey, std::move(g)).first; + } + VocosGraph & g = *it->second; + + // upload mel + compute + if (is_cuda) { + ggml_backend_tensor_set(g.mel, mel_rows.data(), 0, mel_rows.size() * sizeof(float)); + } else { + std::memcpy(g.mel->data, mel_rows.data(), mel_rows.size() * sizeof(float)); + } + const auto status = is_cuda + ? core::compute_backend_graph(backend, g.graph, nullptr, "f5_vocos") + : f5_cpu_graph_compute(g.ctx, g.graph, + dev.threads > 0 ? dev.threads + : static_cast(std::thread::hardware_concurrency())); + if (is_cuda) { + ggml_backend_synchronize(backend); + } + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("vocos graph compute failed"); + } + // spec: ne [1026, T] o-fastest == row t at t*1026 — same as host layout + std::vector spec(static_cast(T) * 2 * n_freqs); + if (is_cuda) { + ggml_backend_tensor_get(g.spec, spec.data(), 0, spec.size() * sizeof(float)); + } else { + std::memcpy(spec.data(), ggml_get_data(g.spec), spec.size() * sizeof(float)); + } + + // host ISTFT tail (identical to vocos_decode) + const int out_len = (T - 1) * kHop + kNfft; + std::vector out(static_cast(out_len), 0.0F); + std::vector wsum(static_cast(out_len), 0.0F); + std::vector hann(kNfft); + for (int i = 0; i < kNfft; ++i) { + hann[i] = 0.5F * (1.0F - std::cos(2.0F * static_cast(kPi) * i / kNfft)); + } + std::vector re(kNfft), im(kNfft); + for (int t = 0; t < T; ++t) { + const float * row = spec.data() + static_cast(t) * (2 * n_freqs); + std::fill(re.begin(), re.end(), 0.0F); + std::fill(im.begin(), im.end(), 0.0F); + for (int f = 0; f < n_freqs; ++f) { + const float mag = std::min(std::exp(row[f]), 100.0F); + const float ph = row[n_freqs + f]; + re[f] = mag * std::cos(ph); + im[f] = mag * std::sin(ph); + if (f > 0 && f < n_freqs - 1) { + re[kNfft - f] = re[f]; + im[kNfft - f] = -im[f]; + } + } + re[n_freqs - 1] = im[n_freqs - 1] = 0.0F; + fft_inplace(re, im, true); + const int start = t * kHop; + for (int i = 0; i < kNfft; ++i) { + out[static_cast(start + i)] += re[i] * hann[i]; + wsum[static_cast(start + i)] += hann[i] * hann[i]; + } + } + std::vector audio(static_cast(T - 1) * kHop + 1); + const int keep = static_cast(audio.size()); + for (int i = 0; i < keep; ++i) { + const size_t idx = static_cast(i) + kNfft / 2; + audio[static_cast(i)] = + idx < out.size() && wsum[idx] > 1e-8F ? out[idx] / wsum[idx] : 0.0F; + } + return audio; +} + +} // namespace + +namespace { + +// UTF-8-aware sentence splitting: returns byte offsets that never split a +// multi-byte character. Break after . ! ? ؛ ، and newlines. +static bool is_break_byte_here(const std::string & text, size_t i, size_t * len) { + const unsigned char c = static_cast(text[i]); + if (c == '.' || c == '!' || c == '?' || c == '\n') { + *len = 1; + return true; + } + // Arabic semicolon U+061B (D8 9B) and comma U+060C (D8 8C) + if (c == 0xD8 && i + 1 < text.size()) { + const unsigned char c2 = static_cast(text[i + 1]); + if (c2 == 0x9B || c2 == 0x8C) { + *len = 2; + return true; + } + } + return false; +} + +static size_t utf8_char_count(const std::string & s) { + size_t n = 0; + for (size_t i = 0; i < s.size();) { + const unsigned char c = static_cast(s[i]); + i += c < 0x80 ? 1 : (c & 0xE0) == 0xC0 ? 2 : (c & 0xF0) == 0xE0 ? 3 : 4; + ++n; + } + return n; +} + +// Split text so each chunk has at most max_chars UTF-8 characters, +// preferring sentence/clause boundaries. +std::vector chunk_text(const std::string & text, size_t max_chars) { + std::vector chunks; + if (utf8_char_count(text) <= max_chars) { + chunks.push_back(text); + return chunks; + } + // collect sentence pieces [start, end) breaking AFTER break chars + std::vector> pieces; + size_t piece_start = 0; + for (size_t i = 0; i < text.size();) { + size_t blen = 0; + if (is_break_byte_here(text, i, &blen)) { + pieces.emplace_back(piece_start, i + blen); + piece_start = i + blen; + i += blen; + continue; + } + const unsigned char c = static_cast(text[i]); + i += c < 0x80 ? 1 : (c & 0xE0) == 0xC0 ? 2 : (c & 0xF0) == 0xE0 ? 3 : 4; + } + if (piece_start < text.size()) { + pieces.emplace_back(piece_start, text.size()); + } + // split every piece into <= max_chars slices (oversize sentences too), + // snapping each cut to a WORD boundary: a mid-word slice makes the next + // chunk start mid-word right after the reference prompt, and the model + // drops the straddled word (observed: "لن |"يحتفظ" split). Falls back to + // a hard cut only when the window contains no space at all. + std::vector> slices; + for (const auto & [ps, pe] : pieces) { + size_t cs = ps; + while (cs < pe) { + size_t cnt = 0, ce = cs; + while (ce < pe && cnt < max_chars) { + const unsigned char c = static_cast(text[ce]); + ce += c < 0x80 ? 1 : (c & 0xE0) == 0xC0 ? 2 : (c & 0xF0) == 0xE0 ? 3 : 4; + ++cnt; + } + if (ce < pe) { + // back off to just after the last space inside the window + for (size_t k = ce; k > cs + 1; --k) { + if (text[k - 1] == ' ') { + ce = k; + break; + } + } + } + slices.emplace_back(cs, ce); + cs = ce; + } + } + for (const auto & [ss, se] : slices) { + const std::string piece = text.substr(ss, se - ss); + const size_t pc = utf8_char_count(piece); + if (!chunks.empty()) { + const std::string & prev = chunks.back(); + const size_t prev_c = utf8_char_count(prev); + // merge when it fits, or when the piece is tiny (< 12 chars): + // tiny chunks destabilize the sampler (observed NaN on 5 chars). + // Absorbing beyond the size budget overflows the duration + // estimate into the 1024-frame cap -> clipped trailing words, so + // the allowance stays within the chunk-sizing safety margin. + const bool tiny = pc < 12; + const bool fits = prev_c + pc <= max_chars; + const bool absorb = tiny && prev_c + pc <= max_chars + 4; + if (fits || absorb) { + chunks.back() = prev + piece; + continue; + } + } + chunks.push_back(piece); + } + return chunks; +} + +struct ChunkResult { + std::vector gen_mel_rows; // [gen][100] + int gen_frames = 0; + int duration_real = 0; +}; + +// Trim head/tail silence from a chunk's generated mel rows. The model parks +// unused duration slack as long pauses at chunk EDGES (observed: 1.5-3s +// pauses mid-text with a fast-paced reference), which lands between random +// word pairs once chunks are concatenated. Speech log-mel row means sit far +// above the silence floor (speech ~-0.5, pauses ~-6..-8), so -4.0 separates +// them robustly. A short head/tail is kept for natural word spacing; a chunk +// ending on sentence-final punctuation keeps a longer tail pause. +void trim_chunk_mel_silence( + std::vector & rows, int & gen_frames, bool sentence_final_tail) { + if (gen_frames <= 0) return; + const float kThresh = -4.0F; + const auto row_mean = [&](int t) { + const float * r = rows.data() + static_cast(t) * kNMel; + float acc = 0.0F; + for (int m = 0; m < kNMel; ++m) acc += r[m]; + return acc / kNMel; + }; + int first = 0; + while (first < gen_frames && row_mean(first) < kThresh) ++first; + first = std::max(0, first - 9); // keep ~0.1s lead-in + int last = gen_frames - 1; + while (last >= first && row_mean(last) < kThresh) --last; + if (last < first) return; // all-silent chunk (pathological): leave as-is + const int tail_keep = sentence_final_tail ? 47 : 24; // ~0.5s / ~0.26s + last = std::min(gen_frames - 1, last + tail_keep); + const int keep = last - first + 1; + if (keep >= gen_frames) return; + std::vector trimmed(static_cast(keep) * kNMel); + std::memcpy(trimmed.data(), rows.data() + static_cast(first) * kNMel, + trimmed.size() * sizeof(float)); + rows = std::move(trimmed); + gen_frames = keep; +} + + +// Total mel-frame budget for one CFM pass (ref + generated). 2048 allows +// sentence-scale chunks (fewer seams, splits land on clause boundaries — +// the 1024 budget forced ~57-char mid-word slices that dropped straddled +// words) at ~6 GiB peak on an RTX 3090 (measured). Override via +// F5SynthesisRequest::frame_budget (session option f5_tts.frame_budget). +int frame_budget(const F5SynthesisRequest & request) { + return request.frame_budget > 0 ? request.frame_budget : 2048; +} + +// One CFM pass for a single chunk: the original pipeline verbatim. +ChunkResult synthesize_chunk( + const std::string & model_path, + const F5SynthesisRequest & request, + const std::vector & ref_mel_cols, // [100][ref_frames] + int ref_frames, + int ref_voiced_frames, + const std::vector & chunk_ids, + const std::string & chunk_text_string, + const std::string & chunk_ref_text, + F5ComputeDevice & dev, + uint32_t seed, + std::vector * out_final_latent_rows, + double duration_slack = 1.0, + int tail_pad = 0) { + const F5Architecture arch; + + // Pacing in CHARACTERS (Arabic is 2 bytes/char; byte-based pacing + // underestimates duration ~1.8x). The reference's speaking rate sets the + // expectation, bounded by a normal-speech ceiling so an unusually slow + // reference cannot drag generated speech into a compressed clamp. + const int gen_chars = static_cast(utf8_char_count(chunk_text_string)); + const int ref_chars = std::max(1, static_cast(utf8_char_count(chunk_ref_text))); + const double ref_rate = static_cast(ref_voiced_frames) / ref_chars; // frames/char + // ceiling only guards pathological refs (e.g. mostly silence); the + // reference's speaking rate drives pacing, so slow refs stay slow + constexpr double kMaxFramesPerChar = 93.75 / 2.5; // >= 2.5 chars/s floor... (frames/char ceiling) + const double rate = std::max(std::min(ref_rate, kMaxFramesPerChar), 93.75 / 14.0); + float local_speed = request.speed; + if (gen_chars < 10) local_speed = 0.3F; + // duration_slack (>1 for chunked long-form): the rate estimate has zero + // tolerance for run-to-run pace variance; mid-text chunk tails end on + // whole words, and any undershoot clips the last word (observed: "النفط", + // "فقط" dropped at chunk tails). Excess frames become a short tail pause. + int duration = ref_frames + static_cast(rate * gen_chars * duration_slack / local_speed); + // tail_pad (>0 for single-chunk short text): without it a slightly slow + // sampled pace runs out of frames and the final phonemes are clipped + // (observed: "أين اللون الأحمر؟" -> "الأخر"). ~0.2s is proportionally + // negligible for long text and ends as silence, not slower speech. + duration += tail_pad; + duration = std::max(duration, gen_chars + 1); + // per-chunk safety cap: a single chunk never exceeds the graph budget; + // longer inputs are split upstream by chunk_text instead of truncated. + const int kChunkFrameCap = frame_budget(request); + if (duration > kChunkFrameCap) duration = kChunkFrameCap; + const int duration_real = duration; + duration = (duration + 63) / 64 * 64; // graph bucket reuse + + std::vector cond(static_cast(duration) * kNMel, 0.0F); + for (int t = 0; t < ref_frames && t < duration; ++t) { + for (int m = 0; m < kNMel; ++m) { + cond[static_cast(t) * kNMel + m] = + ref_mel_cols[static_cast(m) * ref_frames + t]; + } + } + Rng rng(seed); + std::vector y(static_cast(duration) * kNMel); + for (auto & val : y) val = rng.normal(); + + const auto ts = sway_timesteps(request.steps, request.sway_sampling_coef); + for (size_t i = 0; i + 1 < ts.size(); ++i) { + const float t = ts[i]; + const float dt = ts[i + 1] - ts[i]; + const auto pair = f5_dit_forward_cfg( + model_path, y, cond, chunk_ids, t, duration, arch, &dev); + const auto & v_cond = pair.first; + const auto & v_null = pair.second; + for (size_t k = 0; k < y.size(); ++k) { + const float v = v_cond[k] + (v_cond[k] - v_null[k]) * request.cfg_strength; + y[k] += dt * v; + } + } + for (int t = 0; t < ref_frames && t < duration; ++t) { + for (int m = 0; m < kNMel; ++m) { + y[static_cast(t) * kNMel + m] = cond[static_cast(t) * kNMel + m]; + } + } + if (out_final_latent_rows != nullptr) { + *out_final_latent_rows = y; + } + ChunkResult out; + const int gen_frames = std::max(0, duration_real - ref_frames); + out.gen_frames = gen_frames; + out.duration_real = duration_real; + out.gen_mel_rows.resize(static_cast(gen_frames) * kNMel); + for (int t = 0; t < gen_frames; ++t) { + for (int m = 0; m < kNMel; ++m) { + out.gen_mel_rows[static_cast(t) * kNMel + m] = + y[static_cast(ref_frames + t) * kNMel + m]; + } + } + return out; +} + +} // namespace + +F5SynthesisResult f5_synthesize( + const std::string & model_path, + const std::string & vocos_path, + const F5SynthesisRequest & request) { + const auto t0 = std::chrono::steady_clock::now(); + F5SynthesisResult result; + + const std::string dir = std::filesystem::path(model_path).parent_path().string(); + const auto vocab = load_vocab(dir); + + // ref audio -> 24k mono -> normalize to the training RMS -> mel + // (Python F5: audio *= target_rms / rms when rms < target_rms; the + // generated wave is scaled back at the end. Skipping this feeds the model + // a conditioning mel far below its training distribution.) + auto ref24 = resample(request.ref_audio, request.ref_sample_rate, kSampleRate); + double ref_rms = 0.0; + for (const auto v : ref24) ref_rms += double(v) * v; + ref_rms = std::sqrt(ref_rms / std::max(1, ref24.size())); + constexpr double kTargetRms = 0.1; + float ref_gain = 1.0F; + if (ref_rms > 0.0 && ref_rms < kTargetRms) { + ref_gain = static_cast(kTargetRms / ref_rms); + for (auto & v : ref24) v *= ref_gain; + } + auto ref_mel = compute_mel(ref24); + int ref_frames = static_cast(ref_mel.size()) / kNMel; + // The reference may use at most half the frame budget so a meaningful + // generation budget remains. CRITICAL: the ref audio and ref_text must + // stay aligned — truncating the audio while keeping the full transcript + // makes the model speak the UNSAMPLED remainder of the transcript into + // the generated region (observed: EGY ref leaked "استخدمه هيعجبك اوي" + // when its 7.84s audio was cut to 5.46s). + const int kMaxRefFrames = frame_budget(request) / 2; + if (ref_frames > kMaxRefFrames) { + std::fprintf(stderr, + "F5-TTS: reference audio is %.1fs (%d frames) — truncated to %d frames (half the " + "frame budget). The transcript tail will leak into the output; use a reference " + "shorter than %.1fs or raise the frame budget (session option f5_tts.frame_budget).\n", + ref_frames / 93.75, ref_frames, kMaxRefFrames, kMaxRefFrames / 93.75); + std::vector trimmed(static_cast(kMaxRefFrames) * kNMel); + for (int t = 0; t < kMaxRefFrames; ++t) { + for (int m = 0; m < kNMel; ++m) { + trimmed[static_cast(m) * kMaxRefFrames + t] = + ref_mel[static_cast(m) * ref_frames + t]; + } + } + ref_mel = std::move(trimmed); + ref_frames = kMaxRefFrames; + } + + // Voiced ref frames for the pacing rate: ref_frames/ref_chars counts the + // reference's internal pauses as speech time, so a dramatic/paused + // reference (the EGY sample) inflates the frames-per-char rate and the + // model parks the excess duration as long mid-text pauses. Count only + // frames whose log-mel row mean is above the silence floor. + int ref_voiced_frames = 0; + for (int t = 0; t < ref_frames; ++t) { + float acc = 0.0F; + for (int m = 0; m < kNMel; ++m) { + acc += ref_mel[static_cast(m) * ref_frames + t]; + } + if (acc / kNMel > -4.0F) ++ref_voiced_frames; + } + ref_voiced_frames = std::max(1, ref_voiced_frames); + + F5ComputeDevice dev; + dev.use_cuda = request.use_cuda; + dev.device = request.cuda_device; + dev.threads = request.threads; + + // ---- chunk long texts instead of truncating: each chunk is sized by + // the DURATION budget (frames/char from the reference, capped) so no + // chunk ever needs clamping; chunk N+1 is conditioned on the tail of + // chunk N (voice and prosody continuity across seams) ---- + const int ref_chars0 = std::max(1, static_cast(utf8_char_count(request.ref_text))); + const double rate0 = std::max( + std::min(static_cast(ref_voiced_frames) / ref_chars0, 93.75 / 2.5), + 93.75 / 14.0); + const int gen_budget = frame_budget(request) - ref_frames; // frames a chunk may generate + // chars per chunk: budget / rate with a safety margin, so the duration + // estimate NEVER engages the 1024-frame cap (cap = compressed speech). + // The absolute floor is small: a slow reference legitimately means short + // chunks (its 5.5 s reference eats most of the frame budget). + size_t chars_per_chunk = std::max( + 12, static_cast(gen_budget / rate0 * 0.92)); + const std::string gen_text = request.strip_diacritics + ? strip_arabic_diacritics(request.text) + : request.text; + const auto chunks = chunk_text(gen_text, chars_per_chunk); + std::vector all_rows; + const std::string ref_text = apply_ref_trailing_space(request.ref_text); + // Seed semantics match python F5 (seed=None -> fresh randomness every + // run): an unspecified seed draws a random base seed per request, so a + // sampling accident (e.g. a swallowed word like "هرمز") can be re-rolled + // instead of being replayed identically on every request. fixed_seed + // keeps deterministic per-chunk seeds (seed + chunk_index). + const uint32_t base_seed = request.fixed_seed + ? request.seed + : (static_cast(std::random_device{}()) ^ + static_cast(std::chrono::steady_clock::now().time_since_epoch().count())); + // Every chunk is conditioned on the ORIGINAL reference (vanilla F5 + // chunking semantics): chained references drift the voice and decay + // the energy chunk over chunk (observed: two voices + fade to silence). + for (size_t ci = 0; ci < chunks.size(); ++ci) { + const std::string full = assemble_chunk_text(request.dialect, ref_text, chunks[ci]); + const auto chunk_ids = tokenize_text(vocab, full); + auto out = synthesize_chunk( + model_path, request, ref_mel, ref_frames, ref_voiced_frames, chunk_ids, + chunks[ci], ref_text, dev, + base_seed + static_cast(ci), + nullptr, chunks.size() > 1 ? 1.20 : 1.0, + chunks.size() > 1 ? 0 : 40); + if (chunks.size() > 1) { + const char last_ch = chunks[ci].empty() ? ' ' : chunks[ci].back(); + const bool sent_final = last_ch == '.' || last_ch == '!' || last_ch == '?'; + trim_chunk_mel_silence(out.gen_mel_rows, out.gen_frames, sent_final); + } + all_rows.insert(all_rows.end(), out.gen_mel_rows.begin(), out.gen_mel_rows.end()); + } + + result.audio = request.use_cuda + ? vocos_decode_gpu(vocos_path, all_rows, dev) + : vocos_decode(vocos_path, all_rows); + // undo the reference normalization on the output (Python F5 parity) + if (ref_gain != 1.0F) { + double out_rms = 0.0; + for (const auto v : result.audio) out_rms += double(v) * v; + out_rms = std::sqrt(out_rms / std::max(1, result.audio.size())); + if (out_rms > 0.0 && out_rms < kTargetRms) { + const float undo = static_cast(out_rms / kTargetRms); + for (auto & v : result.audio) v *= undo; + } + } + result.sample_rate = kSampleRate; + result.generation_seconds = std::chrono::duration( + std::chrono::steady_clock::now() - t0).count(); + return result; +} + + +#ifdef F5_MEL_TEST +std::vector f5_test_mel(const std::vector & wav) { return compute_mel(wav); } +std::vector f5_test_vocos(const std::string & vp, const std::vector & mel) { return vocos_decode(vp, mel); } +std::vector f5_test_vocos_gpu(const std::string & vp, const std::vector & mel, const F5ComputeDevice & dev) { return vocos_decode_gpu(vp, mel, dev); } +std::vector f5_test_token_ids( + const std::string & model_path, + const std::string & dialect, + const std::string & ref_text, + const std::string & gen_text) { + const std::string dir = std::filesystem::path(model_path).parent_path().string(); + const auto vocab = load_vocab(dir); + const std::string full = assemble_chunk_text( + dialect, apply_ref_trailing_space(ref_text), gen_text); + return tokenize_text(vocab, full); +} +#endif + +} // namespace engine::models::f5_tts diff --git a/src/community_models/f5_tts/weights.cpp b/src/community_models/f5_tts/weights.cpp new file mode 100644 index 00000000..9f0ee7d8 --- /dev/null +++ b/src/community_models/f5_tts/weights.cpp @@ -0,0 +1,84 @@ +#include "engine/community_models/f5_tts/weights.h" + +#include "engine/framework/assets/tensor_source.h" + +#include + +namespace engine::models::f5_tts { +namespace { + +// Strips the "ema_model.transformer." prefix from checkpoint tensor names so +// the framework tensor source resolves them by their torch module paths. +} // namespace + +F5DiTWeights load_dit_weights( + const assets::TensorSource & source, + ggml_backend_t backend, + core::BackendType backend_type) { + F5DiTWeights w; + w.store = std::make_shared( + backend, backend_type, "f5_tts.weights", 2ULL * 1024ULL * 1024ULL * 1024ULL); + + const auto tensor = [&](const std::string & n, + std::initializer_list shape) { + return w.store->load_f32_tensor(source, n, shape); + }; + const auto linear = [&](const std::string & n, int64_t out_f, int64_t in_f) { + modules::LinearWeights lw; + lw.weight = tensor(n + ".weight", {out_f, in_f}); + lw.bias = tensor(n + ".bias", {out_f}); + return lw; + }; + + constexpr int64_t kTextDim = 512; + constexpr int64_t kDim = 1024; + constexpr int64_t kFF = 2048; + constexpr int64_t kMel = 100; + + w.text_embedding = w.store->load_f32_tensor( + source, "text_embed.text_embed.weight", source.require_metadata( + "text_embed.text_embed.weight").shape); // vocab varies per checkpoint + w.vocab_size = w.text_embedding.shape.dims[0]; + w.input_proj = linear("input_embed.proj", kDim, kMel * 2 + kTextDim); + w.cpe0.weight = tensor("input_embed.conv_pos_embed.conv1d.0.weight", {kDim, kDim / 16, 31}); + w.cpe0.bias = tensor("input_embed.conv_pos_embed.conv1d.0.bias", {kDim}); + w.cpe2.weight = tensor("input_embed.conv_pos_embed.conv1d.2.weight", {kDim, kDim / 16, 31}); + w.cpe2.bias = tensor("input_embed.conv_pos_embed.conv1d.2.bias", {kDim}); + w.time0 = linear("time_embed.time_mlp.0", kDim, 256); + w.time2 = linear("time_embed.time_mlp.2", kDim, kDim); + + w.text_blocks.reserve(4); + for (int i = 0; i < 4; ++i) { + const std::string p = "text_embed.text_blocks." + std::to_string(i); + F5TextConvNextWeights b; + b.dwconv.weight = tensor(p + ".dwconv.weight", {kTextDim, 1, 7}); + b.dwconv.bias = tensor(p + ".dwconv.bias", {kTextDim}); + b.norm.weight = tensor(p + ".norm.weight", {kTextDim}); + b.norm.bias = tensor(p + ".norm.bias", {kTextDim}); + b.pw1 = linear(p + ".pwconv1", kDim, kTextDim); + b.pw2 = linear(p + ".pwconv2", kTextDim, kDim); + b.grn_gamma = source.require_f32(p + ".grn.gamma"); + b.grn_beta = source.require_f32(p + ".grn.beta"); + w.text_blocks.push_back(std::move(b)); + } + w.blocks.reserve(22); + for (int i = 0; i < 22; ++i) { + const std::string p = "transformer_blocks." + std::to_string(i); + F5BlockWeights b; + b.attn_norm = linear(p + ".attn_norm.linear", 6 * kDim, kDim); + b.to_q = linear(p + ".attn.to_q", kDim, kDim); + b.to_k = linear(p + ".attn.to_k", kDim, kDim); + b.to_v = linear(p + ".attn.to_v", kDim, kDim); + b.to_out = linear(p + ".attn.to_out.0", kDim, kDim); + b.ff0 = linear(p + ".ff.ff.0.0", kFF, kDim); + b.ff2 = linear(p + ".ff.ff.2", kDim, kFF); + w.blocks.push_back(std::move(b)); + } + w.norm_out = linear("norm_out.linear", 2 * kDim, kDim); + w.proj_out = linear("proj_out", kMel, kDim); + w.store->upload(); + source.release_storage(); + return w; +} + +} // namespace engine::models::f5_tts diff --git a/src/framework/runtime/registry.cpp b/src/framework/runtime/registry.cpp index 7b2623dd..9b7dabcf 100644 --- a/src/framework/runtime/registry.cpp +++ b/src/framework/runtime/registry.cpp @@ -86,6 +86,10 @@ bool ModelRegistry::supports_family(const std::string & family) const noexcept { if (loader->family() == family) { return true; } + const auto aliases = loader->family_aliases(); + if (std::find(aliases.begin(), aliases.end(), family) != aliases.end()) { + return true; + } } return false; } @@ -160,11 +164,15 @@ void ModelRegistry::validate_request(const ModelLoadRequest & request) const { const IVoiceModelLoader * ModelRegistry::find_loader(const ModelLoadRequest & request) const { for (const auto & loader : loaders_) { - if (request.family_hint.has_value() && loader->family() != *request.family_hint) { - continue; - } if (request.family_hint.has_value()) { - return loader.get(); + if (loader->family() == *request.family_hint) { + return loader.get(); + } + const auto aliases = loader->family_aliases(); + if (std::find(aliases.begin(), aliases.end(), *request.family_hint) != aliases.end()) { + return loader.get(); + } + continue; } if (loader->can_load(request)) { return loader.get(); diff --git a/tests/f5_cfg_parity_main.cpp b/tests/f5_cfg_parity_main.cpp new file mode 100644 index 00000000..b5ed184a --- /dev/null +++ b/tests/f5_cfg_parity_main.cpp @@ -0,0 +1,103 @@ +// CFG parity: f5_dit_forward_cfg (batched cond+uncond) must match the real +// f5_tts DiT run with cfg_infer=True — i.e. the null half sees ZEROED cond +// audio and FILLER text ids (python drop_audio_cond=True / drop_text=True). +// Goldens generated by /mnt/ai/f5-parity/gen_golden_cfg.py. +#include "engine/community_models/f5_tts/runtime.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +std::vector load_bin(const std::string & path, size_t expected) { + std::ifstream f(path, std::ios::binary); + if (!f) { + std::fprintf(stderr, "cannot open %s\n", path.c_str()); + std::exit(1); + } + std::vector out(expected); + f.read(reinterpret_cast(out.data()), static_cast(expected * sizeof(float))); + if (static_cast(f.gcount()) != expected * sizeof(float)) { + std::fprintf(stderr, "short read on %s\n", path.c_str()); + std::exit(1); + } + return out; +} + +std::vector load_ids(const std::string & path, size_t expected) { + std::ifstream f(path, std::ios::binary); + if (!f) { + std::fprintf(stderr, "cannot open %s\n", path.c_str()); + std::exit(1); + } + std::vector out(expected); + f.read(reinterpret_cast(out.data()), static_cast(expected * sizeof(int32_t))); + if (static_cast(f.gcount()) != expected * sizeof(int32_t)) { + std::fprintf(stderr, "short read on %s\n", path.c_str()); + std::exit(1); + } + return out; +} + +double cosine(const std::vector & a, const std::vector & b) { + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < a.size(); ++i) { + dot += static_cast(a[i]) * b[i]; + na += static_cast(a[i]) * a[i]; + nb += static_cast(b[i]) * b[i]; + } + return dot / (std::sqrt(na) * std::sqrt(nb) + 1e-12); +} + +double max_abs(const std::vector & a, const std::vector & b) { + double m = 0; + for (size_t i = 0; i < a.size(); ++i) { + m = std::max(m, std::abs(static_cast(a[i]) - b[i])); + } + return m; +} + +} // namespace + +int main(int argc, char ** argv) { + const std::string gold = "/mnt/ai/f5-parity/golden"; + const std::string ckpt = argc > 1 ? argv[1] : "/mnt/ai/models/Habibi-TTS/Unified/model_200000.safetensors"; + constexpr int N = 64, MEL = 100, NT = 24; + + const auto x = load_bin(gold + "/cfg_input_x.bin", N * MEL); + const auto cond = load_bin(gold + "/cfg_input_cond.bin", N * MEL); + const auto ids = load_ids(gold + "/cfg_input_ids.bin", NT); + const auto want_cond = load_bin(gold + "/cfg_out_cond.bin", N * MEL); + const auto want_null = load_bin(gold + "/cfg_out_null.bin", N * MEL); + + engine::models::f5_tts::F5Architecture arch; + engine::models::f5_tts::F5ComputeDevice dev; + dev.use_cuda = std::getenv("F5_CUDA") != nullptr; + dev.device = 1; // GPU 1 is the TTS GPU on this rig + + const auto [got_cond, got_null] = engine::models::f5_tts::f5_dit_forward_cfg( + ckpt, x, cond, ids, 0.42F, N, arch, &dev); + + int failures = 0; + auto check = [&](const char * name, const std::vector & got, const std::vector & want) { + const double c = cosine(got, want); + const double m = max_abs(got, want); + const bool ok = c >= 0.999 && std::isfinite(c); + std::printf("%-14s cosine=%.6f maxabs=%.5f %s\n", name, c, m, ok ? "OK" : "FAIL"); + if (!ok) ++failures; + }; + check("cfg_cond", got_cond, want_cond); + check("cfg_null", got_null, want_null); // fails if the null half sees real cond/text + + if (failures > 0) { + std::printf("CFG PARITY FAILED (%d halves)\n", failures); + return 1; + } + std::printf("CFG PARITY PASSED\n"); + return 0; +} diff --git a/tests/f5_e2e_main.cpp b/tests/f5_e2e_main.cpp new file mode 100644 index 00000000..de0366e1 --- /dev/null +++ b/tests/f5_e2e_main.cpp @@ -0,0 +1,135 @@ +// E2E: Habibi Arabic synthesis via f5_synthesize + vocos, writes WAV. +#include "engine/community_models/f5_tts/synthesize.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +// minimal WAV reader (16-bit PCM mono/stereo) +struct Wav { + int sample_rate = 0; + int channels = 1; + std::vector samples; // mono mixdown +}; + +bool read_wav(const std::string & path, Wav & out) { + std::ifstream f(path, std::ios::binary); + if (!f) return false; + std::vector data((std::istreambuf_iterator(f)), std::istreambuf_iterator()); + if (data.size() < 44 || std::memcmp(data.data(), "RIFF", 4) != 0) return false; + size_t pos = 12; + int bits = 16; + while (pos + 8 <= data.size()) { + const std::string id(data.data() + pos, 4); + const uint32_t sz = *reinterpret_cast(data.data() + pos + 4); + if (id == "fmt ") { + const char * p = data.data() + pos + 8; // chunk payload + out.channels = *reinterpret_cast(p + 2); + out.sample_rate = *reinterpret_cast(p + 4); + bits = *reinterpret_cast(p + 14); + } else if (id == "data") { + if (bits == 0) return false; + const size_t n = sz / (bits / 8); + const auto * p = reinterpret_cast(data.data() + pos + 8); + out.samples.resize(n); + for (size_t i = 0; i < n; ++i) { + out.samples[i] = static_cast(p[i]) / 32768.0F; + } + break; + } + pos += 8 + sz + (sz & 1); + } + return !out.samples.empty(); +} + +void write_wav(const std::string & path, const std::vector & mono, int sr) { + std::ofstream f(path, std::ios::binary); + const uint32_t n = static_cast(mono.size()); + const uint32_t data_bytes = n * 2; + f.write("RIFF", 4); + const uint32_t riff = 36 + data_bytes; + f.write(reinterpret_cast(&riff), 4); + f.write("WAVEfmt ", 8); + const uint32_t fmt = 16; + f.write(reinterpret_cast(&fmt), 4); + const uint16_t pcm = 1, ch = 1; + const uint16_t bps = 16; + const uint32_t br = sr * 2; + const uint16_t align = 2; + f.write(reinterpret_cast(&pcm), 2); + f.write(reinterpret_cast(&ch), 2); + f.write(reinterpret_cast(&sr), 4); + f.write(reinterpret_cast(&br), 4); + f.write(reinterpret_cast(&align), 2); + f.write(reinterpret_cast(&bps), 2); + f.write("data", 4); + f.write(reinterpret_cast(&data_bytes), 4); + for (const float v : mono) { + const auto s = static_cast(std::clamp(v * 32767.0F, -32768.0F, 32767.0F)); + f.write(reinterpret_cast(&s), 2); + } +} + +} // namespace + +int main(int argc, char ** argv) { + std::fprintf(stderr, "start\n"); fflush(stderr); + const std::string model = argc > 1 ? argv[1] : "/mnt/ai/models/Habibi-TTS/Unified/model_200000.safetensors"; + const std::string vocos = argc > 2 ? argv[2] : "/mnt/ai/models/vocos-mel-24khz/vocos.safetensors"; + const std::string ref = argc > 3 ? argv[3] : "/tmp/f5ref/pkg/habibi/habibi_tts/assets/IRQ.wav"; + const std::string out_path = argc > 4 ? argv[4] : "/tmp/habibi_e2e.wav"; + + Wav ref_wav; + if (!read_wav(ref, ref_wav)) { + std::fprintf(stderr, "cannot read ref wav %s\n", ref.c_str()); + return 1; + } + std::printf("ref: %d Hz, %zu samples (%.2fs)\n", ref_wav.sample_rate, ref_wav.samples.size(), + static_cast(ref_wav.samples.size()) / ref_wav.sample_rate); + + engine::models::f5_tts::F5SynthesisRequest req; + // Same sentence as the Python reference (/tmp/python_ref.py). Keep the + // UTF-8 literal (NOT byte escapes): the previous escaped version had + // "habibi" misspelled with ه instead of ح and invalid "\D8"/"\D9" + // sequences that compiled to literal "D8"/"D9" garbage inside تجربة and + // جي بي يو — the source of the mispronounced letters. + req.text = std::getenv("F5_TEXT") ? std::getenv("F5_TEXT") + : "أهلا، هذه تجربة للنطق باللغة العربية، من نموذج حبيبي، داخل أوديو سي بي بي، على مجموعة جي بي يو بالبشرة."; + req.dialect = std::getenv("F5_DIALECT") ? std::getenv("F5_DIALECT") : "UNK"; + req.ref_audio = ref_wav.samples; + req.ref_sample_rate = ref_wav.sample_rate; + req.ref_text = std::getenv("F5_REF_TEXT") ? std::getenv("F5_REF_TEXT") + : "يعني ااا ما نقدر ناخذ وقت أكثر، ااا لأنه شروط كلش يحتاجلها وقت."; + if (std::getenv("F5_LONG") != nullptr) { + // long-text test: ~4x the cap; exercises chunking + chaining + req.text = req.text + " " + req.text + " " + req.text + " " + req.text; + } + req.steps = std::getenv("F5_STEPS") ? std::atoi(std::getenv("F5_STEPS")) : 16; + req.cfg_strength = 2.0F; + req.seed = 42; + req.fixed_seed = std::getenv("F5_RANDSEED") == nullptr; // F5_RANDSEED=1: random per run + req.use_cuda = std::getenv("F5_CUDA") != nullptr; + req.cuda_device = 1; + + std::printf("synthesizing...\n"); + fflush(stdout); + // first call = graph build; second = cached graphs (server steady state) + const auto warm = engine::models::f5_tts::f5_synthesize(model, vocos, req); + (void)warm; + const auto result = engine::models::f5_tts::f5_synthesize(model, vocos, req); + std::printf("generated %.2fs audio in %.2fs wall (%.2fx RTF)\n", + static_cast(result.audio.size()) / result.sample_rate, + result.generation_seconds, + result.generation_seconds / (static_cast(result.audio.size()) / result.sample_rate)); + write_wav(out_path, result.audio, result.sample_rate); + std::printf("written %s\n", out_path.c_str()); + return 0; +} diff --git a/tests/f5_parity_main.cpp b/tests/f5_parity_main.cpp new file mode 100644 index 00000000..d5d22324 --- /dev/null +++ b/tests/f5_parity_main.cpp @@ -0,0 +1,125 @@ +// Parity harness: run f5_dit_forward on golden inputs, compare stages vs goldens. +#include "engine/community_models/f5_tts/runtime.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +std::vector load_bin(const std::string & path, size_t expected) { + std::ifstream f(path, std::ios::binary); + if (!f) { + std::fprintf(stderr, "cannot open %s\n", path.c_str()); + std::exit(1); + } + std::vector out(expected); + f.read(reinterpret_cast(out.data()), static_cast(expected * sizeof(float))); + if (static_cast(f.gcount()) != expected * sizeof(float)) { + std::fprintf(stderr, "short read on %s\n", path.c_str()); + std::exit(1); + } + return out; +} + +double cosine(const std::vector & a, const std::vector & b) { + double dot = 0, na = 0, nb = 0; + for (size_t i = 0; i < a.size(); ++i) { + dot += static_cast(a[i]) * b[i]; + na += static_cast(a[i]) * a[i]; + nb += static_cast(b[i]) * b[i]; + } + return dot / (std::sqrt(na) * std::sqrt(nb) + 1e-12); +} + +double max_abs(const std::vector & a, const std::vector & b) { + double m = 0; + for (size_t i = 0; i < a.size(); ++i) { + m = std::max(m, std::abs(static_cast(a[i]) - b[i])); + } + return m; +} + +// golden tensors are row-major [T, F]; C++ taps are column [F, T] +std::vector col_from_row(const std::vector & row, int T, int F) { + std::vector col(row.size()); + for (int t = 0; t < T; ++t) + for (int f = 0; f < F; ++f) + col[static_cast(f) * T + t] = row[static_cast(t) * F + f]; + return col; +} + +} // namespace + +int main(int argc, char ** argv) { + const std::string gold = "/mnt/ai/f5-parity/golden"; + const std::string ckpt = argc > 1 ? argv[1] : "/mnt/ai/models/Habibi-TTS/Unified/model_200000.safetensors"; + + const auto x = load_bin(gold + "/input_x.bin", 64 * 100); + const auto cond = load_bin(gold + "/input_cond.bin", 64 * 100); + std::vector ids(24); + { + const auto raw = load_bin(gold + "/input_ids.bin", 24); + for (int i = 0; i < 24; ++i) ids[i] = static_cast(raw[i]); + } + + engine::models::f5_tts::F5Architecture arch; + engine::models::f5_tts::F5DebugTaps taps; + std::vector t_text_embed, t_text_convnext, t_text_padded, t_input_embed, t_time_embed, t_block0, t_block21; + taps.text_embed = &t_text_embed; + taps.text_convnext = &t_text_convnext; + taps.text_padded = &t_text_padded; + taps.input_embed = &t_input_embed; + taps.time_embed = &t_time_embed; + taps.block0 = &t_block0; + taps.block21 = &t_block21; + + const bool use_cuda = std::getenv("F5_CUDA") != nullptr; + const bool with_taps = std::getenv("F5_RAW_TAPS") != nullptr; // module graph: taps unwired + engine::models::f5_tts::F5ComputeDevice dev; + dev.use_cuda = use_cuda; + dev.device = 1; // GPU 1 is the TTS GPU on this rig + const auto out = engine::models::f5_tts::f5_dit_forward( + ckpt, x, cond, ids, 0.42F, 64, arch, false, false, (with_taps ? &taps : nullptr), &dev); + + int failures = 0; + auto check = [&](const char * name, const std::vector & mine_col, const std::vector & golden_row, int T, int F) { + (void) T; (void) F; + // Both golden and taps are stored with time as the slow axis ([t][f]); + // no transpose needed. (The F5Runtime returns column layout only for + // the final output; stage taps keep the same [t][f] memory.) + const double c = cosine(mine_col, golden_row); + const double m = max_abs(mine_col, golden_row); + if (mine_col.empty()) { + std::printf("%-18s (tap not wired in module graph; skipped)\n", name); + return; + } + const bool ok = c >= 0.999 && std::isfinite(c); + std::printf("%-18s cosine=%.6f maxabs=%.5f %s\n", name, c, m, ok ? "OK" : "FAIL"); + if (!ok) failures++; + }; + + check("01_text_embed", t_text_embed, load_bin(gold + "/01_text_embed.bin", 24 * 512), 24, 512); + check("02_text_convnext", t_text_convnext, load_bin(gold + "/02_text_after_convnext.bin", 24 * 512), 24, 512); + check("03_text_padded", t_text_padded, load_bin(gold + "/03_text_padded.bin", 64 * 512), 64, 512); + check("04_input_embed", t_input_embed, load_bin(gold + "/04_input_embed.bin", 64 * 1024), 64, 1024); + check("05_time_embed", t_time_embed, load_bin(gold + "/05_time_embed.bin", 1024), 1, 1024); + check("07_block0", t_block0, load_bin(gold + "/07_block0_out.bin", 64 * 1024), 64, 1024); + check("07_block21", t_block21, load_bin(gold + "/07_block21_out.bin", 64 * 1024), 64, 1024); + { + // out tensor ne [MEL, N] => memory (m,n) at n*MEL + m, i.e. raw bytes + // are already golden's [n][m] row-major layout. Direct compare. + const std::vector out_rows = out; + check("08_final_out", out_rows, load_bin(gold + "/08_final_out.bin", 64 * 100), 64, 100); + } + + if (failures > 0) { + std::printf("PARITY FAILED (%d stages)\n", failures); + return 1; + } + std::printf("PARITY PASSED\n"); + return 0; +} diff --git a/tests/f5_tokenizer_main.cpp b/tests/f5_tokenizer_main.cpp new file mode 100644 index 00000000..ad6e1462 --- /dev/null +++ b/tests/f5_tokenizer_main.cpp @@ -0,0 +1,84 @@ +// Tokenizer parity: the C++ text pipeline (dialect wrap + ref trailing-space +// rule + UTF-8 char tokenization) must reproduce python list_str_to_idx ids +// exactly. Goldens generated by /mnt/ai/f5-parity/gen_golden_tokens.py. +#include "engine/community_models/f5_tts/synthesize.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +std::string read_text(const std::string & path) { + std::ifstream f(path, std::ios::binary); + if (!f) { + std::fprintf(stderr, "cannot open %s\n", path.c_str()); + std::exit(1); + } + return std::string((std::istreambuf_iterator(f)), std::istreambuf_iterator()); +} + +std::vector read_ids(const std::string & path) { + std::ifstream f(path, std::ios::binary); + if (!f) { + std::fprintf(stderr, "cannot open %s\n", path.c_str()); + std::exit(1); + } + const std::string raw((std::istreambuf_iterator(f)), std::istreambuf_iterator()); + std::vector ids(raw.size() / 4); + std::memcpy(ids.data(), raw.data(), ids.size() * 4); + return ids; +} + +} // namespace + +int main(int argc, char ** argv) { + const std::string gold = argc > 1 ? argv[1] : "/mnt/ai/f5-parity/golden/tokens"; + const std::string model = argc > 2 + ? argv[2] : "/mnt/ai/models/Habibi-TTS/Unified/model_200000.safetensors"; + + const int count = std::atoi(read_text(gold + "/count").c_str()); + int failures = 0; + for (int i = 0; i < count; ++i) { + const std::string base = gold + "/case" + std::to_string(i); + const std::string dialect = read_text(base + ".dialect"); + const std::string ref = read_text(base + ".ref"); + const std::string gen = read_text(base + ".gen"); + const auto want = read_ids(base + ".ids"); + + const auto got = engine::models::f5_tts::f5_test_token_ids(model, dialect, ref, gen); + bool ok = got.size() == want.size(); + size_t first_bad = 0; + if (ok) { + for (size_t k = 0; k < want.size(); ++k) { + if (got[k] != want[k]) { + ok = false; + first_bad = k; + break; + } + } + } + std::printf("case%d dialect=%-3s ids=%zu %s\n", i, dialect.c_str(), want.size(), + ok ? "OK" : "FAIL"); + if (!ok) { + ++failures; + if (got.size() != want.size()) { + std::printf(" length mismatch: got %zu want %zu\n", got.size(), want.size()); + } else { + std::printf(" first mismatch at %zu: got %d want %d\n", + first_bad, got[first_bad], want[first_bad]); + } + } + } + if (failures > 0) { + std::printf("TOKENIZER PARITY FAILED (%d/%d cases)\n", failures, count); + return 1; + } + std::printf("TOKENIZER PARITY PASSED (%d cases)\n", count); + return 0; +} diff --git a/tools/convert_f5_tts.py b/tools/convert_f5_tts.py new file mode 100644 index 00000000..4e04bdaa --- /dev/null +++ b/tools/convert_f5_tts.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python3 +"""Convert Habibi/F5-TTS safetensors checkpoints to audio.cpp GGUF packages. + +Produces one self-contained GGUF per checkpoint with two tensor namespaces: + transformer.* — the DiT flow-matching transformer (raw EMA torch names) + vocos.* — the Vocos mel vocoder + +and copies vocab.txt alongside it, giving a complete model directory that +audiocpp_cli / audiocpp_server load with --family f5_tts (aliases: habibi, +habibi_tts). A standalone vocoder GGUF can also be produced for use with the +original safetensors checkpoints. + +Examples: + # unified checkpoint (default package) + python3 tools/convert_f5_tts.py \ + --checkpoint /models/Habibi-TTS/Unified \ + --vocos /models/vocos-mel-24khz/vocos.safetensors \ + --converter build/bin/audiocpp_gguf --name habibi-unified + + # every specialized dialect checkpoint under a Habibi-TTS root + python3 tools/convert_f5_tts.py --checkpoint-root /models/Habibi-TTS \ + --vocos /models/vocos-mel-24khz/vocos.safetensors \ + --converter build/bin/audiocpp_gguf + + # standalone vocoder package + python3 tools/convert_f5_tts.py --vocos-only \ + --vocos /models/vocos-mel-24khz/vocos.safetensors \ + --converter build/bin/audiocpp_gguf +""" +import argparse +import shutil +import subprocess +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent.parent +SPEC = REPO_ROOT / "model_specs" / "f5_tts.json" + + +def convert(converter: Path, checkpoint: Path, vocos: Path, output: Path, + quant_type: str, overwrite: bool) -> None: + ckpt = _find_checkpoint(checkpoint) + command = [ + str(converter), + "--input", f"transformer={ckpt}", + "--input", f"vocos={vocos}", + "--root", str(checkpoint), + "--family", "f5_tts", + "--model-spec", str(SPEC), + "--type", quant_type, + "--output", str(output), + ] + if overwrite: + command.append("--overwrite") + print("+", " ".join(command)) + subprocess.run(command, check=True) + shutil.copyfile(checkpoint / "vocab.txt", output.parent / "vocab.txt") + print(f"copied vocab.txt -> {output.parent}") + + +def convert_vocos(converter: Path, vocos: Path, output: Path, + quant_type: str, overwrite: bool) -> None: + command = [ + str(converter), + "--input", f"vocos={vocos}", + "--root", str(vocos.parent), + "--type", quant_type, + "--allow-missing-model-spec", + "--no-sidecars", + "--output", str(output), + ] + if overwrite: + command.append("--overwrite") + print("+", " ".join(command)) + subprocess.run(command, check=True) + + +def _find_checkpoint(directory: Path) -> Path: + candidates = sorted(directory.glob("*.safetensors")) + if not candidates: + raise SystemExit(f"no .safetensors checkpoint in {directory}") + return candidates[-1] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--checkpoint", type=Path, + help="one checkpoint directory (model_*.safetensors + vocab.txt)") + parser.add_argument("--checkpoint-root", type=Path, + help="Habibi-TTS root: converts Unified/ and every Specialized/* dir") + parser.add_argument("--vocos", type=Path, required=True, + help="vocos.safetensors (bundled into each checkpoint GGUF)") + parser.add_argument("--vocos-only", action="store_true", + help="only convert the standalone vocoder package") + parser.add_argument("--converter", type=Path, required=True, + help="path to the audiocpp_gguf binary") + parser.add_argument("--output-dir", type=Path, default=Path("gguf-out")) + parser.add_argument("--type", default="orig", + choices=["orig", "f16", "bf16", "q8_0", "q2_k", "q3_k", + "q4_k", "q5_k", "q6_k"], + help="GGUF storage type (default orig = keep f32)") + parser.add_argument("--name", help="GGUF base name for --checkpoint (default: dir name, lowercased)") + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + + converter = args.converter.resolve() + if not converter.is_file(): + raise SystemExit(f"converter not found: {converter} (build target audiocpp_gguf)") + if not args.vocos.is_file(): + raise SystemExit(f"vocos checkpoint not found: {args.vocos}") + args.output_dir.mkdir(parents=True, exist_ok=True) + + if args.vocos_only: + convert_vocos(converter, args.vocos.resolve(), + args.output_dir / f"vocos-mel-24khz-{args.type}.gguf", + args.type, args.overwrite) + return + + jobs: list[tuple[Path, str]] = [] + if args.checkpoint: + name = args.name or args.checkpoint.name.lower() + jobs.append((args.checkpoint, name)) + elif args.checkpoint_root: + root = args.checkpoint_root + if (root / "Unified").is_dir(): + jobs.append((root / "Unified", "habibi-unified")) + for d in sorted((root / "Specialized").glob("*")): + if d.is_dir(): + jobs.append((d, f"habibi-{d.name.lower()}")) + else: + raise SystemExit("pass --checkpoint or --checkpoint-root (or --vocos-only)") + + for checkpoint, name in jobs: + out_dir = args.output_dir / name + out_dir.mkdir(parents=True, exist_ok=True) + convert(converter, checkpoint, args.vocos.resolve(), + out_dir / f"{name}-{args.type}.gguf", args.type, args.overwrite) + + # standalone vocoder package alongside the checkpoints + convert_vocos(converter, args.vocos.resolve(), + args.output_dir / f"vocos-mel-24khz-{args.type}.gguf", + args.type, args.overwrite) + print("\nDone. Upload the .gguf files + each package's vocab.txt to the hosting repo.") + + +if __name__ == "__main__": + main()