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

shared/api/gemma4_audio_features.hpp

@@ -344,4 +344,79 @@ 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  ->  Gemma4UnifiedAudioFrames
+//
+// Inputs:   float (1, num_samples) — mono PCM at `sampling_rate` Hz
+// Outputs:  float (num_tokens, audio_samples_per_token) — raw waveform frames
+//
+// The frame-level mask is intentionally not emitted here: for the single-clip
+// inference path the downstream processor fills an all-true mask, and the
+// batch framework zero-pads ragged clips when stacking.
+class Gemma4UnifiedAudioFrames {
+ public:
+  Gemma4UnifiedAudioFrames() = default;
+
+  OrtxStatus Compute(const ortc::Tensor<float>& pcm_input,
+                     ortc::Tensor<float>& frames_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});
+    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);
+    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 {
+        return {kOrtxErrorInvalidArgument,
+                "[Gemma4UnifiedAudioFrames]: unknown attribute '" + key + "'"};
+      }
+    }
+    if (audio_samples_per_token_ <= 0) {
+      return {kOrtxErrorInvalidArgument,
+              "[Gemma4UnifiedAudioFrames]: audio_samples_per_token must be positive"};
+    }
+    return {};
+  }
+
+ private:
+  int64_t audio_samples_per_token_ = 640;  // 640 samples = 40 ms @ 16 kHz
+  int64_t sampling_rate_ = 16000;
+  float padding_value_ = 0.0f;
+};
+
 }  // 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); }},
+    {"Gemma4UnifiedAudioFrames", []() { return CreateKernelInstance(&Gemma4UnifiedAudioFrames::Compute); }}};
 
 SpeechFeatureExtractor::SpeechFeatureExtractor() : OrtxObjectImpl(extObjectKind_t::kOrtxKindFeatureExtractor) {}
 

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

@@ -0,0 +1,23 @@
+{
+  "feature_extraction": {
+    "sequence": [
+      {
+        "operation": {
+          "name": "audio_decoder",
+          "type": "AudioDecoder"
+        }
+      },
+      {
+        "operation": {
+          "name": "gemma4_unified_audio_frames",
+          "type": "Gemma4UnifiedAudioFrames",
+          "attrs": {
+            "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/pp_api_test/test_feature_extraction.cc

@@ -446,4 +446,40 @@ 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.  Pipeline: AudioDecoder -> Gemma4UnifiedAudioFrames
+  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
+
+  // All values finite and within the normalized PCM range.
+  for (int64_t i = 0; i < std::min<int64_t>(shape[1] * 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";
+  }
+}

test/pp_api_test/test_processor.cc

@@ -407,6 +407,120 @@ TEST(ProcessorTest, TestGemma4ImageProcessing) {
   EXPECT_EQ(nst_peek[0], 260) << "num_soft_tokens should be 260 for australia.jpg (HF reference)";
 }
 
+TEST(ProcessorTest, TestGemma4UnifiedImageProcessing) {
+  // gemma-4-12B "unified" (encoder-free) vision preprocessing.
+  //
+  // The unified model consumes 48px MERGED patches (patch_dim = 48*48*3 = 6912)
+  // directly, with no SigLIP encoder to pool 3x3 teacher patches. HuggingFace
+  // produces these via (16px patchify -> 3x3 patches_merge). That is provably
+  // identical to a direct 48px patchify, so the unified contract is generated by
+  // reusing Gemma4ImageTransform with patch_size=48, pooling_kernel_size=1.
+  //
+  // This test verifies (a) the 6912-dim contract and (b) that the reused op's
+  // top-left 48px patch matches the 16px teacher patch from the standard config.
+  const char* image_path[] = {"data/processor/australia.jpg"};
+
+  OrtxObjectPtr<OrtxRawImages> raw_images{};
+  extError_t err = OrtxLoadImages(raw_images.ToBeAssigned(), image_path, 1, nullptr);
+  ASSERT_EQ(err, kOrtxOK);
+
+  // --- unified config: 48px merged patches ---
+  OrtxObjectPtr<OrtxProcessor> processor;
+  err = OrtxCreateProcessor(processor.ToBeAssigned(), "data/models/gemma-4-unified/image_processor.json");
+  if (err != kOrtxOK) {
+    std::cout << "Error: " << OrtxGetLastErrorMessage() << std::endl;
+  }
+  ASSERT_EQ(err, kOrtxOK);
+
+  OrtxObjectPtr<OrtxTensorResult> result;
+  err = OrtxImagePreProcess(processor.get(), raw_images.get(), result.ToBeAssigned());
+  ASSERT_EQ(err, kOrtxOK);
+
+  OrtxObjectPtr<OrtxTensor> tensor;
+  err = OrtxTensorResultGetAt(result.get(), 0, tensor.ToBeAssigned());
+  ASSERT_EQ(err, kOrtxOK);
+
+  const float* pv_data{};
+  const int64_t* shape{};
+  size_t num_dims;
+  err = OrtxGetTensorData(tensor.get(), reinterpret_cast<const void**>(&pv_data), &shape, &num_dims);
+  ASSERT_EQ(err, kOrtxOK);
+  ASSERT_EQ(num_dims, 3ULL);
+  ASSERT_EQ(shape[0], 1);
+  constexpr int64_t kMaxSoftTokens = 280;
+  constexpr int64_t kMergedPatchDim = 48 * 48 * 3;  // 6912
+  ASSERT_EQ(shape[1], kMaxSoftTokens);
+  ASSERT_EQ(shape[2], kMergedPatchDim);
+
+  // position_ids — merged 48-grid coordinates.
+  err = OrtxTensorResultGetAt(result.get(), 1, tensor.ToBeAssigned());
+  ASSERT_EQ(err, kOrtxOK);
+  const int64_t* pos_data{};
+  err = OrtxGetTensorData(tensor.get(), reinterpret_cast<const void**>(&pos_data), &shape, &num_dims);
+  ASSERT_EQ(err, kOrtxOK);
+  EXPECT_EQ(pos_data[0], 0);  // patch 0 x
+  EXPECT_EQ(pos_data[1], 0);  // patch 0 y
+
+  // num_soft_tokens: merged count = teacher-grid / 9. For australia.jpg the
+  // teacher grid is 60x39 -> merged 20x13 = 260 (same value as the standard
+  // config, which reports it before pooling).
+  OrtxObjectPtr<OrtxTensor> nst_tensor;
+  err = OrtxTensorResultGetAt(result.get(), 2, nst_tensor.ToBeAssigned());
+  ASSERT_EQ(err, kOrtxOK);
+  const int64_t* nst{};
+  const int64_t* nst_shape{};
+  size_t nst_dims{};
+  err = OrtxGetTensorData(nst_tensor.get(), reinterpret_cast<const void**>(&nst), &nst_shape, &nst_dims);
+  ASSERT_EQ(err, kOrtxOK);
+  EXPECT_EQ(nst[0], 260) << "merged soft-token count for australia.jpg (HF reference)";
+  const int64_t num_merged = nst[0];
+  // Last real merged patch position: (20-1, 13-1) = (19, 12).
+  EXPECT_EQ(pos_data[(num_merged - 1) * 2], 19);
+  EXPECT_EQ(pos_data[(num_merged - 1) * 2 + 1], 12);
+  // Padding beyond real patches is (-1, -1).
+  for (int64_t i = num_merged; i < kMaxSoftTokens; ++i) {
+    EXPECT_EQ(pos_data[i * 2], -1);
+    EXPECT_EQ(pos_data[i * 2 + 1], -1);
+  }
+
+  // Copy merged patch 0 before running the second (teacher) config, which
+  // invalidates pv_data.
+  std::vector<float> merged_patch0(pv_data, pv_data + kMergedPatchDim);
+
+  // --- standard config: 16px teacher patches ---
+  OrtxObjectPtr<OrtxProcessor> teacher_proc;
+  err = OrtxCreateProcessor(teacher_proc.ToBeAssigned(), "data/models/gemma-4/image_processor.json");
+  ASSERT_EQ(err, kOrtxOK);
+  OrtxObjectPtr<OrtxTensorResult> teacher_result;
+  err = OrtxImagePreProcess(teacher_proc.get(), raw_images.get(), teacher_result.ToBeAssigned());
+  ASSERT_EQ(err, kOrtxOK);
+  OrtxObjectPtr<OrtxTensor> teacher_tensor;
+  err = OrtxTensorResultGetAt(teacher_result.get(), 0, teacher_tensor.ToBeAssigned());
+  ASSERT_EQ(err, kOrtxOK);
+  const float* teacher_pv{};
+  const int64_t* teacher_shape{};
+  size_t teacher_dims{};
+  err = OrtxGetTensorData(teacher_tensor.get(), reinterpret_cast<const void**>(&teacher_pv), &teacher_shape,
+                          &teacher_dims);
+  ASSERT_EQ(err, kOrtxOK);
+
+  // Merged patch 0 is the top-left 48x48 image block, HWC. Its top-left 16x16
+  // sub-block must equal teacher patch 0 (also the top-left 16x16 block, HWC).
+  //   merged flat index: ((r*48 + col)*3 + c)   for r,col in [0,16)
+  //   teacher flat index: ((r*16 + col)*3 + c)
+  constexpr int64_t kTeacherPatchDim = 16 * 16 * 3;  // 768
+  for (int64_t r = 0; r < 16; ++r) {
+    for (int64_t col = 0; col < 16; ++col) {
+      for (int64_t c = 0; c < 3; ++c) {
+        float m = merged_patch0[(r * 48 + col) * 3 + c];
+        float t = teacher_pv[(r * 16 + col) * 3 + c];
+        ASSERT_NEAR(m, t, 1e-6f) << "merged vs teacher mismatch at r=" << r << " col=" << col << " c=" << c;
+      }
+    }
+  }
+  (void)kTeacherPatchDim;
+}
+
 TEST(ProcessorTest, TestGemma4ImageProcessingMultiImage) {
   // Verify batched processing works with multiple images of different sizes.
   const char* image_paths[] = {"data/processor/standard_s.jpg", "data/processor/australia.jpg"};