From 840d86ea341ad898b540ea0f70da5fbda776b2d6 Mon Sep 17 00:00:00 2001
From: ruiren_microsoft <ruiren@microsoft.com>
Date: Wed, 22 Apr 2026 16:02:29 -0700
Subject: [PATCH 1/6] add the audio attention

---
 src/mobius/components/_qwen25_omni_audio.py | 115 ++++++++++++++++++++
 src/mobius/models/qwen25_omni.py            |  29 +++++
 2 files changed, 144 insertions(+)
 create mode 100644 src/mobius/components/_qwen25_omni_audio.py
 create mode 100644 src/mobius/models/qwen25_omni.py

diff --git a/src/mobius/components/_qwen25_omni_audio.py b/src/mobius/components/_qwen25_omni_audio.py
new file mode 100644
index 00000000..38556e3a
--- /dev/null
+++ b/src/mobius/components/_qwen25_omni_audio.py
@@ -0,0 +1,115 @@
+"""Qwen25-Omni audio encoder components.
+
+Whisper-inspired audio encoder with 3x Conv1d,
+sinusoidal positional embeddings, and bidirectional transformer
+encoder layers with LayerNorm.
+
+Reference: Transformers
+https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from onnxscript import nn
+from onnxscript._internal import builder
+
+from mobius.components._common import LayerNorm, Linear
+
+
+if TYPE_CHECKING:
+    import onnx_ir as ir
+
+class Qwen25OmniAudioAttention(nn.Module):
+    """Bidirectional multi-head attention for Qwen2_5Omni audio encoder.
+
+    Unlike WhisperAttention, all projections (Q, V, Out) have bias and K does not have bias.
+    No causal masking — the encoder uses full bidirectional attention.
+    """
+    def __init__(self, d_model: int, num_heads: int):
+        super().__init__()
+        self.q_proj = Linear(d_model, d_model, bias=True)
+        self.k_proj = Linear(d_model, d_model, bias=False)
+        self.v_proj = Linear(d_model, d_model, bias=True)
+        self.out_proj = Linear(d_model, d_model, bias=True)
+        self._num_heads = num_heads
+        self._head_dim = d_model // num_heads
+
+    def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
+        """Bidirectional self-attention
+
+        Args:
+            hidden_states: (batch, seq_len, d_model)
+
+        Returns:
+            output: (batch, seq_len, d_model)
+        """
+
+        q = self.q_proj(op, hidden_states)
+        k = self.k_proj(op, hidden_states)
+        v = self.v_proj(op, hidden_states)
+
+        # Use ONNX Attention op (bidirectional: no causal mask)
+        attn_output = op.Attention(
+            q,
+            k,
+            v,
+            q_num_heads = self._num_heads,
+            kv_num_heads = self._num_heads,
+            scale=float(self._head_dim**-0.5),
+        )
+        return self.out_proj(op, attn_output)
+
+
+class Qwen25OmniAudioEncoderLayer(nn.Module):
+    """Qwen25-Omni audio encoder layer.
+
+    Pre-norm pattern:  LayerNorm → self-attn → residual
+    → LayerNorm → FFN → residual.
+    Uses GELU activation in the FFN.
+    
+    Huggingface class: ``Qwen2_5OmniAudioEncoder``
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        ffn_dim: int,
+        eps: float = 1e-5,
+    ):
+        super().__init__()
+        self.self_attn = Qwen25OmniAudioAttention(d_model, num_heads)
+        self.self_attn_layer_norm = LayerNorm(d_model, eps=eps)
+        self.fc1 = Linear(d_model, ffn_dim, bias=True)
+        self.fc2 = Linear(ffn_dim, d_model, bias=True)
+        self.final_layer_norm = LayerNorm(d_model, eps=eps)
+
+    def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
+        """Pre-norm encoder layer with bidirectional attention.
+
+        Args:
+            hidden_states: (batch, seq_len, d_model)
+
+        returns:
+            hidden_states: (batch, seq_len, d_model)
+        """
+        # Self-attention with pre-norm and residual
+        residual = hidden_states
+        hidden_states = self.self_attn_layer_norm(op, hidden_states)
+        hidden_states = self.self_attn(op, hidden_states)
+        hidden_states = op.Add(residual, hidden_states)
+
+        # FFN with pre-norm, GELU, and residual
+        residual = hidden_states
+        hidden_states = self.final_layer_norm(op, hidden_states)
+        hidden_states = self.fc1(op, hidden_states)
+        hidden_states = op.Gelu(hidden_states)
+        hidden_states = self.fc2(op, hidden_states)
+        hidden_states = op.Add(residual, hidden_states)
+
+        return hidden_states
+
+
+
diff --git a/src/mobius/models/qwen25_omni.py b/src/mobius/models/qwen25_omni.py
new file mode 100644
index 00000000..4c55012d
--- /dev/null
+++ b/src/mobius/models/qwen25_omni.py
@@ -0,0 +1,29 @@
+"""Qwen2.5-Omni: Multimodal model with audio + vision + text.
+
+Architecture (Thinker only):
+  - Audio encoder: Conv1d x2 → sinusoidal PE → 32 encoder layers → AvgPool → proj
+  - Vision encoder: Conv3d patch embed → 32 ViT blocks → patch merger
+  - Fusion: Audio/vision features replace placeholder token positions
+  - Text decoder: Qwen2 (no QK norm) + MRoPE
+
+Reference: https://huggingface.co/Qwen/Qwen2.5-Omni-7B
+HuggingFace class: Qwen2_5OmniForConditionalGeneration
+"""
+
+
+class Qwen25OmniAudioEncoder(nn.Module):
+    pass
+
+class Qwen25OmniVisionEncoder(nn.Module):
+    pass
+
+class Qwen25OmniEmbeddingModel(nn.Module):
+    pass
+
+class Qwen25OmniDecoderModel(nn.Module):
+    pass
+
+class Qwen25OmniThinkerForConditionalGeneration(nn.Module):
+    pass
+
+

From 4f676c253dbd5f18c38b36780a910f19a3fbc08a Mon Sep 17 00:00:00 2001
From: ruiren_microsoft <ruiren@microsoft.com>
Date: Wed, 22 Apr 2026 20:44:51 -0700
Subject: [PATCH 2/6] add audio encoder

---
 src/mobius/_registry.py          |   7 ++
 src/mobius/components/_conv.py   |  39 ++++++++
 src/mobius/models/qwen25_omni.py | 152 ++++++++++++++++++++++++++++++-
 3 files changed, 197 insertions(+), 1 deletion(-)

diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py
index 69a93f95..44e3c42a 100644
--- a/src/mobius/_registry.py
+++ b/src/mobius/_registry.py
@@ -115,6 +115,7 @@
 from mobius.models.qwen3_asr import Qwen3ASRForConditionalGeneration
 from mobius.models.qwen3_tts import Qwen3TTSForConditionalGeneration
 from mobius.models.qwen3_tts_tokenizer import Qwen3TTSTokenizerV2Model
+from mobius.models.qwen25_omni import Qwen25OmniThinkerForConditionalGeneration
 from mobius.models.sam2 import Sam2VisionModel
 from mobius.models.segformer import SegformerForSemanticSegmentation
 from mobius.models.starcoder2 import StarCoder2CausalLMModel
@@ -519,6 +520,12 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
         task="speech-to-text",
         config_class=WhisperConfig,
     ),
+
+    # --- Omni ---
+    "qwen2_5_omni": ModelRegistration(
+        Qwen25OmniThinkerForConditionalGeneration,
+        task="speech-language",
+    ),
     # --- Encoder-only ---
     "albert": ModelRegistration(BertModel, task="feature-extraction"),
     "bert": ModelRegistration(BertModel, task="feature-extraction"),
diff --git a/src/mobius/components/_conv.py b/src/mobius/components/_conv.py
index 1bd88b2f..aee0fa8b 100644
--- a/src/mobius/components/_conv.py
+++ b/src/mobius/components/_conv.py
@@ -54,6 +54,45 @@ def forward(self, op: builder.OpBuilder, x: ir.Value):
         )
 
 
+class Conv1d(nn.Module):
+    """1D convolution with bias.
+
+    Matches ``torch.nn.Conv1d`` with ``bias=True``.  The default ``padding=0``
+    follows PyTorch convention; callers should specify padding explicitly.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        padding: int = 0,
+        groups: int = 1,
+    ):
+        super().__init__()
+        self.weight = nn.Parameter(
+            (out_channels, in_channels // groups, kernel_size)
+        )
+        self.bias = nn.Parameter((out_channels))
+        self._kernel_size = kernel_size
+        self._stride = stride
+        self._padding = padding
+        self._groups = groups
+
+    def forward(self, op: builder.OpBuilder, x: ir.Value):
+        p = self._padding
+        return op.Conv(
+            x,
+            self.weight,
+            self.bias,
+            kernel_shape=[self._kernel_size],
+            strides=[self._stride],
+            pads=[p, p],
+            group=self._groups,
+        )
+
+
 class Conv2dNoBias(nn.Module):
     """2D convolution without bias."""
 
diff --git a/src/mobius/models/qwen25_omni.py b/src/mobius/models/qwen25_omni.py
index 4c55012d..61c3cbe4 100644
--- a/src/mobius/models/qwen25_omni.py
+++ b/src/mobius/models/qwen25_omni.py
@@ -10,9 +10,159 @@
 HuggingFace class: Qwen2_5OmniForConditionalGeneration
 """
 
+from __future__ import annotations
+
+import dataclasses
+
+import numpy as np
+import onnx_ir as ir
+import torch
+from onnxscript import nn
+from onnxscript._internal import builder
+
+
+from mobius._configs import ArchitectureConfig
+from mobius.components._common import (
+    Embedding,
+    LayerNorm,
+    Linear,
+    create_attention_bias,
+)
+from mobius.components._conv import Conv1d
+from mobius.components._qwen25_omni_audio import Qwen25OmniAudioEncoderLayer
+
+
+
+def _sinusoidal_position_embedding(max_positions: int, d_model: int) -> np.ndarray:
+    """Compute sinusoidal positional embeddings matching Qwen3-ASR.
+
+    Uses log-timescale increments (different from Whisper which uses
+    alternating sin/cos layout). Layout: [sin_0..sin_n, cos_0..cos_n].
+    """
+    channels = d_model
+    log_timescale_increment = np.log(10000.0) / (channels // 2 - 1)
+    inv_timescales = np.exp(
+        -log_timescale_increment * np.arange(channels // 2, dtype=np.float32)
+    )
+    scaled_time = (
+        np.arange(max_positions, dtype=np.float32)[:, np.newaxis]
+        * inv_timescales[np.newaxis, :]
+    )
+    # Layout: [sin, cos] matching HF SinusoidsPositionEmbedding
+    pe = np.concatenate([np.sin(scaled_time), np.cos(scaled_time)], axis=1).astype(np.float32)
+    return pe
+
 
 class Qwen25OmniAudioEncoder(nn.Module):
-    pass
+    """Qwen25-Omni audio encoder
+    Converts mel spectrogram to audio feature embeddings:
+      mel (batch, num_mel_bins, seq_len)
+      -> 2x Conv1d with GELU
+      -> sinusoidal position embeddings
+      -> N bidirectional encoder layers
+      -> AvgPool1d (2x downsample)
+      -> LayerNorm (ln_post)
+      -> Linear proj (d_model -> output_dim)
+
+    Output: (batch, out_seq_len, output_dim)
+    """
+
+    def __init__(self, config: ArchitectureConfig):
+        super().__init__()
+        audio = config.audio
+        assert audio is not None
+
+        d_model = audio.d_model or 1280
+        num_mel_bin = audio.num_mel_bins or 128
+        encoder_layers = audio.encoder_layers or 32
+        encoder_heads = audio.encoder_attention_heads or 20
+        encoder_ffn = audio.encoder_ffn_dim or 3584
+        max_source_positions = audio.max_source_positions or 1500
+        n_window = audio.n_window or 100
+        output_dim = audio.output_dim or 3584
+
+        # 2x Conv1d: mel -> d_model with GELU between them
+        self.conv1 = Conv1d(
+            num_mel_bin,
+            d_model,
+            kernel_size=3,
+            padding=1,
+        )
+        self.conv2 = Conv1d(
+            d_model,
+            d_model,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+        )
+
+        # Sinusoidal positional embeddings (frozen)
+        pe_data = _sinusoidal_position_embedding(max_source_positions, d_model)
+        self.positional_embedding = nn.Parameter(
+            [max_source_positions, d_model],
+            name="positional_embedding.positional_embedding",
+            data=ir.tensor(pe_data),
+        )
+
+
+        # Encoder transformer layers
+        self.layers = nn.ModuleList(
+            [
+                Qwen25OmniAudioEncoderLayer(d_model, encoder_heads, encoder_ffn)
+                for _ in range(encoder_layers)
+            ]
+        )
+
+        # Post-encoder normalization
+        self.ln_post = LayerNorm(d_model)
+
+        # Output projection: d_model -> output_dim
+        self.proj = Linear(d_model, output_dim)
+
+    def forward(self, op: builder.OpBuilder, input_features: ir.Value):
+        """Encode mel spectrogram to audio features.
+
+        Args:
+            input_features: (batch, num_mel_bins, seq_len) mel spectrogram
+
+        Returns:
+            audio_features: (batch, out_seq_len, output_dim)
+        """
+
+        # 2X Conv1d with GELU: (batch, mel, seq) -> (batch, d_model, seq//2)
+        hidden_states = op.Gelu(self.conv1(op, input_features))
+        hidden_states = op.Gelu(self.conv2(op, hidden_states))
+
+        # Transpose to (batch, seq//2, d_model) for transformer layers
+        hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1])
+
+        # Add sinusoidal positional embeddings
+        seq_len = op.Shape(hidden_states, start=1, end=2)
+        pe_slice = op.Slice(
+            self.positional_embedding,
+            op.Constant(value_ints=[0]),
+            seq_len,
+            op.Constant(value_ints=[0])
+        )
+
+        hidden_states = op.Add(hidden_states, pe_slice)
+
+        # Encoder layer
+        for layer in self.layers:
+            hidden_states = layer(op, hidden_states)
+
+        # AvgPool1d(kernel=2, stride=2): halves sequence length
+        # Transpose to (batch, d_model, seq) for pooling, then back
+        hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1])
+        hidden_states = op.AveragePool(hidden_states, kernel_shape=[2], strides=[2])
+        hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1])
+
+        # ln_post once, then proj
+        hidden_states = self.ln_post(op, hidden_states)
+        hidden_states = self.proj(op, hidden_states)
+
+        return hidden_states
+
 
 class Qwen25OmniVisionEncoder(nn.Module):
     pass

From 4b0254cbe971791072489ad0852165f594d7f1bc Mon Sep 17 00:00:00 2001
From: ruiren_microsoft <ruiren@microsoft.com>
Date: Mon, 27 Apr 2026 16:56:31 -0700
Subject: [PATCH 3/6] add vision encoder

---
 src/mobius/models/qwen25_omni.py | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/mobius/models/qwen25_omni.py b/src/mobius/models/qwen25_omni.py
index 61c3cbe4..9d7d26d5 100644
--- a/src/mobius/models/qwen25_omni.py
+++ b/src/mobius/models/qwen25_omni.py
@@ -30,6 +30,7 @@
 )
 from mobius.components._conv import Conv1d
 from mobius.components._qwen25_omni_audio import Qwen25OmniAudioEncoderLayer
+from mobius.components import Qwen25VLVisionModel
 
 
 
@@ -165,7 +166,31 @@ def forward(self, op: builder.OpBuilder, input_features: ir.Value):
 
 
 class Qwen25OmniVisionEncoder(nn.Module):
-    pass
+    """
+    Qwen2.5-Omni vision encoder - reuses Qwen2.5-VL ViT
+    """
+
+    def __init__(self, config: ArchitectureConfig):
+        super().__init__()
+        vc = config.vision
+        assert vc is not None
+
+        self.visual = Qwen25VLVisionModel(
+            depth=vc.num_hidden_layers or 32,
+            hidden_size = vc.hidden_size or 1280,
+            intermediate_size = vc.intermediate_size or 3420,
+            num_heads = vc.num_attention_heads or 16,
+            patch_size = vc.patch_size or 14,
+            temporal_patch_size = vc.temporal_patch_size or 2,
+            in_channels = vc.in_channels or 3,
+            out_hidden_size = vc.out_hidden_size or 3584,
+            spatial_merge_size = vc.spatial_merge_size or 2,
+            fullatt_block_indexes = vc.fullatt_block_indexes or (7, 15, 23, 31),
+            window_size = vc.window_size or 112
+        )
+
+    def forward(self, op, pixel_values, image_grid_thw):
+        return self.visual(op, pixel_values, image_grid_thw)
 
 class Qwen25OmniEmbeddingModel(nn.Module):
     pass

From 8c3124fbb4171b426c03762637bd99125e3eac41 Mon Sep 17 00:00:00 2001
From: ruiren_microsoft <ruiren@microsoft.com>
Date: Mon, 27 Apr 2026 18:13:50 -0700
Subject: [PATCH 4/6] add embeding and decoder

---
 src/mobius/models/qwen25_omni.py | 107 ++++++++++++++++++++++++++++++-
 1 file changed, 105 insertions(+), 2 deletions(-)

diff --git a/src/mobius/models/qwen25_omni.py b/src/mobius/models/qwen25_omni.py
index 9d7d26d5..e81ec47d 100644
--- a/src/mobius/models/qwen25_omni.py
+++ b/src/mobius/models/qwen25_omni.py
@@ -193,10 +193,113 @@ def forward(self, op, pixel_values, image_grid_thw):
         return self.visual(op, pixel_values, image_grid_thw)
 
 class Qwen25OmniEmbeddingModel(nn.Module):
-    pass
+    """Fuses text embedding with audio and image feature
+    
+    Replaces audio_token_id positions with audio features,
+    and image_token_id positions with image features.
+
+    Inputs:
+        input_ids: (batch, seq_len)
+        audio_features: (num_audio_tokens, hidden_size)
+        image_features: (num_image_tokens, hidden_size)
+
+    Output:
+        inputs_embeds: (batch, seq_len, hidden_size)
+    """
+
+    def __init__(self, config: ArchitectureConfig):
+        super().__init__()
+        self.embed_tokens = Embedding(
+            config.vocab_size, config.hidden_size, config.pad_token_id
+        )
+
+        # Token IDs from Qwen2.5-Omni config
+        audio = config.audio
+        self._audio_token_id = audio.audio_token_id if audio else 151646
+        self._image_token_id = audio.image_token_id or 151655
+
+    
+    def forward(self, op, inputs_ids, audio_features, image_features):
+        inputs_embeds = self.embed_tokens(op, inputs_ids)
+
+        # Fuse audio features at audio token positions
+        inputs_embeds = self._replace_tokens(
+            op, inputs_embeds, inputs_ids, audio_features, self._audio_token_id
+        )
+
+        # Fuse image features at image token position
+        inputs_embeds = self._replace_tokens(
+            op, inputs_embeds, inputs_ids, image_features, self._image_token_ids
+        )
+
+        return input_embeds
+
+    def _replace_tokens(self, op, inputs_embeds, input_ids, features, token_id):
+        """Replace token positions with encoder features (masked_scatter equivalent)."""
+        mask = op.Equal(input_ids, op.Constant(value_int=token_id))
+        mask_3d = op.Unsqueeze(mask, [-1])
+
+        # Pad with zero row for safety (text-only case: no features)
+        feature_dim = op.Shape(features, start=1, end=2)
+        zero_shape = op.Concat(op.Constant(value_ints=[1]), feature_dim, axis=0)
+        zero_row = op.ConstantOfShape(zero_shape, value=ir.tensor(np.zeros(1, dtype=np.float32)))
+        padded = op.Concat(zero_row, features, axis=0)
+
+        # CumSum to get gather indices
+        mask_int = op.Cast(mask, to=7)
+        flat = op.Reshape(mask_int, op.Constant(value_ints=[-1]))
+        indices = op.CumSum(flat, op.Constant(value_int=0))
+        indices = op.Mul(indices, flat)
+        indices = op.Reshape(indices, op.Shape(input_ids))
+
+        gathered = op.Gather(padded, indices, axis=0)
+        return op.Where(mask_3d, gathered, inputs_embeds)
+
 
 class Qwen25OmniDecoderModel(nn.Module):
-    pass
+    """Qwen2.5-Omni text decoder: inputs_embeds → logits + KV cache.
+
+    Standard Qwen2 decoder with MRoPE (3D position_ids).
+    No QK norm (unlike Qwen3-ASR which uses attn_qk_norm=True).
+    """
+
+    def __init__(self, config: ArchitectureConfig):
+        super().__init__()
+        self._dtype = config.dtype
+        self.layers = nn.ModuleList(
+            [DecoderLayer(config) for _ in range(config.num_hidden_layers)]
+        )
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = initialize_rope(config)
+        self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def forward(self, op, inputs_embeds, attention_mask, position_ids, past_key_values=None):
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(op, position_ids)
+
+        attention_bias = create_attention_bias(
+            op,
+            input_ids=inputs_embeds,
+            attention_mask=attention_mask,
+            dtype=self._dtype,
+        )
+
+        present_key_values = []
+        past_kvs = past_key_values or [None] * len(self.layers)
+        for layer, past_kv in zip(self.layers, past_kvs):
+            hidden_states, present_kv = layer(
+                op,
+                hidden_states=hidden_states,
+                attention_bias=attention_bias,
+                position_embeddings=position_embeddings,
+                past_key_value=past_kv,
+            )
+            present_key_values.append(present_kv)
+
+        hidden_states = self.norm(op, hidden_states)
+        logits = self.lm_head(op, hidden_states)
+        return logits, present_key_values
+        )
 
 class Qwen25OmniThinkerForConditionalGeneration(nn.Module):
     pass

From 67f29888c9c4a05449b23921ab2eea7b108e1d4e Mon Sep 17 00:00:00 2001
From: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Date: Tue, 2 Jun 2026 20:26:47 +0000
Subject: [PATCH 5/6] Qwen2.5-Omni: fix syntax + complete composite Thinker
 model
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Pick up the partial Qwen2.5-Omni onboarding (#193) and bring it to a
state where the file imports cleanly, the full pytest suite collects
without errors, and the registry-completeness test passes.

Concrete fixes:

* Syntax: removed stray ")" at end of Qwen25OmniDecoderModel that
  broke every test that imported the module (65 collection errors on
  plain main).
* Missing imports: RMSNorm, DecoderLayer, initialize_rope were
  referenced by Qwen25OmniDecoderModel but not imported. Also
  consolidated the OpBuilder import (was using the internal
  onnxscript._internal.builder module).
* Typos in Qwen25OmniEmbeddingModel:
  - Used audio.image_token_id (image_token_id lives on VisionConfig).
  - self._image_token_ids plural typo where _image_token_id was meant.
  - "return input_embeds" referenced a non-existent local; renamed.
  - Forward arg "inputs_ids" -> "input_ids".
* Composite class implemented: Qwen25OmniThinkerForConditionalGeneration
  was just `pass`. Now mirrors the gemma4 multi-encoder layout — owns
  decoder, embedding, vision_encoder, and audio_tower (named to match
  what SpeechLanguageTask looks up). forward() raises
  NotImplementedError because this is a multi-output split needing a
  dedicated task. preprocess_weights maps HF
  thinker.audio_tower.* / thinker.visual.* / thinker.model.* /
  thinker.lm_head.* onto the new module structure, drops Talker /
  codec-head keys, and applies the standard embed/lm_head weight
  tying.

Test infrastructure:

* Added qwen2_5_omni to _KNOWN_UNTESTED_MODEL_TYPES with a comment
  explaining the situation: the audio + vision + text 4-model split
  needs a dedicated task (like Gemma4Task); the existing
  speech-language task only drives audio + text. When that task is
  added in a follow-up, this entry should move into a dedicated graph
  test (mirroring test_gemma4_any_to_any_graph).

Also cleaned up: unused n_window lookup, unused dataclasses import,
trailing blank lines, and audio-attention docstring punctuation that
ruff flagged.

Tests: full tests/build_graph_test.py + src/mobius/ suite passes
(2767 passed, 43 skipped). Ruff check + format clean on touched files.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
---
 src/mobius/components/_qwen25_omni_audio.py |  17 +-
 src/mobius/models/qwen25_omni.py            | 241 ++++++++++++++------
 tests/build_graph_test.py                   |   7 +
 3 files changed, 184 insertions(+), 81 deletions(-)

diff --git a/src/mobius/components/_qwen25_omni_audio.py b/src/mobius/components/_qwen25_omni_audio.py
index 38556e3a..f2597e13 100644
--- a/src/mobius/components/_qwen25_omni_audio.py
+++ b/src/mobius/components/_qwen25_omni_audio.py
@@ -17,16 +17,17 @@
 
 from mobius.components._common import LayerNorm, Linear
 
-
 if TYPE_CHECKING:
     import onnx_ir as ir
 
+
 class Qwen25OmniAudioAttention(nn.Module):
     """Bidirectional multi-head attention for Qwen2_5Omni audio encoder.
 
     Unlike WhisperAttention, all projections (Q, V, Out) have bias and K does not have bias.
     No causal masking — the encoder uses full bidirectional attention.
     """
+
     def __init__(self, d_model: int, num_heads: int):
         super().__init__()
         self.q_proj = Linear(d_model, d_model, bias=True)
@@ -37,7 +38,7 @@ def __init__(self, d_model: int, num_heads: int):
         self._head_dim = d_model // num_heads
 
     def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
-        """Bidirectional self-attention
+        """Bidirectional self-attention.
 
         Args:
             hidden_states: (batch, seq_len, d_model)
@@ -45,7 +46,6 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
         Returns:
             output: (batch, seq_len, d_model)
         """
-
         q = self.q_proj(op, hidden_states)
         k = self.k_proj(op, hidden_states)
         v = self.v_proj(op, hidden_states)
@@ -55,8 +55,8 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
             q,
             k,
             v,
-            q_num_heads = self._num_heads,
-            kv_num_heads = self._num_heads,
+            q_num_heads=self._num_heads,
+            kv_num_heads=self._num_heads,
             scale=float(self._head_dim**-0.5),
         )
         return self.out_proj(op, attn_output)
@@ -68,7 +68,7 @@ class Qwen25OmniAudioEncoderLayer(nn.Module):
     Pre-norm pattern:  LayerNorm → self-attn → residual
     → LayerNorm → FFN → residual.
     Uses GELU activation in the FFN.
-    
+
     Huggingface class: ``Qwen2_5OmniAudioEncoder``
     """
 
@@ -92,7 +92,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
         Args:
             hidden_states: (batch, seq_len, d_model)
 
-        returns:
+        Returns:
             hidden_states: (batch, seq_len, d_model)
         """
         # Self-attention with pre-norm and residual
@@ -110,6 +110,3 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
         hidden_states = op.Add(residual, hidden_states)
 
         return hidden_states
-
-
-
diff --git a/src/mobius/models/qwen25_omni.py b/src/mobius/models/qwen25_omni.py
index e81ec47d..fcdbdbe2 100644
--- a/src/mobius/models/qwen25_omni.py
+++ b/src/mobius/models/qwen25_omni.py
@@ -12,16 +12,13 @@
 
 from __future__ import annotations
 
-import dataclasses
-
 import numpy as np
 import onnx_ir as ir
 import torch
-from onnxscript import nn
-from onnxscript._internal import builder
-
+from onnxscript import OpBuilder, nn
 
 from mobius._configs import ArchitectureConfig
+from mobius.components import Qwen25VLVisionModel
 from mobius.components._common import (
     Embedding,
     LayerNorm,
@@ -29,9 +26,10 @@
     create_attention_bias,
 )
 from mobius.components._conv import Conv1d
+from mobius.components._decoder import DecoderLayer
 from mobius.components._qwen25_omni_audio import Qwen25OmniAudioEncoderLayer
-from mobius.components import Qwen25VLVisionModel
-
+from mobius.components._rms_norm import RMSNorm
+from mobius.components._rotary_embedding import initialize_rope
 
 
 def _sinusoidal_position_embedding(max_positions: int, d_model: int) -> np.ndarray:
@@ -55,7 +53,8 @@ def _sinusoidal_position_embedding(max_positions: int, d_model: int) -> np.ndarr
 
 
 class Qwen25OmniAudioEncoder(nn.Module):
-    """Qwen25-Omni audio encoder
+    """Qwen25-Omni audio encoder.
+
     Converts mel spectrogram to audio feature embeddings:
       mel (batch, num_mel_bins, seq_len)
       -> 2x Conv1d with GELU
@@ -79,23 +78,11 @@ def __init__(self, config: ArchitectureConfig):
         encoder_heads = audio.encoder_attention_heads or 20
         encoder_ffn = audio.encoder_ffn_dim or 3584
         max_source_positions = audio.max_source_positions or 1500
-        n_window = audio.n_window or 100
         output_dim = audio.output_dim or 3584
 
         # 2x Conv1d: mel -> d_model with GELU between them
-        self.conv1 = Conv1d(
-            num_mel_bin,
-            d_model,
-            kernel_size=3,
-            padding=1,
-        )
-        self.conv2 = Conv1d(
-            d_model,
-            d_model,
-            kernel_size=3,
-            stride=2,
-            padding=1,
-        )
+        self.conv1 = Conv1d(num_mel_bin, d_model, kernel_size=3, padding=1)
+        self.conv2 = Conv1d(d_model, d_model, kernel_size=3, stride=2, padding=1)
 
         # Sinusoidal positional embeddings (frozen)
         pe_data = _sinusoidal_position_embedding(max_source_positions, d_model)
@@ -105,7 +92,6 @@ def __init__(self, config: ArchitectureConfig):
             data=ir.tensor(pe_data),
         )
 
-
         # Encoder transformer layers
         self.layers = nn.ModuleList(
             [
@@ -120,7 +106,7 @@ def __init__(self, config: ArchitectureConfig):
         # Output projection: d_model -> output_dim
         self.proj = Linear(d_model, output_dim)
 
-    def forward(self, op: builder.OpBuilder, input_features: ir.Value):
+    def forward(self, op: OpBuilder, input_features: ir.Value):
         """Encode mel spectrogram to audio features.
 
         Args:
@@ -129,7 +115,6 @@ def forward(self, op: builder.OpBuilder, input_features: ir.Value):
         Returns:
             audio_features: (batch, out_seq_len, output_dim)
         """
-
         # 2X Conv1d with GELU: (batch, mel, seq) -> (batch, d_model, seq//2)
         hidden_states = op.Gelu(self.conv1(op, input_features))
         hidden_states = op.Gelu(self.conv2(op, hidden_states))
@@ -143,32 +128,26 @@ def forward(self, op: builder.OpBuilder, input_features: ir.Value):
             self.positional_embedding,
             op.Constant(value_ints=[0]),
             seq_len,
-            op.Constant(value_ints=[0])
+            op.Constant(value_ints=[0]),
         )
-
         hidden_states = op.Add(hidden_states, pe_slice)
 
-        # Encoder layer
         for layer in self.layers:
             hidden_states = layer(op, hidden_states)
 
-        # AvgPool1d(kernel=2, stride=2): halves sequence length
-        # Transpose to (batch, d_model, seq) for pooling, then back
+        # AvgPool1d(kernel=2, stride=2): halves sequence length.
+        # Transpose to (batch, d_model, seq) for pooling, then back.
         hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1])
         hidden_states = op.AveragePool(hidden_states, kernel_shape=[2], strides=[2])
         hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1])
 
-        # ln_post once, then proj
         hidden_states = self.ln_post(op, hidden_states)
         hidden_states = self.proj(op, hidden_states)
-
         return hidden_states
 
 
 class Qwen25OmniVisionEncoder(nn.Module):
-    """
-    Qwen2.5-Omni vision encoder - reuses Qwen2.5-VL ViT
-    """
+    """Qwen2.5-Omni vision encoder — reuses the Qwen2.5-VL ViT."""
 
     def __init__(self, config: ArchitectureConfig):
         super().__init__()
@@ -177,26 +156,27 @@ def __init__(self, config: ArchitectureConfig):
 
         self.visual = Qwen25VLVisionModel(
             depth=vc.num_hidden_layers or 32,
-            hidden_size = vc.hidden_size or 1280,
-            intermediate_size = vc.intermediate_size or 3420,
-            num_heads = vc.num_attention_heads or 16,
-            patch_size = vc.patch_size or 14,
-            temporal_patch_size = vc.temporal_patch_size or 2,
-            in_channels = vc.in_channels or 3,
-            out_hidden_size = vc.out_hidden_size or 3584,
-            spatial_merge_size = vc.spatial_merge_size or 2,
-            fullatt_block_indexes = vc.fullatt_block_indexes or (7, 15, 23, 31),
-            window_size = vc.window_size or 112
+            hidden_size=vc.hidden_size or 1280,
+            intermediate_size=vc.intermediate_size or 3420,
+            num_heads=vc.num_attention_heads or 16,
+            patch_size=vc.patch_size or 14,
+            temporal_patch_size=vc.temporal_patch_size or 2,
+            in_channels=vc.in_channels or 3,
+            out_hidden_size=vc.out_hidden_size or 3584,
+            spatial_merge_size=vc.spatial_merge_size or 2,
+            fullatt_block_indexes=vc.fullatt_block_indexes or (7, 15, 23, 31),
+            window_size=vc.window_size or 112,
         )
 
-    def forward(self, op, pixel_values, image_grid_thw):
+    def forward(self, op: OpBuilder, pixel_values: ir.Value, image_grid_thw: ir.Value):
         return self.visual(op, pixel_values, image_grid_thw)
 
+
 class Qwen25OmniEmbeddingModel(nn.Module):
-    """Fuses text embedding with audio and image feature
-    
-    Replaces audio_token_id positions with audio features,
-    and image_token_id positions with image features.
+    """Fuses text embedding with audio and image features.
+
+    Replaces ``audio_token_id`` positions with audio features and
+    ``image_token_id`` positions with image features.
 
     Inputs:
         input_ids: (batch, seq_len)
@@ -210,42 +190,53 @@ class Qwen25OmniEmbeddingModel(nn.Module):
     def __init__(self, config: ArchitectureConfig):
         super().__init__()
         self.embed_tokens = Embedding(
-            config.vocab_size, config.hidden_size, config.pad_token_id
+            config.vocab_size,
+            config.hidden_size,
+            config.pad_token_id,
         )
 
-        # Token IDs from Qwen2.5-Omni config
+        # Token IDs default to the values used by Qwen2.5-Omni-7B.
         audio = config.audio
-        self._audio_token_id = audio.audio_token_id if audio else 151646
-        self._image_token_id = audio.image_token_id or 151655
-
-    
-    def forward(self, op, inputs_ids, audio_features, image_features):
-        inputs_embeds = self.embed_tokens(op, inputs_ids)
-
-        # Fuse audio features at audio token positions
+        vision = config.vision
+        self._audio_token_id = (audio.audio_token_id if audio else None) or 151646
+        self._image_token_id = (vision.image_token_id if vision else None) or 151655
+
+    def forward(
+        self,
+        op: OpBuilder,
+        input_ids: ir.Value,
+        audio_features: ir.Value,
+        image_features: ir.Value,
+    ):
+        inputs_embeds = self.embed_tokens(op, input_ids)
+
+        # Fuse audio features at audio token positions.
         inputs_embeds = self._replace_tokens(
-            op, inputs_embeds, inputs_ids, audio_features, self._audio_token_id
+            op, inputs_embeds, input_ids, audio_features, self._audio_token_id
         )
 
-        # Fuse image features at image token position
+        # Fuse image features at image token positions.
         inputs_embeds = self._replace_tokens(
-            op, inputs_embeds, inputs_ids, image_features, self._image_token_ids
+            op, inputs_embeds, input_ids, image_features, self._image_token_id
         )
 
-        return input_embeds
+        return inputs_embeds
 
     def _replace_tokens(self, op, inputs_embeds, input_ids, features, token_id):
         """Replace token positions with encoder features (masked_scatter equivalent)."""
         mask = op.Equal(input_ids, op.Constant(value_int=token_id))
         mask_3d = op.Unsqueeze(mask, [-1])
 
-        # Pad with zero row for safety (text-only case: no features)
+        # Pad with a zero row for safety (text-only case: no features).
         feature_dim = op.Shape(features, start=1, end=2)
         zero_shape = op.Concat(op.Constant(value_ints=[1]), feature_dim, axis=0)
-        zero_row = op.ConstantOfShape(zero_shape, value=ir.tensor(np.zeros(1, dtype=np.float32)))
+        zero_row = op.ConstantOfShape(
+            zero_shape, value=ir.tensor(np.zeros(1, dtype=np.float32))
+        )
         padded = op.Concat(zero_row, features, axis=0)
 
-        # CumSum to get gather indices
+        # CumSum-based per-position gather index. Mask positions get the
+        # next feature row in order; non-mask positions get the zero row.
         mask_int = op.Cast(mask, to=7)
         flat = op.Reshape(mask_int, op.Constant(value_ints=[-1]))
         indices = op.CumSum(flat, op.Constant(value_int=0))
@@ -273,9 +264,18 @@ def __init__(self, config: ArchitectureConfig):
         self.rotary_emb = initialize_rope(config)
         self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False)
 
-    def forward(self, op, inputs_embeds, attention_mask, position_ids, past_key_values=None):
+    def forward(
+        self,
+        op: OpBuilder,
+        inputs_embeds: ir.Value,
+        attention_mask: ir.Value,
+        position_ids: ir.Value,
+        past_key_values=None,
+    ):
         hidden_states = inputs_embeds
-        position_embeddings = self.rotary_emb(op, position_ids)
+        position_embeddings = (
+            self.rotary_emb(op, position_ids) if self.rotary_emb is not None else None
+        )
 
         attention_bias = create_attention_bias(
             op,
@@ -299,9 +299,108 @@ def forward(self, op, inputs_embeds, attention_mask, position_ids, past_key_valu
         hidden_states = self.norm(op, hidden_states)
         logits = self.lm_head(op, hidden_states)
         return logits, present_key_values
-        )
+
 
 class Qwen25OmniThinkerForConditionalGeneration(nn.Module):
-    pass
+    """Qwen2.5-Omni Thinker: composite audio + vision + text model.
+
+    Builds four separate ONNX models:
 
+    - ``decoder``: Qwen2.5 text decoder taking ``inputs_embeds``
+    - ``vision_encoder``: Qwen2.5-VL ViT (pixel_values + grid_thw → image features)
+    - ``audio_tower``: 2x Conv1d + transformer audio tower (mel → audio features)
+    - ``embedding``: word embedding + multimodal feature fusion
 
+    HuggingFace class: ``Qwen2_5OmniForConditionalGeneration`` (Thinker only —
+    the Talker / streaming code generation head is out of scope for now).
+    """
+
+    default_task: str = "speech-language"
+    category: str = "Multimodal"
+    config_class: type = ArchitectureConfig
+
+    def __init__(self, config: ArchitectureConfig):
+        super().__init__()
+        self.config = config
+        self.decoder = Qwen25OmniDecoderModel(config)
+        self.embedding = Qwen25OmniEmbeddingModel(config)
+        self.vision_encoder: Qwen25OmniVisionEncoder | None = (
+            Qwen25OmniVisionEncoder(config) if config.vision is not None else None
+        )
+        # Named ``audio_tower`` (not ``audio_encoder``) to match
+        # ``SpeechLanguageTask`` which looks up ``module.audio_tower``.
+        # The vision sub-module is currently not routed through any
+        # task — wiring a dedicated Qwen25OmniTask that builds all four
+        # ONNX sub-models is a follow-up; today the
+        # ``speech-language`` task entry only drives audio + text.
+        self.audio_tower: Qwen25OmniAudioEncoder | None = (
+            Qwen25OmniAudioEncoder(config) if config.audio is not None else None
+        )
+
+    def forward(self, op: OpBuilder, **kwargs):
+        raise NotImplementedError(
+            "Qwen25OmniThinkerForConditionalGeneration is a multi-model split; the corresponding "
+            "task class builds each sub-module (decoder, embedding, vision_encoder, audio_encoder) "
+            "separately."
+        )
+
+    def preprocess_weights(
+        self, state_dict: dict[str, torch.Tensor]
+    ) -> dict[str, torch.Tensor]:
+        """Map HuggingFace weight names to ONNX module structure.
+
+        HF Qwen2.5-Omni checkpoints prefix every Thinker key with ``thinker.``:
+
+        - ``thinker.audio_tower.*`` → ``audio_tower.*``
+        - ``thinker.visual.*`` → ``vision_encoder.visual.*``
+        - ``thinker.model.embed_tokens.*`` → ``embedding.embed_tokens.*``
+        - ``thinker.model.layers.N.*`` and ``model.norm.*`` → ``decoder.*``
+        - ``thinker.lm_head.*`` → ``decoder.lm_head.*``
+        - ``thinker.model.rotary_emb.*`` → ``decoder.rotary_emb.*``
+
+        The Talker sub-tree (``talker.*``) and the audio-output codec head
+        are not consumed by this model and are silently dropped.
+        """
+        cleaned: dict[str, torch.Tensor] = {}
+        for key, value in state_dict.items():
+            # Strip the thinker. prefix if present.
+            if key.startswith("thinker."):
+                key = key[len("thinker.") :]
+
+            # Drop talker.* and any codec output keys — not part of Thinker.
+            if key.startswith(("talker.", "code_predictor.")):
+                continue
+
+            if key.startswith("audio_tower."):
+                cleaned[key] = value
+                continue
+
+            if key.startswith("visual."):
+                cleaned["vision_encoder." + key] = value
+                continue
+
+            if key.startswith("lm_head."):
+                cleaned["decoder." + key] = value
+                continue
+
+            if key.startswith("model."):
+                inner = key[len("model.") :]
+                if inner.startswith("embed_tokens."):
+                    cleaned["embedding." + inner] = value
+                    continue
+                if inner.startswith(("layers.", "norm.", "rotary_emb.")):
+                    cleaned["decoder." + inner] = value
+                    continue
+
+            cleaned[key] = value
+
+        # Weight tying: ``embedding.embed_tokens.weight`` ↔ ``decoder.lm_head.weight``.
+        embed_key = "embedding.embed_tokens.weight"
+        lm_key = "decoder.lm_head.weight"
+        if getattr(self.config, "tie_word_embeddings", False):
+            if embed_key in cleaned and lm_key not in cleaned:
+                cleaned[lm_key] = cleaned[embed_key]
+            elif lm_key in cleaned and embed_key not in cleaned:
+                cleaned[embed_key] = cleaned[lm_key]
+
+        return cleaned
diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py
index 7c15be06..955537c5 100644
--- a/tests/build_graph_test.py
+++ b/tests/build_graph_test.py
@@ -4288,6 +4288,13 @@ def test_jamba_preprocess_weights_moe_renames(self):
 _KNOWN_UNTESTED_MODEL_TYPES: set[str] = {
     "deepseek_v2_moe",  # Alias for deepseek_v2 — tested via deepseek_v2
     "qwen3_5_vl_text",  # VL text decoder — tested via parent VL model
+    # Qwen2.5-Omni Thinker is a 4-model split (audio + vision + embedding +
+    # decoder). The model graph is implemented but the existing
+    # SpeechLanguageTask only knows audio + text, so wiring a dedicated
+    # task that also drives the vision encoder is a follow-up — tracked
+    # in the Qwen2.5-Omni onboarding PR (#193). When that lands, move
+    # the entry into a dedicated test (mirroring gemma4_any_to_any).
+    "qwen2_5_omni",
 }
 
 

From 5ca65be5665ffc5e25998eea88aceccaf2cc3a12 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 13 Jul 2026 02:47:07 +0000
Subject: [PATCH 6/6] Complete Qwen2.5-Omni Thinker export

Add a dedicated four-model task and align audio, vision, decoder, configuration, and weight routing with the current Hugging Face Thinker implementation. Cover the nested config and complete graph package with focused regression tests.

Signed-off-by: GitHub <noreply@github.com>
---
 src/mobius/_configs/_base.py                  |   2 +
 src/mobius/_configs/_extractors.py            |   3 +-
 src/mobius/_configs/_sub_configs.py           |   1 +
 src/mobius/_configs/per_model/__init__.py     |   1 +
 .../_configs/per_model/_qwen25_omni_vision.py |  39 ++++
 src/mobius/_registry.py                       |   3 +-
 src/mobius/components/__init__.py             |   6 +
 src/mobius/components/_conv.py                |  13 +-
 src/mobius/components/_qwen25_omni_audio.py   |  80 +++++--
 src/mobius/models/__init__.py                 |   2 +
 src/mobius/models/qwen25_omni.py              | 200 ++++++++++++++----
 src/mobius/models/qwen25_omni_test.py         | 106 ++++++++++
 src/mobius/tasks/__init__.py                  |   3 +
 src/mobius/tasks/_qwen25_omni.py              | 123 +++++++++++
 tests/build_graph_test.py                     |  74 ++++++-
 15 files changed, 581 insertions(+), 75 deletions(-)
 create mode 100644 src/mobius/_configs/per_model/_qwen25_omni_vision.py
 create mode 100644 src/mobius/models/qwen25_omni_test.py
 create mode 100644 src/mobius/tasks/_qwen25_omni.py

diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py
index f315cbb9..ba415eb6 100644
--- a/src/mobius/_configs/_base.py
+++ b/src/mobius/_configs/_base.py
@@ -418,6 +418,7 @@ class ArchitectureConfig(BaseModelConfig):
     # Vision shared fields (accessed as top-level config.X by tasks)
     mm_tokens_per_image: int | None = None
     image_token_id: int | None = None
+    video_token_id: int | None = None
     spatial_merge_size: int = 2
     temporal_patch_size: int = 2
     deepstack_visual_indexes: list[int] | None = None
@@ -649,6 +650,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig:
                                         "bloom",
                                         "qwen2",
                                         "qwen2_5_vl_text",
+                                        "qwen2_5_omni_text",
                                         "qwen2_moe",
                                         "qwen2_vl_text",
                                     ),
diff --git a/src/mobius/_configs/_extractors.py b/src/mobius/_configs/_extractors.py
index ac4c6e41..6c36a51a 100644
--- a/src/mobius/_configs/_extractors.py
+++ b/src/mobius/_configs/_extractors.py
@@ -156,7 +156,7 @@ def extract_vision_config(config, parent_config, model_type: str) -> dict:
     Either step can populate ``fields`` (which become kwargs for
     :class:`VisionConfig`), or a per-model hook can return a fully-formed
     dict to short-circuit. The dispatcher also lifts a fixed set of
-    "shared" vision fields (``image_token_id``, ``spatial_merge_size``,
+    "shared" vision fields (``image_token_id``, ``video_token_id``, ``spatial_merge_size``,
     ...) up to the top-level of the returned dict so callers can access
     them as ``config.image_token_id`` directly.
     """
@@ -182,6 +182,7 @@ def extract_vision_config(config, parent_config, model_type: str) -> dict:
     for shared in (
         "mm_tokens_per_image",
         "image_token_id",
+        "video_token_id",
         "spatial_merge_size",
         "temporal_patch_size",
         "deepstack_visual_indexes",
diff --git a/src/mobius/_configs/_sub_configs.py b/src/mobius/_configs/_sub_configs.py
index f872d9a7..d4cfd6de 100644
--- a/src/mobius/_configs/_sub_configs.py
+++ b/src/mobius/_configs/_sub_configs.py
@@ -47,6 +47,7 @@ class VisionConfig:
     norm_eps: float = 1e-6
     mm_tokens_per_image: int | None = None
     image_token_id: int | None = None
+    video_token_id: int | None = None
     # Pixtral / Mistral-3 vision fields
     model_type: str | None = None
     head_dim: int | None = None
diff --git a/src/mobius/_configs/per_model/__init__.py b/src/mobius/_configs/per_model/__init__.py
index fa21685a..44d70cae 100644
--- a/src/mobius/_configs/per_model/__init__.py
+++ b/src/mobius/_configs/per_model/__init__.py
@@ -31,5 +31,6 @@
     _phi4mm_audio,
     _phi4mm_vision,
     _qwen3_asr_audio,
+    _qwen25_omni_vision,
     _sensevoice_audio,
 )
diff --git a/src/mobius/_configs/per_model/_qwen25_omni_vision.py b/src/mobius/_configs/per_model/_qwen25_omni_vision.py
new file mode 100644
index 00000000..a86bf69f
--- /dev/null
+++ b/src/mobius/_configs/per_model/_qwen25_omni_vision.py
@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Qwen2.5-Omni vision extractor (vision config lives under thinker_config)."""
+
+from __future__ import annotations
+
+from mobius._configs._extractors import register_vision_hook
+
+
+@register_vision_hook("qwen2_5_omni_text")
+def _qwen25_omni_vision(config, parent_config, model_type: str, fields: dict):
+    thinker = getattr(parent_config, "thinker_config", None)
+    if thinker is None:
+        return None
+    if isinstance(thinker, dict):
+        thinker = type("ThinkerConfig", (), thinker)()
+    vision = getattr(thinker, "vision_config", None)
+    if vision is None:
+        return None
+    if isinstance(vision, dict):
+        vision = type("VisionConfig", (), vision)()
+
+    fields.update(
+        hidden_size=getattr(vision, "hidden_size", None),
+        intermediate_size=getattr(vision, "intermediate_size", None),
+        num_hidden_layers=getattr(vision, "depth", None),
+        num_attention_heads=getattr(vision, "num_heads", None),
+        patch_size=getattr(vision, "patch_size", None),
+        out_hidden_size=getattr(vision, "out_hidden_size", None),
+        in_channels=getattr(vision, "in_channels", 3),
+        spatial_merge_size=getattr(vision, "spatial_merge_size", 2),
+        temporal_patch_size=getattr(vision, "temporal_patch_size", 2),
+        fullatt_block_indexes=getattr(vision, "fullatt_block_indexes", None),
+        window_size=getattr(vision, "window_size", 112),
+        image_token_id=getattr(thinker, "image_token_id", None),
+    )
+    fields["video_token_id"] = getattr(thinker, "video_token_id", None)
+    return None
diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py
index c8a581ed..55b1251f 100644
--- a/src/mobius/_registry.py
+++ b/src/mobius/_registry.py
@@ -626,11 +626,10 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
         task="speech-to-text",
         config_class=WhisperConfig,
     ),
-
     # --- Omni ---
     "qwen2_5_omni": ModelRegistration(
         Qwen25OmniThinkerForConditionalGeneration,
-        task="speech-language",
+        task="qwen25-omni",
     ),
     # --- Encoder-only ---
     "albert": ModelRegistration(BertModel, task="feature-extraction"),
diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py
index 4ef07186..24eb1d45 100644
--- a/src/mobius/components/__init__.py
+++ b/src/mobius/components/__init__.py
@@ -54,6 +54,8 @@
     "PostNormDecoderLayer",
     "QuantizedEmbedding",
     "QuantizedLinear",
+    "Qwen25OmniAudioAttention",
+    "Qwen25OmniAudioEncoderLayer",
     "RMSNorm",
     "SelectiveScan",
     "SiLU",
@@ -221,6 +223,10 @@
 from mobius.components._qwen3_vl_vision import (
     Qwen3VLVisionRotaryEmbedding as Qwen3VLVisionRotaryEmbedding,
 )
+from mobius.components._qwen25_omni_audio import (
+    Qwen25OmniAudioAttention,
+    Qwen25OmniAudioEncoderLayer,
+)
 from mobius.components._qwen25_vl_vision import (
     Qwen2VLVisionBlock as Qwen2VLVisionBlock,
 )
diff --git a/src/mobius/components/_conv.py b/src/mobius/components/_conv.py
index 80ee70fc..6657841c 100644
--- a/src/mobius/components/_conv.py
+++ b/src/mobius/components/_conv.py
@@ -5,14 +5,9 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 import onnx_ir as ir
 from onnxscript import OpBuilder, nn
 
-if TYPE_CHECKING:
-    pass
-
 
 class Conv2d(nn.Module):
     """2D convolution with bias.
@@ -70,16 +65,14 @@ def __init__(
         groups: int = 1,
     ):
         super().__init__()
-        self.weight = nn.Parameter(
-            (out_channels, in_channels // groups, kernel_size)
-        )
-        self.bias = nn.Parameter((out_channels))
+        self.weight = nn.Parameter((out_channels, in_channels // groups, kernel_size))
+        self.bias = nn.Parameter((out_channels,))
         self._kernel_size = kernel_size
         self._stride = stride
         self._padding = padding
         self._groups = groups
 
-    def forward(self, op: builder.OpBuilder, x: ir.Value):
+    def forward(self, op: OpBuilder, x: ir.Value):
         p = self._padding
         return op.Conv(
             x,
diff --git a/src/mobius/components/_qwen25_omni_audio.py b/src/mobius/components/_qwen25_omni_audio.py
index f2597e13..9d744b5d 100644
--- a/src/mobius/components/_qwen25_omni_audio.py
+++ b/src/mobius/components/_qwen25_omni_audio.py
@@ -1,8 +1,9 @@
-"""Qwen25-Omni audio encoder components.
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
 
-Whisper-inspired audio encoder with 3x Conv1d,
-sinusoidal positional embeddings, and bidirectional transformer
-encoder layers with LayerNorm.
+"""Qwen2.5-Omni audio encoder components.
+
+Packed bidirectional transformer layers with LayerNorm.
 
 Reference: Transformers
 https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py
@@ -10,16 +11,13 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
+import onnx_ir as ir
 from onnxscript import nn
 from onnxscript._internal import builder
 
+from mobius._build_context import get_build_dtype
 from mobius.components._common import LayerNorm, Linear
 
-if TYPE_CHECKING:
-    import onnx_ir as ir
-
 
 class Qwen25OmniAudioAttention(nn.Module):
     """Bidirectional multi-head attention for Qwen2_5Omni audio encoder.
@@ -37,7 +35,12 @@ def __init__(self, d_model: int, num_heads: int):
         self._num_heads = num_heads
         self._head_dim = d_model // num_heads
 
-    def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
+    def forward(
+        self,
+        op: builder.OpBuilder,
+        hidden_states: ir.Value,
+        cu_seqlens: ir.Value,
+    ):
         """Bidirectional self-attention.
 
         Args:
@@ -46,19 +49,53 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
         Returns:
             output: (batch, seq_len, d_model)
         """
-        q = self.q_proj(op, hidden_states)
-        k = self.k_proj(op, hidden_states)
-        v = self.v_proj(op, hidden_states)
+        seq_len = op.Shape(hidden_states, start=0, end=1)
+        packed_shape = op.Concat(seq_len, [self._num_heads, self._head_dim], axis=0)
+        q = op.Reshape(self.q_proj(op, hidden_states), packed_shape)
+        k = op.Reshape(self.k_proj(op, hidden_states), packed_shape)
+        v = op.Reshape(self.v_proj(op, hidden_states), packed_shape)
+
+        # Build the block-diagonal mask represented by HF's cu_seqlens.
+        positions = op.Range(0, op.Squeeze(seq_len, [0]), 1)
+        segment_ids = op.Sub(
+            op.ReduceSum(
+                op.Cast(
+                    op.GreaterOrEqual(
+                        op.Unsqueeze(positions, [1]),
+                        op.Unsqueeze(op.Cast(cu_seqlens, to=7), [0]),
+                    ),
+                    to=7,
+                ),
+                [1],
+                keepdims=False,
+            ),
+            1,
+        )
+        same_segment = op.Equal(
+            op.Unsqueeze(segment_ids, [1]),
+            op.Unsqueeze(segment_ids, [0]),
+        )
+        attention_bias = op.Where(
+            same_segment,
+            op.CastLike(0.0, q),
+            op.CastLike(-1e9, q),
+        )
+        attention_bias = op.Unsqueeze(attention_bias, [0, 1])
 
-        # Use ONNX Attention op (bidirectional: no causal mask)
+        q = op.Unsqueeze(op.Transpose(q, perm=[1, 0, 2]), [0])
+        k = op.Unsqueeze(op.Transpose(k, perm=[1, 0, 2]), [0])
+        v = op.Unsqueeze(op.Transpose(v, perm=[1, 0, 2]), [0])
         attn_output = op.Attention(
             q,
             k,
             v,
+            attention_bias,
             q_num_heads=self._num_heads,
             kv_num_heads=self._num_heads,
             scale=float(self._head_dim**-0.5),
         )
+        attn_output = op.Transpose(op.Squeeze(attn_output, [0]), perm=[1, 0, 2])
+        attn_output = op.Reshape(attn_output, op.Concat(seq_len, [-1], axis=0))
         return self.out_proj(op, attn_output)
 
 
@@ -86,7 +123,12 @@ def __init__(
         self.fc2 = Linear(ffn_dim, d_model, bias=True)
         self.final_layer_norm = LayerNorm(d_model, eps=eps)
 
-    def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
+    def forward(
+        self,
+        op: builder.OpBuilder,
+        hidden_states: ir.Value,
+        cu_seqlens: ir.Value,
+    ):
         """Pre-norm encoder layer with bidirectional attention.
 
         Args:
@@ -98,7 +140,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
         # Self-attention with pre-norm and residual
         residual = hidden_states
         hidden_states = self.self_attn_layer_norm(op, hidden_states)
-        hidden_states = self.self_attn(op, hidden_states)
+        hidden_states = self.self_attn(op, hidden_states, cu_seqlens)
         hidden_states = op.Add(residual, hidden_states)
 
         # FFN with pre-norm, GELU, and residual
@@ -108,5 +150,11 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
         hidden_states = op.Gelu(hidden_states)
         hidden_states = self.fc2(op, hidden_states)
         hidden_states = op.Add(residual, hidden_states)
+        if get_build_dtype() == ir.DataType.FLOAT16:
+            hidden_states = op.Clip(
+                hidden_states,
+                op.CastLike(-64504.0, hidden_states),
+                op.CastLike(64504.0, hidden_states),
+            )
 
         return hidden_states
diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py
index 49c94dd1..1516759f 100644
--- a/src/mobius/models/__init__.py
+++ b/src/mobius/models/__init__.py
@@ -93,6 +93,7 @@
     "Phi4MMMultiModalModel",
     "PhiCausalLMModel",
     "Qwen25VLCausalLMModel",
+    "Qwen25OmniThinkerForConditionalGeneration",
     "Qwen25VLDecoderModel",
     "Qwen25VLEmbeddingModel",
     "Qwen25VLTextModel",
@@ -250,6 +251,7 @@
     Qwen3TTSCodecEncoderModel,
     Qwen3TTSTokenizerV2Model,
 )
+from mobius.models.qwen25_omni import Qwen25OmniThinkerForConditionalGeneration
 from mobius.models.qwen35 import (
     Qwen35CausalLMModel,
     Qwen35MoECausalLMModel,
diff --git a/src/mobius/models/qwen25_omni.py b/src/mobius/models/qwen25_omni.py
index fcdbdbe2..fa54e6d3 100644
--- a/src/mobius/models/qwen25_omni.py
+++ b/src/mobius/models/qwen25_omni.py
@@ -1,4 +1,7 @@
-"""Qwen2.5-Omni: Multimodal model with audio + vision + text.
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Qwen2.5-Omni Thinker: audio + vision + text.
 
 Architecture (Thinker only):
   - Audio encoder: Conv1d x2 → sinusoidal PE → 32 encoder layers → AvgPool → proj
@@ -17,8 +20,16 @@
 import torch
 from onnxscript import OpBuilder, nn
 
+from mobius._build_context import ep_capabilities
 from mobius._configs import ArchitectureConfig
-from mobius.components import Qwen25VLVisionModel
+from mobius.components import (
+    GatedMLP,
+    Qwen25OmniAudioEncoderLayer,
+    Qwen25VLPatchMerger,
+    Qwen25VLVisionAttention,
+    Qwen25VLVisionBlock,
+    Qwen25VLVisionModel,
+)
 from mobius.components._common import (
     Embedding,
     LayerNorm,
@@ -27,7 +38,6 @@
 )
 from mobius.components._conv import Conv1d
 from mobius.components._decoder import DecoderLayer
-from mobius.components._qwen25_omni_audio import Qwen25OmniAudioEncoderLayer
 from mobius.components._rms_norm import RMSNorm
 from mobius.components._rotary_embedding import initialize_rope
 
@@ -73,10 +83,11 @@ def __init__(self, config: ArchitectureConfig):
         assert audio is not None
 
         d_model = audio.d_model or 1280
+        self._d_model = d_model
         num_mel_bin = audio.num_mel_bins or 128
         encoder_layers = audio.encoder_layers or 32
         encoder_heads = audio.encoder_attention_heads or 20
-        encoder_ffn = audio.encoder_ffn_dim or 3584
+        encoder_ffn = audio.encoder_ffn_dim or 5120
         max_source_positions = audio.max_source_positions or 1500
         output_dim = audio.output_dim or 3584
 
@@ -106,20 +117,39 @@ def __init__(self, config: ArchitectureConfig):
         # Output projection: d_model -> output_dim
         self.proj = Linear(d_model, output_dim)
 
-    def forward(self, op: OpBuilder, input_features: ir.Value):
-        """Encode mel spectrogram to audio features.
+    def forward(
+        self,
+        op: OpBuilder,
+        input_features: ir.Value,
+        chunk_lengths: ir.Value,
+        pool_indices: ir.Value,
+    ):
+        """Encode pre-chunked mel spectrograms to packed audio features.
 
         Args:
-            input_features: (batch, num_mel_bins, seq_len) mel spectrogram
+            input_features: (num_chunks, num_mel_bins, max_chunk_len)
+            chunk_lengths: Valid mel-frame count for each chunk.
+            pool_indices: Indices of the first token in each stride-2 pooling pair.
 
         Returns:
-            audio_features: (batch, out_seq_len, output_dim)
+            audio_features: (num_audio_tokens, output_dim)
         """
-        # 2X Conv1d with GELU: (batch, mel, seq) -> (batch, d_model, seq//2)
-        hidden_states = op.Gelu(self.conv1(op, input_features))
+        input_features = op.CastLike(input_features, self.conv1.weight)
+
+        # Match HF's chunk padding mask before the stride-2 convolution.
+        chunk_seq_len = op.Shape(input_features, start=2, end=3)
+        chunk_positions = op.Range(0, op.Squeeze(chunk_seq_len, [0]), 1)
+        chunk_mask = op.Less(
+            op.Unsqueeze(chunk_positions, [0]),
+            op.Unsqueeze(chunk_lengths, [1]),
+        )
+        chunk_mask = op.Unsqueeze(op.CastLike(chunk_mask, input_features), [1])
+
+        # (num_chunks, mel, time) -> (num_chunks, d_model, ceil(time / 2))
+        hidden_states = op.Mul(op.Gelu(self.conv1(op, input_features)), chunk_mask)
         hidden_states = op.Gelu(self.conv2(op, hidden_states))
 
-        # Transpose to (batch, seq//2, d_model) for transformer layers
+        # (num_chunks, d_model, time) -> (num_chunks, time, d_model)
         hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1])
 
         # Add sinusoidal positional embeddings
@@ -132,29 +162,119 @@ def forward(self, op: OpBuilder, input_features: ir.Value):
         )
         hidden_states = op.Add(hidden_states, pe_slice)
 
-        for layer in self.layers:
-            hidden_states = layer(op, hidden_states)
-
-        # AvgPool1d(kernel=2, stride=2): halves sequence length.
-        # Transpose to (batch, d_model, seq) for pooling, then back.
-        hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1])
-        hidden_states = op.AveragePool(hidden_states, kernel_shape=[2], strides=[2])
-        hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1])
+        # Remove per-chunk padding and derive packed-attention boundaries.
+        after_conv_lengths = op.Add(op.Div(op.Sub(chunk_lengths, 1), 2), 1)
+        valid_mask = op.Less(
+            op.Unsqueeze(op.Range(0, op.Squeeze(seq_len, [0]), 1), [0]),
+            op.Unsqueeze(after_conv_lengths, [1]),
+        )
+        valid_indices = op.Squeeze(op.NonZero(op.Reshape(valid_mask, [-1])), [0])
+        hidden_states = op.Gather(
+            op.Reshape(hidden_states, [-1, self._d_model]),
+            valid_indices,
+            axis=0,
+        )
+        cu_seqlens = op.Concat(
+            op.Constant(value_ints=[0]),
+            op.CumSum(after_conv_lengths, op.Constant(value_int=0)),
+            axis=0,
+        )
 
+        for layer in self.layers:
+            hidden_states = layer(op, hidden_states, cu_seqlens)
+
+        # HF pools adjacent valid tokens using indices computed per original audio.
+        pooled_first = op.Gather(hidden_states, pool_indices, axis=0)
+        pooled_second = op.Gather(hidden_states, op.Add(pool_indices, 1), axis=0)
+        hidden_states = op.Mul(
+            op.Add(pooled_first, pooled_second),
+            op.CastLike(0.5, hidden_states),
+        )
         hidden_states = self.ln_post(op, hidden_states)
         hidden_states = self.proj(op, hidden_states)
         return hidden_states
 
 
+class Qwen25OmniVisionAttention(Qwen25VLVisionAttention):
+    """Qwen2.5-Omni vision attention with separate Q/K/V checkpoint weights."""
+
+    def __init__(self, hidden_size: int, num_heads: int):
+        nn.Module.__init__(self)
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.q = Linear(hidden_size, hidden_size, bias=True)
+        self.k = Linear(hidden_size, hidden_size, bias=True)
+        self.v = Linear(hidden_size, hidden_size, bias=True)
+        self.proj = Linear(hidden_size, hidden_size, bias=True)
+
+    def forward(self, op, hidden_states, cu_seqlens, cos, sin):
+        seq_len = op.Shape(hidden_states, start=0, end=1)
+        head_shape = op.Concat(seq_len, [self.num_heads, self.head_dim], axis=0)
+        q = self._apply_rotary(op, op.Reshape(self.q(op, hidden_states), head_shape), cos, sin)
+        k = self._apply_rotary(op, op.Reshape(self.k(op, hidden_states), head_shape), cos, sin)
+        v = op.Reshape(self.v(op, hidden_states), head_shape)
+
+        if ep_capabilities().supports_packed_multi_head_attention:
+            output = self._emit_packed_mha(op, q, k, v, cu_seqlens, seq_len)
+        else:
+            output = self._emit_standard_attention(op, q, k, v, cu_seqlens, seq_len)
+        return self.proj(op, output)
+
+
+class Qwen25OmniVisionBlock(Qwen25VLVisionBlock):
+    """Qwen2.5-Omni vision block with separate attention projections."""
+
+    def __init__(self, hidden_size: int, intermediate_size: int, num_heads: int):
+        super().__init__(hidden_size, intermediate_size, num_heads)
+        self.attn = Qwen25OmniVisionAttention(hidden_size, num_heads)
+        self.mlp = GatedMLP(
+            hidden_size,
+            intermediate_size,
+            activation="silu",
+            bias=True,
+        )
+
+
+class Qwen25OmniVisionModel(Qwen25VLVisionModel):
+    """Qwen2.5-Omni vision tower using the Omni checkpoint parameter layout."""
+
+    def __init__(
+        self,
+        depth: int,
+        hidden_size: int,
+        intermediate_size: int,
+        num_heads: int,
+        **kwargs,
+    ):
+        super().__init__(
+            depth,
+            hidden_size,
+            intermediate_size,
+            num_heads,
+            **kwargs,
+        )
+        self.blocks = nn.ModuleList(
+            [
+                Qwen25OmniVisionBlock(hidden_size, intermediate_size, num_heads)
+                for _ in range(depth)
+            ]
+        )
+        self.merger = Qwen25VLPatchMerger(
+            out_hidden_size=kwargs.get("out_hidden_size") or hidden_size,
+            hidden_size=hidden_size,
+            spatial_merge_size=kwargs.get("spatial_merge_size", 2),
+        )
+
+
 class Qwen25OmniVisionEncoder(nn.Module):
-    """Qwen2.5-Omni vision encoder — reuses the Qwen2.5-VL ViT."""
+    """Qwen2.5-Omni vision encoder."""
 
     def __init__(self, config: ArchitectureConfig):
         super().__init__()
         vc = config.vision
         assert vc is not None
 
-        self.visual = Qwen25VLVisionModel(
+        self.visual = Qwen25OmniVisionModel(
             depth=vc.num_hidden_layers or 32,
             hidden_size=vc.hidden_size or 1280,
             intermediate_size=vc.intermediate_size or 3420,
@@ -175,13 +295,13 @@ def forward(self, op: OpBuilder, pixel_values: ir.Value, image_grid_thw: ir.Valu
 class Qwen25OmniEmbeddingModel(nn.Module):
     """Fuses text embedding with audio and image features.
 
-    Replaces ``audio_token_id`` positions with audio features and
-    ``image_token_id`` positions with image features.
+    Replaces audio, image, and video placeholder tokens with encoder features.
 
     Inputs:
         input_ids: (batch, seq_len)
         audio_features: (num_audio_tokens, hidden_size)
         image_features: (num_image_tokens, hidden_size)
+        video_features: (num_video_tokens, hidden_size)
 
     Output:
         inputs_embeds: (batch, seq_len, hidden_size)
@@ -200,6 +320,7 @@ def __init__(self, config: ArchitectureConfig):
         vision = config.vision
         self._audio_token_id = (audio.audio_token_id if audio else None) or 151646
         self._image_token_id = (vision.image_token_id if vision else None) or 151655
+        self._video_token_id = config.video_token_id or 151656
 
     def forward(
         self,
@@ -207,6 +328,7 @@ def forward(
         input_ids: ir.Value,
         audio_features: ir.Value,
         image_features: ir.Value,
+        video_features: ir.Value,
     ):
         inputs_embeds = self.embed_tokens(op, input_ids)
 
@@ -219,6 +341,9 @@ def forward(
         inputs_embeds = self._replace_tokens(
             op, inputs_embeds, input_ids, image_features, self._image_token_id
         )
+        inputs_embeds = self._replace_tokens(
+            op, inputs_embeds, input_ids, video_features, self._video_token_id
+        )
 
         return inputs_embeds
 
@@ -230,9 +355,7 @@ def _replace_tokens(self, op, inputs_embeds, input_ids, features, token_id):
         # Pad with a zero row for safety (text-only case: no features).
         feature_dim = op.Shape(features, start=1, end=2)
         zero_shape = op.Concat(op.Constant(value_ints=[1]), feature_dim, axis=0)
-        zero_row = op.ConstantOfShape(
-            zero_shape, value=ir.tensor(np.zeros(1, dtype=np.float32))
-        )
+        zero_row = op.Expand(op.CastLike(0.0, features), zero_shape)
         padded = op.Concat(zero_row, features, axis=0)
 
         # CumSum-based per-position gather index. Mask positions get the
@@ -315,7 +438,7 @@ class Qwen25OmniThinkerForConditionalGeneration(nn.Module):
     the Talker / streaming code generation head is out of scope for now).
     """
 
-    default_task: str = "speech-language"
+    default_task: str = "qwen25-omni"
     category: str = "Multimodal"
     config_class: type = ArchitectureConfig
 
@@ -327,20 +450,15 @@ def __init__(self, config: ArchitectureConfig):
         self.vision_encoder: Qwen25OmniVisionEncoder | None = (
             Qwen25OmniVisionEncoder(config) if config.vision is not None else None
         )
-        # Named ``audio_tower`` (not ``audio_encoder``) to match
-        # ``SpeechLanguageTask`` which looks up ``module.audio_tower``.
-        # The vision sub-module is currently not routed through any
-        # task — wiring a dedicated Qwen25OmniTask that builds all four
-        # ONNX sub-models is a follow-up; today the
-        # ``speech-language`` task entry only drives audio + text.
-        self.audio_tower: Qwen25OmniAudioEncoder | None = (
+        self.audio_encoder: Qwen25OmniAudioEncoder | None = (
             Qwen25OmniAudioEncoder(config) if config.audio is not None else None
         )
 
     def forward(self, op: OpBuilder, **kwargs):
         raise NotImplementedError(
             "Qwen25OmniThinkerForConditionalGeneration is a multi-model split; the corresponding "
-            "task class builds each sub-module (decoder, embedding, vision_encoder, audio_encoder) "
+            "Qwen25OmniTask builds each sub-module (decoder, embedding, vision_encoder, "
+            "audio_encoder) "
             "separately."
         )
 
@@ -351,7 +469,7 @@ def preprocess_weights(
 
         HF Qwen2.5-Omni checkpoints prefix every Thinker key with ``thinker.``:
 
-        - ``thinker.audio_tower.*`` → ``audio_tower.*``
+        - ``thinker.audio_tower.*`` → ``audio_encoder.*``
         - ``thinker.visual.*`` → ``vision_encoder.visual.*``
         - ``thinker.model.embed_tokens.*`` → ``embedding.embed_tokens.*``
         - ``thinker.model.layers.N.*`` and ``model.norm.*`` → ``decoder.*``
@@ -368,15 +486,19 @@ def preprocess_weights(
                 key = key[len("thinker.") :]
 
             # Drop talker.* and any codec output keys — not part of Thinker.
-            if key.startswith(("talker.", "code_predictor.")):
+            if key.startswith(("talker.", "token2wav.", "code_predictor.")):
                 continue
 
             if key.startswith("audio_tower."):
-                cleaned[key] = value
+                if ".audio_bos_eos_token." not in key:
+                    cleaned["audio_encoder." + key[len("audio_tower.") :]] = value
                 continue
 
             if key.startswith("visual."):
-                cleaned["vision_encoder." + key] = value
+                new_key = "vision_encoder." + key
+                new_key = new_key.replace(".merger.mlp.0.", ".merger.mlp_0.")
+                new_key = new_key.replace(".merger.mlp.2.", ".merger.mlp_2.")
+                cleaned[new_key] = value
                 continue
 
             if key.startswith("lm_head."):
diff --git a/src/mobius/models/qwen25_omni_test.py b/src/mobius/models/qwen25_omni_test.py
new file mode 100644
index 00000000..990bf74d
--- /dev/null
+++ b/src/mobius/models/qwen25_omni_test.py
@@ -0,0 +1,106 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import torch
+
+from mobius._configs import ArchitectureConfig
+from mobius.models.qwen25_omni import Qwen25OmniThinkerForConditionalGeneration
+
+
+def _hf_config():
+    text = SimpleNamespace(
+        model_type="qwen2_5_omni_text",
+        vocab_size=256,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        hidden_act="silu",
+        rms_norm_eps=1e-6,
+        max_position_embeddings=128,
+        rope_parameters={
+            "rope_type": "default",
+            "rope_theta": 1_000_000.0,
+            "mrope_section": [4, 2, 2],
+        },
+        tie_word_embeddings=False,
+    )
+    thinker = SimpleNamespace(
+        text_config=text,
+        audio_config=SimpleNamespace(
+            d_model=64,
+            encoder_layers=2,
+            encoder_attention_heads=4,
+            encoder_ffn_dim=128,
+            num_mel_bins=32,
+            max_source_positions=128,
+            n_window=8,
+            output_dim=64,
+        ),
+        vision_config=SimpleNamespace(
+            hidden_size=64,
+            intermediate_size=128,
+            depth=2,
+            num_heads=4,
+            patch_size=14,
+            temporal_patch_size=2,
+            in_channels=3,
+            out_hidden_size=64,
+            spatial_merge_size=2,
+            fullatt_block_indexes=[0],
+            window_size=112,
+        ),
+        audio_token_id=100,
+        image_token_id=101,
+        video_token_id=102,
+    )
+    return text, SimpleNamespace(thinker_config=thinker, tie_word_embeddings=False)
+
+
+def test_qwen25_omni_extracts_nested_thinker_config():
+    text, parent = _hf_config()
+    config = ArchitectureConfig.from_transformers(text, parent_config=parent)
+
+    assert config.attn_qkv_bias
+    assert config.audio is not None
+    assert config.audio.encoder_ffn_dim == 128
+    assert config.audio.audio_token_id == 100
+    assert config.vision is not None
+    assert config.vision.hidden_size == 64
+    assert config.image_token_id == 101
+    assert config.video_token_id == 102
+
+
+def test_qwen25_omni_preprocess_weights_routes_thinker_components():
+    text, parent = _hf_config()
+    config = ArchitectureConfig.from_transformers(text, parent_config=parent)
+    model = Qwen25OmniThinkerForConditionalGeneration(config)
+    weight = torch.randn(1)
+
+    processed = model.preprocess_weights(
+        {
+            "thinker.audio_tower.conv1.weight": weight,
+            "thinker.audio_tower.audio_bos_eos_token.weight": weight,
+            "thinker.visual.blocks.0.attn.q.weight": weight,
+            "thinker.visual.merger.mlp.0.weight": weight,
+            "thinker.model.embed_tokens.weight": weight,
+            "thinker.model.layers.0.self_attn.q_proj.bias": weight,
+            "thinker.lm_head.weight": weight,
+            "talker.model.layers.0.weight": weight,
+            "token2wav.dit.weight": weight,
+        }
+    )
+
+    assert set(processed) == {
+        "audio_encoder.conv1.weight",
+        "vision_encoder.visual.blocks.0.attn.q.weight",
+        "vision_encoder.visual.merger.mlp_0.weight",
+        "embedding.embed_tokens.weight",
+        "decoder.layers.0.self_attn.q_proj.bias",
+        "decoder.lm_head.weight",
+    }
diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py
index 2671c100..49e7ba55 100644
--- a/src/mobius/tasks/__init__.py
+++ b/src/mobius/tasks/__init__.py
@@ -31,6 +31,7 @@
     "DFlashDraftTask",
     "Eagle3DraftTask",
     "Qwen35MtpTask",
+    "Qwen25OmniTask",
     "DenoisingTask",
     "FeatureExtractionTask",
     "FunASRSpeechLanguageTask",
@@ -102,6 +103,7 @@
 from mobius.tasks._multimodal import MultiModalTask
 from mobius.tasks._object_detection import ObjectDetectionTask
 from mobius.tasks._phi4mm_multimodal import Phi4MMMultiModalTask
+from mobius.tasks._qwen25_omni import Qwen25OmniTask
 from mobius.tasks._qwen35_mtp import Qwen35MtpTask
 from mobius.tasks._qwen_image_vae import QwenImageVAETask
 from mobius.tasks._rnnt import RNNTTask
@@ -144,6 +146,7 @@
     "dflash-draft": DFlashDraftTask,
     "eagle3-draft": Eagle3DraftTask,
     "qwen35-mtp": Qwen35MtpTask,
+    "qwen25-omni": Qwen25OmniTask,
     "vae": VAETask,
     "qwen-image-vae": QwenImageVAETask,
     "vision-language": VisionLanguageTask,
diff --git a/src/mobius/tasks/_qwen25_omni.py b/src/mobius/tasks/_qwen25_omni.py
new file mode 100644
index 00000000..34e05eb3
--- /dev/null
+++ b/src/mobius/tasks/_qwen25_omni.py
@@ -0,0 +1,123 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Qwen2.5-Omni Thinker four-model split task."""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+import onnx_ir as ir
+from onnxscript import nn
+
+from mobius._configs import ArchitectureConfig
+from mobius._model_package import ModelPackage
+from mobius.tasks._base import (
+    ComponentSpec,
+    _make_graph,
+    _make_model,
+    build_decoder_from_embeds,
+)
+from mobius.tasks._vision_language_3model import QwenVLTask
+
+
+class Qwen25OmniTask(QwenVLTask):
+    """Build the Thinker's audio, vision, embedding, and decoder ONNX models."""
+
+    model_roles: ClassVar[dict[str, str]] = {
+        "audio_encoder": "encoder",
+        "vision_encoder": "encoder",
+        "embedding": "embedding",
+        "decoder": "decoder",
+    }
+    components = ComponentSpec(
+        audio_encoder="audio_encoder",
+        vision_encoder="vision_encoder",
+        embedding="embedding",
+        decoder="decoder",
+    )
+
+    def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage:
+        self._validate_components(module)
+        models = {
+            "audio_encoder": self._build_audio(module.audio_encoder, config),
+            "vision_encoder": self._build_vision(module.vision_encoder, config),
+            "embedding": self._build_embedding(module.embedding, config),
+            "decoder": build_decoder_from_embeds(module.decoder, config, mrope=True),
+        }
+        return ModelPackage(models, config=config)
+
+    def _build_audio(self, audio_encoder: nn.Module, config: ArchitectureConfig) -> ir.Model:
+        """Build packed audio chunks into packed LLM audio tokens."""
+        num_chunks = ir.SymbolicDim("num_audio_chunks")
+        chunk_len = ir.SymbolicDim("audio_chunk_len")
+        num_audio_tokens = ir.SymbolicDim("num_audio_tokens")
+        n_mels = (config.audio.num_mel_bins if config.audio else None) or 128
+
+        graph, builder = _make_graph(name="audio_encoder")
+        input_features = builder.input(
+            "input_features",
+            dtype=ir.DataType.FLOAT,
+            shape=[num_chunks, n_mels, chunk_len],
+        )
+        chunk_lengths = builder.input(
+            "chunk_lengths",
+            dtype=ir.DataType.INT64,
+            shape=[num_chunks],
+        )
+        pool_indices = builder.input(
+            "pool_indices",
+            dtype=ir.DataType.INT64,
+            shape=[num_audio_tokens],
+        )
+        audio_features = audio_encoder(
+            builder.op,
+            input_features,
+            chunk_lengths,
+            pool_indices,
+        )
+        builder.add_output(audio_features, "audio_features")
+        return _make_model(graph)
+
+    def _build_embedding(
+        self,
+        embedding: nn.Module,
+        config: ArchitectureConfig,
+    ) -> ir.Model:
+        """Build text embedding and three-modality feature replacement."""
+        batch = ir.SymbolicDim("batch")
+        seq_len = ir.SymbolicDim("sequence_len")
+        num_audio_tokens = ir.SymbolicDim("num_audio_tokens")
+        num_image_tokens = ir.SymbolicDim("num_image_tokens")
+        num_video_tokens = ir.SymbolicDim("num_video_tokens")
+
+        graph, builder = _make_graph(name="embedding")
+        input_ids = builder.input(
+            "input_ids",
+            dtype=ir.DataType.INT64,
+            shape=[batch, seq_len],
+        )
+        audio_features = builder.input(
+            "audio_features",
+            dtype=config.dtype,
+            shape=[num_audio_tokens, config.hidden_size],
+        )
+        image_features = builder.input(
+            "image_features",
+            dtype=config.dtype,
+            shape=[num_image_tokens, config.hidden_size],
+        )
+        video_features = builder.input(
+            "video_features",
+            dtype=config.dtype,
+            shape=[num_video_tokens, config.hidden_size],
+        )
+        inputs_embeds = embedding(
+            builder.op,
+            input_ids,
+            audio_features,
+            image_features,
+            video_features,
+        )
+        builder.add_output(inputs_embeds, "inputs_embeds")
+        return _make_model(graph)
diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py
index 8b2672cf..1548fbd8 100644
--- a/tests/build_graph_test.py
+++ b/tests/build_graph_test.py
@@ -2862,6 +2862,72 @@ def test_3model_pipeline_runs_with_ort(self):
         assert logits.shape[1] == seq_len
 
 
+class TestBuildGraphQwen25Omni:
+    """Verify the Qwen2.5-Omni Thinker four-model split."""
+
+    def _omni_config(self):
+        return _base_config(
+            model_type="qwen2_5_omni_text",
+            attn_qkv_bias=True,
+            hidden_act="silu",
+            mrope_section=[4, 2, 2],
+            audio=AudioConfig(
+                d_model=64,
+                encoder_layers=2,
+                encoder_attention_heads=4,
+                encoder_ffn_dim=128,
+                num_mel_bins=32,
+                max_source_positions=128,
+                output_dim=64,
+                audio_token_id=100,
+                n_window=8,
+            ),
+            vision=VisionConfig(
+                hidden_size=64,
+                intermediate_size=128,
+                num_hidden_layers=2,
+                num_attention_heads=4,
+                patch_size=14,
+                temporal_patch_size=2,
+                in_channels=3,
+                out_hidden_size=64,
+                spatial_merge_size=2,
+                fullatt_block_indexes=[0],
+                window_size=112,
+                image_token_id=101,
+                video_token_id=102,
+            ),
+            image_token_id=101,
+            video_token_id=102,
+        )
+
+    def test_package_builds_four_models(self):
+        from mobius.models import Qwen25OmniThinkerForConditionalGeneration
+        from mobius.tasks import Qwen25OmniTask
+
+        config = self._omni_config()
+        module = Qwen25OmniThinkerForConditionalGeneration(config)
+        package = build_from_module(module, config, task=Qwen25OmniTask())
+
+        assert set(package) == {
+            "audio_encoder",
+            "vision_encoder",
+            "embedding",
+            "decoder",
+        }
+        assert {value.name for value in package["audio_encoder"].graph.inputs} == {
+            "input_features",
+            "chunk_lengths",
+            "pool_indices",
+        }
+        assert {value.name for value in package["embedding"].graph.inputs} == {
+            "input_ids",
+            "audio_features",
+            "image_features",
+            "video_features",
+        }
+
+
 class TestBuildGraphFunASR:
     """Verify Fun-ASR-Nano 3-model split with FunASRSpeechLanguageTask."""
 
@@ -4938,6 +5004,7 @@ def test_jamba_preprocess_weights_moe_renames(self):
     "mms",
     "qwen3_asr",
     "qwen3_forced_aligner",
+    "qwen2_5_omni",
     "qwen3_tts",
     "qwen3_tts_tokenizer_12hz",
     "whisper",
@@ -4984,13 +5051,6 @@ def test_jamba_preprocess_weights_moe_renames(self):
 _KNOWN_UNTESTED_MODEL_TYPES: set[str] = {
     "deepseek_v2_moe",  # Alias for deepseek_v2 — tested via deepseek_v2
     "qwen3_5_vl_text",  # VL text decoder — tested via parent VL model
-    # Qwen2.5-Omni Thinker is a 4-model split (audio + vision + embedding +
-    # decoder). The model graph is implemented but the existing
-    # SpeechLanguageTask only knows audio + text, so wiring a dedicated
-    # task that also drives the vision encoder is a follow-up — tracked
-    # in the Qwen2.5-Omni onboarding PR (#193). When that lands, move
-    # the entry into a dedicated test (mirroring gemma4_any_to_any).
-    "qwen2_5_omni",
 }