diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index ee63aea29..5ebff24d8 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -338,11 +338,11 @@ def _extract_mrope_fields(config) -> dict: ) if mrope_interleaved: result["mrope_interleaved"] = True - section = rope_scaling.get("mrope_section", None) or rope_parameters.get( - "mrope_section", None - ) - if section is not None: - result["mrope_section"] = section + section = rope_scaling.get("mrope_section", None) or rope_parameters.get( + "mrope_section", None + ) + if section is not None: + result["mrope_section"] = section return result @@ -481,6 +481,7 @@ class BaseModelConfig: vocab_size: int = DEFAULT_INT hidden_size: int = DEFAULT_INT + embedding_size: int | None = None intermediate_size: int = DEFAULT_INT num_hidden_layers: int = DEFAULT_INT num_attention_heads: int = DEFAULT_INT @@ -774,6 +775,9 @@ class ArchitectureConfig(BaseModelConfig): mrope_section: list[int] | None = None mrope_interleaved: bool = False + # Qwen2.5-Omni uses independent Thinker and Talker decoder dimensions. + talker: ArchitectureConfig | None = None + # Standalone vision config image_size: int = 224 patch_size: int = 16 @@ -1016,6 +1020,7 @@ def _per_layer_value(attribute: str) -> int | None: or 0 ), hidden_size=_as_int(hidden_size), + embedding_size=getattr(config, "embedding_size", None), intermediate_size=_as_int( getattr(config, "intermediate_size", None) or getattr(config, "mlp_hidden_size", None) @@ -1090,6 +1095,7 @@ def _per_layer_value(attribute: str) -> int | None: "qwen2", "qwen2_5_vl_text", "qwen2_5_omni_text", + "qwen2_5_omni_talker", "qwen2_moe", "qwen2_vl_text", ), @@ -1502,6 +1508,13 @@ def _per_layer_value(attribute: str) -> int | None: upsampling_ratios=list(getattr(ec, "upsampling_ratios", [8, 6, 5, 4])), ) + if model_type == "qwen2_5_omni_text" and parent_config is not None: + talker_config = getattr(parent_config, "talker_config", None) + if talker_config is not None: + if isinstance(talker_config, dict): + talker_config = type("TalkerConfig", (), talker_config)() + options["talker"] = ArchitectureConfig.from_transformers(talker_config) + # Model dtype resolved = _resolve_dtype(config) if resolved is not None: diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 1cd51caef..cac148176 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -1439,6 +1439,7 @@ def _create_default_registry() -> ModelRegistry: "moonshine": "moonshine-ai/moonshine-tiny", "moonshine_streaming": "moonshine-ai/moonshine-streaming-tiny", "whisper": "openai/whisper-tiny", + "qwen2_5_omni": "Qwen/Qwen2.5-Omni-7B", "qwen3_asr": "Qwen/Qwen3-ASR-0.6B", "fun_asr": "justinchuby/Fun-ASR-Nano-2512", "glmasr": "zai-org/GLM-ASR-Nano-2512", diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index 7e4fa1183..b2fe4cc2c 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -598,7 +598,16 @@ def build_transformers_model( ) is_gptoss_mxfp4_source = _is_native_gptoss_mxfp4(config) if dtype is not None: - config = dataclasses.replace(config, dtype=resolve_dtype(dtype)) + resolved_dtype = resolve_dtype(dtype) + talker_config = getattr(config, "talker", None) + if talker_config is not None: + config = dataclasses.replace( + config, + dtype=resolved_dtype, + talker=dataclasses.replace(talker_config, dtype=resolved_dtype), + ) + else: + config = dataclasses.replace(config, dtype=resolved_dtype) elif compressed_tensors_config is not None and keep_quantized: # The pinned Microsoft block-weight ABI is W4A16/W8A16 with FP16 A/Y. config = dataclasses.replace(config, dtype=ir.DataType.FLOAT16) diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index e4acc012b..a2ce3446e 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -158,6 +158,8 @@ "PhiCausalLMModel", "Qwen25VLCausalLMModel", "Qwen25OmniThinkerForConditionalGeneration", + "Qwen25OmniTalkerForConditionalGeneration", + "Qwen25OmniTalkerModel", "Qwen25VLDecoderModel", "Qwen25VLEmbeddingModel", "Qwen25VLTextModel", @@ -413,7 +415,11 @@ Qwen4ExpCausalLMModel, Qwen4ExpForConditionalGeneration, ) -from mobius.models.qwen25_omni import Qwen25OmniThinkerForConditionalGeneration +from mobius.models.qwen25_omni import ( + Qwen25OmniTalkerForConditionalGeneration, + Qwen25OmniTalkerModel, + 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 f5a00ceb9..bf0530038 100644 --- a/src/mobius/models/qwen25_omni.py +++ b/src/mobius/models/qwen25_omni.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Qwen2.5-Omni Thinker: audio + vision + text. +"""Qwen2.5-Omni: Thinker and Talker models. Architecture (Thinker only): - Audio encoder: Conv1d x2 → sinusoidal PE → 32 encoder layers → AvgPool → proj @@ -419,21 +419,101 @@ def forward( hidden_states = self.norm(op, hidden_states) logits = self.lm_head(op, hidden_states) - return logits, present_key_values + return logits, hidden_states, present_key_values + + +class Qwen25OmniTalkerModel(nn.Module): + """Talker backbone with codec embedding, MRoPE decoder layers, and final norm.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + embedding_size = config.embedding_size or config.hidden_size + self._dtype = config.dtype + self.embed_tokens = Embedding(config.vocab_size, embedding_size) + 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) + + 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) if self.rotary_emb is not None else None + ) + 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) + + return self.norm(op, hidden_states), present_key_values + + +class Qwen25OmniTalkerForConditionalGeneration(nn.Module): + """Generate codec-token logits from Thinker-width input embeddings.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + embedding_size = config.embedding_size or config.hidden_size + self.thinker_to_talker_proj = Linear(embedding_size, config.hidden_size, bias=True) + self.model = Qwen25OmniTalkerModel(config) + self.codec_head = Linear(config.hidden_size, config.vocab_size, bias=False) + + def forward( + self, + op: OpBuilder, + inputs_embeds: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values=None, + ): + # The host combines Thinker reply states, text embeddings, and codec + # embeddings in the shared embedding space before this projection. + hidden_states = self.thinker_to_talker_proj(op, inputs_embeds) + hidden_states, present_key_values = self.model( + op, + inputs_embeds=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + return self.codec_head(op, hidden_states), present_key_values class Qwen25OmniThinkerForConditionalGeneration(nn.Module): - """Qwen2.5-Omni Thinker: composite audio + vision + text model. + """Qwen2.5-Omni composite Thinker and optional Talker model. - Builds four separate ONNX models: + Builds four Thinker models and, when configured, two Talker models: - ``decoder``: Qwen2.5 text decoder taking ``inputs_embeds`` - ``vision_encoder``: Qwen2.5-VL ViT (pixel_values + grid_thw → image features) - ``audio_encoder``: 2x Conv1d + transformer audio tower (mel → audio features) - ``embedding``: word embedding + multimodal feature fusion + - ``talker_embedding``: codec token embedding in the Thinker-width space + - ``talker``: projection + speech-token decoder + codec logits - HuggingFace class: ``Qwen2_5OmniForConditionalGeneration`` (Thinker only — - the Talker / streaming code generation head is out of scope for now). + HuggingFace class: ``Qwen2_5OmniForConditionalGeneration``. """ default_task: str = "qwen25-omni" @@ -451,12 +531,17 @@ def __init__(self, config: ArchitectureConfig): self.audio_encoder: Qwen25OmniAudioEncoder | None = ( Qwen25OmniAudioEncoder(config) if config.audio is not None else None ) + self.talker: Qwen25OmniTalkerForConditionalGeneration | None = ( + Qwen25OmniTalkerForConditionalGeneration(config.talker) + if config.talker is not None + else None + ) def forward(self, op: OpBuilder, **kwargs): raise NotImplementedError( "Qwen25OmniThinkerForConditionalGeneration is a multi-model split; the corresponding " "Qwen25OmniTask builds each sub-module (decoder, embedding, vision_encoder, " - "audio_encoder) " + "audio_encoder, talker, and talker_embedding) " "separately." ) @@ -474,17 +559,27 @@ def preprocess_weights( - ``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. + Talker keys already align with the nested ``talker.*`` module and are + retained when audio output is enabled. """ cleaned: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): + if key == "talker.model.embed_tokens.weight": + # The embedding component traces this nested module directly, + # so its initializer has no outer Talker scopes. + cleaned["embed_tokens.weight"] = value + continue + if key.startswith("talker."): + if self.talker is not None: + cleaned[key] = value + continue + # 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.", "token2wav.", "code_predictor.")): + # Token2wav is exported independently from the Talker. + if key.startswith(("token2wav.", "code_predictor.")): continue if key.startswith("audio_tower."): diff --git a/src/mobius/models/qwen25_omni_test.py b/src/mobius/models/qwen25_omni_test.py index 6c865d41e..852579fdd 100644 --- a/src/mobius/models/qwen25_omni_test.py +++ b/src/mobius/models/qwen25_omni_test.py @@ -8,8 +8,10 @@ import pytest import torch +from mobius import build_from_module from mobius._configs import ArchitectureConfig from mobius.models.qwen25_omni import Qwen25OmniThinkerForConditionalGeneration +from mobius.tasks import Qwen25OmniTask def _hf_config(): @@ -60,7 +62,28 @@ def _hf_config(): image_token_id=101, video_token_id=102, ) - return text, SimpleNamespace(thinker_config=thinker, tie_word_embeddings=False) + talker = SimpleNamespace( + model_type="qwen2_5_omni_talker", + vocab_size=8448, + embedding_size=64, + hidden_size=32, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + hidden_act="silu", + rms_norm_eps=1e-6, + max_position_embeddings=128, + rope_scaling={"rope_type": "default", "mrope_section": [2, 1, 1]}, + rope_theta=1_000_000.0, + ) + return text, SimpleNamespace( + model_type="qwen2_5_omni", + thinker_config=thinker, + talker_config=talker, + tie_word_embeddings=False, + ) def test_qwen25_omni_extracts_nested_thinker_config(): @@ -75,10 +98,17 @@ def test_qwen25_omni_extracts_nested_thinker_config(): assert config.vision.hidden_size == 64 assert config.image_token_id == 101 assert config.video_token_id == 102 + assert config.hidden_size == 64 + assert config.talker is not None + assert config.talker.embedding_size == 64 + assert config.talker.hidden_size == 32 + assert config.talker.vocab_size == 8448 + assert config.talker.attn_qkv_bias + assert config.talker.mrope_section == [2, 1, 1] @pytest.mark.parametrize("raw_json", [False, True]) -def test_qwen25_omni_full_checkpoint_builds_thinker(monkeypatch, raw_json): +def test_qwen25_omni_full_checkpoint_builds_thinker_and_talker(monkeypatch, raw_json): from transformers import Qwen2_5OmniConfig from mobius.integrations.transformers import _builder @@ -100,7 +130,14 @@ def test_qwen25_omni_full_checkpoint_builds_thinker(monkeypatch, raw_json): package = _builder.build_transformers_model("test/qwen25-omni", load_weights=False) - assert set(package) == {"audio_encoder", "vision_encoder", "embedding", "decoder"} + assert set(package) == { + "audio_encoder", + "vision_encoder", + "embedding", + "decoder", + "talker_embedding", + "talker", + } assert package.config.vocab_size == 256 assert package.config.hidden_size == 64 assert package.config.audio.audio_token_id == 100 @@ -123,7 +160,10 @@ def test_qwen25_omni_preprocess_weights_routes_thinker_components(): "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, + "talker.thinker_to_talker_proj.weight": weight, + "talker.model.embed_tokens.weight": weight, + "talker.model.layers.0.self_attn.q_proj.weight": weight, + "talker.codec_head.weight": weight, "token2wav.dit.weight": weight, } ) @@ -135,4 +175,73 @@ def test_qwen25_omni_preprocess_weights_routes_thinker_components(): "embedding.embed_tokens.weight", "decoder.layers.0.self_attn.q_proj.bias", "decoder.lm_head.weight", + "talker.thinker_to_talker_proj.weight", + "embed_tokens.weight", + "talker.model.layers.0.self_attn.q_proj.weight", + "talker.codec_head.weight", + } + + +def test_qwen25_omni_package_builds_talker_models(): + text, parent = _hf_config() + config = ArchitectureConfig.from_transformers(text, parent_config=parent) + module = Qwen25OmniThinkerForConditionalGeneration(config) + package = build_from_module(module, config, task=Qwen25OmniTask()) + + assert set(package) == { + "audio_encoder", + "vision_encoder", + "embedding", + "decoder", + "talker_embedding", + "talker", + } + assert {value.name for value in package["talker_embedding"].graph.inputs} == {"input_ids"} + assert set(package["talker_embedding"].graph.initializers) == {"embed_tokens.weight"} + decoder_outputs = {value.name: value for value in package["decoder"].graph.outputs} + assert set(decoder_outputs) >= { + "logits", + "hidden_states", + "present.0.key", + "present.0.value", + } + assert [str(dim) for dim in decoder_outputs["hidden_states"].shape] == [ + "batch", + "sequence_len", + "64", + ] + assert {value.name for value in package["talker"].graph.inputs} >= { + "inputs_embeds", + "attention_mask", + "position_ids", + "past_key_values.0.key", + "past_key_values.0.value", + } + assert {value.name for value in package["talker"].graph.outputs} >= { + "logits", + "present.0.key", + "present.0.value", + } + talker_outputs = {value.name: value for value in package["talker"].graph.outputs} + for name in ("present.0.key", "present.0.value"): + assert [str(dim) for dim in talker_outputs[name].shape] == [ + "batch", + "2", + "past_sequence_len + sequence_len", + "8", + ] + + +def test_qwen25_omni_package_omits_talker_models_when_disabled(): + text, parent = _hf_config() + config = ArchitectureConfig.from_transformers(text, parent_config=parent) + config.talker = None + module = Qwen25OmniThinkerForConditionalGeneration(config) + package = build_from_module(module, config, task=Qwen25OmniTask()) + + assert set(package) == { + "audio_encoder", + "vision_encoder", + "embedding", + "decoder", } diff --git a/src/mobius/tasks/_qwen25_omni.py b/src/mobius/tasks/_qwen25_omni.py index 5caf41c0e..6c2275917 100644 --- a/src/mobius/tasks/_qwen25_omni.py +++ b/src/mobius/tasks/_qwen25_omni.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Qwen2.5-Omni Thinker four-model split task.""" +"""Qwen2.5-Omni Thinker and Talker split task.""" from __future__ import annotations @@ -16,19 +16,24 @@ ComponentSpec, _make_graph, _make_model, - build_decoder_from_embeds, +) +from mobius.tasks._cache_utils import ( + _make_kv_cache_inputs, + _register_kv_cache_outputs, ) from mobius.tasks._vision_language_3model import QwenVLTask class Qwen25OmniTask(QwenVLTask): - """Build the Thinker's audio, vision, embedding, and decoder ONNX models.""" + """Build the Thinker models and optional Talker ONNX models.""" model_roles: ClassVar[dict[str, str]] = { "audio_encoder": "encoder", "vision_encoder": "encoder", "embedding": "embedding", "decoder": "decoder", + "talker_embedding": "embedding", + "talker": "decoder", } components = ComponentSpec( audio_encoder="audio_encoder", @@ -42,14 +47,148 @@ def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage: for name in ("audio_encoder", "vision_encoder"): if getattr(module, name) is None: raise ValueError(f"Qwen25OmniTask requires a non-None {name}.") - models = { + models: dict[str, ir.Model] = { "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), + "decoder": self._build_decoder(module.decoder, config), } + if module.talker is not None and config.talker is not None: + models["talker_embedding"] = self._build_talker_embedding( + module.talker.model.embed_tokens, config.talker + ) + models["talker"] = self._build_talker(module.talker, config.talker) return ModelPackage(models, config=config) + def _build_decoder( + self, + decoder: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build Thinker logits, normalized hidden states, and KV cache.""" + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + past_seq_len = ir.SymbolicDim("past_sequence_len") + graph, builder = _make_graph(name="decoder") + inputs_embeds = builder.input( + "inputs_embeds", + dtype=config.dtype, + shape=[batch, seq_len, config.hidden_size], + ) + attention_mask = builder.input( + "attention_mask", + dtype=ir.DataType.INT64, + shape=[batch, "past_sequence_len + sequence_len"], + ) + position_ids = builder.input( + "position_ids", + dtype=ir.DataType.INT64, + shape=[3, batch, seq_len], + ) + past_key_values = _make_kv_cache_inputs( + builder, + config.num_hidden_layers, + config.num_key_value_heads, + config.head_dim, + config.dtype, + batch, + past_seq_len, + ) + logits, hidden_states, present_key_values = decoder( + builder.op, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + builder.add_output(logits, "logits") + builder.add_output(hidden_states, "hidden_states") + _register_kv_cache_outputs( + builder, + present_key_values, + batch=batch, + num_kv_heads=config.num_key_value_heads, + key_head_dim=config.head_dim, + value_head_dim=config.head_dim, + total_seq_len="past_sequence_len + sequence_len", + dtype=config.dtype, + ) + return _make_model(graph) + + def _build_talker_embedding( + self, + embedding: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build codec-token embeddings in the shared Thinker-width space.""" + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + graph, builder = _make_graph(name="talker_embedding") + input_ids = builder.input( + "input_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], + ) + codec_embeds = embedding(builder.op, input_ids) + builder.add_output(codec_embeds, "codec_embeds") + return _make_model(graph) + + def _build_talker( + self, + talker: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build speech-token logits from preconstructed shared-space embeddings.""" + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + past_seq_len = ir.SymbolicDim("past_sequence_len") + embedding_size = config.embedding_size or config.hidden_size + + graph, builder = _make_graph(name="talker") + inputs_embeds = builder.input( + "inputs_embeds", + dtype=config.dtype, + shape=[batch, seq_len, embedding_size], + ) + attention_mask = builder.input( + "attention_mask", + dtype=ir.DataType.INT64, + shape=[batch, "past_sequence_len + sequence_len"], + ) + position_ids = builder.input( + "position_ids", + dtype=ir.DataType.INT64, + shape=[3, batch, seq_len], + ) + past_key_values = _make_kv_cache_inputs( + builder, + config.num_hidden_layers, + config.num_key_value_heads, + config.head_dim, + config.dtype, + batch, + past_seq_len, + ) + logits, present_key_values = talker( + builder.op, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + builder.add_output(logits, "logits") + _register_kv_cache_outputs( + builder, + present_key_values, + batch=batch, + num_kv_heads=config.num_key_value_heads, + key_head_dim=config.head_dim, + value_head_dim=config.head_dim, + total_seq_len="past_sequence_len + sequence_len", + dtype=config.dtype, + ) + return _make_model(graph) + 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") diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 885b4caa4..a0b7d8572 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -3625,7 +3625,7 @@ def vl_overrides(model_type: str) -> dict: }, True, ), - # --- Qwen2.5-Omni Thinker (audio + vision + embedding + decoder) --- + # --- Qwen2.5-Omni (audio + vision + text Thinker, speech Talker) --- ( "qwen2_5_omni", { @@ -3660,6 +3660,21 @@ def vl_overrides(model_type: str) -> dict: ), "image_token_id": 101, "video_token_id": 102, + "talker": ArchitectureConfig( + model_type="qwen2_5_omni_talker", + vocab_size=128, + embedding_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + attn_qkv_bias=True, + hidden_act="silu", + mrope_section=[2, 1, 1], + rope_type="default", + ), }, True, ), diff --git a/tests/build_graph/speech_test.py b/tests/build_graph/speech_test.py index d1fd1372b..77f209eac 100644 --- a/tests/build_graph/speech_test.py +++ b/tests/build_graph/speech_test.py @@ -46,7 +46,7 @@ class TestBuildGraphQwen25Omni: - """Verify the Qwen2.5-Omni Thinker four-model split.""" + """Verify the Qwen2.5-Omni Thinker and Talker six-model split.""" def _omni_config(self): overrides = next( @@ -59,7 +59,7 @@ def _omni_config(self): @pytest.mark.parametrize( "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] ) - def test_package_builds_four_models(self, dtype): + def test_package_builds_six_models(self, dtype): from mobius.models import Qwen25OmniThinkerForConditionalGeneration from mobius.tasks import Qwen25OmniTask @@ -73,6 +73,8 @@ def test_package_builds_four_models(self, dtype): "vision_encoder", "embedding", "decoder", + "talker_embedding", + "talker", } assert {value.name for value in package["audio_encoder"].graph.inputs} == { "input_features", @@ -100,7 +102,14 @@ def test_missing_encoder_rejected(self, missing): _SPEECH_TASK_KEYS: dict[str, set[str]] = { - "qwen25-omni": {"audio_encoder", "vision_encoder", "embedding", "decoder"}, + "qwen25-omni": { + "audio_encoder", + "vision_encoder", + "embedding", + "decoder", + "talker_embedding", + "talker", + }, "speech-to-text": {"encoder", "decoder"}, "speech-language": {"audio_encoder", "embedding", "decoder"}, "codec": {"decoder", "encoder"}, diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index 7167e670e..255082ba8 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -294,6 +294,9 @@ def _all_registered_with_test_id() -> dict[str, str]: "molmo": "VL model — requires image inputs", "phi4_multimodal": "VL model (14B) — needs GPU for golden", "phi4mm": "VL model (14B) — needs GPU for golden", + "qwen2_5_omni": "Six-stage audio, vision, Thinker, and Talker pipeline — L1 graph " + "and L3 Thinker parity are covered; the generic L4/L5 runner cannot orchestrate " + "multimodal inputs and codec-token generation", "qwen2_5_vl": "VL model — requires image inputs", "qwen2_vl": "VL model — requires image inputs", "qwen3_5": "VL model — hybrid VL, requires image inputs", diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index b05ad73c4..66840231a 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -63,6 +63,7 @@ def test_qwen2_5_omni_thinker_synthetic_parity(): overrides = next(o for mt, o, _ in SPEECH_CONFIGS if mt == "qwen2_5_omni") config = _base_config(**overrides) + config.talker = None audio, vision = config.audio, config.vision hf_config = Qwen2_5OmniThinkerConfig( audio_config={