Add VibeVoice streaming ASR export

.github/workflows/main.yml

@@ -189,6 +189,58 @@ jobs:
           flags: synthetic-parity
           report-type: test_results
 
+  vibevoice-asr-source-parity:
+    name: VibeVoice ASR L3 Source Parity
+    runs-on: ubuntu-latest
+    needs: [detect-affected]
+    if: >-
+      always() && !cancelled() &&
+      (needs.detect-affected.result == 'failure' ||
+       needs.detect-affected.outputs.run_all == 'true' ||
+       contains(needs.detect-affected.outputs.affected, 'VibeVoiceForASRStreamingTraining'))
+    steps:
+      - uses: actions/checkout@v7
+      - name: Setup Python
+        uses: actions/setup-python@v7
+        with:
+          python-version: "3.12"
+      - name: Cache pip packages
+        uses: actions/cache@v6
+        with:
+          path: ~/.cache/pip
+          key: pip-vibevoice-asr-${{ hashFiles('pyproject.toml', 'requirements/ci/requirements.txt', 'requirements/ci/vibevoice-asr.txt') }}
+          restore-keys: |
+            pip-vibevoice-asr-
+      - name: Install PyTorch CPU
+        run: pip install torch --index-url https://download.pytorch.org/whl/cpu
+      - name: Install pinned ASR reference
+        run: |
+          pip install -r requirements/ci/requirements.txt
+          pip install onnxruntime
+          pip install -e '.[testing]'
+          pip install --force-reinstall --no-deps -r requirements/ci/vibevoice-asr.txt
+      - name: Run pinned VibeVoice ASR source parity
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: |
+          pytest src/mobius/models/vibevoice_asr_test.py -m integration -v --tb=short \
+            -p no:xdist --timeout=300 --junitxml=junit.xml \
+            --cov=src --cov-report=xml --cov-branch
+        timeout-minutes: 20
+      - name: Upload coverage to Codecov
+        if: always()
+        uses: codecov/codecov-action@v7
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          flags: vibevoice-asr-source-parity
+      - name: Upload test results to Codecov
+        if: ${{ !cancelled() }}
+        uses: codecov/codecov-action@v7
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          flags: vibevoice-asr-source-parity
+          report-type: test_results
+
   golden-comparison:
     name: L4 Golden Comparison
     needs: [detect-affected]

docs/model-catalog.md

@@ -189,6 +189,7 @@ Also registered with `T5ForConditionalGeneration` (task: `seq2seq`):
 | `moonshine_streaming` | `MoonshineStreamingForConditionalGeneration` | `speech-to-text` | `moonshine-ai/moonshine-streaming-tiny` |
 | `qwen3_asr` | `Qwen3ASRForConditionalGeneration` | `speech-language` | — |
 | `qwen3_forced_aligner` | `Qwen3ASRForConditionalGeneration` | `speech-language` | — |
+| `VibeVoiceForASRStreamingTraining` | `VibeVoiceASRForConditionalGeneration` | `vibevoice-asr-streaming` | `microsoft/VibeVoice-ASR-Streaming-1.5B`, `microsoft/VibeVoice-ASR-Streaming-7B` |
 
 ### Text-to-Speech
 

requirements/ci/vibevoice-asr.txt

@@ -0,0 +1,6 @@
+# Exact executable source and dependency ABI for VibeVoice streaming ASR parity tests.
+# CI force-installs this file with --no-deps after Mobius's public Transformers>=5 extra.
+transformers==4.51.3
+huggingface-hub==0.31.4
+tokenizers==0.21.4
+vibevoice @ git+https://github.com/microsoft/VibeVoice.git@505653d3873b065a488aea551c6ee3dc51d3062f

src/mobius/_configs/__init__.py

@@ -108,6 +108,7 @@
 )
 from mobius._configs._world_model import WorldModelConfig
 from mobius._configs.vibevoice import (
+    VibeVoiceASRConfig,
     VibeVoiceConfig,
     VibeVoiceDiffusionConfig,
     VibeVoiceTokenizerConfig,
@@ -175,6 +176,7 @@
     "TTSConfig",
     "VisionConfig",
     "VisionLanguageConfig",
+    "VibeVoiceASRConfig",
     "VibeVoiceConfig",
     "VibeVoiceDiffusionConfig",
     "VibeVoiceTokenizerConfig",

src/mobius/_configs/vibevoice.py

@@ -1,13 +1,17 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-"""Configuration for the Transformers-native VibeVoice text-to-speech model."""
+"""Configuration for the Transformers-native VibeVoice TTS and streaming ASR models."""
 
 from __future__ import annotations
 
 import dataclasses
 
-from mobius._configs._base import ArchitectureConfig, _as_attribute_config
+from mobius._configs._base import (
+    ArchitectureConfig,
+    _as_attribute_config,
+    _resolve_dtype_value,
+)
 
 
 @dataclasses.dataclass
@@ -27,6 +31,7 @@ class VibeVoiceTokenizerConfig:
     rms_norm_eps: float = 1e-5
     layer_scale_init_value: float = 1e-6
     vae_std: float = 0.625
+    std_dist_type: str = "gaussian"
 
     @property
     def hop_length(self) -> int:
@@ -68,6 +73,61 @@ def _tokenizer_config(config, *, default_hidden_size: int) -> VibeVoiceTokenizer
     )
 
 
+def _asr_tokenizer_config(config, *, default_hidden_size: int) -> VibeVoiceTokenizerConfig:
+    """Extract the ASR encoder spelling into the shared causal-tokenizer geometry.
+
+    The public ASR implementation reverses ``encoder_ratios`` before constructing
+    its downsampling layers. Keep that conversion at the configuration boundary so
+    the reusable encoder receives its actual execution order.
+    """
+    config = _as_attribute_config(config)
+    encoder_ratios = list(getattr(config, "encoder_ratios", [8, 5, 5, 4, 2, 2]))
+    required_values = {
+        "channels": (getattr(config, "channels", 1), 1),
+        "causal": (getattr(config, "causal", True), True),
+        "mixer_layer": (getattr(config, "mixer_layer", "depthwise_conv"), "depthwise_conv"),
+        "conv_norm": (getattr(config, "conv_norm", "none"), "none"),
+        "pad_mode": (getattr(config, "pad_mode", "constant"), "constant"),
+        "disable_last_norm": (getattr(config, "disable_last_norm", True), True),
+        "layernorm": (getattr(config, "layernorm", "RMSNorm"), "RMSNorm"),
+        "layernorm_elementwise_affine": (
+            getattr(config, "layernorm_elementwise_affine", True),
+            True,
+        ),
+        "conv_bias": (getattr(config, "conv_bias", True), True),
+    }
+    unsupported = [
+        f"{name}={actual!r}"
+        for name, (actual, expected) in required_values.items()
+        if actual != expected
+    ]
+    if unsupported:
+        raise ValueError(
+            "Unsupported VibeVoice ASR tokenizer configuration: "
+            + ", ".join(unsupported)
+            + "."
+        )
+    if not encoder_ratios or any(int(ratio) <= 0 for ratio in encoder_ratios):
+        raise ValueError("VibeVoice ASR encoder_ratios must contain positive integers.")
+    return VibeVoiceTokenizerConfig(
+        channels=int(getattr(config, "channels", 1)),
+        hidden_size=int(getattr(config, "vae_dim", default_hidden_size)),
+        kernel_size=int(getattr(config, "kernel_size", 7)),
+        num_filters=int(getattr(config, "encoder_n_filters", 32)),
+        downsampling_ratios=list(reversed(encoder_ratios)),
+        depths=[
+            int(depth)
+            for depth in str(getattr(config, "encoder_depths", "3-3-3-3-3-3-8")).split("-")
+        ],
+        ffn_expansion=int(getattr(config, "ffn_expansion", 4)),
+        hidden_act=str(getattr(config, "hidden_act", "gelu")),
+        rms_norm_eps=float(getattr(config, "layernorm_eps", 1e-5)),
+        layer_scale_init_value=float(getattr(config, "layer_scale_init_value", 1e-6)),
+        vae_std=float(getattr(config, "fix_std", 0.625)),
+        std_dist_type=str(getattr(config, "std_dist_type", "gaussian")),
+    )
+
+
 @dataclasses.dataclass
 class VibeVoiceConfig(ArchitectureConfig):
     """Mobius configuration for VibeVoice's LM, tokenizers, and diffusion head."""
@@ -149,3 +209,85 @@ def from_transformers(
             pad_token_id=getattr(parent, "pad_token_id", 151643),
             sampling_rate=24_000,
         )
+
+
+@dataclasses.dataclass
+class VibeVoiceASRConfig(ArchitectureConfig):
+    """Mobius configuration for VibeVoice's dual causal tokenizers and Qwen2 ASR decoder."""
+
+    acoustic_tokenizer: VibeVoiceTokenizerConfig = dataclasses.field(
+        default_factory=VibeVoiceTokenizerConfig
+    )
+    semantic_tokenizer: VibeVoiceTokenizerConfig = dataclasses.field(
+        default_factory=lambda: VibeVoiceTokenizerConfig(hidden_size=128, std_dist_type="none")
+    )
+    speech_start_token_id: int = 151646
+    speech_end_token_id: int = 151647
+    speech_placeholder_token_id: int = 151648
+    text_chunk_end_token_id: int = 151665
+    compression_ratio: int = 3200
+    sampling_rate: int = 24_000
+    chunk_frames: int = 22
+    lookahead_frames: int = 4
+
+    @classmethod
+    def from_transformers(
+        cls,
+        config,
+        parent_config=None,
+        *,
+        allow_block_fp8_dense_fallback: bool = False,
+    ) -> VibeVoiceASRConfig:
+        """Extract the public ASR composite while preserving its Qwen2 decoder config."""
+        parent = _as_attribute_config(parent_config or config)
+        result = super().from_transformers(
+            config,
+            parent_config=parent,
+            allow_block_fp8_dense_fallback=allow_block_fp8_dense_fallback,
+        )
+        acoustic = _asr_tokenizer_config(
+            getattr(parent, "acoustic_tokenizer_config", None),
+            default_hidden_size=64,
+        )
+        semantic = _asr_tokenizer_config(
+            getattr(parent, "semantic_tokenizer_config", None),
+            default_hidden_size=128,
+        )
+        if acoustic.std_dist_type != "gaussian" or semantic.std_dist_type != "none":
+            raise ValueError(
+                "VibeVoice streaming ASR requires a gaussian acoustic tokenizer and "
+                "a deterministic semantic tokenizer."
+            )
+        if acoustic.hop_length != int(
+            getattr(parent, "speech_tokenizer_compression_ratio", 3200)
+        ):
+            raise ValueError(
+                "VibeVoice streaming ASR acoustic tokenizer ratios must match the "
+                "configured speech tokenizer compression ratio."
+            )
+        if semantic.hop_length != acoustic.hop_length:
+            raise ValueError(
+                "VibeVoice ASR acoustic and semantic tokenizer hop lengths must match."
+            )
+        # The composite ASR config controls the tokenizer and connector precision,
+        # while decoder_config independently records Qwen's checkpoint storage dtype.
+        pipeline_dtype = _resolve_dtype_value(getattr(parent, "dtype", None))
+        if pipeline_dtype is None:
+            pipeline_dtype = _resolve_dtype_value(getattr(parent, "torch_dtype", None))
+        return dataclasses.replace(
+            result,
+            model_type="vibevoice",
+            acoustic_tokenizer=acoustic,
+            semantic_tokenizer=semantic,
+            speech_start_token_id=int(getattr(parent, "object_ref_start_token_id", 151646)),
+            speech_end_token_id=int(getattr(parent, "object_ref_end_token_id", 151647)),
+            speech_placeholder_token_id=int(getattr(parent, "speech_pad_token_id", 151648)),
+            text_chunk_end_token_id=int(getattr(parent, "text_chunk_end_token_id", 151665)),
+            compression_ratio=acoustic.hop_length,
+            sampling_rate=int(getattr(parent, "target_sample_rate", 24_000)),
+            chunk_frames=int(getattr(parent, "chunk_frames", 22)),
+            lookahead_frames=int(getattr(parent, "lookahead_frames", 4)),
+            eos_token_id=getattr(parent, "eos_token_id", 151643),
+            pad_token_id=getattr(parent, "pad_token_id", 151655),
+            dtype=pipeline_dtype or result.dtype,
+        )

src/mobius/_registry.py

@@ -47,6 +47,7 @@
     Plamo2Config,
     Qwen4ExpConfig,
     SenseNovaU1Config,
+    VibeVoiceASRConfig,
     VibeVoiceConfig,
     WhisperConfig,
     XverseConfig,
@@ -159,6 +160,7 @@
     SmallThinkerGGUFCausalLMModel,
     SmolLM3CausalLMModel,
     SortformerDiarizationModel,
+    VibeVoiceASRForConditionalGeneration,
     VibeVoiceForConditionalGeneration,
     WhisperForConditionalGeneration,
     XverseCausalLMModel,
@@ -905,6 +907,15 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
         test_revision=VIBEVOICE_REVISION,
         family="vibevoice",
     ),
+    "VibeVoiceForASRStreamingTraining": ModelRegistration(
+        VibeVoiceASRForConditionalGeneration,
+        task="vibevoice-asr-streaming",
+        config_class=VibeVoiceASRConfig,
+        test_model_id="microsoft/VibeVoice-ASR-Streaming-1.5B",
+        test_revision="4262d23d8a539a6530cf64fbd0b1751ef9a30853",
+        family="vibevoice",
+        variant="streaming-asr",
+    ),
     "whisper": ModelRegistration(
         WhisperForConditionalGeneration,
         task="speech-to-text",

src/mobius/integrations/onnx_genai/auto_export.py

@@ -769,6 +769,23 @@ def _looks_like_vibevoice_tts(pkg: Any) -> bool:
     } <= names
 
 
+def _looks_like_vibevoice_asr(pkg: Any) -> bool:
+    """Detect the stateful three-stage VibeVoice streaming ASR topology."""
+    try:
+        if set(pkg.keys()) != {"audio_encoder", "embedding", "decoder"}:
+            return False
+        audio_inputs = {value.name for value in pkg["audio_encoder"].graph.inputs}
+    except AttributeError:
+        return False
+    return {
+        "speech_tensors",
+        "speech_masks",
+        "is_final_chunk",
+        "past_acoustic_conv.0",
+        "past_semantic_conv.0",
+    } <= audio_inputs
+
+
 def _looks_like_speculative(pkg: Any) -> bool:
     try:
         return {"proposer", "verifier"} <= set(pkg.keys())
@@ -970,6 +987,47 @@ def write_onnx_genai_config(
         )
         return _write_advisory_component_contract(pkg, output_dir, warning=warning)
     component_names = set(pkg)
+    if _looks_like_vibevoice_asr(pkg):
+        if kv_native_dtype is not None:
+            raise ValueError(
+                "workflow VibeVoice ASR export derives KV and convolution state dtypes "
+                "from ONNX ports; kv_native_dtype overrides are unsupported"
+            )
+        artifacts = _write_text_runtime_assets(output_dir, source, revision=revision)
+        artifacts.update(
+            _copy_runtime_assets(
+                output_dir,
+                source,
+                (
+                    "processor_config.json",
+                    "preprocessor_config.json",
+                    "generation_config.json",
+                ),
+                revision=revision,
+            )
+        )
+        audio_processor_path = _write_hf_audio_processor(
+            output_dir,
+            source,
+            revision=revision,
+        )
+        if audio_processor_path is not None:
+            artifacts["audio_processor"] = audio_processor_path
+        artifacts.update(
+            _write_advisory_component_contract(
+                pkg,
+                output_dir,
+                warning=(
+                    "The tested onnx-genai runtime cannot orchestrate VibeVoice streaming "
+                    "ASR's dual causal convolution states, flattened speech-placeholder "
+                    "replacement with arbitrary left-padded attention masks, forced "
+                    "<|text_chunk_end|> control tokens, or host-side hotword and speaker "
+                    "JSON handling. Exact graph and pinned processor contracts are exported "
+                    "without claiming downstream runtime support."
+                ),
+            )
+        )
+        return artifacts
     if component_names in ({"audio_encoder"}, {"speaker_encoder"}) and getattr(
         pkg, "gguf_projector_type", None
     ):