Add Gemma4 image_processor.json generation with onnxruntime-extension…

src/mobius/integrations/ort_genai/auto_export.py

@@ -158,13 +158,22 @@ def _copy_tokenizer_files_from_local(
     return copied
 
 
-def _write_processor_config(
+def _write_vision_processor_config(
     config: Any,
     output_dir: str,
 ) -> str | None:
-    """Write a minimal processor_config.json for VLM models.
+    """Write the vision processor config file for VLM models.
 
-    Returns the path if written, None otherwise.
+    The output format depends on the model type:
+
+    - **Gemma4** (``gemma4``, ``gemma4_text``): Writes ``image_processor.json``
+      in the onnxruntime-extensions transforms pipeline format required by
+      ``OrtxCreateProcessor``.  The pipeline is
+      ``DecodeImage → Gemma4ImageTransform``.
+    - **Other models**: Writes ``processor_config.json`` with a minimal
+      HuggingFace-style schema (``image_size``, ``patch_size``).
+
+    Returns the written file path, or None if the config has no vision section.
     """
     vision = getattr(config, "vision", None)
     if vision is None:
@@ -173,30 +182,52 @@ def _write_processor_config(
     model_type = getattr(config, "model_type", "")
 
     if model_type in ("gemma4", "gemma4_text"):
-        # Gemma4 needs a processor wrapper with model-specific fields
-        tokens_per_image = (
+        # Gemma4 needs an onnxruntime-extensions format processor config
+        # with a transforms pipeline (DecodeImage -> Gemma4ImageTransform).
+        # The OrtxCreateProcessor API requires this format.
+        #
+        # max_soft_tokens: maps from HF's mm_tokens_per_image (the number of
+        # vision tokens per image after pooling) into the Gemma4ImageTransform's
+        # max_soft_tokens attribute, which controls the padded patch budget.
+        max_soft_tokens = (
             getattr(vision, "mm_tokens_per_image", None)
             or getattr(config, "mm_tokens_per_image", None)
-            or getattr(vision, "max_soft_tokens", None)
             or 280
         )
-        image_size = getattr(vision, "image_size", None) or 448
         patch_size = getattr(vision, "patch_size", None) or 16
-        processor: dict[str, Any] = {
+        pooling_kernel_size = getattr(vision, "pooling_kernel_size", None) or 3
+        processor = {
             "processor": {
-                "name": "gemma4_image_processor",
-                "image_size": image_size,
-                "patch_size": patch_size,
-                "tokens_per_image": tokens_per_image,
+                "name": "gemma_4_image_processing",
+                "transforms": [
+                    {
+                        "operation": {
+                            "name": "decode_image",
+                            "type": "DecodeImage",
+                            "attrs": {"color_space": "RGB"},
+                        }
+                    },
+                    {
+                        "operation": {
+                            "name": "gemma4_image_transform",
+                            "type": "Gemma4ImageTransform",
+                            "attrs": {
+                                "patch_size": patch_size,
+                                "max_soft_tokens": max_soft_tokens,
+                                "pooling_kernel_size": pooling_kernel_size,
+                            },
+                        }
+                    },
+                ],
             }
         }
+        proc_filename = "image_processor.json"
     else:
         processor = {
             "image_size": getattr(vision, "image_size", None) or 448,
             "patch_size": getattr(vision, "patch_size", None) or 14,
         }
-
-    proc_filename = "processor_config.json"
+        proc_filename = "processor_config.json"
 
     path = os.path.join(output_dir, proc_filename)
     with open(path, "w", encoding="utf-8") as f:
@@ -280,6 +311,7 @@ def _write_genai_config(
                 vision_cfg = getattr(config, "vision", None)
                 sms = getattr(vision_cfg, "spatial_merge_size", 2)
                 vision_kwargs["spatial_merge_size"] = sms
+                vision_kwargs["config_filename"] = "image_processor.json"
 
             if vision_input_mapping is not None:
                 vision_kwargs["input_names"] = vision_input_mapping
@@ -450,7 +482,7 @@ def write_ort_genai_config(
             result[tf] = os.path.join(directory, tf)
 
     # Write processor_config.json for VLMs
-    processor_path = _write_processor_config(config, directory)
+    processor_path = _write_vision_processor_config(config, directory)
     if processor_path:
         result["processor_config"] = processor_path
 

src/mobius/integrations/ort_genai/auto_export_test.py

@@ -18,7 +18,7 @@
     _graph_input_names,
     _resolve_ort_genai_model_type,
     _write_genai_config,
-    _write_processor_config,
+    _write_vision_processor_config,
     write_ort_genai_config,
 )
 
@@ -42,7 +42,7 @@ class TestWriteProcessorConfig:
     def test_no_vision_returns_none(self, tmp_path):
         config = mock.MagicMock(spec=[])
         del config.vision  # ensure no vision attribute
-        assert _write_processor_config(config, str(tmp_path)) is None
+        assert _write_vision_processor_config(config, str(tmp_path)) is None
 
     def test_writes_vision_config(self, tmp_path):
         vision = mock.MagicMock()
@@ -51,7 +51,7 @@ def test_writes_vision_config(self, tmp_path):
         config = mock.MagicMock()
         config.vision = vision
 
-        path = _write_processor_config(config, str(tmp_path))
+        path = _write_vision_processor_config(config, str(tmp_path))
         assert path is not None
         with open(path) as f:
             data = json.load(f)
@@ -314,6 +314,69 @@ def test_processor_config_not_written_without_vision(self, tmp_path):
         assert "processor_config" not in result
         assert not os.path.exists(os.path.join(str(tmp_path), "processor_config.json"))
 
+    def test_gemma4_image_processor_json_written(self, tmp_path):
+        """Gemma4 writes image_processor.json with onnxruntime-extensions transforms pipeline."""
+        import dataclasses
+
+        from mobius._model_package import ModelPackage
+        from mobius.integrations.ort_genai.auto_export import write_ort_genai_config
+
+        @dataclasses.dataclass
+        class FakeVision:
+            image_size: int = 448
+            patch_size: int = 16
+            mm_tokens_per_image: int = 260
+            pooling_kernel_size: int = 3
+
+        @dataclasses.dataclass
+        class FakeConfig:
+            model_type: str = "gemma4"
+            vocab_size: int = 262144
+            hidden_size: int = 1536
+            num_hidden_layers: int = 35
+            num_attention_heads: int = 8
+            num_key_value_heads: int = 1
+            head_dim: int = 256
+            vision: FakeVision = dataclasses.field(default_factory=FakeVision)
+
+        pkg = ModelPackage(
+            {
+                "model": mock.MagicMock(),
+                "vision": mock.MagicMock(),
+                "embedding": mock.MagicMock(),
+            },
+            config=FakeConfig(),
+        )
+        result = write_ort_genai_config(pkg, str(tmp_path))
+
+        # Should write image_processor.json, not processor_config.json
+        assert "processor_config" in result
+        proc_path = result["processor_config"]
+        assert proc_path.endswith("image_processor.json")
+        assert os.path.isfile(proc_path)
+        assert not os.path.exists(os.path.join(str(tmp_path), "processor_config.json"))
+
+        with open(proc_path) as f:
+            data = json.load(f)
+
+        # Verify onnxruntime-extensions transforms pipeline structure
+        assert "processor" in data
+        assert "transforms" in data["processor"]
+        transforms = data["processor"]["transforms"]
+        assert len(transforms) == 2
+
+        # First op: DecodeImage
+        op0 = transforms[0]["operation"]
+        assert op0["type"] == "DecodeImage"
+        assert op0["attrs"]["color_space"] == "RGB"
+
+        # Second op: Gemma4ImageTransform with correct attrs from config
+        op1 = transforms[1]["operation"]
+        assert op1["type"] == "Gemma4ImageTransform"
+        assert op1["attrs"]["patch_size"] == 16
+        assert op1["attrs"]["max_soft_tokens"] == 260
+        assert op1["attrs"]["pooling_kernel_size"] == 3
+
     def test_tokenizer_not_copied_without_model_id(self, tmp_path):
         """No tokenizer files copied when hf_model_id=None."""
         from mobius.integrations.ort_genai.auto_export import write_ort_genai_config
@@ -780,7 +843,7 @@ def test_gemma4_genai_config_from_real_model(self, tmp_path):
         # Config-level properties are still present
         assert data["model"]["image_token_id"] == 255999
         assert data["model"]["vision"]["spatial_merge_size"] == 2
-        assert data["model"]["vision"]["config_filename"] == "processor_config.json"
+        assert data["model"]["vision"]["config_filename"] == "image_processor.json"
 
     def test_auto_export_produces_genai_config(self, tmp_path):
         """Mock build() to return a tiny package, verify genai_config."""