diff --git a/src/maxtext/multimodal/processor.py b/src/maxtext/multimodal/processor.py
index b004d91806..64d7f15ded 100644
--- a/src/maxtext/multimodal/processor.py
+++ b/src/maxtext/multimodal/processor.py
@@ -14,12 +14,72 @@
 
 """Multimodal data preprocessor router."""
 
+import functools
+import os
+from maxtext.common.common_types import DecoderBlockType, VisionEncoderBlockType
 from maxtext.multimodal import utils as mm_utils
+from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR
+import omegaconf
+
+
+@functools.lru_cache(maxsize=None)
+def _get_block_name_from_model_yml(model_name: str, block_name: str) -> str | None:
+  """Loads a model YAML from MAXTEXT_CONFIGS_DIR and returns the specified block name.
+
+  Args:
+    model_name: Name of the model configuration file (e.g., 'gemma3-4b').
+    block_name: The architectural attribute to retrieve, can only be
+     'vision_encoder_block' or 'decoder_block'.
+
+  Returns:
+    The string value of the block configuration if defined, otherwise None.
+  """
+  model_yml_path = os.path.join(MAXTEXT_CONFIGS_DIR, "models", f"{model_name}.yml")
+  if os.path.exists(model_yml_path):
+    try:
+      loaded = omegaconf.OmegaConf.load(model_yml_path)
+      if block_name in loaded and loaded[block_name] is not None:
+        return str(loaded[block_name])
+    except Exception:  # pylint: disable=broad-except
+      pass
+  return None
+
+
+def _get_vision_block(config_or_name):
+  """Extract vision encoder architecture from config."""
+  # If input is config, extract vision_encoder name
+  if hasattr(config_or_name, "vision_encoder_block"):
+    block = config_or_name.vision_encoder_block
+    if block != VisionEncoderBlockType.NONE:
+      return block.value.lower()
+  # If input is model_name (backward compatibility), find its corresponding config
+  elif isinstance(config_or_name, str):
+    val = _get_block_name_from_model_yml(config_or_name, "vision_encoder_block")
+    if val is not None and val.lower() != VisionEncoderBlockType.NONE.value:
+      return val.lower()
+  # If non-vision model, return None
+  return None
+
+
+def _get_decoder_block(config_or_name):
+  """Extract decoder architecture from config."""
+  # If input is config, extract decoder name
+  if hasattr(config_or_name, "decoder_block"):
+    block = config_or_name.decoder_block
+    if block != DecoderBlockType.DEFAULT:
+      return block.value.lower()
+  # If input is model_name (backward compatibility), find its corresponding config
+  elif isinstance(config_or_name, str):
+    val = _get_block_name_from_model_yml(config_or_name, "decoder_block")
+    if val is not None and val.lower() != DecoderBlockType.DEFAULT.value:
+      return val.lower()
+  # If decoder_block not found or is default, return model_name/default
+  return str(getattr(config_or_name, "model_name", config_or_name)).lower()
 
 
 def preprocess_mm_data(config):
   """Preprocesses multimodal data based on the provided configuration.
-  Routes to the appropriate preprocessing function based on the model name.
+  Routes to the appropriate preprocessing function based on the vision architecture.
 
   Args:
     config: A `pyconfig.Config` object containing configuration parameters.
@@ -28,90 +88,75 @@ def preprocess_mm_data(config):
     A `PreprocessorOutput` object containing the processed multimodal data.
   """
   processor_outputs = mm_utils.PreprocessorOutput()
+  vision_block = _get_vision_block(config)
 
-  if config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
+  if vision_block in ["gemma3"]:
     from maxtext.multimodal.processor_gemma3 import preprocess_mm_data_gemma3  # pylint: disable=import-outside-toplevel
 
     images = [mm_utils.load_image_from_path(p) for p in config.image_path.split(",")]
     processor_outputs = preprocess_mm_data_gemma3(images)
-  elif config.model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
+  elif vision_block in ["gemma4"]:
     from maxtext.multimodal.processor_gemma4 import preprocess_mm_data_gemma4  # pylint: disable=import-outside-toplevel
 
     images = [mm_utils.load_image_from_path(p) for p in config.image_path.split(",")]
     processor_outputs = preprocess_mm_data_gemma4(images)
-  elif config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
+  elif vision_block in ["llama4"]:
     from maxtext.multimodal.processor_llama4 import preprocess_mm_data_llama4  # pylint: disable=import-outside-toplevel
 
     images = [mm_utils.load_image_from_path(p) for p in config.image_path.split(",")]
     processor_outputs = preprocess_mm_data_llama4(images)
-  elif config.model_name in [
-      "qwen3-omni-30b-a3b",
-      "qwen3-vl-2b",
-      "qwen3-vl-4b",
-      "qwen3-vl-30b-a3b",
-      "qwen3.5-35b-a3b",
-      "qwen3.5-397b-a17b",
-  ]:
+  elif vision_block in ["qwen3_omni", "qwen3_vl", "qwen3_5"]:
     from maxtext.multimodal.processor_qwen3_omni import preprocess_mm_data_qwen3_omni  # pylint: disable=import-outside-toplevel
 
     processor_outputs = preprocess_mm_data_qwen3_omni(config)
   else:
-    raise ValueError(f"Model {config.model_name} not supported for multimodal preprocessing.")
+    raise ValueError(
+        f"Model {config.model_name} (vision block {vision_block}) not supported for multimodal preprocessing."
+    )
 
   return processor_outputs
 
 
 def preprocess_image_for_training(image, config):
-  """Preprocesses a single image for training based on the model name."""
-  if config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
+  """Preprocesses a single image for training based on the vision architecture."""
+  vision_block = _get_vision_block(config)
+  if vision_block in ["gemma3"]:
     from maxtext.multimodal.processor_gemma3 import preprocess_mm_data_gemma3  # pylint: disable=import-outside-toplevel
 
     return preprocess_mm_data_gemma3(image)
-  elif config.model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
+  elif vision_block in ["gemma4"]:
     from maxtext.multimodal.processor_gemma4 import preprocess_mm_data_gemma4  # pylint: disable=import-outside-toplevel
 
     return preprocess_mm_data_gemma4(image)
-  elif config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
+  elif vision_block in ["llama4"]:
     from maxtext.multimodal.processor_llama4 import preprocess_mm_data_llama4  # pylint: disable=import-outside-toplevel
 
     return preprocess_mm_data_llama4(image)
-  elif config.model_name in [
-      "qwen3-omni-30b-a3b",
-      "qwen3-vl-2b",
-      "qwen3-vl-4b",
-      "qwen3-vl-30b-a3b",
-      "qwen3.5-35b-a3b",
-      "qwen3.5-397b-a17b",
-  ]:
+  elif vision_block in ["qwen3_omni", "qwen3_vl", "qwen3_5"]:
     from maxtext.multimodal.processor_qwen3_omni import preprocess_mm_data_qwen3_omni_for_training  # pylint: disable=import-outside-toplevel
 
     return preprocess_mm_data_qwen3_omni_for_training(image, config)
   else:
-    raise ValueError(f"Model {config.model_name} not supported for image preprocessing.")
+    raise ValueError(f"Model {config.model_name} (vision block {vision_block}) not supported for image preprocessing.")
 
 
 def get_image_offsets(config, processor_output: mm_utils.PreprocessorOutput | None):
   """Get the increase in total token count after inserting image token placeholders"""
-  if config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
+  vision_block = _get_vision_block(config)
+
+  if vision_block in ["gemma3"]:
     from maxtext.multimodal.processor_gemma3 import get_image_offsets_gemma3  # pylint: disable=import-outside-toplevel
 
     return get_image_offsets_gemma3(processor_output)
-  elif config.model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
+  elif vision_block in ["gemma4"]:
     from maxtext.multimodal.processor_gemma4 import get_image_offsets_gemma4  # pylint: disable=import-outside-toplevel
 
     return get_image_offsets_gemma4(processor_output)
-  elif config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
+  elif vision_block in ["llama4"]:
     from maxtext.multimodal.processor_llama4 import get_image_offsets_llama4  # pylint: disable=import-outside-toplevel
 
     return get_image_offsets_llama4(processor_output)
-  elif config.model_name in [
-      "qwen3-omni-30b-a3b",
-      "qwen3-vl-2b",
-      "qwen3-vl-4b",
-      "qwen3-vl-30b-a3b",
-      "qwen3.5-35b-a3b",
-      "qwen3.5-397b-a17b",
-  ]:
+  elif vision_block in ["qwen3_omni", "qwen3_vl", "qwen3_5"]:
     from maxtext.multimodal.processor_qwen3_omni import get_mm_offsets_qwen3_omni  # pylint: disable=import-outside-toplevel
 
     return get_mm_offsets_qwen3_omni(config, processor_output)
@@ -121,26 +166,25 @@ def get_image_offsets(config, processor_output: mm_utils.PreprocessorOutput | No
 
 def reformat_prompt(prompt, image_placeholder, model_name, num_images, video_placeholder="<|video|>", num_videos=0):
   """Reformat prompt for different models."""
-  if model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
+  vision_block = _get_vision_block(model_name)
+  if vision_block is None:
+    return prompt
+
+  decoder_block = _get_decoder_block(model_name)
+
+  if decoder_block in ["gemma3"]:
     from maxtext.multimodal.processor_gemma3 import reformat_prompt_gemma3  # pylint: disable=import-outside-toplevel
 
     return reformat_prompt_gemma3(prompt, image_placeholder, num_images)
-  elif model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
+  elif decoder_block in ["gemma4", "gemma4_small"]:
     from maxtext.multimodal.processor_gemma4 import reformat_prompt_gemma4  # pylint: disable=import-outside-toplevel
 
     return reformat_prompt_gemma4(prompt, image_placeholder, num_images)
-  elif model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
+  elif decoder_block in ["llama4"]:
     from maxtext.multimodal.processor_llama4 import reformat_prompt_llama4  # pylint: disable=import-outside-toplevel
 
     return reformat_prompt_llama4(prompt, image_placeholder, num_images)
-  elif model_name in [
-      "qwen3-omni-30b-a3b",
-      "qwen3-vl-2b",
-      "qwen3-vl-4b",
-      "qwen3-vl-30b-a3b",
-      "qwen3.5-35b-a3b",
-      "qwen3.5-397b-a17b",
-  ]:
+  elif decoder_block in ["qwen3", "qwen3_moe", "qwen3_5"]:
     from maxtext.multimodal.processor_qwen3_omni import reformat_prompt_qwen3_omni  # pylint: disable=import-outside-toplevel
 
     return reformat_prompt_qwen3_omni(
@@ -156,23 +200,22 @@ def reformat_prompt(prompt, image_placeholder, model_name, num_images, video_pla
 
 def reformat_response(response, model_name):
   """Reformat response for different models."""
-  if model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
+  vision_block = _get_vision_block(model_name)
+  if vision_block is None:
+    return response
+
+  decoder_block = _get_decoder_block(model_name)
+
+  if decoder_block in ["llama4"]:
     formatted_response = f"{response}<|eot|>"
     return formatted_response
-  elif model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
+  elif decoder_block in ["gemma3"]:
     formatted_response = f"{response}<end_of_turn>"
     return formatted_response
-  elif model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
+  elif decoder_block in ["gemma4", "gemma4_small"]:
     formatted_response = f"{response}<turn|>"
     return formatted_response
-  elif model_name in [
-      "qwen3-omni-30b-a3b",
-      "qwen3-vl-2b",
-      "qwen3-vl-4b",
-      "qwen3-vl-30b-a3b",
-      "qwen3.5-35b-a3b",
-      "qwen3.5-397b-a17b",
-  ]:
+  elif decoder_block in ["qwen3", "qwen3_moe", "qwen3_5"]:
     formatted_response = f"{response}<|im_end|>"
     return formatted_response
   else:
@@ -181,53 +224,48 @@ def reformat_response(response, model_name):
 
 def prepare_text_for_image_fusion(tokens, config, processor_output=None):
   """Prepare text by adding extra tokens for image fusion based on the model."""
-  if config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
+  vision_block = _get_vision_block(config)
+  if vision_block in ["gemma3"]:
     from maxtext.multimodal.processor_gemma3 import add_extra_tokens_for_images_gemma3  # pylint: disable=import-outside-toplevel
 
     return add_extra_tokens_for_images_gemma3(
         tokens, max_num_images=processor_output.num_images  # pyrefly: ignore[missing-attribute]
     )  # pyrefly: ignore[missing-attribute]
-  elif config.model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
+  elif vision_block in ["gemma4"]:
     from maxtext.multimodal.processor_gemma4 import add_extra_tokens_for_images_gemma4  # pylint: disable=import-outside-toplevel
 
     return add_extra_tokens_for_images_gemma4(
         tokens, max_num_images=processor_output.num_images  # pyrefly: ignore[missing-attribute]
     )  # pyrefly: ignore[missing-attribute]
-  elif config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
+  elif vision_block in ["llama4"]:
     from maxtext.multimodal.processor_llama4 import add_extra_tokens_for_images_llama4  # pylint: disable=import-outside-toplevel
-    # pyrefly: ignore[bad-argument-type]
-    return add_extra_tokens_for_images_llama4(tokens, processor_output)
-  elif config.model_name in [
-      "qwen3-omni-30b-a3b",
-      "qwen3-vl-2b",
-      "qwen3-vl-4b",
-      "qwen3-vl-30b-a3b",
-      "qwen3.5-35b-a3b",
-      "qwen3.5-397b-a17b",
-  ]:
+
+    return add_extra_tokens_for_images_llama4(tokens, processor_output)  # pyrefly: ignore[bad-argument-type]
+  elif vision_block in ["qwen3_omni", "qwen3_vl", "qwen3_5"]:
     from maxtext.multimodal.processor_qwen3_omni import add_extra_tokens_for_qwen3_omni  # pylint: disable=import-outside-toplevel
 
     return add_extra_tokens_for_qwen3_omni(tokens, config, processor_output)
   else:
-    raise ValueError(f"Model {config.model_name} does not support multimodal inference.")
+    raise ValueError(f"Model {config.model_name} (vision block {vision_block}) does not support multimodal inference.")
 
 
 def get_dummy_image_shape_for_init(model_name, batch_size=1, num_image_per_sequence=1):
   """Return the shape of the dummy image for specific model's initialization."""
   image_shape = ()
-  if model_name.startswith("gemma3"):
+  vision_block = _get_vision_block(model_name)
+  if vision_block in ["gemma3"]:
     from maxtext.multimodal.processor_gemma3 import get_dummy_image_shape_for_init_gemma3  # pylint: disable=import-outside-toplevel
 
     image_shape = get_dummy_image_shape_for_init_gemma3(batch_size, num_image_per_sequence)
-  elif model_name.startswith("gemma4"):
+  elif vision_block in ["gemma4"]:
     from maxtext.multimodal.processor_gemma4 import get_dummy_image_shape_for_init_gemma4  # pylint: disable=import-outside-toplevel
 
     image_shape = get_dummy_image_shape_for_init_gemma4(batch_size, num_image_per_sequence)
-  elif model_name.startswith("llama4"):
+  elif vision_block in ["llama4"]:
     from maxtext.multimodal.processor_llama4 import get_dummy_image_shape_for_init_llama4  # pylint: disable=import-outside-toplevel
 
     image_shape = get_dummy_image_shape_for_init_llama4(batch_size, num_image_per_sequence)
-  elif model_name.startswith("qwen3-omni") or model_name.startswith("qwen3-vl") or model_name.startswith("qwen3.5"):
+  elif vision_block in ["qwen3_omni", "qwen3_vl", "qwen3_5"]:
     from maxtext.multimodal.processor_qwen3_omni import get_dummy_image_shape_for_init_qwen3_omni  # pylint: disable=import-outside-toplevel
 
     image_shape = get_dummy_image_shape_for_init_qwen3_omni(batch_size)
@@ -256,26 +294,26 @@ def get_dummy_audio_shape_for_init(config):
 def get_bidirectional_mask_vision(config, decoder_input_tokens, is_video: bool = False):
   """Get the bidirectional mask for specific models."""
   bidirectional_mask_vision = None
-  if config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
+
+  vision_block = _get_vision_block(config)
+  if vision_block is None:
+    return bidirectional_mask_vision
+
+  decoder_block = _get_decoder_block(config)
+
+  if decoder_block in ["gemma3"]:
     from maxtext.multimodal.processor_gemma3 import GEMMA_TOKEN_PLACEHOLDER  # pylint: disable=import-outside-toplevel
 
     bidirectional_mask_vision = decoder_input_tokens == GEMMA_TOKEN_PLACEHOLDER
-  elif config.model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
+  elif decoder_block in ["gemma4", "gemma4_small"]:
     from maxtext.multimodal.processor_gemma4 import GEMMA4_TOKEN_PLACEHOLDER  # pylint: disable=import-outside-toplevel
 
     bidirectional_mask_vision = decoder_input_tokens == GEMMA4_TOKEN_PLACEHOLDER
-  elif config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
+  elif decoder_block in ["llama4"]:
     from maxtext.multimodal.processor_llama4 import LLAMA4_PATCH_TOKEN  # pylint: disable=import-outside-toplevel
 
     bidirectional_mask_vision = decoder_input_tokens == LLAMA4_PATCH_TOKEN
-  elif config.model_name in [
-      "qwen3-omni-30b-a3b",
-      "qwen3-vl-2b",
-      "qwen3-vl-4b",
-      "qwen3-vl-30b-a3b",
-      "qwen3.5-35b-a3b",
-      "qwen3.5-397b-a17b",
-  ]:
+  elif decoder_block in ["qwen3", "qwen3_moe", "qwen3_5"]:
     from maxtext.multimodal.processor_qwen3_omni import QwenTokens  # pylint: disable=import-outside-toplevel
 
     tokens = QwenTokens(config)
diff --git a/tests/unit/multimodal_utils_test.py b/tests/unit/multimodal_utils_test.py
index 39eb331d03..336401da58 100644
--- a/tests/unit/multimodal_utils_test.py
+++ b/tests/unit/multimodal_utils_test.py
@@ -14,9 +14,11 @@
 
 """Tests for the common MaxText utilities"""
 import os
+import types
 import unittest
 import numpy as np
 
+from maxtext.common.common_types import DecoderBlockType, VisionEncoderBlockType
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_REPO_ROOT
 from maxtext.multimodal import processor as mm_processor
@@ -375,5 +377,129 @@ def test_merge_mm_embeddings_with_token_level_video_mask(self):
     np.testing.assert_array_equal(merged[0, 5:], text_embeddings[0, 5:])
 
 
+class TestMultimodalProcessorRouting(unittest.TestCase):
+  """Tests for multimodal processor prompt/response formatting and training preprocessing."""
+
+  def test_reformat_response_multimodal_vs_text_only(self):
+    # Multimodal Qwen3 model should append <|im_end|>
+    self.assertEqual(mm_processor.reformat_response("Hello world", "qwen3-vl-2b"), "Hello world<|im_end|>")
+    # Text-only Qwen3 model should return response unchanged
+    self.assertEqual(mm_processor.reformat_response("Hello world", "qwen3-4b"), "Hello world")
+    # Multimodal Gemma 3 model should append <end_of_turn>
+    self.assertEqual(mm_processor.reformat_response("Hello world", "gemma3-4b"), "Hello world<end_of_turn>")
+    # Multimodal Llama 4 model should append <|eot|>
+    self.assertEqual(mm_processor.reformat_response("Hello world", "llama4-17b-16e"), "Hello world<|eot|>")
+
+  def test_reformat_prompt_text_only_fallback(self):
+    # Text-only models should return prompt unchanged
+    self.assertEqual(
+        mm_processor.reformat_prompt("Tell me a story", "<img>", "qwen3-4b", num_images=0), "Tell me a story"
+    )
+    self.assertEqual(
+        mm_processor.reformat_prompt("Tell me a story", "<img>", "llama2-7b", num_images=0), "Tell me a story"
+    )
+
+  def test_get_vision_and_decoder_block_routing(self):
+    # pylint: disable=protected-access
+    # Multimodal Qwen3-VL
+    self.assertEqual(mm_processor._get_vision_block("qwen3-vl-2b"), "qwen3_vl")
+    self.assertEqual(mm_processor._get_decoder_block("qwen3-vl-2b"), "qwen3")
+
+    # Multimodal Qwen3.5
+    self.assertEqual(mm_processor._get_vision_block("qwen3.5-35b-a3b"), "qwen3_5")
+    self.assertEqual(mm_processor._get_decoder_block("qwen3.5-35b-a3b"), "qwen3_5")
+
+    # Small Gemma 4 (tests gemma4_small decoder block)
+    self.assertEqual(mm_processor._get_vision_block("gemma4-e4b"), "gemma4")
+    self.assertEqual(mm_processor._get_decoder_block("gemma4-e4b"), "gemma4_small")
+
+    # Text-only Qwen3-4B should have empty vision block and qwen3 decoder block
+    self.assertIsNone(mm_processor._get_vision_block("qwen3-4b"))
+    self.assertEqual(mm_processor._get_decoder_block("qwen3-4b"), "qwen3")
+
+    # Text-only Llama 2 should have empty vision block and llama2 decoder block
+    self.assertIsNone(mm_processor._get_vision_block("llama2-7b"))
+    self.assertEqual(mm_processor._get_decoder_block("llama2-7b"), "llama2")
+
+  def test_get_vision_and_decoder_block_routing_from_config(self):
+    # pylint: disable=protected-access
+    # Test with pyconfig Config objects
+    base_config_path = os.path.join(MAXTEXT_REPO_ROOT, "src", "maxtext", "configs", "base.yml")
+    config_qwen3_vl = pyconfig.initialize(
+        ["", base_config_path],
+        model_name="qwen3-vl-2b",
+        scan_layers=False,
+        skip_jax_distributed_system=True,
+    )
+    self.assertEqual(mm_processor._get_vision_block(config_qwen3_vl), "qwen3_vl")
+    self.assertEqual(mm_processor._get_decoder_block(config_qwen3_vl), "qwen3")
+
+    config_gemma3 = pyconfig.initialize(
+        ["", base_config_path],
+        model_name="gemma3-4b",
+        skip_jax_distributed_system=True,
+    )
+    self.assertEqual(mm_processor._get_vision_block(config_gemma3), "gemma3")
+    self.assertEqual(mm_processor._get_decoder_block(config_gemma3), "gemma3")
+
+    config_text_only = pyconfig.initialize(
+        ["", base_config_path],
+        model_name="qwen3-4b",
+        skip_jax_distributed_system=True,
+    )
+    self.assertIsNone(mm_processor._get_vision_block(config_text_only))
+    self.assertEqual(mm_processor._get_decoder_block(config_text_only), "qwen3")
+
+    # Test with config-like objects covering DecoderBlockType and VisionEncoderBlockType
+    mock_config = types.SimpleNamespace(
+        decoder_block=DecoderBlockType.GEMMA4_SMALL,
+        vision_encoder_block=VisionEncoderBlockType.GEMMA4,
+    )
+    self.assertEqual(mm_processor._get_vision_block(mock_config), "gemma4")
+    self.assertEqual(mm_processor._get_decoder_block(mock_config), "gemma4_small")
+
+    # Test config with DEFAULT decoder_block and NONE vision_encoder_block
+    mock_default = types.SimpleNamespace(
+        decoder_block=DecoderBlockType.DEFAULT,
+        vision_encoder_block=VisionEncoderBlockType.NONE,
+        model_name="custom_model",
+    )
+    self.assertIsNone(mm_processor._get_vision_block(mock_default))
+    self.assertEqual(mm_processor._get_decoder_block(mock_default), "custom_model")
+
+  def test_get_dummy_image_shape_for_init(self):
+    # Multimodal models should return non-empty dummy shape
+    self.assertGreater(len(mm_processor.get_dummy_image_shape_for_init("gemma3-4b")), 0)
+    self.assertGreater(len(mm_processor.get_dummy_image_shape_for_init("gemma4-26b")), 0)
+    self.assertGreater(len(mm_processor.get_dummy_image_shape_for_init("llama4-17b-16e")), 0)
+    self.assertGreater(len(mm_processor.get_dummy_image_shape_for_init("qwen3-vl-2b")), 0)
+
+    # Text-only models should return empty tuple ()
+    self.assertEqual(mm_processor.get_dummy_image_shape_for_init("qwen3-4b"), ())
+    self.assertEqual(mm_processor.get_dummy_image_shape_for_init("llama2-7b"), ())
+
+  def test_preprocess_image_for_training(self):
+    dummy_image = np.zeros((224, 224, 3), dtype=np.uint8)
+    base_config_path = os.path.join(MAXTEXT_REPO_ROOT, "src", "maxtext", "configs", "base.yml")
+
+    # Multimodal model should return non-empty output
+    config_gemma3 = pyconfig.initialize(
+        ["", base_config_path],
+        model_name="gemma3-4b",
+        skip_jax_distributed_system=True,
+    )
+    output = mm_processor.preprocess_image_for_training([dummy_image], config_gemma3)
+    self.assertIsNotNone(output)
+
+    # Text-only model should raise ValueError
+    config_text_only = pyconfig.initialize(
+        ["", base_config_path],
+        model_name="qwen3-4b",
+        skip_jax_distributed_system=True,
+    )
+    with self.assertRaises(ValueError):
+      mm_processor.preprocess_image_for_training([dummy_image], config_text_only)
+
+
 if __name__ == "__main__":
   unittest.main()