diff --git a/CMakeLists.txt b/CMakeLists.txt index b0cedea3..0f069a85 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1055,6 +1055,26 @@ audiocpp_add_model(magpie_tts engine::models::magpie_tts::make_magpie_tts_loader ) +audiocpp_add_model(bark_tts + SOURCES + src/models/bark_tts/assets.cpp + src/models/bark_tts/codec.cpp + src/models/bark_tts/generator.cpp + src/models/bark_tts/session.cpp + src/models/bark_tts/tokenizer.cpp + src/models/bark_tts/transformer.cpp + INCLUDES + engine/models/bark_tts/assets.h + engine/models/bark_tts/codec.h + engine/models/bark_tts/generator.h + engine/models/bark_tts/session.h + engine/models/bark_tts/tokenizer.h + engine/models/bark_tts/transformer.h + engine/models/bark_tts/types.h + LOADERS + engine::models::bark_tts::make_bark_tts_loader +) + audiocpp_add_model(confucius4_tts SOURCES src/models/confucius4_tts/assets.cpp @@ -2279,6 +2299,9 @@ if (ENGINE_BUILD_TESTS OR ENGINE_BUILD_EXTENDED_TESTS OR ENGINE_BUILD_MODEL_TEST add_engine_unittest(hf_tokenizer_vocab_bounds_test tests/unittests/test_hf_tokenizer_vocab_bounds.cpp) add_test(NAME hf_tokenizer_vocab_bounds_test COMMAND hf_tokenizer_vocab_bounds_test) + add_engine_unittest(bark_tokenizer_test tests/unittests/test_bark_tokenizer.cpp) + add_test(NAME bark_tokenizer_test COMMAND bark_tokenizer_test) + add_engine_unittest(safetensors_offsets_test tests/unittests/test_safetensors_offsets.cpp) add_test(NAME safetensors_offsets_test COMMAND safetensors_offsets_test) @@ -2716,6 +2739,17 @@ if (ENGINE_BUILD_TESTS OR ENGINE_BUILD_EXTENDED_TESTS OR ENGINE_BUILD_MODEL_TEST ) target_link_libraries(miocodec_wavlm_parity PRIVATE engine_runtime ggml) + add_executable(bark_codec_parity EXCLUDE_FROM_ALL tests/bark_tts/bark_codec_parity.cpp) + target_link_libraries(bark_codec_parity PRIVATE engine_runtime ggml) + if (ENGINE_ENABLE_OPENMP) + target_link_libraries(bark_codec_parity PRIVATE OpenMP::OpenMP_CXX) + endif() + add_executable(bark_transformer_parity EXCLUDE_FROM_ALL tests/bark_tts/bark_transformer_parity.cpp) + target_link_libraries(bark_transformer_parity PRIVATE engine_runtime ggml) + if (ENGINE_ENABLE_OPENMP) + target_link_libraries(bark_transformer_parity PRIVATE OpenMP::OpenMP_CXX) + endif() + add_engine_unittest(dots_tts_vocoder_parity tests/dots_tts/dots_tts_vocoder_parity.cpp) # Needs the GGUF and a PyTorch reference dump, so it is driven by hand diff --git a/README.md b/README.md index e6b80999..b2abd1b6 100644 --- a/README.md +++ b/README.md @@ -78,6 +78,7 @@ Runtime tags summarize the supported loading paths. GGUF package precision varie | **confucius4_tts** | Clone | zh, en, ja, ko, de, fr, es, id, it, th, pt, ru, ms, vi | Confucius4-TTS multilingual voice cloning | GGUF F32, Stream | | **cosyvoice3** | TTS, Clone | zh, en, ja, ko, de, es, fr, it, ru, yue | Fun-CosyVoice3 zero-shot, cross-lingual, and instruction-conditioned TTS | GGUF F32/Q8 | | **dots_tts** | TTS, Clone, Edit, Ctrl | multilingual | DotTTS SOAR, MeanFlow, and Edit | GGUF 16/Q8, Stream | +| **bark_tts** | TTS | 13 languages | Bark Small, 261 speaker presets | GGUF 16/Q8 | | **dramabox** | TTS, Clone | en | DramaBox expressive TTS and voice cloning | GGUF Q8 | | **fish_audio** | TTS, Clone, Ctrl | auto, en, zh | Fish Audio S2 Pro | GGUF 16/Q8 | | **firered_audio** | ASR, TTS, Clone, Design, Ctrl | zh, en | FireRedAudio multimodal speech/audio model with ASR, understanding, cloning, design, and edit paths | GGUF original/Q8 | diff --git a/demos/bark-small-fixed-full-sentence.mp4 b/demos/bark-small-fixed-full-sentence.mp4 new file mode 100644 index 00000000..5102f946 Binary files /dev/null and b/demos/bark-small-fixed-full-sentence.mp4 differ diff --git a/docs/models/bark_tts.md b/docs/models/bark_tts.md new file mode 100644 index 00000000..ca2c2b09 --- /dev/null +++ b/docs/models/bark_tts.md @@ -0,0 +1,28 @@ +# Bark Small + +Bark is a multilingual, expressive text-to-audio model. The audio.cpp package +contains the complete Bark Small semantic, coarse, fine, and EnCodec pipeline, +plus all 261 upstream speaker histories. + +```sh +python3 tools/model_manager_v2.py install bark_small_f16 + +audiocpp_cli --task tts --family bark_tts \ + --model Bark-Small-GGUF/bark-small-f16.gguf \ + --text "Hello from Bark. [laughs] This model can be quite expressive." \ + --option history_prompt=v2/en_speaker_6 --output bark.wav +``` + +Use `--option history_prompt=` to select another packaged history. Presets +follow upstream names such as `v2/de_speaker_3`, `v2/ja_speaker_0`, or +`announcer`. `temperature`, `top_k`, `max_tokens`, and `seed` are available as +request options. Bark supports non-speech cues written in brackets, although +the result remains probabilistic. + +F16 is the recommended package. The hybrid Q8 package keeps both autoregressive +semantic/coarse transformers and the complete EnCodec path in F16, and only +quantizes the non-causal fine transformer's dense matrices. Quantizing Bark's +autoregressive stages causes token errors to compound and is not quality-safe. + +Source model: [suno/bark-small](https://huggingface.co/suno/bark-small), pinned +to revision `1dbd7a128513b8ae4a4e2130fed57b7ac9da5bcd`. Bark is MIT licensed. diff --git a/include/engine/models/bark_tts/assets.h b/include/engine/models/bark_tts/assets.h new file mode 100644 index 00000000..c3905480 --- /dev/null +++ b/include/engine/models/bark_tts/assets.h @@ -0,0 +1,22 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/assets/tensor_source.h" +#include "engine/models/bark_tts/types.h" + +#include +#include +#include + +namespace engine::models::bark_tts { + +struct BarkAssets { + engine::assets::ResourceBundle resources; + std::shared_ptr weights; + BarkConfig config; + std::unordered_map presets; +}; + +std::shared_ptr load_bark_assets(const std::filesystem::path & model_path); + +} // namespace engine::models::bark_tts diff --git a/include/engine/models/bark_tts/codec.h b/include/engine/models/bark_tts/codec.h new file mode 100644 index 00000000..be524112 --- /dev/null +++ b/include/engine/models/bark_tts/codec.h @@ -0,0 +1,26 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" + +#include +#include +#include + +namespace engine::core { class ExecutionContext; } +namespace engine::models::bark_tts { + +struct BarkAssets; + +class BarkCodecDecoder { +public: + BarkCodecDecoder(std::shared_ptr assets, + engine::core::ExecutionContext & execution, + engine::assets::TensorStorageType storage_type); + ~BarkCodecDecoder(); + std::vector decode(const std::vector> & codes) const; +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::bark_tts diff --git a/include/engine/models/bark_tts/generator.h b/include/engine/models/bark_tts/generator.h new file mode 100644 index 00000000..e0e9f7ab --- /dev/null +++ b/include/engine/models/bark_tts/generator.h @@ -0,0 +1,31 @@ +#pragma once + +#include "engine/models/bark_tts/codec.h" +#include "engine/models/bark_tts/tokenizer.h" +#include "engine/models/bark_tts/transformer.h" +#include "engine/models/bark_tts/types.h" + +#include +#include +#include + +namespace engine::core { class ExecutionContext; } +namespace engine::models::bark_tts { + +class BarkGenerator { +public: + BarkGenerator(std::shared_ptr assets, + engine::core::ExecutionContext & execution, + engine::assets::TensorStorageType transformer_storage, + engine::assets::TensorStorageType codec_storage); + std::vector synthesize(const std::string & text, const BarkGenerationOptions & options) const; +private: + std::shared_ptr assets_; + BarkTokenizer tokenizer_; + BarkTransformer semantic_; + BarkTransformer coarse_; + BarkTransformer fine_; + BarkCodecDecoder codec_; +}; + +} // namespace engine::models::bark_tts diff --git a/include/engine/models/bark_tts/session.h b/include/engine/models/bark_tts/session.h new file mode 100644 index 00000000..0c8a10d3 --- /dev/null +++ b/include/engine/models/bark_tts/session.h @@ -0,0 +1,34 @@ +#pragma once + +#include "engine/framework/model_spec/metadata.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/models/bark_tts/assets.h" + +#include + +namespace engine::core { class ExecutionContext; } +namespace engine::models::bark_tts { +class BarkGenerator; +std::shared_ptr make_bark_tts_loader(); + +class BarkSession final : public engine::runtime::RuntimeSessionBase, + public engine::runtime::IOfflineVoiceTaskSession { +public: + BarkSession(engine::runtime::TaskSpec task, engine::runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract); + ~BarkSession() override; + std::string family() const override; + engine::runtime::VoiceTaskKind task_kind() const override; + engine::runtime::RunMode run_mode() const override; + void prepare(const engine::runtime::SessionPreparationRequest & request) override; + engine::runtime::TaskResult run(const engine::runtime::TaskRequest & request) override; +private: + engine::runtime::TaskSpec task_; + engine::runtime::SessionOptions options_; + std::shared_ptr assets_; + std::shared_ptr contract_; + std::unique_ptr execution_; + std::unique_ptr generator_; +}; +} // namespace engine::models::bark_tts diff --git a/include/engine/models/bark_tts/tokenizer.h b/include/engine/models/bark_tts/tokenizer.h new file mode 100644 index 00000000..0b7efd10 --- /dev/null +++ b/include/engine/models/bark_tts/tokenizer.h @@ -0,0 +1,20 @@ +#pragma once + +#include +#include +#include +#include + +namespace engine::models::bark_tts { + +class BarkTokenizer { +public: + explicit BarkTokenizer(const std::string & tokenizer_json); + std::vector encode(const std::string & text) const; + +private: + std::unordered_map vocab_; + int32_t unknown_ = 100; +}; + +} // namespace engine::models::bark_tts diff --git a/include/engine/models/bark_tts/transformer.h b/include/engine/models/bark_tts/transformer.h new file mode 100644 index 00000000..19cece7a --- /dev/null +++ b/include/engine/models/bark_tts/transformer.h @@ -0,0 +1,40 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/models/bark_tts/assets.h" + +#include +#include +#include +#include +#include + +namespace engine::core { class ExecutionContext; } + +namespace engine::models::bark_tts { + +class BarkTransformer { +public: + BarkTransformer(std::shared_ptr assets, + engine::core::ExecutionContext & execution, + std::string prefix, + BarkTransformerConfig config, + bool causal, + engine::assets::TensorStorageType storage_type); + ~BarkTransformer(); + + // For causal models, embedding_channels may contain one or more token rows + // whose embeddings are summed. Returns logits for the final position. + std::vector causal_logits(const std::vector> & embedding_channels) const; + + // Fine Bark sums codebook embeddings 0..codebook_idx and returns one logit + // row per sequence position. + std::vector fine_logits(const std::vector> & codes, + int codebook_idx) const; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::bark_tts diff --git a/include/engine/models/bark_tts/types.h b/include/engine/models/bark_tts/types.h new file mode 100644 index 00000000..a83d65a9 --- /dev/null +++ b/include/engine/models/bark_tts/types.h @@ -0,0 +1,43 @@ +#pragma once + +#include +#include +#include +#include + +namespace engine::models::bark_tts { + +struct BarkTransformerConfig { + int64_t hidden = 768; + int64_t layers = 12; + int64_t heads = 12; + int64_t block_size = 1024; + int64_t input_vocab = 0; + int64_t output_vocab = 0; + bool bias = false; +}; + +struct BarkConfig { + BarkTransformerConfig semantic; + BarkTransformerConfig coarse; + BarkTransformerConfig fine; + int64_t codebook_size = 1024; + int64_t codebook_dim = 128; + int64_t sample_rate = 24000; +}; + +struct BarkSpeakerPreset { + std::vector semantic; + std::vector> coarse; + std::vector> fine; +}; + +struct BarkGenerationOptions { + std::string voice_id = "v2/en_speaker_6"; + float temperature = 0.7F; + int top_k = 50; + int64_t max_tokens = 768; + uint64_t seed = 0; +}; + +} // namespace engine::models::bark_tts diff --git a/model_specs/bark_tts.json b/model_specs/bark_tts.json new file mode 100644 index 00000000..8c3c5d84 --- /dev/null +++ b/model_specs/bark_tts.json @@ -0,0 +1,78 @@ +{ + "schema_version": 1, + "family": "bark_tts", + "display_name": "Bark Small", + "description": "Suno Bark is a multilingual text-to-audio model with semantic, coarse-acoustic, and fine-acoustic generation followed by a 24 kHz EnCodec decoder. This package includes Bark Small and its downloaded built-in speaker presets.", + "category": "tts", + "status": "supported", + "tasks": ["tts"], + "modes": ["offline"], + "languages": ["de", "en", "es", "fr", "hi", "it", "ja", "ko", "pl", "pt", "ru", "tr", "zh"], + "runtime": {"tags": ["gguf"]}, + "capabilities": {"tts": ["built_in_voices"]}, + "options": { + "request": [ + {"name": "history_prompt", "type": "string", "description": "Packaged Bark speaker history, for example v2/en_speaker_6.", "required": false, "default": "v2/en_speaker_6"}, + {"name": "temperature", "type": "float", "description": "Semantic and coarse token sampling temperature.", "required": false, "min": 0.01, "default": 0.7}, + {"name": "top_k", "type": "int", "description": "Semantic and coarse token top-k sampling limit.", "required": false, "min": 1, "default": 50}, + {"name": "max_tokens", "type": "int", "description": "Maximum generated semantic token count.", "required": false, "min": 1, "default": 768}, + {"name": "seed", "type": "int", "description": "Sampling seed.", "required": false, "min": 0, "default": 0} + ], + "session": [ + {"name": "weight_type", "type": "enum", "description": "Transformer matmul storage type.", "preset": "weight_type_full", "required": false, "default": "native"}, + {"name": "codec_weight_type", "type": "enum", "description": "EnCodec decoder storage type; native is recommended for audio quality.", "preset": "weight_type_full", "required": false, "default": "native"}, + {"name": "graph_arena_mb", "type": "int", "description": "Reusable graph arena size in MiB.", "required": false, "min": 64, "default": 768} + ], + "load": [] + }, + "package_defaults": { + "download": {"kind": "huggingface_snapshot", "repo": "audio-cpp/audio.cpp-gguf", "revision": "main", "gated": false} + }, + "packages": [ + { + "id": "bark_small_f16", + "display_name": "Bark Small F16 GGUF", + "default": true, + "format": "gguf", + "precision": "f16", + "target_directory": "Bark-Small-GGUF", + "files": ["Bark-Small-GGUF/bark-small-f16.gguf"], + "strip_prefix": "Bark-Small-GGUF" + }, + { + "id": "bark_small_q8_0", + "display_name": "Bark Small Q8_0 GGUF", + "default": false, + "format": "gguf", + "precision": "q8_0", + "target_directory": "Bark-Small-GGUF", + "files": ["Bark-Small-GGUF/bark-small-q8_0.gguf"], + "strip_prefix": "Bark-Small-GGUF" + } + ], + "dependencies": [], + "ui": { + "recommended_package": "bark_small_f16", + "tags": ["TTS", "GGUF"], + "docs": ["docs/models/bark_tts.md", "docs/tts.md", "docs/gguf.md"] + }, + "sources": [ + { + "format": "gguf", + "roots": {"model": ".", "weights": "$gguf"}, + "files": { + "config_json": "model:config.json", + "generation_config_json": "model:generation_config.json", + "tokenizer_json": "model:tokenizer.json", + "speaker_presets_json": "model:speaker_presets.json" + }, + "tensors": {"weights": {"source": "weights:", "prefix": "bark"}} + } + ], + "provenance": { + "model": "suno/bark-small", + "model_revision": "1dbd7a128513b8ae4a4e2130fed57b7ac9da5bcd", + "license": "MIT", + "upstream": "https://github.com/suno-ai/bark" + } +} diff --git a/src/models/bark_tts/assets.cpp b/src/models/bark_tts/assets.cpp new file mode 100644 index 00000000..79bf8213 --- /dev/null +++ b/src/models/bark_tts/assets.cpp @@ -0,0 +1,71 @@ +#include "engine/models/bark_tts/assets.h" + +#include "engine/framework/io/json.h" +#include "engine/framework/model_spec/package.h" + +#include + +namespace engine::models::bark_tts { +namespace { + +BarkTransformerConfig transformer_config(const engine::io::json::Value & value) { + BarkTransformerConfig out; + out.hidden = engine::io::json::require_i64(value, "hidden_size"); + out.layers = engine::io::json::require_i64(value, "num_layers"); + out.heads = engine::io::json::require_i64(value, "num_heads"); + out.block_size = engine::io::json::require_i64(value, "block_size"); + out.input_vocab = engine::io::json::require_i64(value, "input_vocab_size"); + out.output_vocab = engine::io::json::require_i64(value, "output_vocab_size"); + out.bias = engine::io::json::optional_bool(value, "bias", false); + return out; +} + +std::vector> matrix_i32(const engine::io::json::Value & value) { + std::vector> out; + for (const auto & row : value.as_array()) out.push_back(engine::io::json::number_array_as(row)); + return out; +} + +void validate(const BarkAssets & assets) { + if (assets.config.semantic.hidden != 768 || assets.config.coarse.hidden != 768 || + assets.config.fine.hidden != 768 || assets.config.semantic.heads != 12 || + assets.config.sample_rate != 24000 || assets.config.codebook_size != 1024) { + throw std::runtime_error("unsupported Bark architecture; this runtime currently targets suno/bark-small"); + } + engine::assets::require_tensor_shape(*assets.weights, "semantic.input_embeds_layer.weight", + {assets.config.semantic.input_vocab, assets.config.semantic.hidden}); + engine::assets::require_tensor_shape(*assets.weights, "coarse_acoustics.input_embeds_layer.weight", + {assets.config.coarse.input_vocab, assets.config.coarse.hidden}); + engine::assets::require_tensor_shape(*assets.weights, "fine_acoustics.input_embeds_layers.0.weight", + {assets.config.fine.input_vocab, assets.config.fine.hidden}); + engine::assets::require_tensor_shape(*assets.weights, "codec_model.quantizer.layers.0.codebook.embed", + {assets.config.codebook_size, assets.config.codebook_dim}); +} + +} // namespace + +std::shared_ptr load_bark_assets(const std::filesystem::path & model_path) { + auto assets = std::make_shared(); + assets->resources = engine::model_spec::load_resource_bundle_for_family(model_path, "bark_tts"); + const auto config = assets->resources.parse_json("config_json"); + assets->config.semantic = transformer_config(config.require("semantic_config")); + assets->config.coarse = transformer_config(config.require("coarse_acoustics_config")); + assets->config.fine = transformer_config(config.require("fine_acoustics_config")); + const auto & codec = config.require("codec_config"); + assets->config.codebook_size = engine::io::json::require_i64(codec, "codebook_size"); + assets->config.codebook_dim = engine::io::json::require_i64(codec, "codebook_dim"); + assets->config.sample_rate = engine::io::json::require_i64(codec, "sampling_rate"); + assets->weights = assets->resources.open_tensor_source("weights"); + const auto preset_root = assets->resources.parse_json("speaker_presets_json").require("presets"); + for (const auto & [name, value] : preset_root.as_object()) { + BarkSpeakerPreset preset; + preset.semantic = engine::io::json::number_array_as(value.require("semantic")); + preset.coarse = matrix_i32(value.require("coarse")); + preset.fine = matrix_i32(value.require("fine")); + assets->presets.emplace(name, std::move(preset)); + } + validate(*assets); + return assets; +} + +} // namespace engine::models::bark_tts diff --git a/src/models/bark_tts/codec.cpp b/src/models/bark_tts/codec.cpp new file mode 100644 index 00000000..fd3c797c --- /dev/null +++ b/src/models/bark_tts/codec.cpp @@ -0,0 +1,186 @@ +#include "engine/models/bark_tts/codec.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/recurrent_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" +#include "engine/models/bark_tts/assets.h" + +#include +#include + +#include +#include + +namespace engine::models::bark_tts { +namespace { + +struct Conv { engine::modules::Conv1dWeights weights; int64_t in, out, kernel, dilation; }; +struct Up { engine::modules::ConvTranspose1dWeights weights; int64_t in, out, kernel, stride; }; +struct Residual { Conv first; Conv second; Conv shortcut; }; +struct CodecWeights { + std::shared_ptr store; + std::vector codebooks; + Conv initial; + engine::modules::LSTMStackWeights lstm; + std::array ups; + std::array residuals; + Conv final; +}; +struct ContextDelete { void operator()(ggml_context * p) const noexcept { if (p) ggml_free(p); } }; + +Conv load_conv(engine::core::BackendWeightStore & store, const engine::assets::TensorSource & source, + const std::string & prefix, int64_t in, int64_t out, int64_t kernel, int64_t dilation, + engine::assets::TensorStorageType storage) { + Conv result{{}, in, out, kernel, dilation}; + result.weights = engine::modules::binding::conv1d_from_source(store, source, prefix, storage, + out, in, kernel, true); + return result; +} + +Up load_up(engine::core::BackendWeightStore & store, const engine::assets::TensorSource & source, + const std::string & prefix, int64_t in, int64_t out, int64_t kernel, int64_t stride, + engine::assets::TensorStorageType storage) { + Up result{{}, in, out, kernel, stride}; + result.weights = engine::modules::binding::conv_transpose1d_from_source(store, source, prefix, storage, + in, out, kernel, true); + return result; +} + +std::shared_ptr load(const BarkAssets & assets, ggml_backend_t backend, engine::core::BackendType backend_type, + engine::assets::TensorStorageType storage) { + auto out = std::make_shared(); + out->store = std::make_shared(backend, backend_type, + "bark.codec.weights", 384ULL * 1024ULL * 1024ULL); + const auto & source = *assets.weights; + for (int i = 0; i < 8; ++i) out->codebooks.push_back(out->store->load_tensor(source, + "codec_model.quantizer.layers." + std::to_string(i) + ".codebook.embed", storage, {1024, 128})); + out->initial = load_conv(*out->store, source, "codec_model.decoder.layers.0.conv", 128, 512, 7, 1, storage); + for (int i = 0; i < 2; ++i) { + const std::string p = "codec_model.decoder.layers.1.lstm."; + engine::modules::LSTMCellWeights cell; + cell.weight_ih = out->store->load_tensor(source, p + "weight_ih_l" + std::to_string(i), storage, {2048, 512}); + cell.weight_hh = out->store->load_tensor(source, p + "weight_hh_l" + std::to_string(i), storage, {2048, 512}); + cell.bias_ih = out->store->load_tensor(source, p + "bias_ih_l" + std::to_string(i), engine::assets::TensorStorageType::F32, {2048}); + cell.bias_hh = out->store->load_tensor(source, p + "bias_hh_l" + std::to_string(i), engine::assets::TensorStorageType::F32, {2048}); + out->lstm.layers.push_back(std::move(cell)); + } + const int layer[] = {3, 6, 9, 12}; + const int in[] = {512, 256, 128, 64}; + const int channel[] = {256, 128, 64, 32}; + const int ratio[] = {8, 5, 4, 2}; + const int residual_layer[] = {4, 7, 10, 13}; + for (int i = 0; i < 4; ++i) { + const auto base = "codec_model.decoder.layers." + std::to_string(layer[i]) + ".conv"; + out->ups[static_cast(i)] = load_up(*out->store, source, base, in[i], channel[i], ratio[i] * 2, ratio[i], storage); + const auto r = "codec_model.decoder.layers." + std::to_string(residual_layer[i]); + out->residuals[static_cast(i)] = { + load_conv(*out->store, source, r + ".block.1.conv", channel[i], channel[i] / 2, 3, 1, storage), + load_conv(*out->store, source, r + ".block.3.conv", channel[i] / 2, channel[i], 1, 1, storage), + load_conv(*out->store, source, r + ".shortcut.conv", channel[i], channel[i], 1, 1, storage), + }; + } + out->final = load_conv(*out->store, source, "codec_model.decoder.layers.15.conv", 32, 1, 7, 1, storage); + out->store->upload(); + return out; +} + +engine::core::TensorValue causal_conv(engine::core::ModuleBuildContext & ctx, + const engine::core::TensorValue & input, const Conv & conv) { + const int64_t effective = (conv.kernel - 1) * conv.dilation + 1; + const auto contiguous = engine::core::ensure_backend_addressable_layout(ctx, input); + auto padded = engine::modules::ReflectPad1dModule({effective - 1, 0}).build(ctx, contiguous); + return engine::modules::Conv1dModule({conv.in, conv.out, conv.kernel, 1, 0, + static_cast(conv.dilation), true}).build(ctx, padded, conv.weights); +} + +engine::core::TensorValue upsample(engine::core::ModuleBuildContext & ctx, + const engine::core::TensorValue & input, const Up & up) { + auto full = engine::modules::ConvTranspose1dModule({up.in, up.out, up.kernel, + static_cast(up.stride), 0, 1, true}).build(ctx, input, up.weights); + const int64_t frames = full.shape.dims[2] - (up.kernel - up.stride); + return engine::modules::SliceModule({2, 0, frames}).build(ctx, full); +} + +engine::core::TensorValue elu(engine::core::ModuleBuildContext & ctx, const engine::core::TensorValue & input) { + return engine::modules::EluModule{}.build(ctx, input); +} + +engine::core::TensorValue residual(engine::core::ModuleBuildContext & ctx, + const engine::core::TensorValue & input, const Residual & weights) { + auto hidden = causal_conv(ctx, elu(ctx, input), weights.first); + hidden = causal_conv(ctx, elu(ctx, hidden), weights.second); + return engine::modules::AddModule{}.build(ctx, causal_conv(ctx, input, weights.shortcut), hidden); +} + +} // namespace + +struct BarkCodecDecoder::Impl { + std::shared_ptr assets; + engine::core::ExecutionContext & execution; + std::shared_ptr weights; +}; + +BarkCodecDecoder::BarkCodecDecoder(std::shared_ptr assets, + engine::core::ExecutionContext & execution, + engine::assets::TensorStorageType storage) + : impl_(std::make_unique(Impl{assets, execution, load(*assets, execution.backend(), execution.backend_type(), storage)})) {} +BarkCodecDecoder::~BarkCodecDecoder() = default; + +std::vector BarkCodecDecoder::decode(const std::vector> & codes) const { + if (codes.size() != 8 || codes.front().empty()) throw std::runtime_error("Bark codec requires 8 non-empty codebooks"); + const int64_t frames = static_cast(codes.front().size()); + for (const auto & row : codes) if (static_cast(row.size()) != frames) + throw std::runtime_error("Bark codec codebook lengths differ"); + std::unique_ptr context(ggml_init({256ULL * 1024ULL * 1024ULL, nullptr, true})); + engine::core::ModuleBuildContext build{context.get(), "bark.codec", impl_->execution.backend_type()}; + engine::core::TensorValue quantized; + std::vector inputs; + for (size_t i = 0; i < 8; ++i) { + auto * raw = ggml_new_tensor_1d(context.get(), GGML_TYPE_I32, frames); + inputs.push_back(raw); + auto ids = engine::core::wrap_tensor(raw, engine::core::TensorShape::from_dims({frames}), GGML_TYPE_I32); + auto value = engine::modules::EmbeddingModule({1024, 128}).build(build, ids, impl_->weights->codebooks[i]); + value = engine::core::reshape_tensor(build, value, engine::core::TensorShape::from_dims({1, frames, 128})); + value = engine::modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build, value); + quantized = quantized.valid() ? engine::modules::AddModule{}.build(build, quantized, value) : value; + } + auto hidden = causal_conv(build, quantized, impl_->weights->initial); + hidden = engine::modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build, hidden); + hidden = engine::core::ensure_backend_addressable_layout(build, hidden); + hidden = engine::core::reshape_tensor(build, hidden, engine::core::TensorShape::from_dims({frames, 512})); + const auto lstm_residual = hidden; + for (const auto & cell : impl_->weights->lstm.layers) { + auto zero = engine::core::wrap_tensor(ggml_scale(context.get(), engine::modules::SliceModule({0, 0, 1}).build(build, hidden).tensor, 0.0F), + engine::core::TensorShape::from_dims({1, 512}), GGML_TYPE_F32); + hidden = engine::modules::LSTMSequenceModule({512, 512, false}).build(build, hidden, zero, zero, {cell}).sequence; + } + hidden = engine::modules::AddModule{}.build(build, hidden, lstm_residual); + hidden = engine::core::reshape_tensor(build, hidden, engine::core::TensorShape::from_dims({1, frames, 512})); + hidden = engine::modules::TransposeModule({{0, 2, 1, 3}, 3}).build(build, hidden); + for (size_t i = 0; i < 4; ++i) hidden = residual(build, upsample(build, elu(build, hidden), impl_->weights->ups[i]), impl_->weights->residuals[i]); + auto waveform = causal_conv(build, elu(build, hidden), impl_->weights->final); + ggml_set_output(waveform.tensor); + auto * graph = ggml_new_graph_custom(context.get(), 65536, false); + ggml_build_forward_expand(graph, waveform.tensor); + auto * buffer = ggml_backend_alloc_ctx_tensors(context.get(), impl_->execution.backend()); + if (!buffer) throw std::runtime_error("failed to allocate Bark codec graph"); + for (size_t i = 0; i < 8; ++i) ggml_backend_tensor_set(inputs[i], codes[i].data(), 0, codes[i].size() * sizeof(int32_t)); + engine::core::set_backend_threads(impl_->execution.backend(), impl_->execution.config().threads); + const auto status = engine::core::compute_backend_graph(impl_->execution.backend(), graph); + ggml_backend_synchronize(impl_->execution.backend()); + if (status != GGML_STATUS_SUCCESS) throw std::runtime_error("Bark codec graph compute failed"); + std::vector out(static_cast(waveform.shape.dims[2])); + ggml_backend_tensor_get(waveform.tensor, out.data(), 0, out.size() * sizeof(float)); + engine::core::release_backend_graph_resources(impl_->execution.backend(), graph); + ggml_backend_buffer_free(buffer); + return out; +} + +} // namespace engine::models::bark_tts diff --git a/src/models/bark_tts/generator.cpp b/src/models/bark_tts/generator.cpp new file mode 100644 index 00000000..2a43c057 --- /dev/null +++ b/src/models/bark_tts/generator.cpp @@ -0,0 +1,162 @@ +#include "engine/models/bark_tts/generator.h" + +#include +#include +#include +#include +#include +#include + +namespace engine::models::bark_tts { +namespace { + +int32_t sample(std::vector logits, int begin, int end, int top_k, float temperature, std::mt19937_64 & rng) { + if (begin < 0 || end > static_cast(logits.size()) || begin >= end) throw std::runtime_error("invalid Bark sampling range"); + std::vector indices(static_cast(end - begin)); + std::iota(indices.begin(), indices.end(), begin); + const int keep = std::min(top_k, static_cast(indices.size())); + std::partial_sort(indices.begin(), indices.begin() + keep, indices.end(), [&](int32_t a, int32_t b) { return logits[a] > logits[b]; }); + indices.resize(static_cast(keep)); + const float maximum = logits[indices.front()] / temperature; + std::vector weights; + weights.reserve(indices.size()); + for (int32_t id : indices) weights.push_back(std::exp(static_cast(logits[id] / temperature - maximum))); + return indices[std::discrete_distribution(weights.begin(), weights.end())(rng)]; +} + +std::vector semantic_tokens(BarkTransformer & model, const BarkTokenizer & tokenizer, + const BarkSpeakerPreset & preset, const std::string & text, + const BarkGenerationOptions & options, std::mt19937_64 & rng) { + auto text_ids = tokenizer.encode(text); + if (text_ids.size() > 256) text_ids.resize(256); + for (auto & id : text_ids) id += 10048; + text_ids.resize(256, 129595); + std::vector history; + const size_t history_start = preset.semantic.size() > 256 ? preset.semantic.size() - 256 : 0; + history.assign(preset.semantic.begin() + static_cast(history_start), preset.semantic.end()); + history.resize(256, 10000); + // Bark's merge_context path sums the first 256 text embeddings with the + // 256 semantic-history embeddings, then appends SEMANTIC_INFER without a + // history embedding at that position. Two channels reproduce that + // merged 257-position sequence; negative ids are masked by the runtime. + text_ids.push_back(129599); + history.push_back(-1); + std::vector generated; + for (int64_t step = 0; step < options.max_tokens && text_ids.size() < 1024; ++step) { + auto logits = model.causal_logits({text_ids, history}); + const int32_t id = sample(std::move(logits), 0, 10001, options.top_k, options.temperature, rng); + if (id == 10000) break; + generated.push_back(id); + text_ids.push_back(id); + history.push_back(-1); + } + if (generated.empty()) throw std::runtime_error("Bark semantic model generated no speech tokens"); + return generated; +} + +std::vector coarse_tokens(BarkTransformer & model, const BarkSpeakerPreset & preset, + const std::vector & semantic, const BarkGenerationOptions & options, + std::mt19937_64 & rng) { + constexpr double ratio = 75.0 / 49.9 * 2.0; + std::vector sem_history = preset.semantic; + std::vector coarse_history; + if (preset.coarse.size() != 2) throw std::runtime_error("Bark preset coarse history must have two codebooks"); + for (size_t frame = 0; frame < preset.coarse[0].size(); ++frame) + for (int book = 0; book < 2; ++book) + coarse_history.push_back(preset.coarse[book][frame] + book * 1024 + 10000); + const int max_sem_history = static_cast(std::floor(630.0 / ratio)); + int sem_count = std::min({max_sem_history, static_cast(sem_history.size() / 2 * 2), + static_cast(std::floor(coarse_history.size() / ratio))}); + int coarse_count = static_cast(std::round(sem_count * ratio)); + if (sem_count > 0) sem_history.erase(sem_history.begin(), sem_history.end() - sem_count); + if (coarse_count > 0) coarse_history.erase(coarse_history.begin(), coarse_history.end() - coarse_count); + if (coarse_history.size() >= 2) coarse_history.resize(coarse_history.size() - 2); + const size_t original_history = coarse_history.size(); + std::vector all_semantic = sem_history; + all_semantic.insert(all_semantic.end(), semantic.begin(), semantic.end()); + const int generated_length = static_cast(std::round(std::floor(semantic.size() * ratio / 2.0) * 2.0)); + for (int total = 0; total < generated_length;) { + const int semantic_index = sem_count + static_cast(std::round(total / ratio)); + const int begin = std::max(0, semantic_index - max_sem_history); + // This deliberately takes the suffix beginning at semantic_index's + // history boundary and then its first 256 values. It matches Bark's + // reference x_semantic[:, max(0, semantic_idx-max_history):][:, :256]. + std::vector input(all_semantic.begin() + begin, all_semantic.end()); + if (input.size() > 256) input.resize(256); + input.resize(256, 12048); + input.push_back(12050); + const size_t take = std::min(630, coarse_history.size()); + input.insert(input.end(), coarse_history.end() - static_cast(take), coarse_history.end()); + const int window = std::min(60, generated_length - total); + for (int i = 0; i < window; ++i) { + auto logits = model.causal_logits({input}); + const int book = total % 2; + const int32_t id = sample(std::move(logits), 10000 + book * 1024, + 10000 + (book + 1) * 1024, options.top_k, options.temperature, rng); + input.push_back(id); + coarse_history.push_back(id); + ++total; + } + } + return {coarse_history.begin() + static_cast(original_history), coarse_history.end()}; +} + +std::vector> fine_tokens(BarkTransformer & model, const BarkSpeakerPreset & preset, + const std::vector & coarse) { + const int64_t frames = static_cast(coarse.size() / 2); + std::vector> values(8); + for (int64_t frame = 0; frame < frames; ++frame) for (int book = 0; book < 2; ++book) + values[book].push_back((coarse[static_cast(frame * 2 + book)] - 10000) % 1024); + const int64_t history = std::min(512, preset.fine.empty() ? 0 : preset.fine.front().size()); + std::vector> input(8); + for (int book = 0; book < 8; ++book) { + if (history) input[book].insert(input[book].end(), preset.fine[book].end() - history, preset.fine[book].end()); + if (book < 2) input[book].insert(input[book].end(), values[book].begin(), values[book].end()); + else input[book].resize(static_cast(history + frames), 1024); + } + const int64_t padded = std::max(1024, history + frames); + for (auto & row : input) row.resize(static_cast(padded), 1024); + const int loops = std::max(0, static_cast(std::ceil((frames - (1024 - history)) / 512.0))) + 1; + for (int outer = 0; outer < loops; ++outer) { + const int64_t start = std::min(outer * 512, padded - 1024); + const int64_t fill = std::min(history + outer * 512, padded - 512); + const int64_t relative = fill - start; + std::vector> window(8); + for (int book = 0; book < 8; ++book) window[book].assign(input[book].begin() + start, input[book].begin() + start + 1024); + for (int book = 2; book < 8; ++book) { + const auto logits = model.fine_logits(window, book); + for (int64_t position = relative; position < 1024; ++position) { + const auto begin = logits.begin() + position * 1056; + window[book][static_cast(position)] = static_cast(std::distance(begin, + std::max_element(begin, begin + 1024))); + } + } + for (int book = 2; book < 8; ++book) + std::copy(window[book].begin() + relative, window[book].end(), input[book].begin() + fill); + } + for (int book = 0; book < 8; ++book) + values[book].assign(input[book].begin() + history, input[book].begin() + history + frames); + return values; +} + +} // namespace + +BarkGenerator::BarkGenerator(std::shared_ptr assets, engine::core::ExecutionContext & execution, + engine::assets::TensorStorageType transformer_storage, + engine::assets::TensorStorageType codec_storage) + : assets_(std::move(assets)), tokenizer_(assets_->resources.read_text("tokenizer_json")), + semantic_(assets_, execution, "semantic", assets_->config.semantic, true, transformer_storage), + coarse_(assets_, execution, "coarse_acoustics", assets_->config.coarse, true, transformer_storage), + fine_(assets_, execution, "fine_acoustics", assets_->config.fine, false, transformer_storage), + codec_(assets_, execution, codec_storage) {} + +std::vector BarkGenerator::synthesize(const std::string & text, const BarkGenerationOptions & options) const { + const auto found = assets_->presets.find(options.voice_id); + if (found == assets_->presets.end()) throw std::runtime_error("unknown Bark voice_id: " + options.voice_id); + std::mt19937_64 rng(options.seed); + auto semantic = semantic_tokens(const_cast(semantic_), tokenizer_, found->second, text, options, rng); + auto coarse = coarse_tokens(const_cast(coarse_), found->second, semantic, options, rng); + return codec_.decode(fine_tokens(const_cast(fine_), found->second, coarse)); +} + +} // namespace engine::models::bark_tts diff --git a/src/models/bark_tts/session.cpp b/src/models/bark_tts/session.cpp new file mode 100644 index 00000000..87bbd4ff --- /dev/null +++ b/src/models/bark_tts/session.cpp @@ -0,0 +1,70 @@ +#include "engine/models/bark_tts/session.h" + +#include "engine/framework/runtime/options.h" +#include "engine/framework/runtime/spec_backed_model.h" +#include "engine/models/bark_tts/generator.h" + +#include +#include + +namespace engine::models::bark_tts { +namespace { +constexpr const char * kFamily = "bark_tts"; + +engine::assets::TensorStorageType storage(const engine::runtime::SessionOptions & options, const char * name) { + return engine::runtime::parse_tensor_storage_option(options.options, name, + engine::assets::TensorStorageType::Native, {engine::assets::TensorStorageType::Native, + engine::assets::TensorStorageType::F32, engine::assets::TensorStorageType::F16, + engine::assets::TensorStorageType::BF16, engine::assets::TensorStorageType::Q8_0}); +} + +std::unique_ptr create(const engine::runtime::TaskSpec & task, + const engine::runtime::SessionOptions & options, std::shared_ptr assets, + std::shared_ptr contract) { + return std::make_unique(task, options, std::move(assets), std::move(contract)); +} +} // namespace + +BarkSession::BarkSession(engine::runtime::TaskSpec task, engine::runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract) + : RuntimeSessionBase(options), task_(task), options_(std::move(options)), assets_(std::move(assets)), + contract_(std::move(contract)) { + if (!assets_ || !contract_) throw std::runtime_error("Bark session requires assets and model contract"); + engine::runtime::validate_spec_backed_session_options(options_, *contract_, kFamily, "Bark"); + if (task_.task != engine::runtime::VoiceTaskKind::Tts || task_.mode != engine::runtime::RunMode::Offline) + throw std::runtime_error("Bark supports offline TTS"); + execution_ = std::make_unique(options_.backend); + generator_ = std::make_unique(assets_, *execution_, storage(options_, "bark_tts.weight_type"), + storage(options_, "bark_tts.codec_weight_type")); +} +BarkSession::~BarkSession() = default; +std::string BarkSession::family() const { return kFamily; } +engine::runtime::VoiceTaskKind BarkSession::task_kind() const { return task_.task; } +engine::runtime::RunMode BarkSession::run_mode() const { return task_.mode; } +void BarkSession::prepare(const engine::runtime::SessionPreparationRequest & request) { + engine::runtime::validate_spec_backed_request_options(request.options, *contract_, "Bark"); mark_prepared(); +} +engine::runtime::TaskResult BarkSession::run(const engine::runtime::TaskRequest & request) { + require_prepared("Bark run"); + engine::runtime::validate_spec_backed_request_options(request.options, *contract_, "Bark"); + if (!request.text_input.has_value() || request.text_input->text.empty()) throw std::runtime_error("Bark requires text_input"); + BarkGenerationOptions generation; + generation.voice_id = engine::runtime::find_option(request.options, {"history_prompt"}).value_or(generation.voice_id); + generation.temperature = engine::runtime::parse_float_option(request.options, {"temperature"}).value_or(generation.temperature); + generation.top_k = engine::runtime::parse_int_option(request.options, {"top_k"}).value_or(generation.top_k); + generation.max_tokens = engine::runtime::parse_i64_option(request.options, {"max_tokens"}).value_or(generation.max_tokens); + generation.seed = engine::runtime::parse_u64_option(request.options, {"seed"}).value_or(generation.seed); + if (!(generation.temperature > 0.0F) || !std::isfinite(generation.temperature) || generation.top_k <= 0 || generation.max_tokens <= 0) + throw std::runtime_error("invalid Bark generation options"); + engine::runtime::TaskResult result; + result.audio_output = engine::runtime::AudioBuffer{24000, 1, generator_->synthesize(request.text_input->text, generation)}; + return result; +} + +std::shared_ptr make_bark_tts_loader() { + engine::runtime::SpecBackedVoiceModelConfig config; + config.family = kFamily; config.load_assets = load_bark_assets; config.create_session = create; + return engine::runtime::make_spec_backed_voice_loader(std::move(config)); +} +} // namespace engine::models::bark_tts diff --git a/src/models/bark_tts/tokenizer.cpp b/src/models/bark_tts/tokenizer.cpp new file mode 100644 index 00000000..6b0273a5 --- /dev/null +++ b/src/models/bark_tts/tokenizer.cpp @@ -0,0 +1,95 @@ +#include "engine/models/bark_tts/tokenizer.h" + +#include "engine/framework/io/json.h" + +#include +#include + +namespace engine::models::bark_tts { +namespace { + +struct Character { uint32_t codepoint; std::string bytes; }; + +std::vector characters(const std::string & text) { + std::vector out; + for (size_t pos = 0; pos < text.size();) { + const auto first = static_cast(text[pos]); + size_t width = 1; + uint32_t value = first; + if ((first & 0xe0U) == 0xc0U) { width = 2; value = first & 0x1fU; } + else if ((first & 0xf0U) == 0xe0U) { width = 3; value = first & 0x0fU; } + else if ((first & 0xf8U) == 0xf0U) { width = 4; value = first & 0x07U; } + if (pos + width > text.size()) throw std::runtime_error("Bark text contains truncated UTF-8"); + for (size_t i = 1; i < width; ++i) { + const auto next = static_cast(text[pos + i]); + if ((next & 0xc0U) != 0x80U) throw std::runtime_error("Bark text contains invalid UTF-8"); + value = (value << 6U) | (next & 0x3fU); + } + out.push_back({value, text.substr(pos, width)}); + pos += width; + } + return out; +} + +bool is_cjk(uint32_t cp) { + return (cp >= 0x3400 && cp <= 0x4dbf) || (cp >= 0x4e00 && cp <= 0x9fff) || + (cp >= 0xf900 && cp <= 0xfaff) || (cp >= 0x20000 && cp <= 0x2fa1f); +} + +bool is_punctuation(uint32_t cp) { + if (cp < 128) return std::ispunct(static_cast(cp)) != 0; + return (cp >= 0x2000 && cp <= 0x206f) || (cp >= 0x3000 && cp <= 0x303f); +} + +} // namespace + +BarkTokenizer::BarkTokenizer(const std::string & tokenizer_json) { + const auto root = engine::io::json::parse(tokenizer_json); + const auto & values = root.require("model").require("vocab").as_object(); + vocab_.reserve(values.size()); + for (const auto & [token, id] : values) vocab_.emplace(token, static_cast(id.as_i64())); + if (const auto it = vocab_.find("[UNK]"); it != vocab_.end()) unknown_ = it->second; +} + +std::vector BarkTokenizer::encode(const std::string & text) const { + std::vector words; + std::string current; + for (const auto & ch : characters(text)) { + if (ch.codepoint <= 0x20 || ch.codepoint == 0x7f) { + if (!current.empty()) { words.push_back(std::move(current)); current.clear(); } + } else if (is_cjk(ch.codepoint) || is_punctuation(ch.codepoint)) { + if (!current.empty()) { words.push_back(std::move(current)); current.clear(); } + words.push_back(ch.bytes); + } else { + current += ch.bytes; + } + } + if (!current.empty()) words.push_back(std::move(current)); + + // Bark calls BertTokenizer.encode(..., add_special_tokens=False). + // [CLS]/[SEP] are not part of the semantic text conditioning sequence. + std::vector ids; + for (const auto & word : words) { + const auto chars = characters(word); + if (chars.size() > 100) { ids.push_back(unknown_); continue; } + size_t start = 0; + std::vector pieces; + while (start < chars.size()) { + size_t end = chars.size(); + bool found = false; + while (end > start) { + std::string piece = start == 0 ? "" : "##"; + for (size_t i = start; i < end; ++i) piece += chars[i].bytes; + if (const auto it = vocab_.find(piece); it != vocab_.end()) { + pieces.push_back(it->second); start = end; found = true; break; + } + --end; + } + if (!found) { pieces.assign(1, unknown_); break; } + } + ids.insert(ids.end(), pieces.begin(), pieces.end()); + } + return ids; +} + +} // namespace engine::models::bark_tts diff --git a/src/models/bark_tts/transformer.cpp b/src/models/bark_tts/transformer.cpp new file mode 100644 index 00000000..df52e5fd --- /dev/null +++ b/src/models/bark_tts/transformer.cpp @@ -0,0 +1,234 @@ +#include "engine/models/bark_tts/transformer.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/constant_tensor_cache.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/modules/attention/feed_forward.h" +#include "engine/framework/modules/attention/self_attention.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" + +#include +#include + +#include + +namespace engine::models::bark_tts { +namespace { + +namespace binding = engine::modules::binding; + +struct LayerWeights { + engine::modules::NormWeights norm1; + engine::modules::AttentionWeights attention; + engine::modules::NormWeights norm2; + engine::modules::FeedForwardWeights mlp; +}; + +struct Weights { + std::shared_ptr store; + std::vector embeddings; + engine::core::TensorValue positions; + std::vector layers; + engine::modules::NormWeights final_norm; + std::vector heads; +}; + +struct ContextDelete { void operator()(ggml_context * p) const noexcept { if (p) ggml_free(p); } }; + +engine::modules::NormWeights load_norm(engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, int64_t hidden, bool bias) { + engine::modules::NormWeights out; + out.weight = store.load_tensor(source, prefix + ".weight", engine::assets::TensorStorageType::F32, {hidden}); + if (bias) out.bias = store.load_tensor(source, prefix + ".bias", engine::assets::TensorStorageType::F32, {hidden}); + return out; +} + +std::shared_ptr load_weights(const BarkAssets & assets, + ggml_backend_t backend, + engine::core::BackendType backend_type, + const std::string & prefix, + const BarkTransformerConfig & config, + bool causal, + engine::assets::TensorStorageType storage) { + const auto & source = *assets.weights; + auto out = std::make_shared(); + out->store = std::make_shared(backend, backend_type, + "bark.transformer.weights", 768ULL * 1024ULL * 1024ULL); + const int embeddings = causal ? 1 : 8; + for (int i = 0; i < embeddings; ++i) { + const std::string name = causal ? prefix + ".input_embeds_layer.weight" : + prefix + ".input_embeds_layers." + std::to_string(i) + ".weight"; + out->embeddings.push_back(out->store->load_tensor(source, name, storage, {config.input_vocab, config.hidden})); + } + out->positions = out->store->load_tensor(source, prefix + ".position_embeds_layer.weight", storage, + {config.block_size, config.hidden}); + for (int64_t i = 0; i < config.layers; ++i) { + const std::string p = prefix + ".layers." + std::to_string(i); + LayerWeights layer; + layer.norm1 = load_norm(*out->store, source, p + ".layernorm_1", config.hidden, !causal || config.bias); + layer.norm2 = load_norm(*out->store, source, p + ".layernorm_2", config.hidden, !causal || config.bias); + layer.attention.qkv_weight = out->store->load_tensor(source, p + ".attn.att_proj.weight", storage, + {3 * config.hidden, config.hidden}); + if (config.bias) { + layer.attention.qkv_bias = out->store->load_tensor(source, p + ".attn.att_proj.bias", + engine::assets::TensorStorageType::F32, {3 * config.hidden}); + } + layer.attention.out_weight = out->store->load_tensor(source, p + ".attn.out_proj.weight", storage, + {config.hidden, config.hidden}); + if (config.bias) { + layer.attention.out_bias = out->store->load_tensor(source, p + ".attn.out_proj.bias", + engine::assets::TensorStorageType::F32, {config.hidden}); + } + layer.mlp.fc1_weight = out->store->load_tensor(source, p + ".mlp.in_proj.weight", storage, + {4 * config.hidden, config.hidden}); + layer.mlp.fc2_weight = out->store->load_tensor(source, p + ".mlp.out_proj.weight", storage, + {config.hidden, 4 * config.hidden}); + if (config.bias) { + layer.mlp.fc1_bias = out->store->load_tensor(source, p + ".mlp.in_proj.bias", + engine::assets::TensorStorageType::F32, {4 * config.hidden}); + layer.mlp.fc2_bias = out->store->load_tensor(source, p + ".mlp.out_proj.bias", + engine::assets::TensorStorageType::F32, {config.hidden}); + } + out->layers.push_back(std::move(layer)); + } + out->final_norm = load_norm(*out->store, source, prefix + ".layernorm_final", config.hidden, !causal || config.bias); + const int heads = causal ? 1 : 7; + for (int i = 0; i < heads; ++i) { + const std::string name = causal ? prefix + ".lm_head.weight" : prefix + ".lm_heads." + std::to_string(i) + ".weight"; + out->heads.push_back(out->store->load_tensor(source, name, storage, {config.output_vocab, config.hidden})); + } + out->store->upload(); + return out; +} + +engine::core::TensorValue block(engine::core::ModuleBuildContext & ctx, + const engine::core::TensorValue & input, + const LayerWeights & weights, + const BarkTransformerConfig & config, + bool causal) { + auto n1 = engine::modules::LayerNormModule({config.hidden, 1.0e-5F, true, !causal || config.bias}) + .build(ctx, input, weights.norm1); + auto attention = engine::modules::SelfAttentionModule({ + config.hidden, config.heads, config.bias, GGML_PREC_F32, GGML_PREC_F32, + engine::modules::AttentionPrefixCacheLayout::SequenceHeads, true, causal, false, 0, 0, 0, false, + }).build(ctx, n1, weights.attention); + auto hidden = engine::modules::AddModule{}.build(ctx, input, attention); + auto n2 = engine::modules::LayerNormModule({config.hidden, 1.0e-5F, true, !causal || config.bias}) + .build(ctx, hidden, weights.norm2); + auto ff = engine::modules::FeedForwardGeluModule({config.hidden, 4 * config.hidden, + config.bias, engine::modules::GeluApproximation::ExactErf, GGML_PREC_F32}).build(ctx, n2, weights.mlp); + return engine::modules::AddModule{}.build(ctx, hidden, ff); +} + +} // namespace + +struct BarkTransformer::Impl { + std::shared_ptr assets; + engine::core::ExecutionContext & execution; + std::string prefix; + BarkTransformerConfig config; + bool causal; + std::shared_ptr weights; + + std::vector run(const std::vector> & channels, int head, bool last_only) const { + if (channels.empty() || channels.front().empty()) throw std::runtime_error("Bark transformer input is empty"); + const int64_t sequence = static_cast(channels.front().size()); + if (sequence > config.block_size) throw std::runtime_error("Bark transformer input exceeds block size"); + for (const auto & row : channels) if (static_cast(row.size()) != sequence) + throw std::runtime_error("Bark embedding channel lengths differ"); + if (head < 0 || static_cast(head) >= weights->heads.size()) throw std::runtime_error("Bark LM head index is invalid"); + std::unique_ptr context(ggml_init({128ULL * 1024ULL * 1024ULL, nullptr, true})); + if (!context) throw std::runtime_error("failed to create Bark graph context"); + engine::core::ModuleBuildContext build{context.get(), "bark.transformer", execution.backend_type()}; + engine::core::TensorValue hidden; + std::vector inputs; + std::vector masks; + std::vector> safe_channels = channels; + std::vector> mask_values(channels.size(), std::vector(static_cast(sequence), 1.0F)); + for (size_t channel = 0; channel < channels.size(); ++channel) { + for (int64_t step = 0; step < sequence; ++step) { + if (safe_channels[channel][static_cast(step)] < 0) { + safe_channels[channel][static_cast(step)] = 0; + mask_values[channel][static_cast(step)] = 0.0F; + } + } + auto * raw = ggml_new_tensor_1d(context.get(), GGML_TYPE_I32, sequence); + inputs.push_back(raw); + auto ids = engine::core::wrap_tensor(raw, engine::core::TensorShape::from_dims({sequence}), GGML_TYPE_I32); + auto embedded = engine::modules::EmbeddingModule({config.input_vocab, config.hidden}) + .build(build, ids, weights->embeddings.at(causal ? 0 : channel)); + embedded = engine::core::reshape_tensor(build, embedded, + engine::core::TensorShape::from_dims({1, sequence, config.hidden})); + auto * mask_raw = ggml_new_tensor_2d(context.get(), GGML_TYPE_F32, 1, sequence); + masks.push_back(mask_raw); + auto mask = engine::core::wrap_tensor(mask_raw, + engine::core::TensorShape::from_dims({1, sequence, 1}), GGML_TYPE_F32); + embedded = engine::core::wrap_tensor(ggml_mul(context.get(), embedded.tensor, mask.tensor), + embedded.shape, GGML_TYPE_F32); + hidden = hidden.valid() ? engine::modules::AddModule{}.build(build, hidden, embedded) : embedded; + } + std::vector positions(static_cast(sequence)); + for (int32_t i = 0; i < sequence; ++i) positions[static_cast(i)] = i; + auto * pos_raw = ggml_new_tensor_1d(context.get(), GGML_TYPE_I32, sequence); + auto pos_ids = engine::core::wrap_tensor(pos_raw, engine::core::TensorShape::from_dims({sequence}), GGML_TYPE_I32); + auto pos = engine::modules::EmbeddingModule({config.block_size, config.hidden}).build(build, pos_ids, weights->positions); + pos = engine::core::reshape_tensor(build, pos, engine::core::TensorShape::from_dims({1, sequence, config.hidden})); + hidden = engine::modules::AddModule{}.build(build, hidden, pos); + for (const auto & layer : weights->layers) hidden = block(build, hidden, layer, config, causal); + hidden = engine::modules::LayerNormModule({config.hidden, 1.0e-5F, true, !causal || config.bias}) + .build(build, hidden, weights->final_norm); + if (last_only) hidden = engine::modules::SliceModule({1, sequence - 1, 1}).build(build, hidden); + auto logits = engine::modules::LinearModule({config.hidden, config.output_vocab, false, GGML_PREC_F32}) + .build(build, hidden, {weights->heads[static_cast(head)], std::nullopt}); + ggml_set_output(logits.tensor); + auto * graph = ggml_new_graph_custom(context.get(), 65536, false); + ggml_build_forward_expand(graph, logits.tensor); + auto * buffer = ggml_backend_alloc_ctx_tensors(context.get(), execution.backend()); + if (!buffer) throw std::runtime_error("failed to allocate Bark graph tensors"); + for (size_t i = 0; i < inputs.size(); ++i) { + ggml_backend_tensor_set(inputs[i], safe_channels[i].data(), 0, safe_channels[i].size() * sizeof(int32_t)); + ggml_backend_tensor_set(masks[i], mask_values[i].data(), 0, mask_values[i].size() * sizeof(float)); + } + ggml_backend_tensor_set(pos_raw, positions.data(), 0, positions.size() * sizeof(int32_t)); + engine::core::set_backend_threads(execution.backend(), execution.config().threads); + const auto status = engine::core::compute_backend_graph(execution.backend(), graph); + ggml_backend_synchronize(execution.backend()); + if (status != GGML_STATUS_SUCCESS) throw std::runtime_error("Bark transformer graph compute failed"); + std::vector out(static_cast((last_only ? 1 : sequence) * config.output_vocab)); + ggml_backend_tensor_get(logits.tensor, out.data(), 0, out.size() * sizeof(float)); + engine::core::release_backend_graph_resources(execution.backend(), graph); + ggml_backend_buffer_free(buffer); + return out; + } +}; + +BarkTransformer::BarkTransformer(std::shared_ptr assets, + engine::core::ExecutionContext & execution, + std::string prefix, + BarkTransformerConfig config, + bool causal, + engine::assets::TensorStorageType storage) + : impl_(std::make_unique(Impl{assets, execution, std::move(prefix), config, causal, nullptr})) { + impl_->weights = load_weights(*impl_->assets, execution.backend(), execution.backend_type(), impl_->prefix, config, causal, storage); +} + +BarkTransformer::~BarkTransformer() = default; + +std::vector BarkTransformer::causal_logits(const std::vector> & channels) const { + if (!impl_->causal) throw std::runtime_error("fine Bark transformer cannot run causal logits"); + return impl_->run(channels, 0, true); +} + +std::vector BarkTransformer::fine_logits(const std::vector> & codes, int codebook_idx) const { + if (impl_->causal || codebook_idx < 1 || codebook_idx > 7) throw std::runtime_error("invalid Bark fine codebook"); + return impl_->run(std::vector>(codes.begin(), codes.begin() + codebook_idx + 1), codebook_idx - 1, false); +} + +} // namespace engine::models::bark_tts diff --git a/tests/bark_tts/bark_codec_parity.cpp b/tests/bark_tts/bark_codec_parity.cpp new file mode 100644 index 00000000..ef7cfbb9 --- /dev/null +++ b/tests/bark_tts/bark_codec_parity.cpp @@ -0,0 +1,40 @@ +#include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/model_spec/package.h" +#include "engine/models/bark_tts/assets.h" +#include "engine/models/bark_tts/codec.h" + +#include +#include +#include +#include + +int main(int argc, char ** argv) try { + if (argc != 5) { + std::cerr << "usage: bark_codec_parity \n"; + return 2; + } + const int frames = std::stoi(argv[3]); + engine::model_spec::ScopedSpecOverride spec{std::filesystem::path(argv[2]), std::filesystem::path(argv[1])}; + auto assets = engine::models::bark_tts::load_bark_assets(argv[1]); + const auto found = assets->presets.find("v2/en_speaker_6"); + if (found == assets->presets.end()) throw std::runtime_error("missing parity preset"); + std::vector> codes(8); + for (int book = 0; book < 8; ++book) { + codes[book].assign(found->second.fine[book].begin(), + found->second.fine[book].begin() + std::min(frames, found->second.fine[book].size())); + } + engine::core::BackendConfig config; + config.type = engine::core::BackendType::Cpu; + config.threads = 8; + engine::core::ExecutionContext execution(config); + engine::models::bark_tts::BarkCodecDecoder decoder(assets, execution, engine::assets::TensorStorageType::Native); + const auto waveform = decoder.decode(codes); + std::ofstream output(argv[4], std::ios::binary); + output.write(reinterpret_cast(waveform.data()), static_cast(waveform.size() * sizeof(float))); + std::cout << "samples=" << waveform.size() << "\n"; + return 0; +} catch (const std::exception & error) { + std::cerr << error.what() << "\n"; + return 1; +} diff --git a/tests/bark_tts/bark_transformer_parity.cpp b/tests/bark_tts/bark_transformer_parity.cpp new file mode 100644 index 00000000..ca0ed69a --- /dev/null +++ b/tests/bark_tts/bark_transformer_parity.cpp @@ -0,0 +1,84 @@ +#include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/model_spec/package.h" +#include "engine/models/bark_tts/assets.h" +#include "engine/models/bark_tts/tokenizer.h" +#include "engine/models/bark_tts/transformer.h" + +#include +#include +#include +#include + +int main(int argc, char ** argv) try { + if (argc != 4) { + std::cerr << "usage: bark_transformer_parity \n"; + return 2; + } + engine::model_spec::ScopedSpecOverride spec{std::filesystem::path(argv[2]), std::filesystem::path(argv[1])}; + auto assets = engine::models::bark_tts::load_bark_assets(argv[1]); + const auto & preset = assets->presets.at("v2/en_speaker_6"); + engine::models::bark_tts::BarkTokenizer tokenizer(assets->resources.read_text("tokenizer_json")); + auto text = tokenizer.encode("Hello, Bark!"); + for (auto & id : text) id += 10048; + text.resize(256, 129595); + std::vector history(preset.semantic.end() - 256, preset.semantic.end()); + text.push_back(129599); + history.push_back(-1); + engine::core::BackendConfig config; + config.type = engine::core::BackendType::Cpu; + config.threads = 8; + engine::core::ExecutionContext execution(config); + engine::models::bark_tts::BarkTransformer transformer(assets, execution, "semantic", assets->config.semantic, + true, engine::assets::TensorStorageType::Native); + const auto logits = transformer.causal_logits({text, history}); + std::ofstream output(argv[3], std::ios::binary); + output.write(reinterpret_cast(logits.data()), static_cast(logits.size() * sizeof(float))); + output.close(); + text.push_back(147); + history.push_back(-1); + const auto logits2 = transformer.causal_logits({text, history}); + std::ofstream output2(std::string(argv[3]) + ".semantic2", std::ios::binary); + output2.write(reinterpret_cast(logits2.data()), + static_cast(logits2.size() * sizeof(float))); + constexpr double ratio = 75.0 / 49.9 * 2.0; + std::vector semantic_seed{147, 302, 2089, 330, 602, 206, 218, 2009}; + std::vector sem_history = preset.semantic; + std::vector coarse_history; + for (size_t frame = 0; frame < preset.coarse[0].size(); ++frame) + for (int book = 0; book < 2; ++book) + coarse_history.push_back(preset.coarse[book][frame] + book * 1024 + 10000); + const int max_sem_history = static_cast(std::floor(630.0 / ratio)); + const int sem_count = std::min({max_sem_history, static_cast(sem_history.size() / 2 * 2), + static_cast(std::floor(coarse_history.size() / ratio))}); + const int coarse_count = static_cast(std::round(sem_count * ratio)); + sem_history.erase(sem_history.begin(), sem_history.end() - sem_count); + coarse_history.erase(coarse_history.begin(), coarse_history.end() - coarse_count); + coarse_history.resize(coarse_history.size() - 2); + sem_history.insert(sem_history.end(), semantic_seed.begin(), semantic_seed.end()); + if (sem_history.size() > 256) sem_history.resize(256); + sem_history.resize(256, 12048); + sem_history.push_back(12050); + const size_t take = std::min(630, coarse_history.size()); + sem_history.insert(sem_history.end(), coarse_history.end() - static_cast(take), coarse_history.end()); + engine::models::bark_tts::BarkTransformer coarse(assets, execution, "coarse_acoustics", assets->config.coarse, + true, engine::assets::TensorStorageType::Native); + const auto coarse_logits = coarse.causal_logits({sem_history}); + std::ofstream coarse_output(std::string(argv[3]) + ".coarse", std::ios::binary); + coarse_output.write(reinterpret_cast(coarse_logits.data()), + static_cast(coarse_logits.size() * sizeof(float))); + std::vector> fine_codes(8); + for (int book = 0; book < 8; ++book) + fine_codes[book].assign(preset.fine[book].begin(), preset.fine[book].begin() + 64); + engine::models::bark_tts::BarkTransformer fine(assets, execution, "fine_acoustics", assets->config.fine, + false, engine::assets::TensorStorageType::Native); + const auto fine_logits = fine.fine_logits(fine_codes, 2); + std::ofstream fine_output(std::string(argv[3]) + ".fine", std::ios::binary); + fine_output.write(reinterpret_cast(fine_logits.data()), + static_cast(fine_logits.size() * sizeof(float))); + std::cout << "semantic_logits=" << logits.size() << " fine_logits=" << fine_logits.size() << "\n"; + return 0; +} catch (const std::exception & error) { + std::cerr << error.what() << "\n"; + return 1; +} diff --git a/tests/unittests/test_bark_tokenizer.cpp b/tests/unittests/test_bark_tokenizer.cpp new file mode 100644 index 00000000..684fac28 --- /dev/null +++ b/tests/unittests/test_bark_tokenizer.cpp @@ -0,0 +1,20 @@ +#include "engine/models/bark_tts/tokenizer.h" + +#include +#include + +int main() { + const std::string json = R"({"model":{"vocab":{"[UNK]":100,"[CLS]":101,"[SEP]":102,"Hello":31178,",":117,"Bar":20698,"##k":10174,"!":106,"你":2262,"好":3240}}})"; + engine::models::bark_tts::BarkTokenizer tokenizer(json); + const std::vector expected{31178, 117, 20698, 10174, 106}; + if (tokenizer.encode("Hello, Bark!") != expected) { + std::cerr << "Bark WordPiece tokenization mismatch\n"; + return 1; + } + const std::vector chinese{2262, 3240}; + if (tokenizer.encode("你好") != chinese) { + std::cerr << "Bark Chinese-character splitting mismatch\n"; + return 1; + } + return 0; +} diff --git a/tools/community_models/convert_bark.py b/tools/community_models/convert_bark.py new file mode 100644 index 00000000..2a3fad13 --- /dev/null +++ b/tools/community_models/convert_bark.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +"""Convert Hugging Face ``suno/bark-small`` into an audio.cpp GGUF package. + +The generated package contains Bark's semantic, coarse, and fine transformers, +the eight EnCodec codebooks used by Bark, the EnCodec decoder, the tokenizer, +and every downloaded speaker preset. Weight-normalized EnCodec convolutions +are materialized before export so the native runtime does not need PyTorch. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import subprocess +import tempfile +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[2] +SPEC = REPO_ROOT / "model_specs" / "bark_tts.json" + + +def _load_checkpoint(path: Path): + try: + import torch + except ImportError as exc: + raise SystemExit("PyTorch is required: uv run --with torch --with safetensors " + __file__) from exc + checkpoint = torch.load(path, map_location="cpu", weights_only=True) + return checkpoint.get("model", checkpoint) + + +def _materialize_weight_norm(state: dict, prefix: str): + import torch + + g_key, v_key = prefix + ".weight_g", prefix + ".weight_v" + if g_key not in state or v_key not in state: + return None + g, v = state[g_key].float(), state[v_key].float() + # PyTorch weight_norm(..., dim=0): one magnitude per output channel. + dims = tuple(range(1, v.ndim)) + return (v * (g / torch.linalg.vector_norm(v, dim=dims, keepdim=True))).contiguous() + + +def _export_tensors(source: Path, output: Path) -> None: + from safetensors.torch import save_file + + state = _load_checkpoint(source / "pytorch_model.bin") + exported = {} + for name, tensor in state.items(): + if name.endswith(".attn.bias") or name.startswith("codec_model.encoder."): + continue + if name.startswith("codec_model.quantizer.layers."): + parts = name.split(".") + if int(parts[3]) >= 8 or not name.endswith(".codebook.embed"): + continue + if name.endswith(".weight_g") or name.endswith(".weight_v"): + prefix = name.rsplit(".", 1)[0] + weight_name = prefix + ".weight" + if weight_name not in exported: + exported[weight_name] = _materialize_weight_norm(state, prefix) + continue + # Fine-model LM heads are tied to embeddings in the PyTorch object; + # safetensors requires each exported name to own its storage. + exported[name] = tensor.detach().clone().contiguous() + save_file(exported, str(output)) + + +def _collect_presets(source: Path, output: Path) -> int: + import numpy as np + + preset_root = source / "speaker_embeddings" + records = {} + for semantic in sorted(preset_root.rglob("*_semantic_prompt.npy")): + stem = semantic.name.removesuffix("_semantic_prompt.npy") + relative = semantic.parent.relative_to(preset_root) + key = str(relative / stem) + coarse = semantic.with_name(stem + "_coarse_prompt.npy") + fine = semantic.with_name(stem + "_fine_prompt.npy") + if not coarse.is_file() or not fine.is_file(): + raise RuntimeError(f"incomplete Bark preset: {key}") + records[key] = { + "semantic": np.load(semantic).astype("int32").tolist(), + "coarse": np.load(coarse).astype("int32").tolist(), + "fine": np.load(fine).astype("int32").tolist(), + } + output.write_text(json.dumps({"version": 1, "presets": records}, separators=(",", ":"))) + return len(records) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", type=Path, required=True, help="downloaded suno/bark-small snapshot") + parser.add_argument("--converter", type=Path, required=True, help="built audiocpp_gguf executable") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--type", default="q8_0", choices=["orig", "f16", "bf16", "q8_0"]) + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + source, output = args.source.resolve(), args.output.resolve() + required = ["pytorch_model.bin", "config.json", "generation_config.json", "tokenizer.json"] + missing = [name for name in required if not (source / name).is_file()] + if missing: + raise SystemExit(f"Bark source is missing: {', '.join(missing)}") + if not args.converter.is_file(): + raise SystemExit(f"converter not found: {args.converter}") + output.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix="audiocpp-bark-") as temp_string: + temp = Path(temp_string) + tensors = temp / "bark.safetensors" + presets = temp / "speaker_presets.json" + _export_tensors(source, tensors) + preset_count = _collect_presets(source, presets) + for name in ("config.json", "generation_config.json", "tokenizer.json", "tokenizer_config.json"): + candidate = source / name + if candidate.is_file(): + shutil.copy2(candidate, temp / name) + command = [ + str(args.converter), "--input", f"bark={tensors}", "--root", str(temp), + "--family", "bark_tts", "--model-spec", str(SPEC), "--type", args.type, + "--output", str(output), + ] + if args.type == "q8_0": + # Autoregressive errors compound across Bark's semantic and coarse + # stages. Keep both complete causal transformers, plus the waveform + # decoder and fine lookup/output tensors, in F16. Only the fine + # transformer's dense matrices are safe to quantize to Q8_0. + for pattern in ( + "bark/codec_model.*", "bark/semantic.*", "bark/coarse_acoustics.*", + "bark/fine_acoustics.input_embeds_layers.*", + "bark/fine_acoustics.position_embeds_layer.*", "bark/fine_acoustics.layernorm_final.*", + "bark/fine_acoustics.lm_heads.*", + ): + command.extend(["--keep-type", pattern + "=f16"]) + if args.overwrite: + command.append("--overwrite") + print("+", " ".join(command), flush=True) + subprocess.run(command, check=True) + print(f"wrote {output} with {preset_count} speaker presets") + + +if __name__ == "__main__": + main()