diff --git a/pyproject.toml b/pyproject.toml index 2559bcd5b..4af374589 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,8 +70,7 @@ dev = [ "FlagEmbedding==1.4.0", # for hunyuan_image3 tests "kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0) "s3tokenizer==0.3.0", # for tts test - "step-audio2==1.0.0", # Step-Audio2 Token2Wav (flashcosyvoice / hyperpyyaml) - "stepaudio2-minicpmo", # for MiniCPM-o 4.5 talker vocoder (stepaudio2-minicpmo Token2wav) + "step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml) ] demo = [ diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py index f2d96be0d..5ffde9e15 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py @@ -60,9 +60,12 @@ class BatchedToken2Wav(nn.Module): asset loader and prompt feature extractor. """ - def __init__(self, token2wav: Any): + def __init__(self, token2wav: Any, trt_stepper: Any | None = None): super().__init__() self._token2wav = token2wav + # Optional TrtDiTStepper (step_audio2_dit_trt): replaces only the + # per-timestep DiT estimator call; encoder and HiFT stay on torch. + self._trt_stepper = trt_stepper self.flow = token2wav.flow self.hift = token2wav.hift hift_parameter = next(self.hift.parameters(), None) @@ -200,6 +203,17 @@ class BatchedToken2Wav(nn.Module): cnn_cache: torch.Tensor | None, att_cache: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if self._trt_stepper is not None: + out, new_cnn, new_att = self._trt_stepper.step( + x=x, + mu=mu, + t=time, + spks=speakers, + cond=cond, + cnn_cache=cnn_cache, + att_cache=att_cache, + ) + return out.to(mu.dtype), new_cnn, new_att time_embedding = estimator.t_embedder(time).unsqueeze(1) width = int(x.shape[-1]) speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index 92eccb552..019ab0594 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -727,19 +727,17 @@ class MiniCPMO45Code2Wav(nn.Module): if self.backend is not None: return + # In-tree adapter over StepAudio2Token2WavCore on every platform. It + # matches the external `stepaudio2-minicpmo` Token2wav bit-for-bit on + # CUDA (same cosyvoice2 flow/DiT modules and weights) while dropping + # that package's hard-coded `.cuda()` calls, and auto-applies the + # Ascend fixes (HiFT linear downsample, DiT mask expand, MATH SDPA) + # on NPU. + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_token2wav import ( + MiniCPMO45Token2wav as Token2wav, + ) from vllm_omni.platforms import current_omni_platform - if current_omni_platform.is_npu(): - # NPU/Ascend: the external `stepaudio2` package hard-codes `.cuda()`, - # so use the in-tree NPU-aware adapter instead. It delegates to - # StepAudio2Token2WavCore, which auto-applies the Ascend fixes - # (HiFT linear downsample, DiT mask expand, MATH SDPA) on NPU. - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_token2wav import ( - MiniCPMO45Token2wav as Token2wav, - ) - else: - from stepaudio2.token2wav import Token2wav - extra = self._extra_config() # Hub repo ids only need to become local directories once the vocoder # assets are actually read; unit tests construct this model with fake @@ -765,4 +763,26 @@ class MiniCPMO45Code2Wav(nn.Module): ) finally: torch.set_default_dtype(previous_dtype) - self.backend = BatchedToken2Wav(token2wav) + + trt_stepper = None + use_trt = bool(extra.get("token2wav_trt", False)) or os.environ.get("MINICPMO_TOKEN2WAV_TRT", "") == "1" + # TensorRT is CUDA-only; other platforms ignore the toggle. + if use_trt and current_omni_platform.is_cuda(): + from vllm_omni.model_executor.models.step_audio2.step_audio2_dit_trt import build_dit_trt_stepper + + dtype_name = str( + extra.get("token2wav_trt_dtype", os.environ.get("MINICPMO_TOKEN2WAV_TRT_DTYPE", "fp16")) + ).lower() + trt_dtype = torch.float32 if dtype_name in ("fp32", "float32") else torch.float16 + max_batch = int(extra.get("token2wav_trt_max_batch", 16)) + device = next(token2wav.flow.parameters()).device + trt_stepper = build_dit_trt_stepper( + token2wav.flow.decoder.estimator, + device=device, + dtype=trt_dtype, + max_batch=max_batch, + ) + logger.info("MiniCPM-o Code2Wav: DiT estimator running on TensorRT (%s)", dtype_name) + token2wav.enable_trt_spk_embedding() + logger.info("MiniCPM-o Code2Wav: campplus speaker embedding running on TensorRT") + self.backend = BatchedToken2Wav(token2wav, trt_stepper=trt_stepper) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py index c6221cb24..9451c2560 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py @@ -120,6 +120,10 @@ class MiniCPMO45Token2wav: """Delegate prompt feature extraction to the wrapped core.""" return self._core._prepare_prompt(prompt_wav) + def enable_trt_spk_embedding(self) -> None: + """Run the campplus speaker-embedding model on TensorRT.""" + self._core.enable_trt_spk_embedding() + def __call__(self, generated_speech_tokens, prompt_wav) -> bytes: """One-shot tokens → 24 kHz WAV bytes.""" return self._core.forward( diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_dit_trt.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_dit_trt.py new file mode 100644 index 000000000..6feefd128 --- /dev/null +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_dit_trt.py @@ -0,0 +1,375 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TensorRT acceleration for the Step-Audio2 / CosyVoice2 flow DiT estimator. + +Ports stepfun-ai/Step-Audio2#70 into vLLM-Omni. The upstream PR patches +``cosyvoice2.flow.flow_matching`` to dispatch between a torch estimator and a +raw TRT execution context; here the swap happens one level up instead — the +batched Code2Wav driver (``BatchedToken2Wav._estimator_step``) already owns +the per-timestep estimator call, so a :class:`TrtDiTStepper` replaces only +that call and the pip-installed ``cosyvoice2`` package stays untouched. + +Differences from the upstream PR: + +* ONNX export is a function (:func:`export_dit_chunk_onnx`) that wraps the + already-loaded torch ``DiT`` module and runs at engine-build time — no + separate ``tools/export_onnx_*.py`` step. +* The engine executes on the caller's current CUDA stream, so no context + pool, no private stream, and no ``synchronize()`` per call (a private + stream without synchronization produces silent garbage audio; with + synchronization it costs two round-trips per denoise step). +* Cache tensors are exchanged as stacked ``(depth, 2B, ...)`` tensors, the + same layout ``BatchedToken2Wav`` already uses. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import torch +import torch.nn as nn +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_TRT_CACHE_ENV = "MINICPMO_TOKEN2WAV_TRT_CACHE" + + +class _DiTChunkWrapper(nn.Module): + """ONNX-exportable streaming ``forward_chunk`` with explicit stacked caches. + + Mirrors ``DiT.forward_chunk`` (cosyvoice2 pip package) but takes and + returns caches as plain tensors instead of registered buffers, matching + the I/O contract of the TRT engine and of ``BatchedToken2Wav``. + """ + + def __init__(self, dit: nn.Module): + super().__init__() + self.dit = dit + + def forward( + self, + x: torch.Tensor, # (2B, mel, T) + mu: torch.Tensor, # (2B, mel, T) + t: torch.Tensor, # (2B,) + spks: torch.Tensor, # (2B, spk) + cond: torch.Tensor, # (2B, mel, T) + cnn_cache: torch.Tensor, # (depth, 2B, c1+c2, kernel-1) + att_cache: torch.Tensor, # (depth, 2B, heads, L, head_dim*2) + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + t_emb = self.dit.t_embedder(t).unsqueeze(1) + spks_f = spks.unsqueeze(-1).expand(-1, -1, x.shape[-1]) + h = torch.cat((x, mu, spks_f, cond), dim=1).transpose(1, 2) + h = self.dit.in_proj(h) + new_cnn = [] + new_att = [] + for i, block in enumerate(self.dit.blocks): + h, block_cnn, block_att = block.forward_chunk(h, t_emb, cnn_cache[i], att_cache[i], None) + new_cnn.append(block_cnn) + new_att.append(block_att) + h = self.dit.final_layer(h, t_emb) + return h.transpose(1, 2), torch.stack(new_cnn), torch.stack(new_att) + + +def _dit_dims(dit: nn.Module) -> dict[str, int]: + block0 = dit.blocks[0] + return { + "depth": len(dit.blocks), + "mel": int(dit.out_channels), + "spk": int(dit.in_channels) - 3 * int(dit.out_channels), + "heads": int(block0.attn.num_heads), + "att_dim": int(block0.attn.head_dim) * 2, + "cnn_ch": int(block0.conv.in_channels + block0.conv.out_channels), + "cnn_w": int(block0.conv.block[1].causal_padding[0]), + } + + +def export_dit_chunk_onnx( + dit: nn.Module, + onnx_path: str | Path, + device: torch.device, + dtype: torch.dtype = torch.float16, +) -> None: + """Export the streaming DiT chunk graph to ONNX in ``dtype``. + + Fused-in replacement for the upstream ``tools/export_onnx_streaming_token2wav.py``. + TensorRT >= 10 deprecates weakly-typed networks (TRT 11 removes + ``BuilderFlag.FP16`` outright), so the engine precision comes from the + ONNX graph itself. The DiT is deep-copied for the cast — the live torch + estimator keeps its original dtype. + """ + import copy + + dims = _dit_dims(dit) + export_dit = copy.deepcopy(dit).to(device=device, dtype=dtype).eval() + wrapper = _DiTChunkWrapper(export_dit).eval() + batch, chunk, cache_len = 2, 48, 16 + sample = ( + torch.randn(batch, dims["mel"], chunk, device=device, dtype=dtype), + torch.randn(batch, dims["mel"], chunk, device=device, dtype=dtype), + torch.rand(batch, device=device, dtype=dtype), + torch.randn(batch, dims["spk"], device=device, dtype=dtype), + torch.randn(batch, dims["mel"], chunk, device=device, dtype=dtype), + torch.zeros(dims["depth"], batch, dims["cnn_ch"], dims["cnn_w"], device=device, dtype=dtype), + torch.randn(dims["depth"], batch, dims["heads"], cache_len, dims["att_dim"], device=device, dtype=dtype), + ) + input_names = ["x", "mu", "t", "spks", "cond", "cnn_cache", "att_cache"] + output_names = ["out", "new_cnn_cache", "new_att_cache"] + dynamic_axes = { + "x": {0: "b", 2: "t"}, + "mu": {0: "b", 2: "t"}, + "t": {0: "b"}, + "spks": {0: "b"}, + "cond": {0: "b", 2: "t"}, + "cnn_cache": {1: "b"}, + "att_cache": {1: "b", 3: "l"}, + "out": {0: "b", 2: "t"}, + "new_cnn_cache": {1: "b"}, + "new_att_cache": {1: "b", 3: "l_out"}, + } + logger.info("Exporting DiT chunk ONNX to %s", onnx_path) + onnx_path = Path(onnx_path) + onnx_path.parent.mkdir(parents=True, exist_ok=True) + # Process-unique temp name + atomic rename, so concurrent first builds on + # a shared cache dir never see (or clobber) a partially written file. + tmp = onnx_path.with_suffix(f".tmp.{os.getpid()}") + try: + with torch.no_grad(): + torch.onnx.export( + wrapper, + sample, + str(tmp), + input_names=input_names, + output_names=output_names, + dynamic_axes=dynamic_axes, + opset_version=17, + dynamo=False, + ) + os.replace(tmp, onnx_path) + finally: + tmp.unlink(missing_ok=True) + + +def convert_onnx_to_trt( + onnx_path: str | Path, + plan_path: str | Path, + dims: dict[str, int], + *, + max_batch: int = 16, + max_len: int = 3000, + max_cache_len: int = 1000, +) -> None: + """Build a TRT engine from the exported ONNX (fused-in from upstream tools). + + Engine precision is strongly typed from the ONNX graph itself — export + the graph in the desired dtype (TRT 11 removed ``BuilderFlag.FP16``). + """ + import tensorrt as trt + + logger.info("Building DiT TRT engine at %s — this can take a few minutes", plan_path) + trt_logger = trt.Logger(trt.Logger.INFO) + builder = trt.Builder(trt_logger) + # TRT >= 10 removed EXPLICIT_BATCH (networks are always explicit-batch). + flag = getattr(trt.NetworkDefinitionCreationFlag, "EXPLICIT_BATCH", None) + network = builder.create_network(1 << int(flag) if flag is not None else 0) + parser = trt.OnnxParser(network, trt_logger) + with open(onnx_path, "rb") as f: + if not parser.parse(f.read()): + errors = [str(parser.get_error(i)) for i in range(parser.num_errors)] + raise ValueError(f"failed to parse {onnx_path}: {errors}") + + config = builder.create_builder_config() + + depth, mel, spk = dims["depth"], dims["mel"], dims["spk"] + heads, att_dim = dims["heads"], dims["att_dim"] + cnn_ch, cnn_w = dims["cnn_ch"], dims["cnn_w"] + opt_b = min(4, max_batch) + shapes = { + "x": [(2, mel, 2), (opt_b, mel, 500), (max_batch, mel, max_len)], + "mu": [(2, mel, 2), (opt_b, mel, 500), (max_batch, mel, max_len)], + "t": [(2,), (opt_b,), (max_batch,)], + "spks": [(2, spk), (opt_b, spk), (max_batch, spk)], + "cond": [(2, mel, 2), (opt_b, mel, 500), (max_batch, mel, max_len)], + "cnn_cache": [ + (depth, 2, cnn_ch, cnn_w), + (depth, opt_b, cnn_ch, cnn_w), + (depth, max_batch, cnn_ch, cnn_w), + ], + "att_cache": [ + (depth, 2, heads, 0, att_dim), + (depth, opt_b, heads, 500, att_dim), + (depth, max_batch, heads, max_cache_len, att_dim), + ], + } + profile = builder.create_optimization_profile() + for name, (mn, op, mx) in shapes.items(): + profile.set_shape(name, mn, op, mx) + config.add_optimization_profile(profile) + + engine_bytes = builder.build_serialized_network(network, config) + if engine_bytes is None: + raise RuntimeError("TensorRT engine build failed") + plan_path = Path(plan_path) + plan_path.parent.mkdir(parents=True, exist_ok=True) + # Process-unique temp name: concurrent first builds (multiple ranks on a + # shared cache dir) must not clobber each other's partial writes; the + # final os.replace stays atomic and last-writer-wins with equal content. + tmp = plan_path.with_suffix(f".tmp.{os.getpid()}") + try: + with open(tmp, "wb") as f: + f.write(engine_bytes) + os.replace(tmp, plan_path) + finally: + tmp.unlink(missing_ok=True) + logger.info("DiT TRT engine written to %s", plan_path) + + +class TrtDiTStepper: + """Drop-in TRT replacement for one DiT denoise step. + + ``step`` matches the tensors ``BatchedToken2Wav._decode_cfm`` has in hand + per timestep. Executes on the caller's current CUDA stream — same-stream + ordering makes explicit synchronization unnecessary. + """ + + def __init__( + self, + plan_path: str | Path, + dims: dict[str, int], + device: torch.device, + dtype: torch.dtype, + *, + max_batch: int = 16, + max_len: int = 3000, + max_cache_len: int = 1000, + ): + import tensorrt as trt + + self.dims = dims + self.device = device + self.dtype = dtype + # Optimization-profile bounds the engine was built with; step() checks + # them so an out-of-profile call fails with a readable error instead + # of an opaque execute_async_v3 failure. + self.max_batch = int(max_batch) + self.max_len = int(max_len) + self.max_cache_len = int(max_cache_len) + with open(plan_path, "rb") as f: + self.engine = trt.Runtime(trt.Logger(trt.Logger.INFO)).deserialize_cuda_engine(f.read()) + if self.engine is None: + raise RuntimeError(f"failed to deserialize TRT engine {plan_path}") + self.context = self.engine.create_execution_context() + if self.context is None: + raise RuntimeError("failed to create TRT execution context") + # TRT rejects null bindings; zero-length cache tensors share this address. + self._dummy = torch.zeros(1, device=device, dtype=dtype) + + def _addr(self, tensor: torch.Tensor) -> int: + return self._dummy.data_ptr() if tensor.numel() == 0 else tensor.data_ptr() + + def step( + self, + *, + x: torch.Tensor, + mu: torch.Tensor, + t: torch.Tensor, + spks: torch.Tensor, + cond: torch.Tensor, + cnn_cache: torch.Tensor | None, + att_cache: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + dims = self.dims + batch = int(x.shape[0]) + chunk = int(x.shape[2]) + cache_frames = int(att_cache.shape[3]) if att_cache is not None else 0 + if batch > self.max_batch or chunk > self.max_len or cache_frames > self.max_cache_len: + raise ValueError( + f"DiT TRT engine profile exceeded: batch={batch} (max {self.max_batch}, " + f"= token2wav_trt_max_batch, CFG included), chunk={chunk} frames " + f"(max {self.max_len}), att cache={cache_frames} frames (max {self.max_cache_len}). " + "Rebuild the engine with a larger token2wav_trt_max_batch or disable token2wav_trt." + ) + if cnn_cache is None: + cnn_cache = torch.zeros( + (dims["depth"], batch, dims["cnn_ch"], dims["cnn_w"]), device=self.device, dtype=self.dtype + ) + if att_cache is None: + att_cache = torch.zeros( + (dims["depth"], batch, dims["heads"], 0, dims["att_dim"]), device=self.device, dtype=self.dtype + ) + cache_len = int(att_cache.shape[3]) + + inputs = { + "x": x.to(self.dtype).contiguous(), + "mu": mu.to(self.dtype).contiguous(), + "t": t.to(self.dtype).contiguous(), + "spks": spks.to(self.dtype).contiguous(), + "cond": cond.to(self.dtype).contiguous(), + "cnn_cache": cnn_cache.to(self.dtype).contiguous(), + "att_cache": att_cache.to(self.dtype).contiguous(), + } + out = torch.empty((batch, dims["mel"], chunk), device=self.device, dtype=self.dtype) + new_cnn = torch.empty_like(inputs["cnn_cache"]) + new_att = torch.empty( + (dims["depth"], batch, dims["heads"], cache_len + chunk, dims["att_dim"]), + device=self.device, + dtype=self.dtype, + ) + + for name, tensor in inputs.items(): + self.context.set_input_shape(name, tuple(tensor.shape)) + self.context.set_tensor_address(name, self._addr(tensor)) + self.context.set_tensor_address("out", out.data_ptr()) + self.context.set_tensor_address("new_cnn_cache", new_cnn.data_ptr()) + self.context.set_tensor_address("new_att_cache", self._addr(new_att)) + + stream = torch.cuda.current_stream(self.device) + if not self.context.execute_async_v3(stream.cuda_stream): + raise RuntimeError("TRT DiT execution failed") + return out, new_cnn, new_att + + +def _default_cache_dir() -> Path: + env = os.environ.get(_TRT_CACHE_ENV, "").strip() + if env: + return Path(env) + hf_home = os.environ.get("HF_HOME", "").strip() + base = Path(hf_home) if hf_home else Path.home() / ".cache" / "huggingface" + return base / "vllm_omni" / "token2wav_trt" + + +def build_dit_trt_stepper( + dit: nn.Module, + *, + device: torch.device | str = "cuda", + dtype: torch.dtype = torch.float16, + max_batch: int = 16, + cache_dir: str | Path | None = None, +) -> TrtDiTStepper: + """Export (if needed), build (if needed) and load the DiT TRT engine. + + Artifacts are cached under ``cache_dir`` (default: ``$MINICPMO_TOKEN2WAV_TRT_CACHE`` + or ``$HF_HOME/vllm_omni/token2wav_trt``), keyed by dtype/batch/TRT + version/GPU architecture — serialized plans are SM-specific, so a cache + shared across heterogeneous machines must not reuse another arch's plan. + """ + import tensorrt as trt + + device = torch.device(device) + dims = _dit_dims(dit) + cache = Path(cache_dir) if cache_dir is not None else _default_cache_dir() + cache.mkdir(parents=True, exist_ok=True) + dtype_tag = str(dtype).replace("torch.", "") + major, minor = torch.cuda.get_device_capability(device) + sm_tag = f"sm{major}{minor}" + onnx_path = cache / f"dit_chunk.d{dims['depth']}.{dtype_tag}.onnx" + plan_path = cache / (f"dit_chunk.d{dims['depth']}.{dtype_tag}.b{max_batch}.{sm_tag}.trt{trt.__version__}.plan") + + if not plan_path.is_file() or plan_path.stat().st_size == 0: + if not onnx_path.is_file() or onnx_path.stat().st_size == 0: + export_dit_chunk_onnx(dit, onnx_path, device, dtype) + convert_onnx_to_trt(onnx_path, plan_path, dims, max_batch=max_batch) + else: + logger.info("Reusing cached DiT TRT engine %s", plan_path) + return TrtDiTStepper(plan_path, dims, device, dtype, max_batch=max_batch) diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py index 026276c52..3923933f6 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py @@ -178,6 +178,28 @@ class StepAudio2Token2WavCore(nn.Module): self._ensure_models_loaded() return self._hift + def _forward_spk_embedding(self, spk_feat: torch.Tensor) -> torch.Tensor: + """Run campplus on ``[T, 80]`` fbank features; returns ``[1, 192]`` on ``self.device``.""" + spk_model = self.spk_model + if isinstance(spk_model, onnxruntime.InferenceSession): + return torch.tensor( + spk_model.run(None, {spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()})[0], + device=self.device, + ) + # CampplusTRT (or any callable taking [T, 80] and returning [1, 192]). + return spk_model(spk_feat).to(self.device) + + def enable_trt_spk_embedding(self) -> None: + """Swap the onnxruntime campplus session for the shared TensorRT engine. + + Reuses the CosyVoice3 CAM++ engine wrapper — MiniCPM-o / Step-Audio2 + ship the same campplus.onnx architecture and I/O contract. + """ + from vllm_omni.model_executor.models.cosyvoice3.speaker_embedding_trt import get_campplus_trt + + self._ensure_models_loaded() + self._spk_model = get_campplus_trt(f"{self.model_path}/campplus.onnx", device=self.device) + def _prepare_prompt(self, prompt_wav: str): """Prepare prompt audio for conditioning""" # Prefer soundfile/librosa path to avoid torchaudio->torchcodec runtime coupling. @@ -202,10 +224,7 @@ class StepAudio2Token2WavCore(nn.Module): spk_feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000) spk_feat = spk_feat - spk_feat.mean(dim=0, keepdim=True) - spk_emb = torch.tensor( - self.spk_model.run(None, {self.spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()})[0], - device=self.device, - ) + spk_emb = self._forward_spk_embedding(spk_feat) # 24 kHz branch: mel spectrogram for flow model conditioning # Must resample from the ORIGINAL audio, not the 16 kHz version.