Skip to content

Commit 99d3155

Browse files
justinchubyCopilot
andauthored
Add VibeVoice streaming ASR export (#723)
## Implementation - Adds architecture-discriminated `VibeVoiceForASRStreamingTraining` support for both official Microsoft streaming-ASR checkpoints, without changing the existing VibeVoice TTS route: - `microsoft/VibeVoice-ASR-Streaming-1.5B@4262d23d8a539a6530cf64fbd0b1751ef9a30853` - `microsoft/VibeVoice-ASR-Streaming-7B@60d858b518b4e19d404af3737f848fc185b30177` - Exports the executable ASR pipeline as three truthful stages: `audio_encoder` (acoustic and semantic causal tokenizers plus connectors), `embedding` (speech embedding replacement), and `decoder` (Qwen2 with explicit KV caches). - Models both causal convolution cache sets and final-chunk stride padding as ONNX inputs/outputs. Streaming control tokens (`<speech_start>`, `<speech_end>`, forced `<|text_chunk_end|>`), chunk loop/sampling, hotword prompt assembly, and speaker JSON parsing remain host-owned. - Declares the decoder attention mask arbitrary through the generic `mobius.attention.requires_arbitrary_mask` contract, preventing incorrect GQA fusion for left-padded ASR batches. - Routes all 901 executable checkpoint tensors exactly once for each variant and explicitly excludes the 276 `acoustic_tokenizer.decoder.*` training-VAE tensors that are not on the ASR inference path. The 1.5B tied embedding/head checkpoint preserves its explicit `lm_head.weight`. - Emits graph-derived advisory metadata and processor assets, but marks ORT GenAI runtime orchestration unsupported rather than publishing a nonfunctional runnable claim. ## Provenance and runtime boundary - Executable reference: `microsoft/VibeVoice@505653d3873b065a488aea551c6ee3dc51d3062f` with `transformers==4.51.3`; ASR-critical blobs were checked unchanged through upstream `1541f590c7099820f10ea012f48d2399282df69f`. - The dedicated source-parity job force-installs the exact Transformers 4.51 ABI (`huggingface-hub==0.31.4`, `tokenizers==0.21.4`) after the public Transformers 5 environment. - No local real-weight downloads were attempted, so this PR makes no L4/L5 real-transcription claim. `onnxruntime-ep-mlx 0.29.5` / ORT 1.29 cannot execute the complete VibeVoice decoder due to `CumSum cannot infer output shapes`; this blocks local MLX real-runtime evidence. ## Validation - L1 staged graph/cache/role and arbitrary-mask-fusion-exclusion tests. - Full pinned L2 builds and complete 1,177-tensor safetensors-index audits for both variants. - Exact-source synthetic two-chunk prefill/cached-decode parity plus both pinned processor contracts: 24 kHz framing, left-padding, English/Chinese hotwords, and speaker JSON normalization. - Full non-integration suite: `4730 passed, 56 skipped, 1 subtests passed`. - `lintrunner f --output oneline --all-files` and diff checks. - Fresh independent exact-head GPT-5.6 Terra review: no high-confidence findings. --------- Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: d4ef52c4-0459-4efa-bc06-d0ebca2275f2
1 parent 12ed392 commit 99d3155

19 files changed

Lines changed: 1546 additions & 21 deletions

‎docs/model-catalog.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,7 @@ Also registered with `T5ForConditionalGeneration` (task: `seq2seq`):
189189
| `moonshine_streaming` | `MoonshineStreamingForConditionalGeneration` | `speech-to-text` | `moonshine-ai/moonshine-streaming-tiny` |
190190
| `qwen3_asr` | `Qwen3ASRForConditionalGeneration` | `speech-language` | — |
191191
| `qwen3_forced_aligner` | `Qwen3ASRForConditionalGeneration` | `speech-language` | — |
192+
| `VibeVoiceForASRStreamingTraining` | `VibeVoiceASRStreamingForConditionalGeneration` | `vibevoice-asr-streaming` | `microsoft/VibeVoice-ASR-Streaming-1.5B`, `microsoft/VibeVoice-ASR-Streaming-7B` |
192193

193194
### Text-to-Speech
194195

‎src/mobius/_configs/__init__.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@
112112
)
113113
from mobius._configs._world_model import WorldModelConfig
114114
from mobius._configs.vibevoice import (
115+
VibeVoiceASRStreamingConfig,
115116
VibeVoiceConfig,
116117
VibeVoiceDiffusionConfig,
117118
VibeVoiceStreamingConfig,
@@ -184,8 +185,9 @@
184185
"TTSConfig",
185186
"VisionConfig",
186187
"VisionLanguageConfig",
187-
"VibeVoiceConfig",
188188
"VibeVoiceASRConfig",
189+
"VibeVoiceASRStreamingConfig",
190+
"VibeVoiceConfig",
189191
"VibeVoiceDiffusionConfig",
190192
"VibeVoiceStreamingConfig",
191193
"VibeVoiceStreamingDiffusionConfig",

‎src/mobius/_configs/vibevoice.py‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ class VibeVoiceTokenizerConfig:
3131
rms_norm_eps: float = 1e-5
3232
layer_scale_init_value: float = 1e-6
3333
vae_std: float = 0.625
34+
std_dist_type: str = "gaussian"
3435

3536
@property
3637
def hop_length(self) -> int:
@@ -372,3 +373,141 @@ def validate(self) -> None:
372373
raise ValueError(
373374
"VibeVoice Realtime diffusion latent_size must match acoustic vae_dim"
374375
)
376+
377+
378+
def _asr_streaming_tokenizer_config(
379+
config, *, default_hidden_size: int
380+
) -> VibeVoiceTokenizerConfig:
381+
"""Extract the streaming-ASR encoder schema into the shared causal-tokenizer geometry."""
382+
source = _as_attribute_config(config)
383+
encoder_ratios = list(getattr(source, "encoder_ratios", [8, 5, 5, 4, 2, 2]))
384+
required_values = {
385+
"channels": (getattr(source, "channels", 1), 1),
386+
"causal": (getattr(source, "causal", True), True),
387+
"mixer_layer": (getattr(source, "mixer_layer", "depthwise_conv"), "depthwise_conv"),
388+
"conv_norm": (getattr(source, "conv_norm", "none"), "none"),
389+
"pad_mode": (getattr(source, "pad_mode", "constant"), "constant"),
390+
"disable_last_norm": (getattr(source, "disable_last_norm", True), True),
391+
"layernorm": (getattr(source, "layernorm", "RMSNorm"), "RMSNorm"),
392+
"layernorm_elementwise_affine": (
393+
getattr(source, "layernorm_elementwise_affine", True),
394+
True,
395+
),
396+
"conv_bias": (getattr(source, "conv_bias", True), True),
397+
}
398+
unsupported = [
399+
f"{name}={actual!r}"
400+
for name, (actual, expected) in required_values.items()
401+
if actual != expected
402+
]
403+
if unsupported:
404+
raise ValueError(
405+
"Unsupported VibeVoice streaming ASR tokenizer configuration: "
406+
+ ", ".join(unsupported)
407+
+ "."
408+
)
409+
if not encoder_ratios or any(int(ratio) <= 0 for ratio in encoder_ratios):
410+
raise ValueError(
411+
"VibeVoice streaming ASR encoder_ratios must contain positive integers."
412+
)
413+
return VibeVoiceTokenizerConfig(
414+
channels=int(getattr(source, "channels", 1)),
415+
hidden_size=int(getattr(source, "vae_dim", default_hidden_size)),
416+
kernel_size=int(getattr(source, "kernel_size", 7)),
417+
num_filters=int(getattr(source, "encoder_n_filters", 32)),
418+
downsampling_ratios=list(reversed(encoder_ratios)),
419+
depths=[
420+
int(depth)
421+
for depth in str(getattr(source, "encoder_depths", "3-3-3-3-3-3-8")).split("-")
422+
],
423+
ffn_expansion=int(getattr(source, "ffn_expansion", 4)),
424+
hidden_act=str(getattr(source, "hidden_act", "gelu")),
425+
rms_norm_eps=float(getattr(source, "layernorm_eps", 1e-5)),
426+
layer_scale_init_value=float(getattr(source, "layer_scale_init_value", 1e-6)),
427+
vae_std=float(getattr(source, "fix_std", 0.625)),
428+
std_dist_type=str(getattr(source, "std_dist_type", "gaussian")),
429+
)
430+
431+
432+
@dataclasses.dataclass
433+
class VibeVoiceASRStreamingConfig(ArchitectureConfig):
434+
"""Configuration for VibeVoice streaming ASR's causal audio encoders and Qwen2 decoder."""
435+
436+
acoustic_tokenizer: VibeVoiceTokenizerConfig = dataclasses.field(
437+
default_factory=VibeVoiceTokenizerConfig
438+
)
439+
semantic_tokenizer: VibeVoiceTokenizerConfig = dataclasses.field(
440+
default_factory=lambda: VibeVoiceTokenizerConfig(hidden_size=128, std_dist_type="none")
441+
)
442+
speech_start_token_id: int = 151646
443+
speech_end_token_id: int = 151647
444+
speech_placeholder_token_id: int = 151648
445+
text_chunk_end_token_id: int = 151665
446+
compression_ratio: int = 3200
447+
sampling_rate: int = 24_000
448+
chunk_frames: int = 22
449+
lookahead_frames: int = 4
450+
451+
@classmethod
452+
def from_transformers(
453+
cls,
454+
config,
455+
parent_config=None,
456+
*,
457+
allow_block_fp8_dense_fallback: bool = False,
458+
) -> VibeVoiceASRStreamingConfig:
459+
"""Extract the public streaming-ASR composite while preserving Qwen2 decoder config."""
460+
parent = _as_attribute_config(parent_config or config)
461+
decoder = _as_attribute_config(getattr(parent, "decoder_config", config))
462+
result = super().from_transformers(
463+
decoder,
464+
parent_config=parent,
465+
allow_block_fp8_dense_fallback=allow_block_fp8_dense_fallback,
466+
)
467+
acoustic = _asr_streaming_tokenizer_config(
468+
getattr(parent, "acoustic_tokenizer_config", None),
469+
default_hidden_size=64,
470+
)
471+
semantic = _asr_streaming_tokenizer_config(
472+
getattr(parent, "semantic_tokenizer_config", None),
473+
default_hidden_size=128,
474+
)
475+
if acoustic.std_dist_type != "gaussian" or semantic.std_dist_type != "none":
476+
raise ValueError(
477+
"VibeVoice streaming ASR requires a gaussian acoustic tokenizer and "
478+
"a deterministic semantic tokenizer."
479+
)
480+
compression_ratio = int(
481+
getattr(parent, "speech_tok_compress_ratio", acoustic.hop_length)
482+
)
483+
if acoustic.hop_length != compression_ratio:
484+
raise ValueError(
485+
"VibeVoice streaming ASR acoustic tokenizer ratios must match the "
486+
"configured speech tokenizer compression ratio."
487+
)
488+
if semantic.hop_length != acoustic.hop_length:
489+
raise ValueError(
490+
"VibeVoice streaming ASR acoustic and semantic tokenizer hop lengths must match."
491+
)
492+
# The composite ASR config controls the tokenizer and connector precision,
493+
# while decoder_config independently records Qwen's checkpoint storage dtype.
494+
pipeline_dtype = _resolve_dtype_value(getattr(parent, "dtype", None))
495+
if pipeline_dtype is None:
496+
pipeline_dtype = _resolve_dtype_value(getattr(parent, "torch_dtype", None))
497+
return dataclasses.replace(
498+
result,
499+
model_type="vibevoice_asr_streaming",
500+
acoustic_tokenizer=acoustic,
501+
semantic_tokenizer=semantic,
502+
speech_start_token_id=int(getattr(parent, "object_ref_start_token_id", 151646)),
503+
speech_end_token_id=int(getattr(parent, "object_ref_end_token_id", 151647)),
504+
speech_placeholder_token_id=int(getattr(parent, "speech_pad_token_id", 151648)),
505+
text_chunk_end_token_id=int(getattr(parent, "text_chunk_end_token_id", 151665)),
506+
compression_ratio=compression_ratio,
507+
sampling_rate=int(getattr(parent, "target_sample_rate", 24_000)),
508+
chunk_frames=int(getattr(parent, "chunk_frames", 22)),
509+
lookahead_frames=int(getattr(parent, "lookahead_frames", 4)),
510+
eos_token_id=getattr(parent, "eos_token_id", 151643),
511+
pad_token_id=getattr(parent, "pad_token_id", 151655),
512+
dtype=pipeline_dtype or result.dtype,
513+
)

‎src/mobius/_registry.py‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
Qwen4ExpConfig,
4949
SenseNovaU1Config,
5050
VibeVoiceASRConfig,
51+
VibeVoiceASRStreamingConfig,
5152
VibeVoiceConfig,
5253
VibeVoiceStreamingConfig,
5354
WhisperConfig,
@@ -162,6 +163,7 @@
162163
SmolLM3CausalLMModel,
163164
SortformerDiarizationModel,
164165
VibeVoiceASRForConditionalGeneration,
166+
VibeVoiceASRStreamingForConditionalGeneration,
165167
VibeVoiceForConditionalGeneration,
166168
VibeVoiceStreamingForConditionalGeneration,
167169
WhisperForConditionalGeneration,
@@ -231,7 +233,12 @@
231233
from mobius.models.t5 import T5EncoderModel, T5ForConditionalGeneration
232234
from mobius.models.talkie import TalkieForCausalLM
233235
from mobius.models.trocr import TrOCRForConditionalGeneration
234-
from mobius.models.vibevoice import VIBEVOICE_MODEL_ID, VIBEVOICE_REVISION
236+
from mobius.models.vibevoice import (
237+
VIBEVOICE_ASR_STREAMING_MODEL_ID,
238+
VIBEVOICE_ASR_STREAMING_REVISION,
239+
VIBEVOICE_MODEL_ID,
240+
VIBEVOICE_REVISION,
241+
)
235242
from mobius.models.vibevoice_asr import VIBEVOICE_ASR_MODEL_ID, VIBEVOICE_ASR_REVISION
236243
from mobius.models.vibevoice_streaming import (
237244
VIBEVOICE_STREAMING_MODEL_ID,
@@ -941,6 +948,15 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
941948
family="vibevoice",
942949
variant="offline-asr",
943950
),
951+
"VibeVoiceForASRStreamingTraining": ModelRegistration(
952+
VibeVoiceASRStreamingForConditionalGeneration,
953+
task="vibevoice-asr-streaming",
954+
config_class=VibeVoiceASRStreamingConfig,
955+
test_model_id=VIBEVOICE_ASR_STREAMING_MODEL_ID,
956+
test_revision=VIBEVOICE_ASR_STREAMING_REVISION,
957+
family="vibevoice",
958+
variant="streaming-asr",
959+
),
944960
"whisper": ModelRegistration(
945961
WhisperForConditionalGeneration,
946962
task="speech-to-text",

‎src/mobius/integrations/onnx_genai/auto_export.py‎

Lines changed: 56 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -769,7 +769,7 @@ def _looks_like_vibevoice_tts(pkg: Any) -> bool:
769769
} <= names
770770

771771

772-
def _looks_like_vibevoice_asr(pkg: Any) -> bool:
772+
def _looks_like_vibevoice_offline_asr(pkg: Any) -> bool:
773773
"""Detect the offline VibeVoice-ASR dual-encoder component topology."""
774774
try:
775775
names = set(pkg.keys())
@@ -784,6 +784,23 @@ def _looks_like_vibevoice_asr(pkg: Any) -> bool:
784784
} <= names
785785

786786

787+
def _looks_like_vibevoice_streaming_asr(pkg: Any) -> bool:
788+
"""Detect the stateful three-stage VibeVoice streaming-ASR topology."""
789+
try:
790+
if set(pkg.keys()) != {"audio_encoder", "embedding", "decoder"}:
791+
return False
792+
audio_inputs = {value.name for value in pkg["audio_encoder"].graph.inputs}
793+
except AttributeError:
794+
return False
795+
return {
796+
"speech_tensors",
797+
"speech_masks",
798+
"is_final_chunk",
799+
"past_acoustic_conv.0",
800+
"past_semantic_conv.0",
801+
} <= audio_inputs
802+
803+
787804
def _write_vibevoice_asr_processor_contract(output_dir: str, config: Any) -> str:
788805
"""Materialize the source processor protocol absent from the ASR checkpoint."""
789806
os.makedirs(output_dir, exist_ok=True)
@@ -1039,6 +1056,43 @@ def write_onnx_genai_config(
10391056
"runtime configuration is claimed."
10401057
)
10411058
return _write_advisory_component_contract(pkg, output_dir, warning=warning)
1059+
if _looks_like_vibevoice_streaming_asr(pkg):
1060+
if kv_native_dtype is not None:
1061+
raise ValueError(
1062+
"workflow VibeVoice streaming ASR export derives KV and convolution state "
1063+
"dtypes from ONNX ports; kv_native_dtype overrides are unsupported"
1064+
)
1065+
artifacts = _write_text_runtime_assets(output_dir, source, revision=revision)
1066+
artifacts.update(
1067+
_copy_runtime_assets(
1068+
output_dir,
1069+
source,
1070+
(
1071+
"processor_config.json",
1072+
"preprocessor_config.json",
1073+
"generation_config.json",
1074+
),
1075+
revision=revision,
1076+
)
1077+
)
1078+
audio_processor_path = _write_hf_audio_processor(output_dir, source, revision=revision)
1079+
if audio_processor_path is not None:
1080+
artifacts["audio_processor"] = audio_processor_path
1081+
artifacts.update(
1082+
_write_advisory_component_contract(
1083+
pkg,
1084+
output_dir,
1085+
warning=(
1086+
"The tested onnx-genai runtime cannot orchestrate VibeVoice streaming "
1087+
"ASR's dual causal convolution states, flattened speech-placeholder "
1088+
"replacement with arbitrary left-padded attention masks, forced "
1089+
"<|text_chunk_end|> control tokens, or host-side hotword and speaker "
1090+
"JSON handling. Exact graph and pinned processor contracts are exported "
1091+
"without claiming downstream runtime support."
1092+
),
1093+
)
1094+
)
1095+
return artifacts
10421096
decoder = pkg.get("decoder") or pkg.get("model")
10431097
decoder_inputs = (
10441098
{value.name for value in decoder.graph.inputs} if decoder is not None else set()
@@ -1376,7 +1430,7 @@ def write_onnx_genai_config(
13761430
)
13771431
return artifacts
13781432

1379-
if _looks_like_vibevoice_asr(pkg):
1433+
if _looks_like_vibevoice_offline_asr(pkg):
13801434
if kv_native_dtype is not None:
13811435
raise ValueError(
13821436
"VibeVoice-ASR derives KV and convolution state dtypes from ONNX ports; "

‎src/mobius/integrations/onnx_genai/auto_export_test.py‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,77 @@ def test_vibevoice_asr_writes_processor_and_advisory_contract(monkeypatch, tmp_p
178178
assert compatibility["runtime_validation_status"] == "unsupported-by-tested-runtime"
179179

180180

181+
@pytest.mark.parametrize(
182+
("model_id", "revision"),
183+
(
184+
("microsoft/VibeVoice-ASR-Streaming-1.5B", "4262d23d8a539a6530cf64fbd0b1751ef9a30853"),
185+
("microsoft/VibeVoice-ASR-Streaming-7B", "60d858b518b4e19d404af3737f848fc185b30177"),
186+
),
187+
)
188+
def test_vibevoice_streaming_asr_writes_pinned_advisory_contract(
189+
monkeypatch, tmp_path, model_id, revision
190+
):
191+
from mobius.integrations.onnx_genai import auto_export
192+
from mobius.models.vibevoice import (
193+
VibeVoiceASRStreamingForConditionalGeneration,
194+
)
195+
from mobius.models.vibevoice_asr_test import _streaming_config
196+
from mobius.tasks import VibeVoiceASRStreamingTask
197+
198+
config = _streaming_config()
199+
package = VibeVoiceASRStreamingTask().build(
200+
VibeVoiceASRStreamingForConditionalGeneration(config), config
201+
)
202+
calls: list[tuple[str, str | None]] = []
203+
204+
def text_assets(output_dir, source, *, revision=None):
205+
calls.append(("text", revision))
206+
return {"tokenizer": str(Path(output_dir) / "tokenizer.json")}
207+
208+
def runtime_assets(output_dir, source, names, *, revision=None):
209+
calls.append(("runtime", revision))
210+
assert names == (
211+
"processor_config.json",
212+
"preprocessor_config.json",
213+
"generation_config.json",
214+
)
215+
return {"processor_config": str(Path(output_dir) / "processor_config.json")}
216+
217+
def audio_processor(output_dir, source, *, revision=None):
218+
calls.append(("audio", revision))
219+
return str(Path(output_dir) / "audio_processor.json")
220+
221+
monkeypatch.setattr(auto_export, "_write_text_runtime_assets", text_assets)
222+
monkeypatch.setattr(auto_export, "_copy_runtime_assets", runtime_assets)
223+
monkeypatch.setattr(auto_export, "_write_hf_audio_processor", audio_processor)
224+
monkeypatch.setattr(
225+
auto_export,
226+
"_write_advisory_component_contract",
227+
lambda *args, **kwargs: {
228+
"inference_metadata": str(tmp_path / "inference_metadata.yaml")
229+
},
230+
)
231+
232+
artifacts = write_onnx_genai_config(
233+
package,
234+
str(tmp_path),
235+
source=model_id,
236+
revision=revision,
237+
)
238+
239+
assert calls == [
240+
("text", revision),
241+
("runtime", revision),
242+
("audio", revision),
243+
]
244+
assert {
245+
"tokenizer",
246+
"processor_config",
247+
"audio_processor",
248+
"inference_metadata",
249+
} <= set(artifacts)
250+
251+
181252
def _video_diffusion_package() -> ModelPackage:
182253
latent = ["batch", "frames", 4, "height", "width"]
183254
transformer = _model(

0 commit comments

Comments
 (0)