ort_genai: native processor configs for gemma-4-12B unified

src/mobius/integrations/ort_genai/auto_export.py

@@ -75,11 +75,12 @@
     "gemma2": "gemma",
     "gemma4": "gemma4",
     "gemma4_text": "gemma4_text",
-    # gemma-4-12B "unified" (encoder-free) variant reuses the gemma4 ORT GenAI
-    # pipelines: the multimodal package (decoder taking inputs_embeds + vision
-    # embedder + embedding fusion) maps to "gemma4"; the standalone text
-    # backbone maps to "gemma4_text".
-    "gemma4_unified": "gemma4",
+    # gemma-4-12B "unified" (encoder-free) variant. The multimodal package
+    # (decoder taking inputs_embeds + encoder-free vision/audio embedders +
+    # embedding fusion) maps to the dedicated "gemma4_unified" ORT GenAI type,
+    # which drives the unified processor (48px merged patches / raw 640-sample
+    # audio frames). The standalone text backbone maps to "gemma4_text".
+    "gemma4_unified": "gemma4_unified",
     "gemma4_unified_text": "gemma4_text",
     "mistral": "mistral",
     "mistral3": "mistral3",
@@ -99,12 +100,12 @@
 )
 # Encoder-free gemma-4-12B "unified" variants. Their image/audio inputs are raw
 # merged pixel patches (48px, 6912-dim) / raw waveform frames (640-dim), NOT the
-# SigLIP 16px / 128-dim log-mel contract that the ort-extensions
-# ``Gemma4ImageTransform`` / ``Gemma4LogMel`` ops implement. There is no
-# genai-native transform for the unified contract, so we deliberately do NOT
-# emit image_processor.json / audio_processor.json for these models — callers
-# must preprocess with the HuggingFace processor and feed tensors via
-# ``Generator.set_inputs`` (see examples/gemma4_unified_ort_genai.py).
+# SigLIP 16px / 128-dim log-mel contract of the standard gemma4 (E2B/E4B) model.
+# These are produced natively by ort-extensions: ``Gemma4ImageTransform`` with
+# patch_size=48 / pooling_kernel_size=1 (the 48px merged patch is identical to a
+# direct 48px patchify) and the ``Gemma4UnifiedAudioFrames`` raw-framing op. The
+# genai ``gemma4_unified`` processor consumes both (see the companion
+# onnxruntime-genai / onnxruntime-extensions support).
 _GEMMA4_UNIFIED_MODEL_TYPES = frozenset({"gemma4_unified", "gemma4_unified_text"})
 _PIXTRAL_MODEL_TYPES = frozenset({"mistral3"})
 _QWEN_VL_MODEL_TYPES = frozenset(
@@ -427,9 +428,10 @@ def _write_vision_processor_config(
 
     - **Gemma4** (``gemma4``, ``gemma4_text``): Writes ``image_processor.json``
       with a ``DecodeImage → Gemma4ImageTransform`` pipeline.
-    - **Gemma4 unified** (``gemma4_unified*``): Returns ``None`` — the
-      encoder-free model has no matching ort-extensions transform; callers feed
-      HF-preprocessed pixel_values via ``Generator.set_inputs``.
+    - **Gemma4 unified** (``gemma4_unified*``): Writes ``image_processor.json``
+      with a ``DecodeImage → Gemma4ImageTransform`` pipeline configured for the
+      encoder-free 48px merged-patch contract (patch_size=48,
+      pooling_kernel_size=1, patch_dim=6912).
     - **Pixtral / Mistral3**: Writes ``processor_config.json`` with a 7-step
       pipeline (DecodeImage → ConvertRGB → Resize → Rescale → Normalize →
       Permute3D → PixtralImageSizes).
@@ -443,21 +445,52 @@ def _write_vision_processor_config(
         return None
 
     model_type = getattr(config, "model_type", "")
-    if model_type in _GEMMA4_UNIFIED_MODEL_TYPES:
-        # Encoder-free unified model: no ort-extensions transform matches its
-        # raw merged-patch contract. Emit no image_processor.json; callers feed
-        # HF-preprocessed pixel_values via Generator.set_inputs.
-        logger.info(
-            "Skipping image_processor.json for encoder-free %s "
-            "(no native ort-extensions transform; use HF processor + set_inputs)",
-            model_type,
-        )
-        return None
-
     vision_model_type = getattr(vision, "model_type", None)
     is_pixtral = vision_model_type == "pixtral" or model_type in _PIXTRAL_MODEL_TYPES
 
-    if model_type in _GEMMA4_MODEL_TYPES:
+    if model_type in _GEMMA4_UNIFIED_MODEL_TYPES:
+        # Encoder-free unified model: it consumes 48px MERGED patches
+        # (patch_dim = 48*48*3 = 6912) directly, with no SigLIP tower to pool
+        # 3x3 teacher patches. HuggingFace produces these via
+        # (16px patchify -> 3x3 patches_merge), which is provably identical to a
+        # direct 48px patchify. So we reuse the same ort-extensions
+        # ``Gemma4ImageTransform`` op with the merged geometry: patch_size =
+        # patch_size * pooling_kernel_size (48) and pooling_kernel_size = 1.
+        max_soft_tokens = (
+            getattr(vision, "mm_tokens_per_image", None)
+            or getattr(config, "mm_tokens_per_image", None)
+            or 280
+        )
+        patch_size = getattr(vision, "patch_size", None) or 16
+        pooling = getattr(vision, "pooling_kernel_size", None) or 3
+        model_patch_size = patch_size * pooling  # 16 * 3 = 48
+        processor_config: dict[str, Any] = {
+            "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": model_patch_size,
+                                "max_soft_tokens": max_soft_tokens,
+                                "pooling_kernel_size": 1,
+                            },
+                        }
+                    },
+                ],
+            }
+        }
+        path = os.path.join(output_dir, "image_processor.json")
+    elif model_type in _GEMMA4_MODEL_TYPES:
         # Gemma4 needs an onnxruntime-extensions format processor config
         # with a transforms pipeline (DecodeImage -> Gemma4ImageTransform).
         max_soft_tokens = (
@@ -467,7 +500,7 @@ def _write_vision_processor_config(
         )
         patch_size = getattr(vision, "patch_size", None) or 16
         pooling_kernel_size = getattr(vision, "pooling_kernel_size", None) or 3
-        processor_config: dict[str, Any] = {
+        processor_config = {
             "processor": {
                 "name": "gemma_4_image_processing",
                 "transforms": [
@@ -631,18 +664,40 @@ def _write_audio_processor_config(
     model_type = getattr(config, "model_type", "")
 
     if model_type in _GEMMA4_UNIFIED_MODEL_TYPES:
-        # Encoder-free unified model: raw 640-dim waveform frames, not the
-        # 128-dim log-mel Gemma4LogMel contract. Emit no audio_processor.json;
-        # callers feed HF-preprocessed input_features via Generator.set_inputs.
-        logger.info(
-            "Skipping audio_processor.json for encoder-free %s "
-            "(no native ort-extensions transform; use HF processor + set_inputs)",
-            model_type,
-        )
-        return None
-
-    if model_type in _GEMMA4_MODEL_TYPES:
-        # Gemma4 USM-style 128-dim log-mel spectrogram.
+        # Encoder-free unified model: each audio soft token is a raw chunk of the
+        # 16 kHz waveform (audio_samples_per_token = audio_embed_dim, 640), not a
+        # 128-dim log-mel frame. Reproduced natively by the ort-extensions
+        # ``Gemma4Audio`` op with ``type="raw_frames"`` (pad to a whole number of
+        # frames, reshape to (num_tokens, 640)).
+        samples_per_token = getattr(audio, "hidden_size", None) or 640
+        processor = {
+            "feature_extraction": {
+                "sequence": [
+                    {
+                        "operation": {
+                            "name": "audio_decoder",
+                            "type": "AudioDecoder",
+                        }
+                    },
+                    {
+                        "operation": {
+                            "name": "gemma4_audio",
+                            "type": "Gemma4Audio",
+                            "attrs": {
+                                "type": "raw_frames",
+                                "audio_samples_per_token": samples_per_token,
+                                "sampling_rate": 16000,
+                                "padding_value": 0.0,
+                            },
+                        }
+                    },
+                ]
+            }
+        }
+        proc_filename = "audio_feature_extraction.json"
+    elif model_type in _GEMMA4_MODEL_TYPES:
+        # Gemma4 USM-style 128-dim log-mel spectrogram via the ort-extensions
+        # ``Gemma4Audio`` op with ``type="log_mel"``.
         # OrtxCreateSpeechFeatureExtractor requires the feature_extraction.sequence format.
         processor = {
             "feature_extraction": {
@@ -655,9 +710,10 @@ def _write_audio_processor_config(
                     },
                     {
                         "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,

src/mobius/integrations/ort_genai/auto_export_test.py

@@ -89,10 +89,10 @@ def test_phi4mm_model_types(self):
         assert _resolve_ort_genai_model_type("phi") == "phi"
 
     def test_gemma4_unified_model_types(self):
-        # The gemma-4-12B unified checkpoint (model_type "gemma4_unified")
-        # reuses the multimodal "gemma4" ORT GenAI pipeline; its standalone
-        # text decoder ("gemma4_unified_text") maps to "gemma4_text".
-        assert _resolve_ort_genai_model_type("gemma4_unified") == "gemma4"
+        # The gemma-4-12B unified checkpoint (model_type "gemma4_unified") maps
+        # to the dedicated "gemma4_unified" ORT GenAI multimodal pipeline; its
+        # standalone text decoder ("gemma4_unified_text") maps to "gemma4_text".
+        assert _resolve_ort_genai_model_type("gemma4_unified") == "gemma4_unified"
         assert _resolve_ort_genai_model_type("gemma4_unified_text") == "gemma4_text"
         # Released gemma4 mappings remain unchanged.
         assert _resolve_ort_genai_model_type("gemma4") == "gemma4"
@@ -115,12 +115,12 @@ def test_decoder_only_prefers_config_type(self):
     def test_multimodal_keeps_hf_type(self):
         # Full multimodal export: build() unwraps the composite to its text
         # sub-config, so config.model_type may be a text type even though the
-        # package is multimodal. Must keep the HF parent type -> gemma4.
+        # package is multimodal. Must keep the HF parent type -> gemma4_unified.
         assert (
             _select_ort_model_type(
                 "gemma4_unified_text", "gemma4_unified", is_decoder_only=False
             )
-            == "gemma4"
+            == "gemma4_unified"
         )
 
     def test_decoder_only_falls_back_to_hf_when_config_missing(self):
@@ -177,14 +177,29 @@ def test_writes_transform_pipeline(self, tmp_path):
         assert len(norm_attrs["mean"]) == 3
         assert len(norm_attrs["std"]) == 3
 
-    def test_gemma4_unified_skips_image_processor(self, tmp_path):
-        """Encoder-free gemma4_unified has no native transform: no image_processor.json."""
+    def test_gemma4_unified_image_processor(self, tmp_path):
+        """Encoder-free gemma4_unified writes a 48px merged-patch transform."""
         vision = mock.MagicMock()
         vision.model_type = None
+        vision.patch_size = 16
+        vision.pooling_kernel_size = 3
+        vision.mm_tokens_per_image = 280
         config = mock.MagicMock()
         config.vision = vision
         config.model_type = "gemma4_unified"
-        assert _write_vision_processor_config(config, str(tmp_path)) is None
+
+        path = _write_vision_processor_config(config, str(tmp_path))
+        assert path is not None
+        assert path.endswith("image_processor.json")
+        with open(path) as f:
+            data = json.load(f)
+
+        transform = data["processor"]["transforms"][1]["operation"]
+        assert transform["type"] == "Gemma4ImageTransform"
+        # 48px merged patches (16*3), no further pooling.
+        assert transform["attrs"]["patch_size"] == 48
+        assert transform["attrs"]["pooling_kernel_size"] == 1
+        assert transform["attrs"]["max_soft_tokens"] == 280
 
     def test_pixtral_vision_config(self, tmp_path):
         """Generates pixtral-specific processor config with 7 transforms."""
@@ -326,12 +341,25 @@ def test_audio_non_gemma4_returns_none(self, tmp_path):
         config.model_type = "whisper"
         assert _write_audio_processor_config(config, str(tmp_path)) is None
 
-    def test_audio_gemma4_unified_skips_audio_processor(self, tmp_path):
-        """Encoder-free gemma4_unified has no native transform: no audio_processor.json."""
+    def test_audio_gemma4_unified_writes_raw_frames(self, tmp_path):
+        """Encoder-free gemma4_unified writes a raw-waveform-frame extractor."""
         config = mock.MagicMock()
         config.audio = mock.MagicMock()
+        config.audio.hidden_size = 640
         config.model_type = "gemma4_unified"
-        assert _write_audio_processor_config(config, str(tmp_path)) is None
+
+        path = _write_audio_processor_config(config, str(tmp_path))
+        assert path is not None
+        assert path.endswith("audio_feature_extraction.json")
+        with open(path) as f:
+            data = json.load(f)
+
+        seq = data["feature_extraction"]["sequence"]
+        assert seq[0]["operation"]["type"] == "AudioDecoder"
+        op = seq[1]["operation"]
+        assert op["type"] == "Gemma4Audio"
+        assert op["attrs"]["type"] == "raw_frames"
+        assert op["attrs"]["audio_samples_per_token"] == 640
 
     def test_audio_gemma4_writes_feature_extraction_json(self, tmp_path):
         config = mock.MagicMock()
@@ -349,8 +377,9 @@ def test_audio_gemma4_writes_feature_extraction_json(self, tmp_path):
         seq = data["feature_extraction"]["sequence"]
         assert len(seq) == 2
         assert seq[0]["operation"]["type"] == "AudioDecoder"
-        assert seq[1]["operation"]["type"] == "Gemma4LogMel"
+        assert seq[1]["operation"]["type"] == "Gemma4Audio"
         attrs = seq[1]["operation"]["attrs"]
+        assert attrs["type"] == "log_mel"
         assert attrs["feature_size"] == 128
         assert attrs["sampling_rate"] == 16000
         assert attrs["frame_length_ms"] == 20.0  # noqa: RUF069
@@ -720,9 +749,10 @@ class FakeConfig:
         op0 = seq[0]["operation"]
         assert op0["type"] == "AudioDecoder"
 
-        # Second op: Gemma4LogMel with expected attrs
+        # Second op: Gemma4Audio (type=log_mel) with expected attrs
         op1 = seq[1]["operation"]
-        assert op1["type"] == "Gemma4LogMel"
+        assert op1["type"] == "Gemma4Audio"
+        assert op1["attrs"]["type"] == "log_mel"
         assert op1["attrs"]["feature_size"] == 128
         assert op1["attrs"]["sampling_rate"] == 16000
         assert op1["attrs"]["mel_floor"] == 0.001  # noqa: RUF069
@@ -1513,7 +1543,7 @@ def test_text_only_genai_config_is_decoder_only(self, tmp_path):
             data = json.load(f)
 
         # ORT-GenAI type resolved from pkg.config.model_type (gemma4_text),
-        # NOT the multimodal HF gemma4_unified -> gemma4.
+        # NOT the multimodal HF gemma4_unified -> gemma4_unified.
         assert data["model"]["type"] == "gemma4_text"
         # Decoder-only: input_ids decoder, no multimodal sections.
         assert "vision" not in data["model"]