From a630b0c2ec5e8ffd3be4ff490e4a85fe06f46536 Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:13:40 +0800 Subject: [PATCH] [Ming] Cleanup Ming-family shared modules (#6119) Signed-off-by: Yuanheng Signed-off-by: Yuanheng Zhao --- .../text_to_speech/ming_tts/end2end.py | 5 +- recipes/inclusionAI/Ming-omni-tts.md | 19 ++-- .../model_loader/test_resolve_model_path.py | 92 +++++++++++++++++++ .../ming_tts/test_ming_shared_modules.py | 91 +++++++++++++++++- .../model_loader/weight_utils.py | 44 +++++++++ .../models/common/ming/__init__.py | 11 +++ .../{ming_tts => common/ming}/aggregator.py | 2 +- .../fm => common/ming}/cfm_cudagraph.py | 2 +- .../fm/dit.py => common/ming/cfm_dit.py} | 2 +- .../ming/cfm_head.py} | 5 +- .../common/ming/{fm.py => cfm_solver.py} | 16 +++- .../common/ming/{dit.py => dit_blocks.py} | 16 ---- .../models/common/ming/speaker_extractor.py | 56 +++++++++++ .../ming}/spk_embedding.py | 0 .../ming_flash_omni/ming_flash_omni_talker.py | 36 ++++---- .../ming_flash_omni_thinker.py | 7 +- .../models/ming_flash_omni/talker_module.py | 61 ++---------- .../models/ming_flash_omni/voice_presets.py | 34 +++---- .../models/ming_tts/fm/__init__.py | 2 - .../models/ming_tts/ming_tts_llm.py | 13 +-- .../models/ming_tts/prompt_encoder.py | 4 +- .../models/ming_tts/speaker_extractor.py | 54 +---------- 22 files changed, 380 insertions(+), 192 deletions(-) create mode 100644 tests/model_executor/model_loader/test_resolve_model_path.py rename vllm_omni/model_executor/models/{ming_tts => common/ming}/aggregator.py (96%) rename vllm_omni/model_executor/models/{ming_tts/fm => common/ming}/cfm_cudagraph.py (98%) rename vllm_omni/model_executor/models/{ming_tts/fm/dit.py => common/ming/cfm_dit.py} (97%) rename vllm_omni/model_executor/models/{ming_tts/flowloss_head.py => common/ming/cfm_head.py} (94%) rename vllm_omni/model_executor/models/common/ming/{fm.py => cfm_solver.py} (77%) rename vllm_omni/model_executor/models/common/ming/{dit.py => dit_blocks.py} (93%) create mode 100644 vllm_omni/model_executor/models/common/ming/speaker_extractor.py rename vllm_omni/model_executor/models/{ming_flash_omni => common/ming}/spk_embedding.py (100%) delete mode 100644 vllm_omni/model_executor/models/ming_tts/fm/__init__.py diff --git a/examples/offline_inference/text_to_speech/ming_tts/end2end.py b/examples/offline_inference/text_to_speech/ming_tts/end2end.py index 1a6991e40..89d5fda73 100644 --- a/examples/offline_inference/text_to_speech/ming_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/ming_tts/end2end.py @@ -13,6 +13,7 @@ import yaml from transformers import AutoTokenizer from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( KEY_CFG, KEY_MAX_DECODE_STEPS, @@ -21,7 +22,6 @@ from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( SAMPLE_RATE, ) from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt -from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor try: from .runner import ( @@ -142,7 +142,8 @@ def _resolve_speaker_embedding(args, case: dict, paths: list[str]): return _load_speaker_embedding(args.speaker_embedding) if not (case.get("auto_extract_speaker_embeddings") or args.extract_speaker_embeddings) or not paths: return None - embs = MingSpeakerEmbeddingExtractor(args.model).extract_many(paths) + # Runs before the engine is built, so nothing has pre-fetched campplus.onnx. + embs = SpeakerEmbeddingExtractor(args.model, allow_download=True).extract_many(paths) if not embs: raise RuntimeError("Speaker extraction produced no embeddings") return embs[0] if len(embs) == 1 else torch.stack(embs, dim=0) diff --git a/recipes/inclusionAI/Ming-omni-tts.md b/recipes/inclusionAI/Ming-omni-tts.md index 0b3584691..3666dce0f 100644 --- a/recipes/inclusionAI/Ming-omni-tts.md +++ b/recipes/inclusionAI/Ming-omni-tts.md @@ -5,20 +5,21 @@ ## Summary - Vendor: inclusionAI -- Model: `inclusionAI/Ming-omni-tts-0.5B` +- Model: `inclusionAI/Ming-omni-tts-16.8B-A3B`, `inclusionAI/Ming-omni-tts-0.5B` - Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls - Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API - Maintainer: Community ## References -- [Huggingface Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B) +- [Huggingface: inclusionAI/Ming-omni-tts-16.8B-A3B](https://huggingface.co/inclusionAI/Ming-omni-tts-16.8B-A3B) +- [Huggingface: inclusionAI/Ming-omni-tts-0.5B](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B) - Upstream repository [inclusionAI/Ming-omni-tts](https://github.com/inclusionAI/Ming-omni-tts) ## Hardware Support -This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B two-stage TTS pipeline deployment. +This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B and the MoE 16.8B two-stage TTS pipeline deployment. Other hardware is welcome as community validation lands. ## CUDA @@ -39,11 +40,7 @@ Other hardware is welcome as community validation lands. Launch the two-stage talker: ```bash -vllm serve inclusionAI/Ming-omni-tts-0.5B \ - --deploy-config vllm_omni/deploy/ming_tts.yaml \ - --omni \ - --port 8091 \ - --enforce-eager +vllm serve inclusionAI/Ming-omni-tts-16.8B-A3B --omni --port 8091 ``` #### Verification @@ -54,7 +51,7 @@ Basic synthesis (save the WAV bytes): curl -X POST http://localhost:8091/v1/audio/speech \ -H "Content-Type: application/json" \ -d '{ - "model": "inclusionAI/Ming-omni-tts-0.5B", + "model": "inclusionAI/Ming-omni-tts-16.8B-A3B", "input": "你好,这是 Ming 在线语音合成测试。", "response_format": "wav" }' --output ming_tts_basic.wav @@ -74,7 +71,7 @@ Reference-audio zero-shot cloning: REF_AUDIO="$BASE/10002287-00000094.wav" jq -n --arg ref_audio "$REF_AUDIO" \ '{ - model: "inclusionAI/Ming-omni-tts-0.5B", + model: "inclusionAI/Ming-omni-tts-16.8B-A3B", input: "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。", ref_audio: $ref_audio, ref_text: "在此奉劝大家别乱打美白针。", @@ -107,7 +104,7 @@ jq -n \ --arg a "$REF_A" \ --arg b "$REF_B" \ '{ - model: "inclusionAI/Ming-omni-tts-0.5B", + model: "inclusionAI/Ming-omni-tts-16.8B-A3B", input: $input, ref_audio: [$a, $b], ref_text: $ref_text, diff --git a/tests/model_executor/model_loader/test_resolve_model_path.py b/tests/model_executor/model_loader/test_resolve_model_path.py new file mode 100644 index 000000000..996313b06 --- /dev/null +++ b/tests/model_executor/model_loader/test_resolve_model_path.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for resolve_model_to_local_path.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import huggingface_hub +import pytest +from huggingface_hub.errors import LocalEntryNotFoundError + +from vllm_omni.model_executor.model_loader import weight_utils +from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Both resolution paths are patched in every test, so nothing is ever fetched +REPO_ID = "test-org/test-model" +RESOLVED = "/cache/snapshots/deadbeef" + + +@pytest.fixture +def paths(mocker): + """The resolver's two ways out, both patched.""" + return SimpleNamespace( + cache_only=mocker.patch.object(huggingface_hub, "snapshot_download", return_value=RESOLVED), + download=mocker.patch.object(weight_utils, "download_weights_from_hf_specific", return_value=RESOLVED), + ) + + +@pytest.mark.parametrize("allow_download", [False, True]) +def test_local_directory_short_circuits(tmp_path, paths, allow_download): + """A local dir resolves to itself without either path running.""" + assert resolve_model_to_local_path(str(tmp_path), allow_download=allow_download) == str(tmp_path) + + paths.cache_only.assert_not_called() + paths.download.assert_not_called() + + +@pytest.mark.parametrize( + ("kwargs", "allow_patterns", "cache_dir"), + [ + ({}, None, None), + ({"allow_patterns": ["subdir/*"], "cache_dir": "/cache"}, ["subdir/*"], "/cache"), + ], +) +def test_repo_id_is_cache_only_by_default(paths, kwargs, allow_patterns, cache_dir): + """Without `allow_download` nothing is fetched, and narrowing is forwarded.""" + assert resolve_model_to_local_path(REPO_ID, **kwargs) == RESOLVED + + paths.cache_only.assert_called_once_with( + REPO_ID, + cache_dir=cache_dir, + allow_patterns=allow_patterns, + local_files_only=True, + ) + paths.download.assert_not_called() + + +def test_cache_miss_propagates(paths): + """A failed resolution raises instead of degrading into the raw repo id.""" + paths.cache_only.side_effect = LocalEntryNotFoundError("cold cache") + + with pytest.raises(LocalEntryNotFoundError): + resolve_model_to_local_path(REPO_ID) + + paths.download.assert_not_called() + + +@pytest.mark.parametrize( + ("kwargs", "allow_patterns"), + [ + ({}, ["*"]), + ({"allow_patterns": ["subdir/model*.safetensors"]}, ["subdir/model*.safetensors"]), + ], +) +def test_download_routes_through_the_locked_helper(paths, kwargs, allow_patterns): + """Downloads go through the vLLM helper holding the node-wide lock. + + That helper is also what honors `VLLM_USE_MODELSCOPE`; + calling `snapshot_download` directly would drop both. + """ + assert resolve_model_to_local_path(REPO_ID, allow_download=True, cache_dir="/cache", **kwargs) == RESOLVED + + paths.download.assert_called_once_with( + model_name_or_path=REPO_ID, + cache_dir="/cache", + allow_patterns=allow_patterns, + require_all=True, + ) + paths.cache_only.assert_not_called() diff --git a/tests/model_executor/models/ming_tts/test_ming_shared_modules.py b/tests/model_executor/models/ming_tts/test_ming_shared_modules.py index 7a0ad68de..21c8dfd14 100644 --- a/tests/model_executor/models/ming_tts/test_ming_shared_modules.py +++ b/tests/model_executor/models/ming_tts/test_ming_shared_modules.py @@ -7,8 +7,14 @@ import pytest import torch from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig -from vllm_omni.model_executor.models.common.ming.fm import Solver +from vllm_omni.model_executor.models.common.ming.cfm_cudagraph import CFMGraphExecutor, CFMSampler +from vllm_omni.model_executor.models.common.ming.cfm_dit import DiT +from vllm_omni.model_executor.models.common.ming.cfm_head import FlowLoss +from vllm_omni.model_executor.models.common.ming.cfm_solver import Solver, get_epss_timesteps +from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor +from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor from vllm_omni.model_executor.models.ming_tts.constants import ( AGGREGATOR_HIDDEN_SIZE, HISTORY_PATCH_SIZE, @@ -24,7 +30,6 @@ from vllm_omni.model_executor.models.ming_tts.constants import ( SPEAKER_EMBEDDING_DIM, VAE_PATCH_SIZE, ) -from vllm_omni.model_executor.models.ming_tts.flowloss_head import FlowLoss from vllm_omni.model_executor.models.ming_tts.ming_tts_llm import MingLLMModel from vllm_omni.model_executor.models.ming_tts.speaker_extractor import _resolve_speaker_embeddings from vllm_omni.model_executor.models.ming_tts.validation import validate_ming_tts_config @@ -44,7 +49,87 @@ def test_ming_tts_audio_vae_uses_common_config(): def test_ming_tts_cfm_solver_uses_common_implementation(): """Ming dense imports the shared solver implementation directly.""" - assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.fm" + assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.cfm_solver" + + +@pytest.mark.parametrize( + ("symbol", "expected_module"), + [ + (Aggregator, "vllm_omni.model_executor.models.common.ming.aggregator"), + (DiT, "vllm_omni.model_executor.models.common.ming.cfm_dit"), + (FlowLoss, "vllm_omni.model_executor.models.common.ming.cfm_head"), + (CFMGraphExecutor, "vllm_omni.model_executor.models.common.ming.cfm_cudagraph"), + (CFMSampler, "vllm_omni.model_executor.models.common.ming.cfm_cudagraph"), + (SpeakerEmbeddingExtractor, "vllm_omni.model_executor.models.common.ming.speaker_extractor"), + (SpkembExtractor, "vllm_omni.model_executor.models.common.ming.spk_embedding"), + ], +) +def test_ming_tts_reusable_modules_live_in_common(symbol, expected_module): + """The pieces a second Ming-family TTS model reuses live under common/ming. + + Guards the promotion: if an implementation drifts back into `ming_tts` + (or into another model package), a second model would have to reach across + packages to get it. + """ + assert symbol.__module__ == expected_module + + +def test_ming_flash_talker_uses_the_shared_aggregator(): + """Ming flash no longer carries its own copy of the aggregator. + + Both Ming packages ran near-identical aggregators one directory apart, + which is the duplication this module layout exists to remove. + """ + from vllm_omni.model_executor.models.ming_flash_omni import talker_module + + assert talker_module.Aggregator is Aggregator + assert not any( + isinstance(obj, type) and obj.__module__ == talker_module.__name__ and name == "Aggregator" + for name, obj in vars(talker_module).items() + ) + + +def test_shared_aggregator_matches_the_ming_flash_construction(): + """The shared aggregator reproduces Ming flash's former default geometry. + + `talker_module.Aggregator` defaulted to `in_channels=64` while the + shared class defaults to 4, so the call site now passes `latent_dim` + explicitly. This pins the resulting shapes rather than the default. + """ + latent_dim, llm_hidden = 64, 32 + aggregator = Aggregator( + in_channels=latent_dim, + llm_input_dim=llm_hidden, + hidden_size=16, + depth=1, + num_heads=2, + ) + + assert aggregator.x_embedder.in_features == latent_dim + assert aggregator.final_layer.linear.out_features == llm_hidden + + out = aggregator(torch.randn(1, 3, latent_dim)) + # Keeps only the CLS projection: [B, T, D] -> [B, 1, llm_hidden]. + assert out.shape == (1, 1, llm_hidden) + + +def test_shared_aggregator_rejects_mismatched_latent_width(): + """Wrong-width input fails loudly instead of silently projecting incorrect.""" + aggregator = Aggregator(in_channels=64, llm_input_dim=32, hidden_size=16, depth=1, num_heads=2) + + with pytest.raises(ValueError, match="feature dim mismatch"): + aggregator(torch.randn(1, 3, 4)) + + +def test_cfm_solver_owns_the_timestep_schedule(): + """get_epss_timesteps sits with the solver, not with the DiT blocks.""" + from vllm_omni.model_executor.models.common.ming import dit_blocks + + assert not hasattr(dit_blocks, "get_epss_timesteps") + + t = get_epss_timesteps(10, device="cpu", dtype=torch.float32) + assert t.shape == (11,) + assert float(t[0]) == 0.0 and float(t[-1]) == 1.0 def test_ming_tts_flowloss_preserves_checkpoint_prefix(): diff --git a/vllm_omni/model_executor/model_loader/weight_utils.py b/vllm_omni/model_executor/model_loader/weight_utils.py index c5225ec02..59c2cd860 100644 --- a/vllm_omni/model_executor/model_loader/weight_utils.py +++ b/vllm_omni/model_executor/model_loader/weight_utils.py @@ -1,3 +1,4 @@ +import os import time from pathlib import Path @@ -85,3 +86,46 @@ def download_weights_from_hf_specific( time_taken, ) return hf_folder + + +def resolve_model_to_local_path( + model: str, + *, + allow_download: bool = False, + allow_patterns: list[str] | None = None, + cache_dir: str | None = None, +) -> str: + """Resolve a model reference to a local directory. + Failure is always raised: whether a failure is fatal is the caller's decision. + + Args: + model: A local directory or an HF Hub repo ID. + allow_download: Fetch missing files instead of requiring a warm cache. + Downloads go through :func:`download_weights_from_hf_specific` + allow_patterns: Restrict the fetch to these glob patterns. + Defaults to the whole repository. Ignored when model is already a directory. + cache_dir: Download cache location; None uses the HF default. + + Returns: + Path to a local directory containing the model files. + """ + if os.path.isdir(model): + return model + + if allow_download: + return download_weights_from_hf_specific( + model_name_or_path=model, + cache_dir=cache_dir, + allow_patterns=allow_patterns or ["*"], + require_all=True, + ) + + # Cache-only lookups stay on huggingface_hub even under VLLM_USE_MODELSCOPE: + # ModelScope's snapshot_download has no local_files_only equivalent, + # so a ModelScope deployment that needs resolution here must pass allow_download. + return huggingface_hub.snapshot_download( + model, + cache_dir=cache_dir, + allow_patterns=allow_patterns, + local_files_only=True, + ) diff --git a/vllm_omni/model_executor/models/common/ming/__init__.py b/vllm_omni/model_executor/models/common/ming/__init__.py index 208f01a7c..5134c7cf9 100644 --- a/vllm_omni/model_executor/models/common/ming/__init__.py +++ b/vllm_omni/model_executor/models/common/ming/__init__.py @@ -1,2 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Components shared by the Ming-family model packages. + +Naming: The flow-matching pieces all use the cfm_ prefix, +because "flow matching", "fm" and "cfm" were three names for one concept here: + +* dit_blocks - DiT building blocks (DiTBlock, FinalLayer, CondEmbedder) +* cfm_dit - the DiT used as the CFM velocity model +* cfm_solver - timestep schedules and the ODE/SDE integrator +* cfm_head - inference-only CFM wrapper preserving cfm.model.* names +* cfm_cudagraph - batch-1 CUDA-graph executor for the sampling loop +""" diff --git a/vllm_omni/model_executor/models/ming_tts/aggregator.py b/vllm_omni/model_executor/models/common/ming/aggregator.py similarity index 96% rename from vllm_omni/model_executor/models/ming_tts/aggregator.py rename to vllm_omni/model_executor/models/common/ming/aggregator.py index 3233c9659..e32e8b47b 100644 --- a/vllm_omni/model_executor/models/ming_tts/aggregator.py +++ b/vllm_omni/model_executor/models/common/ming/aggregator.py @@ -7,7 +7,7 @@ import torch import torch.nn as nn from x_transformers.x_transformers import RotaryEmbedding -from vllm_omni.model_executor.models.common.ming.dit import DiTBlock, FinalLayer +from vllm_omni.model_executor.models.common.ming.dit_blocks import DiTBlock, FinalLayer class Aggregator(nn.Module): diff --git a/vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py b/vllm_omni/model_executor/models/common/ming/cfm_cudagraph.py similarity index 98% rename from vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py rename to vllm_omni/model_executor/models/common/ming/cfm_cudagraph.py index 7a7280034..1eb3a9db0 100644 --- a/vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py +++ b/vllm_omni/model_executor/models/common/ming/cfm_cudagraph.py @@ -17,7 +17,7 @@ from __future__ import annotations import torch import torch.nn as nn -from vllm_omni.model_executor.models.common.ming.fm import build_timesteps +from vllm_omni.model_executor.models.common.ming.cfm_solver import build_timesteps class CFMSampler(nn.Module): diff --git a/vllm_omni/model_executor/models/ming_tts/fm/dit.py b/vllm_omni/model_executor/models/common/ming/cfm_dit.py similarity index 97% rename from vllm_omni/model_executor/models/ming_tts/fm/dit.py rename to vllm_omni/model_executor/models/common/ming/cfm_dit.py index 3eb13e987..c29ddf48d 100644 --- a/vllm_omni/model_executor/models/ming_tts/fm/dit.py +++ b/vllm_omni/model_executor/models/common/ming/cfm_dit.py @@ -7,7 +7,7 @@ import torch.nn as nn from x_transformers.x_transformers import RotaryEmbedding from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding -from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer +from vllm_omni.model_executor.models.common.ming.dit_blocks import CondEmbedder, DiTBlock, FinalLayer class DiT(nn.Module): diff --git a/vllm_omni/model_executor/models/ming_tts/flowloss_head.py b/vllm_omni/model_executor/models/common/ming/cfm_head.py similarity index 94% rename from vllm_omni/model_executor/models/ming_tts/flowloss_head.py rename to vllm_omni/model_executor/models/common/ming/cfm_head.py index 3f300cec0..14e63557e 100644 --- a/vllm_omni/model_executor/models/ming_tts/flowloss_head.py +++ b/vllm_omni/model_executor/models/common/ming/cfm_head.py @@ -5,9 +5,8 @@ import torch import torch.nn as nn -from vllm_omni.model_executor.models.common.ming.fm import Solver, build_timesteps - -from .fm.dit import DiT +from vllm_omni.model_executor.models.common.ming.cfm_dit import DiT +from vllm_omni.model_executor.models.common.ming.cfm_solver import Solver, build_timesteps class FlowLoss(nn.Module): diff --git a/vllm_omni/model_executor/models/common/ming/fm.py b/vllm_omni/model_executor/models/common/ming/cfm_solver.py similarity index 77% rename from vllm_omni/model_executor/models/common/ming/fm.py rename to vllm_omni/model_executor/models/common/ming/cfm_solver.py index d2cbc6211..5d9db4634 100644 --- a/vllm_omni/model_executor/models/common/ming/fm.py +++ b/vllm_omni/model_executor/models/common/ming/cfm_solver.py @@ -3,7 +3,21 @@ import torch -from .dit import get_epss_timesteps + +def get_epss_timesteps(n, device, dtype): + dt = 1 / 32 + predefined_timesteps = { + 5: [0, 2, 4, 8, 16, 32], + 6: [0, 2, 4, 6, 8, 16, 32], + 7: [0, 2, 4, 6, 8, 16, 24, 32], + 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32], + 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32], + 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32], + } + t = predefined_timesteps.get(n, []) + if not t: + return torch.linspace(0, 1, n + 1, device=device, dtype=dtype) + return dt * torch.tensor(t, device=device, dtype=dtype) def apply_sway_sampling(t, sway_sampling_coef): diff --git a/vllm_omni/model_executor/models/common/ming/dit.py b/vllm_omni/model_executor/models/common/ming/dit_blocks.py similarity index 93% rename from vllm_omni/model_executor/models/common/ming/dit.py rename to vllm_omni/model_executor/models/common/ming/dit_blocks.py index 56e65f47f..8a093811e 100644 --- a/vllm_omni/model_executor/models/common/ming/dit.py +++ b/vllm_omni/model_executor/models/common/ming/dit_blocks.py @@ -216,19 +216,3 @@ class CondEmbedder(nn.Module): def forward(self, llm_cond): return self.cond_embedder(llm_cond) - - -def get_epss_timesteps(n, device, dtype): - dt = 1 / 32 - predefined_timesteps = { - 5: [0, 2, 4, 8, 16, 32], - 6: [0, 2, 4, 6, 8, 16, 32], - 7: [0, 2, 4, 6, 8, 16, 24, 32], - 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32], - 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32], - 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32], - } - t = predefined_timesteps.get(n, []) - if not t: - return torch.linspace(0, 1, n + 1, device=device, dtype=dtype) - return dt * torch.tensor(t, device=device, dtype=dtype) diff --git a/vllm_omni/model_executor/models/common/ming/speaker_extractor.py b/vllm_omni/model_executor/models/common/ming/speaker_extractor.py new file mode 100644 index 000000000..bb06a3b58 --- /dev/null +++ b/vllm_omni/model_executor/models/common/ming/speaker_extractor.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/spkemb_extractor.py +"""CAM++ speaker-embedding extraction, shared by the Ming model packages.""" + +from __future__ import annotations + +import os + +import torch +import torchaudio +from vllm.multimodal.media.audio import load_audio + +from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path +from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor + + +class SpeakerEmbeddingExtractor: + """CAM++ ONNX speaker embeddings, resampled to the extractor's rate.""" + + def __init__(self, model, target_sr=16000, *, allow_download: bool = False): + local_model_path = resolve_model_to_local_path( + model, + allow_download=allow_download, + allow_patterns=["campplus.onnx"], + ) + campplus_path = os.path.join(local_model_path, "campplus.onnx") + if not os.path.exists(campplus_path): + raise RuntimeError(f"Missing Ming speaker extractor model: {campplus_path}") + + self.target_sr = int(target_sr) + self._core = SpkembExtractor(campplus_path, target_sr=self.target_sr) + + def extract_from_waveform(self, waveform, sample_rate): + if not isinstance(waveform, torch.Tensor): + waveform = torch.as_tensor(waveform) + + tensor = waveform.detach().to(torch.float32) + if tensor.ndim == 1: + tensor = tensor.unsqueeze(0) + if int(sample_rate) != self.target_sr: + tensor = torchaudio.transforms.Resample(orig_freq=int(sample_rate), new_freq=self.target_sr)(tensor) + + embedding = self._core._extract_spk_embedding(tensor) + return embedding.squeeze(0).to(dtype=torch.float32) + + def extract_from_file(self, audio_path): + audio, sample_rate = load_audio(audio_path, sr=None, mono=False) + waveform = torch.from_numpy(audio).contiguous() + return self.extract_from_waveform(waveform, sample_rate) + + def extract_many(self, audio_paths): + return [self.extract_from_file(path) for path in audio_paths] + + +__all__ = ["SpeakerEmbeddingExtractor"] diff --git a/vllm_omni/model_executor/models/ming_flash_omni/spk_embedding.py b/vllm_omni/model_executor/models/common/ming/spk_embedding.py similarity index 100% rename from vllm_omni/model_executor/models/ming_flash_omni/spk_embedding.py rename to vllm_omni/model_executor/models/common/ming/spk_embedding.py diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py index 8868e9d2a..d7e31beb3 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py @@ -25,14 +25,15 @@ from vllm.model_executor.models.utils import AutoWeightsLoader from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin -from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific +from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path +from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE, AudioVAEConfig from vllm_omni.model_executor.models.ming_tts.constants import SPEAKER_EMBEDDING_DIM from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.transformers_utils.configs.ming_flash_omni import MingFlashOmniTalkerConfig from .prompt_utils import DEFAULT_PROMPT as MING_DEFAULT_PROMPT -from .talker_module import CFM, Aggregator, DiT, MingAudioGenerator, build_tts_input +from .talker_module import CFM, DiT, MingAudioGenerator, build_tts_input from .text_processing import segment_and_normalize from .voice_presets import VoicePresetRegistry @@ -119,7 +120,11 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin) DiT(llm_input_dim=self.hidden_size, **config.flowmodel), steps=config.steps, ) - self.aggregator = Aggregator(llm_input_dim=self.hidden_size, **config.aggregator) + # config.aggregator still wins `in_channels` if the checkpoint states it. + self.aggregator = Aggregator( + llm_input_dim=self.hidden_size, + **{"in_channels": self.latent_dim, **config.aggregator}, + ) self.stop_head = nn.Linear(self.hidden_size, 2, bias=True) self.spk_head = nn.Linear(SPEAKER_EMBEDDING_DIM, self.hidden_size, bias=True) @@ -525,27 +530,22 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin) def _iter_talker_safetensors(self) -> Iterable[tuple[str, torch.Tensor]]: """Yield (name, tensor) pairs from talker/model*.safetensors.""" - model_path = self._model_path - # Try local path first - for candidate in (os.path.join(model_path, "talker"), model_path): + model_root = resolve_model_to_local_path( + self._model_path, + allow_download=True, + allow_patterns=["talker/model*.safetensors"], + cache_dir=self.vllm_config.load_config.download_dir, + ) + + # Nested talker/ layout first, then a flat checkpoint at the root. + for candidate in (os.path.join(model_root, "talker"), model_root): sf_files = sorted(glob_module.glob(os.path.join(candidate, "model*.safetensors"))) if sf_files: for sf_path in sf_files: yield from load_file(sf_path, device="cpu").items() return - # HF hub fallback: download only the talker checkpoint files - model_root = download_weights_from_hf_specific( - model_path, - self.vllm_config.load_config.download_dir, - allow_patterns=["talker/model*.safetensors"], - ) - talker_dir = os.path.join(model_root, "talker") - sf_files = sorted(glob_module.glob(os.path.join(talker_dir, "model*.safetensors"))) - if not sf_files: - raise RuntimeError(f"No talker safetensors found under {model_root}. Expected talker/model*.safetensors.") - for sf_path in sf_files: - yield from load_file(sf_path, device="cpu").items() + raise RuntimeError(f"No talker safetensors found under {model_root}. Expected talker/model*.safetensors.") def _load_vae_weights(self) -> set[str]: """Load AudioVAE weights from talker/vae/model.safetensors.""" diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py index 8c527b8d9..03cf87abb 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py @@ -64,6 +64,7 @@ from vllm.sequence import IntermediateTensors from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin +from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config from vllm_omni.transformers_utils.processors.ming import ( @@ -715,12 +716,8 @@ class MingFlashOmniThinkerForConditionalGeneration( p = Path(model_path) / "mlp" / "model.safetensors" if not p.exists() and not os.path.exists(model_path): - from vllm_omni.model_executor.model_loader.weight_utils import ( - download_weights_from_hf_specific, - ) - try: - local_path = download_weights_from_hf_specific(model_path, None, ["mlp/*"]) + local_path = resolve_model_to_local_path(model_path, allow_download=True, allow_patterns=["mlp/*"]) p = Path(local_path) / "mlp" / "model.safetensors" except Exception: logger.exception("[MingFlashOmniThinker] failed to download mlp/ from %s", model_path) diff --git a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py index cfb1090c2..8e3e55611 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py @@ -33,9 +33,14 @@ from vllm.platforms import current_platform from x_transformers.x_transformers import RotaryEmbedding from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding +from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE -from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer, get_epss_timesteps -from vllm_omni.model_executor.models.common.ming.fm import apply_sway_sampling, integrate_cfm_steps +from vllm_omni.model_executor.models.common.ming.cfm_solver import ( + apply_sway_sampling, + get_epss_timesteps, + integrate_cfm_steps, +) +from vllm_omni.model_executor.models.common.ming.dit_blocks import CondEmbedder, DiTBlock, FinalLayer logger = init_logger(__name__) @@ -458,58 +463,6 @@ def silence_holder( return speech_out, sil_cache -######################################################################## -# Audio Postprocess -# Ported from: -# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/aggregator.py -######################################################################## - - -class Aggregator(nn.Module): - """Maps generated audio latent patches back to LLM embedding space.""" - - def __init__( - self, - in_channels: int = 64, - hidden_size: int = 1152, - depth: int = 28, - num_heads: int = 16, - mlp_ratio: float = 4.0, - llm_input_dim: int = 896, - **kwargs, - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = in_channels - self.num_heads = num_heads - - self.word_embedder = nn.Embedding(1, hidden_size) - self.x_embedder = nn.Linear(in_channels, hidden_size) - self.hidden_size = hidden_size - - self.rotary_embed = RotaryEmbedding(hidden_size // num_heads) - - self.blocks = nn.ModuleList( - [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)] - ) - self.final_layer = FinalLayer(hidden_size, llm_input_dim) - - def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: - x = self.x_embedder(x) - cls_embed = self.word_embedder(torch.zeros((x.shape[0], 1), dtype=torch.long, device=x.device)) - x = torch.cat([cls_embed, x], dim=1) - - rope = self.rotary_embed.forward_from_seq_len(x.shape[1]) - if mask is not None: - mask_pad = mask.clone().detach()[:, :1] - mask = torch.cat([mask_pad, mask], dim=-1) - for block in self.blocks: - x = block(x, mask, rope) - x = self.final_layer(x) - x = x[:, :1, :] - return x - - ######################################################################## # Prompt Builder # Adapted from: diff --git a/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py b/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py index 0a096df86..8f6db1afe 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py @@ -16,16 +16,15 @@ import torch from transformers.utils.hub import cached_file from vllm.logger import init_logger -from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific +from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path +from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor -from .spk_embedding import SpkembExtractor from .talker_module import resample if TYPE_CHECKING: + from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE - from .talker_module import Aggregator - logger = init_logger(__name__) @@ -273,18 +272,19 @@ class VoicePresetRegistry: if os.path.isfile(path): return path, candidate - if not os.path.isdir(self._model_path): - try: - hf_root = download_weights_from_hf_specific( - self._model_path, - self._download_dir, - allow_patterns=["talker/data/**"], - require_all=True, - ) - candidate = os.path.join(hf_root, "talker", "data", "voice_name.json") - if os.path.isfile(candidate): - return candidate, os.path.join(hf_root, "talker") - except Exception as e: # pragma: no cover - logger.info("Could not download voice presets from HF: %s", e) + try: + hf_root = resolve_model_to_local_path( + self._model_path, + allow_download=True, + allow_patterns=["talker/data/**"], + cache_dir=self._download_dir, + ) + except Exception as e: + logger.info("Could not download voice presets from HF: %s", e) + return None, None + + candidate = os.path.join(hf_root, "talker", "data", "voice_name.json") + if os.path.isfile(candidate): + return candidate, os.path.join(hf_root, "talker") return None, None diff --git a/vllm_omni/model_executor/models/ming_tts/fm/__init__.py b/vllm_omni/model_executor/models/ming_tts/fm/__init__.py deleted file mode 100644 index 208f01a7c..000000000 --- a/vllm_omni/model_executor/models/ming_tts/fm/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py index c0ec5a513..b42c4e9cb 100644 --- a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py +++ b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py @@ -11,15 +11,21 @@ import torch import torch.nn as nn from vllm.config import VllmConfig from vllm.logger import init_logger +from vllm.model_executor.models.bailing_moe import BailingMoeModel from vllm.model_executor.models.qwen2 import Qwen2Model from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper, maybe_prefix from vllm.sequence import IntermediateTensors from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata +from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator +from vllm_omni.model_executor.models.common.ming.cfm_cudagraph import ( + CFMGraphExecutor, + CFMSampler, +) +from vllm_omni.model_executor.models.common.ming.cfm_head import FlowLoss from vllm_omni.model_executor.models.output_templates import OmniOutput -from .aggregator import Aggregator from .config_ming_tts import ( KEY_CFG, KEY_DECODE_STEP, @@ -34,7 +40,6 @@ from .config_ming_tts import ( MingTTSConfig, ) from .constants import SPEAKER_EMBEDDING_DIM -from .flowloss_head import FlowLoss from .patch_emission import ( MING_STOP_REASON_CODES, MING_STOP_REASON_CONTINUE, @@ -82,8 +87,6 @@ class MingLLMModel(nn.Module): self.prefix = prefix self.fm_dtype = _resolve_ming_runtime_dtype(vllm_config) if self.ming_config.model_variant == "moe": - from vllm.model_executor.models.bailing_moe import BailingMoeModel - # BailingMoeModel reads ``vllm_config.model_config.hf_config`` directly # (no get_text_config()), so re-wrap with the nested bailing_moe config. llm_config = vllm_config.model_config.hf_config.get_text_config() @@ -335,8 +338,6 @@ class MingLLMModel(nn.Module): if self._cfm_graph is not None: return self._cfm_graph try: - from .fm.cfm_cudagraph import CFMGraphExecutor, CFMSampler - # TODO(perf): # (1) compile per-block (dit_model.blocks) instead of the # whole forward for tighter fusion / fewer graph breaks; diff --git a/vllm_omni/model_executor/models/ming_tts/prompt_encoder.py b/vllm_omni/model_executor/models/ming_tts/prompt_encoder.py index 251b2e0f9..355bd8873 100644 --- a/vllm_omni/model_executor/models/ming_tts/prompt_encoder.py +++ b/vllm_omni/model_executor/models/ming_tts/prompt_encoder.py @@ -12,7 +12,7 @@ import torch from safetensors import safe_open from vllm.logger import init_logger -from vllm_omni.engine.stage_init_utils import _resolve_model_to_local_path +from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE from .audio_prep import ( @@ -78,7 +78,7 @@ def _load_prompt_encoder(wrapper: Any) -> AudioVAE: loaded_encoder_params = set() with torch.no_grad(): for shard_path in _iter_model_safetensors( - _resolve_model_to_local_path(str(wrapper.vllm_config.model_config.model)) + resolve_model_to_local_path(str(wrapper.vllm_config.model_config.model)) ): with safe_open(str(shard_path), framework="pt", device="cpu") as handle: for key in handle.keys(): diff --git a/vllm_omni/model_executor/models/ming_tts/speaker_extractor.py b/vllm_omni/model_executor/models/ming_tts/speaker_extractor.py index 8e9f69068..6b6da71fe 100644 --- a/vllm_omni/model_executor/models/ming_tts/speaker_extractor.py +++ b/vllm_omni/model_executor/models/ming_tts/speaker_extractor.py @@ -3,16 +3,13 @@ # Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/spkemb_extractor.py from __future__ import annotations -import os import threading from typing import TYPE_CHECKING, Any import torch -import torchaudio from vllm.logger import init_logger -from vllm.multimodal.media.audio import load_audio -from vllm_omni.model_executor.models.ming_flash_omni.spk_embedding import SpkembExtractor +from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor from .audio_prep import coerce_prompt_waveform, coerce_speaker_embeddings from .constants import ( @@ -29,48 +26,9 @@ logger = init_logger(__name__) _EXTRACTOR_LOAD_LOCK = threading.Lock() -def resolve_model_to_local_path(model): - if os.path.isdir(model): - return model - - from huggingface_hub import snapshot_download - - return snapshot_download(model) - - -class MingSpeakerEmbeddingExtractor: - def __init__(self, model, target_sr=16000): - local_model_path = resolve_model_to_local_path(model) - campplus_path = os.path.join(local_model_path, "campplus.onnx") - if not os.path.exists(campplus_path): - raise RuntimeError(f"Missing Ming speaker extractor model: {campplus_path}") - - self.target_sr = int(target_sr) - self._core = SpkembExtractor(campplus_path, target_sr=self.target_sr) - - def extract_from_waveform(self, waveform, sample_rate): - if not isinstance(waveform, torch.Tensor): - waveform = torch.as_tensor(waveform) - - tensor = waveform.detach().to(torch.float32) - if tensor.ndim == 1: - tensor = tensor.unsqueeze(0) - if int(sample_rate) != self.target_sr: - tensor = torchaudio.transforms.Resample(orig_freq=int(sample_rate), new_freq=self.target_sr)(tensor) - - embedding = self._core._extract_spk_embedding(tensor) - return embedding.squeeze(0).to(dtype=torch.float32) - - def extract_from_file(self, audio_path): - audio, sample_rate = load_audio(audio_path, sr=None, mono=False) - waveform = torch.from_numpy(audio).contiguous() - return self.extract_from_waveform(waveform, sample_rate) - - def extract_many(self, audio_paths): - return [self.extract_from_file(path) for path in audio_paths] - - -def _load_speaker_extractor(wrapper: MingTTSForConditionalGeneration) -> MingSpeakerEmbeddingExtractor: +# Ming's bundle ships the CAM++ artifact as `campplus.onnx` in the model +# directory, which is the shared extractor's default. +def _load_speaker_extractor(wrapper: MingTTSForConditionalGeneration) -> SpeakerEmbeddingExtractor: extractor = getattr(wrapper, "_speaker_extractor", None) if extractor is not None: return extractor @@ -79,7 +37,7 @@ def _load_speaker_extractor(wrapper: MingTTSForConditionalGeneration) -> MingSpe extractor = getattr(wrapper, "_speaker_extractor", None) if extractor is None: model_path = wrapper.vllm_config.model_config.model - extractor = MingSpeakerEmbeddingExtractor(model_path, target_sr=16000) + extractor = SpeakerEmbeddingExtractor(model_path, target_sr=16000) wrapper._speaker_extractor = extractor return extractor @@ -174,8 +132,6 @@ def _resolve_speaker_embeddings( __all__ = [ - "MingSpeakerEmbeddingExtractor", "_load_speaker_extractor", "_resolve_speaker_embeddings", - "resolve_model_to_local_path", ]