From dc1a845595e6f2baa80fa39cf017d147a2873f20 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:18:21 -0400 Subject: [PATCH] Add LoRA support and GGUF merge tools --- app/gguf/main.cpp | 43 +++ docs/asr.md | 2 + docs/lora.md | 141 ++++++++ docs/models/cosyvoice3.md | 2 + docs/models/f5_tts.md | 8 + docs/models/fish_audio.md | 11 + docs/models/outetts.md | 2 + docs/models/qwen3_asr.md | 2 + docs/models/qwen3_tts.md | 2 + docs/models/voxcpm1.md | 2 + docs/tts.md | 4 + .../engine/community_models/f5_tts/runtime.h | 7 + .../engine/community_models/f5_tts/session.h | 3 + .../community_models/f5_tts/synthesize.h | 1 + .../engine/framework/assets/tensor_source.h | 8 + include/engine/models/qwen3_tts/assets.h | 4 + model_specs/cosyvoice3.json | 13 + model_specs/f5_tts.json | 13 + model_specs/fish_audio.json | 13 + model_specs/outetts.json | 13 + model_specs/qwen3_asr.json | 13 + model_specs/qwen3_tts.json | 13 + model_specs/voxcpm1.json | 13 + model_specs/voxcpm2.json | 13 + model_specs/yue2.json | 9 +- src/community_models/f5_tts/runtime.cpp | 176 ++++++---- src/community_models/f5_tts/session.cpp | 12 + src/community_models/f5_tts/synthesize.cpp | 1 + src/community_models/outetts/session.cpp | 27 +- src/community_models/voxcpm1/session.cpp | 76 +++- src/framework/assets/tensor_source.cpp | 56 ++- src/framework/assets/weight_adapter.cpp | 7 +- src/models/cosyvoice3/session.cpp | 79 ++++- src/models/fish_audio/loader.cpp | 2 + src/models/fish_audio/session.cpp | 42 ++- src/models/qwen3_asr/loader.cpp | 2 + src/models/qwen3_asr/session.cpp | 38 +- src/models/qwen3_tts/assets.cpp | 62 ++++ src/models/qwen3_tts/loader.cpp | 2 + src/models/qwen3_tts/session.cpp | 4 +- src/models/vibevoice/lora.cpp | 25 +- src/models/voxcpm2/loader.cpp | 2 + src/models/voxcpm2/session.cpp | 50 ++- tests/lora/test_merge_lora.py | 326 ++++++++++++++++++ tests/unittests/test_gguf_tensor_source.cpp | 59 ++++ .../test_torch_bin_storage_offset.cpp | 26 ++ tools/lora/bindings.py | 257 ++++++++++++++ tools/lora/checkpoint.py | 36 ++ tools/lora/merge_lora.py | 317 +++++++++++++++++ tools/lora/prepare_fish_audio_adapter.py | 77 +++++ 50 files changed, 2008 insertions(+), 108 deletions(-) create mode 100644 docs/lora.md create mode 100644 tests/lora/test_merge_lora.py create mode 100644 tools/lora/bindings.py create mode 100644 tools/lora/checkpoint.py create mode 100644 tools/lora/merge_lora.py create mode 100644 tools/lora/prepare_fish_audio_adapter.py diff --git a/app/gguf/main.cpp b/app/gguf/main.cpp index 55c285bfc..1bb2ba956 100644 --- a/app/gguf/main.cpp +++ b/app/gguf/main.cpp @@ -1,6 +1,9 @@ #include "engine/framework/assets/tensor_source.h" +#include "engine/framework/assets/weight_adapter.h" +#include "engine/framework/debug/trace.h" #include "engine/framework/io/filesystem.h" #include "engine/framework/io/json.h" +#include "engine/framework/runtime/options.h" #include #include @@ -578,6 +581,9 @@ void print_usage() { "[--bnb-nf4-type q8_0] [--exclude-prefix ] " "[--keep-type *=] " "[--fold-weight-norm *] " + "[--weight-adapter ] [--adapter-strength ] " + "[--adapter-merge-mode delta_f32|rankwise_f32|rounded_bf16_delta] " + "[--log] " "[--overwrite] [--no-sidecars] " "[--allow-missing-model-spec]\n" << " audiocpp_gguf --inspect \n"; @@ -596,6 +602,11 @@ int main(int argc, char ** argv) { std::optional family; std::string type; engine::assets::GgufConversionOptions conversion_options; + std::filesystem::path adapter_path; + engine::assets::WeightAdapterReadOptions adapter_options; + adapter_options.require_scaling_metadata = true; + auto adapter_mode = engine::assets::WeightMergeMode::DeltaF32; + bool adapter_options_set = false; bool overwrite = false; bool embed_sidecars = true; bool allow_missing_model_spec = false; @@ -605,6 +616,10 @@ int main(int argc, char ** argv) { print_usage(); return 0; } + if (arg == "--log") { + engine::debug::configure_logging({true}); + continue; + } if (arg == "--overwrite") { overwrite = true; continue; @@ -620,6 +635,7 @@ int main(int argc, char ** argv) { if ((arg == "--input" || arg == "--output" || arg == "--type" || arg == "--inspect" || arg == "--root" || arg == "--sidecar" || arg == "--family" || arg == "--model-spec" || arg == "--model-spec-override" || arg == "--bnb-nf4-type" || arg == "--exclude-prefix" || arg == "--keep-type" || + arg == "--weight-adapter" || arg == "--adapter-strength" || arg == "--adapter-merge-mode" || arg == "--fold-weight-norm") && i + 1 < argc) { const std::string value = argv[++i]; @@ -633,6 +649,19 @@ int main(int argc, char ** argv) { output = value; else if (arg == "--type") type = value; + else if (arg == "--weight-adapter") + adapter_path = std::filesystem::u8path(value); + else if (arg == "--adapter-strength") { + adapter_options.strength = *engine::runtime::parse_finite_float_option( + {{"adapter_strength", value}}, {"adapter_strength"}); + adapter_options_set = true; + } else if (arg == "--adapter-merge-mode") { + if (value == "delta_f32") adapter_mode = engine::assets::WeightMergeMode::DeltaF32; + else if (value == "rankwise_f32") adapter_mode = engine::assets::WeightMergeMode::RankWiseF32; + else if (value == "rounded_bf16_delta") adapter_mode = engine::assets::WeightMergeMode::RoundedBF16Delta; + else throw std::runtime_error("unsupported adapter merge mode: " + value); + adapter_options_set = true; + } else if (arg == "--bnb-nf4-type") conversion_options.bnb_nf4_type = engine::assets::parse_tensor_storage_type(value); else if (arg == "--exclude-prefix") @@ -737,6 +766,20 @@ int main(int argc, char ** argv) { *conversion_options.bnb_nf4_type != engine::assets::TensorStorageType::Q8_0) { throw std::runtime_error("--bnb-nf4-type currently supports q8_0"); } + if (adapter_path.empty() && adapter_options_set) { + throw std::runtime_error("adapter strength/mode requires --weight-adapter"); + } + if (!adapter_path.empty()) { + auto adapter = engine::assets::read_weight_adapter(adapter_path, adapter_options); + if (adapter.updates.empty() && adapter.replacements.empty()) { + throw std::runtime_error("weight adapter contains no updates or replacements"); + } + for (auto & entry : adapter.updates) { + entry.second.merge_mode = adapter_mode; + } + conversion_options.weight_adapter = std::make_shared(std::move(adapter)); + conversion_options.weight_adapter_merge_mode = adapter_mode; + } engine::assets::convert_tensor_sources_to_gguf(inputs, output, storage_type, overwrite, embed_sidecars, resolved_sidecar_root, sidecars, embedded_model_spec, conversion_options); diff --git a/docs/asr.md b/docs/asr.md index a861501d1..508bbd0ba 100644 --- a/docs/asr.md +++ b/docs/asr.md @@ -1,5 +1,7 @@ # ASR Models +For Qwen3-ASR adapter loading and merged GGUF export, see the [LoRA guide](lora.md). + | Model | Family | Mode(s) | Quick Start | |---|---|---|---| | Audio Flamingo 3 / Next | `audio_flamingo` | offline | [Audio Flamingo](models/audio_flamingo.md) | diff --git a/docs/lora.md b/docs/lora.md new file mode 100644 index 000000000..01913453c --- /dev/null +++ b/docs/lora.md @@ -0,0 +1,141 @@ +# LoRA + +LoRA adapters must match the base checkpoint's architecture, size, and tensor +layout. An adapter trained for another release of the same family is not +automatically compatible. + +## Supported Models + +| Model | Family | Session options | Adapter format and target | +|---|---|---|---| +| Qwen3-TTS | `qwen3_tts` | `qwen3_tts.lora`, `qwen3_tts.lora_scale` | PEFT Talker/Subtalker adapters. Custom-voice packages can also include `speaker_embedding.safetensors` and `tts_config.json`. | +| Qwen3-ASR | `qwen3_asr` | `qwen3_asr.lora`, `qwen3_asr.lora_scale` | PEFT encoder/decoder adapters; both original and Transformers weight layouts are supported. | +| VoxCPM2 | `voxcpm2` | `voxcpm2.lora`, `voxcpm2.lora_scale` | Native `lora_weights.safetensors` with adjacent `lora_config.json`; generator only, not AudioVAE. | +| VoxCPM-0.5B | `voxcpm1` | `voxcpm1.lora`, `voxcpm1.lora_scale` | Native VoxCPM adapter and `lora_config.json`; upstream tensor names are mapped to the GGUF layout. | +| Fish Audio S2 Pro | `fish_audio` | `fish_audio.lora`, `fish_audio.lora_scale` | Prepared Fish Speech adapters for slow/fast decoder projections and embeddings; codec weights are unchanged. | +| CosyVoice3 | `cosyvoice3` | `cosyvoice3.lora`, `cosyvoice3.lora_scale` | PEFT LLM adapters, or a native checkpoint with `lora_weights.pt`, `config.json`, `speech_embedding.pt`, and `llm_decoder.pt`. | +| F5-TTS | `f5_tts` | `f5_tts.lora`, `f5_tts.lora_scale` | PEFT DiT adapters. Reconvert the matching original checkpoint with the current converter; do not use the currently published GGUF packages as the base. Vocos is unchanged. | +| OuteTTS | `outetts` | `outetts.lora`, `outetts.lora_scale` | PEFT Llama adapters for the matching OuteTTS checkpoint. DAC is unchanged. | +| YuE2 | `yue2` | `yue2.ar_lora`, `yue2.ar_lora_scale`, `yue2.nar_lora`, `yue2.nar_lora_scale` | Separate AR/NAR adapters, including NAR projection replacements. | +| ACE-Step | `ace_step` | `ace_step.lora`, `ace_step.lora_scale` | DiT matrix LoRA/LoKr in PEFT or ComfyUI/LyCORIS format. | +| VibeVoice | `vibevoice` | `vibevoice.lora`, `vibevoice.lora_scale` | PEFT language-model adapters with optional diffusion-head and connector replacements. Also accepted as load options. | + +## Session Loading + +Container format and adapter layout are independent. Shared tensor I/O accepts +safetensors and PyTorch checkpoints (`.pt`, `.pth`, `.bin`, `.ckpt`); the suffix +does not select PEFT, native LoRA, or LoKr semantics. Runtime loading expects a +tensor state dictionary, not an arbitrary training checkpoint. + +Directory discovery uses logical names such as `adapter_model`, `lora_weights`, +and `speaker_embedding`, with the supported extensions. It prefers safetensors, +then `.pt`, `.pth`, `.bin`, and `.ckpt`. A malformed selected file is an error; +the loader does not silently try another file. The filenames below show the +usual upstream packaging, not a restriction to that container. + +CosyVoice3's native `lora_weights` package uses `config.json` and requires +speech-embedding/output-head sidecars. PEFT uses `adapter_config.json`. +When both layouts are present, directory loading and export select the native +package. An explicit file selects that file; `lora_weights` retains native +configuration and required sidecars, while `adapter_model` uses PEFT metadata. + +Use `--session-option` on the CLI, or the same keys in a server model's +`session_options`. Adapters are applied once while loading session weights; +they do not add inference graph operations or modify the base files. +Pre-v1 GGUF packages still require a current external or compiled model spec. + +```bash +audiocpp_cli --family qwen3_asr --task asr --model base.gguf \ + --backend cuda --audio speech.wav \ + --session-option qwen3_asr.lora=adapters/domain \ + --session-option qwen3_asr.lora_scale=1 --log +``` + +Keep PEFT's `adapter_config.json` next to `adapter_model.safetensors`. For native +VoxCPM, point to the directory containing `lora_weights.safetensors` and +`lora_config.json`. Fish Speech Lightning checkpoints need preparation: + +```bash +python tools/lora/prepare_fish_audio_adapter.py \ + --input lora.ckpt --alpha 16 --output adapter.safetensors +``` + +Use the training alpha, not alpha/r. The exporter handles Fish embedding-factor +orientation and records scaling in the output. + +For Qwen3-TTS custom voices, keep the trained speaker embedding and configuration +with the adapter and select the configured name with `voice_id`. A sidecar-backed +custom voice uses the TTS route, not the Base checkpoint's reference-clone prompt. + +## Scaling + +The new integrations use `lora_scale=1` as a multiplier for the adapter's +alpha/r coefficient; PEFT rsLoRA uses alpha/sqrt(rank). Full-weight replacements +are not scaled. To restore the unmodified model, omit the adapter rather than +setting its strength to zero when it includes replacements. + +Existing controls retain their behavior: VibeVoice's optional `lora_scale` +overrides the effective coefficient, YuE2 native factors use their direct +AR/NAR strengths, and ACE-Step uses an adapter-strength multiplier with zero +disabling the adapter. + +YuE2 export supports custom strengths for A/B adapters; LoKr export currently +requires unit strength. + +## Merged GGUF + +Use the shared model-aware merge tool for every family in the table: + +The tool requires PyTorch, `safetensors`, `gguf`, and a built `audiocpp_gguf`. +Fish Lightning preparation also requires `omegaconf`. + +```bash +python tools/lora/merge_lora.py --family qwen3_tts \ + --base models/qwen3-tts-base-bf16.gguf --adapter adapters/my-voice \ + --type q8_0 --output models/qwen3-tts-my-voice-q8_0.gguf +``` + +Original converter-compatible safetensors and shard indexes can be passed directly, +with the same component namespaces as `audiocpp_gguf`. For Qwen3-TTS: + +```bash +python tools/lora/merge_lora.py --family qwen3_tts \ + --input model_weights=models/Qwen3-TTS/model.safetensors \ + --input speech_tokenizer_weights=models/Qwen3-TTS/speech_tokenizer/model.safetensors \ + --root models/Qwen3-TTS --adapter adapters/my-voice \ + --type bf16 --output models/qwen3-tts-my-voice-bf16.gguf +``` + +Checkpoints requiring model-specific preparation still use their existing +conversion tool first. The merge tool does not guess or rewrite an unsupported +base tensor layout. + +| Model | Additional merge arguments | +|---|---| +| YuE2 | `--adapter ` and/or `--nar-adapter `; strengths are `--scale` and `--nar-scale`. | +| ACE-Step | Select `--component dit_turbo_weights` (or the matching DiT component) when the base bundles multiple variants. LoRA/LoKr parsing and merging use the shared framework. | +| Fish Audio S2 Pro | A prepared adapter, or a Lightning `.ckpt` plus its training `--alpha`. | +| Qwen3-TTS | Speaker embedding/config sidecars are applied and embedded automatically. | +| CosyVoice3 | Native speech-embedding and output-head sidecars are merged automatically. | +| VibeVoice | Diffusion-head and connector sidecars are merged when included. | +| Other listed models | Use `--adapter` and optional `--scale`; default strength is `1`. | + +`--scale` is an adapter-strength multiplier, including for VibeVoice; it is not +VibeVoice's runtime coefficient override. `--keep-type '='` can be +repeated. `--converter ` selects an existing `audiocpp_gguf` executable. +Use `--overwrite` to replace an existing output. The base files are never changed. + +The script stages only adapter factors and any changed sidecars. It does not +write a full intermediate merged checkpoint, and removes its staging directory +after conversion. Run the resulting GGUF without passing the adapter again. + +The GGUF converter can merge a normalized, fully bound adapter before applying +the requested weight type. Use a 16-bit or original-dtype base where available; +merging into quantized weights first dequantizes them, then requantizes the +result. Quantization can change output independently of LoRA. + +The converter's `--weight-adapter` input uses complete GGUF tensor names, including +component namespaces. `--adapter-merge-mode` selects `delta_f32`, `rankwise_f32` +(VibeVoice), or `rounded_bf16_delta` (YuE2). Model-specific adapter names must be +bound before using this low-level interface. Quantization and `--keep-type` +follow the normal [GGUF conversion rules](gguf.md). diff --git a/docs/models/cosyvoice3.md b/docs/models/cosyvoice3.md index c388c4e9a..201eafd95 100644 --- a/docs/models/cosyvoice3.md +++ b/docs/models/cosyvoice3.md @@ -88,6 +88,8 @@ audiocpp_cli \ | Option | Values | Default | Meaning | |---|---|---:|---| | `cosyvoice3.reference_cache_slots` | integer >= 0 | `4` | Prepared reference-audio cache slots. | +| `cosyvoice3.lora` | adapter directory or file | not set | PEFT LoRA, or a native checkpoint containing `lora_weights.pt`, `config.json`, `speech_embedding.pt`, and `llm_decoder.pt`. Only LLM weights are adapted. | +| `cosyvoice3.lora_scale` | finite float | `1` | Multiplier for alpha/r. Full embedding/head replacements are not scaled. | | `cosyvoice3.mem_saver` | bool | `false` | Release cached runtime graphs after request phases. | | `cosyvoice3.exec_plan` | JSON object | `{}` | Component execution plans. The only component is `flow`; see the RKNN section. | diff --git a/docs/models/f5_tts.md b/docs/models/f5_tts.md index ee454ed7f..18c655dbf 100644 --- a/docs/models/f5_tts.md +++ b/docs/models/f5_tts.md @@ -86,6 +86,14 @@ leaks into the output. ### Converting checkpoints to GGUF +For LoRA support, reconvert the matching original checkpoint with the current +converter; do not use the currently published F5-TTS GGUF packages as the base. +PEFT DiT adapters can then be loaded with `--session-option f5_tts.lora=`. +Keep `adapter_config.json` beside the adapter weights. The optional +`f5_tts.lora_scale` multiplies alpha/r and defaults to `1`. Adapted weights and +DiT graph caches are session-owned; the Vocos vocoder is not adapted. +See the [LoRA routing table](../lora.md). + `tools/community_models/convert_f5_tts.py` wraps `audiocpp_gguf` and produces one self-contained GGUF per checkpoint (`transformer` + `vocos` namespaces) plus the standalone vocoder package: diff --git a/docs/models/fish_audio.md b/docs/models/fish_audio.md index aae7de9c7..d136e3417 100644 --- a/docs/models/fish_audio.md +++ b/docs/models/fish_audio.md @@ -95,6 +95,8 @@ this is output streaming, not incremental text input. | Option | Values | Default | Meaning | |---|---|---:|---| | `fish_audio.mem_saver` | bool | `false` | Release cached AR runtime graphs after each request. | +| `fish_audio.lora` | checkpoint file | not set | Prepared Fish Speech adapter. Slow/fast decoder projections and embeddings are supported; codec weights are unchanged. | +| `fish_audio.lora_scale` | finite float | `1` | Multiplier for the adapter's alpha/r scaling. | | `fish_audio.reference_cache_slots` | integer | `1` | Prepared reference-audio cache slots. | | `fish_audio.weight_type` | `native`, `f32`, `f16`, `bf16`, `q8_0` | `native` | AR matmul weight storage type. | | `fish_audio.codec_weight_type` | `native`, `f32`, `f16`, `q8_0` | `native` | Codec conv/matmul weight storage type. | @@ -104,6 +106,15 @@ this is output streaming, not incremental text input. | `fish_audio.codec_graph_arena_mb` | integer MiB >= 1 | `512` | Codec graph arena size. | | `fish_audio.codec_weight_context_mb` | integer MiB >= 1 | `32` | Codec weight context size. | +To prepare a Fish Speech Lightning checkpoint, use its training alpha (not alpha/r): + +```bash +python tools/lora/prepare_fish_audio_adapter.py --input lora.ckpt --alpha 16 --output adapter.safetensors +``` + +The exporter transposes embedding factors to the runtime's matrix convention and +stores alpha with every target. Use an adapter trained for the exact base checkpoint. + ## Multi-Reference Conditioning `multi_reference_cond` is a Fish Audio request option. It takes an ordered JSON diff --git a/docs/models/outetts.md b/docs/models/outetts.md index ede268610..c2deff912 100644 --- a/docs/models/outetts.md +++ b/docs/models/outetts.md @@ -114,6 +114,8 @@ audiocpp_cli --task tts --family outetts \ The legacy session keys `outetts.aligner_model_path` and `outetts.dac_graph_context_mb` remain accepted for backward compatibility. | `--session-option outetts.reference_cache_slots=` | integer | `1` | LRU slots for prepared reference profiles (alignment, DAC codes, and word features). Set `0` to disable reuse. | +| `--session-option outetts.lora=` | directory or checkpoint file | not set | PEFT LoRA for the same Llama-OuteTTS checkpoint. Keep `adapter_config.json` beside its weights. DAC weights are unchanged. | +| `--session-option outetts.lora_scale=` | finite float | `1` | Multiplier for the adapter's alpha/r scaling. | | `--session-option outetts.mem_saver=true\|false` | bool | `false` | Release the reusable Llama cached-step graph after each generated chunk and release the aligner runtime after preparing a reference. Model and DAC weights stay resident; later requests rebuild released state. | With logging enabled, OuteTTS reports framework chunk count and token budget, per-chunk word/character counts, recommended and effective generation limits, the natural stop reason, reference-profile cache hits/evictions, Llama runtime and step-graph rebuild/reuse, released cache capacity, and timings for reference alignment, DAC encode/decode, prompt construction, generation, and the complete session request. See [OuteTTS validation](../reports/outetts_validation.md) for the reproducible long-lived session and memory test. diff --git a/docs/models/qwen3_asr.md b/docs/models/qwen3_asr.md index d262bc0c6..6bdaa0bb7 100644 --- a/docs/models/qwen3_asr.md +++ b/docs/models/qwen3_asr.md @@ -78,6 +78,8 @@ audiocpp_gguf --input models/QwenCleo-ASR/model.safetensors --root models/QwenCl | Option | Values | Default | Meaning | |---|---|---:|---| +| `qwen3_asr.lora` | adapter directory or checkpoint file | not set | Load a PEFT LoRA at session creation. Encoder and decoder targets are supported; use an adapter trained for the same checkpoint size. Keep `adapter_config.json` beside the weights. | +| `qwen3_asr.lora_scale` | finite float | `1` | Strength multiplier applied after the adapter's alpha/r scaling. | | `--audio` | WAV path | required | Speech audio; use 16 kHz WAV for the examples. | | `--text` | text | empty string | Context prompt. | | `--language` | language code | empty string | Recognition language hint. | diff --git a/docs/models/qwen3_tts.md b/docs/models/qwen3_tts.md index 32ff59a32..040d9e6ec 100644 --- a/docs/models/qwen3_tts.md +++ b/docs/models/qwen3_tts.md @@ -155,6 +155,8 @@ constant-context sizes are in MiB. | Option | Values | Default | Meaning | |---|---|---:|---| | `qwen3_tts.mem_saver` | bool | `false` | Release the talker cached-step graph after each request. Voice-prompt, prefill, code-predictor, and speech-decoder caches remain reusable. | +| `qwen3_tts.lora` | adapter directory or checkpoint file | not set | Load a PEFT Talker/Subtalker adapter at session creation. Keep `adapter_config.json` beside its weights. A trained custom voice may also provide `speaker_embedding` and `tts_config.json` sidecars; select its name with `voice_id`. | +| `qwen3_tts.lora_scale` | finite float | `1` | Strength multiplier applied after the adapter's alpha/r scaling. Speaker embedding replacements are not scaled. | | `qwen3_tts.perf_mode` | `off`, `flash_attention` | `off` | Opt-in attention optimization; `flash_attention` requires Q8_0 weights in both the talker and speech tokenizer. | | `qwen3_tts.stitch_decode` | bool | `false` | Voice the concatenated speech codes of all text chunks in one decoder pass instead of voicing each chunk separately. Per-chunk voicing voices the voice-clone reference once per chunk and zeroes the stateful vocoder cache at every chunk boundary (measured on RK3588, 300 characters, 22.5 s reference: 7 vocoder runs for 159.9 s against 1 run for 57.5 s); the stitched pass costs device memory proportional to the utterance, since one decode call builds a vocoder graph sized by its frame count. | | `qwen3_tts.voice_prompt_cache_slots` | nonnegative integer | `1` | Cached voice-clone prompts; `0` disables prompt caching. | diff --git a/docs/models/voxcpm1.md b/docs/models/voxcpm1.md index 81b8bd304..2217a42a5 100644 --- a/docs/models/voxcpm1.md +++ b/docs/models/voxcpm1.md @@ -107,6 +107,8 @@ Measured on Ubuntu 24.04 with OpenMP optimization (CPU backend): | `--session-option voxcpm1.mem_saver` | `true`, `false` | `false` | Use tighter graph workspaces and release request runtime graphs to reduce resident VRAM. | | `--session-option voxcpm1.prompt_cache_slots` | integer | `1` | Prompt and prompt-audio embedding cache slots. Set `0` to disable prompt caching. | | `--session-option voxcpm1.weight_type` | `native`, `f32`, `f16`, `bf16`, `q8_0` | `native` | Generator weight storage for the session. | +| `--session-option voxcpm1.lora` | directory or checkpoint file | not set | Native LoRA weights with adjacent `lora_config.json`; merged at session creation. | +| `--session-option voxcpm1.lora_scale` | finite float | `1` | Multiplier for the configured alpha/r scaling. | | `--session-option voxcpm1.audiovae_weight_type` | `native`, `f32`, `f16`, `bf16`, `q8_0` | `native` | AudioVAE weight storage for the session. | | `--max-tokens` | integer | `4096` | Maximum MiniCPM output tokens; default 4096. Zero selects model max_length; negative counts are rejected. | | `--num-inference-steps` | integer | `10` | CFM diffusion sampling steps; default 10. Must be positive. | diff --git a/docs/tts.md b/docs/tts.md index 0397eb664..74a7e0909 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -1,5 +1,7 @@ # TTS Models +For adapter loading and merged GGUF export, see the [LoRA guide](lora.md). + | Model | Family | Task(s) | Quick Start | |---|---|---|---| | AuK (experimental) | `auk` | `tts`, `clon`, `gen` editing | [AuK](models/auk.md) | @@ -1052,6 +1054,8 @@ audiocpp_cli --task tts --family voxcpm2 --model models/VoxCPM2 --backend cuda - | `voxcpm2.audiovae_weight_type` | `native`, `f32`, `f16`, `bf16`, `q8_0` | not set | AudioVAE weight storage; defaults independently to native, not the shared weight_type. | | `voxcpm2.prompt_cache_slots` | integer >= 0 | `1` | Cached encoded prompt slots; zero disables reuse. | | `voxcpm2.mem_saver` | bool | `true` | Release staged runtime graphs. | +| `voxcpm2.lora` | directory or checkpoint file | not set | Native LoRA weights with adjacent `lora_config.json`; merged at session creation without changing the AudioVAE. | +| `voxcpm2.lora_scale` | finite float | `1` | Multiplier for the adapter's configured alpha/r scaling. | | `voxcpm2.weight_context_mb` | integer >= 1 | `3072` | Context or graph arena size in MiB. | | `voxcpm2.text_embedding_graph_arena_mb` | integer >= 1 | `64` | Host graph-metadata arena size in MiB; legacy `text_embedding_graph_context_mb` accepted with the same family prefix. | | `voxcpm2.lm_step_graph_arena_mb` | integer >= 1 | `1024` | Host graph-metadata arena size in MiB; legacy `lm_step_graph_context_mb` accepted with the same family prefix. | diff --git a/include/engine/community_models/f5_tts/runtime.h b/include/engine/community_models/f5_tts/runtime.h index 90a43a327..37498a3ad 100644 --- a/include/engine/community_models/f5_tts/runtime.h +++ b/include/engine/community_models/f5_tts/runtime.h @@ -11,6 +11,8 @@ namespace engine::models::f5_tts { +struct F5RuntimeState; + // 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). @@ -42,8 +44,13 @@ struct F5ComputeDevice { bool use_cuda = false; int device = 0; // CUDA device index int threads = 0; // CPU threads; 0 = hardware concurrency + std::shared_ptr state = nullptr; }; +std::shared_ptr create_f5_lora_state( + const std::string & weights_path, const std::filesystem::path & adapter_path, + float strength, const F5ComputeDevice & device); + // Debug taps for parity testing: when non-null, intermediate stage outputs are // appended (column layout, [features, T] flattened feature-major). struct F5DebugTaps { diff --git a/include/engine/community_models/f5_tts/session.h b/include/engine/community_models/f5_tts/session.h index 6cbfe642c..6b26230fc 100644 --- a/include/engine/community_models/f5_tts/session.h +++ b/include/engine/community_models/f5_tts/session.h @@ -11,6 +11,8 @@ namespace engine::models::f5_tts { +struct F5RuntimeState; + // F5-TTS community model assets: resource bundle + resolved checkpoint path. // Weights are loaded lazily by the runtime on first synthesis (graph cache). struct F5TTSAssets { @@ -44,6 +46,7 @@ class F5TTSSession final : public runtime::IOfflineVoiceTaskSession { bool use_cuda_ = false; int cuda_device_ = 0; int threads_ = 0; + std::shared_ptr runtime_state_; }; std::shared_ptr load_f5_tts_assets( diff --git a/include/engine/community_models/f5_tts/synthesize.h b/include/engine/community_models/f5_tts/synthesize.h index bfcf3f416..e77b9bfed 100644 --- a/include/engine/community_models/f5_tts/synthesize.h +++ b/include/engine/community_models/f5_tts/synthesize.h @@ -27,6 +27,7 @@ struct F5SynthesisRequest { int threads = 0; // 0 = hardware concurrency bool use_cuda = false; int cuda_device = 0; + std::shared_ptr runtime_state; }; struct F5SynthesisResult { diff --git a/include/engine/framework/assets/tensor_source.h b/include/engine/framework/assets/tensor_source.h index 764e77996..f6a5516b4 100644 --- a/include/engine/framework/assets/tensor_source.h +++ b/include/engine/framework/assets/tensor_source.h @@ -170,6 +170,9 @@ void set_backend_tensor_from_f32_parallel( ggml_type type, std::vector * retained_bytes = nullptr); std::shared_ptr open_tensor_source(const std::filesystem::path & path); +// Finds a logical checkpoint by stem, preferring safetensors over PyTorch and GGUF. +std::optional find_tensor_file( + const std::filesystem::path & directory, std::string_view stem); std::shared_ptr open_tensor_source( const std::filesystem::path & path, TensorSourceOptions options); std::shared_ptr open_tensor_source( @@ -199,11 +202,16 @@ struct GgufTensorTypeOverride { std::string pattern; TensorStorageType storage_type = TensorStorageType::Native; }; +struct WeightAdapter; +enum class WeightMergeMode; struct GgufConversionOptions { std::optional bnb_nf4_type; std::vector excluded_tensor_prefixes; std::vector type_overrides; std::vector folded_weight_norm_patterns; + // Bound tensor names; applied before conversion and quantization. + std::shared_ptr weight_adapter; + std::optional weight_adapter_merge_mode; }; void convert_tensor_sources_to_gguf(const std::vector & inputs, const std::filesystem::path & output_path, TensorStorageType weight_type, diff --git a/include/engine/models/qwen3_tts/assets.h b/include/engine/models/qwen3_tts/assets.h index bd75e86a1..3bd7a0a73 100644 --- a/include/engine/models/qwen3_tts/assets.h +++ b/include/engine/models/qwen3_tts/assets.h @@ -90,4 +90,8 @@ struct Qwen3TTSAssets { std::shared_ptr load_qwen3_tts_assets(const std::filesystem::path & model_path); +std::shared_ptr apply_qwen3_tts_lora( + std::shared_ptr assets, + const std::unordered_map & options); + } // namespace engine::models::qwen3_tts diff --git a/model_specs/cosyvoice3.json b/model_specs/cosyvoice3.json index 6908f0947..bd6841f42 100644 --- a/model_specs/cosyvoice3.json +++ b/model_specs/cosyvoice3.json @@ -128,6 +128,19 @@ } ], "session": [ + { + "name": "lora", + "type": "path", + "description": "PEFT adapter with adapter_config.json, or native lora_weights with config.json and speech_embedding/llm_decoder sidecars. Checkpoints may use safetensors or PyTorch containers.", + "required": false + }, + { + "name": "lora_scale", + "type": "float", + "description": "Multiplier for the adapter's alpha/r scaling; default 1. Full-weight replacements are not scaled.", + "required": false, + "default": 1.0 + }, { "name": "weight_type", "type": "enum", diff --git a/model_specs/f5_tts.json b/model_specs/f5_tts.json index 0fe6f71e1..df05caa8f 100644 --- a/model_specs/f5_tts.json +++ b/model_specs/f5_tts.json @@ -92,6 +92,19 @@ } ], "session": [ + { + "name": "lora", + "type": "path", + "description": "PEFT DiT LoRA adapter directory or checkpoint file; merged into session-owned weights and graphs. Vocos is unchanged.", + "required": false + }, + { + "name": "lora_scale", + "type": "float", + "description": "Multiplier for the adapter's alpha/r scaling; default 1.", + "required": false, + "default": 1.0 + }, { "name": "vocos_path", "type": "string", diff --git a/model_specs/fish_audio.json b/model_specs/fish_audio.json index 7694e8f53..d48ebff20 100644 --- a/model_specs/fish_audio.json +++ b/model_specs/fish_audio.json @@ -98,6 +98,19 @@ } ], "session": [ + { + "name": "lora", + "type": "path", + "description": "Prepared Fish Speech adapter checkpoint; merged into slow/fast decoder and embedding weights at session creation. Codec weights are unchanged.", + "required": false + }, + { + "name": "lora_scale", + "type": "float", + "description": "Multiplier for the adapter's alpha/r scaling; default 1.", + "required": false, + "default": 1.0 + }, { "name": "mem_saver", "type": "bool", diff --git a/model_specs/outetts.json b/model_specs/outetts.json index be77cb483..66c35c4ab 100644 --- a/model_specs/outetts.json +++ b/model_specs/outetts.json @@ -135,6 +135,19 @@ } ], "session": [ + { + "name": "lora", + "type": "path", + "description": "PEFT LoRA adapter directory or checkpoint file for the same OuteTTS checkpoint; merged into the language model at session creation.", + "required": false + }, + { + "name": "lora_scale", + "type": "float", + "description": "Multiplier for the adapter's alpha/r scaling; default 1.", + "required": false, + "default": 1.0 + }, { "name": "weight_type", "type": "enum", diff --git a/model_specs/qwen3_asr.json b/model_specs/qwen3_asr.json index 8b977bf16..d2b32ae7b 100644 --- a/model_specs/qwen3_asr.json +++ b/model_specs/qwen3_asr.json @@ -109,6 +109,19 @@ } ], "session": [ + { + "name": "lora", + "type": "path", + "description": "PEFT LoRA adapter directory or checkpoint file; merged into matching encoder/decoder weights at session creation.", + "required": false + }, + { + "name": "lora_scale", + "type": "float", + "description": "Multiplier for the adapter's alpha/r scaling; default 1.", + "required": false, + "default": 1.0 + }, { "name": "forced_aligner_path", "type": "path", diff --git a/model_specs/qwen3_tts.json b/model_specs/qwen3_tts.json index 1ac511e0a..c6b715cec 100644 --- a/model_specs/qwen3_tts.json +++ b/model_specs/qwen3_tts.json @@ -194,6 +194,19 @@ } ], "session": [ + { + "name": "lora", + "type": "path", + "description": "PEFT LoRA adapter directory or checkpoint file; optional speaker_embedding and tts_config.json sidecars provide a trained custom voice.", + "required": false + }, + { + "name": "lora_scale", + "type": "float", + "description": "Multiplier for the adapter's alpha/r scaling; default 1.", + "required": false, + "default": 1.0 + }, { "name": "mem_saver", "type": "bool", diff --git a/model_specs/voxcpm1.json b/model_specs/voxcpm1.json index e67c9a941..cf83beec1 100644 --- a/model_specs/voxcpm1.json +++ b/model_specs/voxcpm1.json @@ -116,6 +116,19 @@ } ], "session": [ + { + "name": "lora", + "type": "path", + "description": "Native LoRA directory or checkpoint file with adjacent lora_config.json; merged into generator weights at session creation.", + "required": false + }, + { + "name": "lora_scale", + "type": "float", + "description": "Multiplier for the adapter's configured alpha/r scaling; default 1.", + "required": false, + "default": 1.0 + }, { "name": "mem_saver", "type": "bool", diff --git a/model_specs/voxcpm2.json b/model_specs/voxcpm2.json index 511dd0eca..3421bd9f0 100644 --- a/model_specs/voxcpm2.json +++ b/model_specs/voxcpm2.json @@ -181,6 +181,19 @@ } ], "session": [ + { + "name": "lora", + "type": "path", + "description": "Native LoRA directory or checkpoint file with adjacent lora_config.json; merged into generator weights at session creation.", + "required": false + }, + { + "name": "lora_scale", + "type": "float", + "description": "Multiplier for the adapter's configured alpha/r scaling; default 1.", + "required": false, + "default": 1.0 + }, { "name": "weight_type", "type": "enum", diff --git a/model_specs/yue2.json b/model_specs/yue2.json index ffe8b9090..9cbd7acd2 100644 --- a/model_specs/yue2.json +++ b/model_specs/yue2.json @@ -478,13 +478,20 @@ { "format": "gguf", "roots": { - "model": "." + "model": ".", + "weights": "$gguf" }, "files": { "model_config": "model:sidecars/yue2-model-config.json", "generation_config": "model:sidecars/yue2-generation-config.json", "tiktoken": "model:sidecars/yue2-qwen.tiktoken", "vae_config": "model:sidecars/yue2-vae-config.json" + }, + "tensors": { + "model_weights": { + "source": "weights:", + "prefix": "model_weights" + } } }, { diff --git a/src/community_models/f5_tts/runtime.cpp b/src/community_models/f5_tts/runtime.cpp index 74ac870e4..9557a70dc 100644 --- a/src/community_models/f5_tts/runtime.cpp +++ b/src/community_models/f5_tts/runtime.cpp @@ -6,7 +6,9 @@ #include "engine/community_models/f5_tts/weights.h" #include "engine/framework/assets/tensor_source.h" +#include "engine/framework/assets/weight_adapter.h" #include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" #include "engine/framework/core/backend_weight_store.h" #include "engine/framework/core/module.h" #include "engine/framework/modules/conv_modules.h" @@ -120,6 +122,82 @@ class StrippedView final : public engine::assets::TensorSource { std::unordered_map routes_; }; +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) + // The staged graph constants this build created (pe table, masks, + // per-block ones rows). On CUDA each carries a private backend buffer + // allocated in const_stage_bind: the eviction path destroys the + // unique_ptr, and nothing else frees them. + std::vector * stage = nullptr; + ~CfgGraph() { + // Cache eviction (more than two distinct (N, NT) keys in one process, + // i.e. every chunk past the second of a chunked request) destroys the + // unique_ptr, and without this the context pool, the gallocr arena and + // the io buffer of every evicted graph were leaked: measured 24.5 GiB + // peak RSS for a four-chunk CPU request, unchanged when the cache bound + // was lowered from 2 to 1 (more evictions, same leaks). + if (gallocr != nullptr) ggml_gallocr_free(gallocr); + if (io_buffer != nullptr) ggml_backend_buffer_free(io_buffer); + if (ctx != nullptr) ggml_free(ctx); + const_stage_free(stage); + } +}; + +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; + // The staged graph constants this build created (pe table, per-block ones + // rows). On CUDA each carries a private backend buffer from + // const_stage_bind, so the vector has to die with the graph. + std::vector * stage = nullptr; + ~DiTGraph() { + // Session-owned graphs are destroyed before their backend. + if (gallocr != nullptr) ggml_gallocr_free(gallocr); + if (io_buffer != nullptr) ggml_backend_buffer_free(io_buffer); + if (ctx != nullptr) ggml_free(ctx); + const_stage_free(stage); + } +}; + +} // namespace + +struct F5RuntimeState { + core::ExecutionContext execution; + LoadedModel model; + F5DiTWeights weights; + std::map, std::unique_ptr> cfg_graphs; + std::map, std::unique_ptr> dit_graphs; + + explicit F5RuntimeState(const core::BackendConfig & config) + : execution(config), model{execution.backend(), config.type} {} +}; + +namespace { + const LoadedModel & load_model_once(const std::string & path, const F5ComputeDevice & dev); // Open the DiT checkpoint as a stripped-name tensor source. Safetensors @@ -140,6 +218,9 @@ std::shared_ptr open_dit_source(const std::s // 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) { + if (dev.state) { + return dev.state->weights; + } struct Entry { F5DiTWeights w; }; static auto * cache = new std::unordered_map(); const std::string key = @@ -156,6 +237,9 @@ const F5DiTWeights & load_dit_weights_once( } const LoadedModel & load_model_once(const std::string & path, const F5ComputeDevice & dev) { + if (dev.state) { + return dev.state->model; + } // 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. @@ -183,6 +267,31 @@ const LoadedModel & load_model_once(const std::string & path, const F5ComputeDev } // namespace +std::shared_ptr create_f5_lora_state( + const std::string & weights_path, const std::filesystem::path & adapter_path, + float strength, const F5ComputeDevice & device) { + auto source = open_dit_source(weights_path); + engine::assets::WeightAdapterReadOptions options; + options.strength = strength; + options.require_scaling_metadata = true; + auto adapter = engine::assets::read_weight_adapter(adapter_path, options); + engine::assets::WeightAdapterBindings bindings; + for (const auto & tensor : source->tensors()) { + bindings.add(tensor.name, tensor.name); + bindings.add("transformer." + tensor.name, tensor.name); + } + auto plan = engine::assets::bind_weight_adapter( + *source, std::move(adapter), bindings, engine::assets::WeightMergeMode::DeltaF32); + source = engine::assets::make_weight_adapter_source(std::move(source), std::move(plan), "f5_tts.lora"); + const core::BackendConfig config{ + device.use_cuda ? core::BackendType::Cuda : core::BackendType::Cpu, + device.use_cuda ? device.device : 0, + device.use_cuda ? 1 : std::max(1, device.threads)}; + auto state = std::make_shared(config); + state->weights = load_dit_weights(*source, state->model.backend, state->model.backend_type); + return state; +} + // 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). @@ -206,36 +315,8 @@ std::pair, std::vector> f5_dit_forward_cfg( 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) - // The staged graph constants this build created (pe table, masks, - // per-block ones rows). On CUDA each carries a private backend buffer - // allocated in const_stage_bind: the eviction path destroys the - // unique_ptr, and nothing else frees them. - std::vector * stage = nullptr; - ~CfgGraph() { - // Cache eviction (more than two distinct (N, NT) keys in one process, - // i.e. every chunk past the second of a chunked request) destroys the - // unique_ptr, and without this the context pool, the gallocr arena and - // the io buffer of every evicted graph were leaked: measured 24.5 GiB - // peak RSS for a four-chunk CPU request, unchanged when the cache bound - // was lowered from 2 to 1 (more evictions, same leaks). - if (gallocr != nullptr) ggml_gallocr_free(gallocr); - if (io_buffer != nullptr) ggml_backend_buffer_free(io_buffer); - if (ctx != nullptr) ggml_free(ctx); - const_stage_free(stage); - } - }; - static auto * cache = new std::map, std::unique_ptr>(); + static auto * default_cache = new std::map, std::unique_ptr>(); + auto * cache = dev.state ? &dev.state->cfg_graphs : default_cache; const auto ckey = std::make_tuple(&model, N, NT); auto it = cache->find(ckey); if (it == cache->end()) { @@ -417,42 +498,11 @@ std::vector f5_dit_forward( // 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; - // The staged graph constants this build created (pe table, per-block ones - // rows). On CUDA each carries a private backend buffer from - // const_stage_bind, so the vector has to die with the graph. - std::vector * stage = 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); - const_stage_free(stage); - } - }; - // Graph cache is process-lifetime (CUDA buffers cannot be safely freed // after driver shutdown in static destruction). - static auto * graph_cache = + static auto * default_graph_cache = new std::map, std::unique_ptr>(); + auto * graph_cache = dev.state ? &dev.state->dit_graphs : default_graph_cache; 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); diff --git a/src/community_models/f5_tts/session.cpp b/src/community_models/f5_tts/session.cpp index 2ba32d7a1..f3722394c 100644 --- a/src/community_models/f5_tts/session.cpp +++ b/src/community_models/f5_tts/session.cpp @@ -3,6 +3,7 @@ #include "engine/community_models/f5_tts/synthesize.h" #include "engine/framework/assets/tensor_source.h" +#include "engine/framework/debug/trace.h" #include "engine/framework/runtime/options.h" #include "engine/framework/runtime/spec_backed_model.h" @@ -139,6 +140,15 @@ F5TTSSession::F5TTSSession( use_cuda_ = options.backend.type == core::BackendType::Cuda; cuda_device_ = options.backend.device; threads_ = options.backend.threads; + if (const auto path = runtime::find_option(options.options, {"f5_tts.lora"}); path && !path->empty()) { + F5ComputeDevice device; + device.use_cuda = use_cuda_; + device.device = cuda_device_; + device.threads = threads_; + const float strength = runtime::parse_finite_float_option( + options.options, {"f5_tts.lora_scale"}).value_or(1.0F); + runtime_state_ = create_f5_lora_state(assets_->checkpoint.string(), std::filesystem::u8path(*path), strength, device); + } } std::string F5TTSSession::family() const noexcept { @@ -209,9 +219,11 @@ runtime::TaskResult F5TTSSession::run(const runtime::TaskRequest & request) { req.frame_budget = frame_budget_; req.cuda_device = cuda_device_; req.threads = threads_; + req.runtime_state = runtime_state_; const auto out = f5_synthesize( assets_->checkpoint.string(), vocos_path_, req); + engine::debug::timing_log_scalar("session.wall_ms", out.generation_seconds * 1000.0); runtime::TaskResult result; runtime::AudioBuffer audio; diff --git a/src/community_models/f5_tts/synthesize.cpp b/src/community_models/f5_tts/synthesize.cpp index 1b70e1578..301a4024a 100644 --- a/src/community_models/f5_tts/synthesize.cpp +++ b/src/community_models/f5_tts/synthesize.cpp @@ -1232,6 +1232,7 @@ F5SynthesisResult f5_synthesize( dev.use_cuda = request.use_cuda; dev.device = request.cuda_device; dev.threads = request.threads; + dev.state = request.runtime_state; // ---- chunk long texts instead of truncating: each chunk is sized by // the DURATION budget (frames/char from the reference, capped) so no diff --git a/src/community_models/outetts/session.cpp b/src/community_models/outetts/session.cpp index 3fdfdf4d7..87feefba5 100644 --- a/src/community_models/outetts/session.cpp +++ b/src/community_models/outetts/session.cpp @@ -1,5 +1,7 @@ #include "engine/community_models/outetts/session.h" +#include "engine/framework/assets/weight_adapter.h" + #include "engine/framework/audio/fft.h" #include "engine/framework/runtime/options.h" #include "engine/framework/runtime/spec_backed_model.h" @@ -39,10 +41,28 @@ constexpr size_t kDefaultLlamaWeightContextBytes = 4096ull * 1024ull * 1024ull; #endif std::shared_ptr require_assets( - std::shared_ptr assets) { + std::shared_ptr assets, const runtime::SessionOptions & options) { if (assets == nullptr) throw std::runtime_error("OuteTTS session requires assets"); - return assets; + const auto path = runtime::find_option(options.options, {"outetts.lora"}); + if (!path || path->empty()) { + return assets; + } + engine::assets::WeightAdapterReadOptions read_options; + read_options.require_scaling_metadata = true; + read_options.strength = runtime::parse_finite_float_option( + options.options, {"outetts.lora_scale"}).value_or(1.0F); + auto adapter = engine::assets::read_weight_adapter(std::filesystem::u8path(*path), read_options); + engine::assets::WeightAdapterBindings bindings; + for (const auto & tensor : assets->model_weights->tensors()) { + bindings.add(tensor.name, tensor.name); + } + auto plan = engine::assets::bind_weight_adapter( + *assets->model_weights, std::move(adapter), bindings, engine::assets::WeightMergeMode::DeltaF32); + auto adapted = std::make_shared(*assets); + adapted->model_weights = engine::assets::make_weight_adapter_source( + assets->model_weights, std::move(plan), "outetts.lora"); + return adapted; } std::shared_ptr require_contract( @@ -69,6 +89,7 @@ void validate_session_option_keys( if (key.rfind(family_prefix, 0) == 0 && contract.session_option_keys.find(key) == contract.session_option_keys.end() && + key != "outetts.lora" && key != "outetts.lora_scale" && !is_legacy_session_option(key)) { throw std::runtime_error( "unknown OuteTTS session option: " + key); @@ -528,7 +549,7 @@ OuteTTSSession::OuteTTSSession(runtime::TaskSpec task, std::shared_ptr contract) : RuntimeSessionBase(options), task_(task), - assets_(require_assets(std::move(assets))), + assets_(require_assets(std::move(assets), options)), contract_(require_contract(std::move(contract))), tokenizer_(assets_), dac_(assets_, execution_context(), diff --git a/src/community_models/voxcpm1/session.cpp b/src/community_models/voxcpm1/session.cpp index dd77903d1..ac3e3c062 100644 --- a/src/community_models/voxcpm1/session.cpp +++ b/src/community_models/voxcpm1/session.cpp @@ -1,5 +1,8 @@ #include "engine/community_models/voxcpm1/session.h" +#include "engine/framework/assets/weight_adapter.h" +#include "engine/framework/io/filesystem.h" +#include "engine/framework/io/json.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/model_spec/metadata.h" #include "engine/framework/model_spec/package.h" @@ -13,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -27,11 +31,72 @@ using Clock = std::chrono::steady_clock; constexpr int64_t kDefaultTextChunkSize = 2048; std::shared_ptr -require_assets(std::shared_ptr assets) { +require_assets(std::shared_ptr assets, + const runtime::SessionOptions & options) { if (assets == nullptr) { throw std::runtime_error("VoxCPM1 session requires assets"); } - return assets; + const auto path = runtime::find_option(options.options, {"voxcpm1.lora"}); + if (!path || path->empty()) { + return assets; + } + const auto adapter_path = std::filesystem::u8path(*path); + auto file = adapter_path; + if (io::is_existing_directory(adapter_path)) { + const auto resolved = engine::assets::find_tensor_file(adapter_path, "lora_weights"); + if (!resolved) throw std::runtime_error("VoxCPM1 lora_weights checkpoint is missing in " + adapter_path.string()); + file = *resolved; + } + const auto config = io::json::parse_file(file.parent_path() / "lora_config.json"); + const auto & lora = config.require("lora_config"); + const auto rank = io::json::require_i64(lora, "r"); + if (rank <= 0) { + throw std::runtime_error("VoxCPM1 LoRA rank must be positive"); + } + engine::assets::WeightAdapterReadOptions read_options; + read_options.scale_override = io::json::require_f32(lora, "alpha") / static_cast(rank); + read_options.strength = runtime::parse_finite_float_option( + options.options, {"voxcpm1.lora_scale"}).value_or(1.0F); + auto adapter = engine::assets::read_weight_adapter(file, read_options); + for (const auto & [name, update] : adapter.updates) { + const auto * factors = std::get_if(&update.factors); + if (factors == nullptr || factors->r != rank) { + throw std::runtime_error("VoxCPM1 LoRA factors do not match lora_config.json: " + name); + } + } + engine::assets::WeightAdapterBindings bindings; + const std::regex layer_pattern(R"(((residual_lm|locenc|locdit)\.)?blk\.([0-9]+)\.(attn_q|attn_k|attn_v|attn_output|ffn_gate|ffn_up|ffn_down)\.weight)"); + const std::unordered_map prefixes{ + {"", "base_lm"}, {"residual_lm", "residual_lm"}, + {"locenc", "feat_encoder.encoder"}, {"locdit", "feat_decoder.estimator.decoder"}}; + const std::unordered_map projections{ + {"attn_q", "self_attn.q_proj"}, {"attn_k", "self_attn.k_proj"}, + {"attn_v", "self_attn.v_proj"}, {"attn_output", "self_attn.o_proj"}, + {"ffn_gate", "mlp.gate_proj"}, {"ffn_up", "mlp.up_proj"}, {"ffn_down", "mlp.down_proj"}}; + for (const auto & tensor : assets->model_weights->tensors()) { + bindings.add(tensor.name, tensor.name); + std::smatch match; + if (std::regex_match(tensor.name, match, layer_pattern)) { + bindings.add(prefixes.at(match[2].str()) + ".layers." + match[3].str() + "." + + projections.at(match[4].str()), tensor.name); + } + } + for (const auto & entry : std::initializer_list>{ + {"enc_to_lm_proj", "proj.enc_to_lm.weight"}, + {"lm_to_dit_proj", "proj.lm_to_dit.weight"}, + {"res_to_dit_proj", "proj.res_to_dit.weight"}, + {"feat_encoder.in_proj", "locenc.in_proj.weight"}, + {"feat_decoder.estimator.in_proj", "locdit.in_proj.weight"}, + {"feat_decoder.estimator.cond_proj", "locdit.cond_proj.weight"}, + {"feat_decoder.estimator.out_proj", "locdit.out_proj.weight"}}) { + bindings.add(entry.first, entry.second); + } + auto plan = engine::assets::bind_weight_adapter( + *assets->model_weights, std::move(adapter), bindings, engine::assets::WeightMergeMode::DeltaF32); + auto adapted = std::make_shared(*assets); + adapted->model_weights = engine::assets::make_weight_adapter_source( + assets->model_weights, std::move(plan), "voxcpm1.lora"); + return adapted; } void reject_enabled_denoise( @@ -116,7 +181,8 @@ void validate_session_options( if (key.rfind("voxcpm1.", 0) != 0) { continue; } - if (key == "voxcpm1.weight_context_mb" || + if (key == "voxcpm1.lora" || key == "voxcpm1.lora_scale" || + key == "voxcpm1.weight_context_mb" || key == "voxcpm1.text_embedding_graph_context_mb" || key == "voxcpm1.text_embedding_graph_arena_mb" || key == "voxcpm1.lm_step_graph_context_mb" || @@ -184,7 +250,7 @@ VoxCPM1SessionBase::VoxCPM1SessionBase(runtime::TaskSpec task, runtime::SessionOptions options, std::shared_ptr assets) : RuntimeSessionBase(options), task_(task), - assets_(require_assets(std::move(assets))), + assets_(require_assets(std::move(assets), options)), encoded_prompt_cache_(prompt_cache_slots_from_options(options.options)) { if (task_.mode != runtime::RunMode::Offline && task_.mode != runtime::RunMode::Streaming) { @@ -878,6 +944,8 @@ runtime::ModelCliInterface cli_v1(const VoxCPM1Assets &) { {"text_chunk_size", "int", "Maximum characters per offline text chunk; must be positive. Does not control streaming chunk cadence. Default: 2048.", false, "2048", "1"}, }; out.session_options = { + {"voxcpm1.lora", "path", "Native LoRA directory or checkpoint file with adjacent lora_config.json; merged into generator weights at session creation.", false}, + {"voxcpm1.lora_scale", "float", "Multiplier for the adapter's configured alpha/r scaling; default 1.", false, "1"}, {"voxcpm1.mem_saver", "bool", "Use tighter graph workspaces and release request runtime graphs; default false.", false, "false"}, {"voxcpm1.prompt_cache_slots", "int", "Prompt and prompt-audio embedding cache slots; default 1.", false, "1"}, {"voxcpm1.weight_type", "native|f32|f16|bf16|q8_0", "Model weight storage type.", false, "native"}, diff --git a/src/framework/assets/tensor_source.cpp b/src/framework/assets/tensor_source.cpp index c7175c7d4..c8a2d6713 100644 --- a/src/framework/assets/tensor_source.cpp +++ b/src/framework/assets/tensor_source.cpp @@ -1,4 +1,6 @@ #include "engine/framework/assets/tensor_source.h" +#include "engine/framework/assets/weight_adapter.h" +#include "engine/framework/assets/torch_bin.h" #include "engine/framework/io/binary.h" #include "engine/framework/io/filesystem.h" @@ -22,10 +24,21 @@ #include #include #include +#include namespace engine::assets { namespace { +enum class TensorContainer { Safetensors, Torch, Gguf }; +constexpr std::pair kTensorContainers[] = { + {".safetensors", TensorContainer::Safetensors}, + {".pt", TensorContainer::Torch}, + {".pth", TensorContainer::Torch}, + {".bin", TensorContainer::Torch}, + {".ckpt", TensorContainer::Torch}, + {".gguf", TensorContainer::Gguf}, +}; + constexpr int64_t kParallelF32ConvertElements = 1ll << 20; constexpr int64_t kF32ConvertChunkElements = 1ll << 16; // Quantization chunks are whole ROWS, so this is a budget rounded down to a row @@ -1650,6 +1663,15 @@ std::shared_ptr open_tensor_source(const std::filesystem::pa return open_tensor_source(path, TensorSourceOptions{}); } +std::optional find_tensor_file( + const std::filesystem::path & directory, std::string_view stem) { + for (const auto & entry : kTensorContainers) { + const auto candidate = directory / (std::string(stem) + std::string(entry.first)); + if (io::is_existing_file(candidate)) return candidate; + } + return std::nullopt; +} + std::shared_ptr open_tensor_source( const std::filesystem::path & path, TensorSourceOptions options) { const std::string file_name = lower_ascii(path.filename().string()); @@ -1659,11 +1681,13 @@ std::shared_ptr open_tensor_source( return open_indexed_tensor_source(path, path.parent_path(), options); } const std::string extension = lower_ascii(path.extension().string()); - if (extension == ".safetensors") { - return std::make_shared(path, options); - } - if (extension == ".gguf") { - return std::make_shared(path, options); + for (const auto & entry : kTensorContainers) { + if (extension != entry.first) continue; + switch (entry.second) { + case TensorContainer::Safetensors: return std::make_shared(path, options); + case TensorContainer::Torch: return open_torch_bin_tensor_source(path); + case TensorContainer::Gguf: return std::make_shared(path, options); + } } throw std::runtime_error("unsupported tensor source format: " + path.string()); } @@ -1958,6 +1982,25 @@ void convert_tensor_sources_to_gguf(const std::vector & input std::move(source), options.folded_weight_norm_patterns); } + std::unordered_set adapted_tensors; + if (options.weight_adapter) { + const auto & adapter = *options.weight_adapter; + const auto mode = options.weight_adapter_merge_mode.value_or(adapter.updates.empty() + ? WeightMergeMode::DeltaF32 : adapter.updates.begin()->second.merge_mode); + WeightAdapterBindings bindings; + for (const auto & tensor : source->tensors()) { + bindings.add(tensor.name, tensor.name); + } + for (const auto & entry : adapter.updates) { + if (entry.second.merge_mode != mode) { + throw std::runtime_error("GGUF conversion requires one adapter merge policy"); + } + } + auto plan = bind_weight_adapter(*source, adapter, bindings, mode); + for (const auto & entry : plan.updates) adapted_tensors.insert(entry.first); + for (const auto & entry : plan.replacements) adapted_tensors.insert(entry.first); + source = make_weight_adapter_source(std::move(source), std::move(plan), "gguf.lora"); + } const auto metadata = source->tensors(); if (metadata.empty()) { throw std::runtime_error("cannot convert an empty tensor source to GGUF"); @@ -2281,7 +2324,8 @@ void convert_tensor_sources_to_gguf(const std::vector & input RawTensorData raw; if (convert_bnb_nf4 && source->has_tensor(bnb_nf4_quant_state_name(item.metadata.name))) { raw = convert_bnb_nf4_weight(*source, item.metadata, item.storage_type); - } else if (item.storage_type == TensorStorageType::Native) { + } else if (item.storage_type == TensorStorageType::Native && + adapted_tensors.count(item.metadata.name) == 0) { raw = source->require_tensor_data(item.metadata.name); } else { const auto tensor = source->require_tensor( diff --git a/src/framework/assets/weight_adapter.cpp b/src/framework/assets/weight_adapter.cpp index 4a06f00d8..0cabbef30 100644 --- a/src/framework/assets/weight_adapter.cpp +++ b/src/framework/assets/weight_adapter.cpp @@ -128,7 +128,12 @@ WeightAdapter read_weight_adapter(const std::filesystem::path & path, const Weig if (!std::isfinite(options.strength) || (options.scale_override && !std::isfinite(*options.scale_override))) { throw std::runtime_error("adapter strength and scale must be finite"); } - const auto file = io::is_existing_directory(path) ? path / "adapter_model.safetensors" : path; + auto file = path; + if (io::is_existing_directory(path)) { + const auto resolved = find_tensor_file(path, "adapter_model"); + if (!resolved) throw std::runtime_error("missing adapter_model checkpoint in " + path.string()); + file = *resolved; + } const auto config_path = file.parent_path() / "adapter_config.json"; const bool has_config = io::is_existing_file(config_path); const auto source = open_tensor_source(file); diff --git a/src/models/cosyvoice3/session.cpp b/src/models/cosyvoice3/session.cpp index 03d3221d6..7e5c6bc48 100644 --- a/src/models/cosyvoice3/session.cpp +++ b/src/models/cosyvoice3/session.cpp @@ -1,5 +1,7 @@ #include "engine/models/cosyvoice3/session.h" +#include "engine/framework/assets/weight_adapter.h" +#include "engine/framework/io/filesystem.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/io/json.h" #include "engine/framework/rknn/engine.h" @@ -28,11 +30,80 @@ constexpr const char * kModelName = "CosyVoice3"; constexpr int64_t kDefaultTextChunkSize = 600; constexpr size_t kDefaultReferenceCacheSlots = 4; -std::shared_ptr require_assets(std::shared_ptr assets) { +std::shared_ptr require_assets( + std::shared_ptr assets, const runtime::SessionOptions & options) { if (assets == nullptr) { throw std::runtime_error("CosyVoice3 session requires assets"); } - return assets; + const auto path = runtime::find_option(options.options, {"cosyvoice3.lora"}); + if (!path || path->empty()) { + return assets; + } + const auto adapter_path = std::filesystem::u8path(*path); + const auto directory = io::is_existing_directory(adapter_path) ? adapter_path : adapter_path.parent_path(); + const bool native_package = io::is_existing_directory(adapter_path) + ? engine::assets::find_tensor_file(directory, "lora_weights").has_value() + : adapter_path.stem() == "lora_weights"; + const bool native = io::is_existing_file(directory / "config.json") && + (native_package || !io::is_existing_file(directory / "adapter_config.json")); + auto file = adapter_path; + if (io::is_existing_directory(adapter_path)) { + const auto resolved = engine::assets::find_tensor_file(directory, native ? "lora_weights" : "adapter_model"); + if (!resolved) throw std::runtime_error("CosyVoice3 adapter checkpoint is missing in " + directory.string()); + file = *resolved; + } + engine::assets::WeightAdapterReadOptions read_options; + read_options.require_scaling_metadata = true; + read_options.strength = runtime::parse_finite_float_option( + options.options, {"cosyvoice3.lora_scale"}).value_or(1.0F); + int64_t native_rank = 0; + if (native) { + const auto config = io::json::parse_file(directory / "config.json"); + const auto & training = config.require("config"); + native_rank = io::json::require_i64(training, "lora_r"); + if (native_rank <= 0) { + throw std::runtime_error("CosyVoice3 LoRA rank must be positive"); + } + read_options.scale_override = io::json::require_f32(training, "lora_alpha") / static_cast(native_rank); + } + auto adapter = engine::assets::read_weight_adapter(file, read_options); + if (native) { + for (const auto & [name, update] : adapter.updates) { + const auto * factors = std::get_if(&update.factors); + if (factors == nullptr || factors->r != native_rank) { + throw std::runtime_error("CosyVoice3 LoRA factors do not match configured lora_r=" + + std::to_string(native_rank) + ": " + name); + } + } + } + engine::assets::WeightAdapterBindings bindings; + for (const auto & tensor : assets->llm_weights->tensors()) { + bindings.add(tensor.name, tensor.name); + if (tensor.name.rfind("llm.model.", 0) == 0) { + bindings.add(tensor.name.substr(10), tensor.name); + } + } + auto plan = engine::assets::bind_weight_adapter( + *assets->llm_weights, std::move(adapter), bindings, engine::assets::WeightMergeMode::DeltaF32); + if (native) { + for (const auto & module : {"speech_embedding", "llm_decoder"}) { + const auto file = engine::assets::find_tensor_file(directory, module); + if (!file) throw std::runtime_error("CosyVoice3 adapter is missing " + std::string(module)); + const auto source = engine::assets::open_tensor_source(*file); + for (const auto & tensor : source->tensors()) { + const auto name = std::string(module) + "." + tensor.name; + engine::assets::require_tensor_shape(*assets->llm_weights, name, tensor.shape); + if (!plan.replacements.emplace(name, engine::assets::WeightReplacement{ + tensor.shape, source->require_f32(tensor.name)}).second) { + throw std::runtime_error("CosyVoice3 adapter supplies duplicate replacement: " + name); + } + } + } + } + auto adapted = std::make_shared(*assets); + adapted->llm_weights = engine::assets::make_weight_adapter_source( + assets->llm_weights, std::move(plan), "cosyvoice3.lora"); + return adapted; } std::shared_ptr require_contract( @@ -106,12 +177,14 @@ CosyVoice3Session::CosyVoice3Session( std::shared_ptr contract) : runtime::RuntimeSessionBase(options), task_(task), - assets_(require_assets(std::move(assets))), + assets_(require_assets(std::move(assets), options)), contract_(require_contract(std::move(contract))), tokenizer_(std::make_unique(assets_)) { auto validation_options = options; // This previously advertised option never selected convolution storage. validation_options.options.erase("cosyvoice3.conv_weight_type"); + validation_options.options.erase("cosyvoice3.lora"); + validation_options.options.erase("cosyvoice3.lora_scale"); // Older embedded contracts omit exec_plan; parse_exec_plan validates its value. if (contract_->session_option_keys.find("cosyvoice3.exec_plan") == contract_->session_option_keys.end()) { validation_options.options.erase("cosyvoice3.exec_plan"); diff --git a/src/models/fish_audio/loader.cpp b/src/models/fish_audio/loader.cpp index 9fbe45455..5076d5784 100644 --- a/src/models/fish_audio/loader.cpp +++ b/src/models/fish_audio/loader.cpp @@ -45,6 +45,8 @@ runtime::ModelCliInterface cli(const FishAudioAssets &) { {"seed", "N", "Sampling seed for reproducible output; omitted uses a random seed.", false, std::nullopt, "0"}, }; out.session_options = { + {"fish_audio.lora", "path", "Prepared Fish Speech adapter checkpoint; merged into slow/fast decoder and embedding weights at session creation. Codec weights are unchanged.", false}, + {"fish_audio.lora_scale", "float", "Multiplier for the adapter's alpha/r scaling; default 1.", false, "1"}, {"fish_audio.mem_saver", "true|false", "Release cached AR runtime graphs after each request; default false.", false, "false"}, {"fish_audio.reference_cache_slots", "n", "Prepared reference-audio cache slots; default 1.", false, "1", "0"}, {"fish_audio.weight_type", "native|f32|f16|bf16|q8_0", "AR matmul weight storage type; default native.", false, "native"}, diff --git a/src/models/fish_audio/session.cpp b/src/models/fish_audio/session.cpp index 4525bb7a6..568b2abe9 100644 --- a/src/models/fish_audio/session.cpp +++ b/src/models/fish_audio/session.cpp @@ -1,5 +1,6 @@ #include "engine/models/fish_audio/session.h" +#include "engine/framework/assets/weight_adapter.h" #include "engine/framework/audio/wav_reader.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/core/execution_context.h" @@ -35,11 +36,46 @@ constexpr int64_t kDefaultReferenceCacheSlots = 1; constexpr const char * kReferenceTextOption = "reference_text"; constexpr const char * kMultiReferenceCondOption = "multi_reference_cond"; -std::shared_ptr require_assets(std::shared_ptr assets) { +std::shared_ptr require_assets( + std::shared_ptr assets, const runtime::SessionOptions & options) { if (assets == nullptr) { throw std::runtime_error("Fish Audio session requires assets"); } - return assets; + const auto path = runtime::find_option(options.options, {"fish_audio.lora"}); + if (!path || path->empty()) { + return assets; + } + engine::assets::WeightAdapterReadOptions read_options; + read_options.require_scaling_metadata = true; + read_options.strength = runtime::parse_finite_float_option( + options.options, {"fish_audio.lora_scale"}).value_or(1.0F); + auto adapter = engine::assets::read_weight_adapter(std::filesystem::u8path(*path), read_options); + engine::assets::WeightAdapterBindings bindings; + const std::string suffix = ".q_proj.weight"; + for (const auto & tensor : assets->model_weights->tensors()) { + bindings.add(tensor.name, tensor.name); + bindings.add("model." + tensor.name, tensor.name); + if (tensor.name.size() < suffix.size() || + tensor.name.compare(tensor.name.size() - suffix.size(), suffix.size(), suffix) != 0) { + continue; + } + const auto prefix = tensor.name.substr(0, tensor.name.size() - suffix.size()); + const auto q_rows = tensor.shape.at(0); + const auto k_name = prefix + ".k_proj.weight"; + const auto v_name = prefix + ".v_proj.weight"; + const auto k_rows = assets->model_weights->require_metadata(k_name).shape.at(0); + const auto v_rows = assets->model_weights->require_metadata(v_name).shape.at(0); + const std::vector targets{ + {tensor.name, 0, q_rows}, {k_name, q_rows, k_rows}, {v_name, q_rows + k_rows, v_rows}}; + bindings.add(prefix + ".wqkv", targets); + bindings.add("model." + prefix + ".wqkv", targets); + } + auto plan = engine::assets::bind_weight_adapter( + *assets->model_weights, std::move(adapter), bindings, engine::assets::WeightMergeMode::DeltaF32); + auto adapted = std::make_shared(*assets); + adapted->model_weights = engine::assets::make_weight_adapter_source( + assets->model_weights, std::move(plan), "fish_audio.lora"); + return adapted; } assets::TensorStorageType option_weight_type( @@ -305,7 +341,7 @@ FishAudioSession::FishAudioSession( std::shared_ptr assets) : RuntimeSessionBase(options), task_(task), - assets_(require_assets(std::move(assets))), + assets_(require_assets(std::move(assets), options)), reference_cache_(resolve_reference_cache_slots(this->options())) { if ((task_.task != runtime::VoiceTaskKind::Tts && task_.task != runtime::VoiceTaskKind::VoiceCloning) || diff --git a/src/models/qwen3_asr/loader.cpp b/src/models/qwen3_asr/loader.cpp index 1962e21c6..f1f56cb55 100644 --- a/src/models/qwen3_asr/loader.cpp +++ b/src/models/qwen3_asr/loader.cpp @@ -40,6 +40,8 @@ runtime::ModelCliInterface cli() { {"preserve_punctuation", "bool", "Opt-in text output mode for timestamped chunked ASR: preserve ASR punctuation in text_output instead of rebuilding text from aligned words.", false, "false"}, }; out.session_options = { + {"qwen3_asr.lora", "path", "PEFT LoRA adapter directory or checkpoint file; merged into matching encoder/decoder weights at session creation.", false}, + {"qwen3_asr.lora_scale", "float", "Multiplier for the adapter's alpha/r scaling; default 1.", false, "1"}, {"qwen3_asr.forced_aligner_path", "path", "Optional Qwen3 ForcedAligner model path required for return_timestamps.", false}, {"qwen3_asr.vad_model_path", "path", "Optional Silero VAD model path for VAD chunking; default assets/framework/models/silero_vad.", false, "assets/framework/models/silero_vad"}, {"qwen3_asr.weight_type", "native|f32|f16|bf16|q8_0", "Shared text decoder weight storage fallback; default native.", false, "native"}, diff --git a/src/models/qwen3_asr/session.cpp b/src/models/qwen3_asr/session.cpp index f8e6f98f4..5fef78a74 100644 --- a/src/models/qwen3_asr/session.cpp +++ b/src/models/qwen3_asr/session.cpp @@ -1,6 +1,7 @@ #include "engine/models/qwen3_asr/session.h" #include "engine/framework/audio/chunking.h" +#include "engine/framework/assets/weight_adapter.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/runtime/options.h" #include "engine/models/qwen3_forced_aligner/session.h" @@ -22,11 +23,40 @@ constexpr double kTimestampFixedChunkContextSeconds = 1.0; constexpr double kPreferredStreamingFeedSeconds = 1.0; constexpr float kDefaultStreamingWindowSeconds = 30.0F; -std::shared_ptr require_assets(std::shared_ptr assets) { +std::shared_ptr require_assets( + std::shared_ptr assets, const runtime::SessionOptions & options) { if (assets == nullptr) { throw std::runtime_error("Qwen3 ASR session requires assets"); } - return assets; + const auto path = runtime::find_option(options.options, {"qwen3_asr.lora"}); + if (!path || path->empty()) { + return assets; + } + engine::assets::WeightAdapterReadOptions read_options; + read_options.require_scaling_metadata = true; + read_options.strength = runtime::parse_finite_float_option( + options.options, {"qwen3_asr.lora_scale"}).value_or(1.0F); + auto adapter = engine::assets::read_weight_adapter(std::filesystem::u8path(*path), read_options); + engine::assets::WeightAdapterBindings bindings; + for (const auto & tensor : assets->model_weights->tensors()) { + bindings.add(tensor.name, tensor.name); + if (tensor.name.rfind("thinker.", 0) == 0) { + bindings.add(tensor.name.substr(8), tensor.name); + } + if (assets->config.hf_transformers_layout) { + if (tensor.name.rfind("model.language_model.", 0) == 0) { + bindings.add("model." + tensor.name.substr(21), tensor.name); + } else if (tensor.name.rfind("model.audio_tower.", 0) == 0) { + bindings.add(tensor.name.substr(6), tensor.name); + } + } + } + auto plan = engine::assets::bind_weight_adapter( + *assets->model_weights, std::move(adapter), bindings, engine::assets::WeightMergeMode::DeltaF32); + auto adapted = std::make_shared(*assets); + adapted->model_weights = engine::assets::make_weight_adapter_source( + assets->model_weights, std::move(plan), "qwen3_asr.lora"); + return adapted; } void validate_matmul_weight_storage(engine::assets::TensorStorageType storage_type, const char * option_name) { @@ -185,7 +215,7 @@ Qwen3ASRSession::Qwen3ASRSession( std::shared_ptr assets) : RuntimeSessionBase(options), task_(task), - assets_(require_assets(std::move(assets))), + assets_(require_assets(std::move(assets), options)), audio_encoder_graph_arena_bytes_(runtime::parse_size_mb_option(options.options, {"qwen3_asr.audio_encoder_graph_arena_mb"}, 128ull * 1024ull * 1024ull)), thinker_prefill_graph_arena_bytes_(runtime::parse_size_mb_option(options.options, {"qwen3_asr.text_decoder_prefill_graph_arena_mb", "qwen3_asr.thinker_prefill_graph_arena_mb"}, 256ull * 1024ull * 1024ull)), thinker_decode_graph_arena_bytes_(runtime::parse_size_mb_option(options.options, {"qwen3_asr.text_decoder_decode_graph_arena_mb", "qwen3_asr.thinker_decode_graph_arena_mb"}, 256ull * 1024ull * 1024ull)), @@ -234,6 +264,8 @@ Qwen3ASRSession::Qwen3ASRSession( key != "qwen3_asr.forced_aligner_path" && key != "qwen3_asr.vad_path" && key != "qwen3_asr.weight_type" && + key != "qwen3_asr.lora" && + key != "qwen3_asr.lora_scale" && key != "qwen3_asr.forced_aligner_model_path" && key != "qwen3_asr.aligner_model_path" && key != "qwen3_asr.vad_model_path") { diff --git a/src/models/qwen3_tts/assets.cpp b/src/models/qwen3_tts/assets.cpp index 63ca5bb0f..7d1efa25e 100644 --- a/src/models/qwen3_tts/assets.cpp +++ b/src/models/qwen3_tts/assets.cpp @@ -1,8 +1,12 @@ #include "engine/models/qwen3_tts/assets.h" +#include "engine/framework/assets/weight_adapter.h" #include "engine/framework/model_spec/package.h" +#include "engine/framework/io/filesystem.h" #include "engine/framework/io/json.h" +#include "engine/framework/runtime/options.h" +#include #include #include @@ -151,6 +155,64 @@ void validate_config(const Qwen3TTSConfig & config) { } // namespace +std::shared_ptr apply_qwen3_tts_lora( + std::shared_ptr assets, + const std::unordered_map & options) { + const auto path = runtime::find_option(options, {"qwen3_tts.lora"}); + if (!path || path->empty()) { + return assets; + } + engine::assets::WeightAdapterReadOptions read_options; + read_options.require_scaling_metadata = true; + read_options.strength = runtime::parse_finite_float_option(options, {"qwen3_tts.lora_scale"}).value_or(1.0F); + const auto adapter_path = std::filesystem::u8path(*path); + auto adapter = engine::assets::read_weight_adapter(adapter_path, read_options); + engine::assets::WeightAdapterBindings bindings; + for (const auto & tensor : assets->model_weights->tensors()) { + bindings.add(tensor.name, tensor.name); + if (tensor.name.rfind("talker.", 0) == 0) { + bindings.add(tensor.name.substr(7), tensor.name); + } + } + auto plan = engine::assets::bind_weight_adapter( + *assets->model_weights, std::move(adapter), bindings, engine::assets::WeightMergeMode::DeltaF32); + auto adapted = std::make_shared(*assets); + const auto directory = io::is_existing_directory(adapter_path) ? adapter_path : adapter_path.parent_path(); + const auto speaker_file = engine::assets::find_tensor_file(directory, "speaker_embedding"); + if (speaker_file) { + const auto config = json::parse_file(directory / "tts_config.json"); + if (json::require_string(config, "tts_model_type") != "custom_voice") { + throw std::runtime_error("Qwen3 TTS adapter speaker embedding requires a custom_voice tts_config.json"); + } + const auto talker = parse_talker_config(config.require("talker_config")); + if (talker.hidden_size != assets->config.talker.hidden_size || + talker.vocab_size != assets->config.talker.vocab_size || talker.speaker_id.size() != 1) { + throw std::runtime_error("Qwen3 TTS adapter speaker embedding requires one speaker and matching base dimensions"); + } + const auto id = talker.speaker_id.begin()->second; + if (id < 0 || id >= talker.vocab_size) { + throw std::runtime_error("Qwen3 TTS adapter speaker ID is outside the codec embedding vocabulary"); + } + const auto speaker_source = engine::assets::open_tensor_source(*speaker_file); + const auto speaker = speaker_source->require_f32("speaker_embedding", {talker.hidden_size}); + const std::string name = "talker.model.codec_embedding.weight"; + if (plan.updates.count(name) || plan.replacements.count(name)) { + throw std::runtime_error("Qwen3 TTS adapter supplies both codec embedding updates and a speaker sidecar"); + } + auto embedding = assets->model_weights->require_f32(name, {talker.vocab_size, talker.hidden_size}); + std::copy(speaker.begin(), speaker.end(), embedding.begin() + id * talker.hidden_size); + plan.replacements.emplace(name, engine::assets::WeightReplacement{ + {talker.vocab_size, talker.hidden_size}, std::move(embedding)}); + adapted->config.variant = Qwen3TTSVariant::CustomVoice; + adapted->config.tts_model_type = "custom_voice"; + adapted->config.talker.speaker_id = talker.speaker_id; + adapted->config.talker.speaker_dialect = talker.speaker_dialect; + } + adapted->model_weights = engine::assets::make_weight_adapter_source( + assets->model_weights, std::move(plan), "qwen3_tts.lora"); + return adapted; +} + std::shared_ptr load_qwen3_tts_assets(const std::filesystem::path & model_path) { auto resources = engine::model_spec::load_resource_bundle( model_path, diff --git a/src/models/qwen3_tts/loader.cpp b/src/models/qwen3_tts/loader.cpp index 54d5f830d..697560e7d 100644 --- a/src/models/qwen3_tts/loader.cpp +++ b/src/models/qwen3_tts/loader.cpp @@ -81,6 +81,8 @@ runtime::ModelCliInterface cli(const Qwen3TTSAssets &) { {"text_chunk_size", "int", "Maximum UTF-8 codepoints per long-form text chunk; default 8192.", false, "8192", "1"}, }; out.session_options = { + {"qwen3_tts.lora", "path", "PEFT LoRA adapter directory or checkpoint file; optional speaker_embedding and tts_config.json sidecars provide a trained custom voice.", false}, + {"qwen3_tts.lora_scale", "float", "Multiplier for the adapter's alpha/r scaling; default 1.", false, "1"}, {"qwen3_tts.mem_saver", "bool", "Release the talker cached-step graph after each request; default false.", false, "false"}, {"qwen3_tts.stitch_decode", "bool", "Voice the concatenated speech codes of all text chunks in one decoder pass instead of voicing each chunk separately; default false. Per-chunk voicing re-voices the voice-clone reference once per chunk and restarts the stateful vocoder cache at every chunk boundary; the stitched pass costs device memory proportional to the utterance.", false, "false"}, {"qwen3_tts.perf_mode", "off|flash_attention", "Q8_0-only attention performance mode; default off keeps the exact-safe path.", false, "off"}, diff --git a/src/models/qwen3_tts/session.cpp b/src/models/qwen3_tts/session.cpp index 2e41b735c..706e6c917 100644 --- a/src/models/qwen3_tts/session.cpp +++ b/src/models/qwen3_tts/session.cpp @@ -288,7 +288,7 @@ Qwen3TTSSession::Qwen3TTSSession( std::shared_ptr assets) : RuntimeSessionBase(options), task_(task), - assets_(require_assets(std::move(assets))), + assets_(apply_qwen3_tts_lora(require_assets(std::move(assets)), options.options)), mem_saver_(mem_saver_from_options(options)), stitch_decode_(stitch_decode_from_options(options)), perf_mode_(perf_mode_from_options(options)), @@ -362,6 +362,8 @@ Qwen3TTSSession::Qwen3TTSSession( key != "qwen3_tts.code_predictor_constant_context_mb" && key != "qwen3_tts.speech_decoder_constant_context_mb" && key != "qwen3_tts.weight_type" && + key != "qwen3_tts.lora" && + key != "qwen3_tts.lora_scale" && key != "qwen3_tts.conv_weight_type" && key != "qwen3_tts.talker_weight_type" && key != "qwen3_tts.speech_encoder_weight_type" && diff --git a/src/models/vibevoice/lora.cpp b/src/models/vibevoice/lora.cpp index c6d2b5cd6..723ae9786 100644 --- a/src/models/vibevoice/lora.cpp +++ b/src/models/vibevoice/lora.cpp @@ -1,6 +1,5 @@ #include "engine/models/vibevoice/lora.h" -#include "engine/framework/assets/torch_bin.h" #include "engine/framework/assets/weight_adapter.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/debug/trace.h" @@ -35,19 +34,13 @@ std::filesystem::path resolve_adapter_dir(const std::filesystem::path & adapter_ return adapter_path.parent_path(); } -// Opens the diffusion-head fine-tune weights, preferring safetensors over the pickled variants. +// Preserve the upstream package locations; shared I/O selects the container. std::shared_ptr open_diffusion_head_source(const std::filesystem::path & adapter_dir) { - const auto safetensors = adapter_dir / "diffusion_head" / "model.safetensors"; - if (engine::io::is_existing_file(safetensors)) { - return assets::open_tensor_source(safetensors); - } - const auto bin = adapter_dir / "diffusion_head_full.bin"; - if (engine::io::is_existing_file(bin)) { - return assets::open_torch_bin_tensor_source(bin); - } - const auto nested_bin = adapter_dir / "diffusion_head" / "diffusion_head_full.bin"; - if (engine::io::is_existing_file(nested_bin)) { - return assets::open_torch_bin_tensor_source(nested_bin); + for (const auto * name : {"diffusion_head/model", "diffusion_head_full", "diffusion_head/diffusion_head_full"}) { + const auto logical = adapter_dir / name; + if (const auto file = assets::find_tensor_file(logical.parent_path(), logical.filename().string())) { + return assets::open_tensor_source(*file); + } } return nullptr; } @@ -87,12 +80,12 @@ void add_connector_overrides( std::string_view base_prefix, const char * label, std::map & overrides) { - const auto path = adapter_dir / subdir / "pytorch_model.bin"; - if (!engine::io::is_existing_file(path)) { + const auto path = assets::find_tensor_file(adapter_dir / subdir, "pytorch_model"); + if (!path) { log_line(std::string(label) + ": not present, skipped"); return; } - const auto source = assets::open_torch_bin_tensor_source(path); + const auto source = assets::open_tensor_source(*path); const int count = add_full_overrides(base, *source, base_prefix, overrides); log_line(std::string(label) + ": overrode " + std::to_string(count) + " tensors"); } diff --git a/src/models/voxcpm2/loader.cpp b/src/models/voxcpm2/loader.cpp index 5bd2e89d3..4a159704e 100644 --- a/src/models/voxcpm2/loader.cpp +++ b/src/models/voxcpm2/loader.cpp @@ -53,6 +53,8 @@ runtime::ModelCliInterface cli(const VoxCPM2Assets &) { {"text_chunk_mode", "default|tag_aware|japanese|endline", "Framework text chunking mode. Offline only. Default: tag_aware.", false, "tag_aware"}, }; out.session_options = { + {"voxcpm2.lora", "path", "Native LoRA directory or checkpoint file with adjacent lora_config.json; merged into generator weights at session creation.", false}, + {"voxcpm2.lora_scale", "float", "Multiplier for the adapter's configured alpha/r scaling; default 1.", false, "1"}, {"voxcpm2.weight_type", "native|f32|f16|bf16|q8_0", "Language-model weight storage; does not override AudioVAE storage.", false}, {"voxcpm2.audiovae_weight_type", "native|f32|f16|bf16|q8_0", "AudioVAE weight storage; defaults independently to native, not the shared weight_type."}, {"voxcpm2.prompt_cache_slots", "int", "Prompt and prompt-audio embedding cache slots; default 1.", false, "1", "0"}, diff --git a/src/models/voxcpm2/session.cpp b/src/models/voxcpm2/session.cpp index 785925ec8..64b724e15 100644 --- a/src/models/voxcpm2/session.cpp +++ b/src/models/voxcpm2/session.cpp @@ -1,5 +1,8 @@ #include "engine/models/voxcpm2/session.h" +#include "engine/framework/assets/weight_adapter.h" +#include "engine/framework/io/filesystem.h" +#include "engine/framework/io/json.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/io/text.h" #include "engine/framework/runtime/errors.h" @@ -62,11 +65,49 @@ const char *chunk_strategy_name(ChunkStrategy strategy) { } std::shared_ptr -require_assets(std::shared_ptr assets) { +require_assets(std::shared_ptr assets, + const runtime::SessionOptions & options) { if (assets == nullptr) { throw std::runtime_error("VoxCPM2 session requires assets"); } - return assets; + const auto path = runtime::find_option(options.options, {"voxcpm2.lora"}); + if (!path || path->empty()) { + return assets; + } + const auto adapter_path = std::filesystem::u8path(*path); + auto file = adapter_path; + if (io::is_existing_directory(adapter_path)) { + const auto resolved = engine::assets::find_tensor_file(adapter_path, "lora_weights"); + if (!resolved) throw std::runtime_error("VoxCPM2 lora_weights checkpoint is missing in " + adapter_path.string()); + file = *resolved; + } + const auto config = io::json::parse_file(file.parent_path() / "lora_config.json"); + const auto & lora = config.require("lora_config"); + const auto rank = io::json::require_i64(lora, "r"); + if (rank <= 0) { + throw std::runtime_error("VoxCPM2 LoRA rank must be positive"); + } + engine::assets::WeightAdapterReadOptions read_options; + read_options.scale_override = io::json::require_f32(lora, "alpha") / static_cast(rank); + read_options.strength = runtime::parse_finite_float_option( + options.options, {"voxcpm2.lora_scale"}).value_or(1.0F); + auto adapter = engine::assets::read_weight_adapter(file, read_options); + for (const auto & [name, update] : adapter.updates) { + const auto * factors = std::get_if(&update.factors); + if (factors == nullptr || factors->r != rank) { + throw std::runtime_error("VoxCPM2 LoRA factors do not match lora_config.json: " + name); + } + } + engine::assets::WeightAdapterBindings bindings; + for (const auto & tensor : assets->model_weights->tensors()) { + bindings.add(tensor.name, tensor.name); + } + auto plan = engine::assets::bind_weight_adapter( + *assets->model_weights, std::move(adapter), bindings, engine::assets::WeightMergeMode::DeltaF32); + auto adapted = std::make_shared(*assets); + adapted->model_weights = engine::assets::make_weight_adapter_source( + assets->model_weights, std::move(plan), "voxcpm2.lora"); + return adapted; } void reject_enabled_denoise( @@ -196,7 +237,8 @@ void validate_session_options( if (key.rfind("voxcpm2.", 0) != 0) { continue; } - if (key == "voxcpm2.weight_context_mb" || + if (key == "voxcpm2.lora" || key == "voxcpm2.lora_scale" || + key == "voxcpm2.weight_context_mb" || key == "voxcpm2.text_embedding_graph_context_mb" || key == "voxcpm2.text_embedding_graph_arena_mb" || key == "voxcpm2.lm_step_graph_context_mb" || @@ -251,7 +293,7 @@ VoxCPM2SessionBase::VoxCPM2SessionBase(runtime::TaskSpec task, runtime::SessionOptions options, std::shared_ptr assets) : RuntimeSessionBase(options), task_(task), - assets_(require_assets(std::move(assets))), + assets_(require_assets(std::move(assets), options)), encoded_prompt_cache_(prompt_cache_slots_from_options(options.options)) { if (task_.mode != runtime::RunMode::Offline && task_.mode != runtime::RunMode::Streaming) { diff --git a/tests/lora/test_merge_lora.py b/tests/lora/test_merge_lora.py new file mode 100644 index 000000000..68587aab2 --- /dev/null +++ b/tests/lora/test_merge_lora.py @@ -0,0 +1,326 @@ +import json +from pathlib import Path +import sys +import tempfile +import unittest +from unittest.mock import patch + +import torch +import gguf +import numpy as np +from safetensors.torch import save_file +from safetensors import SafetensorError + +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tools/lora")) +from bindings import BaseGGUF, BaseSources, Bindings +from merge_lora import main, normalize, read_adapter, sidecars +from checkpoint import EXTENSIONS, find_checkpoint, load_checkpoint + + +class MergeLoraTests(unittest.TestCase): + def test_cosy_mixed_package_selection_and_sidecars(self): + name = "model.layers.0.q_proj" + bindings = Bindings("cosyvoice3", { + "llm/llm.model." + name + ".weight": (4, 4), + "llm/speech_embedding.weight": (4, 4), + "llm/llm_decoder.weight": (4, 4), + }) + native = {name + ".lora_A": torch.ones(2, 4), name + ".lora_B": torch.ones(4, 2)} + peft = {name + ".lora_A.weight": torch.full((1, 4), 3.), + name + ".lora_B.weight": torch.full((4, 1), 7.)} + for extension in EXTENSIONS: + with self.subTest(extension=extension), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "config.json").write_text(json.dumps({"config": {"lora_r": 2, "lora_alpha": 4}})) + (root / "adapter_config.json").write_text(json.dumps({"r": 1, "lora_alpha": 1})) + for stem, state in (("lora_weights", native), ("adapter_model", peft), + ("speech_embedding", {"weight": torch.full((4, 4), 5.)}), + ("llm_decoder", {"weight": torch.full((4, 4), 6.)})): + file = root / (stem + extension) + if extension == ".safetensors": + save_file(state, file) + else: + torch.save(state, file) + for path, is_native in ((root, True), (root / ("lora_weights" + extension), True), + (root / ("adapter_model" + extension), False)): + loaded, config, _, native_config = read_adapter(path, "cosyvoice3", root, None) + self.assertEqual(native_config is not None, is_native) + output = {} + normalize(loaded, config, bindings, output, 1, {}, native_scaling=is_native) + target = "llm/llm.model." + name + ".weight" + self.assertEqual(output[target + ".alpha"].item(), 4 if is_native else 1) + expected_a = native[name + ".lora_A"] if is_native else peft[name + ".lora_A.weight"] + self.assertTrue(torch.equal(output[target + ".lora_A"], expected_a)) + sidecars("cosyvoice3", root, native_config, None, bindings, output, root) + self.assertEqual("llm/speech_embedding.weight" in output, is_native) + self.assertEqual("llm/llm_decoder.weight" in output, is_native) + (root / ("speech_embedding" + extension)).unlink() + for path in (root, root / ("lora_weights" + extension)): + _, _, _, native_config = read_adapter(path, "cosyvoice3", root, None) + with self.assertRaises(FileNotFoundError): + sidecars("cosyvoice3", root, native_config, None, bindings, {}, root) + + def test_cosy_native_rank_mismatch_is_rejected(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "config.json").write_text(json.dumps({"config": {"lora_r": 3, "lora_alpha": 6}})) + name = "model.layers.0.q_proj" + save_file({name + ".lora_A": torch.ones(2, 4), name + ".lora_B": torch.ones(4, 2)}, root / "lora_weights.safetensors") + loaded, config, _, native_config = read_adapter(root, "cosyvoice3", root, None) + bindings = Bindings("cosyvoice3", {"llm/llm.model." + name + ".weight": (4, 4)}) + with self.assertRaisesRegex(ValueError, "Configured rank differs"): + normalize(loaded, config, bindings, {}, 1, {}, native_scaling=native_config is not None) + + def test_adapter_layout_is_independent_of_container(self): + for native in (False, True): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + if native: + (root / "config.json").write_text(json.dumps({"config": {"lora_r": 2, "lora_alpha": 4}})) + else: + (root / "adapter_config.json").write_text(json.dumps({"r": 2, "lora_alpha": 4})) + stem = "lora_weights" if native else "adapter_model" + suffix = "" if native else ".weight" + name = "model.layers.0.q_proj" + state = {name + ".lora_A" + suffix: torch.arange(8.).reshape(2, 4), + name + ".lora_B" + suffix: torch.arange(8.).reshape(4, 2)} + bindings = Bindings("cosyvoice3", {"llm/llm.model." + name + ".weight": (4, 4)}) + expected = None + for extension in EXTENSIONS: + with self.subTest(native=native, extension=extension): + path = root / (stem + extension) + if extension == ".safetensors": + save_file(state, path) + else: + torch.save(state, path) + self.assertEqual(find_checkpoint(root, stem), path) + for supplied_path in (root, path): + loaded, config, _, native_config = read_adapter(supplied_path, "cosyvoice3", root, None) + self.assertEqual(native_config is not None, native) + output = {} + normalize(loaded, config, bindings, output, 1, {}, native_scaling=native) + if expected is None: + expected = output + self.assertEqual(output.keys(), expected.keys()) + for key in output: + self.assertTrue(torch.equal(output[key], expected[key]), key) + path.unlink() + + def test_checkpoint_discovery_does_not_hide_corruption(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + torch.save({"weight": torch.ones(2)}, root / "adapter_model.pt") + invalid = root / "adapter_model.safetensors" + invalid.write_bytes(b"invalid") + self.assertEqual(find_checkpoint(root, "adapter_model"), invalid) + with self.assertRaises(SafetensorError): + load_checkpoint(invalid) + self.assertIsNone(find_checkpoint(root, "missing", required=False)) + with self.assertRaises(FileNotFoundError): + find_checkpoint(root, "missing") + + def test_fish_training_layout_is_independent_of_container(self): + state = {"model.layers.0.attention.wqkv.lora_A": torch.ones(2, 4), + "model.layers.0.attention.wqkv.lora_B": torch.ones(8, 2)} + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + for extension in EXTENSIONS: + with self.subTest(extension=extension): + path = root / ("training" + extension) + if extension == ".safetensors": + save_file(state, path) + else: + torch.save({"state_dict": state, "epoch": 1}, path) + loaded, _, _, _ = read_adapter(path, "fish_audio", root, 4) + self.assertTrue(torch.equal(loaded["layers.0.attention.wqkv.lora_A"], state["model.layers.0.attention.wqkv.lora_A"])) + self.assertEqual(loaded["layers.0.attention.wqkv.alpha"].item(), 4) + + def test_prepared_fish_scaling_is_not_overwritten(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + path = root / "prepared.safetensors" + state = {"layers.0.attention.wqkv.lora_A": torch.ones(2, 4), + "layers.0.attention.wqkv.lora_B": torch.ones(8, 2), + "layers.0.attention.wqkv.alpha": torch.tensor(4.)} + save_file(state, path) + loaded, _, _, _ = read_adapter(path, "fish_audio", root, 16) + self.assertEqual(loaded.keys(), state.keys()) + for name in state: + self.assertTrue(torch.equal(loaded[name], state[name])) + + def test_embedded_byte_array_is_read_in_bulk(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + file = root / "base.gguf" + writer = gguf.GGUFWriter(file, "test") + writer.add_array("audiocpp.embedded_files.names", ["vocab.txt"]) + writer.add_key_value("audiocpp.embedded_files.offsets", [0, 3], + gguf.GGUFValueType.ARRAY, gguf.GGUFValueType.UINT64) + writer.add_array("audiocpp.embedded_files.data", b"ABC") + writer.add_tensor("projection.weight", np.ones((2, 32), dtype=np.float32)) + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + base = BaseGGUF(file) + self.assertEqual(base.reader.fields["audiocpp.embedded_files.data"].contents(), [65, 66, 67]) + self.assertEqual(base.shapes["projection.weight"], (2, 32)) + self.assertTrue(torch.equal(base.read_f32("projection.weight"), torch.ones(2, 32))) + destination = root / "sidecars" + base.extract_sidecars(destination) + self.assertEqual((destination / "vocab.txt").read_bytes(), b"ABC") + + def test_all_family_bindings(self): + cases = [ + ("qwen3_tts", "model_weights/talker.model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.q_proj"), + ("qwen3_asr", "model.language_model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.q_proj"), + ("voxcpm2", "weights/base_lm.layers.0.self_attn.q_proj.weight", "base_lm.layers.0.self_attn.q_proj"), + ("voxcpm1", "blk.0.attn_q.weight", "base_lm.layers.0.self_attn.q_proj"), + ("fish_audio", "model_weights/fast_output.weight", "model.fast_output"), + ("cosyvoice3", "llm/llm.model.model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.q_proj"), + ("f5_tts", "transformer/ema_model.transformer.transformer_blocks.0.attn.to_q.weight", "transformer.transformer_blocks.0.attn.to_q"), + ("outetts", "model_weights/model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.q_proj"), + ("yue2", "model.layers.0.self_attn.q_proj.weight", "layers.0.self_attn.q_proj"), + ("ace_step", "dit_turbo_weights/decoder.layers.0.q_proj.weight", "layers.0.q_proj"), + ("vibevoice", "model.language_model.layers.0.self_attn.q_proj.weight", "layers.0.self_attn.q_proj"), + ] + for family, name, external in cases: + with self.subTest(family=family): + shapes = {name: (4, 4)} + if family == "yue2": + shapes.update({name.replace("q_proj", part): (4, 4) for part in ("k_proj", "v_proj")}) + bindings = Bindings(family, shapes) + self.assertEqual(bindings.resolve("base_model.model." + external)[0].name, name) + + def test_fused_fish_slices_keep_all_rows(self): + shapes = {f"model_weights/layers.0.attention.{part}_proj.weight": (rows, 4) + for part, rows in (("q", 4), ("k", 2), ("v", 2))} + bindings = Bindings("fish_audio", shapes) + a = torch.arange(8, dtype=torch.float32).reshape(2, 4) + b = torch.arange(16, dtype=torch.float32).reshape(8, 2) + state = {"model.layers.0.attention.wqkv.lora_A": a, + "model.layers.0.attention.wqkv.lora_B": b, + "model.layers.0.attention.wqkv.alpha": torch.tensor(4.)} + output = {} + normalize(state, {}, bindings, output, 1, {}) + offset = 0 + for part, rows in (("q", 4), ("k", 2), ("v", 2)): + name = f"model_weights/layers.0.attention.{part}_proj.weight" + self.assertTrue(torch.equal(output[name + ".lora_A"], a)) + self.assertTrue(torch.equal(output[name + ".lora_B"], b[offset:offset + rows])) + self.assertEqual(output[name + ".alpha"].item(), 4) + offset += rows + with tempfile.TemporaryDirectory() as directory: + save_file(output, Path(directory) / "normalized.safetensors") + + def test_fused_slices_reject_missing_and_excess_rows(self): + for family, prefix, fused in ( + ("fish_audio", "model_weights/layers.0.attention", "model.layers.0.attention.wqkv"), + ("yue2", "model.layers.0.self_attn", "layers.0.self_attn.qkv_proj"), + ): + shapes = {f"{prefix}.{part}_proj.weight": (rows, 4) + for part, rows in (("q", 4), ("k", 2), ("v", 2))} + bindings = Bindings(family, shapes) + for rows in (7, 9): + with self.subTest(family=family, rows=rows): + state = {fused + ".lora_A": torch.ones(2, 4), + fused + ".lora_B": torch.ones(rows, 2), + fused + ".alpha": torch.tensor(4.)} + output = {} + with self.assertRaisesRegex(ValueError, "output row"): + normalize(state, {}, bindings, output, 1, {}) + self.assertEqual(output, {}) + + def test_mixed_lora_and_lokr_factors_are_rejected(self): + name = "model.layers.0.q_proj" + bindings = Bindings("qwen3_asr", {name + ".weight": (4, 4)}) + for factor in ("lokr_w1", "lokr_w2", "lokr_w1_a", "lokr_w1_b", "lokr_w2_a", "lokr_w2_b"): + with self.subTest(factor=factor): + state = {name + ".lora_A": torch.ones(2, 4), + name + ".lora_B": torch.ones(4, 2), + name + ".alpha": torch.tensor(2.), + name + "." + factor: torch.ones(2, 2)} + output = {} + with self.assertRaisesRegex(ValueError, "Mixed adapter algorithms"): + normalize(state, {}, bindings, output, 1, {}) + self.assertEqual(output, {}) + + def test_peft_patterns_and_rslora_are_preserved(self): + name = "model.layers.0.q_proj.weight" + bindings = Bindings("qwen3_asr", {name: (4, 4)}) + source = "base_model.model.model.layers.0.q_proj" + state = {source + ".lora_A.weight": torch.ones(2, 4), source + ".lora_B.weight": torch.ones(4, 2)} + config = {"r": 8, "lora_alpha": 16, "rank_pattern": {"q_proj": 2}, "alpha_pattern": {"q_proj": 6}, "use_rslora": True} + rs = {"rank_pattern": {}, "alpha_pattern": {}} + output = {} + normalize(state, config, bindings, output, 1, rs) + self.assertEqual(list(rs["rank_pattern"].values()), [2]) + self.assertEqual(list(rs["alpha_pattern"].values()), [6]) + self.assertIn(name + ".lora_A.weight", output) + + def test_duplicate_aliases_cannot_overwrite_rslora(self): + name = "model.layers.0.q_proj.weight" + bindings = Bindings("qwen3_asr", {name: (4, 4)}) + source = "base_model.model.model.layers.0.q_proj" + state = {source + ".lora_A.weight": torch.ones(2, 4), + source + ".lora_B.weight": torch.ones(4, 2)} + config = {"r": 2, "lora_alpha": 2, "use_rslora": True} + output = {} + normalize(state, config, bindings, output, 1, {"rank_pattern": {}, "alpha_pattern": {}}) + with self.assertRaisesRegex(ValueError, "Multiple updates"): + normalize(state, config, bindings, output, 1, {"rank_pattern": {}, "alpha_pattern": {}}) + + def test_native_does_not_inherit_neighbor_peft_config(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + file = root / "native.safetensors" + save_file({"layers.0.self_attn.q_proj.lora_A": torch.ones(2, 4), + "layers.0.self_attn.q_proj.lora_B": torch.ones(4, 2)}, file) + (root / "adapter_config.json").write_text(json.dumps({"r": 128, "lora_alpha": 256, "use_dora": True})) + state, config, _, _ = read_adapter(file, "yue2", root, None) + self.assertEqual(config, {}) + output = {} + shapes = {f"model.layers.0.self_attn.{part}_proj.weight": (4, 4) for part in ("q", "k", "v")} + normalize(state, config, Bindings("yue2", shapes), + output, 1, {}, require_scaling=False) + self.assertEqual(output["model.layers.0.self_attn.q_proj.weight.alpha"].item(), 2) + + def test_missing_peft_metadata_and_empty_checkpoint_are_rejected(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + file = root / "adapter.safetensors" + save_file({"base_model.model.proj.lora_A.weight": torch.ones(2, 4), + "base_model.model.proj.lora_B.weight": torch.ones(4, 2)}, file) + with self.assertRaisesRegex(ValueError, "adapter_config.json"): + read_adapter(file, "qwen3_asr", root, None) + save_file({}, file) + with self.assertRaisesRegex(ValueError, "empty"): + read_adapter(file, "qwen3_asr", root, None) + + def test_component_ambiguity_is_explicit(self): + shapes = {f"{component}/decoder.proj.weight": (4, 4) for component in ("dit_base_weights", "dit_turbo_weights")} + with self.assertRaisesRegex(ValueError, "Ambiguous"): + Bindings("ace_step", shapes) + bindings = Bindings("ace_step", shapes, "dit_turbo_weights") + self.assertEqual(bindings.resolve("proj")[0].name, "dit_turbo_weights/decoder.proj.weight") + + def test_original_sharded_inputs(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + save_file({"model.proj.weight": torch.ones(4, 4)}, root / "part.safetensors") + index = root / "model.safetensors.index.json" + index.write_text(json.dumps({"weight_map": {"model.proj.weight": "part.safetensors"}})) + base = BaseSources([f"model_weights={index}"], root) + self.assertEqual(base.shapes, {"model_weights/model.proj.weight": (4, 4)}) + self.assertTrue(torch.equal(base.read_f32("model_weights/model.proj.weight"), torch.ones(4, 4))) + original = (root / "part.safetensors").read_bytes() + arguments = ["merge_lora.py", "--family", "qwen3_asr", "--input", str(index), + "--output", str(root / "part.safetensors"), "--overwrite", "--adapter", "unused"] + with patch.object(sys, "argv", arguments), self.assertRaises(SystemExit) as error: + main() + self.assertEqual(error.exception.code, 2) + self.assertEqual((root / "part.safetensors").read_bytes(), original) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittests/test_gguf_tensor_source.cpp b/tests/unittests/test_gguf_tensor_source.cpp index a3d9715f1..5e436ad2d 100644 --- a/tests/unittests/test_gguf_tensor_source.cpp +++ b/tests/unittests/test_gguf_tensor_source.cpp @@ -1,5 +1,6 @@ #include "engine/framework/model_spec/package.h" #include "engine/framework/assets/tensor_source.h" +#include "engine/framework/assets/weight_adapter.h" #include "engine/framework/assets/resource_bundle.h" #include "engine/framework/io/filesystem.h" #include "engine/framework/io/safetensors.h" @@ -7,6 +8,7 @@ #include "test_assert.h" #include +#include #include #include #include @@ -532,6 +534,62 @@ void test_malformed_embedded_sidecars_are_rejected_by_the_name_reader() { std::filesystem::remove_all(root); } +void test_weight_adapter_conversion() { + using namespace engine; + const auto root = std::filesystem::temp_directory_path() / "audiocpp_gguf_adapter_test"; + std::filesystem::remove_all(root); + std::filesystem::create_directories(root); + const auto original = root / "original.safetensors"; + io::write_safetensors_file(original, { + {"projection.weight", "F32", {2, 32}, bytes_for(std::vector(64, 0.25F))}, + {"untouched.weight", "F32", {2, 32}, bytes_for(std::vector(64, 0.5F))}, + {"step", "F32", {}, bytes_for(std::vector{42.0F})}, + }); + const auto base_path = root / "base.gguf"; + assets::convert_tensor_sources_to_gguf({{original, ""}}, base_path, assets::TensorStorageType::F16, + false, false); + const auto base = assets::open_tensor_source(base_path); + const auto original_bytes = base->require_tensor_data("projection.weight").bytes; + auto adapter = std::make_shared(); + adapter->updates.emplace("projection.weight", assets::WeightUpdate{ + assets::LowRankWeightUpdate{std::vector(32, 0.5F), {1.0F, -1.0F}, 1, 32, 2}, + 0.5F, assets::WeightMergeMode::DeltaF32}); + const auto overlay = assets::make_weight_adapter_source(base, *adapter); + int index = 0; + for (auto type : {assets::TensorStorageType::Native, assets::TensorStorageType::F16, + assets::TensorStorageType::BF16, assets::TensorStorageType::Q8_0}) { + assets::GgufConversionOptions options; + options.weight_adapter = adapter; + const auto output = root / (std::to_string(index++) + ".gguf"); + assets::convert_tensor_sources_to_gguf({{base_path, ""}}, output, type, false, false, + {}, {}, std::nullopt, options); + const auto actual = assets::open_tensor_source(output); + test::require(actual->require_tensor_data("step").bytes == base->require_tensor_data("step").bytes, + "Adapter conversion must preserve untouched scalar metadata"); + for (const auto & name : {"projection.weight", "untouched.weight"}) { + const auto expected = overlay->require_tensor(name, type, {2, 32}); + test::require(actual->require_tensor_data(name).bytes == expected.bytes, + "GGUF adapter must merge before dtype conversion, including orig"); + } + } + test::require(base->require_tensor_data("projection.weight").bytes == original_bytes, + "GGUF adapter conversion changed base weights"); + auto replacement = std::make_shared(); + replacement->replacements.emplace("projection.weight", + assets::WeightReplacement{{2, 32}, std::vector(64, 0.1001F)}); + assets::GgufConversionOptions replacement_options; + replacement_options.weight_adapter = replacement; + replacement_options.weight_adapter_merge_mode = assets::WeightMergeMode::RoundedBF16Delta; + const auto replacement_path = root / "replacement.gguf"; + assets::convert_tensor_sources_to_gguf({{base_path, ""}}, replacement_path, + assets::TensorStorageType::F32, false, false, {}, {}, std::nullopt, replacement_options); + const auto rounded = assets::open_tensor_source(replacement_path)->require_f32("projection.weight"); + const float want = ggml_bf16_to_fp32(ggml_fp32_to_bf16(0.1001F)); + test::require(std::all_of(rounded.begin(), rounded.end(), [want](float value) { return value == want; }), + "Replacement-only export must retain the selected BF16 merge policy"); + std::filesystem::remove_all(root); +} + } // namespace int main() { @@ -546,6 +604,7 @@ int main() { test_release_named_gguf_directory_selects_the_gguf_source(); test_package_spec_errors_name_selected_spec(); test_malformed_embedded_sidecars_are_rejected_by_the_name_reader(); + test_weight_adapter_conversion(); } catch (const std::exception & error) { std::cerr << error.what() << '\n'; return 1; diff --git a/tests/unittests/test_torch_bin_storage_offset.cpp b/tests/unittests/test_torch_bin_storage_offset.cpp index 4ca44b8c2..fa8e29c84 100644 --- a/tests/unittests/test_torch_bin_storage_offset.cpp +++ b/tests/unittests/test_torch_bin_storage_offset.cpp @@ -23,6 +23,7 @@ // rather than read in storage order. #include "engine/framework/assets/torch_bin.h" +#include "engine/framework/assets/weight_adapter.h" #include #include @@ -223,6 +224,30 @@ void test_views_honour_storage_offset() { format(tail)); } +void test_container_extensions_do_not_select_adapter_layout() { + const auto original = write_checkpoint({ + {"projection.lora_A", 0, {1, 4}, {4, 1}}, + {"projection.lora_B", 4, {3, 1}, {1, 1}}}, 7); + const auto root = std::filesystem::temp_directory_path() / "audiocpp_tensor_container_test"; + std::filesystem::create_directories(root); + for (const auto * extension : {".pt", ".pth", ".bin", ".ckpt"}) { + const auto file = root / (std::string("adapter_model") + extension); + std::filesystem::copy_file(original, file, std::filesystem::copy_options::overwrite_existing); + require(engine::assets::find_tensor_file(root, "adapter_model") == file, + "logical checkpoint resolution differs by container extension"); + const auto source = engine::assets::open_tensor_source(file); + require(source->require_f32("projection.lora_B") == iota(4, 3), + "generic tensor reader lost storage offsets"); + const auto adapter = engine::assets::read_weight_adapter(root); + const auto & factors = std::get(adapter.updates.at("projection").factors); + require(factors.a == iota(0, 4) && factors.b == iota(4, 3), + "adapter factors differ by container extension"); + std::filesystem::remove(file); + } + std::filesystem::remove_all(root); + std::filesystem::remove(original); +} + void test_chunked_matrix_honours_storage_offset() { // w = torch.arange(18.).reshape(6, 3); q, k = w.chunk(2, 0); torch.save({"q": q, "k": k}, path) // k is w[3:]: storage_offset 9, shape (3, 3), stride (3, 1). @@ -270,6 +295,7 @@ void test_unsupported_views_are_rejected() { int main() { try { test_views_honour_storage_offset(); + test_container_extensions_do_not_select_adapter_layout(); test_chunked_matrix_honours_storage_offset(); test_size_one_dims_ignore_stride(); test_unsupported_views_are_rejected(); diff --git a/tools/lora/bindings.py b/tools/lora/bindings.py new file mode 100644 index 000000000..deafc8122 --- /dev/null +++ b/tools/lora/bindings.py @@ -0,0 +1,257 @@ +"""Model-local adapter names bound to complete audio.cpp GGUF tensor names.""" + +from dataclasses import dataclass +import json +from pathlib import Path +import re +import shutil + +import gguf +import numpy as np +import torch +from safetensors import safe_open + +from checkpoint import EXTENSIONS + + +FAMILIES = ( + "qwen3_tts", "qwen3_asr", "voxcpm2", "voxcpm1", "fish_audio", + "cosyvoice3", "f5_tts", "outetts", "yue2", "ace_step", "vibevoice", +) + +MERGE_MODES = {"yue2": "rounded_bf16_delta", "vibevoice": "rankwise_f32"} + +COMPONENTS = { + "qwen3_tts": "model_weights", "voxcpm2": "weights", + "fish_audio": "model_weights", "cosyvoice3": "llm", + "f5_tts": "transformer", "outetts": "model_weights", +} + + +@dataclass(frozen=True) +class Target: + name: str + first_row: int = 0 + rows: int = 0 + + +class Bindings: + def __init__(self, family, shapes, component=None): + self.targets = {} + component = component or COMPONENTS.get(family) + names = list(shapes) + if component: + names = [name for name in names if name.startswith(component + "/")] + if not names: + raise ValueError(f"Base GGUF has no {component!r} component") + for name in names: + logical = name.split("/", 1)[-1] + target = (Target(name),) + self.add(name, target) + self.add(logical, target) + if family == "qwen3_tts" and logical.startswith("talker."): + self.add(logical.removeprefix("talker."), target) + elif family == "qwen3_asr": + if logical.startswith("thinker."): + self.add(logical.removeprefix("thinker."), target) + elif logical.startswith("model.language_model."): + self.add("model." + logical.removeprefix("model.language_model."), target) + elif logical.startswith("model.audio_tower."): + self.add(logical.removeprefix("model."), target) + elif family == "cosyvoice3" and logical.startswith("llm.model."): + self.add(logical.removeprefix("llm.model."), target) + elif family == "vibevoice" and logical.startswith("model.language_model."): + self.add(logical.removeprefix("model.language_model."), target) + elif family == "yue2" and logical.startswith("model."): + self.add(logical.removeprefix("model."), target) + elif family == "ace_step" and logical.startswith("decoder."): + self.add(logical.removeprefix("decoder."), target) + elif family == "fish_audio": + self.add("model." + logical, target) + elif family == "f5_tts": + stripped = logical.removeprefix("ema_model.transformer.") + self.add(stripped, target) + self.add("transformer." + stripped, target) + elif family == "voxcpm1": + self.add_voxcpm1(logical, target) + + if family in {"fish_audio", "yue2"}: + for name in names: + if name.endswith(".q_proj.weight"): + prefix = name.removesuffix(".q_proj.weight") + projections = [prefix + "." + part + ".weight" for part in ("q_proj", "k_proj", "v_proj")] + fused = ".wqkv" if family == "fish_audio" else ".qkv_proj" + elif family == "yue2" and name.endswith(".gate_proj.weight"): + prefix = name.removesuffix(".gate_proj.weight") + projections = [prefix + "." + part + ".weight" for part in ("gate_proj", "up_proj")] + fused = ".gate_up_proj" + else: + continue + offset = 0 + targets = [] + for projection in projections: + rows = shapes[projection][0] + targets.append(Target(projection, offset, rows)) + offset += rows + logical = (prefix + fused).split("/", 1)[-1] + self.add(logical, tuple(targets)) + if family == "fish_audio": + self.add("model." + logical, tuple(targets)) + else: + self.add(logical.removeprefix("model."), tuple(targets)) + + def add(self, name, targets): + name = name.removesuffix(".weight") + for alias in {name, name.replace(".", "_")}: + if alias in self.targets and self.targets[alias] != targets: + raise ValueError(f"Ambiguous adapter binding {alias!r}; select --component") + self.targets[alias] = targets + + def resolve(self, name): + for prefix in ("base_model.model.", "text_encoders.", "lora_unet_", "lora_transformer_", "lycoris_"): + if name.startswith(prefix): + name = name.removeprefix(prefix) + break + name = name.replace(".modules_to_save.default", "") + key = name.removesuffix(".weight") + if key not in self.targets: + raise ValueError(f"Adapter target has no model binding: {name}") + return self.targets[key] + + def add_voxcpm1(self, name, target): + match = re.fullmatch( + r"(?:(residual_lm|locenc|locdit)\.)?blk\.(\d+)\." + r"(attn_q|attn_k|attn_v|attn_output|attn_norm|ffn_norm|ffn_gate|ffn_up|ffn_down)\.weight", name) + if match: + prefix = {None: "base_lm", "residual_lm": "residual_lm", "locenc": "feat_encoder.encoder", + "locdit": "feat_decoder.estimator.decoder"}[match[1]] + projection = { + "attn_q": "self_attn.q_proj", "attn_k": "self_attn.k_proj", "attn_v": "self_attn.v_proj", + "attn_output": "self_attn.o_proj", "attn_norm": "input_layernorm", + "ffn_norm": "post_attention_layernorm", "ffn_gate": "mlp.gate_proj", + "ffn_up": "mlp.up_proj", "ffn_down": "mlp.down_proj", + }[match[3]] + self.add(f"{prefix}.layers.{match[2]}.{projection}", target) + return + replacements = { + "token_embd": "base_lm.embed_tokens", "output_norm": "base_lm.norm", + "residual_lm.output_norm": "residual_lm.norm", "locenc.output_norm": "feat_encoder.encoder.norm", + "locdit.output_norm": "feat_decoder.estimator.decoder.norm", + "proj.enc_to_lm": "enc_to_lm_proj", "proj.lm_to_dit": "lm_to_dit_proj", + "proj.res_to_dit": "res_to_dit_proj", "stop.stop_proj": "stop_proj", "stop.stop_head": "stop_head", + } + for suffix in (".weight", ".bias"): + if not name.endswith(suffix): + continue + module = name.removesuffix(suffix) + if module in replacements: + self.add(replacements[module] + suffix, target) + for internal, external in (("locenc.", "feat_encoder."), ("locdit.", "feat_decoder.estimator."), + ("fsq.", "fsq_layer.")): + if module.startswith(internal): + self.add(external + module.removeprefix(internal) + suffix, target) + + +class GGUFReader(gguf.GGUFReader): + def _get_field_parts(self, offset, raw_type): + # Embedded sidecars are byte arrays; one NumPy object per byte is costly. + if raw_type == gguf.GGUFValueType.ARRAY: + item_type = self._get(offset, np.uint32) + if item_type[0] == gguf.GGUFValueType.UINT8: + count = self._get(offset + 4, np.uint64) + values = self._get(offset + 12, np.uint8, int(count[0])) + return 12 + values.nbytes, [item_type, count, values], [2], [ + gguf.GGUFValueType.ARRAY, gguf.GGUFValueType.UINT8] + return super()._get_field_parts(offset, raw_type) + + +class BaseGGUF: + def __init__(self, path): + self.reader = GGUFReader(path) + fields = self.reader.fields + if "audiocpp.tensor_names" in fields: + names = fields["audiocpp.tensor_names"].contents() + ranks = fields["audiocpp.tensor_ranks"].contents() + dimensions = fields["audiocpp.tensor_shapes"].contents() + self.shapes = {} + offset = 0 + for name, rank in zip(names, ranks, strict=True): + self.shapes[name] = tuple(dimensions[offset:offset + rank]) + offset += rank + if offset != len(dimensions): + raise ValueError("Base GGUF logical shape metadata is inconsistent") + else: + names = [tensor.name for tensor in self.reader.tensors] + self.shapes = {tensor.name: tuple(reversed(tensor.shape.tolist())) for tensor in self.reader.tensors} + self.tensors = dict(zip(names, self.reader.tensors, strict=True)) + + def read_f32(self, name): + tensor = self.tensors[name] + values = gguf.quants.dequantize(tensor.data, tensor.tensor_type) + return torch.from_numpy(np.asarray(values, dtype=np.float32).reshape(self.shapes[name]).copy()) + + def extract_sidecars(self, destination): + fields = self.reader.fields + names = fields["audiocpp.embedded_files.names"].contents() + if "audiocpp.embedded_files.data" in fields: + field = fields["audiocpp.embedded_files.data"] + first = field.parts[field.data[0]] + offset = first.__array_interface__["data"][0] - self.reader.data.__array_interface__["data"][0] + data = self.reader.data[offset:offset + int(field.parts[4][0])] + offsets = fields["audiocpp.embedded_files.offsets"].contents() + contents = (memoryview(data)[offsets[i]:offsets[i + 1]] for i in range(len(names))) + else: + contents = (value.encode("utf-8") for value in fields["audiocpp.embedded_files.contents"].contents()) + for name, content in zip(names, contents, strict=True): + target = destination / name + if not target.resolve().is_relative_to(destination.resolve()): + raise ValueError(f"Unsafe embedded file name: {name}") + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(content) + + +class BaseSources: + def __init__(self, inputs, root): + self.root = root + self.shapes = {} + self.tensors = {} + for value in inputs: + prefix, separator, filename = value.partition("=") + path = Path(filename if separator else prefix) + prefix = prefix if separator else "" + if path.name.endswith(".safetensors.index.json"): + mapping = json.loads(path.read_text())["weight_map"] + shards = {name: path.parent / shard for name, shard in mapping.items()} + elif path.suffix == ".safetensors": + with safe_open(path, framework="pt", device="cpu") as source: + shards = {name: path for name in source.keys()} + else: + raise ValueError(f"Tensor inputs must be safetensors or a shard index: {path}") + for shard in sorted(set(shards.values())): + with safe_open(shard, framework="pt", device="cpu") as source: + for name, file in shards.items(): + if file != shard: + continue + logical = prefix + "/" + name if prefix else name + if logical in self.shapes: + raise ValueError(f"Duplicate input tensor: {logical}") + self.shapes[logical] = tuple(source.get_slice(name).get_shape()) + self.tensors[logical] = (shard, name) + + def read_f32(self, name): + path, key = self.tensors[name] + with safe_open(path, framework="pt", device="cpu") as source: + return source.get_tensor(key).float().clone() + + def extract_sidecars(self, destination): + for path in self.root.rglob("*"): + relative = path.relative_to(self.root) + if any(part in {".git", ".cache"} for part in relative.parts): + continue + if not path.is_file() or path.suffix.lower() in (*EXTENSIONS, ".gguf"): + continue + if path.stat().st_size > 64 * 1024 * 1024: + continue + target = destination / relative + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(path, target) diff --git a/tools/lora/checkpoint.py b/tools/lora/checkpoint.py new file mode 100644 index 000000000..18ef58bf8 --- /dev/null +++ b/tools/lora/checkpoint.py @@ -0,0 +1,36 @@ +"""Checkpoint containers and logical-file discovery, independent of adapter layout.""" + +from pathlib import Path + +import torch +from safetensors.torch import load_file + + +TORCH_EXTENSIONS = (".pt", ".pth", ".bin", ".ckpt") +EXTENSIONS = (".safetensors", *TORCH_EXTENSIONS) + + +def find_checkpoint(directory, stem, required=True): + for extension in EXTENSIONS: + path = directory / (stem + extension) + if path.is_file(): + return path + if required: + raise FileNotFoundError(f"No {stem} checkpoint in {directory}; expected {', '.join(EXTENSIONS)}") + return None + + +def load_checkpoint(path): + path = Path(path) + extension = path.suffix.lower() + if extension == ".safetensors": + state = load_file(path, device="cpu") + elif extension in TORCH_EXTENSIONS: + state = torch.load(path, map_location="cpu", weights_only=True) + if isinstance(state, dict) and "state_dict" in state: + state = state["state_dict"] + else: + raise ValueError(f"Unsupported checkpoint container: {path}") + if not isinstance(state, dict) or not all(isinstance(value, torch.Tensor) for value in state.values()): + raise ValueError(f"Expected a tensor state dictionary: {path}") + return state diff --git a/tools/lora/merge_lora.py b/tools/lora/merge_lora.py new file mode 100644 index 000000000..c619803e2 --- /dev/null +++ b/tools/lora/merge_lora.py @@ -0,0 +1,317 @@ +#!/usr/bin/env python3 +"""Bind a model's LoRA to its base GGUF and merge with audiocpp_gguf.""" + +import argparse +import json +from pathlib import Path +import re +import subprocess +import tempfile + +import torch +from safetensors.torch import save_file + +from bindings import BaseGGUF, BaseSources, Bindings, FAMILIES, MERGE_MODES +from checkpoint import find_checkpoint, load_checkpoint +from prepare_fish_audio_adapter import prepare as prepare_fish + + +SUFFIXES = ( + (".lora_A.default.weight", "a"), (".lora_B.default.weight", "b"), + (".lora_A.weight", "a"), (".lora_B.weight", "b"), + (".lora_A", "a"), (".lora_B", "b"), + (".lora_down.weight", "a"), (".lora_up.weight", "b"), + (".lokr_w1", "lokr_w1"), (".lokr_w2", "lokr_w2"), + (".lokr_w1_a", "lokr_w1_a"), (".lokr_w1_b", "lokr_w1_b"), + (".lokr_w2_a", "lokr_w2_a"), (".lokr_w2_b", "lokr_w2_b"), + (".alpha", "alpha"), +) + + +def read_adapter(path, family, work, alpha): + directory = path if path.is_dir() else path.parent + native_config = None + native_cosy = False + if family == "cosyvoice3": + native_package = ( + find_checkpoint(path, "lora_weights", required=False) is not None + if path.is_dir() else path.stem == "lora_weights") + native_cosy = (directory / "config.json").is_file() and ( + native_package or not (directory / "adapter_config.json").is_file()) + if path.is_dir(): + peft_file = find_checkpoint(path, "adapter_model", required=False) if not native_cosy else None + if native_cosy: + path = find_checkpoint(path, "lora_weights") + elif peft_file is not None: + path = peft_file + else: + stem = "lora_weights" if family in {"voxcpm1", "voxcpm2"} else "adapter_model" + path = find_checkpoint(path, stem) + if family == "fish_audio" and alpha is not None: + prepared = work / "fish.safetensors" + prepare_fish(path, prepared, alpha) + path = prepared + if family in {"voxcpm1", "voxcpm2"}: + native_config = json.loads((directory / "lora_config.json").read_text())["lora_config"] + config = {"r": native_config["r"], "lora_alpha": native_config["alpha"]} + elif native_cosy: + native_config = json.loads((directory / "config.json").read_text())["config"] + config = {"r": native_config["lora_r"], "lora_alpha": native_config["lora_alpha"]} + else: + config_file = directory / "adapter_config.json" + config = json.loads(config_file.read_text()) if config_file.is_file() else {} + state = load_checkpoint(path) + if not state: + raise ValueError(f"Adapter checkpoint is empty: {path}") + peft_suffixes = (".lora_A.weight", ".lora_B.weight", ".lora_A.default.weight", ".lora_B.default.weight") + has_peft = any(name.endswith(peft_suffixes) for name in state) + if has_peft and native_config is None and not (directory / "adapter_config.json").is_file(): + for name in state: + for suffix in peft_suffixes: + if name.endswith(suffix) and name.removesuffix(suffix) + ".alpha" not in state: + raise ValueError("PEFT adapters require adapter_config.json or per-target alpha tensors") + if not has_peft and native_config is None: + config = {} + return state, config, directory, native_config + + +def pattern_value(config, field, name, default): + matches = [value for pattern, value in config.get(field, {}).items() + if re.fullmatch(r"(.*\.)?(" + pattern + ")", name)] + if matches and any(value != matches[0] for value in matches): + raise ValueError(f"Ambiguous {field} for {name}") + return matches[0] if matches else default + + +def normalize(state, config, bindings, output, scale, rs_config, native_scaling=False, require_scaling=True): + groups = {} + peft_targets = set() + replacements = {} + for name, tensor in state.items(): + for suffix, part in SUFFIXES: + if name.endswith(suffix): + target = name.removesuffix(suffix) + group = groups.setdefault(target, {}) + if part in group: + raise ValueError(f"Duplicate adapter factor: {name}") + group[part] = tensor + if suffix in {".lora_A.weight", ".lora_B.weight", ".lora_A.default.weight", ".lora_B.default.weight"}: + peft_targets.add(target) + break + else: + if not name.endswith((".weight", ".bias")): + raise ValueError(f"Unsupported adapter tensor: {name}") + targets = bindings.resolve(name) + if len(targets) != 1 or targets[0].rows: + raise ValueError(f"Cannot slice a full replacement: {name}") + if targets[0].name in replacements: + raise ValueError(f"Multiple replacements target {targets[0].name}") + replacements[targets[0].name] = tensor.contiguous() + for name, factors in groups.items(): + targets = bindings.resolve(name) + group_config = config if name in peft_targets or native_scaling else {} + if "a" in factors or "b" in factors: + if factors.keys() - {"a", "b", "alpha"}: + raise ValueError(f"Mixed adapter algorithms: {name}") + a, b = factors["a"], factors["b"] + if a.ndim != 2 or b.ndim != 2 or a.shape[0] != b.shape[1]: + raise ValueError(f"Invalid A/B shapes: {name}") + rank = a.shape[0] + if require_scaling and not group_config and "alpha" not in factors: + raise ValueError(f"Adapter has no scaling metadata: {name}") + if pattern_value(group_config, "rank_pattern", name.removeprefix("base_model.model."), group_config.get("r", rank)) != rank: + raise ValueError(f"Configured rank differs from factors: {name}") + coefficient = pattern_value(group_config, "alpha_pattern", name.removeprefix("base_model.model."), + group_config.get("lora_alpha", rank)) + coefficient = float(factors.get("alpha", torch.tensor(coefficient)).item()) + next_row = 0 + for target in sorted(targets, key=lambda target: target.first_row): + rows = target.rows or b.shape[0] + if target.first_row != next_row or rows <= 0 or target.first_row + rows > b.shape[0]: + raise ValueError(f"Invalid adapter output row slices: {name}") + next_row += rows + if next_row != b.shape[0]: + raise ValueError(f"Adapter output rows were not fully bound: {name}") + for target in targets: + rows = target.rows or b.shape[0] + prefix = target.name + if any(key == prefix or key.startswith(prefix + ".") for key in output): + raise ValueError(f"Multiple updates target {prefix}") + output[prefix + ".lora_A"] = a.contiguous().clone() + output[prefix + ".lora_B"] = b[target.first_row:target.first_row + rows].contiguous().clone() + effective_alpha = coefficient * scale + if group_config.get("use_rslora", False): + # Keep rsLoRA scaling in the framework importer, not in the factors. + output[prefix + ".lora_A.weight"] = output.pop(prefix + ".lora_A") + output[prefix + ".lora_B.weight"] = output.pop(prefix + ".lora_B") + rs_config["rank_pattern"][re.escape(prefix)] = rank + rs_config["alpha_pattern"][re.escape(prefix)] = effective_alpha + else: + output[prefix + ".alpha"] = torch.tensor(effective_alpha, dtype=torch.float32) + else: + if len(targets) != 1 or targets[0].rows: + raise ValueError(f"Cannot slice LoKr factors: {name}") + prefix = targets[0].name + if any(key == prefix or key.startswith(prefix + ".") for key in output): + raise ValueError(f"Multiple updates target {prefix}") + for part, tensor in factors.items(): + if part == "alpha": + output[prefix + ".alpha"] = tensor.float() * scale + else: + output[prefix + "." + part] = tensor.contiguous() + if scale != 1: + raise ValueError("Non-unit LoKr export strength is not supported yet") + for name, tensor in replacements.items(): + if name in output: + raise ValueError(f"Duplicate replacement: {name}") + output[name] = tensor + + +def sidecars(family, directory, native_config, base, bindings, output, work): + root = None + speaker_file = find_checkpoint(directory, "speaker_embedding", required=False) if family == "qwen3_tts" else None + if speaker_file is not None: + config_file = directory / "tts_config.json" + config = json.loads(config_file.read_text()) + talker = config["talker_config"] + if config["tts_model_type"] != "custom_voice" or len(talker["spk_id"]) != 1: + raise ValueError("Speaker sidecar requires a single custom_voice speaker") + target = bindings.resolve("talker.model.codec_embedding.weight")[0].name + if any(key == target or key.startswith(target + ".") for key in output): + raise ValueError("Speaker sidecar conflicts with codec embedding updates") + shape = base.shapes[target] + if shape != (talker["vocab_size"], talker["hidden_size"]): + raise ValueError("Speaker sidecar configuration differs from the base checkpoint") + speaker_id = next(iter(talker["spk_id"].values())) + if not 0 <= speaker_id < shape[0]: + raise ValueError("Speaker ID is outside the codec vocabulary") + speaker = load_checkpoint(speaker_file)["speaker_embedding"] + if tuple(speaker.shape) != (shape[1],): + raise ValueError("Speaker embedding width differs from the base checkpoint") + embedding = base.read_f32(target) + embedding[speaker_id] = speaker.float() + output[target] = embedding + root = work / "model" + root.mkdir() + base.extract_sidecars(root) + (root / "config.json").write_text(json.dumps(config, ensure_ascii=False), encoding="utf-8") + if family == "cosyvoice3" and native_config is not None: + for module in ("speech_embedding", "llm_decoder"): + state = load_checkpoint(find_checkpoint(directory, module)) + for name, tensor in state.items(): + target = bindings.resolve(module + "." + name)[0].name + if target in output: + raise ValueError(f"Duplicate sidecar replacement: {target}") + output[target] = tensor.contiguous() + if family == "vibevoice": + candidates = ( + ("diffusion_head/model", "model.prediction_head."), + ("diffusion_head_full", "model.prediction_head."), + ("diffusion_head/diffusion_head_full", "model.prediction_head."), + ("acoustic_connector/pytorch_model", "model.acoustic_connector."), + ("semantic_connector/pytorch_model", "model.semantic_connector."), + ) + loaded = set() + for filename, prefix in candidates: + if prefix in loaded: + continue + logical = directory / filename + path = find_checkpoint(logical.parent, logical.name, required=False) + if path is None: + continue + for name, tensor in load_checkpoint(path).items(): + target = bindings.resolve(prefix + name)[0].name + output[target] = tensor.contiguous() + loaded.add(prefix) + return root + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--family", choices=FAMILIES, required=True) + base_options = parser.add_mutually_exclusive_group(required=True) + base_options.add_argument("--base", type=Path, help="Unmerged base GGUF, preferably original dtype or 16-bit") + base_options.add_argument("--input", action="append", help="Original converter-compatible [namespace=]safetensors or shard index; repeat for components") + parser.add_argument("--root", type=Path, help="Original checkpoint directory containing config/tokenizer files") + parser.add_argument("--adapter", type=Path, help="Adapter file or directory; AR adapter for YuE2") + parser.add_argument("--nar-adapter", type=Path, help="Additional YuE2 NAR adapter") + parser.add_argument("--scale", type=float, default=1.0, help="Adapter strength multiplier") + parser.add_argument("--nar-scale", type=float, default=1.0) + parser.add_argument("--alpha", type=float, help="Training alpha for a Fish Lightning checkpoint") + parser.add_argument("--component", help="GGUF component namespace, e.g. dit_turbo_weights for ACE-Step") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--type", default="orig", choices=("orig", "f32", "f16", "bf16", "q8_0", "q4_0", "q2_k", "q3_k", "q4_k", "q5_k", "q6_k")) + parser.add_argument("--keep-type", action="append", default=[], help="Converter tensor-pattern dtype override") + parser.add_argument("--converter", type=Path, default=Path(__file__).resolve().parents[2] / "build/debug/bin/audiocpp_gguf") + parser.add_argument("--model-spec", type=Path, help="Override the repository's current model spec") + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + if args.base and args.base.resolve() == args.output.resolve(): + parser.error("--output must not replace the base GGUF") + for source in args.input or []: + path = Path(source.partition("=")[2] if "=" in source else source) + if path.resolve() == args.output.resolve(): + parser.error("--output must not replace an original weight source") + if args.output.exists() and not args.overwrite: + parser.error("Output exists; pass --overwrite to replace it") + if args.nar_adapter and args.family != "yue2": + parser.error("--nar-adapter is only valid for YuE2") + if args.adapter is None and args.nar_adapter is None: + parser.error("Provide --adapter, or --nar-adapter for a NAR-only YuE2 merge") + if args.alpha is not None and args.family != "fish_audio": + parser.error("--alpha is only used to prepare Fish Speech Lightning checkpoints") + if not torch.isfinite(torch.tensor([args.scale, args.nar_scale])).all(): + parser.error("Adapter strengths must be finite") + if args.base: + base = BaseGGUF(args.base) + inputs = [str(args.base)] + else: + inputs = args.input + first = Path(inputs[0].partition("=")[2] if "=" in inputs[0] else inputs[0]) + base = BaseSources(inputs, args.root or first.parent) + if any(path.resolve() == args.output.resolve() for path, _ in base.tensors.values()): + parser.error("--output must not replace an original weight shard") + bindings = Bindings(args.family, base.shapes, args.component) + with tempfile.TemporaryDirectory(prefix="audiocpp-lora-") as temporary: + work = Path(temporary) + output = {} + rs_config = {"r": 1, "lora_alpha": 1, "use_rslora": True, "rank_pattern": {}, "alpha_pattern": {}} + root = args.root + for adapter, strength in ((args.adapter, args.scale), (args.nar_adapter, args.nar_scale)): + if adapter is None or (strength == 0 and args.family in {"yue2", "ace_step"}): + continue + state, config, directory, native = read_adapter(adapter, args.family, work, args.alpha) + for option in ("use_dora", "use_qalora", "fan_in_fan_out", "lora_bias"): + if config.get(option, False): + raise ValueError(f"Unsupported adapter option: {option}") + normalize(state, config, bindings, output, strength if args.family == "yue2" else 1.0, + rs_config, native_scaling=native is not None, require_scaling=args.family not in {"yue2", "ace_step"}) + model_root = sidecars(args.family, directory, native, base, bindings, output, work) + if model_root is not None: + root = model_root + command = [str(args.converter), "--output", str(args.output), + "--type", args.type, "--family", args.family, "--model-spec", + str(args.model_spec or Path(__file__).resolve().parents[2] / "model_specs" / (args.family + ".json")), "--log"] + for value in inputs: + command.extend(("--input", value)) + if output: + normalized = work / "adapter.safetensors" + save_file(output, normalized) + if rs_config["rank_pattern"]: + (work / "adapter_config.json").write_text(json.dumps(rs_config)) + command.extend(("--weight-adapter", str(normalized), "--adapter-merge-mode", + MERGE_MODES.get(args.family, "delta_f32"))) + if args.family != "yue2": + command.extend(("--adapter-strength", str(args.scale))) + if root is not None: + command.extend(("--root", str(root))) + for pattern in args.keep_type: + command.extend(("--keep-type", pattern)) + if args.overwrite: + command.append("--overwrite") + subprocess.run(command, check=True) + print(f"adapter_tensors={len(output)}") + + +if __name__ == "__main__": + main() diff --git a/tools/lora/prepare_fish_audio_adapter.py b/tools/lora/prepare_fish_audio_adapter.py new file mode 100644 index 000000000..5d7a6db58 --- /dev/null +++ b/tools/lora/prepare_fish_audio_adapter.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +"""Export Fish Speech Lightning LoRA factors as runtime-ready safetensors.""" + +import argparse +import collections +from pathlib import Path +import typing + +import torch +from safetensors.torch import save_file + +from checkpoint import load_checkpoint + + +def prepare(input_path: Path, output_path: Path, alpha: float) -> None: + # Lightning stores scheduler settings as OmegaConf objects, not just tensors. + from omegaconf import ListConfig + from omegaconf.base import ContainerMetadata, Metadata + from omegaconf.nodes import AnyNode + + allowed = [ + dict, list, int, collections.defaultdict, typing.Any, + ListConfig, ContainerMetadata, Metadata, AnyNode, + ] + with torch.serialization.safe_globals(allowed): + state = load_checkpoint(input_path) + if (input_path.parent / "adapter_config.json").is_file() or any(name.endswith(".alpha") for name in state): + # Already-prepared adapters retain their own scaling metadata. + output_path.parent.mkdir(parents=True, exist_ok=True) + save_file(state, output_path, metadata={"format": "pt", "family": "fish_audio"}) + return + tensors = {} + targets = set() + for name, value in state.items(): + if not isinstance(value, torch.Tensor) or not name.startswith("model."): + raise ValueError(f"Unexpected Fish adapter entry: {name}") + name = name.removeprefix("model.") + if not name.endswith((".lora_A", ".lora_B")): + raise ValueError(f"Expected LoRA factor, got {name}") + target, factor = name.rsplit(".", 1) + if value.ndim != 2: + raise ValueError(f"Expected matrix factor: {name}") + if target in {"embeddings", "codebook_embeddings", "fast_embeddings"}: + # loralib.Embedding adds (B @ A).T to its [vocabulary, width] weight. + factor = "lora_B" if factor == "lora_A" else "lora_A" + value = value.T + tensors[f"{target}.{factor}"] = value.contiguous() + targets.add(target) + if not targets: + raise ValueError("Checkpoint contains no LoRA targets") + for target in targets: + a = tensors[f"{target}.lora_A"] + b = tensors[f"{target}.lora_B"] + if a.shape[0] != b.shape[1] or a.shape[0] <= 0: + raise ValueError(f"Inconsistent LoRA rank: {target}") + tensors[f"{target}.alpha"] = torch.tensor(alpha, dtype=torch.float32) + output_path.parent.mkdir(parents=True, exist_ok=True) + save_file(tensors, output_path, metadata={"format": "pt", "family": "fish_audio"}) + print(f"adapter={output_path}") + print(f"targets={len(targets)}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", type=Path, required=True, help="Fish Speech Lightning .ckpt") + parser.add_argument("--output", type=Path, required=True, help="Output adapter .safetensors") + parser.add_argument("--alpha", type=float, required=True, help="Training LoRA alpha, not alpha/r") + args = parser.parse_args() + if args.output.exists(): + parser.error(f"Output already exists: {args.output}") + if not torch.isfinite(torch.tensor(args.alpha)): + parser.error("--alpha must be finite") + prepare(args.input, args.output, args.alpha) + + +if __name__ == "__main__": + main()