Onboard Qwen2.5-Omni model

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

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

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
+
+
+

src/mobius/models/qwen25_omni.py

@@ -0,0 +1,179 @@
+"""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
+"""
+
+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):
+    """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
+
+class Qwen25OmniEmbeddingModel(nn.Module):
+    pass
+
+class Qwen25OmniDecoderModel(nn.Module):
+    pass
+
+class Qwen25OmniThinkerForConditionalGeneration(nn.Module):
+    pass
+
+