Skip to content

Commit d3f772f

Browse files
Copilotxiaoyu-work
Copilot
and
xiaoyu-work
authored
Merge remote-tracking branch 'origin/main' into feat/per-component-quantization-unified
Signed-off-by: GitHub <noreply@github.com> # Conflicts: # src/mobius/integrations/transformers/_builder.py Co-authored-by: xiaoyu-work <85524621+xiaoyu-work@users.noreply.github.com>
2 parents 1a4e390 + 45a477d commit d3f772f

4 files changed

Lines changed: 141 additions & 32 deletions

File tree

‎src/mobius/integrations/transformers/_builder.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -404,7 +404,6 @@ def build_transformers_model(
404404
_config_from_hf,
405405
_default_task_for_model,
406406
)
407-
from mobius.models.vibevoice import VIBEVOICE_ASR_STREAMING_MODEL_REVISIONS
408407

409408
detection_revision = revision
410409
vibevoice_sources = None
@@ -439,9 +438,6 @@ def build_transformers_model(
439438
# processor contract. Keep config detection and weight loading pinned.
440439
revision = VIBEVOICE_ASR_REVISION
441440
detection_revision = VIBEVOICE_ASR_REVISION
442-
if model_id in VIBEVOICE_ASR_STREAMING_MODEL_REVISIONS and detection_revision is None:
443-
revision = VIBEVOICE_ASR_STREAMING_MODEL_REVISIONS[model_id]
444-
detection_revision = revision
445441
if model_id == "nvidia/RE-USE" and detection_revision is None:
446442
# Pin the very first AutoConfig/raw-JSON probe, not only the later
447443
# bespoke loader. Otherwise mutable Hub main could change dispatch
@@ -768,7 +764,11 @@ def build_transformers_model(
768764
"memory. The default native MXFP4 streaming path is bounded."
769765
)
770766
state_dict = _download_weights(model_id, revision=revision)
771-
if vibevoice_sources is not None and hasattr(model_module, "preprocess_weights"):
767+
if (
768+
vibevoice_sources is not None
769+
and vibevoice_sources.weight_layout is not None
770+
and hasattr(model_module, "preprocess_weights")
771+
):
772772
state_dict = model_module.preprocess_weights(
773773
state_dict,
774774
checkpoint_layout=vibevoice_sources.weight_layout,

‎src/mobius/integrations/transformers/_builder_test.py‎

Lines changed: 79 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1481,11 +1481,88 @@ def preprocess_weights(self, state_dict, *, checkpoint_layout):
14811481
}
14821482

14831483

1484+
@pytest.mark.parametrize(
1485+
("model_id", "revision"),
1486+
(
1487+
("microsoft/VibeVoice-ASR-Streaming-1.5B", "4262d23d8a539a6530cf64fbd0b1751ef9a30853"),
1488+
("microsoft/VibeVoice-ASR-Streaming-7B", "60d858b518b4e19d404af3737f848fc185b30177"),
1489+
),
1490+
)
1491+
def test_streaming_asr_uses_pinned_sources_and_native_weight_preprocessing(
1492+
monkeypatch, model_id, revision
1493+
) -> None:
1494+
"""The shared resolver must preserve streaming ASR's preprocessing signature."""
1495+
parent = SimpleNamespace(model_type="vibevoice", architectures=[])
1496+
package = ModelPackage(
1497+
{
1498+
"audio_encoder": ir.Model(
1499+
ir.Graph([], [], nodes=[], name="audio_encoder"), ir_version=11
1500+
)
1501+
}
1502+
)
1503+
calls = []
1504+
1505+
class StreamingAsrModule(_DummyModule):
1506+
def preprocess_weights(self, state_dict):
1507+
calls.append(("preprocess", state_dict))
1508+
return state_dict
1509+
1510+
monkeypatch.setattr(
1511+
transformers_builder,
1512+
"_load_transformers_config",
1513+
lambda *args, **kwargs: calls.append(("config", args, kwargs)) or (parent, False),
1514+
)
1515+
monkeypatch.setattr(
1516+
transformers_builder,
1517+
"_select_primary_config",
1518+
lambda value: (value, value, "VibeVoiceForASRStreamingTraining"),
1519+
)
1520+
monkeypatch.setattr(
1521+
transformers_builder,
1522+
"_resolve_module_class",
1523+
lambda *args, **kwargs: (
1524+
StreamingAsrModule,
1525+
"vibevoice-asr-streaming",
1526+
"VibeVoiceForASRStreamingTraining",
1527+
),
1528+
)
1529+
monkeypatch.setattr(
1530+
_config_resolver,
1531+
"_config_from_hf",
1532+
lambda *args, **kwargs: make_config(model_type="vibevoice"),
1533+
)
1534+
monkeypatch.setattr(
1535+
transformers_builder, "build_from_module", lambda *args, **kwargs: package
1536+
)
1537+
state_dict = {"weight": torch.ones(())}
1538+
monkeypatch.setattr(
1539+
transformers_builder,
1540+
"_download_weights",
1541+
lambda *args, **kwargs: calls.append(("weights", args, kwargs)) or state_dict,
1542+
)
1543+
1544+
result = transformers_builder.build_transformers_model(model_id)
1545+
1546+
assert result is package
1547+
assert calls == [
1548+
(
1549+
"config",
1550+
(model_id,),
1551+
{"revision": revision, "trust_remote_code": False},
1552+
),
1553+
("weights", (model_id,), {"revision": revision}),
1554+
("preprocess", state_dict),
1555+
]
1556+
assert package["audio_encoder"].metadata_props == {
1557+
"mobius.source_revision": revision,
1558+
"mobius.executable_source": f"{model_id}@{revision}",
1559+
"mobius.processor_source": f"{model_id}@{revision}",
1560+
}
1561+
1562+
14841563
@pytest.mark.parametrize(
14851564
"model_id",
14861565
[
1487-
"microsoft/VibeVoice-ASR-Streaming-7B",
1488-
"microsoft/VibeVoice-ASR-Streaming-1.5B",
14891566
"microsoft/VibeVoice-ASR-BitNet",
14901567
"microsoft/VibeVoice-AcousticTokenizer",
14911568
"Microsoft/VibeVoice-ASR-BitNet",

‎src/mobius/models/vibevoice.py‎

Lines changed: 33 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -46,16 +46,24 @@
4646
VIBEVOICE_EXECUTABLE_MODEL_ID = "vibevoice/VibeVoice-1.5B-hf"
4747
VIBEVOICE_EXECUTABLE_REVISION = "edc39f80f5cae656da37baf8faa8f5502bf7081f"
4848
VIBEVOICE_MICROSOFT_PROVENANCE_REVISION = VIBEVOICE_REVISION
49+
VIBEVOICE_ASR_STREAMING_MODEL_REVISIONS = {
50+
"microsoft/VibeVoice-ASR-Streaming-1.5B": "4262d23d8a539a6530cf64fbd0b1751ef9a30853",
51+
"microsoft/VibeVoice-ASR-Streaming-7B": "60d858b518b4e19d404af3737f848fc185b30177",
52+
}
53+
VIBEVOICE_ASR_STREAMING_MODEL_ID = "microsoft/VibeVoice-ASR-Streaming-1.5B"
54+
VIBEVOICE_ASR_STREAMING_REVISION = VIBEVOICE_ASR_STREAMING_MODEL_REVISIONS[
55+
VIBEVOICE_ASR_STREAMING_MODEL_ID
56+
]
57+
VIBEVOICE_ASR_STREAMING_SOURCE_REVISION = "505653d3873b065a488aea551c6ee3dc51d3062f"
4958

5059

5160
@dataclasses.dataclass(frozen=True)
5261
class VibeVoiceSources:
53-
"""Immutable provenance for one VibeVoice TTS build.
62+
"""Immutable provenance for one supported VibeVoice checkpoint.
5463
55-
``model_id`` and ``weight_revision`` are always the user's checkpoint. The
56-
official 1.5B release predates Transformers-native VibeVoice metadata, so
57-
its executable config and processor are resolved from the pinned conversion
58-
mirror while its official weights remain the only downloaded weights.
64+
``model_id`` and ``weight_revision`` are always the user's checkpoint.
65+
``weight_layout`` is set only when the module accepts an explicit checkpoint
66+
layout during weight preprocessing.
5967
"""
6068

6169
model_id: str
@@ -64,18 +72,10 @@ class VibeVoiceSources:
6472
config_revision: str
6573
processor_model_id: str
6674
processor_revision: str
67-
weight_layout: str
75+
weight_layout: str | None
6876

6977

7078
_UNSUPPORTED_VIBEVOICE_MODELS = {
71-
"microsoft/VibeVoice-ASR-Streaming-7B": (
72-
"VibeVoice ASR Streaming requires the VibeVoice-ASR streaming task, "
73-
"which Mobius does not export yet."
74-
),
75-
"microsoft/VibeVoice-ASR-Streaming-1.5B": (
76-
"VibeVoice ASR Streaming requires the VibeVoice-ASR streaming task, "
77-
"which Mobius does not export yet."
78-
),
7979
"microsoft/VibeVoice-ASR-BitNet": (
8080
"VibeVoice ASR BitNet requires the VibeVoice-ASR task and BitNet "
8181
"weight loader, which Mobius does not export yet."
@@ -104,6 +104,25 @@ def resolve_vibevoice_sources(model_id: str, revision: str | None) -> VibeVoiceS
104104
raise NotImplementedError(
105105
f"{model_id} is unsupported: {unsupported[canonical_model_id]}"
106106
)
107+
streaming_revisions = {
108+
known_model_id.casefold(): pinned_revision
109+
for known_model_id, pinned_revision in VIBEVOICE_ASR_STREAMING_MODEL_REVISIONS.items()
110+
}
111+
if canonical_model_id in streaming_revisions:
112+
pinned_revision = streaming_revisions[canonical_model_id]
113+
if revision not in {None, pinned_revision}:
114+
raise ValueError(
115+
f"{model_id} is only verified at revision {pinned_revision}; got {revision}."
116+
)
117+
return VibeVoiceSources(
118+
model_id=model_id,
119+
weight_revision=pinned_revision,
120+
config_model_id=model_id,
121+
config_revision=pinned_revision,
122+
processor_model_id=model_id,
123+
processor_revision=pinned_revision,
124+
weight_layout=None,
125+
)
107126
if canonical_model_id == VIBEVOICE_MODEL_ID.casefold():
108127
if revision not in {None, VIBEVOICE_REVISION}:
109128
raise ValueError(
@@ -1233,17 +1252,6 @@ def preprocess_weights(
12331252
return routed
12341253

12351254

1236-
VIBEVOICE_ASR_STREAMING_MODEL_REVISIONS = {
1237-
"microsoft/VibeVoice-ASR-Streaming-1.5B": "4262d23d8a539a6530cf64fbd0b1751ef9a30853",
1238-
"microsoft/VibeVoice-ASR-Streaming-7B": "60d858b518b4e19d404af3737f848fc185b30177",
1239-
}
1240-
VIBEVOICE_ASR_STREAMING_MODEL_ID = "microsoft/VibeVoice-ASR-Streaming-1.5B"
1241-
VIBEVOICE_ASR_STREAMING_REVISION = VIBEVOICE_ASR_STREAMING_MODEL_REVISIONS[
1242-
VIBEVOICE_ASR_STREAMING_MODEL_ID
1243-
]
1244-
VIBEVOICE_ASR_STREAMING_SOURCE_REVISION = "505653d3873b065a488aea551c6ee3dc51d3062f"
1245-
1246-
12471255
class VibeVoiceASRStreamingAudioEncoder(nn.Module):
12481256
"""Encode waveform chunks with streaming ASR's two causal tokenizers."""
12491257

‎src/mobius/models/vibevoice_test.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from mobius._registry import registry
2727
from mobius._testing.ort_inference import OnnxModelSession
2828
from mobius.models.vibevoice import (
29+
VIBEVOICE_ASR_STREAMING_MODEL_REVISIONS,
2930
VIBEVOICE_EXECUTABLE_MODEL_ID,
3031
VIBEVOICE_EXECUTABLE_REVISION,
3132
VIBEVOICE_MODEL_ID,
@@ -127,6 +128,29 @@ def test_official_vibevoice_sources_normalize_hub_id_case():
127128
resolve_vibevoice_sources("Microsoft/VibeVoice-ASR-BitNet", None)
128129

129130

131+
@pytest.mark.parametrize(
132+
("model_id", "revision"),
133+
tuple(VIBEVOICE_ASR_STREAMING_MODEL_REVISIONS.items()),
134+
)
135+
def test_streaming_asr_sources_pin_checkpoint_assets(model_id, revision):
136+
sources = resolve_vibevoice_sources(model_id, None)
137+
138+
assert sources is not None
139+
assert sources.model_id == model_id
140+
assert sources.weight_revision == revision
141+
assert sources.config_model_id == model_id
142+
assert sources.config_revision == revision
143+
assert sources.processor_model_id == model_id
144+
assert sources.processor_revision == revision
145+
assert sources.weight_layout is None
146+
147+
148+
@pytest.mark.parametrize("model_id", VIBEVOICE_ASR_STREAMING_MODEL_REVISIONS)
149+
def test_streaming_asr_sources_reject_unverified_revisions(model_id):
150+
with pytest.raises(ValueError, match="only verified"):
151+
resolve_vibevoice_sources(model_id, "different-streaming-asr-revision")
152+
153+
130154
@pytest.mark.parametrize(
131155
("source", "expected"),
132156
[

0 commit comments

Comments
 (0)