diff --git a/docs/models/omnivoice.md b/docs/models/omnivoice.md index 4de5a5828..3ccb4cf6a 100644 --- a/docs/models/omnivoice.md +++ b/docs/models/omnivoice.md @@ -145,6 +145,7 @@ Arena and weight-context sizes are in MiB. | `omnivoice.weight_type` | `native`, `f32`, `f16`, `bf16`, `q8_0` | `native` | Generator weight storage fallback. | | `omnivoice.generator_weight_type` | `native`, `f32`, `f16`, `bf16`, `q8_0` | inherits `weight_type` | Explicit generator weight storage override. | | `omnivoice.audio_tokenizer_weight_type` | `native`, `f32`, `f16`, `bf16`, `q8_0` | `native` | Runtime storage type for audio tokenizer weights. | +| `omnivoice.audio_tokenizer_codebook_weight_type` | `native`, `f32`, `f16`, `bf16` | `f32` | Storage type for the audio tokenizer codebooks, codebook score matrix, semantic feature projection, and `fc2`. These ignore `audio_tokenizer_weight_type`. `native` keeps the packaged type and is rejected when that type is quantized, as in the Q8_0 GGUF. | | `omnivoice.audio_tokenizer_graph_arena_mb` | positive integer | `128` | Audio tokenizer graph arena. | | `omnivoice.generator_prefill_graph_arena_mb` | positive integer | `256` | Generator prefill graph arena. | | `omnivoice.generator_decode_graph_arena_mb` | positive integer | `256` | Generator decode graph arena. | diff --git a/include/engine/models/omnivoice/audio_tokenizer.h b/include/engine/models/omnivoice/audio_tokenizer.h index 29b4f9f73..f985a7e7c 100644 --- a/include/engine/models/omnivoice/audio_tokenizer.h +++ b/include/engine/models/omnivoice/audio_tokenizer.h @@ -33,7 +33,8 @@ class OmniVoiceHiggsAudioV2TokenizerRuntime { core::ExecutionContext & execution_context, size_t graph_arena_bytes, size_t weight_context_bytes, - engine::assets::TensorStorageType weight_storage_type); + engine::assets::TensorStorageType weight_storage_type, + engine::assets::TensorStorageType codebook_weight_storage_type = engine::assets::TensorStorageType::F32); ~OmniVoiceHiggsAudioV2TokenizerRuntime(); OmniVoiceAudioTokens encode_reference_audio( diff --git a/include/engine/models/omnivoice/session.h b/include/engine/models/omnivoice/session.h index 3b12bf87c..87eb91c5c 100644 --- a/include/engine/models/omnivoice/session.h +++ b/include/engine/models/omnivoice/session.h @@ -85,6 +85,7 @@ class OmniVoiceSession final size_t audio_tokenizer_weight_context_bytes_ = 128ull * 1024ull * 1024ull; size_t generator_weight_context_bytes_ = 256ull * 1024ull * 1024ull; engine::assets::TensorStorageType audio_tokenizer_weight_storage_type_ = engine::assets::TensorStorageType::Native; + engine::assets::TensorStorageType audio_tokenizer_codebook_weight_storage_type_ = engine::assets::TensorStorageType::F32; engine::assets::TensorStorageType generator_weight_storage_type_ = engine::assets::TensorStorageType::Native; bool mem_saver_ = false; OmniVoiceGeneratorPerfMode generator_perf_mode_ = OmniVoiceGeneratorPerfMode::Standard; diff --git a/model_specs/omnivoice.json b/model_specs/omnivoice.json index fee751d03..76b5e62df 100644 --- a/model_specs/omnivoice.json +++ b/model_specs/omnivoice.json @@ -187,6 +187,14 @@ "required": false, "default": "native" }, + { + "name": "audio_tokenizer_codebook_weight_type", + "type": "enum", + "description": "Storage type for the audio tokenizer codebooks, codebook score matrix, semantic feature projection, and fc2; default f32.", + "values": ["native", "f32", "f16", "bf16"], + "required": false, + "default": "f32" + }, { "name": "audio_tokenizer_graph_arena_mb", "type": "int", diff --git a/src/models/omnivoice/audio_tokenizer.cpp b/src/models/omnivoice/audio_tokenizer.cpp index 0996346f0..cf69d2704 100644 --- a/src/models/omnivoice/audio_tokenizer.cpp +++ b/src/models/omnivoice/audio_tokenizer.cpp @@ -70,6 +70,41 @@ void validate_weight_storage_type(assets_ns::TensorStorageType storage_type) { } } +// Quantized types are rejected: these tensors are what broke runtime q8_0 tokenizer loading. +void validate_codebook_weight_storage_type(assets_ns::TensorStorageType storage_type) { + switch (storage_type) { + case assets_ns::TensorStorageType::Native: + case assets_ns::TensorStorageType::F32: + case assets_ns::TensorStorageType::F16: + case assets_ns::TensorStorageType::BF16: + return; + default: + throw std::runtime_error( + "OmniVoice audio tokenizer codebook_weight_type supports only native, f32, f16, and bf16"); + } +} + +// native keeps the packaged type, so a quantized package (e.g. the Q8_0 GGUF) would load these tensors quantized. +void validate_native_codebook_tensor_type( + const assets_ns::TensorSource & source, + const std::string & name, + assets_ns::TensorStorageType storage_type) { + if (storage_type != assets_ns::TensorStorageType::Native) { + return; + } + const auto dtype = source.require_metadata(name).dtype; + switch (assets_ns::tensor_storage_type_for_dtype(dtype)) { + case assets_ns::TensorStorageType::F32: + case assets_ns::TensorStorageType::F16: + case assets_ns::TensorStorageType::BF16: + return; + default: + throw std::runtime_error( + "OmniVoice audio tokenizer codebook_weight_type=native requires f32, f16, or bf16 tensors, but " + + name + " is " + dtype + " in this package; use codebook_weight_type=f32, f16, or bf16"); + } +} + struct NormalizedReferenceAudio { std::vector acoustic_samples_24k; std::vector semantic_samples_16k_padded; @@ -970,9 +1005,18 @@ std::shared_ptr load_weights( ggml_backend_t backend, core::BackendType backend_type, size_t weight_context_bytes, - assets_ns::TensorStorageType storage_type) { + assets_ns::TensorStorageType storage_type, + assets_ns::TensorStorageType codebook_storage_type) { const auto & config = assets.config.audio_tokenizer; const auto & source = *assets.audio_tokenizer_weights; + validate_native_codebook_tensor_type(source, "semantic_model.feature_projection.projection.weight", codebook_storage_type); + for (int64_t quantizer_index = 0; quantizer_index < config.num_codebooks; ++quantizer_index) { + validate_native_codebook_tensor_type( + source, + "quantizer.quantizers." + std::to_string(quantizer_index) + ".codebook.embed", + codebook_storage_type); + } + validate_native_codebook_tensor_type(source, "fc2.weight", codebook_storage_type); auto weights = std::make_shared(); weights->store = std::make_shared( backend, @@ -1027,7 +1071,7 @@ std::shared_ptr load_weights( "semantic_model.feature_projection.projection", semantic.hidden_size, semantic.conv_dim.back(), - assets_ns::TensorStorageType::F32, + codebook_storage_type, true); if (semantic.do_stable_layer_norm) { @@ -1247,7 +1291,7 @@ std::shared_ptr load_weights( quantizer.score = { weights->store->make_from_f32( core::TensorShape::from_dims({config.codebook_size, config.codebook_dim}), - assets_ns::TensorStorageType::F32, + codebook_storage_type, score_weight), weights->store->make_from_f32( core::TensorShape::from_dims({config.codebook_size}), @@ -1257,7 +1301,7 @@ std::shared_ptr load_weights( quantizer.codebook = weights->store->load_tensor( source, prefix + ".codebook.embed", - assets_ns::TensorStorageType::F32, + codebook_storage_type, {config.codebook_size, config.codebook_dim}); quantizer.project_out = load_linear( *weights->store, @@ -1276,7 +1320,7 @@ std::shared_ptr load_weights( "fc2", acoustic.hidden_size, config.hidden_size, - assets_ns::TensorStorageType::F32, + codebook_storage_type, true); weights->acoustic_decoder.conv1 = load_flat_conv1d( @@ -2252,12 +2296,14 @@ OmniVoiceHiggsAudioV2TokenizerRuntime::OmniVoiceHiggsAudioV2TokenizerRuntime( core::ExecutionContext & execution_context, size_t graph_arena_bytes, size_t weight_context_bytes, - engine::assets::TensorStorageType weight_storage_type) + engine::assets::TensorStorageType weight_storage_type, + engine::assets::TensorStorageType codebook_weight_storage_type) : impl_(std::make_unique()) { if (assets == nullptr) { throw std::runtime_error("OmniVoice audio tokenizer requires assets"); } validate_weight_storage_type(weight_storage_type); + validate_codebook_weight_storage_type(codebook_weight_storage_type); impl_->assets = std::move(assets); impl_->execution_context = &execution_context; impl_->graph_arena_bytes = graph_arena_bytes; @@ -2266,7 +2312,8 @@ OmniVoiceHiggsAudioV2TokenizerRuntime::OmniVoiceHiggsAudioV2TokenizerRuntime( execution_context.backend(), execution_context.backend_type(), weight_context_bytes, - weight_storage_type); + weight_storage_type, + codebook_weight_storage_type); } OmniVoiceHiggsAudioV2TokenizerRuntime::~OmniVoiceHiggsAudioV2TokenizerRuntime() = default; diff --git a/src/models/omnivoice/loader.cpp b/src/models/omnivoice/loader.cpp index cd01b8e10..9bc46c423 100644 --- a/src/models/omnivoice/loader.cpp +++ b/src/models/omnivoice/loader.cpp @@ -59,6 +59,7 @@ runtime::ModelCliInterface cli(const OmniVoiceAssets &) { {"omnivoice.weight_type", "native|f32|f16|bf16|q8_0", "Shared generator weight storage type; default native.", false, "native"}, {"omnivoice.generator_weight_type", "native|f32|f16|bf16|q8_0", "Generator weight storage type; defaults to weight_type, which defaults to native."}, {"omnivoice.audio_tokenizer_weight_type", "native|f32|f16|bf16|q8_0", "Audio tokenizer weight storage type; default native.", false, "native"}, + {"omnivoice.audio_tokenizer_codebook_weight_type", "native|f32|f16|bf16", "Storage type for the audio tokenizer codebooks, codebook score matrix, semantic feature projection, and fc2; default f32.", false, "f32"}, {"omnivoice.audio_tokenizer_graph_arena_mb", "int", "Audio tokenizer graph arena size in MiB; default 128.", false, "128", "1"}, {"omnivoice.generator_prefill_graph_arena_mb", "int", "Generator prefill graph arena size in MiB; default 256.", false, "256", "1"}, {"omnivoice.generator_decode_graph_arena_mb", "int", "Generator decode graph arena size in MiB; default 256.", false, "256", "1"}, diff --git a/src/models/omnivoice/session.cpp b/src/models/omnivoice/session.cpp index 83e24101d..7d67c4526 100644 --- a/src/models/omnivoice/session.cpp +++ b/src/models/omnivoice/session.cpp @@ -231,6 +231,8 @@ OmniVoiceSession::OmniVoiceSession( runtime::parse_size_mb_option(options.options, {"omnivoice.generator_weight_context_mb"}, kDefaultGeneratorWeightContextBytes)), audio_tokenizer_weight_storage_type_( option_weight_type(options, "omnivoice.audio_tokenizer_weight_type", engine::assets::TensorStorageType::Native)), + audio_tokenizer_codebook_weight_storage_type_( + option_weight_type(options, "omnivoice.audio_tokenizer_codebook_weight_type", engine::assets::TensorStorageType::F32)), generator_weight_storage_type_( option_weight_type( options, @@ -244,7 +246,8 @@ OmniVoiceSession::OmniVoiceSession( execution_context(), audio_tokenizer_graph_arena_bytes_, audio_tokenizer_weight_context_bytes_, - audio_tokenizer_weight_storage_type_), + audio_tokenizer_weight_storage_type_, + audio_tokenizer_codebook_weight_storage_type_), prompt_builder_(assets_, tokenizer_), generator_( assets_,