Onboard Qwen2.5-Omni Thinker model

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",
                                     ),

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

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

src/mobius/_configs/per_model/__init__.py

@@ -31,5 +31,6 @@
     _phi4mm_audio,
     _phi4mm_vision,
     _qwen3_asr_audio,
+    _qwen25_omni_vision,
     _sensevoice_audio,
 )

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

src/mobius/_registry.py

@@ -127,6 +127,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.sensevoice_small import SenseVoiceSmallModel
@@ -625,6 +626,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",
+    ),
     # --- Encoder-only ---
     "albert": ModelRegistration(BertModel, task="feature-extraction"),
     "bert": ModelRegistration(BertModel, task="feature-extraction"),

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,
 )

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.
@@ -53,6 +48,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,160 @@
+# 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 nn
+from onnxscript._internal import builder
+
+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: builder.OpBuilder,
+        hidden_states: ir.Value,
+        cu_seqlens: ir.Value,
+    ):
+        """Bidirectional self-attention.
+
+        Args:
+            hidden_states: (batch, seq_len, d_model)
+
+        Returns:
+            output: (batch, seq_len, d_model)
+        """
+        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])
+
+        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)
+
+
+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,
+        cu_seqlens: 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, 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/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,