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>

Author: Copilot

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

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

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.
@@ -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,

src/mobius/components/_qwen25_omni_audio.py

@@ -1,25 +1,23 @@
-"""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
 """
 
 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

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,