diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index a04d16ed4..60ca4b54b 100644 --- a/.github/workflows/main.yml +++ b/.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] diff --git a/docs/model-catalog.md b/docs/model-catalog.md index 94c983b54..c1afbb134 100644 --- a/docs/model-catalog.md +++ b/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 diff --git a/requirements/ci/vibevoice-asr.txt b/requirements/ci/vibevoice-asr.txt new file mode 100644 index 000000000..e05b9cbcb --- /dev/null +++ b/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 diff --git a/src/mobius/_configs/__init__.py b/src/mobius/_configs/__init__.py index bd51d6c8d..e6d54a496 100644 --- a/src/mobius/_configs/__init__.py +++ b/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", diff --git a/src/mobius/_configs/vibevoice.py b/src/mobius/_configs/vibevoice.py index 04f1dfb57..b510bed68 100644 --- a/src/mobius/_configs/vibevoice.py +++ b/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, + ) diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index f3d186f0a..35cbd7fb2 100644 --- a/src/mobius/_registry.py +++ b/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", diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 686e8a6b3..4c0fecb55 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/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 ): diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 3cf67a13a..a9b893b85 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -142,6 +142,77 @@ def audio_processor(output_dir, source, *, revision=None): } <= set(artifacts) +@pytest.mark.parametrize( + ("model_id", "revision"), + ( + ("microsoft/VibeVoice-ASR-Streaming-1.5B", "4262d23d8a539a6530cf64fbd0b1751ef9a30853"), + ("microsoft/VibeVoice-ASR-Streaming-7B", "60d858b518b4e19d404af3737f848fc185b30177"), + ), +) +def test_vibevoice_asr_writes_pinned_advisory_contract( + monkeypatch, tmp_path, model_id, revision +): + from mobius.integrations.onnx_genai import auto_export + from mobius.models.vibevoice import ( + VibeVoiceASRForConditionalGeneration, + ) + from mobius.models.vibevoice_asr_test import _config + from mobius.tasks import VibeVoiceASRStreamingTask + + config = _config() + package = VibeVoiceASRStreamingTask().build( + VibeVoiceASRForConditionalGeneration(config), config + ) + calls: list[tuple[str, str | None]] = [] + warnings: list[str] = [] + + def text_assets(output_dir, source, *, revision=None): + calls.append(("text", revision)) + return {"tokenizer": str(Path(output_dir) / "tokenizer.json")} + + def runtime_assets(output_dir, source, names, *, revision=None): + calls.append(("runtime", revision)) + assert names == ( + "processor_config.json", + "preprocessor_config.json", + "generation_config.json", + ) + return {"processor_config": str(Path(output_dir) / "processor_config.json")} + + def audio_processor(output_dir, source, *, revision=None): + calls.append(("audio", revision)) + return str(Path(output_dir) / "audio_processor.json") + + def advisory(*args, warning, **kwargs): + warnings.append(warning) + return {"inference_metadata": str(tmp_path / "inference_metadata.yaml")} + + monkeypatch.setattr(auto_export, "_write_text_runtime_assets", text_assets) + monkeypatch.setattr(auto_export, "_copy_runtime_assets", runtime_assets) + monkeypatch.setattr(auto_export, "_write_hf_audio_processor", audio_processor) + monkeypatch.setattr(auto_export, "_write_advisory_component_contract", advisory) + + artifacts = write_onnx_genai_config( + package, + str(tmp_path), + source=model_id, + revision=revision, + ) + + assert calls == [ + ("text", revision), + ("runtime", revision), + ("audio", revision), + ] + assert "arbitrary left-padded attention masks" in warnings[0] + assert { + "tokenizer", + "processor_config", + "audio_processor", + "inference_metadata", + } <= set(artifacts) + + def _video_diffusion_package() -> ModelPackage: latent = ["batch", "frames", 4, "height", "width"] transformer = _model( diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index d45495a5b..d96af0522 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -127,15 +127,21 @@ def _load_transformers_config( ) -> tuple[object | None, bool]: """Load a Transformers config and report whether raw JSON was required.""" import transformers + from huggingface_hub import errors as hub_errors from mobius.integrations.transformers._config_resolver import _try_load_config_json + strict_validation_error = getattr( + hub_errors, + "StrictDataclassClassValidationError", + ValueError, + ) try: kwargs = {"trust_remote_code": trust_remote_code} if revision is not None: kwargs["revision"] = revision return transformers.AutoConfig.from_pretrained(model_id, **kwargs), False - except (ValueError, KeyError, OSError): + except (strict_validation_error, ValueError, KeyError, OSError): return _try_load_config_json(model_id, revision=revision), True @@ -148,7 +154,14 @@ def _select_primary_config(hf_config): parent_config = hf_config model_type = hf_config.model_type - if hasattr(hf_config, "talker_config"): + if model_type == "vibevoice" and hasattr(hf_config, "decoder_config"): + # TTS uses ``text_config`` while the distinct streaming ASR checkpoint + # stores its Qwen2 config under ``decoder_config``. + decoder = hf_config.decoder_config + if isinstance(decoder, dict): + decoder = _dict_to_pretrained_config(decoder) + hf_config = decoder + elif hasattr(hf_config, "talker_config"): hf_config = hf_config.talker_config elif hasattr(hf_config, "thinker_config"): thinker = hf_config.thinker_config @@ -189,6 +202,24 @@ def _resolve_module_class( ) -> tuple[type[nn.Module], str | ModelTask | None, str]: """Resolve architecture aliases and structural fallback registrations.""" architectures = getattr(parent_config, "architectures", None) or [] + if model_type == "vibevoice": + if len(architectures) != 1: + raise ValueError( + "VibeVoice checkpoints must declare exactly one recognized architecture; " + "Mobius will not guess between the incompatible TTS and streaming ASR " + "pipelines." + ) + architecture = architectures[0] + if architecture == "VibeVoiceForConditionalGeneration": + pass + elif architecture == "VibeVoiceForASRStreamingTraining": + model_type = architecture + else: + raise ValueError( + f"Unsupported VibeVoice architecture {architecture!r}; supported architectures " + "are 'VibeVoiceForConditionalGeneration' and " + "'VibeVoiceForASRStreamingTraining'." + ) if allow_parent_architecture_override and architectures and architectures[0] in registry: architecture_key = architectures[0] model_type_class = registry.get(model_type) if model_type in registry else None @@ -300,15 +331,17 @@ def build_transformers_model( _config_from_hf, _default_task_for_model, ) + from mobius.models.vibevoice import VIBEVOICE_ASR_MODEL_REVISIONS, VIBEVOICE_REVISION detection_revision = revision - if model_id == "vibevoice/VibeVoice-1.5B-hf" and detection_revision is None: - from mobius.models.vibevoice import VIBEVOICE_REVISION - + if ( + model_id in {"vibevoice/VibeVoice-1.5B-hf", *VIBEVOICE_ASR_MODEL_REVISIONS} + and detection_revision is None + ): # The native conversion is the executable source of truth. Pin the # first config probe and every later processor/weight call together. - revision = VIBEVOICE_REVISION - detection_revision = VIBEVOICE_REVISION + revision = VIBEVOICE_ASR_MODEL_REVISIONS.get(model_id, VIBEVOICE_REVISION) + detection_revision = revision if model_id == "nvidia/RE-USE" and detection_revision is None: # Pin the very first AutoConfig/raw-JSON probe, not only the later # bespoke loader. Otherwise mutable Hub main could change dispatch @@ -525,7 +558,10 @@ def build_transformers_model( ) for name, model in package.items(): model.graph.name = f"{model_id}/{name}" - if model_type in _QWEN4_MODEL_TYPES | {"vibevoice"}: + if model_type in _QWEN4_MODEL_TYPES | { + "vibevoice", + "VibeVoiceForASRStreamingTraining", + }: model.metadata_props["mobius.source_revision"] = revision or "unpinned" if load_weights: diff --git a/src/mobius/integrations/transformers/_builder_test.py b/src/mobius/integrations/transformers/_builder_test.py index 85a72542d..ae3980b4b 100644 --- a/src/mobius/integrations/transformers/_builder_test.py +++ b/src/mobius/integrations/transformers/_builder_test.py @@ -272,6 +272,46 @@ def stop_after_config(model_id, **kwargs): ] +@pytest.mark.parametrize( + ("model_id", "revision"), + ( + ("microsoft/VibeVoice-ASR-Streaming-1.5B", "4262d23d8a539a6530cf64fbd0b1751ef9a30853"), + ("microsoft/VibeVoice-ASR-Streaming-7B", "60d858b518b4e19d404af3737f848fc185b30177"), + ), +) +def test_vibevoice_asr_none_revision_pins_first_config_probe( + monkeypatch, model_id, revision +) -> None: + """Both official ASR variants must pin their initial config probe.""" + calls = [] + + def stop_after_config(called_model_id, **kwargs): + calls.append((called_model_id, kwargs)) + raise RuntimeError("stop after revision assertion") + + monkeypatch.setattr( + transformers_builder, + "_load_transformers_config", + stop_after_config, + ) + + with pytest.raises(RuntimeError, match="stop after revision assertion"): + transformers_builder.build_transformers_model( + model_id, + load_weights=False, + ) + + assert calls == [ + ( + model_id, + { + "revision": revision, + "trust_remote_code": False, + }, + ) + ] + + @pytest.mark.parametrize("revision", [None, "feature/revision"]) def test_transformers_config_forwards_only_explicit_revision(monkeypatch, revision) -> None: import transformers diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 6cbef77d2..ee0879307 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -222,6 +222,7 @@ "PLMCausalLMModel", "TalkieForCausalLM", "Qwen4ExpForConditionalGeneration", + "VibeVoiceASRForConditionalGeneration", ] from mobius.models.adapters import IPAdapterModel, T2IAdapterModel @@ -453,7 +454,10 @@ remap_diffusers_unet_lora, ) from mobius.models.vae import AutoencoderKLModel -from mobius.models.vibevoice import VibeVoiceForConditionalGeneration +from mobius.models.vibevoice import ( + VibeVoiceASRForConditionalGeneration, + VibeVoiceForConditionalGeneration, +) from mobius.models.video_vae import VideoAutoencoderModel from mobius.models.vit import ViTModel from mobius.models.wav2vec2 import Wav2Vec2Model diff --git a/src/mobius/models/vibevoice.py b/src/mobius/models/vibevoice.py index 3841aaab4..485abffd7 100644 --- a/src/mobius/models/vibevoice.py +++ b/src/mobius/models/vibevoice.py @@ -19,6 +19,7 @@ from onnxscript import OpBuilder, nn from mobius._configs import ( + VibeVoiceASRConfig, VibeVoiceConfig, VibeVoiceDiffusionConfig, VibeVoiceTokenizerConfig, @@ -40,6 +41,11 @@ VIBEVOICE_MODEL_ID = "vibevoice/VibeVoice-1.5B-hf" VIBEVOICE_REVISION = "edc39f80f5cae656da37baf8faa8f5502bf7081f" VIBEVOICE_MICROSOFT_PROVENANCE_REVISION = "c00898d257e6b46004e3e2866a47534085fb685a" +VIBEVOICE_ASR_MODEL_REVISIONS = { + "microsoft/VibeVoice-ASR-Streaming-1.5B": "4262d23d8a539a6530cf64fbd0b1751ef9a30853", + "microsoft/VibeVoice-ASR-Streaming-7B": "60d858b518b4e19d404af3737f848fc185b30177", +} +VIBEVOICE_ASR_SOURCE_REVISION = "505653d3873b065a488aea551c6ee3dc51d3062f" class _CacheAllocator: @@ -71,13 +77,23 @@ def prepend( ) -> ir.Value: past = self._past[index] padded = op.Concat(past, hidden_states, axis=2) + self.update(op, padded, index=index, left_pad=left_pad) + return padded + + def update( + self, + op: OpBuilder, + hidden_states: ir.Value, + *, + index: int, + left_pad: int, + ) -> None: self._present[index] = op.Slice( - padded, + hidden_states, op.Constant(value_ints=[-left_pad]), op.Constant(value_ints=[2**63 - 1]), op.Constant(value_ints=[2]), ) - return padded def outputs(self) -> list[ir.Value]: if any(value is None for value in self._present): @@ -160,6 +176,7 @@ def __init__( self._left_pad = (kernel_size - 1) * dilation - (stride - 1) if self._left_pad < 0: raise ValueError("VibeVoice causal convolution padding must be non-negative") + self._stride = stride self._cache_index = allocator.add(in_channels, self._left_pad) def forward( @@ -167,6 +184,7 @@ def forward( op: OpBuilder, hidden_states: ir.Value, state: _ConvState | None = None, + is_final_chunk: ir.Value | None = None, ): if state is None: hidden_states = op.Pad( @@ -180,6 +198,36 @@ def forward( index=self._cache_index, left_pad=self._left_pad, ) + if is_final_chunk is not None: + # HF pads causal chunks to the next stride boundary only on the final + # call. For its kernel=2*stride convolutions this is ``(-length) % stride``. + input_length = op.Shape(hidden_states, start=2, end=3) + extra_padding = op.Mod( + op.Sub( + op.Constant(value_int=self._stride), + op.Mod(input_length, op.Constant(value_int=self._stride)), + ), + op.Constant(value_int=self._stride), + ) + final_padding = op.Mul( + extra_padding, + op.Cast(is_final_chunk, to=ir.DataType.INT64), + ) + hidden_states = op.Pad( + hidden_states, + op.Concat( + op.Constant(value_ints=[0, 0, 0, 0, 0]), + final_padding, + axis=0, + ), + ) + if state is not None: + state.update( + op, + hidden_states, + index=self._cache_index, + left_pad=self._left_pad, + ) return self.conv(op, hidden_states) @@ -206,6 +254,7 @@ def forward( op: OpBuilder, hidden_states: ir.Value, state: _ConvState | None = None, + is_final_chunk: ir.Value | None = None, ): input_length = op.Shape(hidden_states, start=2, end=3) if state is not None: @@ -278,12 +327,13 @@ def forward( op: OpBuilder, hidden_states: ir.Value, state: _ConvState | None = None, + is_final_chunk: ir.Value | None = None, ): residual = hidden_states mixed = op.Transpose(hidden_states, perm=[0, 2, 1]) mixed = self.norm(op, mixed) mixed = op.Transpose(mixed, perm=[0, 2, 1]) - mixed = self.mixer(op, mixed, state) + mixed = self.mixer(op, mixed, state, is_final_chunk) mixed = op.Mul(mixed, op.Unsqueeze(self.gamma, [-1])) hidden_states = op.Add(residual, mixed) # (batch, channels, frames) @@ -312,10 +362,16 @@ def __init__(self, config: VibeVoiceTokenizerConfig, allocator: _CacheAllocator) ] ) - def forward(self, op: OpBuilder, hidden_states: ir.Value, state: _ConvState | None): - hidden_states = self.conv(op, hidden_states, state) + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + state: _ConvState | None, + is_final_chunk: ir.Value | None = None, + ): + hidden_states = self.conv(op, hidden_states, state, is_final_chunk) for block in self.stage: - hidden_states = block(op, hidden_states, state) + hidden_states = block(op, hidden_states, state, is_final_chunk) return hidden_states @@ -342,10 +398,16 @@ def __init__( ] ) - def forward(self, op: OpBuilder, hidden_states: ir.Value, state: _ConvState | None): - hidden_states = self.conv(op, hidden_states, state) + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + state: _ConvState | None, + is_final_chunk: ir.Value | None = None, + ): + hidden_states = self.conv(op, hidden_states, state, is_final_chunk) for block in self.stage: - hidden_states = block(op, hidden_states, state) + hidden_states = block(op, hidden_states, state, is_final_chunk) return hidden_states @@ -375,12 +437,13 @@ def forward( op: OpBuilder, input_values: ir.Value, past_conv_states: Sequence[ir.Value] | None = None, + is_final_chunk: ir.Value | None = None, ): state = _ConvState(past_conv_states) if past_conv_states is not None else None - hidden_states = self.stem(op, input_values, state) + hidden_states = self.stem(op, input_values, state, is_final_chunk) for layer in self.conv_layers: - hidden_states = layer(op, hidden_states, state) - hidden_states = self.head(op, hidden_states, state) + hidden_states = layer(op, hidden_states, state, is_final_chunk) + hidden_states = self.head(op, hidden_states, state, is_final_chunk) latents = op.Transpose(hidden_states, perm=[0, 2, 1]) # (batch, frames, latent) return latents, state.outputs() if state is not None else [] @@ -530,19 +593,56 @@ def forward( return op.GatherND(latents, valid_indices) # (valid_audio_frames, latent_size) -class VibeVoiceMultiModalProjector(nn.Module): - """Linear -> RMSNorm -> Linear connector used by both continuous tokenizers.""" +class _VibeVoiceConnector(nn.Module): + """Linear -> RMSNorm -> Linear connector with checkpoint-specific field names.""" - def __init__(self, input_dim: int, output_dim: int): + def __init__( + self, + input_dim: int, + output_dim: int, + *, + first_name: str, + norm_name: str, + second_name: str, + ): super().__init__() - self.linear_1 = Linear(input_dim, output_dim) - self.act = RMSNorm(output_dim, eps=1e-6) - self.linear_2 = Linear(output_dim, output_dim) + setattr(self, first_name, Linear(input_dim, output_dim)) + setattr(self, norm_name, RMSNorm(output_dim, eps=1e-6)) + setattr(self, second_name, Linear(output_dim, output_dim)) + self._first_name = first_name + self._norm_name = norm_name + self._second_name = second_name def forward(self, op: OpBuilder, audio_features: ir.Value): - hidden_states = self.linear_1(op, audio_features) - hidden_states = self.act(op, hidden_states) - return self.linear_2(op, hidden_states) + hidden_states = getattr(self, self._first_name)(op, audio_features) + hidden_states = getattr(self, self._norm_name)(op, hidden_states) + return getattr(self, self._second_name)(op, hidden_states) + + +class VibeVoiceMultiModalProjector(_VibeVoiceConnector): + """TTS connector retaining the native ``linear_1``/``act``/``linear_2`` names.""" + + def __init__(self, input_dim: int, output_dim: int): + super().__init__( + input_dim, + output_dim, + first_name="linear_1", + norm_name="act", + second_name="linear_2", + ) + + +class VibeVoiceSpeechConnector(_VibeVoiceConnector): + """ASR connector matching the public source ``fc1``/``norm``/``fc2`` weights.""" + + def __init__(self, input_dim: int, output_dim: int): + super().__init__( + input_dim, + output_dim, + first_name="fc1", + norm_name="norm", + second_name="fc2", + ) class VibeVoiceAcousticProjector(nn.Module): @@ -601,7 +701,7 @@ def forward( class VibeVoiceDecoderModel(nn.Module): """Qwen2 decoder returning vocabulary logits and post-norm hidden states.""" - def __init__(self, config: VibeVoiceConfig): + def __init__(self, config: VibeVoiceConfig | VibeVoiceASRConfig): super().__init__() self.layers = nn.ModuleList( [DecoderLayer(config) for _ in range(config.num_hidden_layers)] @@ -947,3 +1047,232 @@ def preprocess_weights( elif key == "lm_head.weight": routed["decoder.lm_head.weight"] = value return routed + + +class VibeVoiceASRAudioEncoder(nn.Module): + """Encode waveform chunks with VibeVoice ASR's two causal tokenizers. + + The public model samples acoustic latents from a host-provided pair of random + tensors, while semantic latents use their mean directly. Both tokenizer cache + sets are explicit graph state so the host can preserve the source's long-audio + causal encoding semantics. + """ + + def __init__(self, config: VibeVoiceASRConfig): + super().__init__() + self.acoustic_tokenizer = VibeVoiceTokenizerEncoder(config.acoustic_tokenizer) + self.semantic_tokenizer = VibeVoiceTokenizerEncoder(config.semantic_tokenizer) + self.acoustic_connector = VibeVoiceSpeechConnector( + config.acoustic_tokenizer.hidden_size, + config.hidden_size, + ) + self.semantic_connector = VibeVoiceSpeechConnector( + config.semantic_tokenizer.hidden_size, + config.hidden_size, + ) + self.acoustic_cache_specs = self.acoustic_tokenizer.cache_specs + self.semantic_cache_specs = self.semantic_tokenizer.cache_specs + self._acoustic_std_scale = config.acoustic_tokenizer.vae_std / 0.8 + self._dtype = config.dtype + + def forward( + self, + op: OpBuilder, + speech_tensors: ir.Value, + speech_masks: ir.Value, + acoustic_sample_noise: ir.Value, + acoustic_latent_noise: ir.Value, + acoustic_past_conv_states: Sequence[ir.Value], + semantic_past_conv_states: Sequence[ir.Value], + is_final_chunk: ir.Value, + ): + # The processor emits raw [batch, samples] speech_tensors; the reference + # adds the tokenizer's mono channel dimension immediately before encoding. + waveform = op.Cast(op.Unsqueeze(speech_tensors, [1]), to=self._dtype) + acoustic_mean, acoustic_present = self.acoustic_tokenizer( + op, + waveform, + acoustic_past_conv_states, + is_final_chunk, + ) + # The source samples one standard deviation per batch item, then applies + # independent elementwise noise. Keeping both tensors host-owned makes + # this stage deterministic and exactly replayable. + acoustic_std = op.Mul(acoustic_sample_noise, self._acoustic_std_scale) + acoustic_latents = op.Add( + acoustic_mean, + op.Mul(op.Unsqueeze(acoustic_std, [1, 2]), acoustic_latent_noise), + ) + semantic_latents, semantic_present = self.semantic_tokenizer( + op, + waveform, + semantic_past_conv_states, + is_final_chunk, + ) + acoustic_embeds = self.acoustic_connector(op, acoustic_latents) + semantic_embeds = self.semantic_connector(op, semantic_latents) + speech_embeds = op.Add(acoustic_embeds, semantic_embeds) # (batch, frames, hidden) + speech_indices = op.Transpose(op.NonZero(speech_masks), perm=[1, 0]) + return ( + op.GatherND(speech_embeds, speech_indices), # (valid_speech_frames, hidden) + acoustic_present, + semantic_present, + ) + + +class VibeVoiceASREmbeddingModel(nn.Module): + """Replace ASR speech-placeholder embeddings with flattened encoded speech frames.""" + + def __init__(self, config: VibeVoiceASRConfig): + super().__init__() + self.embed_tokens = Embedding(config.vocab_size, config.hidden_size) + + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + speech_embeds: ir.Value, + acoustic_input_mask: ir.Value, + ): + inputs_embeds = self.embed_tokens(op, input_ids) + placeholder_indices = op.Transpose(op.NonZero(acoustic_input_mask), perm=[1, 0]) + return op.ScatterND(inputs_embeds, placeholder_indices, speech_embeds) + + +class VibeVoiceASRForConditionalGeneration(nn.Module): + """VibeVoice streaming ASR with explicit encoder and decoder state. + + Mirrors the executable reference at + ``microsoft/VibeVoice@505653d3873b065a488aea551c6ee3dc51d3062f`` with + ``transformers==4.51.3`` for the pinned Microsoft 1.5B and 7B streaming + ASR checkpoints. + + ```mermaid + flowchart LR + W[24 kHz waveform chunk] --> A[Acoustic causal tokenizer] + W --> S[Semantic causal tokenizer] + A --> AN[Gaussian acoustic sample] + AN --> AC[acoustic connector] + S --> SC[semantic connector] + AC --> SUM[Add and retain valid speech frames] + SC --> SUM + TOK[Prompt + speech placeholders] --> E[Token embedding replacement] + SUM --> E + E --> Q[Qwen2 decoder + KV cache] + Q --> TXT[Text / speaker JSON tokens] + ``` + + ``audio_encoder`` accepts the processor's 24 kHz ``speech_tensors`` [B, S] + and ``speech_masks`` [B, F] directly; it returns speech embeddings in + row-major mask order. ``embedding`` accepts the processor's + ``acoustic_input_mask`` directly. The host owns prompt formatting, hotword + text, sampling, forced ``<|text_chunk_end|>`` insertion, cache lifetime, + and speaker-JSON parsing. + """ + + default_task = "vibevoice-asr-streaming" + category = "Audio" + config_class = VibeVoiceASRConfig + + # Acoustic tokenizer decoder tensors are part of the VAE training checkpoint, + # but the public ASR ``encode_speech`` path never invokes them. + INTENTIONALLY_UNUSED_WEIGHT_PREFIXES: ClassVar[tuple[str, ...]] = ( + "model.acoustic_tokenizer.decoder.", + ) + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "audio_encoder": ( + "model.acoustic_tokenizer.encoder", + "model.semantic_tokenizer.encoder", + "model.acoustic_connector", + "model.semantic_connector", + ), + "embedding": ("model.language_model.embed_tokens",), + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "lm_head", + ), + } + + def __init__(self, config: VibeVoiceASRConfig): + super().__init__() + self.config = config + self.audio_encoder = VibeVoiceASRAudioEncoder(config) + self.embedding = VibeVoiceASREmbeddingModel(config) + self.decoder = VibeVoiceDecoderModel(config) + + def forward(self, op: OpBuilder, *args, **kwargs): + raise NotImplementedError( + "VibeVoiceASRStreamingTask exports each ASR stage independently" + ) + + @staticmethod + def _encoder_weight_suffix(suffix: str) -> str: + """Map executable ASR encoder hierarchy to the shared tokenizer hierarchy.""" + if suffix.startswith("downsample_layers."): + _, index_text, zero, remainder = suffix.split(".", maxsplit=3) + if zero != "0": + raise ValueError(f"Unsupported VibeVoice ASR downsample path: {suffix}") + index = int(index_text) + prefix = "stem" if index == 0 else f"conv_layers.{index - 1}" + return f"{prefix}.{remainder}" + if suffix.startswith("stages."): + _, index_text, remainder = suffix.split(".", maxsplit=2) + index = int(index_text) + prefix = "stem.stage" if index == 0 else f"conv_layers.{index - 1}.stage" + return f"{prefix}.{remainder}".replace( + ".mixer.conv.conv.conv.", + ".mixer.conv.", + ) + if suffix.startswith("head."): + return suffix.replace("head.conv.conv.", "head.conv.") + raise ValueError(f"Unsupported VibeVoice ASR encoder weight: {suffix}") + + def preprocess_weights( + self, + state_dict: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Route every inference-used tensor and explicitly exclude the unused VAE decoder.""" + routed: dict[str, torch.Tensor] = {} + has_explicit_lm_head = "lm_head.weight" in state_dict + for key, value in state_dict.items(): + if key.startswith("model.acoustic_tokenizer.encoder."): + suffix = key.removeprefix("model.acoustic_tokenizer.encoder.") + routed[ + f"audio_encoder.acoustic_tokenizer.{self._encoder_weight_suffix(suffix)}" + ] = value + elif key.startswith("model.semantic_tokenizer.encoder."): + suffix = key.removeprefix("model.semantic_tokenizer.encoder.") + routed[ + f"audio_encoder.semantic_tokenizer.{self._encoder_weight_suffix(suffix)}" + ] = value + elif key.startswith("model.acoustic_connector."): + suffix = key.removeprefix("model.acoustic_connector.") + routed[f"audio_encoder.acoustic_connector.{suffix}"] = value + elif key.startswith("model.semantic_connector."): + suffix = key.removeprefix("model.semantic_connector.") + routed[f"audio_encoder.semantic_connector.{suffix}"] = value + elif key.startswith("model.language_model.embed_tokens."): + suffix = key.removeprefix("model.language_model.embed_tokens.") + routed[f"embedding.embed_tokens.{suffix}"] = value + # Some tied checkpoints retain lm_head.weight in their index. + # Route that trained tensor below rather than silently overwriting it. + if ( + suffix == "weight" + and self.config.tie_word_embeddings + and not has_explicit_lm_head + ): + routed["decoder.lm_head.weight"] = value + elif key.startswith("model.language_model.layers."): + suffix = key.removeprefix("model.language_model.") + routed[f"decoder.{suffix}"] = value + elif key.startswith("model.language_model.norm."): + suffix = key.removeprefix("model.language_model.") + routed[f"decoder.{suffix}"] = value + elif key == "lm_head.weight": + routed["decoder.lm_head.weight"] = value + elif key.startswith(self.INTENTIONALLY_UNUSED_WEIGHT_PREFIXES): + continue + else: + raise ValueError(f"Unexpected VibeVoice ASR checkpoint tensor: {key}") + return routed diff --git a/src/mobius/models/vibevoice_asr_test.py b/src/mobius/models/vibevoice_asr_test.py new file mode 100644 index 000000000..88e361b7f --- /dev/null +++ b/src/mobius/models/vibevoice_asr_test.py @@ -0,0 +1,614 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Graph-contract and checkpoint-routing tests for VibeVoice streaming ASR.""" + +from __future__ import annotations + +import dataclasses +import json +import subprocess +from importlib import metadata +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import onnx_ir as ir +import pytest +import torch + +from mobius._builder import build_from_module +from mobius._configs import VibeVoiceASRConfig, VibeVoiceTokenizerConfig +from mobius._pipeline_contract import ( + optional_input_contract, + requires_arbitrary_attention_mask, +) +from mobius._registry import registry +from mobius._testing.ort_inference import OnnxModelSession +from mobius.models.vibevoice import ( + VIBEVOICE_ASR_MODEL_REVISIONS, + VIBEVOICE_ASR_SOURCE_REVISION, + VibeVoiceASRForConditionalGeneration, +) +from mobius.tasks import VibeVoiceASRStreamingTask + + +def _config() -> VibeVoiceASRConfig: + tokenizer = VibeVoiceTokenizerConfig( + hidden_size=4, + kernel_size=3, + num_filters=4, + downsampling_ratios=[2, 2], + depths=[1, 1, 1], + ffn_expansion=2, + vae_std=0.5, + ) + return VibeVoiceASRConfig( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=8, + vocab_size=64, + max_position_embeddings=128, + rms_norm_eps=1e-6, + hidden_act="silu", + attn_qkv_bias=True, + rope_type="default", + acoustic_tokenizer=tokenizer, + semantic_tokenizer=dataclasses.replace( + tokenizer, + hidden_size=6, + std_dist_type="none", + ), + ) + + +def _source_key_for_parameter(parameter_name: str) -> str: + """Invert the public ASR source-to-package mapping for exhaustive routing tests.""" + if parameter_name.startswith("audio_encoder.acoustic_tokenizer."): + tower = "acoustic" + suffix = parameter_name.removeprefix("audio_encoder.acoustic_tokenizer.") + elif parameter_name.startswith("audio_encoder.semantic_tokenizer."): + tower = "semantic" + suffix = parameter_name.removeprefix("audio_encoder.semantic_tokenizer.") + else: + tower = None + suffix = parameter_name + + if tower is not None: + if suffix.startswith("stem.conv."): + suffix = f"downsample_layers.0.0.{suffix.removeprefix('stem.')}" + elif suffix.startswith("conv_layers."): + _, index_text, remainder = suffix.split(".", maxsplit=2) + if remainder.startswith("stage."): + suffix = f"stages.{int(index_text) + 1}.{remainder.removeprefix('stage.')}" + else: + suffix = f"downsample_layers.{int(index_text) + 1}.0.{remainder}" + elif suffix.startswith("stem.stage."): + suffix = f"stages.0.{suffix.removeprefix('stem.stage.')}" + elif not suffix.startswith("head."): + raise AssertionError(f"Unexpected tokenizer parameter {parameter_name}") + suffix = suffix.replace(".mixer.conv.", ".mixer.conv.conv.conv.") + suffix = suffix.replace("head.conv.", "head.conv.conv.") + return f"model.{tower}_tokenizer.encoder.{suffix}" + + if parameter_name.startswith("audio_encoder.acoustic_connector."): + return "model.acoustic_connector." + parameter_name.removeprefix( + "audio_encoder.acoustic_connector." + ) + if parameter_name.startswith("audio_encoder.semantic_connector."): + return "model.semantic_connector." + parameter_name.removeprefix( + "audio_encoder.semantic_connector." + ) + if parameter_name.startswith("embedding.embed_tokens."): + return "model.language_model.embed_tokens." + parameter_name.removeprefix( + "embedding.embed_tokens." + ) + if parameter_name.startswith(("decoder.layers.", "decoder.norm.")): + return "model.language_model." + parameter_name.removeprefix("decoder.") + if parameter_name == "decoder.lm_head.weight": + return "lm_head.weight" + raise AssertionError(f"Unexpected ASR package parameter {parameter_name}") + + +class TestVibeVoiceASR: + """Exercise every stage, runtime contract, and trained-weight route at tiny scale.""" + + def test_stage_contract_and_explicit_state(self): + config = _config() + module = VibeVoiceASRForConditionalGeneration(config) + package = VibeVoiceASRStreamingTask().build(module, config) + + assert set(package) == {"audio_encoder", "embedding", "decoder"} + audio_inputs = {value.name for value in package["audio_encoder"].graph.inputs} + audio_outputs = {value.name for value in package["audio_encoder"].graph.outputs} + assert { + "speech_tensors", + "speech_masks", + "acoustic_sample_noise", + "acoustic_latent_noise", + "is_final_chunk", + } <= audio_inputs + for prefix, names in ( + ("past_acoustic_conv", audio_inputs), + ("present_acoustic_conv", audio_outputs), + ("past_semantic_conv", audio_inputs), + ("present_semantic_conv", audio_outputs), + ): + assert {f"{prefix}.{index}" for index in range(7)} <= names + + speech_embeds = next( + value + for value in package["embedding"].graph.inputs + if value.name == "speech_embeds" + ) + assert optional_input_contract(speech_embeds) == { + "presence": "audio", + "absent": {"kind": "zeros", "shape": [0, config.hidden_size]}, + } + assert requires_arbitrary_attention_mask(package["decoder"].graph) + + def test_cuda_fp16_excludes_prefix_only_gqa_fusion(self): + config = dataclasses.replace(_config(), dtype=ir.DataType.FLOAT16) + package = build_from_module( + VibeVoiceASRForConditionalGeneration(config), + config, + task=VibeVoiceASRStreamingTask(), + execution_provider="cuda", + ) + decoder_nodes = list(package["decoder"].graph.all_nodes()) + assert any(node.op_type == "Attention" for node in decoder_nodes) + assert not any(node.op_type == "GroupQueryAttention" for node in decoder_nodes) + + def test_checkpoint_routes_every_inference_parameter_once(self): + config = _config() + module = VibeVoiceASRForConditionalGeneration(config) + package = VibeVoiceASRStreamingTask().build(module, config) + parameter_names = { + value.name + for model in package.values() + for value in model.graph.initializers.values() + if value.const_value is None + } + state_dict = { + _source_key_for_parameter(name): torch.zeros(1) for name in parameter_names + } + # The acoustic VAE decoder is present in the publication checkpoint but + # is provably not on the executable ASR encode_speech path. + state_dict["model.acoustic_tokenizer.decoder.head.conv.conv.weight"] = torch.zeros(1) + routed = module.preprocess_weights(state_dict) + + assert set(routed) == parameter_names + assert len(routed) == len(parameter_names) + assert not any("acoustic_tokenizer.decoder" in name for name in routed) + + def test_registration_is_pinned_and_architecture_specific(self): + registration = registry.get_registration("VibeVoiceForASRStreamingTraining") + assert registration.module_class is VibeVoiceASRForConditionalGeneration + assert registration.task == "vibevoice-asr-streaming" + assert registration.config_class is VibeVoiceASRConfig + assert registration.test_model_id in VIBEVOICE_ASR_MODEL_REVISIONS + assert ( + registration.test_revision + == VIBEVOICE_ASR_MODEL_REVISIONS[registration.test_model_id] + ) + + def test_tied_checkpoint_lm_head_preserves_its_explicit_tensor(self): + """A tied checkpoint's explicit LM head must not be overwritten by an embedding fallback.""" + module = VibeVoiceASRForConditionalGeneration( + dataclasses.replace(_config(), tie_word_embeddings=True) + ) + embedding_weight = torch.zeros(1) + lm_head_weight = torch.ones(1) + + routed = module.preprocess_weights( + { + "model.language_model.embed_tokens.weight": embedding_weight, + "lm_head.weight": lm_head_weight, + } + ) + + assert set(routed) == {"embedding.embed_tokens.weight", "decoder.lm_head.weight"} + assert routed["embedding.embed_tokens.weight"] is embedding_weight + assert routed["decoder.lm_head.weight"] is lm_head_weight + + def test_vibevoice_dispatch_rejects_ambiguous_or_unknown_architectures(self): + from mobius.integrations.transformers._builder import _resolve_module_class + + ambiguous = SimpleNamespace(model_type="vibevoice", architectures=[]) + with pytest.raises(ValueError, match="exactly one recognized architecture"): + _resolve_module_class("vibevoice", ambiguous, None, None) + + unknown = SimpleNamespace( + model_type="vibevoice", + architectures=["VibeVoiceForSomethingElse"], + ) + with pytest.raises(ValueError, match="Unsupported VibeVoice architecture"): + _resolve_module_class("vibevoice", unknown, None, None) + + tts = SimpleNamespace( + model_type="vibevoice", + architectures=["VibeVoiceForConditionalGeneration"], + ) + module_class, task, model_type = _resolve_module_class("vibevoice", tts, None, None) + assert module_class.__name__ == "VibeVoiceForConditionalGeneration" + assert task is None + assert model_type == "vibevoice" + + def test_asr_config_rejects_non_executable_tokenizer_variants(self): + from mobius._configs.vibevoice import _asr_tokenizer_config + + with pytest.raises(ValueError, match="mixer_layer='conv'"): + _asr_tokenizer_config( + SimpleNamespace(mixer_layer="conv"), + default_hidden_size=64, + ) + + +def _require_pinned_reference(): + """Load the independently executable reference only from a verified local checkout.""" + import transformers + + if transformers.__version__ != "4.51.3": + pytest.skip("Synthetic VibeVoice ASR parity requires transformers==4.51.3.") + try: + direct_url = metadata.distribution("vibevoice").read_text("direct_url.json") + except metadata.PackageNotFoundError: + direct_url = None + source_is_pinned = False + if direct_url is not None: + vcs_info = json.loads(direct_url).get("vcs_info", {}) + if vcs_info.get("commit_id") != VIBEVOICE_ASR_SOURCE_REVISION: + pytest.skip( + "Synthetic VibeVoice ASR parity requires " + f"VibeVoice@{VIBEVOICE_ASR_SOURCE_REVISION}." + ) + source_is_pinned = True + modeling = pytest.importorskip("vibevoice.modular.modeling_vibevoice_asr") + configuration = pytest.importorskip("vibevoice.modular.configuration_vibevoice") + tokenizer = pytest.importorskip("vibevoice.modular.modular_vibevoice_tokenizer") + if source_is_pinned: + return modeling, configuration, tokenizer + source_root = Path(modeling.__file__).parents[2] + try: + source_revision = subprocess.run( + ["git", "-C", str(source_root), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + except (OSError, subprocess.CalledProcessError): + pytest.skip( + "Synthetic VibeVoice ASR parity requires installed VCS metadata or " + f"VibeVoice@{VIBEVOICE_ASR_SOURCE_REVISION} in direct_url.json." + ) + if source_revision != VIBEVOICE_ASR_SOURCE_REVISION: + pytest.skip( + "Synthetic VibeVoice ASR parity requires " + f"VibeVoice@{VIBEVOICE_ASR_SOURCE_REVISION}." + ) + return modeling, configuration, tokenizer + + +def _run_component(package, name: str, feeds: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + session = OnnxModelSession(package[name]) + try: + return session.run(feeds) + finally: + session.close() + + +@pytest.mark.integration +def test_vibevoice_asr_synthetic_two_chunk_prefill_and_cached_decode_parity(): + """Compare all exported stages to the pinned executable source with random tiny weights.""" + modeling, configuration, tokenizer_module = _require_pinned_reference() + source_tokenizer = { + "channels": 1, + "vae_dim": 4, + "fix_std": 0.5, + "encoder_n_filters": 4, + "encoder_ratios": [2, 2], + "encoder_depths": "1-1-1", + "mixer_layer": "depthwise_conv", + "conv_norm": "none", + "pad_mode": "constant", + "disable_last_norm": True, + "layernorm": "RMSNorm", + "layernorm_eps": 1e-5, + "conv_bias": True, + } + decoder = { + "model_type": "qwen2", + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + "vocab_size": 64, + "max_position_embeddings": 128, + "rms_norm_eps": 1e-6, + "hidden_act": "silu", + "rope_theta": 10_000.0, + "tie_word_embeddings": False, + } + torch.manual_seed(7) + reference = ( + modeling.VibeVoiceASRForConditionalGeneration( + configuration.VibeVoiceASRConfig( + acoustic_tokenizer_config={**source_tokenizer, "std_dist_type": "gaussian"}, + semantic_tokenizer_config={ + **source_tokenizer, + "vae_dim": 6, + "std_dist_type": "none", + }, + decoder_config=decoder, + ) + ) + .float() + .eval() + ) + + mobius_tokenizer = VibeVoiceTokenizerConfig( + hidden_size=4, + kernel_size=7, + num_filters=4, + downsampling_ratios=[2, 2], + depths=[1, 1, 1], + ffn_expansion=4, + vae_std=0.5, + std_dist_type="gaussian", + ) + config = VibeVoiceASRConfig( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=8, + vocab_size=64, + max_position_embeddings=128, + rms_norm_eps=1e-6, + hidden_act="silu", + attn_qkv_bias=True, + rope_type="default", + rope_theta=10_000.0, + acoustic_tokenizer=mobius_tokenizer, + semantic_tokenizer=dataclasses.replace( + mobius_tokenizer, + hidden_size=6, + std_dist_type="none", + ), + ) + module = VibeVoiceASRForConditionalGeneration(config) + package = VibeVoiceASRStreamingTask().build(module, config) + package.apply_weights(module.preprocess_weights(reference.state_dict())) + + waveforms = ( + torch.linspace(-0.5, 0.5, 16).reshape(1, 16).repeat(2, 1), + torch.linspace(0.25, -0.25, 12).reshape(1, 12).repeat(2, 1), + ) + sample_noise = torch.tensor([-0.4, 0.7]) + latent_noises = ( + torch.linspace(-1, 1, 2 * 4 * 4).reshape(2, 4, 4), + torch.linspace(1, -1, 2 * 3 * 4).reshape(2, 3, 4), + ) + reference_acoustic_cache = tokenizer_module.VibeVoiceTokenizerStreamingCache() + reference_semantic_cache = tokenizer_module.VibeVoiceTokenizerStreamingCache() + + def reference_audio_chunk(waveform, latent_noise, is_final_chunk): + sample_indices = torch.arange(waveform.shape[0]) + acoustic_mean = reference.model.acoustic_tokenizer.encode( + waveform.unsqueeze(1), + cache=reference_acoustic_cache, + sample_indices=sample_indices, + use_cache=True, + is_final_chunk=is_final_chunk, + ).mean + semantic_latents = reference.model.semantic_tokenizer.encode( + waveform.unsqueeze(1), + cache=reference_semantic_cache, + sample_indices=sample_indices, + use_cache=True, + is_final_chunk=is_final_chunk, + ).mean + acoustic_latents = ( + acoustic_mean + sample_noise[:, None, None] * (0.5 / 0.8) * latent_noise + ) + return reference.model.acoustic_connector( + acoustic_latents + ) + reference.model.semantic_connector(semantic_latents) + + with torch.no_grad(): + expected_audio = [ + reference_audio_chunk(waveforms[0], latent_noises[0], False), + reference_audio_chunk(waveforms[1], latent_noises[1], True), + ] + input_ids = torch.tensor([[3, 4, 5, 6, 7], [3, 4, 5, 6, 7]]) + acoustic_input_mask = torch.tensor( + [[False, True, True, True, True], [False, True, True, True, True]] + ) + reference_embeds = reference.get_input_embeddings()(input_ids) + reference_embeds[acoustic_input_mask] = expected_audio[0].reshape(-1, 16) + reference_prefill = reference( + inputs_embeds=reference_embeds, + attention_mask=torch.ones(2, 5, dtype=torch.long), + position_ids=torch.arange(5).repeat(2, 1), + use_cache=True, + return_dict=True, + ) + decode_ids = torch.tensor([[8], [9]]) + reference_decode = reference( + inputs_embeds=reference.get_input_embeddings()(decode_ids), + attention_mask=torch.ones(2, 6, dtype=torch.long), + position_ids=torch.full((2, 1), 5), + past_key_values=reference_prefill.past_key_values, + use_cache=True, + return_dict=True, + ) + + def initial_cache( + prefix: str, specs: tuple[tuple[int, int], ...] + ) -> dict[str, np.ndarray]: + return { + f"{prefix}.{index}": np.zeros((2, channels, left_pad), dtype=np.float32) + for index, (channels, left_pad) in enumerate(specs) + } + + first = _run_component( + package, + "audio_encoder", + { + "speech_tensors": waveforms[0].numpy(), + "speech_masks": np.ones((2, 4), dtype=bool), + "acoustic_sample_noise": sample_noise.numpy(), + "acoustic_latent_noise": latent_noises[0].numpy(), + "is_final_chunk": np.asarray(False), + **initial_cache("past_acoustic_conv", module.audio_encoder.acoustic_cache_specs), + **initial_cache("past_semantic_conv", module.audio_encoder.semantic_cache_specs), + }, + ) + second_feeds = { + "speech_tensors": waveforms[1].numpy(), + "speech_masks": np.ones((2, 3), dtype=bool), + "acoustic_sample_noise": sample_noise.numpy(), + "acoustic_latent_noise": latent_noises[1].numpy(), + "is_final_chunk": np.asarray(True), + } + for index in range(len(module.audio_encoder.acoustic_cache_specs)): + second_feeds[f"past_acoustic_conv.{index}"] = first[f"present_acoustic_conv.{index}"] + for index in range(len(module.audio_encoder.semantic_cache_specs)): + second_feeds[f"past_semantic_conv.{index}"] = first[f"present_semantic_conv.{index}"] + second = _run_component(package, "audio_encoder", second_feeds) + np.testing.assert_allclose( + first["speech_embeds"], + expected_audio[0].reshape(-1, 16).numpy(), + rtol=1e-5, + atol=1e-5, + ) + np.testing.assert_allclose( + second["speech_embeds"], + expected_audio[1].reshape(-1, 16).numpy(), + rtol=1e-5, + atol=1e-5, + ) + + embeds = _run_component( + package, + "embedding", + { + "input_ids": input_ids.numpy(), + "speech_embeds": first["speech_embeds"], + "acoustic_input_mask": acoustic_input_mask.numpy(), + }, + ) + prefill = _run_component( + package, + "decoder", + { + "inputs_embeds": embeds["inputs_embeds"], + "attention_mask": np.ones((2, 5), dtype=np.int64), + "position_ids": np.tile(np.arange(5, dtype=np.int64), (2, 1)), + "past_key_values.0.key": np.zeros((2, 1, 0, 8), dtype=np.float32), + "past_key_values.0.value": np.zeros((2, 1, 0, 8), dtype=np.float32), + }, + ) + decode_embeds = _run_component( + package, + "embedding", + { + "input_ids": decode_ids.numpy(), + "speech_embeds": np.zeros((0, 16), dtype=np.float32), + "acoustic_input_mask": np.zeros((2, 1), dtype=bool), + }, + ) + decode = _run_component( + package, + "decoder", + { + "inputs_embeds": decode_embeds["inputs_embeds"], + "attention_mask": np.ones((2, 6), dtype=np.int64), + "position_ids": np.full((2, 1), 5, dtype=np.int64), + "past_key_values.0.key": prefill["present.0.key"], + "past_key_values.0.value": prefill["present.0.value"], + }, + ) + np.testing.assert_allclose( + prefill["logits"], + reference_prefill.logits.numpy(), + rtol=1e-5, + atol=1e-5, + ) + np.testing.assert_allclose( + decode["logits"], + reference_decode.logits.numpy(), + rtol=1e-5, + atol=1e-5, + ) + + +@pytest.mark.integration +@pytest.mark.parametrize( + ("model_id", "revision"), + tuple(VIBEVOICE_ASR_MODEL_REVISIONS.items()), +) +def test_vibevoice_asr_pinned_processor_contract_for_hotwords_and_speakers(model_id, revision): + """Validate processor rows, left padding, bilingual hotword prompts, and speaker JSON.""" + _require_pinned_reference() + processor_module = pytest.importorskip("vibevoice.processor.vibevoice_asr_processor") + processor = processor_module.VibeVoiceASRProcessor.from_pretrained( + model_id, + revision=revision, + ) + english = np.linspace(-0.1, 0.1, 3200, dtype=np.float32) + chinese = np.linspace(0.1, -0.1, 6500, dtype=np.float32) + plain = processor(english, sampling_rate=24_000, return_tensors="pt") + hotwords = processor( + english, + sampling_rate=24_000, + return_tensors="pt", + context_info="hotwords: Mobius, 你好", + ) + batch = processor( + [english, chinese], + sampling_rate=24_000, + return_tensors="pt", + context_info="hotwords: Mobius, 你好", + ) + + assert set(batch) == { + "input_ids", + "attention_mask", + "acoustic_input_mask", + "speech_tensors", + "speech_masks", + } + assert batch["speech_tensors"].shape == (2, 6500) + assert batch["speech_masks"].sum(dim=1).tolist() == [1, 3] + assert batch["attention_mask"][0, 0].item() == 0 + assert batch["attention_mask"][1, 0].item() == 1 + assert plain["input_ids"].tolist() != hotwords["input_ids"].tolist() + assert { + token: processor.tokenizer.convert_tokens_to_ids(token) + for token in ( + "<|object_ref_start|>", + "<|object_ref_end|>", + "<|box_start|>", + "<|text_chunk_end|>", + ) + } == { + "<|object_ref_start|>": 151646, + "<|object_ref_end|>": 151647, + "<|box_start|>": 151648, + "<|text_chunk_end|>": 151665, + } + assert processor.post_process_transcription( + '[{"Start time": 0.0, "End time": 1.2, "Speaker ID": "spk-1", ' + '"Content": "hello"}, {"Start": 1.2, "End": 2.0, "Speaker": "说话人2", ' + '"Content": "你好"}]' + ) == [ + {"start_time": 0.0, "end_time": 1.2, "speaker_id": "spk-1", "text": "hello"}, + {"start_time": 1.2, "end_time": 2.0, "speaker_id": "说话人2", "text": "你好"}, + ] diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index b82e691c5..91a02304c 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -105,6 +105,7 @@ "TASK_REGISTRY", "TTSTask", "VibeVoiceTask", + "VibeVoiceASRStreamingTask", "T5TextEncoderTask", "VAETask", "VideoDenoisingTask", @@ -206,6 +207,7 @@ from mobius.tasks._tts import TTSTask from mobius.tasks._vae import VAETask from mobius.tasks._vibevoice import VibeVoiceTask +from mobius.tasks._vibevoice_asr import VibeVoiceASRStreamingTask from mobius.tasks._video_denoising import VideoDenoisingTask from mobius.tasks._video_vae import VideoVAETask from mobius.tasks._vision_encoder_decoder import VisionEncoderDecoderTask @@ -312,6 +314,7 @@ "ssm2-text-generation": SSM2CausalLMTask, "tts": TTSTask, "vibevoice-tts": VibeVoiceTask, + "vibevoice-asr-streaming": VibeVoiceASRStreamingTask, "video-denoising": VideoDenoisingTask, "video-vae": VideoVAETask, "world-model": WorldModelTask, diff --git a/src/mobius/tasks/_streaming_convolution.py b/src/mobius/tasks/_streaming_convolution.py new file mode 100644 index 000000000..df5ca4993 --- /dev/null +++ b/src/mobius/tasks/_streaming_convolution.py @@ -0,0 +1,38 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Explicit graph I/O helpers for causal convolution streaming state.""" + +from __future__ import annotations + +import onnx_ir as ir + + +def make_conv_cache_inputs( + builder, + specs: tuple[tuple[int, int], ...], + batch: ir.SymbolicDim, + dtype: ir.DataType, + *, + prefix: str = "past_conv", +) -> list[ir.Value]: + """Register one channels-first causal cache input for every convolution.""" + return [ + builder.input( + f"{prefix}.{index}", + dtype=dtype, + shape=[batch, channels, left_pad], + ) + for index, (channels, left_pad) in enumerate(specs) + ] + + +def register_conv_cache_outputs( + builder, + values: list[ir.Value], + *, + prefix: str = "present_conv", +) -> None: + """Register causal convolution state outputs under a stable component prefix.""" + for index, value in enumerate(values): + builder.add_output(value, f"{prefix}.{index}") diff --git a/src/mobius/tasks/_vibevoice.py b/src/mobius/tasks/_vibevoice.py index 721c3186a..e45f40507 100644 --- a/src/mobius/tasks/_vibevoice.py +++ b/src/mobius/tasks/_vibevoice.py @@ -18,33 +18,13 @@ declare_optional_input, ) from mobius.tasks._base import ComponentSpec, ModelTask, _make_graph, _make_model -from mobius.tasks._cache_utils import ( - _make_kv_cache_inputs, - _register_kv_cache_outputs, +from mobius.tasks._cache_utils import _make_kv_cache_inputs, _register_kv_cache_outputs +from mobius.tasks._streaming_convolution import ( + make_conv_cache_inputs, + register_conv_cache_outputs, ) -def _make_conv_cache_inputs( - builder, - specs: tuple[tuple[int, int], ...], - batch: ir.SymbolicDim, - dtype: ir.DataType, -) -> list[ir.Value]: - return [ - builder.input( - f"past_conv.{index}", - dtype=dtype, - shape=[batch, channels, left_pad], - ) - for index, (channels, left_pad) in enumerate(specs) - ] - - -def _register_conv_cache_outputs(builder, values: list[ir.Value]) -> None: - for index, value in enumerate(values): - builder.add_output(value, f"present_conv.{index}") - - class VibeVoiceTask(ModelTask): """Build all neural stages needed by VibeVoice's continuous-token TTS loop.""" @@ -269,7 +249,7 @@ def _build_audio_decoder( dtype=config.dtype, shape=[batch, "audio_frames", config.acoustic_tokenizer.hidden_size], ) - past = _make_conv_cache_inputs( + past = make_conv_cache_inputs( builder, module.cache_specs, batch, @@ -278,7 +258,7 @@ def _build_audio_decoder( waveform, present = module(builder.op, scaled_latents, past) waveform.shape = ir.Shape([batch, config.acoustic_tokenizer.channels, "audio_samples"]) builder.add_output(waveform, "waveform") - _register_conv_cache_outputs(builder, present) + register_conv_cache_outputs(builder, present) return _make_model(graph) def _build_semantic_encoder( @@ -293,7 +273,7 @@ def _build_semantic_encoder( dtype=config.dtype, shape=[batch, config.semantic_tokenizer.channels, "audio_samples"], ) - past = _make_conv_cache_inputs( + past = make_conv_cache_inputs( builder, module.cache_specs, batch, @@ -301,7 +281,7 @@ def _build_semantic_encoder( ) latents, present = module(builder.op, waveform, past) builder.add_output(latents, "semantic_latents") - _register_conv_cache_outputs(builder, present) + register_conv_cache_outputs(builder, present) return _make_model(graph) def _build_semantic_projection( diff --git a/src/mobius/tasks/_vibevoice_asr.py b/src/mobius/tasks/_vibevoice_asr.py new file mode 100644 index 000000000..8f173c12e --- /dev/null +++ b/src/mobius/tasks/_vibevoice_asr.py @@ -0,0 +1,191 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Graph contracts for VibeVoice streaming ASR's staged inference pipeline.""" + +from __future__ import annotations + +from typing import ClassVar + +import onnx_ir as ir +from onnxscript import nn + +from mobius._configs import VibeVoiceASRConfig +from mobius._model_package import ModelPackage +from mobius._pipeline_contract import ( + declare_arbitrary_attention_mask, + declare_component_presence, + declare_optional_input, +) +from mobius.tasks._base import ComponentSpec, ModelTask, _make_graph, _make_model +from mobius.tasks._cache_utils import _make_kv_cache_inputs, _register_kv_cache_outputs +from mobius.tasks._streaming_convolution import ( + make_conv_cache_inputs, + register_conv_cache_outputs, +) + + +class VibeVoiceASRStreamingTask(ModelTask): + """Build the executable ASR stages while retaining host-owned stream orchestration.""" + + model_roles: ClassVar[dict[str, str]] = { + "audio_encoder": "encoder", + "embedding": "embedding", + "decoder": "decoder", + } + components: ClassVar[ComponentSpec] = ComponentSpec( + audio_encoder="audio_encoder", + embedding="embedding", + decoder="decoder", + ) + + def build(self, module: nn.Module, config: VibeVoiceASRConfig) -> ModelPackage: + self._validate_components(module) + return ModelPackage( + { + "audio_encoder": self._build_audio_encoder(module.audio_encoder, config), + "embedding": self._build_embedding(module.embedding, config), + "decoder": self._build_decoder(module.decoder, config), + }, + config=config, + ) + + def _build_audio_encoder(self, module: nn.Module, config: VibeVoiceASRConfig) -> ir.Model: + batch = ir.SymbolicDim("batch") + samples = ir.SymbolicDim("audio_samples") + frames = ir.SymbolicDim("audio_frames") + graph, builder = _make_graph(name="vibevoice_asr_audio_encoder") + speech_tensors = builder.input( + "speech_tensors", + dtype=ir.DataType.FLOAT, + shape=[batch, samples], + ) + speech_masks = builder.input( + "speech_masks", + dtype=ir.DataType.BOOL, + shape=[batch, frames], + ) + acoustic_sample_noise = builder.input( + "acoustic_sample_noise", + dtype=config.dtype, + shape=[batch], + ) + acoustic_latent_noise = builder.input( + "acoustic_latent_noise", + dtype=config.dtype, + shape=[batch, frames, config.acoustic_tokenizer.hidden_size], + ) + is_final_chunk = builder.input("is_final_chunk", dtype=ir.DataType.BOOL, shape=[]) + acoustic_past = make_conv_cache_inputs( + builder, + module.acoustic_cache_specs, + batch, + config.dtype, + prefix="past_acoustic_conv", + ) + semantic_past = make_conv_cache_inputs( + builder, + module.semantic_cache_specs, + batch, + config.dtype, + prefix="past_semantic_conv", + ) + speech_embeds, acoustic_present, semantic_present = module( + builder.op, + speech_tensors, + speech_masks, + acoustic_sample_noise, + acoustic_latent_noise, + acoustic_past, + semantic_past, + is_final_chunk, + ) + speech_embeds.shape = ir.Shape(["valid_speech_frames", config.hidden_size]) + builder.add_output(speech_embeds, "speech_embeds") + register_conv_cache_outputs( + builder, + acoustic_present, + prefix="present_acoustic_conv", + ) + register_conv_cache_outputs( + builder, + semantic_present, + prefix="present_semantic_conv", + ) + declare_component_presence(graph, "audio") + return _make_model(graph) + + def _build_embedding(self, module: nn.Module, config: VibeVoiceASRConfig) -> ir.Model: + graph, builder = _make_graph(name="vibevoice_asr_embedding") + input_ids = builder.input( + "input_ids", + dtype=ir.DataType.INT64, + shape=["batch", "sequence_length"], + ) + speech_embeds = builder.input( + "speech_embeds", + dtype=config.dtype, + shape=["valid_speech_frames", config.hidden_size], + ) + declare_optional_input( + speech_embeds, + presence="audio", + absent_shape=[0, config.hidden_size], + ) + acoustic_input_mask = builder.input( + "acoustic_input_mask", + dtype=ir.DataType.BOOL, + shape=["batch", "sequence_length"], + ) + inputs_embeds = module( + builder.op, + input_ids, + speech_embeds, + acoustic_input_mask, + ) + builder.add_output(inputs_embeds, "inputs_embeds") + return _make_model(graph) + + def _build_decoder(self, module: nn.Module, config: VibeVoiceASRConfig) -> ir.Model: + batch = ir.SymbolicDim("batch") + sequence = ir.SymbolicDim("sequence_length") + past_sequence = ir.SymbolicDim("past_sequence_length") + graph, builder = _make_graph(name="vibevoice_asr_decoder") + # The processor left-pads batches, making valid tokens a suffix. The + # generic GQA ABI only represents valid prefixes, so it must stay unfused. + declare_arbitrary_attention_mask(graph) + inputs_embeds = builder.input( + "inputs_embeds", + dtype=config.dtype, + shape=[batch, sequence, config.hidden_size], + ) + attention_mask = builder.input( + "attention_mask", + dtype=ir.DataType.INT64, + shape=[batch, "past_sequence_length + sequence_length"], + ) + position_ids = builder.input( + "position_ids", + dtype=ir.DataType.INT64, + shape=[batch, sequence], + ) + past = _make_kv_cache_inputs( + builder, + config.num_hidden_layers, + config.num_key_value_heads, + config.head_dim, + config.dtype, + batch, + past_sequence, + ) + logits, hidden_states, present = module( + builder.op, + inputs_embeds, + attention_mask, + position_ids, + past, + ) + builder.add_output(logits, "logits") + builder.add_output(hidden_states, "last_hidden_state") + _register_kv_cache_outputs(builder, present) + return _make_model(graph) diff --git a/tests/_test_configs.py b/tests/_test_configs.py index f9aa0d59e..8e3888fa9 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -59,6 +59,7 @@ Sam2Config, SegformerConfig, SenseNovaU1Config, + VibeVoiceASRConfig, VibeVoiceConfig, VibeVoiceDiffusionConfig, VibeVoiceTokenizerConfig, @@ -3733,6 +3734,49 @@ def vl_overrides(model_type: str) -> dict: }, True, ), + # --- VibeVoice streaming ASR (three-stage causal tokenizer + Qwen2 split) --- + ( + "VibeVoiceForASRStreamingTraining", + { + "_config_cls": VibeVoiceASRConfig, + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + "vocab_size": 64, + "max_position_embeddings": 128, + "rms_norm_eps": 1e-6, + "hidden_act": "silu", + "attn_qkv_bias": True, + "rope_type": "default", + "speech_start_token_id": 60, + "speech_end_token_id": 61, + "speech_placeholder_token_id": 62, + "text_chunk_end_token_id": 63, + "acoustic_tokenizer": VibeVoiceTokenizerConfig( + hidden_size=4, + kernel_size=3, + num_filters=4, + downsampling_ratios=[2, 2], + depths=[1, 1, 1], + ffn_expansion=2, + vae_std=0.5, + std_dist_type="gaussian", + ), + "semantic_tokenizer": VibeVoiceTokenizerConfig( + hidden_size=6, + kernel_size=3, + num_filters=4, + downsampling_ratios=[2, 2], + depths=[1, 1, 1], + ffn_expansion=2, + std_dist_type="none", + ), + }, + True, + ), # --- Qwen3-TTS Codec Tokenizer (codec, 2-model split) --- ( "qwen3_tts_tokenizer_12hz", diff --git a/tests/arch_validation_test.py b/tests/arch_validation_test.py index aec99859c..4c7e488cf 100644 --- a/tests/arch_validation_test.py +++ b/tests/arch_validation_test.py @@ -28,6 +28,7 @@ import logging +import onnx_ir as ir import pytest from mobius._registry import registry @@ -104,14 +105,20 @@ def _load_hf_config(model_id: str, revision: str | None = None): model type isn't registered in transformers. """ import transformers + from huggingface_hub import errors as hub_errors + strict_validation_error = getattr( + hub_errors, + "StrictDataclassClassValidationError", + ValueError, + ) try: return transformers.AutoConfig.from_pretrained( model_id, revision=revision, trust_remote_code=False, ) - except (ValueError, OSError): + except (strict_validation_error, ValueError, OSError): return _try_load_config_json(model_id, revision=revision) @@ -138,7 +145,14 @@ def _resolve_hf_config(hf_config, registration=None): owns_composite = ( registration is not None and getattr(registration, "config_class", None) is not None ) - if hasattr(hf_config, "talker_config"): + architectures = getattr(hf_config, "architectures", None) or [] + if ( + getattr(hf_config, "model_type", None) == "vibevoice" + and architectures == ["VibeVoiceForASRStreamingTraining"] + and hasattr(hf_config, "decoder_config") + ): + hf_config = hf_config.decoder_config + elif hasattr(hf_config, "talker_config"): talker = hf_config.talker_config # Qwen3-Omni talker nests the real model config under text_config if hasattr(talker, "text_config"): @@ -162,7 +176,7 @@ def _resolve_hf_config(hf_config, registration=None): return hf_config, parent_config -def _build_graph(model_type: str, model_id: str): +def _build_graph(model_type: str, model_id: str, *, revision: str | None = None): """Download config, build ONNX graph, return ``(ModelPackage, task)``. Uses get_task().build() directly (same pattern as the L1 graph tests) @@ -170,7 +184,10 @@ def _build_graph(model_type: str, model_id: str): (e.g. vision models with vocab_size=0). """ registration = registry.get_registration(model_type) - hf_config = _load_hf_config(model_id, revision=registration.test_revision) + hf_config = _load_hf_config( + model_id, + revision=registration.test_revision if revision is None else revision, + ) if hf_config is None: pytest.skip( f"Cannot download config for {model_id} (gated/private model or network error)" @@ -274,6 +291,60 @@ def test_graph_shapes_consistent(self, model_type: str, model_id: str): del pkg + @pytest.mark.parametrize( + ("model_id", "revision", "hidden_size", "q_heads", "kv_heads", "tied_embeddings"), + ( + ( + "microsoft/VibeVoice-ASR-Streaming-1.5B", + "4262d23d8a539a6530cf64fbd0b1751ef9a30853", + 1536, + 12, + 2, + True, + ), + ( + "microsoft/VibeVoice-ASR-Streaming-7B", + "60d858b518b4e19d404af3737f848fc185b30177", + 3584, + 28, + 4, + False, + ), + ), + ) + def test_vibevoice_asr_pinned_variants_build( + self, + model_id: str, + revision: str, + hidden_size: int, + q_heads: int, + kv_heads: int, + tied_embeddings: bool, + ): + """Build both official streaming-ASR variants from their pinned configs.""" + pkg, task = _build_graph( + "VibeVoiceForASRStreamingTraining", + model_id, + revision=revision, + ) + + config = pkg.config + assert config.hidden_size == hidden_size + assert config.num_attention_heads == q_heads + assert config.num_key_value_heads == kv_heads + assert config.tie_word_embeddings is tied_embeddings + assert config.dtype == ir.DataType.FLOAT + assert config.acoustic_tokenizer.hidden_size == 64 + assert config.semantic_tokenizer.hidden_size == 128 + assert config.compression_ratio == 3200 + assert config.sampling_rate == 24_000 + assert set(pkg) == {"audio_encoder", "embedding", "decoder"} + assert task.model_roles == { + "audio_encoder": "encoder", + "embedding": "embedding", + "decoder": "decoder", + } + class TestRegistryConsistency: """Verify registry and model class metadata are consistent.""" diff --git a/tests/build_graph/speech_test.py b/tests/build_graph/speech_test.py index 4e17e8d88..f542f3a54 100644 --- a/tests/build_graph/speech_test.py +++ b/tests/build_graph/speech_test.py @@ -59,6 +59,7 @@ "semantic_encoder", "semantic_projection", }, + "vibevoice-asr-streaming": {"audio_encoder", "embedding", "decoder"}, } diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index a93772804..efd9938cc 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -306,6 +306,11 @@ def _all_registered_with_test_id() -> dict[str, str]: "mms": "CTC ASR model — tested via TestBuildMMSGraph", "fastconformer_rnnt": "NeMo .nemo RNN-T ASR — tested via tests/nemo_rnnt_integration_test.py", "sortformer": "NeMo .nemo speaker diarization — tested via tests/sortformer_integration_test.py", + "VibeVoiceForASRStreamingTraining": "Streaming ASR has host-owned dual-convolution " + "state, arbitrary-mask decoder, hotword, and speaker-attribution orchestration that " + "the generic L4/L5 runner cannot drive. Pinned L1-L3 graph/config/source-parity and " + "complete checkpoint-index routing are covered for the 1.5B and 7B checkpoints; " + "real-weight goldens require a dedicated GPU workflow.", # --- Models requiring trust_remote_code --- "chatglm": "Requires trust_remote_code (custom HF modeling code)", "dots1": "Requires trust_remote_code (custom HF modeling code)", diff --git a/tests/weight_alignment_test.py b/tests/weight_alignment_test.py index 27521cd39..d46cd3f03 100644 --- a/tests/weight_alignment_test.py +++ b/tests/weight_alignment_test.py @@ -87,6 +87,8 @@ def _build_identity_state_dict(pkg: dict, param_names: set[str]) -> dict[str, to "opt", # ModernBert decoder: expects model.layers.* HF format with renames "modernbert-decoder", + # VibeVoice ASR translates its executable checkpoint hierarchy into split-stage names. + "VibeVoiceForASRStreamingTraining", } @@ -170,6 +172,49 @@ def test_vibevoice_native_hf_weights_cover_every_stage_parameter(): assert parameter_names == set(routed) +@pytest.mark.arch_validation +@pytest.mark.parametrize( + ("model_id", "revision"), + ( + ("microsoft/VibeVoice-ASR-Streaming-1.5B", "4262d23d8a539a6530cf64fbd0b1751ef9a30853"), + ("microsoft/VibeVoice-ASR-Streaming-7B", "60d858b518b4e19d404af3737f848fc185b30177"), + ), +) +def test_vibevoice_asr_pinned_weight_index_routes_every_inference_tensor_once( + model_id, revision +): + """Audit each pinned 1,177-tensor ASR checkpoint without downloading its weights.""" + import json + + from huggingface_hub import hf_hub_download + + from mobius import build + from mobius.models.vibevoice import VibeVoiceASRForConditionalGeneration + + with open( + hf_hub_download( + model_id, + "model.safetensors.index.json", + revision=revision, + ) + ) as handle: + checkpoint_names = set(json.load(handle)["weight_map"]) + + package = build(model_id, revision=revision, load_weights=False) + module = VibeVoiceASRForConditionalGeneration(package.config) + routed = module.preprocess_weights({name: torch.zeros(1) for name in checkpoint_names}) + parameter_names = _collect_parameter_names(package) + intentionally_unused = { + name + for name in checkpoint_names + if name.startswith(module.INTENTIONALLY_UNUSED_WEIGHT_PREFIXES) + } + + assert len(checkpoint_names) == len(routed) + len(intentionally_unused) + assert parameter_names == set(routed) + assert len(intentionally_unused) == 276 + + # --------------------------------------------------------------------------- # Causal LM weight alignment # ---------------------------------------------------------------------------