fix(gemma4): close text-only image modality

Resolve Joi gap #1 by emitting an optional image_features fallback and gating vision_encoder on the generic image presence key.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: justinchuby

src/mobius/integrations/onnx_genai/inference_metadata_test.py

@@ -121,12 +121,20 @@ class _VlmConfig:
 def _embedding_model(
     outputs: list[tuple[str, ir.DataType, list[int | str]]],
     *,
+    optional_image: bool = False,
     include_audio: bool = False,
     optional_audio: bool = False,
 ) -> ir.Model:
+    image_features = _value("image_features", ir.DataType.FLOAT, ["image_tokens", 64])
+    if optional_image:
+        declare_optional_input(
+            image_features,
+            presence="image",
+            absent_shape=[0, 64],
+        )
     inputs = [
         _value("input_ids", ir.DataType.INT64, ["batch", "sequence"]),
-        _value("image_features", ir.DataType.FLOAT, ["image_tokens", 64]),
+        image_features,
     ]
     if include_audio:
         audio_features = _value("audio_features", ir.DataType.FLOAT, ["audio_tokens", 64])
@@ -359,6 +367,7 @@ def test_gemma4_routes_all_embedding_outputs(self, tmp_path):
                             ["batch", "sequence", 128],
                         ),
                     ],
+                    optional_image=True,
                     include_audio=True,
                     optional_audio=True,
                 ),
@@ -392,6 +401,7 @@ def test_gemma4_routes_all_embedding_outputs(self, tmp_path):
             },
             config=config,
         )
+        declare_component_presence(package["vision_encoder"].graph, "image")
         declare_component_presence(package["audio_encoder"].graph, "audio")
 
         source = tmp_path / "gemma"
@@ -456,15 +466,29 @@ def test_gemma4_routes_all_embedding_outputs(self, tmp_path):
             "from": "audio_encoder.audio_features",
         }
         assert emitted_yaml["pipeline"]["models"]["embedding"]["io"]["optional_inputs"] == {
+            "image_features": {
+                "presence": "image",
+                "absent": {"kind": "zeros", "shape": [0, 64]},
+            },
             "audio_features": {
                 "presence": "audio",
                 "absent": {"kind": "zeros", "shape": [0, 64]},
-            }
+            },
+        }
+        assert emitted_yaml["pipeline"]["phases"]["vision_encoder"] == {
+            "run_on": "prompt_only",
+            "when_present": "image",
         }
         assert emitted_yaml["pipeline"]["phases"]["audio_encoder"] == {
             "run_on": "prompt_only",
             "when_present": "audio",
         }
+        vision_stage = next(
+            stage
+            for stage in metadata["pipeline"]["strategy"]["stages"]
+            if stage["strategy"].get("model") == "vision_encoder"
+        )
+        assert vision_stage["run_on"] == "prompt_only"
         audio_stage = next(
             stage
             for stage in metadata["pipeline"]["strategy"]["stages"]

src/mobius/tasks/_gemma4.py

@@ -362,6 +362,7 @@ def _build_vision(
 
         builder.add_output(image_features, "image_features")
 
+        declare_component_presence(graph, "image")
         return _make_model(graph)
 
     def _build_audio(
@@ -455,6 +456,11 @@ def _build_embedding(
             dtype=config.dtype,
             shape=[num_image_tokens, config.hidden_size],
         )
+        declare_optional_input(
+            image_features,
+            presence="image",
+            absent_shape=[0, config.hidden_size],
+        )
 
         audio_features_val: ir.Value | None = None
 
@@ -547,6 +553,7 @@ def _build_vision(
             pixel_position_ids=pixel_position_ids,
         )
         builder.add_output(image_features, "image_features")
+        declare_component_presence(graph, "image")
         return _make_model(graph)
 
     def _build_audio(

tests/build_graph_test.py

@@ -1255,12 +1255,18 @@ def test_gemma4_multimodal_graph(self):
         assert "pixel_values" in vision_input_names
         assert "pixel_position_ids" in vision_input_names
         assert "image_features" in {o.name for o in vision.graph.outputs}
+        assert component_presence(vision.graph) == "image"
         # Embedding: input_ids + image_features (no audio) -> inputs_embeds
         embedding = pkg["embedding"]
         emb_input_names = {i.name for i in embedding.graph.inputs}
         assert "input_ids" in emb_input_names
         assert "image_features" in emb_input_names
         assert "audio_features" not in emb_input_names
+        embedding_image = next(i for i in embedding.graph.inputs if i.name == "image_features")
+        assert optional_input_contract(embedding_image) == {
+            "presence": "image",
+            "absent": {"kind": "zeros", "shape": [0, config.hidden_size]},
+        }
         assert "inputs_embeds" in {o.name for o in embedding.graph.outputs}
 
     def test_gemma4_kv_shared_fallback_attention_is_causal_zero(self):
@@ -1657,6 +1663,7 @@ def test_gemma4_any_to_any_graph(self):
         assert "pixel_values" in vision_input_names
         assert "pixel_position_ids" in vision_input_names
         assert "image_features" in {o.name for o in vision.graph.outputs}
+        assert component_presence(vision.graph) == "image"
         # Audio encoder
         audio = pkg["audio_encoder"]
         audio_input_names = {i.name for i in audio.graph.inputs}
@@ -1672,6 +1679,11 @@ def test_gemma4_any_to_any_graph(self):
         assert "input_ids" in emb_input_names
         assert "image_features" in emb_input_names
         assert "audio_features" in emb_input_names
+        embedding_image = next(i for i in embedding.graph.inputs if i.name == "image_features")
+        assert optional_input_contract(embedding_image) == {
+            "presence": "image",
+            "absent": {"kind": "zeros", "shape": [0, config.hidden_size]},
+        }
         embedding_audio = next(i for i in embedding.graph.inputs if i.name == "audio_features")
         assert optional_input_contract(embedding_audio) == {
             "presence": "audio",
@@ -1796,6 +1808,7 @@ def test_gemma4_unified_multimodal_graph(self):
         v_inputs = {i.name for i in vision.graph.inputs}
         assert v_inputs == {"pixel_values", "pixel_position_ids"}
         assert "image_features" in {o.name for o in vision.graph.outputs}
+        assert component_presence(vision.graph) == "image"
 
         # Audio embedder: raw frames + mask → audio_features
         audio = pkg["audio_encoder"]
@@ -1809,6 +1822,11 @@ def test_gemma4_unified_multimodal_graph(self):
         embedding = pkg["embedding"]
         e_inputs = {i.name for i in embedding.graph.inputs}
         assert {"input_ids", "image_features", "audio_features"} <= e_inputs
+        embedding_image = next(i for i in embedding.graph.inputs if i.name == "image_features")
+        assert optional_input_contract(embedding_image) == {
+            "presence": "image",
+            "absent": {"kind": "zeros", "shape": [0, config.hidden_size]},
+        }
         embedding_audio = next(i for i in embedding.graph.inputs if i.name == "audio_features")
         assert optional_input_contract(embedding_audio) == {
             "presence": "audio",