[OpenVINO] Support Qwen3-Omni MoE with full Talker speech stack (#1700)

* Qwen 3 Omni MoE support

* CI-related fixes

* Test fixes

* latest fixes

* Fixes

* Polishing

* Remaining fixes

* CI fix

Author: sgonorov

docs/source/openvino/models.mdx

@@ -139,6 +139,7 @@ Here is the list of the supported architectures :
 - Qwen3
 - Qwen3MoE
 - Qwen3-VL (including Qwen3-VL-Embedding)
+- Qwen3-Omni-MoE
 - Qwen3.5
 - Qwen3.5-MoE
 - Qwen3.6

optimum/exporters/openvino/__main__.py

@@ -32,6 +32,7 @@
 from optimum.exporters.tasks import TasksManager
 from optimum.intel.utils.import_utils import (
     DIFFUSERS_IMPORT_ERROR,
+    _transformers_version,
     is_diffusers_available,
     is_nncf_available,
     is_openvino_tokenizers_available,
@@ -176,6 +177,38 @@ def infer_task(
     return task
 
 
+def _ensure_qwen3_omni_rope_scaling(config):
+    """Fill in a valid ``rope_scaling`` for Qwen3-Omni-MoE sub-configs that are missing it."""
+    thinker_config = getattr(config, "thinker_config", None)
+    talker_config = getattr(config, "talker_config", None)
+    sub_configs = [
+        getattr(thinker_config, "text_config", None),
+        getattr(talker_config, "text_config", None),
+        getattr(talker_config, "code_predictor_config", None),
+    ]
+
+    for sub_config in sub_configs:
+        if sub_config is None or getattr(sub_config, "rope_scaling", None) is not None:
+            continue
+
+        head_dim = getattr(sub_config, "head_dim", None)
+        if head_dim is None:
+            hidden_size = getattr(sub_config, "hidden_size", None)
+            num_heads = getattr(sub_config, "num_attention_heads", None)
+            if not hidden_size or not num_heads:
+                continue
+            head_dim = hidden_size // num_heads
+
+        # mrope splits head_dim // 2 across the (temporal, height, width) axes; distribute the
+        # remainder to the leading axis, matching how real checkpoints are laid out.
+        half = head_dim // 2
+        base, remainder = divmod(half, 3)
+        mrope_section = [base + remainder, base, base]
+        sub_config.rope_scaling = {"rope_type": "default", "type": "default", "mrope_section": mrope_section}
+
+    return config
+
+
 def update_config_for_eagle3(config):
     moduler_name = "optimum.exporters.openvino.model_patcher"
     spec = importlib.util.find_spec(moduler_name)
@@ -373,6 +406,20 @@ def main_export(
             patch_qwenvl_configs()
 
         model_type = config.model_type
+
+        if model_type == "qwen3_omni_moe":
+            if is_transformers_version("<", "5.0"):
+                raise ValueError(
+                    "Exporting Qwen3-Omni-MoE models requires transformers >= 5.0, "
+                    f"but found {_transformers_version}. Please upgrade transformers."
+                )
+            loading_kwargs["config"] = _ensure_qwen3_omni_rope_scaling(config)
+
+        if original_task == "auto" and model_type in {"phi4mm", "phi4_multimodal"}:
+            task = "image-text-to-text"
+        elif model_type == "qwen3_omni_moe":
+            task = "image-text-to-text"
+
         if model_type not in TasksManager._SUPPORTED_MODEL_TYPE:
             if custom_export_configs is None:
                 raise ValueError(
@@ -687,11 +734,14 @@ def _main_quantize(
             token=token,
         )
 
-    # NOTE: The Phi-4-multimodal-instruct model card contains a pipeline_tag set to automatic-speech-recognition,
-    # which is returned as the inferred task. As a result, we try to load the exported model using the
-    # OVModelForSpeechSeq2Seq class instead of the OVModelForVisualCausalLM class when the task is not specified
-    # explicitly. Because of this, we get an error.
-    if original_task == "auto" and library_name == "transformers":
+    # Some multimodal models must always be reloaded through OVModelForVisualCausalLM, but their task
+    # would otherwise route to the wrong OV class:
+    #   - Phi-4-multimodal-instruct: its model card pins pipeline_tag=automatic-speech-recognition, so the
+    #     *inferred* (auto) task is wrong and would pick OVModelForSpeechSeq2Seq.
+    #   - Qwen3-Omni-MoE: registers several tasks (text-to-audio, ASR, image-text-to-text); any of them,
+    #     whether inferred or passed explicitly, must still load through OVModelForVisualCausalLM.
+    # AutoConfig is cheap (cached) so we always read it to decide on the redirect.
+    if library_name == "transformers":
         config = AutoConfig.from_pretrained(
             model_name_or_path,
             subfolder=subfolder,
@@ -701,7 +751,10 @@ def _main_quantize(
             trust_remote_code=trust_remote_code,
         )
         model_type = config.model_type
-        if model_type in ["phi4mm", "phi4_multimodal"]:
+        # phi4mm is only misrouted by task inference, so keep its redirect limited to the auto case.
+        if original_task == "auto" and model_type in ["phi4mm", "phi4_multimodal"]:
+            task = "image-text-to-text"
+        elif model_type == "qwen3_omni_moe":
             task = "image-text-to-text"
 
     # Step 1. Obtain the correct OpenVINO model class
@@ -777,7 +830,12 @@ def maybe_convert_tokenizers(library_name: str, output: Path, model=None, prepro
                         processor_chat_template = processor.chat_template
             if tokenizer:
                 try:
-                    export_tokenizer(tokenizer, output, task=task, processor_chat_template=processor_chat_template)
+                    export_tokenizer(
+                        tokenizer,
+                        output,
+                        task=task,
+                        processor_chat_template=processor_chat_template,
+                    )
                 except Exception as exception:
                     logger.warning(
                         "Could not load tokenizer using specified model ID or path. OpenVINO tokenizer/detokenizer "

optimum/exporters/openvino/input_generators.py

@@ -1799,3 +1799,139 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int
             return self.random_int_tensor(shape=[1], min_value=1, max_value=10, framework=framework, dtype=float_dtype)
         else:
             raise ValueError(f"Unsupported input {input_name} for DummyKokoroInputGenerator")
+
+
+class DummyQwen3OmniMoeAudioInputGenerator(DummyInputGenerator):
+    SUPPORTED_INPUT_NAMES = ("padded_feature", "padded_mask_after_cnn", "aftercnn_lens", "cu_seqlens")
+
+    def __init__(
+        self,
+        task: str,
+        normalized_config: NormalizedVisionConfig,
+        batch_size: int = 3,
+        **kwargs,
+    ):
+        self.batch_size = batch_size
+        audio_config = normalized_config.config
+        self.num_mels = getattr(audio_config, "num_mel_bins", 128)
+        self.time_in = 200  # Mel spectrogram frames (~12.5s audio)
+        self.aftercnn_time = self.time_in // 8  # After 3 stride-2 conv layers
+        n_window = audio_config.n_window
+        n_window_infer = audio_config.n_window_infer
+        window_aftercnn = self.aftercnn_time * (n_window_infer // (n_window * 2))
+        num_full = self.aftercnn_time // window_aftercnn
+        single_batch_chunks = [window_aftercnn] * num_full
+        remainder = self.aftercnn_time % window_aftercnn
+        if remainder != 0:
+            single_batch_chunks.append(remainder)
+        cu_chunk_lens = single_batch_chunks * self.batch_size
+        cumsum = [0]
+        for length in cu_chunk_lens:
+            cumsum.append(cumsum[-1] + length)
+        self._cu_seqlens = cumsum
+
+    def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
+        if input_name == "padded_feature":
+            return self.random_float_tensor([self.batch_size, self.num_mels, self.time_in], framework=framework)
+        if input_name == "padded_mask_after_cnn":
+            return self.constant_tensor(
+                [self.batch_size, self.aftercnn_time], framework=framework, value=1, dtype=DTYPE_MAPPER.pt("bool")
+            )
+        if input_name == "aftercnn_lens":
+            return self.constant_tensor(
+                [self.batch_size], framework=framework, value=self.aftercnn_time, dtype=DTYPE_MAPPER.pt(int_dtype)
+            )
+        if input_name == "cu_seqlens":
+            import torch
+
+            return torch.tensor(self._cu_seqlens, dtype=torch.int32)
+        return super().generate(input_name, framework, int_dtype, float_dtype)
+
+
+class DummyQwen3OmniMoeLMInputGenerator(DummyQwen3VLLMInputGenerator):
+    def generate(
+        self,
+        input_name: str,
+        framework: str = "pt",
+        int_dtype: str = "int64",
+        float_dtype: str = "fp32",
+        bool_dtype: str = "bool",
+    ):
+        if input_name == "position_ids":
+            base = DummyTextInputGenerator.generate(self, input_name, framework, int_dtype, float_dtype)
+            return base.unsqueeze(0).expand(4, -1, -1)
+        return super().generate(input_name, framework, int_dtype, float_dtype, bool_dtype)
+
+
+class DummyQwen3OmniMoeCode2WavInputGenerator(DummyInputGenerator):
+    SUPPORTED_INPUT_NAMES = ("codes",)
+
+    def __init__(self, task: str, normalized_config: NormalizedVisionConfig, batch_size: int = 1, **kwargs):
+        self.batch_size = batch_size
+        code2wav_config = normalized_config.config
+        self.num_quantizers = getattr(code2wav_config, "num_quantizers", 16)
+        self.codebook_size = getattr(code2wav_config, "codebook_size", 2048)
+        self.seq_len = 10
+
+    def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
+        if input_name == "codes":
+            return self.random_int_tensor(
+                [self.batch_size, self.num_quantizers, self.seq_len],
+                min_value=0,
+                max_value=self.codebook_size,
+                framework=framework,
+            )
+        return super().generate(input_name, framework, int_dtype, float_dtype)
+
+
+class DummyQwen3OmniMoeProjectionInputGenerator(DummyInputGenerator):
+    SUPPORTED_INPUT_NAMES = ("hidden_state",)
+
+    def __init__(
+        self,
+        task: str,
+        normalized_config: NormalizedVisionConfig,
+        batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"],
+        sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"],
+        **kwargs,
+    ):
+        self.batch_size = batch_size
+        self.sequence_length = sequence_length
+        config = normalized_config.config
+        # The projection (Qwen3OmniMoeTalkerResizeMLP.linear_fc1) consumes thinker hidden states whose width is
+        # talker_config.thinker_hidden_size; fall back to text_config.hidden_size only if it's absent.
+        text_config = getattr(config, "text_config", config)
+        self.hidden_size = getattr(config, "thinker_hidden_size", None) or text_config.hidden_size
+
+    def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
+        if input_name == "hidden_state":
+            return self.random_float_tensor(
+                [self.batch_size, self.sequence_length, self.hidden_size], framework=framework
+            )
+        return super().generate(input_name, framework, int_dtype, float_dtype)
+
+
+class DummyQwen3OmniMoeVisionInputGenerator(DummyQwen3VLVisionEmbedInputGenerator):
+    SUPPORTED_INPUT_NAMES = ("hidden_states", "pos_embeds", "attention_mask", "rotary_pos_emb", "input")
+
+    def __init__(self, task, normalized_config, batch_size=1, **kwargs):
+        super().__init__(task, normalized_config, batch_size=batch_size, **kwargs)
+        self.patch_channels = (
+            normalized_config.config.in_channels
+            * normalized_config.config.temporal_patch_size
+            * normalized_config.config.patch_size
+            * normalized_config.config.patch_size
+        )
+
+    def generate(self, input_name, framework="pt", int_dtype="int64", float_dtype="fp32"):
+        grid_h, grid_w = self.height // self.patch_size, self.width // self.patch_size
+        seq_len = self.batch_size * grid_h * grid_w
+
+        if input_name == "hidden_states":
+            # Raw patch data (patch_channels), not embeddings: the vision model embeds patches internally.
+            return self.random_float_tensor([seq_len, self.patch_channels], framework=framework, dtype=float_dtype)
+
+        if input_name == "pos_embeds":
+            return self.random_float_tensor([seq_len, self.embed_dim], framework=framework, dtype=float_dtype)
+
+        return super().generate(input_name, framework, int_dtype, float_dtype)