Add gemma-4-12B unified (encoder-free) preprocessing support by justinchuby · Pull Request #1091 · microsoft/onnxruntime-extensions

justinchuby · 2026-07-09T17:58:01Z

What

Adds the onnxruntime-extensions preprocessing needed for the gemma-4-12B "unified" model, which is encoder-free: it consumes raw merged pixel patches and raw waveform frames directly, rather than the SigLIP image / USM log-mel audio contract of the standard gemma-4 (E2B/E4B) processor.

Today onnxruntime-genai can only run the unified model by preprocessing inputs with the HuggingFace AutoProcessor and feeding tensors via Generator.set_inputs, because there is no native transform matching the unified contract. This PR provides the native building blocks so genai can preprocess without HuggingFace.

Changes

Audio — new op Gemma4UnifiedAudioFrames (shared/api/gemma4_audio_features.hpp, registered in speech_extractor.cc).
Chunks a decoded 16 kHz waveform into fixed audio_samples_per_token (default 640) frames, reproducing HF Gemma4UnifiedAudioFeatureExtractor._extract_waveform_features exactly: zero-pad to a whole number of frames, reshape to (num_tokens, 640). No mel / windowing / normalization.
Pipeline: AudioDecoder -> Gemma4UnifiedAudioFrames.

Image — no new op required.
HF produces the unified 48px merged patches (patch_dim = 48*48*3 = 6912) via (16px patchify → 3×3 patches_merge). That is provably identical to a direct 48px patchify (verified: both pixel values and position IDs match, and the aspect-ratio resize budget is identical). So the existing Gemma4ImageTransform yields the unified contract when configured with patch_size=48, pooling_kernel_size=1, max_soft_tokens=280. A new test locks this equivalence in.

Tests / data

test/data/models/gemma-4-unified/audio_feature_extraction.json and image_processor.json.
ExtractorTest.TestGemma4UnifiedAudioFrames — verifies the (num_tokens, 640) framing contract.
ProcessorTest.TestGemma4UnifiedImageProcessing — verifies the 6912-dim contract, merged position IDs / padding, and cross-checks that the reused op's top-left 48px patch matches the standard config's 16px teacher patch on the same image.

Testing

Built with ./build.sh -DOCOS_ENABLE_C_API=ON -DOCOS_ENABLE_CTEST=ON. All 9 *Gemma4* tests pass (2 new + 7 existing, no regressions).

The gemma-4-12B "unified" model is encoder-free: it consumes raw merged pixel patches and raw waveform frames directly, not the SigLIP/log-mel contract of the standard gemma-4 processor. Add the two ort-extensions pieces needed so onnxruntime-genai can preprocess natively (no external HuggingFace processor): * Audio: new Gemma4UnifiedAudioFrames feature op that chunks a decoded 16 kHz waveform into fixed audio_samples_per_token (640) frames, reproducing Gemma4UnifiedAudioFeatureExtractor exactly (pad to a whole number of frames, reshape to (num_tokens, 640)). * Image: no new op required. The unified 48px merged patch (patch_dim 6912) produced by HuggingFace via (16px patchify -> 3x3 patches_merge) is provably identical to a direct 48px patchify, so the existing Gemma4ImageTransform yields the unified contract when configured with patch_size=48, pooling_kernel_size=1. A test locks this equivalence in. Adds test configs under test/data/models/gemma-4-unified/ and gtests that verify the 640-sample audio framing and the 6912-dim image contract (cross-checked against the standard config's teacher patches). Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Copilot

Pull request overview

Adds native onnxruntime-extensions preprocessing components and tests for the gemma-4-12B “unified” (encoder-free) model contract, enabling image and audio inputs to be preprocessed without relying on HuggingFace AutoProcessor.

Changes:

Registers a new speech feature extraction op Gemma4UnifiedAudioFrames to chunk raw 16 kHz PCM into fixed 640-sample frames.
Adds unified vision processor configuration that reuses Gemma4ImageTransform with patch_size=48 to emit 6912-dim merged patches.
Introduces new processor / extractor tests and model JSON configs under test/data/models/gemma-4-unified/.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
`shared/api/gemma4_audio_features.hpp`	Adds the `Gemma4UnifiedAudioFrames` kernel implementation and attribute parsing.
`shared/api/speech_extractor.cc`	Registers the new `Gemma4UnifiedAudioFrames` op in the speech extractor kernel registry.
`test/pp_api_test/test_feature_extraction.cc`	Adds a framing-contract test for unified audio output shape / sanity checks.
`test/pp_api_test/test_processor.cc`	Adds a unified vision preprocessing test validating 48px merged-patch contract and equivalence checks.
`test/data/models/gemma-4-unified/audio_feature_extraction.json`	Adds a unified audio extraction pipeline config (AudioDecoder → UnifiedAudioFrames).
`test/data/models/gemma-4-unified/image_processor.json`	Adds a unified image processor config using `Gemma4ImageTransform` with 48px patchify.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

* Gemma4UnifiedAudioFrames: document that the op frames by sample count and does not resample (so sampling_rate is informational, matching the decoder's expected output rate), and reject non-positive sampling_rate. * Unified image test: assert the teacher tensor's rank/shape (using kTeacherPatchDim) before indexing to avoid a potential OOB read, and reword the pv_data snapshot comment to describe tensor-buffer rebinding rather than invalidation. Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

sayanshaw24 · 2026-07-10T00:51:59Z

    {"Phi4AudioEmbed", []() { return CreateKernelInstance(&Phi4AudioEmbed::Compute); }},
-    {"Gemma4LogMel", []() { return CreateKernelInstance(&Gemma4LogMel::Compute); }}};
+    {"Gemma4LogMel", []() { return CreateKernelInstance(&Gemma4LogMel::Compute); }},
+    {"Gemma4UnifiedAudioFrames", []() { return CreateKernelInstance(&Gemma4UnifiedAudioFrames::Compute); }}};


do you think we can have a more generic audio processor for gemma? so say instead of having both Gemma4LogMel and Gemma4UnifiedAudioFrames kernels, we maybe have a Gemma4Audio kernel (or similar), that can take JSON options in order to determine which pass to use?

for instance, with the Qwen-2.5VL image processor, instead of adding a new op for Smart Resize, which the model uses, we added smart resize parameters to finetune the pre-existing Resize op.

Great suggestion — done in 2f42800. Consolidated into a single Gemma4Audio op that selects the front-end via a type attribute:

type="log_mel" (default) → 128-dim USM log-mel spectrogram (E2B/E4B)

type="raw_frames" → raw 640-sample waveform frames (12B unified)

Both branches now share the same (features: float, mask: bool) output signature (the raw-frames path emits an all-true frame mask, which also matches HF's Gemma4UnifiedAudioFeatureExtractor returning input_features + input_features_mask). Gemma4Audio internally delegates to the unchanged Gemma4LogMel DSP and the raw-frames impl, so there's no duplicated logic.

One back-compat note: Gemma4LogMel already exists on main and is referenced by shipped gemma-4 configs, so I kept it registered as a thin alias (identical code path) rather than removing it — existing/deployed configs keep loading. The repo's own gemma-4 and gemma-4-unified test configs now both use Gemma4Audio. All 9 Gemma4 tests pass (log-mel via type=log_mel, raw frames via type=raw_frames incl. mask assertions).

I also propagated the op rename to the companion onnxruntime-genai (microsoft/onnxruntime-genai#2286) and mobius PRs. Let me know if you'd prefer I drop the Gemma4LogMel alias entirely (would be a breaking change for existing gemma-4 configs) or rename the attribute (type → e.g. mode).

Per review feedback, instead of shipping a second kernel (Gemma4UnifiedAudioFrames) alongside Gemma4LogMel, expose one generic Gemma4Audio op that selects the front-end via a 'type' attribute: type="log_mel" (default) -> 128-dim USM log-mel spectrogram (E2B/E4B) type="raw_frames" -> raw 640-sample waveform frames (12B unified) Both branches now share the (features: float, mask: bool) output signature; the raw-frames path emits an all-true frame mask, matching HuggingFace Gemma4UnifiedAudioFeatureExtractor (input_features + input_features_mask). Gemma4Audio internally delegates to the unchanged Gemma4LogMel DSP and the raw-frames implementation. Gemma4LogMel stays registered as a backward-compat alias so already-deployed gemma-4 configs keep loading. The repo's own gemma-4 and gemma-4-unified configs now use Gemma4Audio. All 9 Gemma4 tests pass (log-mel via Gemma4Audio type=log_mel, raw frames via type=raw_frames incl. mask assertions). Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

sayanshaw24 · 2026-07-10T17:32:48Z

+    // 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);


nice consolidation into Gemma4Audio! one robustness concern: std::get here (and similar std::get calls in the audio sub-paths) can throw bad_variant_access on malformed JSON, which bypasses OrtxStatus error handling. could we switch to a checked extraction path and return kOrtxErrorInvalidArgument instead?

Fixed in 920f54b. Added a checked GetTypedAttr helper (std::get_if) and routed all attribute reads in Gemma4LogMel, Gemma4UnifiedAudioFrames, and Gemma4Audio through it, so a wrong-typed config value now returns kOrtxErrorInvalidArgument instead of throwing std::bad_variant_access past the OrtxStatus boundary.

sayanshaw24 · 2026-07-10T17:34:17Z

+  template <typename DictT>
+  OrtxStatus Init(const DictT& attrs) {
+    for (const auto& [key, value] : attrs) {
+      if (key == "audio_samples_per_token" || key == "feature_size") {


for audio_samples_per_token and feature_size aliasing: if both keys are provided with different values, the current behavior silently takes whichever appears last. can we explicitly detect and reject conflicting values? silent override here could hide config mistakes and produce unexpected frame shapes.

Fixed in 920f54b. Gemma4UnifiedAudioFrames::Init now tracks audio_samples_per_token and feature_size separately and, if both are provided with different values, returns kOrtxErrorInvalidArgument ("they are aliases and must match") instead of silently taking the last one.

sayanshaw24 · 2026-07-10T17:36:37Z

-}
+}
+
+TEST(ExtractorTest, TestGemma4UnifiedAudioFrames) {


could we add a multi-file raw_frames test (e.g. two clips with different lengths) for Gemma4Audio with type=raw_frames? this would lock in batch stacking + mask behavior for the unified path, similar to the existing multi-file coverage on the log-mel side.

Added TestGemma4UnifiedAudioFramesMultiFile in 920f54b: two clips of different lengths (jfk.flac + 1272-141231-0002.wav). It asserts batch stacking to (2, max_tokens, 640), and that each row's mask is a contiguous true-prefix / false-suffix — i.e. the shorter clip's padded tail is marked invalid while its real frames stay valid. This locks in the batch + mask behavior for the raw-frames path, mirroring the log-mel multi-file coverage.

Follow-up to the Gemma4Audio consolidation: * Use checked std::get_if via a GetTypedAttr helper across Gemma4LogMel, Gemma4UnifiedAudioFrames and Gemma4Audio Init so a malformed config value (wrong JSON type) returns kOrtxErrorInvalidArgument instead of throwing std::bad_variant_access past the OrtxStatus boundary. * Reject conflicting 'audio_samples_per_token' vs 'feature_size' aliases when both are provided with different values, rather than silently taking the last. * Add TestGemma4UnifiedAudioFramesMultiFile: two clips of different lengths, asserting batch stacking pads the shorter clip and the frame mask marks the padded tail invalid (true-prefix / false-suffix per row). All Gemma4 tests pass. Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Copilot

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.

justinchuby · 2026-07-10T23:20:35Z

+  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


Fixed in ff3d01b. Added rank/shape asserts on the position_ids tensor (rank 3, shape (1, max_soft_tokens, 2)) before any pos_data indexing.

justinchuby · 2026-07-10T23:20:36Z

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


Fixed in ff3d01b. The test now asserts the num_soft_tokens tensor rank/shape and that the value is within [1, max_soft_tokens] before it's used to index pos_data.

justinchuby · 2026-07-10T23:20:38Z

+  // 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";
+  }


Fixed in ff3d01b. Tightened the bound to the actual normalized PCM range: decoded 16-bit flac/wav samples are in [-1, 1], so the assertion is now <= 1.0001 (small epsilon for full-scale float rounding) and the comment matches.

justinchuby · 2026-07-10T23:20:39Z

 // the HuggingFace Gemma4AudioFeatureExtractor exactly.
 //
-// Pipeline:  AudioDecoder  ->  Gemma4LogMel
+// Pipeline:  AudioDecoder  ->  Gemma4Audio (type="log_mel")


Fixed in ff3d01b. Reworded the doc block to say this is the log-mel implementation, kept registered under its own Gemma4LogMel name for backward compatibility and reused internally by Gemma4Audio with type="log_mel"; the pipeline line now reads AudioDecoder -> Gemma4LogMel (or Gemma4Audio type="log_mel").

* Unified image test: assert position_ids rank/shape and that num_soft_tokens is within [1, max_soft_tokens] before indexing pos_data, so a changed output layout fails cleanly instead of reading out of bounds. * Unified raw-frames test: tighten the amplitude bound to the normalized PCM range (<= 1.0001, epsilon for full-scale rounding) so it matches the comment. * Gemma4LogMel doc block: clarify it is the log-mel implementation, registered under its own name for back-compat and reused by Gemma4Audio type="log_mel" (was ambiguously labeled as the Gemma4Audio pipeline). All 10 Gemma4 tests pass. Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Copilot AI review requested due to automatic review settings July 9, 2026 17:58

justinchuby requested a review from a team as a code owner July 9, 2026 17:58

Copilot started reviewing on behalf of justinchuby July 9, 2026 17:58 View session

Copilot AI reviewed Jul 9, 2026

View reviewed changes

Comment thread shared/api/gemma4_audio_features.hpp

Comment thread test/pp_api_test/test_processor.cc

Comment thread test/pp_api_test/test_processor.cc Outdated

justinchuby mentioned this pull request Jul 9, 2026

Support gemma-4-12B unified (encoder-free) multimodal model microsoft/onnxruntime-genai#2286

Open

sayanshaw24 reviewed Jul 10, 2026

View reviewed changes

justinchuby requested a review from Copilot July 10, 2026 23:04

Copilot started reviewing on behalf of justinchuby July 10, 2026 23:04 View session

Copilot AI reviewed Jul 10, 2026

View reviewed changes

Uh oh!

Conversation

justinchuby commented Jul 9, 2026

What

Changes

Testing

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants