Add gemma-4-12B unified (encoder-free) preprocessing support

shared/api/gemma4_audio_features.hpp

@@ -294,6 +294,9 @@ class Gemma4LogMel {
       } else if (key == "per_bin_stddev") {
         auto& v = std::get<std::vector<double>>(value);
         per_bin_stddev_.assign(v.begin(), v.end());
+      } else if (key == "type") {
+        // Consumed by the Gemma4Audio dispatcher (selects this log-mel path);
+        // ignored here so a forwarded attribute dict does not error.
       } else {
         return {kOrtxErrorInvalidArgument,
                 "[Gemma4LogMel]: unknown attribute '" + key + "'"};
@@ -344,4 +347,156 @@ class Gemma4LogMel {
   std::vector<float> mel_filters_;  // (n_freq x feature_size), row-major
 };
 
+// Gemma 4 *unified* (encoder-free, gemma-4-12B) audio feature extraction.
+//
+// Unlike Gemma4LogMel (128-dim USM log-mel), the unified model has no audio
+// encoder: each audio soft token is simply a fixed-length chunk of the raw
+// 16 kHz waveform. This op reproduces HuggingFace
+// ``Gemma4UnifiedAudioFeatureExtractor._extract_waveform_features`` exactly:
+// zero-pad the waveform to a multiple of ``audio_samples_per_token`` and
+// reshape it into ``(num_tokens, audio_samples_per_token)`` frames.
+//
+// Pipeline:  AudioDecoder  ->  Gemma4Audio (type="raw_frames")
+//
+// Inputs:   float (1, num_samples) — mono PCM at `sampling_rate` Hz
+// Outputs:  float (num_tokens, audio_samples_per_token) — raw waveform frames
+//           bool  (num_tokens,)                          — frame-level mask (all true)
+//
+// The mask is emitted (all-true) so that this path shares the (features, mask)
+// output signature of Gemma4LogMel, letting a single Gemma4Audio op cover both.
+// It matches HuggingFace ``Gemma4UnifiedAudioFeatureExtractor``, which returns
+// ``input_features`` and ``input_features_mask``; ragged clips are zero-padded
+// by the batch framework when stacking, so per-clip frames are all valid.
+class Gemma4UnifiedAudioFrames {
+ public:
+  Gemma4UnifiedAudioFrames() = default;
+
+  OrtxStatus Compute(const ortc::Tensor<float>& pcm_input,
+                     ortc::Tensor<float>& frames_out,
+                     ortc::Tensor<bool>& mask_out) {
+    const auto& pcm_shape = pcm_input.Shape();
+    if (pcm_shape.size() != 2 || pcm_shape[0] != 1) {
+      return {kOrtxErrorInvalidArgument,
+              "[Gemma4UnifiedAudioFrames]: expected (1, num_samples) float input"};
+    }
+
+    const int64_t num_samples = pcm_shape[1];
+    const int64_t spt = audio_samples_per_token_;
+    // Zero-pad to a whole number of frames (ceil division), matching HF's
+    // ``pad_len = (-len(waveform)) % audio_samples_per_token``.
+    const int64_t num_tokens = (num_samples + spt - 1) / spt;
+
+    float* out = frames_out.Allocate({num_tokens, spt});
+    bool* mask = mask_out.Allocate({num_tokens});
+    if (num_tokens == 0) {
+      return {};
+    }
+    // Fill the (possibly padded) tail of the last frame with the padding value,
+    // then copy the real samples over the front.
+    std::fill(out, out + static_cast<size_t>(num_tokens) * spt, padding_value_);
+    std::copy(pcm_input.Data(), pcm_input.Data() + num_samples, out);
+    // Every frame of a single clip is valid (padding lives within the last frame).
+    std::fill(mask, mask + num_tokens, true);
+    return {};
+  }
+
+  template <typename DictT>
+  OrtxStatus Init(const DictT& attrs) {
+    for (const auto& [key, value] : attrs) {
+      if (key == "audio_samples_per_token" || key == "feature_size") {
+        // ``feature_size`` is accepted as an alias: HF sets feature_size ==
+        // audio_samples_per_token (both default to 640).
+        audio_samples_per_token_ = std::get<int64_t>(value);
+      } else if (key == "sampling_rate") {
+        sampling_rate_ = std::get<int64_t>(value);
+      } else if (key == "padding_value") {
+        padding_value_ = static_cast<float>(std::get<double>(value));
+      } else if (key == "type") {
+        // Consumed by the Gemma4Audio dispatcher (selects this raw-frames path).
+      } else {
+        return {kOrtxErrorInvalidArgument,
+                "[Gemma4UnifiedAudioFrames]: unknown attribute '" + key + "'"};
+      }
+    }
+    if (audio_samples_per_token_ <= 0) {
+      return {kOrtxErrorInvalidArgument,
+              "[Gemma4UnifiedAudioFrames]: audio_samples_per_token must be positive"};
+    }
+    // The op frames whatever PCM the upstream AudioDecoder produces; the frame
+    // size is defined in samples, so this op is intrinsically sample-rate
+    // agnostic. ``sampling_rate`` therefore only documents the rate the decoder
+    // is expected to output (16 kHz for the gemma-4 contract, where 640 samples
+    // == 40 ms). Reject non-positive values so a misconfiguration is loud
+    // rather than silently producing frames at an unintended rate.
+    if (sampling_rate_ <= 0) {
+      return {kOrtxErrorInvalidArgument,
+              "[Gemma4UnifiedAudioFrames]: sampling_rate must be positive"};
+    }
+    return {};
+  }
+
+ private:
+  int64_t audio_samples_per_token_ = 640;  // 640 samples = 40 ms @ 16 kHz
+  // Expected decoder output rate. Informational only: the op frames by sample
+  // count and does not resample (see the note in Init()).
+  int64_t sampling_rate_ = 16000;
+  float padding_value_ = 0.0f;
+};
+
+// Unified Gemma 4 audio feature extraction op.
+//
+// A single registered op that dispatches, via the ``type`` attribute, to one of
+// the gemma-4 audio front-ends rather than exposing a separate kernel per model
+// variant:
+//
+//   type = "log_mel"     (default) -> 128-dim USM log-mel spectrogram (E2B/E4B)
+//   type = "raw_frames"            -> raw 640-sample waveform frames  (12B unified)
+//
+// Both branches share the (features: float, mask: bool) output signature.
+//
+// Pipeline:  AudioDecoder  ->  Gemma4Audio
+class Gemma4Audio {
+ public:
+  Gemma4Audio() = default;
+
+  OrtxStatus Compute(const ortc::Tensor<float>& pcm_input,
+                     ortc::Tensor<float>& features_out,
+                     ortc::Tensor<bool>& mask_out) {
+    if (mode_ == Mode::kRawFrames) {
+      return raw_frames_.Compute(pcm_input, features_out, mask_out);
+    }
+    return log_mel_.Compute(pcm_input, features_out, mask_out);
+  }
+
+  template <typename DictT>
+  OrtxStatus Init(const DictT& attrs) {
+    // Select the front-end from the ``type`` attribute, then forward the full
+    // attribute dict to the chosen implementation (each ignores the ``type``
+    // key). Defaults to log-mel for backward compatibility.
+    for (const auto& [key, value] : attrs) {
+      if (key == "type") {
+        const std::string& type = std::get<std::string>(value);
+        if (type == "raw_frames") {
+          mode_ = Mode::kRawFrames;
+        } else if (type == "log_mel") {
+          mode_ = Mode::kLogMel;
+        } else {
+          return {kOrtxErrorInvalidArgument,
+                  "[Gemma4Audio]: unknown type '" + type + "' (expected 'log_mel' or 'raw_frames')"};
+        }
+      }
+    }
+    if (mode_ == Mode::kRawFrames) {
+      return raw_frames_.Init(attrs);
+    }
+    return log_mel_.Init(attrs);
+  }
+
+ private:
+  enum class Mode { kLogMel, kRawFrames };
+  Mode mode_ = Mode::kLogMel;
+  Gemma4LogMel log_mel_;
+  Gemma4UnifiedAudioFrames raw_frames_;
+};
+
 }  // namespace ort_extensions

shared/api/speech_extractor.cc

@@ -17,7 +17,8 @@ Operation::KernelRegistry SpeechFeatureExtractor::kernel_registry_ = {
     {"NemoLogMel", []() { return CreateKernelInstance(&NemoLogMel::Compute); }},
     {"PerFeatureNormalize", []() { return CreateKernelInstance(&PerFeatureNormalize::Compute); }},
     {"Phi4AudioEmbed", []() { return CreateKernelInstance(&Phi4AudioEmbed::Compute); }},
-    {"Gemma4LogMel", []() { return CreateKernelInstance(&Gemma4LogMel::Compute); }}};
+    {"Gemma4LogMel", []() { return CreateKernelInstance(&Gemma4LogMel::Compute); }},
+    {"Gemma4Audio", []() { return CreateKernelInstance(&Gemma4Audio::Compute); }}};
 
 SpeechFeatureExtractor::SpeechFeatureExtractor() : OrtxObjectImpl(extObjectKind_t::kOrtxKindFeatureExtractor) {}
 

test/data/models/gemma-4-unified/audio_feature_extraction.json

@@ -0,0 +1,24 @@
+{
+  "feature_extraction": {
+    "sequence": [
+      {
+        "operation": {
+          "name": "audio_decoder",
+          "type": "AudioDecoder"
+        }
+      },
+      {
+        "operation": {
+          "name": "gemma4_audio",
+          "type": "Gemma4Audio",
+          "attrs": {
+            "type": "raw_frames",
+            "audio_samples_per_token": 640,
+            "sampling_rate": 16000,
+            "padding_value": 0.0
+          }
+        }
+      }
+    ]
+  }
+}

test/data/models/gemma-4-unified/image_processor.json

@@ -0,0 +1,27 @@
+{
+  "processor": {
+    "name": "gemma_4_unified_image_processing",
+    "transforms": [
+      {
+        "operation": {
+          "name": "decode_image",
+          "type": "DecodeImage",
+          "attrs": {
+            "color_space": "RGB"
+          }
+        }
+      },
+      {
+        "operation": {
+          "name": "gemma4_image_transform",
+          "type": "Gemma4ImageTransform",
+          "attrs": {
+            "patch_size": 48,
+            "max_soft_tokens": 280,
+            "pooling_kernel_size": 1
+          }
+        }
+      }
+    ]
+  }
+}

test/data/models/gemma-4/audio_feature_extraction.json

@@ -9,9 +9,10 @@
       },
       {
         "operation": {
-          "name": "gemma4_log_mel",
-          "type": "Gemma4LogMel",
+          "name": "gemma4_audio",
+          "type": "Gemma4Audio",
           "attrs": {
+            "type": "log_mel",
             "feature_size": 128,
             "sampling_rate": 16000,
             "frame_length_ms": 20.0,

test/pp_api_test/test_feature_extraction.cc

@@ -344,7 +344,7 @@ TEST(ExtractorTest, TestSplitSignalSegments) {
 
 TEST(ExtractorTest, TestGemma4AudioFeatureExtraction) {
   // Use existing test audio files to verify the Gemma 4 USM-style log-mel pipeline:
-  // AudioDecoder -> Gemma4LogMel
+  // AudioDecoder -> Gemma4Audio (type="log_mel")
   const char* audio_path[] = {"data/jfk.flac"};
   OrtxObjectPtr<OrtxRawAudios> raw_audios;
   extError_t err = OrtxLoadAudios(raw_audios.ToBeAssigned(), audio_path, 1);
@@ -446,4 +446,58 @@ TEST(ExtractorTest, TestGemma4AudioFeatureExtractionMultiFile) {
   ASSERT_EQ(err, kOrtxOK);
   ASSERT_EQ(mask_dims, 2ULL);
   ASSERT_EQ(mask_shape[0], 2);
-}
+}
+
+TEST(ExtractorTest, TestGemma4UnifiedAudioFrames) {
+  // gemma-4-12B "unified" (encoder-free) audio: raw 16 kHz waveform chunked
+  // into fixed 640-sample frames via the generic Gemma4Audio op with
+  // type="raw_frames".  Pipeline: AudioDecoder -> Gemma4Audio
+  const char* audio_path[] = {"data/jfk.flac"};
+  OrtxObjectPtr<OrtxRawAudios> raw_audios;
+  extError_t err = OrtxLoadAudios(raw_audios.ToBeAssigned(), audio_path, 1);
+  ASSERT_EQ(err, kOrtxOK);
+
+  OrtxObjectPtr<OrtxFeatureExtractor> feature_extractor(
+      OrtxCreateSpeechFeatureExtractor, "data/models/gemma-4-unified/audio_feature_extraction.json");
+  OrtxObjectPtr<OrtxTensorResult> result;
+  err = OrtxFeatureExtraction(feature_extractor.get(), raw_audios.get(), result.ToBeAssigned());
+  ASSERT_EQ(err, kOrtxOK);
+
+  // Output 0: raw waveform frames — float (batch, num_tokens, 640)
+  OrtxObjectPtr<OrtxTensor> tensor;
+  err = OrtxTensorResultGetAt(result.get(), 0, tensor.ToBeAssigned());
+  ASSERT_EQ(err, kOrtxOK);
+
+  const float* data{};
+  const int64_t* shape{};
+  size_t num_dims;
+  err = OrtxGetTensorData(tensor.get(), reinterpret_cast<const void**>(&data), &shape, &num_dims);
+  ASSERT_EQ(err, kOrtxOK);
+  ASSERT_EQ(num_dims, 3ULL);  // (batch, num_tokens, samples_per_token)
+  ASSERT_EQ(shape[0], 1);     // single audio
+  ASSERT_EQ(shape[2], 640);   // 640 raw samples per token
+  EXPECT_GT(shape[1], 0);     // at least one frame
+  const int64_t num_tokens = shape[1];
+
+  // All values finite and within the normalized PCM range.
+  for (int64_t i = 0; i < std::min<int64_t>(num_tokens * 640, 5000); ++i) {
+    ASSERT_TRUE(std::isfinite(data[i])) << "frame value at index " << i << " is not finite";
+    ASSERT_LE(std::abs(data[i]), 4.0f) << "frame value at index " << i << " out of range";
+  }
+
+  // Output 1: frame mask — bool (batch, num_tokens), all true for a single clip.
+  err = OrtxTensorResultGetAt(result.get(), 1, tensor.ToBeAssigned());
+  ASSERT_EQ(err, kOrtxOK);
+  const bool* mask_data{};
+  const int64_t* mask_shape{};
+  size_t mask_dims;
+  err = OrtxGetTensorData(tensor.get(), reinterpret_cast<const void**>(&mask_data), &mask_shape, &mask_dims);
+  ASSERT_EQ(err, kOrtxOK);
+  ASSERT_EQ(mask_dims, 2ULL);            // (batch, num_tokens)
+  ASSERT_EQ(mask_shape[0], 1);
+  ASSERT_EQ(mask_shape[1], num_tokens);  // same frame count as features
+  for (int64_t i = 0; i < num_tokens; ++i) {
+    EXPECT_TRUE(mask_data[i]) << "single-clip frame " << i << " should be valid";
+  }
+}
+