OmniVoice: make the audio tokenizer codebook weight type configurable

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. |

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(

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;

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",

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<float> acoustic_samples_24k;
     std::vector<float> semantic_samples_16k_padded;
@@ -970,9 +1005,18 @@ std::shared_ptr<const HiggsAudioV2TokenizerWeights> 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<HiggsAudioV2TokenizerWeights>();
     weights->store = std::make_shared<core::BackendWeightStore>(
         backend,
@@ -1027,7 +1071,7 @@ std::shared_ptr<const HiggsAudioV2TokenizerWeights> 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<const HiggsAudioV2TokenizerWeights> 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<const HiggsAudioV2TokenizerWeights> 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<const HiggsAudioV2TokenizerWeights> 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<Impl>()) {
     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;

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"},

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_,