Onboard Qwen2.5-Omni model

src/mobius/_configs/_base.py

@@ -1089,6 +1089,7 @@ def _per_layer_value(attribute: str) -> int | None:
                                         "bloom",
                                         "qwen2",
                                         "qwen2_5_vl_text",
+                                        "qwen2_5_omni_text",
                                         "qwen2_moe",
                                         "qwen2_vl_text",
                                     ),

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.
     """

src/mobius/_configs/per_model/__init__.py

@@ -41,6 +41,7 @@
     _phi4mm_vision,
     _phi_vision,
     _qwen3_asr_audio,
+    _qwen25_omni_vision,
     _sensenova_u1_vision,
     _sensevoice_audio,
 )

src/mobius/_configs/per_model/_qwen25_omni_vision.py

@@ -0,0 +1,43 @@
+# 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", getattr(thinker, "image_token_index", None)
+        ),
+    )
+    fields["video_token_id"] = getattr(
+        thinker, "video_token_id", getattr(thinker, "video_token_index", None)
+    )
+    return None

src/mobius/_configs/per_model/_qwen3_asr_audio.py

@@ -41,7 +41,9 @@ def _qwen3_asr_audio(config, parent_config, model_type: str, fields: dict):
             n_window_infer=getattr(ac, "n_window_infer", None),
         )
     # Special tokens from thinker config
-    fields["audio_token_id"] = getattr(tc, "audio_token_id", None)
+    fields["audio_token_id"] = getattr(
+        tc, "audio_token_id", getattr(tc, "audio_token_index", None)
+    )
     fields["audio_start_token_id"] = getattr(tc, "audio_start_token_id", None)
     fields["audio_end_token_id"] = getattr(tc, "audio_end_token_id", None)
     fields["classify_num"] = getattr(tc, "classify_num", None)

src/mobius/_registry.py

@@ -225,6 +225,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.sensenova_u1 import SenseNovaU1Model
@@ -962,6 +963,11 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
         task="speech-to-text",
         config_class=WhisperConfig,
     ),
+    # --- Omni ---
+    "qwen2_5_omni": ModelRegistration(
+        Qwen25OmniThinkerForConditionalGeneration,
+        task="qwen25-omni",
+    ),
     "moonshine": ModelRegistration(
         MoonshineForConditionalGeneration,
         task="speech-to-text",

src/mobius/components/__init__.py

@@ -111,6 +111,8 @@
     "PixtralProjector",
     "QuantizedEmbedding",
     "QuantizedLinear",
+    "Qwen25OmniAudioAttention",
+    "Qwen25OmniAudioEncoderLayer",
     "RadioVisionModel",
     "RMSNorm",
     "RMSNormBias",
@@ -416,6 +418,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,
 )

src/mobius/components/_conv.py

@@ -6,14 +6,10 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
 
 import onnx_ir as ir
 from onnxscript import OpBuilder, nn
 
-if TYPE_CHECKING:
-    pass
-
 
 def _pair(value: int | Sequence[int], name: str) -> tuple[int, int]:
     if isinstance(value, int):
@@ -103,6 +99,43 @@ def forward(self, op: 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: 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.
 

src/mobius/components/_qwen25_omni_audio.py

@@ -0,0 +1,157 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""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
+"""
+
+from __future__ import annotations
+
+import onnx_ir as ir
+from onnxscript import OpBuilder, nn
+
+from mobius._build_context import get_build_dtype
+from mobius.components._common import LayerNorm, Linear
+
+
+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: OpBuilder,
+        hidden_states: ir.Value,
+        cu_seqlens: ir.Value,
+    ):
+        """Bidirectional self-attention.
+
+        Args:
+            hidden_states: (total_tokens, d_model)
+            cu_seqlens: Cumulative lengths of the independent audio chunks.
+
+        Returns:
+            output: (total_tokens, d_model)
+        """
+        seq_len = op.Shape(hidden_states, start=0, end=1)
+        # Rank-3 Attention splits the projected hidden dimension into heads.
+        q = op.Unsqueeze(self.q_proj(op, hidden_states), [0])
+        k = op.Unsqueeze(self.k_proj(op, hidden_states), [0])
+        v = op.Unsqueeze(self.v_proj(op, hidden_states), [0])
+
+        # 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])
+
+        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.Squeeze(attn_output, [0])
+        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: OpBuilder,
+        hidden_states: ir.Value,
+        cu_seqlens: ir.Value,
+    ):
+        """Pre-norm encoder layer with bidirectional attention.
+
+        Args:
+            hidden_states: (total_tokens, d_model)
+            cu_seqlens: Cumulative lengths of the independent audio chunks.
+
+        Returns:
+            hidden_states: (total_tokens, 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, cu_seqlens)
+        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)
+        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

src/mobius/integrations/transformers/_builder.py

@@ -221,7 +221,8 @@ def _select_primary_config(hf_config):
     parent_config = hf_config
     model_type = hf_config.model_type
 
-    if hasattr(hf_config, "talker_config"):
+    # Omni exports Thinker; its sibling Talker has a separate decoder config.
+    if hasattr(hf_config, "talker_config") and model_type != "qwen2_5_omni":
         hf_config = hf_config.talker_config
     elif hasattr(hf_config, "thinker_config"):
         thinker = hf_config.thinker_config

src/mobius/integrations/transformers/_builder_test.py

@@ -391,6 +391,17 @@ def build_module(_module, config, *args, **kwargs):
     assert built_configs[0].image_token_id == 248056
 
 
+def test_qwen3_tts_primary_config_remains_talker() -> None:
+    talker = SimpleNamespace(model_type="qwen3_tts_talker")
+    hf_config = SimpleNamespace(model_type="qwen3_tts", talker_config=talker)
+
+    primary, parent, model_type = transformers_builder._select_primary_config(hf_config)
+
+    assert primary is talker
+    assert parent is hf_config
+    assert model_type == "qwen3_tts"
+
+
 def test_vibevoice_architecture_dispatch_keeps_native_asr_separate() -> None:
     """Native ASR has its own model_type and cannot fall through to VibeVoice TTS."""
     from mobius.models import (

src/mobius/models/__init__.py

@@ -157,6 +157,7 @@
     "Phi4SigLIPModel",
     "PhiCausalLMModel",
     "Qwen25VLCausalLMModel",
+    "Qwen25OmniThinkerForConditionalGeneration",
     "Qwen25VLDecoderModel",
     "Qwen25VLEmbeddingModel",
     "Qwen25VLTextModel",
@@ -412,6 +413,7 @@
     Qwen4ExpCausalLMModel,
     Qwen4ExpForConditionalGeneration,
 )
+from mobius.models.qwen25_omni import Qwen25OmniThinkerForConditionalGeneration
 from mobius.models.qwen35 import (
     Qwen35CausalLMModel,
     Qwen35MoECausalLMModel,