[Ming] Cleanup Ming-family shared modules (#6119)

Signed-off-by: Yuanheng <jonathan.zhaoyh@gmail.com>
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>

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)

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,

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()

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():

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,
+    )

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
+"""

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):

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):

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):

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):

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):

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)

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"]

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."""

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)

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:

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

vllm_omni/model_executor/models/ming_tts/fm/__init__.py

@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

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;

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():

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",
 ]