custom processor for omni poc

src/maxtext/experimental/omni_poc/omni-gemma3-qwen3.yml

@@ -1,12 +1,15 @@
 # model config for omni-gemma3-qwen3
 # Combines LLM backbone from Qwen 3 4B with Vision features from Gemma 3 4B
 
-# NOTE: model_name is set to "gemma3-4b" so encoders.py loads Gemma 3's Vision Tower without modification.
-# Meanwhile, decoders.py can build Qwen 3 LLM decoder directly from reading `decoder_block: "qwen3"` below.
-model_name: "gemma3-4b"
+# NOTE: vision_encoder_block is set to "gemma3" so encoders.py loads Gemma 3's Vision Tower without modification.
+# Meanwhile, decoders.py builds Qwen 3 LLM decoder directly from reading `decoder_block: "qwen3"` below.
+# base_config is currently a placeholder flag for future SFT runs.
+base_config: "base.yml"
+model_name: "omni-gemma3-qwen3"
 use_multimodal: true
 
-# Multimodal config for vision model (from gemma3-4b.yml)
+# Vision encoder config (from gemma3-4b.yml)
+vision_encoder_block: "gemma3"
 image_size_for_vit: 896
 num_channels_for_vit: 3
 patch_size_for_vit: 14
@@ -15,7 +18,14 @@ hidden_size_for_vit: 1152
 intermediate_size_for_vit: 4304
 num_hidden_layers_for_vit: 27
 num_attention_heads_for_vit: 16
-image_placeholder: <start_of_image>
+image_placeholder: "<image>"
+
+# Vision projector config
+vision_projector_type: "customized_vision_projector"
+vision_connector_hidden_size: 2560
+vision_connector_num_layers: 2
+vision_connector_activation: "gelu"
+vision_connector_use_bias: true
 
 # LLM backbone config (from qwen3-4b.yml)
 base_emb_dim: 2560
@@ -34,6 +44,7 @@ logits_via_embedding: true
 normalize_embedding_logits: false
 enable_dropout: false
 tokenizer_type: "huggingface"
+tokenizer_path: "Qwen/Qwen2.5-7B-Instruct"
 
 # Ensure no pretrained weights or checkpoints are loaded initially
 load_parameters_path: ""

src/maxtext/experimental/omni_poc/tests/processor_omni_gemma3_qwen3_test.py

@@ -0,0 +1,67 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for Omni multimodal processor."""
+
+from types import SimpleNamespace
+import unittest
+import numpy as np
+
+from maxtext.experimental.omni_poc.utils import processor_omni_gemma3_qwen3 as omni_processor
+
+
+class TestProcessorOmniGemma3Qwen3(unittest.TestCase):
+  """Concise test suite for processor_omni_gemma3_qwen3."""
+
+  def test_constants_and_offsets(self):
+    self.assertEqual(omni_processor.IMAGE_PAD_ID, 151655)
+    self.assertEqual(omni_processor.get_image_offsets_omni(None), 255)
+    self.assertEqual(
+        omni_processor.get_image_offsets_omni(SimpleNamespace(pixel_values=np.zeros((3, 1, 1, 3)))),
+        255 * 3,
+    )
+
+  def test_add_extra_tokens_placeholder_expansion(self):
+    pad_id = omni_processor.IMAGE_PAD_ID
+    # Expand single placeholder with custom count
+    tokens = [10, pad_id, 20]
+    out = omni_processor.add_extra_tokens_for_omni(tokens, num_tokens_per_image=3)
+    np.testing.assert_array_equal(out, [10, pad_id, pad_id, pad_id, 20])
+
+    # 2D array and dtype preservation
+    tokens_i64 = np.array([[pad_id], [99]], dtype=np.int64)
+    out_i64 = omni_processor.add_extra_tokens_for_omni(tokens_i64, num_tokens_per_image=2)
+    self.assertEqual(out_i64.dtype, np.int64)
+    np.testing.assert_array_equal(out_i64, [pad_id, pad_id, 99])
+
+  def test_add_extra_tokens_prepend_and_text_only(self):
+    # Prepend vision block when pixel values exist
+    proc_out = SimpleNamespace(pixel_values=np.zeros((1, 1, 1, 3)))
+    out = omni_processor.add_extra_tokens_for_omni([1, 2], processor_output=proc_out, num_tokens_per_image=2)
+    expected = [
+        omni_processor.VISION_START_ID,
+        omni_processor.IMAGE_PAD_ID,
+        omni_processor.IMAGE_PAD_ID,
+        omni_processor.VISION_END_ID,
+        1,
+        2,
+    ]
+    np.testing.assert_array_equal(out, expected)
+
+    # Text-only returns original
+    np.testing.assert_array_equal(omni_processor.add_extra_tokens_for_omni([1, 2]), [1, 2])
+
+
+if __name__ == "__main__":
+  unittest.main()

src/maxtext/experimental/omni_poc/tests/stitch_checkpoint_test.py

@@ -52,12 +52,13 @@
 
 from maxtext.configs import pyconfig
 from maxtext.experimental.omni_poc.utils import stitch_checkpoint
+from maxtext.layers import encoders
 from maxtext.layers.embeddings import Embed
 from maxtext.models import gemma3
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
 from maxtext.utils import model_creation_utils
-from maxtext.utils.globals import MAXTEXT_REPO_ROOT
+from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR, MAXTEXT_PKG_DIR
 
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 warnings.filterwarnings("ignore", category=FutureWarning)
@@ -120,9 +121,7 @@ def tearDown(self):
   def _create_config(self):
     """Creates a MaxText config from the Omni config file and defaults."""
     omni_config_path = os.path.join(
-        MAXTEXT_REPO_ROOT,
-        "src",
-        "maxtext",
+        MAXTEXT_PKG_DIR,
         "experimental",
         "omni_poc",
         "omni-gemma3-qwen3.yml",
@@ -136,10 +135,12 @@ def _create_config(self):
         "stitched_output_path",
         "vision_model_name",
         "llm_model_name",
+        "base_config",
+        "model_name",
     }
     yaml_overrides = {k: v for k, v in custom_cfg.items() if k not in omni_keys}
 
-    base_config_path = os.path.join(MAXTEXT_REPO_ROOT, "src", "maxtext", "configs", "base.yml")
+    base_config_path = os.path.join(MAXTEXT_CONFIGS_DIR, "base.yml")
     config = pyconfig.initialize(
         ["", base_config_path],
         override_model_config=True,
@@ -151,6 +152,7 @@ def _create_config(self):
         **yaml_overrides,
     )
     # Use object.__setattr__ to bypass the read-only check on _HyperParameters
+    object.__setattr__(config, "model_name", custom_cfg.get("model_name", "omni-gemma3-qwen3"))
     object.__setattr__(config, "vision_model_name", "gemma3-4b")
     object.__setattr__(config, "llm_model_name", "qwen3-4b")
     object.__setattr__(config, "ici_context_autoregressive_parallelism", 1)
@@ -206,7 +208,11 @@ def to_concrete_zeros(leaf):
       stitched_inner = stitch_checkpoint._restore_subtrees_from_path(self.output_ckpt_dir, concrete_template, ckptr)
 
       # Load weights from the ORIGINAL Vision checkpoint
-      vision_abstract = concrete_template["vision_encoder"]
+      vision_abstract = {
+          k: v
+          for k, v in concrete_template["vision_encoder"].items()
+          if not ("projector" in k.lower() or "embedder" in k.lower())
+      }
       vision_restored = stitch_checkpoint._restore_subtrees_from_path(
           self.vision_ckpt_dir, {"vision_encoder": vision_abstract}, ckptr
       )
@@ -323,18 +329,14 @@ def test_2_stage_outputs_original_vs_stitched(self):
         "stitched llm token embedder first 10 values:",
         stitched_inner["token_embedder"]["embedding"].reshape(-1)[:10],
     )
-    print(
-        "orig vision projector weights first 10 values:",
-        vision_restored["vision_encoder"]["VisionEmbedder_0"]["mm_input_projection"]["w"].reshape(-1)[:10],
-    )
     print(
         "stitched vision projector weights first 10 values:",
-        stitched_inner["vision_encoder"]["VisionEmbedder_0"]["mm_input_projection"]["w"].reshape(-1)[:10],
+        stitched_inner["vision_encoder"]["VisionEmbedder_0"]["custom_linear_0"]["kernel"].reshape(-1)[:10],
     )
 
     # Instantiate and update original vision & token embedder modules
     orig_vision_tower = gemma3.Gemma3VisionEncoderLayer(config, self.test_mesh, rngs=rngs)
-    orig_projector = gemma3.VisionEmbedder(config, self.test_mesh, rngs=rngs)
+    orig_projector = encoders.MultimodalMLPProjector(config, self.test_mesh, rngs=nnx.Rngs(1))
     # Manually define qwen3's token embedder to match MaxText's structures
     orig_token_embedder = Embed(
         num_embeddings=config.vocab_size,
@@ -348,18 +350,14 @@ def test_2_stage_outputs_original_vs_stitched(self):
         orig_vision_tower,
         vision_restored["vision_encoder"]["Gemma3VisionEncoderLayer_0"],
     )
-    nnx.update(
-        orig_projector,
-        vision_restored["vision_encoder"]["VisionEmbedder_0"],
-    )
     nnx.update(
         orig_token_embedder,
         {"embedding": llm_restored["token_embedder"]["embedding"]},
     )
 
     # Instantiate and update stitched vision & token embedder modules
     stitched_vision_tower = gemma3.Gemma3VisionEncoderLayer(config, self.test_mesh, rngs=rngs)
-    stitched_projector = gemma3.VisionEmbedder(config, self.test_mesh, rngs=rngs)
+    stitched_projector = encoders.MultimodalMLPProjector(config, self.test_mesh, rngs=rngs)
     stitched_token_embedder = Embed(
         num_embeddings=config.vocab_size,
         num_features=config.emb_dim,