diff --git a/docs/model-catalog.md b/docs/model-catalog.md index 94c983b54..0534df209 100644 --- a/docs/model-catalog.md +++ b/docs/model-catalog.md @@ -196,7 +196,13 @@ Also registered with `T5ForConditionalGeneration` (task: `seq2seq`): |---|---|---|---| | `qwen3_tts` | `Qwen3TTSForConditionalGeneration` | `tts` | — | | `qwen3_tts_tokenizer_12hz` | `Qwen3TTSTokenizerV2Model` | `codec` | — | -| `vibevoice` | `VibeVoiceForConditionalGeneration` | `vibevoice-tts` | `vibevoice/VibeVoice-1.5B-hf` | +| `vibevoice` | `VibeVoiceForConditionalGeneration` | `vibevoice-tts` | `microsoft/VibeVoice-1.5B` | + +`microsoft/VibeVoice-1.5B@c00898d` is tested with its official weights. Its legacy +release has no Transformers-native tokenizer or processor assets, so Mobius pins only +those executable sidecars to `vibevoice/VibeVoice-1.5B-hf@edc39f8`; it never substitutes +mirror weights. The Realtime and ASR entries in the Microsoft VibeVoice collection are +explicitly rejected until their distinct streaming or ASR architectures are implemented. ### Audio Feature Extraction diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 8a5a86a78..a8eec40a2 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -328,6 +328,14 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: from mobius.models.reuse import REUSE_REVISION revision = REUSE_REVISION + if args.model: + from mobius.models.vibevoice import resolve_vibevoice_sources + + vibevoice_sources = resolve_vibevoice_sources(args.model, revision) + if vibevoice_sources is not None: + # Pin the early Diffusers probe to the official checkpoint. The + # builder separately resolves its modern executable sidecars. + revision = vibevoice_sources.weight_revision output_dir = args.output_dir os.makedirs(output_dir, exist_ok=True) dtype_override = resolve_dtype(args.dtype) @@ -616,6 +624,30 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: _save_package(pkg, output_dir, args, optimize, component_filter) +def _runtime_asset_source( + pkg, + source: str | None, + explicit_revision: str | None, +) -> tuple[str | None, str | None]: + """Select a pinned processor source without changing checkpoint provenance.""" + processor_sources = { + processor_source + for model in pkg.values() + if ( + processor_source := getattr(model, "metadata_props", {}).get( + "mobius.processor_source" + ) + ) + is not None + } + if len(processor_sources) == 1: + processor_source = processor_sources.pop() + model_id, separator, revision = processor_source.rpartition("@") + if separator and model_id and revision: + return model_id, revision + return source, _runtime_source_revision(pkg, explicit_revision) + + def _runtime_source_revision(pkg, explicit_revision: str | None) -> str | None: """Recover the effective build revision for runtime asset downloads.""" if explicit_revision is not None: @@ -675,15 +707,15 @@ def _save_package( if runtime == "ort-genai": from mobius.integrations.ort_genai import write_ort_genai_config - hf_model_id = getattr(args, "model", None) + hf_model_id, runtime_revision = _runtime_asset_source( + pkg, + getattr(args, "model", None), + getattr(args, "revision", None), + ) ep = getattr(args, "execution_provider", "cpu") # When --config (local dir) is used instead of --model, copy tokenizer # files from the local directory rather than downloading from HF. local_config_dir = getattr(args, "config", None) - runtime_revision = _runtime_source_revision( - pkg, - getattr(args, "revision", None), - ) artifacts = write_ort_genai_config( pkg, output_dir, @@ -705,8 +737,11 @@ def _save_package( ) config = getattr(pkg, "config", None) - source = getattr(args, "config", None) or getattr(args, "model", None) - revision = _runtime_source_revision(pkg, getattr(args, "revision", None)) + source, revision = _runtime_asset_source( + pkg, + getattr(args, "config", None) or getattr(args, "model", None), + getattr(args, "revision", None), + ) if is_native_vlm_package(pkg): try: artifacts = write_native_vlm_package_metadata( diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index d45495a5b..394fdf37b 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -300,15 +300,18 @@ def build_transformers_model( _config_from_hf, _default_task_for_model, ) + from mobius.models.vibevoice import resolve_vibevoice_sources detection_revision = revision - if model_id == "vibevoice/VibeVoice-1.5B-hf" and detection_revision is None: - from mobius.models.vibevoice import VIBEVOICE_REVISION - - # 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 + vibevoice_sources = resolve_vibevoice_sources(model_id, revision) + config_model_id = model_id + if vibevoice_sources is not None: + # The original official checkpoint has legacy metadata and no tokenizer + # assets. Resolve only executable inputs from the pinned conversion; + # model_id remains the graph identity and official weights stay official. + revision = vibevoice_sources.weight_revision + config_model_id = vibevoice_sources.config_model_id + detection_revision = vibevoice_sources.config_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 @@ -318,7 +321,7 @@ def build_transformers_model( detection_revision = REUSE_REVISION hf_config, loaded_from_raw_json = _load_transformers_config( - model_id, + config_model_id, revision=detection_revision, trust_remote_code=trust_remote_code, ) @@ -527,6 +530,13 @@ def build_transformers_model( model.graph.name = f"{model_id}/{name}" if model_type in _QWEN4_MODEL_TYPES | {"vibevoice"}: model.metadata_props["mobius.source_revision"] = revision or "unpinned" + if vibevoice_sources is not None: + model.metadata_props["mobius.executable_source"] = ( + f"{vibevoice_sources.config_model_id}@{vibevoice_sources.config_revision}" + ) + model.metadata_props["mobius.processor_source"] = ( + f"{vibevoice_sources.processor_model_id}@{vibevoice_sources.processor_revision}" + ) if load_weights: _reject_unsupported_affine_qwen4(model_type, config) @@ -586,7 +596,13 @@ def build_transformers_model( else: state_dict = _download_weights(model_id, revision=revision) if hasattr(model_module, "preprocess_weights"): - state_dict = model_module.preprocess_weights(state_dict) + if vibevoice_sources is not None: + state_dict = model_module.preprocess_weights( + state_dict, + checkpoint_layout=vibevoice_sources.weight_layout, + ) + else: + state_dict = model_module.preprocess_weights(state_dict) state_dict = preprocess_component_quantized_state_dict( state_dict, model_module, diff --git a/src/mobius/integrations/transformers/_builder_test.py b/src/mobius/integrations/transformers/_builder_test.py index 85a72542d..a7cb66424 100644 --- a/src/mobius/integrations/transformers/_builder_test.py +++ b/src/mobius/integrations/transformers/_builder_test.py @@ -10,6 +10,7 @@ import onnx_ir as ir import pytest +import torch from onnxscript import nn from mobius._configs import QuantizationConfig, QuantizationOverride @@ -241,7 +242,11 @@ def stop_after_config(model_id, **kwargs): def test_vibevoice_none_revision_pins_first_config_probe(monkeypatch) -> None: - from mobius.models.vibevoice import VIBEVOICE_MODEL_ID, VIBEVOICE_REVISION + from mobius.models.vibevoice import ( + VIBEVOICE_EXECUTABLE_MODEL_ID, + VIBEVOICE_EXECUTABLE_REVISION, + VIBEVOICE_MODEL_ID, + ) calls = [] @@ -263,9 +268,9 @@ def stop_after_config(model_id, **kwargs): assert calls == [ ( - VIBEVOICE_MODEL_ID, + VIBEVOICE_EXECUTABLE_MODEL_ID, { - "revision": VIBEVOICE_REVISION, + "revision": VIBEVOICE_EXECUTABLE_REVISION, "trust_remote_code": False, }, ) @@ -1009,6 +1014,97 @@ def fake_build_diffusers(*args, **kwargs): ] +def test_official_vibevoice_uses_pinned_sidecars_and_official_weights(monkeypatch) -> None: + """The official legacy config must never replace the requested checkpoint.""" + from mobius.models.vibevoice import ( + VIBEVOICE_EXECUTABLE_MODEL_ID, + VIBEVOICE_EXECUTABLE_REVISION, + VIBEVOICE_MODEL_ID, + VIBEVOICE_REVISION, + ) + + parent = SimpleNamespace(model_type="vibevoice", architectures=[]) + package = ModelPackage( + { + "audio_encoder": ir.Model( + ir.Graph([], [], nodes=[], name="audio_encoder"), ir_version=11 + ) + } + ) + calls = [] + + class OfficialVibeVoiceModule(_DummyModule): + def preprocess_weights(self, state_dict, *, checkpoint_layout): + calls.append(("preprocess", state_dict, checkpoint_layout)) + return state_dict + + monkeypatch.setattr( + transformers_builder, + "_load_transformers_config", + lambda *args, **kwargs: calls.append(("config", args, kwargs)) or (parent, False), + ) + monkeypatch.setattr( + transformers_builder, + "_select_primary_config", + lambda value: (value, value, "vibevoice"), + ) + monkeypatch.setattr( + transformers_builder, + "_resolve_module_class", + lambda *args, **kwargs: (OfficialVibeVoiceModule, "vibevoice-tts", "vibevoice"), + ) + monkeypatch.setattr( + _config_resolver, + "_config_from_hf", + lambda *args, **kwargs: make_config(model_type="vibevoice"), + ) + monkeypatch.setattr( + transformers_builder, "build_from_module", lambda *args, **kwargs: package + ) + state_dict = {"weight": torch.ones(())} + monkeypatch.setattr( + transformers_builder, + "_download_weights", + lambda *args, **kwargs: calls.append(("weights", args, kwargs)) or state_dict, + ) + + result = transformers_builder.build_transformers_model(VIBEVOICE_MODEL_ID) + + assert result is package + assert calls == [ + ( + "config", + (VIBEVOICE_EXECUTABLE_MODEL_ID,), + {"revision": VIBEVOICE_EXECUTABLE_REVISION, "trust_remote_code": False}, + ), + ("weights", (VIBEVOICE_MODEL_ID,), {"revision": VIBEVOICE_REVISION}), + ("preprocess", state_dict, "official"), + ] + assert package["audio_encoder"].metadata_props == { + "mobius.source_revision": VIBEVOICE_REVISION, + "mobius.executable_source": f"{VIBEVOICE_EXECUTABLE_MODEL_ID}@{VIBEVOICE_EXECUTABLE_REVISION}", + "mobius.processor_source": f"{VIBEVOICE_EXECUTABLE_MODEL_ID}@{VIBEVOICE_EXECUTABLE_REVISION}", + } + + +@pytest.mark.parametrize( + "model_id", + [ + "microsoft/VibeVoice-Realtime-0.5B", + "microsoft/VibeVoice-ASR", + "microsoft/VibeVoice-ASR-Streaming-7B", + "microsoft/VibeVoice-ASR-Streaming-1.5B", + "microsoft/VibeVoice-ASR-BitNet", + "microsoft/VibeVoice-ASR-HF", + "microsoft/VibeVoice-AcousticTokenizer", + "Microsoft/VibeVoice-ASR", + ], +) +def test_unimplemented_official_vibevoice_collection_ids_fail_closed(model_id): + with pytest.raises(NotImplementedError, match="unsupported"): + transformers_builder.build_transformers_model(model_id, load_weights=False) + + def test_glm_full_attention_rejects_diffusers_dispatch(monkeypatch) -> None: """``--glm-full-attention`` must raise on the diffusers-dispatch branch. diff --git a/src/mobius/models/vibevoice.py b/src/mobius/models/vibevoice.py index 3841aaab4..fc7a0db70 100644 --- a/src/mobius/models/vibevoice.py +++ b/src/mobius/models/vibevoice.py @@ -10,7 +10,9 @@ from __future__ import annotations +import dataclasses import math +import re from typing import TYPE_CHECKING, ClassVar import numpy as np @@ -37,9 +39,207 @@ from collections.abc import Sequence -VIBEVOICE_MODEL_ID = "vibevoice/VibeVoice-1.5B-hf" -VIBEVOICE_REVISION = "edc39f80f5cae656da37baf8faa8f5502bf7081f" -VIBEVOICE_MICROSOFT_PROVENANCE_REVISION = "c00898d257e6b46004e3e2866a47534085fb685a" +VIBEVOICE_MODEL_ID = "microsoft/VibeVoice-1.5B" +VIBEVOICE_REVISION = "c00898d257e6b46004e3e2866a47534085fb685a" +# TODO(#727): Switch to Microsoft's pinned HF-native sidecars once published. +VIBEVOICE_EXECUTABLE_MODEL_ID = "vibevoice/VibeVoice-1.5B-hf" +VIBEVOICE_EXECUTABLE_REVISION = "edc39f80f5cae656da37baf8faa8f5502bf7081f" +VIBEVOICE_MICROSOFT_PROVENANCE_REVISION = VIBEVOICE_REVISION + + +@dataclasses.dataclass(frozen=True) +class VibeVoiceSources: + """Immutable provenance for one VibeVoice TTS build. + + ``model_id`` and ``weight_revision`` are always the user's checkpoint. The + official 1.5B release predates Transformers-native VibeVoice metadata, so + its executable config and processor are resolved from the pinned conversion + mirror while its official weights remain the only downloaded weights. + """ + + model_id: str + weight_revision: str + config_model_id: str + config_revision: str + processor_model_id: str + processor_revision: str + weight_layout: str + + +_UNSUPPORTED_VIBEVOICE_MODELS = { + "microsoft/VibeVoice-Realtime-0.5B": ( + "VibeVoice Realtime requires its streaming backbone and scheduler, " + "which Mobius does not export yet." + ), + "microsoft/VibeVoice-ASR": ( + "VibeVoice ASR requires the VibeVoice-ASR encoder-decoder task, " + "which Mobius does not export yet." + ), + "microsoft/VibeVoice-ASR-Streaming-7B": ( + "VibeVoice ASR Streaming requires the VibeVoice-ASR streaming task, " + "which Mobius does not export yet." + ), + "microsoft/VibeVoice-ASR-Streaming-1.5B": ( + "VibeVoice ASR Streaming requires the VibeVoice-ASR streaming task, " + "which Mobius does not export yet." + ), + "microsoft/VibeVoice-ASR-BitNet": ( + "VibeVoice ASR BitNet requires the VibeVoice-ASR task and BitNet " + "weight loader, which Mobius does not export yet." + ), + "microsoft/VibeVoice-ASR-HF": ( + "VibeVoice ASR requires the VibeVoice-ASR encoder-decoder task, " + "which Mobius does not export yet." + ), + "microsoft/VibeVoice-AcousticTokenizer": ( + "VibeVoice Acoustic Tokenizer requires a standalone codec task, " + "which Mobius does not export yet." + ), +} + + +def resolve_vibevoice_sources(model_id: str, revision: str | None) -> VibeVoiceSources | None: + """Resolve pinned config, processor, and weight sources for supported VibeVoice IDs. + + This fail-closed resolver separates executable dependencies from checkpoint + provenance. It recognizes the current official collection entries so their + shared ``model_type="vibevoice"`` cannot accidentally route ASR weights + into the TTS graph. + """ + canonical_model_id = model_id.casefold() + unsupported = { + known_model_id.casefold(): reason + for known_model_id, reason in _UNSUPPORTED_VIBEVOICE_MODELS.items() + } + if canonical_model_id in unsupported: + raise NotImplementedError( + f"{model_id} is unsupported: {unsupported[canonical_model_id]}" + ) + if canonical_model_id == VIBEVOICE_MODEL_ID.casefold(): + if revision not in {None, VIBEVOICE_REVISION}: + raise ValueError( + f"{model_id} is only verified at revision {VIBEVOICE_REVISION}; " + f"got {revision}. Refusing to pair it with a different executable dependency." + ) + return VibeVoiceSources( + model_id=model_id, + weight_revision=VIBEVOICE_REVISION, + config_model_id=VIBEVOICE_EXECUTABLE_MODEL_ID, + config_revision=VIBEVOICE_EXECUTABLE_REVISION, + processor_model_id=VIBEVOICE_EXECUTABLE_MODEL_ID, + processor_revision=VIBEVOICE_EXECUTABLE_REVISION, + weight_layout="official", + ) + if canonical_model_id == VIBEVOICE_EXECUTABLE_MODEL_ID.casefold(): + if revision not in {None, VIBEVOICE_EXECUTABLE_REVISION}: + raise ValueError( + f"{model_id} is only verified at revision {VIBEVOICE_EXECUTABLE_REVISION}; " + f"got {revision}." + ) + return VibeVoiceSources( + model_id=model_id, + weight_revision=VIBEVOICE_EXECUTABLE_REVISION, + config_model_id=model_id, + config_revision=VIBEVOICE_EXECUTABLE_REVISION, + processor_model_id=model_id, + processor_revision=VIBEVOICE_EXECUTABLE_REVISION, + weight_layout="transformers", + ) + return None + + +_OFFICIAL_WEIGHT_NAME_MAPPING = ( + ( + r"semantic_tokenizer\.encoder\.downsample_layers\.0\.0\.conv\.", + r"semantic_tokenizer_encoder.stem.conv.conv.", + ), + (r"semantic_tokenizer\.encoder\.stages\.0\.", r"semantic_tokenizer_encoder.stem.stage."), + ( + r"semantic_tokenizer\.encoder\.downsample_layers\.(\d+)\.0\.conv\.", + r"semantic_tokenizer_encoder.conv_layers.PLACEHOLDER.conv.conv.", + ), + ( + r"semantic_tokenizer\.encoder\.stages\.(\d+)\.", + r"semantic_tokenizer_encoder.conv_layers.PLACEHOLDER.stage.", + ), + (r"semantic_tokenizer\.encoder\.head\.conv\.", r"semantic_tokenizer_encoder.head."), + ( + r"acoustic_tokenizer\.encoder\.downsample_layers\.0\.0\.conv\.", + r"audio_tower.encoder.stem.conv.conv.", + ), + (r"acoustic_tokenizer\.encoder\.stages\.0\.", r"audio_tower.encoder.stem.stage."), + ( + r"acoustic_tokenizer\.encoder\.downsample_layers\.(\d+)\.0\.conv\.", + r"audio_tower.encoder.conv_layers.PLACEHOLDER.conv.conv.", + ), + ( + r"acoustic_tokenizer\.encoder\.stages\.(\d+)\.", + r"audio_tower.encoder.conv_layers.PLACEHOLDER.stage.", + ), + (r"acoustic_tokenizer\.encoder\.head\.conv\.", r"audio_tower.encoder.head."), + ( + r"acoustic_tokenizer\.decoder\.upsample_layers\.0\.0\.conv\.conv\.", + r"audio_tower.decoder.stem.conv.conv.", + ), + (r"acoustic_tokenizer\.decoder\.stages\.0\.", r"audio_tower.decoder.stem.stage."), + ( + r"acoustic_tokenizer\.decoder\.upsample_layers\.(\d+)\.0\.convtr\.convtr\.", + r"audio_tower.decoder.conv_layers.PLACEHOLDER.convtr.convtr.", + ), + ( + r"acoustic_tokenizer\.decoder\.stages\.(\d+)\.", + r"audio_tower.decoder.conv_layers.PLACEHOLDER.stage.", + ), + (r"acoustic_tokenizer\.decoder\.head\.conv\.", r"audio_tower.decoder.head."), + (r"acoustic_tokenizer\.", r"audio_tower."), + (r"prediction_head\.t_embedder\.mlp\.0\.", r"diffusion_head.timestep_proj.fc1."), + (r"prediction_head\.t_embedder\.mlp\.2\.", r"diffusion_head.timestep_proj.fc2."), + ( + r"prediction_head\.layers\.(\d+)\.adaLN_modulation\.1\.", + r"diffusion_head.layers.\1.linear.", + ), + ( + r"prediction_head\.final_layer\.adaLN_modulation\.1\.", + r"diffusion_head.final_layer.linear_1.", + ), + (r"prediction_head\.final_layer\.linear\.", r"diffusion_head.final_layer.linear_2."), + (r"prediction_head\.", r"diffusion_head."), + (r"acoustic_connector\.fc1\.", r"multi_modal_projector.linear_1."), + (r"acoustic_connector\.norm\.", r"multi_modal_projector.act."), + (r"acoustic_connector\.fc2\.", r"multi_modal_projector.linear_2."), + (r"semantic_connector\.fc1\.", r"semantic_connector.linear_1."), + (r"semantic_connector\.norm\.", r"semantic_connector.act."), + (r"semantic_connector\.fc2\.", r"semantic_connector.linear_2."), + (r"^model\.speech_scaling_factor", r"model.latent_scaling_factor"), + (r"^model\.speech_bias_factor", r"model.latent_bias_factor"), + (r"mixer\.conv\.conv\.conv\.", r"mixer.conv."), + (r"\.conv\.conv\.conv\.", r".conv.conv."), +) + + +def _transform_official_weight_name(name: str) -> str: + """Map one original Microsoft checkpoint key to the pinned HF-native layout.""" + result = name + for pattern, replacement in _OFFICIAL_WEIGHT_NAME_MAPPING: + match = re.search(pattern, result) + if match: + if "PLACEHOLDER" in replacement: + replacement = replacement.replace("PLACEHOLDER", str(int(match.group(1)) - 1)) + result = re.sub(pattern, replacement, result) + return result + + +def _convert_official_weights(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Convert the original Microsoft key layout with collision protection.""" + converted: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + converted_key = _transform_official_weight_name(key) + if converted_key in converted: + raise ValueError( + f"Official VibeVoice weight conversion maps multiple tensors to {converted_key!r}." + ) + converted[converted_key] = value + return converted class _CacheAllocator: @@ -904,8 +1104,14 @@ def forward(self, op: OpBuilder, *args, **kwargs): def preprocess_weights( self, state_dict: dict[str, torch.Tensor], + *, + checkpoint_layout: str = "transformers", ) -> dict[str, torch.Tensor]: - """Route the native HF composite checkpoint to standardized package stages.""" + """Route an official or Transformers-native checkpoint to package stages.""" + if checkpoint_layout == "official": + state_dict = _convert_official_weights(state_dict) + elif checkpoint_layout != "transformers": + raise ValueError(f"Unknown VibeVoice checkpoint layout: {checkpoint_layout!r}") routed: dict[str, torch.Tensor] = {} stage_prefixes = tuple(f"{name}." for name in self.HF_COMPONENT_SOURCES) for key, value in state_dict.items(): diff --git a/src/mobius/models/vibevoice_test.py b/src/mobius/models/vibevoice_test.py index c1cab6c46..a93eb675f 100644 --- a/src/mobius/models/vibevoice_test.py +++ b/src/mobius/models/vibevoice_test.py @@ -26,9 +26,14 @@ from mobius._registry import registry from mobius._testing.ort_inference import OnnxModelSession from mobius.models.vibevoice import ( + VIBEVOICE_EXECUTABLE_MODEL_ID, + VIBEVOICE_EXECUTABLE_REVISION, VIBEVOICE_MODEL_ID, VIBEVOICE_REVISION, VibeVoiceForConditionalGeneration, + _convert_official_weights, + _transform_official_weight_name, + resolve_vibevoice_sources, ) from mobius.tasks import VibeVoiceTask @@ -94,6 +99,83 @@ def _make_tiny_hf_config(): ) +def test_official_vibevoice_sources_preserve_weight_identity(): + """The legacy official release uses only pinned executable sidecar assets.""" + sources = resolve_vibevoice_sources(VIBEVOICE_MODEL_ID, VIBEVOICE_REVISION) + + assert sources is not None + assert sources.model_id == VIBEVOICE_MODEL_ID + assert sources.weight_revision == VIBEVOICE_REVISION + assert sources.config_model_id == VIBEVOICE_EXECUTABLE_MODEL_ID + assert sources.config_revision == VIBEVOICE_EXECUTABLE_REVISION + assert sources.processor_model_id == VIBEVOICE_EXECUTABLE_MODEL_ID + assert sources.processor_revision == VIBEVOICE_EXECUTABLE_REVISION + assert sources.weight_layout == "official" + + +def test_official_vibevoice_rejects_unverified_revisions(): + with pytest.raises(ValueError, match="only verified"): + resolve_vibevoice_sources(VIBEVOICE_MODEL_ID, "different-official-revision") + + +def test_official_vibevoice_sources_normalize_hub_id_case(): + sources = resolve_vibevoice_sources("Microsoft/VibeVoice-1.5B", None) + + assert sources is not None + assert sources.model_id == "Microsoft/VibeVoice-1.5B" + with pytest.raises(NotImplementedError, match="unsupported"): + resolve_vibevoice_sources("Microsoft/VibeVoice-ASR", None) + + +@pytest.mark.parametrize( + ("source", "expected"), + [ + ( + "model.acoustic_tokenizer.encoder.downsample_layers.1.0.conv.conv.weight", + "model.audio_tower.encoder.conv_layers.0.conv.conv.weight", + ), + ( + "model.acoustic_tokenizer.decoder.upsample_layers.3.0.convtr.convtr.bias", + "model.audio_tower.decoder.conv_layers.2.convtr.convtr.bias", + ), + ( + "model.semantic_tokenizer.encoder.stages.2.1.mixer.conv.conv.conv.weight", + "model.semantic_tokenizer_encoder.conv_layers.1.stage.1.mixer.conv.weight", + ), + ( + "model.prediction_head.layers.3.adaLN_modulation.1.weight", + "model.diffusion_head.layers.3.linear.weight", + ), + ( + "model.acoustic_connector.fc1.weight", + "model.multi_modal_projector.linear_1.weight", + ), + ("model.speech_scaling_factor", "model.latent_scaling_factor"), + ], +) +def test_official_vibevoice_weight_names_match_transformers_conversion(source, expected): + """Keep every legacy layout rename synchronized with the upstream converter.""" + assert _transform_official_weight_name(source) == expected + + +def test_official_vibevoice_weight_conversion_preserves_all_entries(): + """Reject collisions rather than silently dropping checkpoint tensors.""" + source = { + "model.acoustic_connector.fc1.weight": torch.ones((2, 2)), + "model.semantic_connector.fc1.weight": torch.ones((2, 2)), + "model.speech_bias_factor": torch.ones(()), + } + + converted = _convert_official_weights(source) + + assert len(converted) == len(source) + assert set(converted) == { + "model.multi_modal_projector.linear_1.weight", + "model.semantic_connector.linear_1.weight", + "model.latent_bias_factor", + } + + def _make_tiny_models(): modeling = pytest.importorskip("transformers.models.vibevoice.modeling_vibevoice") torch.manual_seed(7) diff --git a/tests/cli_test.py b/tests/cli_test.py index 4281ca7ac..7ff9133e3 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -162,6 +162,42 @@ def test_standard_build_dispatches_personaplex_through_public_build(self): assert build_model.call_args.kwargs["dtype"] == ir.DataType.FLOAT save_package.assert_called_once() + def test_official_vibevoice_build_keeps_the_requested_model_identity(self): + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch("mobius.__main__.build", return_value=mock.MagicMock()) as build_model, + mock.patch( + "mobius.integrations.diffusers._builder._load_diffusers_pipeline_index", + return_value=None, + ) as pipeline_probe, + mock.patch("mobius.__main__._save_package"), + ): + main(["build", "--model", "microsoft/VibeVoice-1.5B", tmpdir, "--no-weights"]) + + assert build_model.call_args.args == ("microsoft/VibeVoice-1.5B",) + assert ( + build_model.call_args.kwargs["revision"] + == "c00898d257e6b46004e3e2866a47534085fb685a" + ) + assert build_model.call_args.kwargs["load_weights"] is False + assert pipeline_probe.call_args.kwargs["revision"] == ( + "c00898d257e6b46004e3e2866a47534085fb685a" + ) + + def test_runtime_assets_use_vibevoice_processor_source(self): + from mobius.__main__ import _runtime_asset_source + from mobius._model_package import ModelPackage + + model = ir.Model(ir.Graph([], [], nodes=[], name="vibevoice"), ir_version=11) + model.metadata_props["mobius.source_revision"] = "official-revision" + model.metadata_props["mobius.processor_source"] = "vibevoice/processor@native-revision" + + assert _runtime_asset_source( + ModelPackage({"model": model}), + "microsoft/VibeVoice-1.5B", + "official-revision", + ) == ("vibevoice/processor", "native-revision") + def test_local_personaplex_config_bypasses_transformers(self): with ( tempfile.TemporaryDirectory() as checkpoint, diff --git a/tests/vibevoice_golden_test.py b/tests/vibevoice_golden_test.py index a1bb632d1..0dfeca5c2 100644 --- a/tests/vibevoice_golden_test.py +++ b/tests/vibevoice_golden_test.py @@ -27,9 +27,12 @@ ) from mobius.integrations.transformers._config_resolver import _config_from_hf from mobius.models.vibevoice import ( + VIBEVOICE_EXECUTABLE_MODEL_ID, + VIBEVOICE_EXECUTABLE_REVISION, VIBEVOICE_MODEL_ID, VIBEVOICE_REVISION, VibeVoiceForConditionalGeneration, + resolve_vibevoice_sources, ) _ROOT = Path(__file__).parents[1] @@ -62,9 +65,11 @@ def _selected_cases(cases: list[str]) -> list[str]: def _config(): + sources = resolve_vibevoice_sources(VIBEVOICE_MODEL_ID, VIBEVOICE_REVISION) + assert sources is not None hf_config, _ = _load_transformers_config( - VIBEVOICE_MODEL_ID, - revision=VIBEVOICE_REVISION, + sources.config_model_id, + revision=sources.config_revision, trust_remote_code=False, ) primary, parent, _ = _select_primary_config(hf_config) @@ -104,9 +109,11 @@ def vibevoice_package_dir(tmp_path_factory) -> Path: @pytest.fixture(scope="session") def vibevoice_processor(): transformers = pytest.importorskip("transformers") + sources = resolve_vibevoice_sources(VIBEVOICE_MODEL_ID, VIBEVOICE_REVISION) + assert sources is not None return transformers.AutoProcessor.from_pretrained( - VIBEVOICE_MODEL_ID, - revision=VIBEVOICE_REVISION, + sources.processor_model_id, + revision=sources.processor_revision, ) @@ -194,7 +201,8 @@ def test_vibevoice_l4_real_weight_prefill( device="cuda" if torch.cuda.is_available() else "cpu", max_new_tokens=1, ) - assert expected["revision"] == VIBEVOICE_REVISION + assert expected["model_id"] == VIBEVOICE_EXECUTABLE_MODEL_ID + assert expected["revision"] == VIBEVOICE_EXECUTABLE_REVISION np.testing.assert_allclose( actual.prefill_control_logits, expected["prefill_control_logits"], @@ -314,9 +322,11 @@ def test_vibevoice_real_weight_stage_parity( ): """Compare every real-weight stage, wiring ONNX outputs into the next stage.""" transformers = pytest.importorskip("transformers") + sources = resolve_vibevoice_sources(VIBEVOICE_MODEL_ID, VIBEVOICE_REVISION) + assert sources is not None model = transformers.AutoModelForTextToWaveform.from_pretrained( - VIBEVOICE_MODEL_ID, - revision=VIBEVOICE_REVISION, + sources.config_model_id, + revision=sources.config_revision, dtype=torch.float16, low_cpu_mem_usage=True, ).eval() diff --git a/tests/weight_alignment_test.py b/tests/weight_alignment_test.py index 27521cd39..7f8aeb60d 100644 --- a/tests/weight_alignment_test.py +++ b/tests/weight_alignment_test.py @@ -25,6 +25,8 @@ from __future__ import annotations +import json + import pytest import torch from _test_configs import ( @@ -170,6 +172,37 @@ def test_vibevoice_native_hf_weights_cover_every_stage_parameter(): assert parameter_names == set(routed) +@pytest.mark.integration +def test_official_vibevoice_weight_index_matches_native_conversion(): + """All official index keys map one-to-one to the pinned native checkpoint.""" + from huggingface_hub import hf_hub_download + + from mobius.models.vibevoice import ( + VIBEVOICE_EXECUTABLE_MODEL_ID, + VIBEVOICE_EXECUTABLE_REVISION, + VIBEVOICE_MODEL_ID, + VIBEVOICE_REVISION, + _transform_official_weight_name, + ) + + def weight_names(model_id: str, revision: str) -> set[str]: + path = hf_hub_download( + repo_id=model_id, + filename="model.safetensors.index.json", + revision=revision, + ) + with open(path, encoding="utf-8") as file: + return set(json.load(file)["weight_map"]) + + official = weight_names(VIBEVOICE_MODEL_ID, VIBEVOICE_REVISION) + native = weight_names(VIBEVOICE_EXECUTABLE_MODEL_ID, VIBEVOICE_EXECUTABLE_REVISION) + converted = {_transform_official_weight_name(name) for name in official} + + assert len(official) == 1204 + assert len(converted) == len(official) + assert converted == native + + # --------------------------------------------------------------------------- # Causal LM weight alignment # ---------------------------------------------------------------------------