diff --git a/src/maxtext/experimental/omni_poc/checkpoint_stitcher/stitch.py b/src/maxtext/experimental/omni_poc/checkpoint_stitcher/stitch.py
new file mode 100644
index 0000000000..68654667f2
--- /dev/null
+++ b/src/maxtext/experimental/omni_poc/checkpoint_stitcher/stitch.py
@@ -0,0 +1,278 @@
+"""Utility to stitch vision and LLM checkpoints into a single unified model.
+
+This script initializes a target multimodal model and restores subtrees from separate checkpoints:
+- restore vision encoder (excluding projector) from a vision checkpoint,
+- restore the decoder and token embedders from a different LLM checkpoint,
+- merge them and save as a single unified MaxText checkpoint.
+
+
+Example usage:
+python -m maxtext.experimental.omni_poc.checkpoint_stitcher.stitch \
+    --vision_load_path=gs://YOUR_BUCKET_NAME/checkpoints/gemma3-4b_converted/0/items \
+    --llm_load_path=gs://YOUR_BUCKET_NAME/checkpoints/qwen3-4b_converted/0/items \
+    --stitched_output_path=gs://YOUR_BUCKET_NAME/checkpoints/omni-gemma3-qwen3-4b/0/items
+"""
+
+import os
+from typing import Any, Dict
+
+from absl import app
+from etils import epath
+from flax import nnx
+import jax
+import omegaconf
+from orbax import checkpoint as ocp
+
+from maxtext.common import checkpointing
+from maxtext.configs import pyconfig as pyconfig_mod
+from maxtext.trainers.pre_train.train import initialize
+from maxtext.utils import max_logging
+from maxtext.utils import max_utils
+from maxtext.utils import maxtext_utils
+from maxtext.utils import maxtext_utils_nnx
+from maxtext.utils import model_creation_utils
+
+
+def _unwrap_var(v):
+  """Unwraps flax nnx Variable instances (e.g. nnx.Param) to raw JAX arrays."""
+  return v.get_value() if hasattr(v, "get_value") else v
+
+
+def _restore_subtrees_from_path(
+    ckpt_path: str, subtrees_abstract: Dict[str, Any], ckptr: ocp.Checkpointer
+) -> Dict[str, Any]:
+  """Restores subtrees from a checkpoint.
+
+  Determines if the source checkpoint uses double params wrapper, single params wrapper,
+  base wrapper, or a flat structure, and restores the matching subtree layout.
+  """
+  metadata = ckptr.metadata(epath.Path(ckpt_path))
+  tree = metadata.item_metadata.tree
+  has_params_params = "params" in tree and isinstance(tree.get("params"), dict) and "params" in tree["params"]
+  has_params = "params" in tree and not has_params_params
+  has_base = "base" in tree
+
+  if has_params_params:
+    item = {"params": {"params": subtrees_abstract}}
+    restore_args = {"params": {"params": ocp.checkpoint_utils.construct_restore_args(subtrees_abstract)}}
+  elif has_params:
+    item = {"params": subtrees_abstract}
+    restore_args = {"params": ocp.checkpoint_utils.construct_restore_args(subtrees_abstract)}
+  elif has_base:
+    item = {"base": subtrees_abstract}
+    restore_args = {"base": ocp.checkpoint_utils.construct_restore_args(subtrees_abstract)}
+  else:
+    item = subtrees_abstract
+    restore_args = ocp.checkpoint_utils.construct_restore_args(subtrees_abstract)
+
+  restored = ckptr.restore(
+      epath.Path(ckpt_path),
+      item=item,
+      transforms={},
+      restore_args=restore_args,
+  )
+  if has_params_params:
+    return restored["params"]["params"]
+  elif has_params:
+    return restored["params"]
+  elif has_base:
+    return restored["base"]
+  else:
+    return restored
+
+
+def _assemble(k: str, v: Any, stitched_subtrees: Dict[str, Any]) -> Any:
+  """Merges restored parameter subtrees with fresh target model initial values.
+
+  If a sub-module name is present in stitched_subtrees, we restore it.
+  If it's missing (e.g. the new vision projector), we keep the fresh random initialization.
+  """
+  if k in stitched_subtrees:  # e.g., k is vision_encoder, decoder, or token_embedder
+    restored_val = stitched_subtrees[k]
+    if isinstance(restored_val, dict) and isinstance(v, dict):
+      # Merge sub-modules inside this namespace
+      merged_module = {}
+      for sub_module_name, fresh_weights in v.items():
+        if sub_module_name in restored_val:
+          # If the sub-module (e.g. Gemma3VisionEncoderLayer_0) exists in the checkpoint, load it
+          merged_module[sub_module_name] = restored_val[sub_module_name]
+        else:
+          # If the sub-module is missing from the checkpoint (e.g. the new projector), keep its fresh random weights
+          merged_module[sub_module_name] = fresh_weights
+      return merged_module
+    # k is pointing to a single tensor
+    return restored_val
+
+  # k is a new layer (not in stitched_subtrees), keep fresh random init
+  max_logging.log(f"Keeping fresh random normal initialization for new layer: '{k}'")
+  return v
+
+
+def stitch_and_save_checkpoints(
+    config: Any,
+    vision_checkpoint_path: str,
+    llm_checkpoint_path: str,
+    output_checkpoint_path: str,
+):
+  """Stitches vision model weights and LLM weights into one MaxText checkpoint.
+
+  Args:
+    config: The MaxText target model configuration.
+    vision_checkpoint_path: Path to the vision model checkpoint.
+    llm_checkpoint_path: Path to the LLM checkpoint.
+    output_checkpoint_path: Path to save the stitched checkpoint.
+  """
+  max_logging.log("=" * 60)
+  max_logging.log("Starting Omni Multi-Directory Checkpoint Stitching...")
+
+  vision_model_name = getattr(config, "model_name", None)
+  llm_model_name = getattr(config, "decoder_block", None)
+  assert vision_model_name, "model_name must be configured for vision component."
+  assert llm_model_name, "decoder_block must be configured for LLM component."
+
+  max_logging.log(f"  Vision (Model {vision_model_name}) Path: {vision_checkpoint_path}")
+  max_logging.log(f"  LLM (Model {llm_model_name}) Path:    {llm_checkpoint_path}")
+  max_logging.log(f"  Output Stitched Path:  {output_checkpoint_path}")
+  max_logging.log("=" * 60)
+
+  mesh = maxtext_utils.get_mesh_from_config(config)
+  init_rng = jax.random.PRNGKey(config.init_weights_seed)
+
+  # 1. Generate full target model with initial random weights
+  max_logging.log("Generating target omni model from config with initial random weights...")
+  with jax.set_mesh(mesh):
+    if config.pure_nnx:
+      rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng)
+      model = model_creation_utils.from_config(config, mesh=mesh, rngs=rngs)
+      init_params = nnx.state(model, nnx.Param)
+    else:
+      model = model_creation_utils.from_config(config, jax.devices())
+      _, _, init_params = maxtext_utils.init_initial_state(model, None, config, is_training=False, init_rng=init_rng)
+
+  # Convert to pure pytree for easier processing
+  is_nnx = isinstance(init_params, nnx.State)
+  params_dict = init_params.to_pure_dict() if is_nnx else init_params
+  inner_params = params_dict.get("params", params_dict)
+
+  inner_params = jax.tree.map(_unwrap_var, inner_params, is_leaf=lambda n: isinstance(n, nnx.Variable))
+
+  # Wrap pytree into a Checkpointer object for partial checkpoint restoration
+  ckptr = ocp.Checkpointer(
+      ocp.PyTreeCheckpointHandler(
+          restore_concurrent_gb=config.checkpoint_storage_concurrent_gb,
+          save_concurrent_gb=config.checkpoint_storage_concurrent_gb,
+          use_ocdbt=config.checkpoint_storage_use_ocdbt,
+          use_zarr3=config.checkpoint_storage_use_zarr3,
+      )
+  )
+
+  stitched_subtrees = {}
+
+  # 2. Restore Vision Encoder subtree from Model A
+  if "vision_encoder" in inner_params and vision_checkpoint_path:
+    max_logging.log(f"Restoring 'vision_encoder' from {vision_checkpoint_path}...")
+    # Filter out projector/embedder keys from the abstract state so they are not loaded from disk
+    vision_encoder_abstract = {
+        k: v
+        for k, v in inner_params["vision_encoder"].items()
+        if not ("projector" in k.lower() or "embedder" in k.lower())
+    }
+    vision_abstract = {"vision_encoder": vision_encoder_abstract}
+    vision_restored = _restore_subtrees_from_path(vision_checkpoint_path, vision_abstract, ckptr)
+    stitched_subtrees["vision_encoder"] = vision_restored["vision_encoder"]
+
+  # 3. Restore LLM Decoder subtrees from Model B
+  llm_keys = [k for k in ["decoder", "token_embedder"] if k in inner_params]
+  if llm_keys and llm_checkpoint_path:
+    max_logging.log(f"Restoring LLM subtrees ({llm_keys}) from {llm_checkpoint_path}...")
+    llm_abstract = {k: inner_params[k] for k in llm_keys}
+    llm_restored = _restore_subtrees_from_path(llm_checkpoint_path, llm_abstract, ckptr)
+    for k in llm_keys:
+      stitched_subtrees[k] = llm_restored[k]
+
+  # 4. Assemble: Vision (Model A) + LLM (Model B) + Random Init Projector
+  stitched_inner = {k: _assemble(k, v, stitched_subtrees) for k, v in inner_params.items()}
+  final_params = (
+      {"params": stitched_inner}
+      if "params" in params_dict and isinstance(params_dict["params"], dict)
+      else stitched_inner
+  )
+
+  # 5. Save unified parameter tree to output_checkpoint_path
+  max_logging.log(f"Saving stitched checkpoint to: {output_checkpoint_path}")
+  checkpointing.save_params_to_path(
+      output_checkpoint_path,
+      final_params,
+      use_ocdbt=config.checkpoint_storage_use_ocdbt,
+      use_zarr3=config.checkpoint_storage_use_zarr3,
+  )
+  total_params = max_utils.calculate_num_params_from_pytree(final_params)
+  max_logging.log(f"Total Stitched Model Parameters: {total_params:,} (~{total_params/1e9:.3f}B)")
+  max_logging.log("Checkpoint stitching complete!")
+
+
+def _load_custom_yaml_overrides(yaml_path: str, omni_keys: set[str]):
+  """Loads a custom YAML config and splits it into omni-specific keys and MaxText overrides."""
+  custom_cfg = omegaconf.OmegaConf.to_container(omegaconf.OmegaConf.load(yaml_path), resolve=True)
+
+  omni_yaml_args = {}
+  maxtext_overrides = {}
+  for key, value in custom_cfg.items():
+    if key in omni_keys:
+      omni_yaml_args[key] = value
+    else:
+      maxtext_overrides[key] = value
+
+  return omni_yaml_args, maxtext_overrides
+
+
+def main(argv):
+  # Extract omni stitching arguments directly from argv before passing to pyconfig.initialize
+  omni_keys = {"vision_load_path", "llm_load_path", "stitched_output_path", "vision_model_name", "llm_model_name"}
+  omni_kwargs = {}
+  cleaned_argv = []
+  for arg in argv:
+    if "=" in arg and arg.split("=", 1)[0] in omni_keys:
+      k, v = arg.split("=", 1)
+      omni_kwargs[k] = v
+    else:
+      cleaned_argv.append(arg)
+
+  # To populate all system-wide defaults, MaxText requires base.yml as argv[1].
+  # To apply our custom overrides on top of these defaults, we convert the custom config
+  # overrides into cleaned_argv for initialization.
+  if len(cleaned_argv) >= 2 and cleaned_argv[1].endswith(".yml") and not cleaned_argv[1].endswith("base.yml"):
+    custom_yaml_path = cleaned_argv[1]
+
+    # Load and split custom settings
+    yaml_omni_args, yaml_overrides = _load_custom_yaml_overrides(custom_yaml_path, omni_keys)
+
+    # Merge settings
+    for k, v in yaml_omni_args.items():
+      omni_kwargs.setdefault(k, v)
+
+    # Convert YAML overrides to CLI-style arguments for standard initialize()
+    for k, v in yaml_overrides.items():
+      if isinstance(v, str):
+        cleaned_argv.append(f"{k}='{v}'")
+      else:
+        cleaned_argv.append(f"{k}={v}")
+    cleaned_argv.append("override_model_config=True")
+
+    cleaned_argv[1] = os.path.join(pyconfig_mod.MAXTEXT_CONFIGS_DIR, "base.yml")
+
+  # Initialize MaxText config using standard train.initialize
+  config, _ = initialize(cleaned_argv)
+  # Extract paths from command-line arguments
+  vision_path = omni_kwargs.get("vision_load_path")
+  llm_path = omni_kwargs.get("llm_load_path")
+  output_path = omni_kwargs.get("stitched_output_path")
+  assert (
+      vision_path and llm_path and output_path
+  ), "Must specify vision_load_path, llm_load_path, and stitched_output_path"
+
+  stitch_and_save_checkpoints(config, vision_path, llm_path, output_path)
+
+
+if __name__ == "__main__":
+  app.run(main)
diff --git a/src/maxtext/experimental/omni_poc/checkpoint_stitcher/unittest_stitch.py b/src/maxtext/experimental/omni_poc/checkpoint_stitcher/unittest_stitch.py
new file mode 100644
index 0000000000..f8030b59f5
--- /dev/null
+++ b/src/maxtext/experimental/omni_poc/checkpoint_stitcher/unittest_stitch.py
@@ -0,0 +1,448 @@
+"""Unit tests for checkpoint stitching of Vision and LLM backbones.
+
+This test suite verifies the correctness of `stitch.py` by
+stitching separate Vision (e.g., Gemma 3) and LLM (e.g., Qwen 3) checkpoints
+into a combined multimodal model checkpoint. 
+
+Two test cases:
+
+1. Model layer count matching:
+   Validates that the stitched checkpoint's structural layer counts for
+   both the restored vision encoder blocks and LLM blocks exactly
+   match the counts from the original source checkpoints.
+
+2. Forward logits comparison:
+   Performs forward passes on individual model components (Vision Tower,
+   Token Embedder, and Projector) and compares output representations between
+   weights restored from the original checkpoints vs. the stitched checkpoint.
+   It asserts that:
+   - Vision encoder outputs match exactly;
+   - Text embedder outputs match exactly;
+   - Projector outputs differ (since the new projector is
+     intentionally initialized with random weights).
+"""
+
+import gc
+import os
+import unittest
+import warnings
+
+from etils import epath
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+import numpy as np
+import omegaconf
+from orbax import checkpoint as ocp
+
+from maxtext.configs import pyconfig
+from maxtext.experimental.omni_poc.checkpoint_stitcher import stitch
+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
+
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+
+def _unwrap_params(restored_dict):
+  """Fully unwrap any top-level 'params' / 'base' dictionary wrappers."""
+  while isinstance(restored_dict, dict) and len(restored_dict) == 1:
+    if "params" in restored_dict:
+      restored_dict = restored_dict["params"]
+    elif "base" in restored_dict:
+      restored_dict = restored_dict["base"]
+    else:
+      break
+  return restored_dict
+
+
+def _get_vision_layer_count(vision_state):
+  """Counts unique encoder block names within the vision encoder state."""
+  blocks = set()
+
+  def recurse(d):
+    if isinstance(d, dict):
+      for k, v in d.items():
+        if k.startswith("encoderblock_"):
+          blocks.add(k)
+        recurse(v)
+
+  recurse(vision_state)
+  return len(blocks)
+
+
+class TestOmniCheckpointStitcher(unittest.TestCase):
+  """Unit tests for Omni checkpoint stitching logic.
+
+  Verifies correct weight assembly and projector alignment on full models.
+  """
+
+  def setUp(self):
+    super().setUp()
+    # Read paths from environment variables or default to standard GCS bucket
+    default_dir = "gs://yuchenhou-maxtext-logs/omni_checkpoints"
+    base_dir = os.environ.get("OMNI_TEST_BASE_DIR", default_dir)
+    self.vision_ckpt_dir = f"{base_dir}/gemma3-4b_converted/0/items"
+    self.llm_ckpt_dir = f"{base_dir}/qwen3-4b_converted/0/items"
+    self.output_ckpt_dir = f"{base_dir}/omni_stitched_gemma3-4b_qwen3-4b/0/items"
+
+    # Run tests on CPU mesh to avoid running out of TPU device memory
+    cpu_device = jax.devices("cpu")[0]
+    self.test_mesh = Mesh(np.array([cpu_device]), axis_names=("data",))
+
+  def tearDown(self):
+    super().tearDown()
+    gc.collect()
+
+  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",
+        "experimental",
+        "omni_poc",
+        "omni-gemma3-qwen3.yml",
+    )
+    custom_cfg = omegaconf.OmegaConf.to_container(omegaconf.OmegaConf.load(omni_config_path), resolve=True)
+
+    # Filter out safeguard/stitching-specific keys
+    omni_keys = {
+        "vision_load_path",
+        "llm_load_path",
+        "stitched_output_path",
+        "vision_model_name",
+        "llm_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")
+    config = pyconfig.initialize(
+        ["", base_config_path],
+        override_model_config=True,
+        skip_jax_distributed_system=True,
+        log_config=False,
+        pure_nnx=False,
+        attention="dot_product",
+        dtype="float32",
+        **yaml_overrides,
+    )
+    # Use object.__setattr__ to bypass the read-only check on _HyperParameters
+    object.__setattr__(config, "vision_model_name", "gemma3-4b")
+    object.__setattr__(config, "llm_model_name", "qwen3-4b")
+    object.__setattr__(config, "ici_context_autoregressive_parallelism", 1)
+    return config
+
+  def _ensure_stitched_checkpoint(self):
+    """Runs stitching if the stitched output checkpoint does not exist."""
+    if not epath.Path(self.output_ckpt_dir).exists():
+      print(f"Stitched checkpoint not found at {self.output_ckpt_dir}. " "Running stitching process...")
+      config = self._create_config()
+      stitch.stitch_and_save_checkpoints(
+          config=config,
+          vision_checkpoint_path=self.vision_ckpt_dir,
+          llm_checkpoint_path=self.llm_ckpt_dir,
+          output_checkpoint_path=self.output_ckpt_dir,
+      )
+
+  def _restore_checkpoints(self, config):
+    """Restores stitched, vision, and LLM checkpoints under the mesh."""
+    # pylint: disable=protected-access
+    ckptr = ocp.Checkpointer(
+        ocp.PyTreeCheckpointHandler(
+            restore_concurrent_gb=config.checkpoint_storage_concurrent_gb,
+            save_concurrent_gb=config.checkpoint_storage_concurrent_gb,
+            use_ocdbt=config.checkpoint_storage_use_ocdbt,
+            use_zarr3=config.checkpoint_storage_use_zarr3,
+        )
+    )
+
+    with jax.set_mesh(self.test_mesh):
+      # Create the target model structure
+      Linen_model = model_creation_utils.from_config(config, mesh=self.test_mesh)
+
+      # Get model info (shapes, data types, and sharding layouts)
+      abstract_vars = maxtext_utils.get_abstract_param(Linen_model, config)
+      target_params_abstract = abstract_vars["params"]
+
+      # Extract the raw PyTree of ShapeDtypeStruct nodes
+      target_params_abstract = max_utils.unbox_logicallypartioned(target_params_abstract)
+
+      # Create a concrete "zero-filled template"
+      # Orbax checkpointer requires target JAX arrays to know what shapes and
+      # memory allocations to load checkpoints into. We map the abstract
+      # shape/dtype structs to concrete zero arrays on our mesh.
+      def to_concrete_zeros(leaf):
+        if isinstance(leaf, jax.ShapeDtypeStruct):
+          return jnp.zeros(leaf.shape, dtype=leaf.dtype)
+        return leaf
+
+      concrete_template = jax.tree.map(to_concrete_zeros, target_params_abstract)
+
+      # Load weights from the STITCHED checkpoint path
+      stitched_inner = stitch._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_restored = stitch._restore_subtrees_from_path(
+          self.vision_ckpt_dir, {"vision_encoder": vision_abstract}, ckptr
+      )
+
+      # Load weights from the ORIGINAL LLM checkpoint
+      llm_abstract = {
+          "decoder": concrete_template["decoder"],
+          "token_embedder": concrete_template["token_embedder"],
+      }
+      llm_restored = stitch._restore_subtrees_from_path(self.llm_ckpt_dir, llm_abstract, ckptr)
+    return (
+        stitched_inner,
+        vision_restored,
+        llm_restored,
+        Linen_model,
+        concrete_template,
+    )
+
+  def test_1_stitch_and_assemble_correctness(self):
+    """Test 1: Verifies model component layer counts match."""
+    self._ensure_stitched_checkpoint()
+    config = self._create_config()
+
+    (
+        stitched_inner,
+        vision_restored,
+        llm_restored,
+        Linen_model,
+        concrete_template,
+    ) = self._restore_checkpoints(config)
+
+    # Verify layer counts
+    orig_decoder_layers = llm_restored["decoder"]["layers"]["mlp"]["wo"]["kernel"].shape[1]
+    stitched_decoder_layers = stitched_inner["decoder"]["layers"]["mlp"]["wo"]["kernel"].shape[1]
+    self.assertEqual(stitched_decoder_layers, orig_decoder_layers)
+
+    orig_vision_layers = _get_vision_layer_count(vision_restored["vision_encoder"]["Gemma3VisionEncoderLayer_0"])
+    stitched_vision_layers = _get_vision_layer_count(stitched_inner["vision_encoder"]["Gemma3VisionEncoderLayer_0"])
+    self.assertEqual(stitched_vision_layers, orig_vision_layers)
+
+    print("\n" + "=" * 80)
+    print("TEST 1: MODEL LAYER COUNT MATCHING")
+    print("=" * 80)
+    print(
+        f"  - Verified LLM layers count ({config.llm_model_name}): "
+        f"orig={orig_decoder_layers} and stitched={stitched_decoder_layers}"
+    )
+    print(
+        f"  - Verified Vision layers count ({config.vision_model_name}): "
+        f"orig={orig_vision_layers} and stitched={stitched_vision_layers}"
+    )
+    print("=" * 80 + "\n")
+
+    del (
+        stitched_inner,
+        vision_restored,
+        llm_restored,
+        Linen_model,
+        concrete_template,
+    )
+    gc.collect()
+
+  def test_2_stage_outputs_original_vs_stitched(self):
+    """Test 2: Verifies output logits match original components."""
+    self._ensure_stitched_checkpoint()
+    config = self._create_config()
+    rngs = nnx.Rngs(0)
+
+    (
+        stitched_inner,
+        vision_restored,
+        llm_restored,
+        Linen_model,
+        concrete_template,
+    ) = self._restore_checkpoints(config)
+
+    print("\n" + "=" * 80)
+    print("TEST 2: FORWARD LOGITS COMPARISON")
+    print("=" * 80)
+
+    v_layer = vision_restored["vision_encoder"]["Gemma3VisionEncoderLayer_0"]
+    s_layer = stitched_inner["vision_encoder"]["Gemma3VisionEncoderLayer_0"]
+
+    print(
+        "orig vision embedding kernel shape:",
+        v_layer["embedding"]["kernel"].shape,
+    )
+    print(
+        "orig vision embedding kernel first 10 values:",
+        v_layer["embedding"]["kernel"].reshape(-1)[:10],
+    )
+    print(
+        "stitched vision embedding kernel shape:",
+        s_layer["embedding"]["kernel"].shape,
+    )
+    print(
+        "stitched vision embedding kernel first 10 values:",
+        s_layer["embedding"]["kernel"].reshape(-1)[:10],
+    )
+
+    print(
+        "orig llm token embedder shape:",
+        llm_restored["token_embedder"]["embedding"].shape,
+    )
+    print(
+        "orig llm token embedder first 10 values:",
+        llm_restored["token_embedder"]["embedding"].reshape(-1)[:10],
+    )
+    print(
+        "stitched llm token embedder shape:",
+        stitched_inner["token_embedder"]["embedding"].shape,
+    )
+    print(
+        "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],
+    )
+
+    # 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)
+    # Manually define qwen3's token embedder to match MaxText's structures
+    orig_token_embedder = Embed(
+        num_embeddings=config.vocab_size,
+        num_features=config.emb_dim,
+        config=config,
+        mesh=self.test_mesh,
+        rngs=rngs,
+    )
+
+    nnx.update(
+        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_token_embedder = Embed(
+        num_embeddings=config.vocab_size,
+        num_features=config.emb_dim,
+        config=config,
+        mesh=self.test_mesh,
+        rngs=rngs,
+    )
+
+    nnx.update(
+        stitched_vision_tower,
+        stitched_inner["vision_encoder"]["Gemma3VisionEncoderLayer_0"],
+    )
+    nnx.update(
+        stitched_projector,
+        stitched_inner["vision_encoder"]["VisionEmbedder_0"],
+    )
+    nnx.update(
+        stitched_token_embedder,
+        {"embedding": stitched_inner["token_embedder"]["embedding"]},
+    )
+
+    # Inputs
+    batch_size, prompt_len = 2, 128
+    key1, key2 = jax.random.split(jax.random.PRNGKey(42))
+    image_size = config.image_size_for_vit
+    input_images = jax.random.normal(key1, (batch_size, image_size, image_size, 3), dtype=jnp.float32)
+    input_tokens = jax.random.randint(key2, (batch_size, prompt_len), 0, config.vocab_size)
+
+    with jax.set_mesh(self.test_mesh):
+      # Forward pass: Vision tower only (without projector)
+      orig_vision_features = orig_vision_tower(input_images, deterministic=True)
+      stitched_vision_features = stitched_vision_tower(input_images, deterministic=True)
+
+      # Forward pass: Projector
+      orig_projected_tokens = orig_projector(orig_vision_features)
+      if orig_projected_tokens.ndim == 4 and orig_projected_tokens.shape[1] == 1:
+        orig_projected_tokens = jnp.squeeze(orig_projected_tokens, axis=1)
+
+      stitched_projected_tokens = stitched_projector(stitched_vision_features)
+      if stitched_projected_tokens.ndim == 4 and stitched_projected_tokens.shape[1] == 1:
+        stitched_projected_tokens = jnp.squeeze(stitched_projected_tokens, axis=1)
+
+      # Forward pass: Text embedder
+      orig_text_embeddings = orig_token_embedder(input_tokens)
+      stitched_text_embeddings = stitched_token_embedder(input_tokens)
+
+    # Vision output logits comparison
+    np.testing.assert_allclose(orig_vision_features, stitched_vision_features, atol=1e-5)
+    # Text embeddings output logits comparison
+    np.testing.assert_allclose(orig_text_embeddings, stitched_text_embeddings, atol=1e-5)
+    # Projector outputs logits comparison
+    diff_stage3 = jnp.max(jnp.abs(orig_projected_tokens - stitched_projected_tokens))
+    self.assertGreater(diff_stage3, 1.0)
+
+    # Calculate statistics for logging
+    def get_stats(orig, stitched):
+      abs_diff = jnp.abs(orig - stitched)
+      return (
+          jnp.mean(jnp.abs(orig)),
+          jnp.mean(jnp.abs(stitched)),
+          jnp.mean(abs_diff),
+          jnp.max(abs_diff),
+      )
+
+    orig_v_mean, stitched_v_mean, mean_v_diff, max_v_diff = get_stats(orig_vision_features, stitched_vision_features)
+    orig_emb_mean, stitched_emb_mean, mean_emb_diff, max_emb_diff = get_stats(
+        orig_text_embeddings, stitched_text_embeddings
+    )
+    orig_proj_mean, stitched_proj_mean, mean_proj_diff, max_proj_diff = get_stats(
+        orig_projected_tokens, stitched_projected_tokens
+    )
+
+    print(
+        f"  - Vision encoder output logits mean of abs: orig={orig_v_mean:.2f}, "
+        f"stitched={stitched_v_mean:.2f} (Mean abs diff = {mean_v_diff:.2e}, max = {max_v_diff:.2e})"
+    )
+    print(
+        f"  - Token embedder output logits mean of abs: "
+        f"orig={orig_emb_mean:.2f}, stitched={stitched_emb_mean:.2f} "
+        f"(Mean abs diff = {mean_emb_diff:.2e}, max = {max_emb_diff:.2e})"
+    )
+    print(
+        f"  - Projector output logits mean of abs:     "
+        f"orig={orig_proj_mean:.2f}, stitched={stitched_proj_mean:.2f} "
+        f"(Mean abs diff = {mean_proj_diff:.2e}, max = {max_proj_diff:.2e})"
+    )
+    print("=" * 80 + "\n")
+
+    del (
+        stitched_inner,
+        vision_restored,
+        llm_restored,
+        Linen_model,
+        concrete_template,
+    )
+    del orig_vision_tower, orig_projector, orig_token_embedder
+    del stitched_vision_tower, stitched_projector, stitched_token_embedder
+    gc.collect()
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/src/maxtext/experimental/omni_poc/omni-gemma3-qwen3.yml b/src/maxtext/experimental/omni_poc/omni-gemma3-qwen3.yml
new file mode 100644
index 0000000000..284c8abcbf
--- /dev/null
+++ b/src/maxtext/experimental/omni_poc/omni-gemma3-qwen3.yml
@@ -0,0 +1,29 @@
+# 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"
+use_multimodal: true
+
+# Load LLM Backbone Config (from qwen3-4b.yml)
+base_emb_dim: 2560
+base_num_query_heads: 32
+base_num_kv_heads: 8
+base_mlp_dim: 9728
+base_num_decoder_layers: 36
+head_dim: 128
+mlp_activations: ["silu", "linear"]
+vocab_size: 151936
+decoder_block: "qwen3"
+normalization_layer_epsilon: 1.0e-6
+rope_max_timescale: 1000000
+use_qk_norm: true
+logits_via_embedding: true
+normalize_embedding_logits: false
+enable_dropout: false
+tokenizer_type: "huggingface"
+
+# Ensure no pretrained weights or checkpoints are loaded initially
+load_parameters_path: ""
+
diff --git a/src/maxtext/experimental/omni_poc/prepare_checkpoint.py b/src/maxtext/experimental/omni_poc/prepare_checkpoint.py
new file mode 100644
index 0000000000..5346917f84
--- /dev/null
+++ b/src/maxtext/experimental/omni_poc/prepare_checkpoint.py
@@ -0,0 +1,147 @@
+# -*- coding: utf-8 -*-
+"""Step 1 of 5 maxtext multimodal alignment proof of concept project.
+
+This file reads the omni-gemma3-qwen3.yml configuration and 
+  1. download and convert a Vision-Language model (e.g. gemma3-4b) from hugging face to maxtext format
+  2. download and convert an text-only model (e.g. qwen3-4b) from hugging face to maxtext format
+  3. stitch the vision component and llm checkpoints into a single omni checkpoint
+  4. save the stitched checkpoints to a output directory
+
+Example usage:
+python maxtext/experimental/omni_poc/prepare_checkpoint.py
+"""
+import os
+import subprocess
+import sys
+from etils import epath
+from maxtext.utils.globals import MAXTEXT_PKG_DIR
+
+
+def main():
+
+  # Hugging Face Token & Login. Can be set via HF_TOKEN environment variable.
+  hf_token = os.environ.get("HF_TOKEN", "")
+
+  # Base GCS or local directory where converted and stitched checkpoints will be stored.
+  base_output_directory = "gs://YOUR_BUCKET_NAME/omni_checkpoints"
+
+  if hf_token:
+    try:
+      from huggingface_hub import login  # pylint: disable=import-outside-toplevel
+
+      login(token=hf_token)
+    except ImportError:
+      print("huggingface_hub not installed. Skipping login.")
+    except Exception as e:  # pylint: disable=broad-exception-caught
+      print(f"Hugging Face login failed: {e}")
+      print("Ensure HF_TOKEN is valid if needed.\n")
+
+  # Configuration & Paths
+  vision_maxtext_model = "gemma3-4b"
+  vision_hf_repo = "google/gemma-3-4b-it"
+
+  llm_maxtext_model = "qwen3-4b"
+  llm_hf_repo = "Qwen/Qwen3-4B"
+
+  # Path to combined YAML configuration inside experimental/omni_poc
+  omni_config_path = os.path.join(MAXTEXT_PKG_DIR, "experimental", "omni_poc", "omni-gemma3-qwen3.yml")
+
+  vision_ckpt_dir = f"{base_output_directory}/{vision_maxtext_model}_converted"
+  llm_ckpt_dir = f"{base_output_directory}/{llm_maxtext_model}_converted"
+  stitched_ckpt_dir = f"{base_output_directory}/omni_stitched_{vision_maxtext_model}_{llm_maxtext_model}"
+
+  vision_items_path = epath.Path(vision_ckpt_dir) / "0/items"
+  llm_items_path = epath.Path(llm_ckpt_dir) / "0/items"
+  stitched_items_path = epath.Path(stitched_ckpt_dir) / "0/items"
+
+  print(f"Base Output Directory:  {base_output_directory}")
+  print(f"Vision Converted Path:  {vision_items_path}")
+  print(f"LLM Converted Path:     {llm_items_path}")
+  print(f"Stitched Target Path:   {stitched_items_path}\n")
+
+  env = os.environ.copy()
+  env["JAX_PLATFORMS"] = "cpu"  # Conversion and stitching run smoothly on CPU/TPU
+
+  # Step 1: Download & Convert Vision Model from Hugging Face -> MaxText
+  print("=" * 60)
+  if not vision_items_path.exists():
+    print(f"Converting Vision Model ({vision_maxtext_model}) from Hugging Face ({vision_hf_repo})...")
+    try:
+      subprocess.run(
+          [
+              sys.executable,
+              "-m",
+              "maxtext.checkpoint_conversion.to_maxtext",
+              os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml"),
+              f"model_name={vision_maxtext_model}",
+              f"base_output_directory={vision_ckpt_dir}",
+              "use_multimodal=True",
+              "scan_layers=True",
+              "skip_jax_distributed_system=True",
+              "--eager_load_method=transformers",
+              "--lazy_load_tensors=False",
+              "log_config=False",
+          ],
+          check=True,
+          env=env,
+      )
+      print("Vision checkpoint conversion successful!\n")
+    except subprocess.CalledProcessError as e:
+      print(f"Failed to convert Vision checkpoint: {e}")
+      sys.exit(1)
+  else:
+    print(f"Step 1: Vision checkpoint already exists at {vision_items_path}")
+
+  # Step 2: Download & Convert Language Model from Hugging Face -> MaxText
+  print("=" * 60)
+  if not llm_items_path.exists():
+    print(f"Converting Language Model ({llm_maxtext_model}) from Hugging Face ({llm_hf_repo})...")
+    try:
+      subprocess.run(
+          [
+              sys.executable,
+              "-m",
+              "maxtext.checkpoint_conversion.to_maxtext",
+              os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml"),
+              f"model_name={llm_maxtext_model}",
+              f"base_output_directory={llm_ckpt_dir}",
+              "scan_layers=True",
+              "skip_jax_distributed_system=True",
+              "--eager_load_method=transformers",
+              "--lazy_load_tensors=False",
+              "log_config=False",
+          ],
+          check=True,
+          env=env,
+      )
+      print("LLM checkpoint conversion successful!\n")
+    except subprocess.CalledProcessError as e:
+      print(f"Failed to convert LLM checkpoint: {e}")
+      sys.exit(1)
+  else:
+    print(f"Step 2: LLM checkpoint already exists at {llm_items_path}")
+
+  # Step 3: Checkpoint Stitching (Vision Tower + LLM Decoder + Fresh Projector)
+  print("=" * 60)
+  print("Stitching Vision and LLM subtrees into unified Omni checkpoint...")
+  try:
+    subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "maxtext.experimental.omni_poc.checkpoint_stitcher.stitch",
+            omni_config_path,
+            f"vision_load_path={str(vision_items_path)}",
+            f"llm_load_path={str(llm_items_path)}",
+            f"stitched_output_path={str(stitched_items_path)}",
+        ],
+        check=True,
+        env=env,
+    )
+  except subprocess.CalledProcessError as e:
+    print(f"Failed during checkpoint stitching: {e}")
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+  main()