fix: Gemma4 genai config - vision inputs, decoder input_ids, processor

examples/gemma4/ort_genai/vlm/genai_config.json

@@ -20,6 +20,7 @@
             "num_hidden_layers": 15,
             "inputs": {
                 "inputs_embeds": "inputs_embeds",
+                "input_ids": "input_ids",
                 "attention_mask": "attention_mask",
                 "position_ids": "position_ids",
                 "past_key_names": "past_key_values.%d.key",

src/mobius/integrations/ort_genai/auto_export.py

@@ -150,10 +150,27 @@ def _write_processor_config(
     if vision is None:
         return None
 
-    processor: dict[str, Any] = {
-        "image_size": getattr(vision, "image_size", 448),
-        "patch_size": getattr(vision, "patch_size", 14),
-    }
+    model_type = getattr(config, "model_type", "")
+
+    if model_type in ("gemma4", "gemma4_text"):
+        # Gemma4 needs a processor wrapper with model-specific fields
+        max_soft_tokens = getattr(vision, "max_soft_tokens", 280)
+        processor: dict[str, Any] = {
+            "processor": {
+                "name": "gemma4_image_processor",
+                "image_size": getattr(vision, "image_size", 448),
+                "patch_size": getattr(vision, "patch_size", 16),
+                "tokens_per_image": max_soft_tokens,
+                "mean": [0.5, 0.5, 0.5],
+                "std": [0.5, 0.5, 0.5],
+            }
+        }
+    else:
+        processor = {
+            "image_size": getattr(vision, "image_size", 448),
+            "patch_size": getattr(vision, "patch_size", 14),
+        }
+
     path = os.path.join(output_dir, "processor_config.json")
     with open(path, "w", encoding="utf-8") as f:
         json.dump(processor, f, indent=4)
@@ -201,8 +218,28 @@ def _write_genai_config(
                     "pixel_values": "pixel_values",
                     "image_sizes": "image_sizes",
                 }
+            elif getattr(config, "model_type", "") in (
+                "gemma4",
+                "gemma4_text",
+            ):
+                # Gemma4 uses pixel_values + pixel_position_ids
+                # (not image_grid_thw)
+                vision_kwargs["spatial_merge_size"] = None
+                vision_kwargs["config_filename"] = "processor_config.json"
+                vision_kwargs["input_names"] = {
+                    "pixel_values": "pixel_values",
+                    "pixel_position_ids": "pixel_position_ids",
+                }
             generator.with_vision(image_token_id=image_token_id, **vision_kwargs)
 
+    # Gemma4 decoders need input_ids alongside inputs_embeds for
+    # per-layer token embeddings (E2B architecture).
+    if is_vlm and getattr(config, "model_type", "") in (
+        "gemma4",
+        "gemma4_text",
+    ):
+        generator.with_extra_decoder_inputs(input_ids="input_ids")
+
     if has_speech:
         audio_config = getattr(config, "audio", None)
         audio_token_id = (

src/mobius/integrations/ort_genai/genai_config.py

@@ -159,6 +159,9 @@ def __init__(
         self._embedding: dict[str, Any] | None = None
         self._vlm_token_ids: dict[str, int] = {}
 
+        # Extra decoder inputs (set via with_extra_decoder_inputs())
+        self._extra_decoder_inputs: dict[str, str] = {}
+
         # Optional speech fields (set via with_speech())
         self._speech: dict[str, Any] | None = None
 
@@ -277,6 +280,24 @@ def with_vision(
             self._vlm_token_ids["video_token_id"] = video_token_id
         return self
 
+    def with_extra_decoder_inputs(
+        self,
+        **inputs: str,
+    ) -> GenaiConfigGenerator:
+        """Add extra inputs to the decoder section.
+
+        For example, Gemma4 decoders need ``input_ids`` alongside
+        ``inputs_embeds`` for per-layer token embeddings (E2B).
+
+        Args:
+            **inputs: Mapping of input names to ONNX tensor names,
+                e.g. ``input_ids="input_ids"``.
+
+        Returns self for chaining.
+        """
+        self._extra_decoder_inputs.update(inputs)
+        return self
+
     def with_speech(
         self,
         *,
@@ -329,12 +350,14 @@ def generate(self) -> dict[str, Any]:
         is_multimodal = self._vision is not None or self._speech is not None
 
         # Decoder section
+        decoder_inputs = _default_decoder_inputs(is_vlm=is_multimodal)
+        decoder_inputs.update(self._extra_decoder_inputs)
         decoder: dict[str, Any] = {
             "session_options": _make_session_options(self.ep),
             "filename": "model.onnx",
             "head_size": self.head_dim,
             "hidden_size": self.hidden_size,
-            "inputs": _default_decoder_inputs(is_vlm=is_multimodal),
+            "inputs": decoder_inputs,
             "outputs": _default_decoder_outputs(),
             "num_attention_heads": self.num_attention_heads,
             "num_hidden_layers": self.num_hidden_layers,