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

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

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

@@ -0,0 +1,67 @@
+# 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/fm/dit.py
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from x_transformers.x_transformers import RotaryEmbedding
+
+from vllm_omni.model_executor.models.common.ming.dit_blocks import DiTBlock, FinalLayer
+
+
+class Aggregator(nn.Module):
+    def __init__(
+        self,
+        in_channels=4,
+        hidden_size=1152,
+        depth=28,
+        num_heads=16,
+        mlp_ratio=4.0,
+        llm_input_dim=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, mask=None):
+        if x.ndim != 3:
+            raise ValueError(f"Expected x rank-3 [Batch, Time, Dimension], got {tuple(x.shape)}")
+        if x.shape[-1] != self.in_channels:
+            raise ValueError(f"x feature dim mismatch: got {x.shape[-1]}, expected {self.in_channels}")
+
+        # [Batch, Time, Dimension] -> [Batch, Time, Hidden].
+        x = self.x_embedder(x)
+        cls_embed = self.word_embedder(torch.zeros((x.shape[0], 1), dtype=torch.long, device=x.device))
+        # Prepend a learned CLS token: [Batch, Time, Hidden] -> [Batch, Time + 1, Hidden].
+        x = torch.cat([cls_embed, x], dim=1)
+
+        rope = self.rotary_embed.forward_from_seq_len(x.shape[1])
+        if mask is not None:
+            if mask.ndim != 2:
+                raise ValueError(f"Expected mask rank-2 [Batch, Time], got {tuple(mask.shape)}")
+            if mask.shape[0] != x.shape[0] or mask.shape[1] != x.shape[1] - 1:
+                raise ValueError(
+                    f"Mask shape mismatch: got {tuple(mask.shape)}, expected {(x.shape[0], x.shape[1] - 1)}"
+                )
+            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)
+        # Keep the CLS projection only: [Batch, Time + 1, Hidden] -> [Batch, 1, Hidden].
+        return x[:, :1, :]
+
+
+__all__ = ["Aggregator"]

vllm_omni/model_executor/models/common/ming/cfm_cudagraph.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CUDAGraph-accelerated CFM (flow-matching) diffusion head for Ming-TTS.
+
+The per-AR-step flow head runs ``steps`` ODE iterations, each a tiny DiT forward
+(history+patch tokens, CFG cond+null batched). In eager mode these are dozens of
+micro kernels per step and the wall time is dominated by kernel-launch / Python
+overhead, not GPU math. Capturing the whole sampling loop + Aggregator + stop
+head into one CUDA graph removes that overhead.
+
+Adapted from adaspeech/ming_flash ``CFMGraphExecutor``. Captured for batch=1
+(Ming Stage-0 runs ``max_num_seqs: 1``); other batch sizes fall back to eager.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.common.ming.cfm_solver import build_timesteps
+
+
+class CFMSampler(nn.Module):
+    """Unrolled CFM ODE solver, graph-capturable (no ``torch.randn``/``.item()``).
+
+    All randomness (``y0``, ``sde_rnd``) and the timestep schedule ``t`` are
+    passed in as float32 tensors so the loop is deterministic during replay and
+    the integration matches the eager float32 path. Equivalent to eager
+    ``Solver.integrate`` for ``temperature==0`` (the TTS default).
+    """
+
+    def __init__(self, dit_model: nn.Module, steps: int = 10):
+        super().__init__()
+        self.model = dit_model
+        self.steps = steps
+
+    @torch.no_grad()
+    def forward(self, llm_cond, lat_cond, y0, t, sde_args, sde_rnd):
+        # y0 / t / sde_rnd / sde_args are float32; integrate in float32 and cast
+        # only the DiT inputs to the model dtype, mirroring eager FlowLoss.sample.
+        model_dtype = llm_cond.dtype
+        for step in range(self.steps):
+            dt = t[step + 1] - t[step]
+            pred_cfg = self.model.forward_with_cfg(
+                x=y0.to(model_dtype),
+                t=t[step].to(model_dtype),
+                c=llm_cond,
+                latent_history=lat_cond,
+                cfg_scale=sde_args[0],
+                patch_size=y0.shape[1],
+            )
+            pred, null_pred = torch.chunk(pred_cfg, 2, dim=0)
+            flow = (pred + (pred - null_pred) * sde_args[0]).float()
+            y0 = y0 + flow * dt
+            if step < self.steps - 1:
+                y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step]
+        return y0
+
+
+class CFMGraphExecutor:
+    """Capture CFM sampling + Aggregator + stop head into one CUDA graph.
+
+    First ``execute`` captures the graph; later calls copy inputs into the
+    static placeholders and ``replay()``. Runtime controls (cfg/sigma/temp) are
+    fed via the ``sde_args`` placeholder, so they are tunable without recapture.
+    """
+
+    def __init__(
+        self,
+        cfm_sampler: CFMSampler,
+        aggregator: nn.Module,
+        stop_head: nn.Module,
+        *,
+        patch_size: int,
+        latent_dim: int,
+        steps: int = 10,
+    ):
+        self.cfm_sampler = cfm_sampler
+        self.aggregator = aggregator
+        self.stop_head = stop_head
+        self.patch_size = int(patch_size)
+        self.latent_dim = int(latent_dim)
+        self.steps = int(steps)
+        self.initialized = False
+        self.graph: torch.cuda.CUDAGraph | None = None
+
+    def _timesteps(self, device):
+        # Raw EPSS schedule + sway in float32, identical to eager build_timesteps.
+        return build_timesteps(self.steps, device=device, dtype=torch.float32, use_epss=True, sway_sampling_coef=-1.0)
+
+    @torch.no_grad()
+    def execute(self, llm_cond, his_lat, cfg_strength, sigma, temperature):
+        bsz = llm_cond.shape[0]
+        device = llm_cond.device
+        # ODE state / noise / timesteps are float32 to match the eager path.
+        y0 = torch.randn((bsz, self.patch_size, self.latent_dim), device=device, dtype=torch.float32)
+        t = self._timesteps(device)
+        sde_rnd = torch.randn((self.steps, bsz, self.patch_size, self.latent_dim), device=device, dtype=torch.float32)
+
+        if not self.initialized:
+            self._initialize_graph(llm_cond, his_lat, y0, t, sde_rnd)
+
+        self._llm_cond_ph.copy_(llm_cond)
+        self._his_lat_ph.copy_(his_lat)
+        self._y0_ph.copy_(y0)
+        self._t_ph.copy_(t)
+        self._sde_args_ph[0] = float(cfg_strength)
+        self._sde_args_ph[1] = float(sigma)
+        self._sde_args_ph[2] = float(temperature)
+        self._sde_rnd_ph.copy_(sde_rnd)
+
+        self.graph.replay()
+
+        return self._gen_lat_ph.clone(), self._embeds_ph.clone(), self._stop_ph.clone()
+
+    def _initialize_graph(self, llm_cond, his_lat, y0, t, sde_rnd):
+        model_dtype = llm_cond.dtype
+        self._llm_cond_ph = torch.empty_like(llm_cond)
+        self._his_lat_ph = torch.empty_like(his_lat)
+        self._y0_ph = torch.empty_like(y0)  # float32
+        self._t_ph = torch.empty_like(t)  # float32
+        self._sde_args_ph = torch.empty(3, device=llm_cond.device, dtype=torch.float32)
+        self._sde_rnd_ph = torch.empty_like(sde_rnd)  # float32
+
+        self._llm_cond_ph.copy_(llm_cond)
+        self._his_lat_ph.copy_(his_lat)
+        self._y0_ph.copy_(y0)
+        self._t_ph.copy_(t)
+        self._sde_args_ph[0], self._sde_args_ph[1], self._sde_args_ph[2] = 2.0, 0.25, 0.0
+        self._sde_rnd_ph.copy_(sde_rnd)
+
+        def _run():
+            # Solver runs in float32; cast the final latent to the model dtype
+            # for the Aggregator / stop head (matches eager ``.to(dtype)``).
+            gen_lat = self.cfm_sampler(
+                self._llm_cond_ph, self._his_lat_ph, self._y0_ph, self._t_ph, self._sde_args_ph, self._sde_rnd_ph
+            ).to(model_dtype)
+            embeds = self.aggregator(gen_lat)
+            stop = self.stop_head(self._llm_cond_ph[:, -1, :]).softmax(dim=-1)
+            return gen_lat, embeds, stop
+
+        # Warm up on a side stream to stabilize allocations before capture.
+        warmup = torch.cuda.Stream()
+        warmup.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(warmup):
+            for _ in range(3):
+                _run()
+        torch.cuda.current_stream().wait_stream(warmup)
+
+        self.graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.graph):
+            self._gen_lat_ph, self._embeds_ph, self._stop_ph = _run()
+
+        self.initialized = True

vllm_omni/model_executor/models/common/ming/cfm_dit.py

@@ -0,0 +1,103 @@
+# 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/fm/dit.py
+
+import torch
+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_blocks import CondEmbedder, DiTBlock, FinalLayer
+
+
+class DiT(nn.Module):
+    def __init__(
+        self,
+        in_channels=4,
+        hidden_size=1024,
+        depth=28,
+        num_heads=16,
+        mlp_ratio=4.0,
+        llm_cond_dim=896,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.num_heads = num_heads
+        self.t_embedder = DiTTimestepEmbedding(hidden_size)
+        self.x_embedder = nn.Linear(in_channels, hidden_size)
+        self.c_embedder = CondEmbedder(llm_cond_dim, 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, self.out_channels)
+
+    def forward(self, x, t, c, latent_history, mask=None):
+        if x.ndim != 3:
+            raise ValueError(f"Expected x rank-3 [Batch, Time, Dimension], got {tuple(x.shape)}")
+        if latent_history.ndim != 3:
+            raise ValueError(
+                f"Expected latent_history rank-3 [Batch, Time, Dimension], got {tuple(latent_history.shape)}"
+            )
+        if c.ndim != 3:
+            raise ValueError(f"Expected conditioning rank-3 [Batch, Time, Dimension], got {tuple(c.shape)}")
+        if x.shape[0] != latent_history.shape[0] or x.shape[0] != c.shape[0]:
+            raise ValueError(
+                "Batch mismatch across x, conditioning, and latent_history: "
+                f"{x.shape[0]}, {c.shape[0]}, {latent_history.shape[0]}"
+            )
+        if x.shape[-1] != self.in_channels:
+            raise ValueError(f"x feature dim mismatch: got {x.shape[-1]}, expected {self.in_channels}")
+        if latent_history.shape[-1] != self.in_channels:
+            raise ValueError(
+                f"latent_history feature dim mismatch: got {latent_history.shape[-1]}, expected {self.in_channels}"
+            )
+        if t.ndim == 0:
+            t = t.reshape(1)
+        if t.ndim != 1:
+            raise ValueError(f"Expected timestep rank-1 [Batch], got {tuple(t.shape)}")
+        if t.shape[0] != x.shape[0]:
+            raise ValueError(f"Timestep batch mismatch: got {t.shape[0]}, expected {x.shape[0]}")
+
+        t = self.t_embedder(t).unsqueeze(1)
+        x_now = self.x_embedder(x)
+        x_history = self.x_embedder(latent_history)
+        x = torch.cat([x_history, x_now], dim=1)
+        c = self.c_embedder(c)
+        y = t + c
+        x = torch.cat([y, x], dim=1)
+        rope = self.rotary_embed.forward_from_seq_len(x.shape[1])
+
+        if mask is not None:
+            if mask.ndim != 2:
+                raise ValueError(f"Expected mask rank-2 [Batch, Time], got {tuple(mask.shape)}")
+            if mask.shape[0] != x_now.shape[0] or mask.shape[1] != x_now.shape[1]:
+                raise ValueError(
+                    f"Mask shape mismatch: got {tuple(mask.shape)}, expected {(x_now.shape[0], x_now.shape[1])}"
+                )
+            mask_pad = mask.clone().detach()[:, :1].expand(-1, x_history.shape[1] + c.shape[1])
+            mask = torch.cat([mask_pad, mask], dim=-1)
+        for block in self.blocks:
+            x = block(x, mask, rope)
+        x = self.final_layer(x)
+        return x
+
+    def forward_with_cfg(self, x, t, c, cfg_scale, latent_history, patch_size):
+        # Always run the classifier-free-guidance path (cond + uncond). TTS
+        # inference always uses cfg>1, and branching on ``cfg_scale`` is avoided
+        # because under CUDA-graph capture it is a CUDA tensor — ``bool(cfg_scale)``
+        # would force a device sync and break capture. ``expand`` (not ``repeat``)
+        # avoids a fresh allocation each step.
+        x = torch.cat([x, x], dim=0)
+        latent_history = torch.cat([latent_history, latent_history], dim=0)
+        fake_latent = torch.zeros_like(c)
+        c = torch.cat([c, fake_latent], dim=0)
+        if t.ndim == 0:
+            t = t.reshape(1)
+        t = t.expand(x.shape[0])
+        model_out = self.forward(x, t, c, latent_history)
+        return model_out[:, -patch_size:, :]

vllm_omni/model_executor/models/common/ming/cfm_head.py

@@ -0,0 +1,62 @@
+# 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/fm/flowloss.py
+
+import torch
+import torch.nn as nn
+
+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):
+    """Diffusion Loss"""
+
+    def __init__(self, z_channels, llm_cond_dim, **kwargs):
+        super().__init__()
+        self.z_channels = z_channels
+        # Preserve checkpoint parameter names under flowloss.cfm.model.* while
+        # keeping the inference-only CFM sampling logic local to FlowLoss.
+        self.cfm = nn.Module()
+        self.cfm.add_module("model", DiT(in_channels=z_channels, llm_cond_dim=llm_cond_dim, **kwargs))
+
+    @torch.no_grad()
+    def sample(self, z, latent_history, cfg=2.0, patch_size=1, sigma=0.25, temperature=0):
+        # z: [Batch, Time=1, Dimension]
+        # latent_history: [Batch, History_Time, Dimension]
+        if z.ndim != 3:
+            raise ValueError(f"Expected conditioning rank-3 [Batch, Time, Dimension], got {tuple(z.shape)}")
+        if latent_history.ndim != 3:
+            raise ValueError(
+                f"Expected latent_history rank-3 [Batch, Time, Dimension], got {tuple(latent_history.shape)}"
+            )
+        if z.shape[0] != latent_history.shape[0]:
+            raise ValueError(
+                f"Batch mismatch across conditioning and latent_history: {z.shape[0]}, {latent_history.shape[0]}"
+            )
+        dtype = z.dtype
+        noise = torch.randn(z.shape[0], self.z_channels, patch_size, device=z.device, dtype=torch.float32)
+
+        def fn(t, x):
+            x = x.to(dtype)
+            t = t.to(dtype)
+            if cfg < 1e-5:
+                if t.ndim == 0:
+                    t = t.repeat(x.shape[0])
+                pred = self.cfm.model(x=x, t=t, c=torch.zeros_like(z), latent_history=latent_history)
+                return pred[:, -patch_size:, :].float()
+
+            pred_cfg = self.cfm.model.forward_with_cfg(
+                x=x, t=t, c=z, latent_history=latent_history, cfg_scale=cfg, patch_size=patch_size
+            )
+            pred, null_pred = torch.chunk(pred_cfg, 2, dim=0)
+            return (pred + (pred - null_pred) * cfg).float()
+
+        t = build_timesteps(
+            10,
+            device=noise.device,
+            dtype=noise.dtype,
+            use_epss=True,
+            sway_sampling_coef=-1.0,
+        )
+        return Solver(fn, noise.transpose(1, 2), sigma=sigma, temperature=temperature).integrate(t)[-1].to(dtype)

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/common/ming/spk_embedding.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The vLLM-Omni team.
+# Copyright (c) Ant Group. All rights reserved.
+# Ported from:
+# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/modeling_bailing_talker.py
+
+from __future__ import annotations
+
+import torch
+
+
+class SpkembExtractor:
+    """CAMPPlus ONNX-based speaker embedding extractor (runs on CPU)."""
+
+    def __init__(self, campplus_model: str, target_sr: int = 16000):
+        import onnxruntime
+        import torchaudio.compliance.kaldi as kaldi
+
+        self.kaldi = kaldi
+        option = onnxruntime.SessionOptions()
+        option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+        option.intra_op_num_threads = 2
+        self.campplus_session = onnxruntime.InferenceSession(
+            campplus_model, sess_options=option, providers=["CPUExecutionProvider"]
+        )
+        self.target_sr = target_sr
+
+    def _extract_spk_embedding(self, speech):
+        feat = self.kaldi.fbank(speech, num_mel_bins=80, dither=0, sample_frequency=16000)
+        feat = feat - feat.mean(dim=0, keepdim=True)
+        embedding = (
+            self.campplus_session.run(
+                None,
+                {self.campplus_session.get_inputs()[0].name: feat.unsqueeze(dim=0).cpu().numpy()},
+            )[0]
+            .flatten()
+            .tolist()
+        )
+        embedding = torch.tensor([embedding])
+        return embedding
+
+    def __call__(self, waveform, **kwargs) -> torch.Tensor | None:
+        spk_emb = self._extract_spk_embedding(waveform)
+        return spk_emb