diff --git a/shared/api/gemma4_audio_features.hpp b/shared/api/gemma4_audio_features.hpp index 1f90ad57f..cf72770cc 100644 --- a/shared/api/gemma4_audio_features.hpp +++ b/shared/api/gemma4_audio_features.hpp @@ -18,10 +18,32 @@ namespace ort_extensions { +namespace gemma4_audio_detail { + +// Checked attribute extraction: returns kOrtxErrorInvalidArgument instead of +// throwing std::bad_variant_access when a config value has an unexpected type +// (e.g. a string where a number is required, or an int where a float is). +template +OrtxStatus GetTypedAttr(const VariantT& value, const char* op_name, const std::string& key, T& out) { + const T* ptr = std::get_if(&value); + if (ptr == nullptr) { + return {kOrtxErrorInvalidArgument, + std::string("[") + op_name + "]: attribute '" + key + "' has an unexpected value type"}; + } + out = *ptr; + return {}; +} + +} // namespace gemma4_audio_detail + // Gemma 4 audio feature extraction: USM-style log-mel spectrogram that matches // the HuggingFace Gemma4AudioFeatureExtractor exactly. // -// Pipeline: AudioDecoder -> Gemma4LogMel +// This is the log-mel implementation. It stays registered under its own +// "Gemma4LogMel" name (backward compatibility) and is also used internally by +// the generic Gemma4Audio op with type="log_mel". +// +// Pipeline: AudioDecoder -> Gemma4LogMel (or Gemma4Audio type="log_mel") // // Inputs: float (1, num_samples) — mono PCM at `sampling_rate` Hz // Outputs: float (num_frames, feature_size) — log-mel features @@ -267,33 +289,48 @@ class Gemma4LogMel { template OrtxStatus Init(const DictT& attrs) { + using gemma4_audio_detail::GetTypedAttr; + constexpr const char* kOp = "Gemma4LogMel"; for (const auto& [key, value] : attrs) { if (key == "feature_size") { - feature_size_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, feature_size_); !st.IsOk()) return st; } else if (key == "sampling_rate") { - sampling_rate_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, sampling_rate_); !st.IsOk()) return st; } else if (key == "frame_length_ms") { - frame_length_ms_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, frame_length_ms_); !st.IsOk()) return st; } else if (key == "hop_length_ms") { - hop_length_ms_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, hop_length_ms_); !st.IsOk()) return st; } else if (key == "min_frequency") { - min_frequency_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, min_frequency_); !st.IsOk()) return st; } else if (key == "max_frequency") { - max_frequency_ = std::get(value); + if (auto st = GetTypedAttr(value, kOp, key, max_frequency_); !st.IsOk()) return st; } else if (key == "preemphasis") { - preemphasis_ = static_cast(std::get(value)); + double tmp = 0.0; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + preemphasis_ = static_cast(tmp); } else if (key == "preemphasis_htk_flavor") { - preemphasis_htk_flavor_ = std::get(value) != 0; + int64_t tmp = 0; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + preemphasis_htk_flavor_ = tmp != 0; } else if (key == "fft_overdrive") { - fft_overdrive_ = std::get(value) != 0; + int64_t tmp = 0; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + fft_overdrive_ = tmp != 0; } else if (key == "mel_floor") { - mel_floor_ = static_cast(std::get(value)); + double tmp = 0.0; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + mel_floor_ = static_cast(tmp); } else if (key == "per_bin_mean") { - auto& v = std::get>(value); - per_bin_mean_.assign(v.begin(), v.end()); + std::vector tmp; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + per_bin_mean_.assign(tmp.begin(), tmp.end()); } else if (key == "per_bin_stddev") { - auto& v = std::get>(value); - per_bin_stddev_.assign(v.begin(), v.end()); + std::vector tmp; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + per_bin_stddev_.assign(tmp.begin(), tmp.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 +381,182 @@ class Gemma4LogMel { std::vector 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& pcm_input, + ortc::Tensor& frames_out, + ortc::Tensor& 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(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 + OrtxStatus Init(const DictT& attrs) { + using gemma4_audio_detail::GetTypedAttr; + constexpr const char* kOp = "Gemma4UnifiedAudioFrames"; + // Track the two aliases separately so conflicting values are rejected rather + // than silently taking whichever key appears last. + bool samples_set = false, feature_size_set = false; + int64_t samples_val = 0, feature_size_val = 0; + for (const auto& [key, value] : attrs) { + if (key == "audio_samples_per_token") { + if (auto st = GetTypedAttr(value, kOp, key, samples_val); !st.IsOk()) return st; + samples_set = true; + } else if (key == "feature_size") { + // ``feature_size`` is accepted as an alias: HF sets feature_size == + // audio_samples_per_token (both default to 640). + if (auto st = GetTypedAttr(value, kOp, key, feature_size_val); !st.IsOk()) return st; + feature_size_set = true; + } else if (key == "sampling_rate") { + if (auto st = GetTypedAttr(value, kOp, key, sampling_rate_); !st.IsOk()) return st; + } else if (key == "padding_value") { + double tmp = 0.0; + if (auto st = GetTypedAttr(value, kOp, key, tmp); !st.IsOk()) return st; + padding_value_ = static_cast(tmp); + } else if (key == "type") { + // Consumed by the Gemma4Audio dispatcher (selects this raw-frames path). + } else { + return {kOrtxErrorInvalidArgument, + "[Gemma4UnifiedAudioFrames]: unknown attribute '" + key + "'"}; + } + } + if (samples_set && feature_size_set && samples_val != feature_size_val) { + return {kOrtxErrorInvalidArgument, + "[Gemma4UnifiedAudioFrames]: conflicting 'audio_samples_per_token' (" + + std::to_string(samples_val) + ") and 'feature_size' (" + std::to_string(feature_size_val) + + "); they are aliases and must match"}; + } + if (samples_set) { + audio_samples_per_token_ = samples_val; + } else if (feature_size_set) { + audio_samples_per_token_ = feature_size_val; + } + 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& pcm_input, + ortc::Tensor& features_out, + ortc::Tensor& 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 + 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") { + std::string type; + if (auto st = gemma4_audio_detail::GetTypedAttr(value, "Gemma4Audio", key, type); !st.IsOk()) { + return st; + } + 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 diff --git a/shared/api/speech_extractor.cc b/shared/api/speech_extractor.cc index c32100ddd..ffd3c6960 100644 --- a/shared/api/speech_extractor.cc +++ b/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) {} diff --git a/test/data/models/gemma-4-unified/audio_feature_extraction.json b/test/data/models/gemma-4-unified/audio_feature_extraction.json new file mode 100644 index 000000000..db1de383e --- /dev/null +++ b/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 + } + } + } + ] + } +} diff --git a/test/data/models/gemma-4-unified/image_processor.json b/test/data/models/gemma-4-unified/image_processor.json new file mode 100644 index 000000000..ac9252e7b --- /dev/null +++ b/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 + } + } + } + ] + } +} diff --git a/test/data/models/gemma-4/audio_feature_extraction.json b/test/data/models/gemma-4/audio_feature_extraction.json index 699156dc4..b070f01cc 100644 --- a/test/data/models/gemma-4/audio_feature_extraction.json +++ b/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, diff --git a/test/pp_api_test/test_feature_extraction.cc b/test/pp_api_test/test_feature_extraction.cc index 7acb1497e..fb694b847 100644 --- a/test/pp_api_test/test_feature_extraction.cc +++ b/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 raw_audios; extError_t err = OrtxLoadAudios(raw_audios.ToBeAssigned(), audio_path, 1); @@ -446,4 +446,126 @@ TEST(ExtractorTest, TestGemma4AudioFeatureExtractionMultiFile) { ASSERT_EQ(err, kOrtxOK); ASSERT_EQ(mask_dims, 2ULL); ASSERT_EQ(mask_shape[0], 2); -} \ No newline at end of file +} + +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 raw_audios; + extError_t err = OrtxLoadAudios(raw_audios.ToBeAssigned(), audio_path, 1); + ASSERT_EQ(err, kOrtxOK); + + OrtxObjectPtr feature_extractor( + OrtxCreateSpeechFeatureExtractor, "data/models/gemma-4-unified/audio_feature_extraction.json"); + OrtxObjectPtr 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 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(&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]; + + // Raw waveform frames are the decoded PCM samples, which the AudioDecoder + // normalizes to [-1, 1]; a small epsilon covers float rounding at full scale. + for (int64_t i = 0; i < std::min(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]), 1.0001f) << "frame value at index " << i << " out of normalized PCM 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(&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"; + } +} + +TEST(ExtractorTest, TestGemma4UnifiedAudioFramesMultiFile) { + // Two clips of different lengths: verify batch stacking pads the shorter clip's + // frames and that the frame mask marks the padded tail invalid (false), while + // the real frames of each clip are valid (true). Locks in the unified batch + + // mask behavior, mirroring the log-mel multi-file coverage. + const char* audio_path[] = {"data/jfk.flac", "data/1272-141231-0002.wav"}; + OrtxObjectPtr raw_audios; + extError_t err = OrtxLoadAudios(raw_audios.ToBeAssigned(), audio_path, 2); + ASSERT_EQ(err, kOrtxOK); + + OrtxObjectPtr feature_extractor( + OrtxCreateSpeechFeatureExtractor, "data/models/gemma-4-unified/audio_feature_extraction.json"); + OrtxObjectPtr result; + err = OrtxFeatureExtraction(feature_extractor.get(), raw_audios.get(), result.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + + // Output 0: frames — float (2, max_tokens, 640) + OrtxObjectPtr 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(&data), &shape, &num_dims); + ASSERT_EQ(err, kOrtxOK); + ASSERT_EQ(num_dims, 3ULL); + ASSERT_EQ(shape[0], 2); // batch of 2 clips + ASSERT_EQ(shape[2], 640); // raw samples per token + const int64_t max_tokens = shape[1]; + + // Output 1: mask — bool (2, max_tokens) + 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(&mask_data), &mask_shape, &mask_dims); + ASSERT_EQ(err, kOrtxOK); + ASSERT_EQ(mask_dims, 2ULL); + ASSERT_EQ(mask_shape[0], 2); + ASSERT_EQ(mask_shape[1], max_tokens); + + // Each row's mask must be a contiguous true-prefix (real frames) followed by a + // false-suffix (batch padding). Count valid frames per clip. + int64_t valid_counts[2] = {0, 0}; + for (int64_t b = 0; b < 2; ++b) { + const bool* row = mask_data + b * max_tokens; + bool seen_false = false; + for (int64_t i = 0; i < max_tokens; ++i) { + if (row[i]) { + ASSERT_FALSE(seen_false) << "clip " << b << " mask must not have a true frame after padding"; + ++valid_counts[b]; + } else { + seen_false = true; + } + } + EXPECT_GT(valid_counts[b], 0) << "clip " << b << " should have at least one valid frame"; + } + + // The two clips have different lengths, so exactly one clip fills all max_tokens + // and the shorter clip has a padded (false) tail. + EXPECT_NE(valid_counts[0], valid_counts[1]) << "test clips should differ in length"; + EXPECT_EQ(std::max(valid_counts[0], valid_counts[1]), max_tokens); + const int64_t shorter = std::min(valid_counts[0], valid_counts[1]); + EXPECT_LT(shorter, max_tokens) << "shorter clip should be zero-padded in the batch"; +} + diff --git a/test/pp_api_test/test_processor.cc b/test/pp_api_test/test_processor.cc index fea0d19f9..52c4dd390 100644 --- a/test/pp_api_test/test_processor.cc +++ b/test/pp_api_test/test_processor.cc @@ -407,6 +407,134 @@ 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 raw_images{}; + extError_t err = OrtxLoadImages(raw_images.ToBeAssigned(), image_path, 1, nullptr); + ASSERT_EQ(err, kOrtxOK); + + // --- unified config: 48px merged patches --- + OrtxObjectPtr 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 result; + err = OrtxImagePreProcess(processor.get(), raw_images.get(), result.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + + OrtxObjectPtr 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(&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(&pos_data), &shape, &num_dims); + ASSERT_EQ(err, kOrtxOK); + // Guard the raw pos_data indexing below against an unexpected layout. + ASSERT_EQ(num_dims, 3ULL); // (batch, max_soft_tokens, 2) + ASSERT_EQ(shape[0], 1); + ASSERT_EQ(shape[1], kMaxSoftTokens); + ASSERT_EQ(shape[2], 2); + 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 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(&nst), &nst_shape, &nst_dims); + ASSERT_EQ(err, kOrtxOK); + ASSERT_EQ(nst_dims, 2ULL); // (batch, 1) + ASSERT_EQ(nst_shape[0], 1); + EXPECT_EQ(nst[0], 260) << "merged soft-token count for australia.jpg (HF reference)"; + const int64_t num_merged = nst[0]; + // Value must be within the padded grid before it indexes pos_data below. + ASSERT_GT(num_merged, 0); + ASSERT_LE(num_merged, kMaxSoftTokens); + // 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 out now: the tensor's backing buffer is owned by the + // first result, and reusing `tensor`/running the teacher config below would + // rebind pv_data, so snapshot it into an independent buffer first. + std::vector merged_patch0(pv_data, pv_data + kMergedPatchDim); + + // --- standard config: 16px teacher patches --- + OrtxObjectPtr teacher_proc; + err = OrtxCreateProcessor(teacher_proc.ToBeAssigned(), "data/models/gemma-4/image_processor.json"); + ASSERT_EQ(err, kOrtxOK); + OrtxObjectPtr teacher_result; + err = OrtxImagePreProcess(teacher_proc.get(), raw_images.get(), teacher_result.ToBeAssigned()); + ASSERT_EQ(err, kOrtxOK); + OrtxObjectPtr 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(&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 + // Guard the raw indexing below against an unexpected teacher tensor layout. + ASSERT_EQ(teacher_dims, 3ULL); // (batch, num_patches, patch_dim) + ASSERT_GE(teacher_shape[1], 1); // at least patch 0 + ASSERT_EQ(teacher_shape[2], kTeacherPatchDim); + 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; + } + } + } +} + 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"};