Merge branch 'main' into lance-model-support
This commit is contained in:
224 files changed
+19202
-7026
No files matched your search
@@ -1590,6 +1590,10 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
|
||||
|
||||
self._sampler: Sampler | None = None
|
||||
self._eos_token_id: int = tokenizer.eos_token_id
|
||||
# Lazily built on first sample() call so we can pick up logits.device
|
||||
# without guessing during init. See `sample()` for the comprehension
|
||||
# fast path that uses this.
|
||||
self._blocked_token_ids_tensor: torch.Tensor | None = None
|
||||
|
||||
self._replace_rotary_embeddings()
|
||||
self._patch_moe_blocks()
|
||||
@@ -1992,25 +1996,39 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
|
||||
|
||||
min_score = torch.finfo(logits.dtype).min
|
||||
|
||||
if self._is_comprehension:
|
||||
# Comprehension path is stateless: we only need to mask a fixed
|
||||
# set of blocked token ids on every step. Do it in one batched
|
||||
# index_fill_ instead of a per-(req, id) Python loop of scalar
|
||||
# GPU writes, and skip reading `output_token_ids` from CPU
|
||||
# (unused here) so this branch never forces a D2H sync.
|
||||
if self._blocked_token_ids_tensor is None and self._blocked_token_ids:
|
||||
self._blocked_token_ids_tensor = torch.tensor(
|
||||
sorted(self._blocked_token_ids),
|
||||
dtype=torch.long,
|
||||
device=logits.device,
|
||||
)
|
||||
if self._blocked_token_ids_tensor is not None:
|
||||
logits.index_fill_(-1, self._blocked_token_ids_tensor, min_score)
|
||||
return self._sampler(logits=logits, sampling_metadata=sampling_metadata)
|
||||
|
||||
# Generation path retains the per-request stateful logic for forced
|
||||
# stage-transition tokens and ratio-restriction.
|
||||
for req_idx in range(logits.shape[0]):
|
||||
decoded_tokens: list[int] = (
|
||||
sampling_metadata.output_token_ids[req_idx] if req_idx < len(sampling_metadata.output_token_ids) else []
|
||||
)
|
||||
last_token = decoded_tokens[-1] if decoded_tokens else -1
|
||||
|
||||
if self._is_comprehension:
|
||||
for tid in self._blocked_token_ids:
|
||||
logits[req_idx, tid] = min_score
|
||||
else:
|
||||
forced = self._get_forced_token(decoded_tokens)
|
||||
if forced is not None:
|
||||
logits[req_idx].fill_(min_score)
|
||||
logits[req_idx, forced] = 0
|
||||
elif last_token == self._size_token_id:
|
||||
self._apply_ratio_restriction(logits, req_idx, min_score)
|
||||
elif last_token in self._all_ratio_ids:
|
||||
logits[req_idx].fill_(min_score)
|
||||
logits[req_idx, self._eos_token_id] = 0
|
||||
forced = self._get_forced_token(decoded_tokens)
|
||||
if forced is not None:
|
||||
logits[req_idx].fill_(min_score)
|
||||
logits[req_idx, forced] = 0
|
||||
elif last_token == self._size_token_id:
|
||||
self._apply_ratio_restriction(logits, req_idx, min_score)
|
||||
elif last_token in self._all_ratio_ids:
|
||||
logits[req_idx].fill_(min_score)
|
||||
logits[req_idx, self._eos_token_id] = 0
|
||||
|
||||
return self._sampler(logits=logits, sampling_metadata=sampling_metadata)
|
||||
|
||||
|
||||
@@ -26,6 +26,13 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Minimum safe values for codec streaming parameters. Mirrors the constants
|
||||
# in stage_input_processors/mimo_audio.py — keep in sync.
|
||||
_MIN_CODEC_CHUNK_FRAMES = 3
|
||||
_MIN_CODEC_LEFT_CONTEXT_FRAMES = 40 # must cover vocoder_attn_window_size[0]
|
||||
_DEFAULT_CODEC_CHUNK_FRAMES = 10
|
||||
_DEFAULT_CODEC_LEFT_CONTEXT_FRAMES = 40
|
||||
|
||||
|
||||
def flat_codec_group_element_count(group_size: int, audio_channels: int) -> int:
|
||||
"""Flat token count for one MiMo talker codec group on the code2wav wire.
|
||||
@@ -485,12 +492,37 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
|
||||
if connector_cfg
|
||||
else None
|
||||
)
|
||||
self._codec_chunk_frames = int(extra_cfg.get("codec_chunk_frames", 3)) if isinstance(extra_cfg, dict) else 3
|
||||
if self._codec_chunk_frames <= 0:
|
||||
raise ValueError(f"codec_chunk_frames must be positive, got {self._codec_chunk_frames}")
|
||||
self._codec_left_context_frames = (
|
||||
int(extra_cfg.get("codec_left_context_frames", 3)) if isinstance(extra_cfg, dict) else 3
|
||||
raw_chunk = (
|
||||
int(extra_cfg.get("codec_chunk_frames", _DEFAULT_CODEC_CHUNK_FRAMES))
|
||||
if isinstance(extra_cfg, dict)
|
||||
else _DEFAULT_CODEC_CHUNK_FRAMES
|
||||
)
|
||||
if raw_chunk < _MIN_CODEC_CHUNK_FRAMES:
|
||||
logger.warning(
|
||||
"codec_chunk_frames=%d is below minimum %d; falling back to %d.",
|
||||
raw_chunk,
|
||||
_MIN_CODEC_CHUNK_FRAMES,
|
||||
_DEFAULT_CODEC_CHUNK_FRAMES,
|
||||
)
|
||||
raw_chunk = _DEFAULT_CODEC_CHUNK_FRAMES
|
||||
self._codec_chunk_frames = raw_chunk
|
||||
|
||||
raw_left = (
|
||||
int(extra_cfg.get("codec_left_context_frames", _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES))
|
||||
if isinstance(extra_cfg, dict)
|
||||
else _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES
|
||||
)
|
||||
if raw_left < _MIN_CODEC_LEFT_CONTEXT_FRAMES:
|
||||
logger.warning(
|
||||
"codec_left_context_frames=%d is below minimum %d (must cover vocoder attention "
|
||||
"window %s); falling back to %d to prevent voice instability.",
|
||||
raw_left,
|
||||
_MIN_CODEC_LEFT_CONTEXT_FRAMES,
|
||||
getattr(self.config, "vocoder_attn_window_size", [40, 10]),
|
||||
_DEFAULT_CODEC_LEFT_CONTEXT_FRAMES,
|
||||
)
|
||||
raw_left = _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES
|
||||
self._codec_left_context_frames = raw_left
|
||||
|
||||
def load_weights(
|
||||
self,
|
||||
|
||||
@@ -542,8 +542,20 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
|
||||
)
|
||||
|
||||
self.device = current_omni_platform.get_torch_device()
|
||||
# global_sampler MUST stay greedy (do_sample=False) so its token decision
|
||||
# matches vLLM's external sampler (SamplingParams temperature=0.0). Both
|
||||
# run argmax on the same logits, so they always agree on whether the next
|
||||
# token is <|empty|> (audio step) or a real text token. Enabling
|
||||
# do_sample=True here without also routing vLLM's sampled token back into
|
||||
# this gate check would cause the two to diverge and corrupt KV-cache state.
|
||||
self.global_sampler = MiMoSampler(do_sample=False, temperature=0.6, top_p=0.95)
|
||||
self.local_sampler = MiMoSampler(do_sample=False, temperature=0.9, top_p=0.95)
|
||||
# local_sampler drives audio-code generation inside local_forward. It is
|
||||
# entirely internal and NOT subject to vLLM's SamplingParams, so stochastic
|
||||
# sampling is safe here and is required to produce natural, varied speech.
|
||||
# Setting do_sample=True also disables the CUDA-graph path (use_cg gate in
|
||||
# local_forward checks `do_sample is False`), preventing MiMoLocalSamplerTensor
|
||||
# from silently forcing argmax even when temperature > 0.
|
||||
self.local_sampler = MiMoSampler(do_sample=True, temperature=0.9, top_p=0.95)
|
||||
self.removed_tokens = None
|
||||
|
||||
self.speech_vocab_sizes = config.parsed_speech_vocab_sizes()
|
||||
@@ -806,7 +818,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
|
||||
device=tokens_device,
|
||||
)
|
||||
if local_sampler is None:
|
||||
local_sampler = MiMoSampler(do_sample=False, temperature=0.6, top_p=0.9)
|
||||
local_sampler = MiMoSampler(do_sample=True, temperature=0.9, top_p=0.95)
|
||||
|
||||
past_key_values = DynamicCache()
|
||||
for t in range(delay_iters):
|
||||
@@ -852,7 +864,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
|
||||
local_sampler: MiMoSampler | None = None,
|
||||
):
|
||||
if local_sampler is None:
|
||||
local_sampler = MiMoSampler(do_sample=False, temperature=0.6, top_p=0.9)
|
||||
local_sampler = MiMoSampler(do_sample=True, temperature=0.9, top_p=0.95)
|
||||
|
||||
b = int(local_embeds.shape[0])
|
||||
use_cg = (local_sampler.do_sample is None or local_sampler.do_sample is False) and bool(
|
||||
|
||||
@@ -30,6 +30,7 @@ import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from transformers import PreTrainedTokenizerBase, Qwen2Config, Qwen2Model, StaticCache
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb
|
||||
|
||||
from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding
|
||||
@@ -479,7 +480,7 @@ class CFMGraphExecutor:
|
||||
self.sde_rnd_placeholder = torch.empty_like(sde_rnd)
|
||||
|
||||
self.graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(self.graph):
|
||||
with torch.cuda.graph(self.graph, pool=current_platform.get_global_graph_pool()):
|
||||
self.gen_lat_placeholder = self.cfm.sample(
|
||||
self.last_hidden_state_placeholder,
|
||||
self.his_lat_placeholder,
|
||||
|
||||
@@ -81,7 +81,6 @@ class Qwen2_5OmniForConditionalGeneration(
|
||||
# Initialize thinker model (multimodal processing)
|
||||
self.thinker = init_vllm_registered_model(
|
||||
vllm_config=vllm_config,
|
||||
prefix=maybe_prefix(prefix, "thinker"),
|
||||
hf_config=thinker_config,
|
||||
# Use registry architecture key
|
||||
architectures=["Qwen2_5OmniThinkerModel"],
|
||||
|
||||
@@ -324,6 +324,16 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
"attn_qkv_proj": [
|
||||
"attn_q_proj",
|
||||
"attn_k_proj",
|
||||
"attn_v_proj",
|
||||
],
|
||||
"qkv": [
|
||||
"q",
|
||||
"k",
|
||||
"v",
|
||||
],
|
||||
}
|
||||
|
||||
@classmethod
|
||||
|
||||
@@ -554,10 +554,10 @@ class CUDAGraphDecoderWrapper:
|
||||
chunk_size: int = 300,
|
||||
left_context_size: int = 25,
|
||||
) -> torch.Tensor:
|
||||
wavs = []
|
||||
start_index = 0
|
||||
total_len = codes.shape[-1]
|
||||
total_upsample = self.decoder.total_upsample
|
||||
wav_out = None
|
||||
|
||||
while start_index < total_len:
|
||||
end_index = min(start_index + chunk_size, total_len)
|
||||
@@ -566,21 +566,16 @@ class CUDAGraphDecoderWrapper:
|
||||
codes_chunk = codes[..., start_index - context_size : end_index]
|
||||
wav_chunk = self._decode(codes_chunk, clone_graph_output=False)
|
||||
|
||||
if wav_out is None:
|
||||
wav_out = torch.empty(
|
||||
(*wav_chunk.shape[:-1], total_len * total_upsample),
|
||||
dtype=wav_chunk.dtype,
|
||||
device=wav_chunk.device,
|
||||
)
|
||||
src_start = context_size * total_upsample
|
||||
dst_start = start_index * total_upsample
|
||||
dst_end = end_index * total_upsample
|
||||
wav_out[..., dst_start:dst_end].copy_(wav_chunk[..., src_start:])
|
||||
# Keep origin/main's concat semantics: Qwen3-Omni can return a chunk
|
||||
# that is shorter than the nominal code_len * total_upsample length.
|
||||
# Clone each slice because graph outputs are static buffers that later
|
||||
# replays may overwrite.
|
||||
wavs.append(wav_chunk[..., context_size * total_upsample :].clone())
|
||||
start_index = end_index
|
||||
|
||||
if wav_out is None:
|
||||
if not wavs:
|
||||
return self.decoder(codes)
|
||||
return wav_out
|
||||
return torch.cat(wavs, dim=-1)
|
||||
|
||||
def batched_chunked_decode_with_cudagraph(
|
||||
self,
|
||||
|
||||
@@ -156,12 +156,6 @@ _OMNI_MODELS = {
|
||||
"fish_speech_dac_decoder",
|
||||
"FishSpeechDACDecoder",
|
||||
),
|
||||
## VoxCPM
|
||||
"VoxCPMForConditionalGeneration": (
|
||||
"voxcpm",
|
||||
"voxcpm",
|
||||
"VoxCPMForConditionalGeneration",
|
||||
),
|
||||
## VoxCPM2
|
||||
"VoxCPM2TalkerForConditionalGeneration": (
|
||||
"voxcpm2",
|
||||
|
||||
@@ -1,5 +0,0 @@
|
||||
from .voxcpm import VoxCPMForConditionalGeneration
|
||||
|
||||
__all__ = [
|
||||
"VoxCPMForConditionalGeneration",
|
||||
]
|
||||
@@ -1,886 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import warnings
|
||||
import wave
|
||||
from collections.abc import Callable, Generator, Iterable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from tqdm import tqdm
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
from .voxcpm_loader import (
|
||||
_build_prompt_cache_with_soundfile,
|
||||
_device_to_string,
|
||||
_force_cuda_available_for_npu,
|
||||
_import_voxcpm_audio_vae_classes,
|
||||
_import_voxcpm_base_model_class,
|
||||
_is_torchcodec_load_error,
|
||||
_normalize_dtype_name,
|
||||
_prepare_runtime_model_dir,
|
||||
_resolve_runtime_device,
|
||||
)
|
||||
from .voxcpm_runtime_utils import resolve_voxcpm_model_dir
|
||||
from .voxcpm_stage_wrappers import _DirectVoxCPMAudioVAE, _DirectVoxCPMLatentGenerator
|
||||
|
||||
logger = init_logger(__name__)
|
||||
_VOXCPM_LATENT_MAGIC = 131071
|
||||
|
||||
|
||||
def _make_voxcpm_model_for_omni(base: type[Any]) -> type[Any]:
|
||||
"""Subclass upstream VoxCPMModel: local ``_inference`` + ``latents_only`` prompt-cache generation."""
|
||||
|
||||
from voxcpm.model.utils import get_dtype
|
||||
|
||||
class VoxCPMModelForOmni(base):
|
||||
@torch.inference_mode()
|
||||
def build_prompt_cache(self, *args: Any, **kwargs: Any):
|
||||
try:
|
||||
return super().build_prompt_cache(*args, **kwargs)
|
||||
except (ImportError, ModuleNotFoundError, RuntimeError) as exc:
|
||||
if not _is_torchcodec_load_error(exc):
|
||||
raise
|
||||
return _build_prompt_cache_with_soundfile(self, *args, **kwargs)
|
||||
|
||||
@torch.inference_mode()
|
||||
def _inference(
|
||||
self,
|
||||
text: torch.Tensor,
|
||||
text_mask: torch.Tensor,
|
||||
feat: torch.Tensor,
|
||||
feat_mask: torch.Tensor,
|
||||
min_len: int = 2,
|
||||
max_len: int = 2000,
|
||||
inference_timesteps: int = 10,
|
||||
cfg_value: float = 2.0,
|
||||
streaming: bool = False,
|
||||
streaming_prefix_len: int = 3,
|
||||
) -> Generator[tuple[torch.Tensor, torch.Tensor | list[torch.Tensor]], None, None]:
|
||||
B, _, _, _ = feat.shape
|
||||
|
||||
feat_embed = self.feat_encoder(feat)
|
||||
feat_embed = self.enc_to_lm_proj(feat_embed)
|
||||
|
||||
scale_emb = self.config.lm_config.scale_emb if self.config.lm_config.use_mup else 1.0
|
||||
text_embed = self.base_lm.embed_tokens(text) * scale_emb
|
||||
combined_embed = text_mask.unsqueeze(-1) * text_embed + feat_mask.unsqueeze(-1) * feat_embed
|
||||
|
||||
prefix_feat_cond = feat[:, -1, ...]
|
||||
pred_feat_seq: list[torch.Tensor] = []
|
||||
|
||||
audio_patch_count = int(feat_mask.sum().item())
|
||||
if audio_patch_count > 0:
|
||||
context_len = min(streaming_prefix_len - 1, audio_patch_count)
|
||||
prompt_context_patches = list(feat[:, -context_len:, :, :].split(1, dim=1))
|
||||
pred_feat_seq = prompt_context_patches + pred_feat_seq
|
||||
|
||||
enc_outputs, kv_cache_tuple = self.base_lm(
|
||||
inputs_embeds=combined_embed,
|
||||
is_causal=True,
|
||||
)
|
||||
self.base_lm.kv_cache.fill_caches(kv_cache_tuple)
|
||||
|
||||
enc_outputs = self.fsq_layer(enc_outputs) * feat_mask.unsqueeze(-1) + enc_outputs * text_mask.unsqueeze(-1)
|
||||
lm_hidden = enc_outputs[:, -1, :]
|
||||
|
||||
residual_enc_outputs, residual_kv_cache_tuple = self.residual_lm(
|
||||
inputs_embeds=enc_outputs + feat_mask.unsqueeze(-1) * feat_embed,
|
||||
is_causal=True,
|
||||
)
|
||||
self.residual_lm.kv_cache.fill_caches(residual_kv_cache_tuple)
|
||||
residual_hidden = residual_enc_outputs[:, -1, :]
|
||||
|
||||
for step_idx in tqdm(range(max_len)):
|
||||
dit_hidden = self.lm_to_dit_proj(lm_hidden) + self.res_to_dit_proj(residual_hidden)
|
||||
pred_feat = self.feat_decoder(
|
||||
mu=dit_hidden,
|
||||
patch_size=self.patch_size,
|
||||
cond=prefix_feat_cond.transpose(1, 2).contiguous(),
|
||||
n_timesteps=inference_timesteps,
|
||||
cfg_value=cfg_value,
|
||||
).transpose(1, 2)
|
||||
|
||||
curr_embed = self.enc_to_lm_proj(self.feat_encoder(pred_feat.unsqueeze(1)))
|
||||
pred_feat_seq.append(pred_feat.unsqueeze(1))
|
||||
prefix_feat_cond = pred_feat
|
||||
|
||||
if streaming:
|
||||
pred_feat_chunk = torch.cat(pred_feat_seq[-streaming_prefix_len:], dim=1)
|
||||
feat_pred = rearrange(pred_feat_chunk, "b t p d -> b d (t p)", b=B, p=self.patch_size)
|
||||
yield feat_pred, pred_feat_seq
|
||||
|
||||
stop_flag = self.stop_head(self.stop_actn(self.stop_proj(lm_hidden))).argmax(dim=-1)[0].cpu().item()
|
||||
if step_idx > min_len and stop_flag == 1:
|
||||
break
|
||||
|
||||
lm_hidden = self.base_lm.forward_step(
|
||||
curr_embed[:, 0, :],
|
||||
torch.tensor([self.base_lm.kv_cache.step()], device=curr_embed.device),
|
||||
).clone()
|
||||
lm_hidden = self.fsq_layer(lm_hidden)
|
||||
residual_hidden = self.residual_lm.forward_step(
|
||||
lm_hidden + curr_embed[:, 0, :],
|
||||
torch.tensor([self.residual_lm.kv_cache.step()], device=curr_embed.device),
|
||||
).clone()
|
||||
|
||||
if not streaming:
|
||||
pred_feat_seq_cat = torch.cat(pred_feat_seq, dim=1)
|
||||
feat_pred = rearrange(pred_feat_seq_cat, "b t p d -> b d (t p)", b=B, p=self.patch_size)
|
||||
yield feat_pred, pred_feat_seq_cat.squeeze(0).cpu()
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate_latents_with_prompt_cache(
|
||||
self,
|
||||
target_text: str,
|
||||
prompt_cache: dict,
|
||||
min_len: int = 2,
|
||||
max_len: int = 2000,
|
||||
inference_timesteps: int = 10,
|
||||
cfg_value: float = 2.0,
|
||||
retry_badcase: bool = False,
|
||||
retry_badcase_max_times: int = 3,
|
||||
retry_badcase_ratio_threshold: float = 6.0,
|
||||
streaming_prefix_len: int = 3,
|
||||
) -> tuple[None, torch.Tensor, torch.Tensor]:
|
||||
return next(
|
||||
self._generate_with_prompt_cache(
|
||||
target_text=target_text,
|
||||
prompt_cache=prompt_cache,
|
||||
min_len=min_len,
|
||||
max_len=max_len,
|
||||
inference_timesteps=inference_timesteps,
|
||||
cfg_value=cfg_value,
|
||||
retry_badcase=retry_badcase,
|
||||
retry_badcase_max_times=retry_badcase_max_times,
|
||||
retry_badcase_ratio_threshold=retry_badcase_ratio_threshold,
|
||||
streaming=False,
|
||||
streaming_prefix_len=streaming_prefix_len,
|
||||
latents_only=True,
|
||||
)
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate_latents_with_prompt_cache_streaming(
|
||||
self,
|
||||
target_text: str,
|
||||
prompt_cache: dict,
|
||||
min_len: int = 2,
|
||||
max_len: int = 2000,
|
||||
inference_timesteps: int = 10,
|
||||
cfg_value: float = 2.0,
|
||||
retry_badcase: bool = False,
|
||||
retry_badcase_max_times: int = 3,
|
||||
retry_badcase_ratio_threshold: float = 6.0,
|
||||
streaming_prefix_len: int = 3,
|
||||
) -> Generator[tuple[None, torch.Tensor, torch.Tensor], None, None]:
|
||||
return self._generate_with_prompt_cache(
|
||||
target_text=target_text,
|
||||
prompt_cache=prompt_cache,
|
||||
min_len=min_len,
|
||||
max_len=max_len,
|
||||
inference_timesteps=inference_timesteps,
|
||||
cfg_value=cfg_value,
|
||||
retry_badcase=retry_badcase,
|
||||
retry_badcase_max_times=retry_badcase_max_times,
|
||||
retry_badcase_ratio_threshold=retry_badcase_ratio_threshold,
|
||||
streaming=True,
|
||||
streaming_prefix_len=streaming_prefix_len,
|
||||
latents_only=True,
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def _generate_with_prompt_cache(
|
||||
self,
|
||||
target_text: str,
|
||||
prompt_cache: dict,
|
||||
min_len: int = 2,
|
||||
max_len: int = 2000,
|
||||
inference_timesteps: int = 10,
|
||||
cfg_value: float = 2.0,
|
||||
retry_badcase: bool = False,
|
||||
retry_badcase_max_times: int = 3,
|
||||
retry_badcase_ratio_threshold: float = 6.0,
|
||||
streaming: bool = False,
|
||||
streaming_prefix_len: int = 3,
|
||||
latents_only: bool = False,
|
||||
) -> Generator[tuple[torch.Tensor | None, torch.Tensor, torch.Tensor | list[torch.Tensor]], None, None]:
|
||||
if retry_badcase and streaming:
|
||||
warnings.warn("Retry on bad cases is not supported in streaming mode, setting retry_badcase=False.")
|
||||
retry_badcase = False
|
||||
if prompt_cache is None:
|
||||
prompt_audio_feat = torch.empty((0, self.patch_size, self.audio_vae.latent_dim), dtype=torch.float32)
|
||||
text = target_text
|
||||
else:
|
||||
prompt_audio_feat = prompt_cache["audio_feat"]
|
||||
prompt_text = prompt_cache["prompt_text"]
|
||||
text = prompt_text + target_text
|
||||
|
||||
text_token = torch.LongTensor(self.text_tokenizer(text))
|
||||
text_token = torch.cat(
|
||||
[
|
||||
text_token,
|
||||
torch.tensor([self.audio_start_token], dtype=torch.int32, device=text_token.device),
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
target_text_token = torch.LongTensor(self.text_tokenizer(target_text))
|
||||
|
||||
audio_length = prompt_audio_feat.size(0)
|
||||
text_length = text_token.shape[0]
|
||||
text_pad_token = torch.zeros(audio_length, dtype=torch.int32, device=text_token.device)
|
||||
audio_pad_feat = torch.zeros(
|
||||
(text_token.shape[0], self.patch_size, self.audio_vae.latent_dim),
|
||||
dtype=torch.float32,
|
||||
device=text_token.device,
|
||||
)
|
||||
text_token = torch.cat([text_token, text_pad_token])
|
||||
audio_feat = torch.cat([audio_pad_feat, prompt_audio_feat], dim=0)
|
||||
text_mask = (
|
||||
torch.cat([torch.ones(text_length), torch.zeros(audio_length)]).type(torch.int32).to(text_token.device)
|
||||
)
|
||||
audio_mask = (
|
||||
torch.cat([torch.zeros(text_length), torch.ones(audio_length)]).type(torch.int32).to(text_token.device)
|
||||
)
|
||||
|
||||
text_token = text_token.unsqueeze(0).to(self.device)
|
||||
text_mask = text_mask.unsqueeze(0).to(self.device)
|
||||
audio_feat = audio_feat.unsqueeze(0).to(self.device).to(get_dtype(self.config.dtype))
|
||||
audio_mask = audio_mask.unsqueeze(0).to(self.device)
|
||||
|
||||
target_text_length = len(self.text_tokenizer(target_text))
|
||||
retry_badcase_times = 0
|
||||
while retry_badcase_times < retry_badcase_max_times:
|
||||
inference_result = self._inference(
|
||||
text_token,
|
||||
text_mask,
|
||||
audio_feat,
|
||||
audio_mask,
|
||||
min_len=min_len,
|
||||
max_len=min(int(target_text_length * retry_badcase_ratio_threshold + 10), max_len),
|
||||
inference_timesteps=inference_timesteps,
|
||||
cfg_value=cfg_value,
|
||||
streaming=streaming,
|
||||
streaming_prefix_len=streaming_prefix_len,
|
||||
)
|
||||
if streaming:
|
||||
patch_len = self.patch_size * self.chunk_size
|
||||
for latent_pred, pred_audio_feat in inference_result:
|
||||
if latents_only:
|
||||
decode_audio = None
|
||||
yield (decode_audio, target_text_token, latent_pred)
|
||||
else:
|
||||
decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32))
|
||||
decode_audio = decode_audio[..., -patch_len:].squeeze(1).cpu()
|
||||
yield (decode_audio, target_text_token, pred_audio_feat)
|
||||
break
|
||||
|
||||
latent_pred, pred_audio_feat = next(inference_result)
|
||||
if retry_badcase and pred_audio_feat.shape[0] >= target_text_length * retry_badcase_ratio_threshold:
|
||||
ratio = pred_audio_feat.shape[0] / target_text_length
|
||||
print(f" Badcase detected, audio_text_ratio={ratio}, retrying...", file=sys.stderr)
|
||||
retry_badcase_times += 1
|
||||
continue
|
||||
break
|
||||
|
||||
if not streaming:
|
||||
if latents_only:
|
||||
decode_audio = None
|
||||
else:
|
||||
decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32))
|
||||
patch_len = self.patch_size * self.chunk_size
|
||||
if audio_mask.sum().item() > 0:
|
||||
decode_audio = decode_audio[..., patch_len * (streaming_prefix_len - 1) :].squeeze(1).cpu()
|
||||
else:
|
||||
decode_audio = decode_audio[..., :].squeeze(1).cpu()
|
||||
yield (decode_audio, target_text_token, pred_audio_feat)
|
||||
|
||||
VoxCPMModelForOmni.__name__ = "VoxCPMModelForOmni"
|
||||
VoxCPMModelForOmni.__qualname__ = "VoxCPMModelForOmni"
|
||||
return VoxCPMModelForOmni
|
||||
|
||||
|
||||
def _import_voxcpm_model_class() -> type[Any]:
|
||||
base = _import_voxcpm_base_model_class()
|
||||
return _make_voxcpm_model_for_omni(base)
|
||||
|
||||
|
||||
def _load_native_voxcpm_model(
|
||||
model_path: str,
|
||||
*,
|
||||
device: torch.device,
|
||||
dtype: str | None,
|
||||
):
|
||||
VoxCPMModel = _import_voxcpm_model_class()
|
||||
model_dir = resolve_voxcpm_model_dir(model_path)
|
||||
runtime_model_path = _prepare_runtime_model_dir(model_dir, target_device=device, target_dtype=dtype)
|
||||
|
||||
if device.type == "npu" and hasattr(torch, "npu"):
|
||||
torch.npu.set_device(device)
|
||||
|
||||
with _force_cuda_available_for_npu(device):
|
||||
return VoxCPMModel.from_local(
|
||||
runtime_model_path,
|
||||
optimize=device.type == "cuda",
|
||||
)
|
||||
|
||||
|
||||
def _load_native_voxcpm_latent_generator(
|
||||
model_path: str,
|
||||
*,
|
||||
device: torch.device,
|
||||
dtype: str | None,
|
||||
) -> _DirectVoxCPMLatentGenerator:
|
||||
return _DirectVoxCPMLatentGenerator(_load_native_voxcpm_model(model_path, device=device, dtype=dtype))
|
||||
|
||||
|
||||
def _load_native_voxcpm_audio_vae(
|
||||
model_path: str,
|
||||
*,
|
||||
device: torch.device,
|
||||
) -> _DirectVoxCPMAudioVAE:
|
||||
AudioVAE, AudioVAEConfig = _import_voxcpm_audio_vae_classes()
|
||||
model_dir = resolve_voxcpm_model_dir(model_path)
|
||||
runtime_model_path = _prepare_runtime_model_dir(model_dir, target_device=device, target_dtype="float32")
|
||||
config_dict = json.loads((Path(runtime_model_path) / "config.json").read_text())
|
||||
audio_vae_config = config_dict.get("audio_vae_config")
|
||||
audio_vae = AudioVAE(config=AudioVAEConfig(**audio_vae_config)) if audio_vae_config is not None else AudioVAE()
|
||||
|
||||
state_dict = torch.load(
|
||||
Path(runtime_model_path) / "audiovae.pth",
|
||||
map_location="cpu",
|
||||
weights_only=True,
|
||||
)["state_dict"]
|
||||
audio_vae.load_state_dict(state_dict, strict=True)
|
||||
audio_vae = audio_vae.to(device=device, dtype=torch.float32).eval()
|
||||
if device.type == "npu" and hasattr(torch, "npu"):
|
||||
torch.npu.set_device(device)
|
||||
patch_size = int(config_dict.get("patch_size", 2))
|
||||
return _DirectVoxCPMAudioVAE(audio_vae, patch_size=patch_size)
|
||||
|
||||
|
||||
class VoxCPMForConditionalGeneration(nn.Module):
|
||||
input_modalities = "audio"
|
||||
_LATENT_STAGES = {"latent_generator", "latent", "ar_dit"}
|
||||
_VAE_STAGES = {"vae", "audio_vae"}
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
del prefix
|
||||
self.vllm_config = vllm_config
|
||||
self.model_path = vllm_config.model_config.model
|
||||
self.model_stage = getattr(vllm_config.model_config, "model_stage", "latent_generator")
|
||||
self.have_multimodal_outputs = True
|
||||
self.has_preprocess = False
|
||||
self.has_postprocess = False
|
||||
self.enable_update_additional_information = True
|
||||
self.requires_raw_input_tokens = True
|
||||
self.inject_omni_request_id_into_runtime_info = True
|
||||
self._pipeline = None
|
||||
self._latent_stream_gens: dict[str, Any] = {}
|
||||
self._latent_stream_terminal_pending: dict[str, int] = {}
|
||||
self._latent_stream_completed: set[str] = set()
|
||||
self._next_local_stream_key = 0
|
||||
self._ar_emit_stop_token = True
|
||||
|
||||
def _runner_hidden_device_dtype(self) -> tuple[torch.device, torch.dtype]:
|
||||
device = _resolve_runtime_device(self.vllm_config)
|
||||
model_config = getattr(self.vllm_config, "model_config", None)
|
||||
dtype = getattr(model_config, "dtype", torch.float32) if model_config is not None else torch.float32
|
||||
return device, dtype
|
||||
|
||||
def _ensure_model_loaded(self):
|
||||
if self._pipeline is not None:
|
||||
return
|
||||
|
||||
target_device = _resolve_runtime_device(self.vllm_config)
|
||||
model_dtype = getattr(self.vllm_config.model_config, "dtype", None)
|
||||
normalized_dtype = _normalize_dtype_name(model_dtype)
|
||||
if self.model_stage in self._LATENT_STAGES:
|
||||
self._pipeline = _load_native_voxcpm_latent_generator(
|
||||
self.model_path,
|
||||
device=target_device,
|
||||
dtype=normalized_dtype,
|
||||
)
|
||||
elif self.model_stage in self._VAE_STAGES:
|
||||
self._pipeline = _load_native_voxcpm_audio_vae(
|
||||
self.model_path,
|
||||
device=target_device,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported VoxCPM model_stage: {self.model_stage}. "
|
||||
"pure_voxcpm only supports split-stage latent_generator/vae inference."
|
||||
)
|
||||
|
||||
logger.info("Loaded VoxCPM stage '%s' on %s", self.model_stage, _device_to_string(target_device))
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
del weights
|
||||
self._ensure_model_loaded()
|
||||
return set()
|
||||
|
||||
@staticmethod
|
||||
def _extract_val(info: dict[str, Any], key: str, default: Any) -> Any:
|
||||
value = info.get(key, default)
|
||||
if isinstance(value, list):
|
||||
return value[0] if value else default
|
||||
return value
|
||||
|
||||
def _resolve_stream_request_key(self, info: dict[str, Any]) -> str:
|
||||
request_key = info.get("__voxcpm_stream_key")
|
||||
if request_key is not None:
|
||||
return str(request_key)
|
||||
|
||||
request_key = info.get("_omni_req_id")
|
||||
if request_key is not None:
|
||||
request_key = str(request_key)
|
||||
info["__voxcpm_stream_key"] = request_key
|
||||
return request_key
|
||||
|
||||
request_key = f"voxcpm-local-{self._next_local_stream_key}"
|
||||
self._next_local_stream_key += 1
|
||||
info["__voxcpm_stream_key"] = request_key
|
||||
return str(request_key)
|
||||
|
||||
def _recover_latent_from_input_ids(self, input_ids: torch.Tensor | None) -> torch.Tensor | None:
|
||||
if input_ids is None or input_ids.numel() == 0:
|
||||
return None
|
||||
flat_ids = input_ids.detach().reshape(-1).to("cpu")
|
||||
if flat_ids.numel() < 4 or int(flat_ids[0].item()) != _VOXCPM_LATENT_MAGIC:
|
||||
return None
|
||||
latent_dim = int(flat_ids[1].item())
|
||||
time_dim = int(flat_ids[2].item())
|
||||
payload = flat_ids[3:]
|
||||
expected = latent_dim * time_dim
|
||||
if latent_dim <= 0 or time_dim <= 0:
|
||||
raise ValueError(f"Invalid VoxCPM latent header: latent_dim={latent_dim}, time_dim={time_dim}")
|
||||
if int(payload.numel()) != expected:
|
||||
raise ValueError(
|
||||
"Invalid VoxCPM latent payload size: "
|
||||
f"expected={expected}, actual={int(payload.numel())}, "
|
||||
f"latent_dim={latent_dim}, time_dim={time_dim}"
|
||||
)
|
||||
packed = payload.to(dtype=torch.int32).to(torch.uint16)
|
||||
return packed.view(torch.bfloat16).to(torch.float32).reshape(1, latent_dim, time_dim)
|
||||
|
||||
def _maybe_recover_vae_infos(
|
||||
self,
|
||||
infos: list[dict[str, Any]],
|
||||
input_ids: torch.Tensor | None,
|
||||
*,
|
||||
async_chunk: bool,
|
||||
) -> list[dict[str, Any]]:
|
||||
if not async_chunk:
|
||||
return infos
|
||||
if any(self._extract_val(info, "latent_audio_feat", None) is not None for info in infos):
|
||||
return infos
|
||||
recovered = self._recover_latent_from_input_ids(input_ids)
|
||||
if recovered is None:
|
||||
return infos
|
||||
return [{"latent_audio_feat": recovered}]
|
||||
|
||||
@staticmethod
|
||||
def _normalize_audio_samples(samples: Any) -> np.ndarray:
|
||||
if isinstance(samples, torch.Tensor):
|
||||
return samples.detach().cpu().float().reshape(-1).numpy()
|
||||
return np.asarray(samples, dtype=np.float32).reshape(-1)
|
||||
|
||||
@classmethod
|
||||
def _normalize_ref_audio(cls, ref_audio: Any) -> tuple[np.ndarray, int]:
|
||||
if isinstance(ref_audio, str):
|
||||
raise TypeError("String ref_audio should be handled as a path before waveform normalization.")
|
||||
|
||||
if isinstance(ref_audio, dict):
|
||||
sample_rate = ref_audio.get("sample_rate") or ref_audio.get("sampling_rate") or ref_audio.get("sr")
|
||||
samples = None
|
||||
for key in ("audio", "wav", "samples", "array", "waveform"):
|
||||
if key in ref_audio and ref_audio[key] is not None:
|
||||
samples = ref_audio[key]
|
||||
break
|
||||
if sample_rate is None or samples is None:
|
||||
raise ValueError("ref_audio dict must contain waveform data and sample rate.")
|
||||
return cls._normalize_audio_samples(samples), int(sample_rate)
|
||||
|
||||
if isinstance(ref_audio, (list, tuple)):
|
||||
if len(ref_audio) == 1:
|
||||
return cls._normalize_ref_audio(ref_audio[0])
|
||||
if len(ref_audio) == 2 and np.isscalar(ref_audio[1]):
|
||||
return cls._normalize_audio_samples(ref_audio[0]), int(ref_audio[1])
|
||||
|
||||
raise TypeError(f"Unsupported ref_audio format: {type(ref_audio)!r}")
|
||||
|
||||
@staticmethod
|
||||
def _write_temp_prompt_wav(waveform: np.ndarray, sample_rate: int) -> str:
|
||||
prompt_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
|
||||
prompt_file.close()
|
||||
|
||||
wav = np.asarray(waveform, dtype=np.float32).reshape(-1)
|
||||
wav = np.clip(wav, -1.0, 1.0)
|
||||
pcm16 = (wav * 32767.0).astype(np.int16)
|
||||
with wave.open(prompt_file.name, "wb") as wav_file:
|
||||
wav_file.setnchannels(1)
|
||||
wav_file.setsampwidth(2)
|
||||
wav_file.setframerate(int(sample_rate))
|
||||
wav_file.writeframes(pcm16.tobytes())
|
||||
|
||||
return prompt_file.name
|
||||
|
||||
@classmethod
|
||||
def _resolve_prompt_inputs(cls, info: dict[str, Any]) -> tuple[str | None, str | None, str | None]:
|
||||
prompt_text = cls._extract_val(info, "prompt_text", None)
|
||||
prompt_wav_path = cls._extract_val(info, "prompt_wav_path", None)
|
||||
if prompt_wav_path:
|
||||
if prompt_text is None:
|
||||
prompt_text = cls._extract_val(info, "ref_text", None)
|
||||
return prompt_wav_path, prompt_text, None
|
||||
|
||||
ref_audio = cls._extract_val(info, "ref_audio", None)
|
||||
ref_text = cls._extract_val(info, "ref_text", None)
|
||||
if ref_audio is None or ref_text is None:
|
||||
return None, None, None
|
||||
if isinstance(ref_audio, str):
|
||||
return ref_audio, ref_text, None
|
||||
|
||||
waveform, sample_rate = cls._normalize_ref_audio(ref_audio)
|
||||
temp_prompt_wav = cls._write_temp_prompt_wav(waveform, sample_rate)
|
||||
return temp_prompt_wav, ref_text, temp_prompt_wav
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
|
||||
if input_ids.numel() == 0:
|
||||
return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32)
|
||||
return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
|
||||
|
||||
def _get_vocab_size(self) -> int:
|
||||
model_config = getattr(self.vllm_config, "model_config", None)
|
||||
if model_config is not None:
|
||||
getter = getattr(model_config, "get_vocab_size", None)
|
||||
if callable(getter):
|
||||
try:
|
||||
return int(getter())
|
||||
except Exception:
|
||||
pass
|
||||
hf_config = getattr(model_config, "hf_text_config", None)
|
||||
if hf_config is not None and hasattr(hf_config, "vocab_size"):
|
||||
return int(hf_config.vocab_size)
|
||||
return 32000
|
||||
|
||||
def _make_empty_output(
|
||||
self,
|
||||
*,
|
||||
output_key: str,
|
||||
payload_factory: Callable[[], torch.Tensor],
|
||||
infos: list[dict[str, Any]],
|
||||
sample_rate: int,
|
||||
out_device: torch.device,
|
||||
out_dtype: torch.dtype,
|
||||
hidden_rows: int | None = None,
|
||||
) -> OmniOutput:
|
||||
if hidden_rows is None:
|
||||
hidden_rows = len(infos)
|
||||
return OmniOutput(
|
||||
text_hidden_states=torch.zeros((hidden_rows, 1), device=out_device, dtype=out_dtype),
|
||||
multimodal_outputs={
|
||||
output_key: [payload_factory() for _ in infos],
|
||||
"sr": [torch.tensor(sample_rate, dtype=torch.int32) for _ in infos],
|
||||
},
|
||||
)
|
||||
|
||||
def _finalize_stage_output(
|
||||
self,
|
||||
*,
|
||||
output_key: str,
|
||||
outputs: list[torch.Tensor],
|
||||
sample_rates: list[torch.Tensor],
|
||||
out_device: torch.device,
|
||||
out_dtype: torch.dtype,
|
||||
hidden_rows: int | None = None,
|
||||
) -> OmniOutput:
|
||||
multimodal_outputs: dict[str, Any] = {output_key: outputs, "sr": sample_rates}
|
||||
if hidden_rows is not None:
|
||||
text_hidden_states = torch.zeros((hidden_rows, 1), device=out_device, dtype=out_dtype)
|
||||
elif outputs:
|
||||
outputs_tensor = torch.stack(outputs)
|
||||
text_hidden_states = (
|
||||
outputs_tensor.unsqueeze(-1)
|
||||
if outputs_tensor.ndim == 1
|
||||
else outputs_tensor.reshape(-1, outputs_tensor.shape[-1])
|
||||
)
|
||||
else:
|
||||
text_hidden_states = torch.zeros((0, 1), device=out_device, dtype=out_dtype)
|
||||
text_hidden_states = text_hidden_states.to(device=out_device, dtype=out_dtype)
|
||||
return OmniOutput(
|
||||
text_hidden_states=text_hidden_states,
|
||||
multimodal_outputs=multimodal_outputs,
|
||||
)
|
||||
|
||||
def _forward_vae_stage(
|
||||
self,
|
||||
infos: list[dict[str, Any]],
|
||||
*,
|
||||
sample_rate: int,
|
||||
async_chunk: bool,
|
||||
out_device: torch.device,
|
||||
out_dtype: torch.dtype,
|
||||
) -> OmniOutput:
|
||||
if all(self._extract_val(info, "latent_audio_feat", None) is None for info in infos):
|
||||
self._ar_emit_stop_token = True
|
||||
return self._make_empty_output(
|
||||
output_key="model_outputs",
|
||||
payload_factory=lambda: torch.zeros((0,), dtype=torch.float32),
|
||||
infos=infos,
|
||||
sample_rate=sample_rate,
|
||||
out_device=out_device,
|
||||
out_dtype=out_dtype,
|
||||
)
|
||||
|
||||
outputs: list[torch.Tensor] = []
|
||||
sample_rates: list[torch.Tensor] = []
|
||||
for info in infos:
|
||||
latent_audio_feat = self._extract_val(info, "latent_audio_feat", None)
|
||||
audio_tensor = self._pipeline.decode(latent_audio_feat, trim_streaming_patch=async_chunk)
|
||||
outputs.append(audio_tensor.float().cpu())
|
||||
sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32))
|
||||
|
||||
self._ar_emit_stop_token = True
|
||||
return self._finalize_stage_output(
|
||||
output_key="model_outputs",
|
||||
outputs=outputs,
|
||||
sample_rates=sample_rates,
|
||||
out_device=out_device,
|
||||
out_dtype=out_dtype,
|
||||
)
|
||||
|
||||
def _forward_latent_stage(
|
||||
self,
|
||||
infos: list[dict[str, Any]],
|
||||
*,
|
||||
sample_rate: int,
|
||||
async_chunk: bool,
|
||||
out_device: torch.device,
|
||||
out_dtype: torch.dtype,
|
||||
hidden_rows: int,
|
||||
) -> OmniOutput:
|
||||
texts = [self._extract_val(info, "text", "") for info in infos]
|
||||
if all(not text for text in texts):
|
||||
self._ar_emit_stop_token = True
|
||||
return self._make_empty_output(
|
||||
output_key="latent_audio_feat",
|
||||
payload_factory=lambda: torch.zeros((0,), dtype=torch.float32),
|
||||
infos=infos,
|
||||
sample_rate=sample_rate,
|
||||
out_device=out_device,
|
||||
out_dtype=out_dtype,
|
||||
hidden_rows=hidden_rows,
|
||||
)
|
||||
|
||||
outputs: list[torch.Tensor] = []
|
||||
sample_rates: list[torch.Tensor] = []
|
||||
last_chunk_flags: list[bool] | None = [] if async_chunk else None
|
||||
payload_finished_flags: list[bool] | None = [] if async_chunk else None
|
||||
for info in infos:
|
||||
text = self._extract_val(info, "text", "")
|
||||
cfg_value = float(self._extract_val(info, "cfg_value", 2.0))
|
||||
inference_timesteps = int(self._extract_val(info, "inference_timesteps", 10))
|
||||
min_len = int(self._extract_val(info, "min_len", 2))
|
||||
max_len = int(self._extract_val(info, "max_len", self._extract_val(info, "max_new_tokens", 4096)))
|
||||
retry_badcase = bool(self._extract_val(info, "retry_badcase", True))
|
||||
retry_badcase_max_times = int(self._extract_val(info, "retry_badcase_max_times", 3))
|
||||
retry_badcase_ratio_threshold = float(self._extract_val(info, "retry_badcase_ratio_threshold", 6.0))
|
||||
streaming_prefix_len = int(self._extract_val(info, "streaming_prefix_len", 3))
|
||||
|
||||
request_key = self._resolve_stream_request_key(info)
|
||||
created_temp: str | None = None
|
||||
|
||||
if async_chunk:
|
||||
terminal_pending = self._latent_stream_terminal_pending.get(request_key, 0)
|
||||
if terminal_pending > 0:
|
||||
outputs.append(torch.zeros((0,), dtype=torch.float32))
|
||||
assert last_chunk_flags is not None
|
||||
last_chunk_flags.append(True)
|
||||
assert payload_finished_flags is not None
|
||||
payload_finished_flags.append(terminal_pending == 1)
|
||||
if terminal_pending == 1:
|
||||
self._latent_stream_terminal_pending.pop(request_key, None)
|
||||
else:
|
||||
self._latent_stream_terminal_pending[request_key] = terminal_pending - 1
|
||||
sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32))
|
||||
continue
|
||||
|
||||
if request_key in self._latent_stream_completed:
|
||||
outputs.append(torch.zeros((0,), dtype=torch.float32))
|
||||
assert last_chunk_flags is not None
|
||||
last_chunk_flags.append(True)
|
||||
assert payload_finished_flags is not None
|
||||
payload_finished_flags.append(False)
|
||||
sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32))
|
||||
continue
|
||||
|
||||
if request_key not in self._latent_stream_gens:
|
||||
prompt_wav_path, prompt_text, temp_prompt_wav = self._resolve_prompt_inputs(info)
|
||||
created_temp = temp_prompt_wav
|
||||
self._latent_stream_gens[request_key] = self._pipeline.iter_latent_chunks_streaming(
|
||||
text=text,
|
||||
prompt_wav_path=prompt_wav_path,
|
||||
prompt_text=prompt_text,
|
||||
cfg_value=cfg_value,
|
||||
inference_timesteps=inference_timesteps,
|
||||
min_len=min_len,
|
||||
max_len=max_len,
|
||||
streaming_prefix_len=streaming_prefix_len,
|
||||
retry_badcase=False,
|
||||
retry_badcase_max_times=retry_badcase_max_times,
|
||||
retry_badcase_ratio_threshold=retry_badcase_ratio_threshold,
|
||||
)
|
||||
generator = self._latent_stream_gens[request_key]
|
||||
try:
|
||||
chunk_latent, is_last = next(generator)
|
||||
except StopIteration:
|
||||
self._latent_stream_gens.pop(request_key, None)
|
||||
self._latent_stream_terminal_pending[request_key] = 1
|
||||
self._latent_stream_completed.add(request_key)
|
||||
outputs.append(torch.zeros((0,), dtype=torch.float32))
|
||||
assert last_chunk_flags is not None
|
||||
last_chunk_flags.append(True)
|
||||
assert payload_finished_flags is not None
|
||||
payload_finished_flags.append(True)
|
||||
else:
|
||||
if is_last:
|
||||
self._latent_stream_gens.pop(request_key, None)
|
||||
self._latent_stream_terminal_pending[request_key] = 1
|
||||
self._latent_stream_completed.add(request_key)
|
||||
outputs.append(chunk_latent.detach().float().cpu())
|
||||
assert last_chunk_flags is not None
|
||||
last_chunk_flags.append(bool(is_last))
|
||||
assert payload_finished_flags is not None
|
||||
payload_finished_flags.append(False)
|
||||
finally:
|
||||
if created_temp is not None and os.path.exists(created_temp):
|
||||
os.unlink(created_temp)
|
||||
sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32))
|
||||
continue
|
||||
|
||||
prompt_wav_path, prompt_text, temp_prompt_wav = self._resolve_prompt_inputs(info)
|
||||
try:
|
||||
latent_audio_feat = self._pipeline.generate_latents(
|
||||
text=text,
|
||||
prompt_wav_path=prompt_wav_path,
|
||||
prompt_text=prompt_text,
|
||||
cfg_value=cfg_value,
|
||||
inference_timesteps=inference_timesteps,
|
||||
min_len=min_len,
|
||||
max_len=max_len,
|
||||
retry_badcase=retry_badcase,
|
||||
retry_badcase_max_times=retry_badcase_max_times,
|
||||
retry_badcase_ratio_threshold=retry_badcase_ratio_threshold,
|
||||
)
|
||||
outputs.append(latent_audio_feat.float().cpu())
|
||||
finally:
|
||||
if temp_prompt_wav is not None and os.path.exists(temp_prompt_wav):
|
||||
os.unlink(temp_prompt_wav)
|
||||
|
||||
sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32))
|
||||
|
||||
self._ar_emit_stop_token = all(last_chunk_flags) if async_chunk and last_chunk_flags else True
|
||||
output = self._finalize_stage_output(
|
||||
output_key="latent_audio_feat",
|
||||
outputs=outputs,
|
||||
sample_rates=sample_rates,
|
||||
out_device=out_device,
|
||||
out_dtype=out_dtype,
|
||||
hidden_rows=hidden_rows,
|
||||
)
|
||||
if async_chunk and payload_finished_flags is not None:
|
||||
output.multimodal_outputs["finished"] = [
|
||||
torch.tensor(flag, dtype=torch.bool) for flag in payload_finished_flags
|
||||
]
|
||||
return output
|
||||
|
||||
def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> torch.Tensor:
|
||||
del sampling_metadata
|
||||
if isinstance(hidden_states, OmniOutput):
|
||||
hidden_states = hidden_states.text_hidden_states
|
||||
if hidden_states is None:
|
||||
device, dtype = self._runner_hidden_device_dtype()
|
||||
hidden_states = torch.zeros((0, 1), device=device, dtype=dtype)
|
||||
if hidden_states.ndim == 1:
|
||||
hidden_states = hidden_states.unsqueeze(-1)
|
||||
elif hidden_states.ndim > 2:
|
||||
hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
|
||||
|
||||
vocab_size = self._get_vocab_size()
|
||||
num_rows = int(hidden_states.shape[0])
|
||||
logits = torch.zeros((num_rows, vocab_size), dtype=torch.float32, device=hidden_states.device)
|
||||
eos_id = 2 if vocab_size > 2 else 0
|
||||
safe_id = 1 if vocab_size > 1 and 1 != eos_id else 0
|
||||
emit_stop = getattr(self, "_ar_emit_stop_token", True)
|
||||
if num_rows > 0:
|
||||
if emit_stop:
|
||||
logits[:, eos_id] = 1.0e6
|
||||
else:
|
||||
logits[:, eos_id] = -1.0e9
|
||||
logits[:, safe_id] = 1.0e6
|
||||
return logits
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
positions: torch.Tensor | None = None,
|
||||
intermediate_tensors: Any = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
runtime_additional_information: list[dict[str, Any]] | None = None,
|
||||
model_intermediate_buffer: list[dict[str, Any]] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> OmniOutput:
|
||||
del positions, intermediate_tensors, inputs_embeds, kwargs
|
||||
self._ensure_model_loaded()
|
||||
out_device, out_dtype = self._runner_hidden_device_dtype()
|
||||
if input_ids is not None and input_ids.device.type == out_device.type:
|
||||
out_device = input_ids.device
|
||||
|
||||
infos = model_intermediate_buffer or runtime_additional_information or [{}]
|
||||
hidden_rows = len(infos)
|
||||
if input_ids is not None and len(input_ids.shape) > 0:
|
||||
hidden_rows = max(hidden_rows, int(input_ids.shape[0]))
|
||||
sample_rate = int(getattr(self._pipeline, "sample_rate", 24000))
|
||||
async_chunk = bool(getattr(self.vllm_config.model_config, "async_chunk", False))
|
||||
if self.model_stage in self._VAE_STAGES:
|
||||
infos = self._maybe_recover_vae_infos(infos, input_ids, async_chunk=async_chunk)
|
||||
return self._forward_vae_stage(
|
||||
infos,
|
||||
sample_rate=sample_rate,
|
||||
async_chunk=async_chunk,
|
||||
out_device=out_device,
|
||||
out_dtype=out_dtype,
|
||||
)
|
||||
if self.model_stage in self._LATENT_STAGES:
|
||||
return self._forward_latent_stage(
|
||||
infos,
|
||||
sample_rate=sample_rate,
|
||||
async_chunk=async_chunk,
|
||||
out_device=out_device,
|
||||
out_dtype=out_dtype,
|
||||
hidden_rows=hidden_rows,
|
||||
)
|
||||
raise ValueError(f"Unsupported VoxCPM model_stage at runtime: {self.model_stage}")
|
||||
|
||||
def make_empty_intermediate_tensors(
|
||||
self, batch_size: int, dtype: torch.dtype, device: torch.device
|
||||
) -> IntermediateTensors:
|
||||
del batch_size, dtype, device
|
||||
return {}
|
||||
|
||||
|
||||
__all__ = ["VoxCPMForConditionalGeneration"]
|
||||
@@ -1,247 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
from contextlib import contextmanager
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _iter_voxcpm_src_candidates() -> list[Path]:
|
||||
candidates: list[Path] = []
|
||||
env_path = os.environ.get("VLLM_OMNI_VOXCPM_CODE_PATH")
|
||||
if env_path:
|
||||
candidates.append(Path(env_path).expanduser())
|
||||
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
candidates.append(repo_root.parent / "VoxCPM" / "src")
|
||||
|
||||
unique_candidates: list[Path] = []
|
||||
seen: set[str] = set()
|
||||
for candidate in candidates:
|
||||
candidate_key = str(candidate)
|
||||
if candidate_key in seen:
|
||||
continue
|
||||
seen.add(candidate_key)
|
||||
unique_candidates.append(candidate)
|
||||
return unique_candidates
|
||||
|
||||
|
||||
def _prepend_voxcpm_src(candidate: Path) -> None:
|
||||
candidate_str = str(candidate)
|
||||
if candidate_str not in sys.path:
|
||||
sys.path.insert(0, candidate_str)
|
||||
|
||||
|
||||
def _import_voxcpm_attrs(module_name: str, *attr_names: str) -> tuple[Any, ...]:
|
||||
last_exc: ImportError | None = None
|
||||
for candidate in _iter_voxcpm_src_candidates():
|
||||
if not candidate.exists():
|
||||
continue
|
||||
_prepend_voxcpm_src(candidate)
|
||||
try:
|
||||
module = importlib.import_module(module_name)
|
||||
return tuple(getattr(module, attr_name) for attr_name in attr_names)
|
||||
except ImportError as exc:
|
||||
last_exc = exc
|
||||
|
||||
try:
|
||||
module = importlib.import_module(module_name)
|
||||
return tuple(getattr(module, attr_name) for attr_name in attr_names)
|
||||
except ImportError as exc:
|
||||
last_exc = exc
|
||||
|
||||
raise ImportError(f"Failed to import {module_name}.") from last_exc
|
||||
|
||||
|
||||
def _import_voxcpm_base_model_class():
|
||||
"""Import upstream ``VoxCPMModel`` from ``VoxCPM/src/voxcpm`` (env, sibling tree, or pip)."""
|
||||
try:
|
||||
(VoxCPMModel,) = _import_voxcpm_attrs("voxcpm.model.voxcpm", "VoxCPMModel")
|
||||
return VoxCPMModel
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"Failed to import VoxCPMModel. Install the `voxcpm` package or set "
|
||||
"`VLLM_OMNI_VOXCPM_CODE_PATH` to the VoxCPM repository `src` directory "
|
||||
"(the parent of the `voxcpm` package that contains `model/` and `modules/`)."
|
||||
) from exc
|
||||
|
||||
|
||||
def _import_voxcpm_audio_vae_classes():
|
||||
try:
|
||||
return _import_voxcpm_attrs("voxcpm.modules.audiovae", "AudioVAE", "AudioVAEConfig")
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"Failed to import VoxCPM AudioVAE. Install the `voxcpm` package or set "
|
||||
"`VLLM_OMNI_VOXCPM_CODE_PATH` to the VoxCPM repository `src` directory."
|
||||
) from exc
|
||||
|
||||
|
||||
def _device_to_string(device: torch.device) -> str:
|
||||
if device.index is None:
|
||||
return device.type
|
||||
return f"{device.type}:{device.index}"
|
||||
|
||||
|
||||
def _normalize_dtype_name(dtype: Any) -> str | None:
|
||||
if dtype is None:
|
||||
return None
|
||||
if isinstance(dtype, torch.dtype):
|
||||
mapping = {
|
||||
torch.bfloat16: "bfloat16",
|
||||
torch.float16: "float16",
|
||||
torch.float32: "float32",
|
||||
}
|
||||
return mapping.get(dtype, str(dtype).removeprefix("torch."))
|
||||
dtype_str = str(dtype)
|
||||
return dtype_str.removeprefix("torch.")
|
||||
|
||||
|
||||
def _resolve_runtime_device(vllm_config: VllmConfig) -> torch.device:
|
||||
try:
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
return current_omni_platform.get_torch_device()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
device = getattr(getattr(vllm_config, "device_config", None), "device", None)
|
||||
if isinstance(device, torch.device):
|
||||
return device
|
||||
if device:
|
||||
return torch.device(device)
|
||||
return torch.device("cpu")
|
||||
|
||||
|
||||
def _prepare_runtime_model_dir(
|
||||
model_path: str | Path,
|
||||
*,
|
||||
target_device: torch.device,
|
||||
target_dtype: str | None,
|
||||
) -> str:
|
||||
source_dir = Path(model_path)
|
||||
config_path = source_dir / "config.json"
|
||||
if not config_path.exists():
|
||||
return str(source_dir)
|
||||
|
||||
config_text = config_path.read_text()
|
||||
config_dict = json.loads(config_text)
|
||||
desired_device = target_device.type
|
||||
desired_dtype = target_dtype or config_dict.get("dtype")
|
||||
|
||||
if config_dict.get("device") == desired_device and config_dict.get("dtype") == desired_dtype:
|
||||
return str(source_dir)
|
||||
|
||||
digest = sha256(f"{source_dir.resolve()}:{config_text}:{desired_device}:{desired_dtype}".encode()).hexdigest()[:16]
|
||||
runtime_dir = Path(tempfile.gettempdir()) / "vllm_omni_voxcpm_runtime" / digest
|
||||
runtime_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for entry in source_dir.iterdir():
|
||||
target = runtime_dir / entry.name
|
||||
if entry.name == "config.json" or target.exists():
|
||||
continue
|
||||
try:
|
||||
target.symlink_to(entry, target_is_directory=entry.is_dir())
|
||||
except OSError as exc:
|
||||
logger.warning(
|
||||
"Falling back to copying VoxCPM runtime artifact %s into %s because symlink creation failed: %s",
|
||||
entry,
|
||||
runtime_dir,
|
||||
exc,
|
||||
)
|
||||
if entry.is_dir():
|
||||
shutil.copytree(entry, target, dirs_exist_ok=True)
|
||||
else:
|
||||
shutil.copy2(entry, target)
|
||||
|
||||
patched_config = dict(config_dict)
|
||||
patched_config["device"] = desired_device
|
||||
if desired_dtype is not None:
|
||||
patched_config["dtype"] = desired_dtype
|
||||
(runtime_dir / "config.json").write_text(json.dumps(patched_config, indent=2, sort_keys=True))
|
||||
return str(runtime_dir)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _force_cuda_available_for_npu(device: torch.device):
|
||||
if device.type != "npu":
|
||||
yield
|
||||
return
|
||||
|
||||
with patch("torch.cuda.is_available", return_value=True):
|
||||
yield
|
||||
|
||||
|
||||
def _is_torchcodec_load_error(exc: BaseException) -> bool:
|
||||
message = str(exc).lower()
|
||||
return "torchcodec" in message or "load_with_torchcodec" in message
|
||||
|
||||
|
||||
def _load_audio_with_soundfile(
|
||||
prompt_wav_path: str,
|
||||
*,
|
||||
sample_rate: int,
|
||||
) -> torch.Tensor:
|
||||
try:
|
||||
import soundfile as sf
|
||||
except ImportError:
|
||||
raise
|
||||
|
||||
audio_np, source_sr = sf.read(prompt_wav_path, dtype="float32", always_2d=True)
|
||||
audio = torch.from_numpy(np.ascontiguousarray(audio_np.T))
|
||||
|
||||
if audio.size(0) > 1:
|
||||
audio = audio.mean(dim=0, keepdim=True)
|
||||
|
||||
if int(source_sr) != int(sample_rate):
|
||||
try:
|
||||
import torchaudio
|
||||
except ImportError as exc:
|
||||
raise ImportError("torchaudio is required for resampling prompt audio.") from exc
|
||||
audio = torchaudio.functional.resample(audio, int(source_sr), int(sample_rate))
|
||||
|
||||
return audio
|
||||
|
||||
|
||||
def _build_prompt_cache_with_soundfile(model: Any, *args: Any, **kwargs: Any) -> dict[str, Any]:
|
||||
if args:
|
||||
prompt_text = args[0]
|
||||
prompt_wav_path = args[1] if len(args) > 1 else kwargs.get("prompt_wav_path")
|
||||
else:
|
||||
prompt_text = kwargs.get("prompt_text")
|
||||
prompt_wav_path = kwargs.get("prompt_wav_path")
|
||||
|
||||
if not prompt_text or not prompt_wav_path:
|
||||
raise ValueError("prompt_text and prompt_wav_path are required")
|
||||
|
||||
audio = _load_audio_with_soundfile(prompt_wav_path, sample_rate=int(model.sample_rate))
|
||||
|
||||
patch_len = model.patch_size * model.chunk_size
|
||||
if audio.size(1) % patch_len != 0:
|
||||
padding_size = patch_len - audio.size(1) % patch_len
|
||||
audio = torch.nn.functional.pad(audio, (padding_size, 0))
|
||||
|
||||
audio_feat = model.audio_vae.encode(audio.to(model.device), model.sample_rate).cpu()
|
||||
audio_feat = audio_feat.view(
|
||||
model.audio_vae.latent_dim,
|
||||
-1,
|
||||
model.patch_size,
|
||||
).permute(1, 2, 0)
|
||||
|
||||
return {
|
||||
"prompt_text": prompt_text,
|
||||
"audio_feat": audio_feat,
|
||||
}
|
||||
@@ -1,44 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def resolve_voxcpm_model_dir(model: str) -> Path:
|
||||
model_path = Path(model).expanduser()
|
||||
if model_path.exists():
|
||||
return model_path
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
return Path(snapshot_download(repo_id=model))
|
||||
|
||||
|
||||
def prepare_voxcpm_hf_config_dir(model_dir: str | Path, hf_config_dir: str | Path) -> Path:
|
||||
model_dir = Path(model_dir).expanduser()
|
||||
hf_config_dir = Path(hf_config_dir).expanduser()
|
||||
hf_config_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
source_config_path = model_dir / "config.json"
|
||||
if not source_config_path.exists():
|
||||
raise FileNotFoundError(f"VoxCPM config.json not found under {model_dir}")
|
||||
|
||||
config_path = hf_config_dir / "config.json"
|
||||
shutil.copy2(source_config_path, config_path)
|
||||
|
||||
source_generation_config_path = model_dir / "generation_config.json"
|
||||
if source_generation_config_path.exists():
|
||||
shutil.copy2(source_generation_config_path, hf_config_dir / "generation_config.json")
|
||||
|
||||
config_dict = json.loads(config_path.read_text(encoding="utf-8"))
|
||||
config_dict["model_type"] = "voxcpm"
|
||||
config_dict.setdefault("architectures", ["VoxCPMForConditionalGeneration"])
|
||||
config_path.write_text(json.dumps(config_dict, indent=2, ensure_ascii=False), encoding="utf-8")
|
||||
return hf_config_dir
|
||||
|
||||
|
||||
__all__ = [
|
||||
"prepare_voxcpm_hf_config_dir",
|
||||
"resolve_voxcpm_model_dir",
|
||||
]
|
||||
@@ -1,185 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Generator
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
class _DirectVoxCPMLatentGenerator:
|
||||
def __init__(self, tts_model: Any):
|
||||
self.tts_model = tts_model
|
||||
self.sample_rate = int(getattr(tts_model, "sample_rate", 24000))
|
||||
|
||||
def generate_latents(
|
||||
self,
|
||||
*,
|
||||
text: str,
|
||||
prompt_wav_path: str | None = None,
|
||||
prompt_text: str | None = None,
|
||||
cfg_value: float = 2.0,
|
||||
inference_timesteps: int = 10,
|
||||
min_len: int = 2,
|
||||
max_len: int = 4096,
|
||||
retry_badcase: bool = True,
|
||||
retry_badcase_max_times: int = 3,
|
||||
retry_badcase_ratio_threshold: float = 6.0,
|
||||
) -> torch.Tensor:
|
||||
if not isinstance(text, str) or not text.strip():
|
||||
raise ValueError("target text must be a non-empty string")
|
||||
if (prompt_wav_path is None) != (prompt_text is None):
|
||||
raise ValueError("prompt_wav_path and prompt_text must both be provided or both be None")
|
||||
if prompt_wav_path is not None and not os.path.exists(prompt_wav_path):
|
||||
raise FileNotFoundError(f"prompt_wav_path does not exist: {prompt_wav_path}")
|
||||
|
||||
prompt_cache = None
|
||||
if prompt_wav_path is not None and prompt_text is not None:
|
||||
prompt_cache = self.tts_model.build_prompt_cache(
|
||||
prompt_text=prompt_text,
|
||||
prompt_wav_path=prompt_wav_path,
|
||||
)
|
||||
|
||||
gen_kw = dict(
|
||||
target_text=" ".join(text.split()),
|
||||
prompt_cache=prompt_cache,
|
||||
min_len=min_len,
|
||||
max_len=max_len,
|
||||
inference_timesteps=inference_timesteps,
|
||||
cfg_value=cfg_value,
|
||||
retry_badcase=retry_badcase,
|
||||
retry_badcase_max_times=retry_badcase_max_times,
|
||||
retry_badcase_ratio_threshold=retry_badcase_ratio_threshold,
|
||||
)
|
||||
latent_entry = getattr(self.tts_model, "generate_latents_with_prompt_cache", None)
|
||||
if latent_entry is not None:
|
||||
_, _, pred_audio_feat = latent_entry(**gen_kw)
|
||||
else:
|
||||
try:
|
||||
_, _, pred_audio_feat = self.tts_model.generate_with_prompt_cache(
|
||||
**gen_kw,
|
||||
latents_only=True,
|
||||
)
|
||||
except TypeError:
|
||||
_, _, pred_audio_feat = self.tts_model.generate_with_prompt_cache(**gen_kw)
|
||||
return pred_audio_feat.detach().cpu().to(torch.float32)
|
||||
|
||||
def iter_latent_chunks_streaming(
|
||||
self,
|
||||
*,
|
||||
text: str,
|
||||
prompt_wav_path: str | None = None,
|
||||
prompt_text: str | None = None,
|
||||
cfg_value: float = 2.0,
|
||||
inference_timesteps: int = 10,
|
||||
min_len: int = 2,
|
||||
max_len: int = 4096,
|
||||
streaming_prefix_len: int = 3,
|
||||
retry_badcase: bool = False,
|
||||
retry_badcase_max_times: int = 3,
|
||||
retry_badcase_ratio_threshold: float = 6.0,
|
||||
) -> Generator[tuple[torch.Tensor, bool], None, None]:
|
||||
"""Yield ``(latent_window, is_last_chunk)`` for Omni async_chunk latent to VAE."""
|
||||
if not isinstance(text, str) or not text.strip():
|
||||
raise ValueError("target text must be a non-empty string")
|
||||
if (prompt_wav_path is None) != (prompt_text is None):
|
||||
raise ValueError("prompt_wav_path and prompt_text must both be provided or both be None")
|
||||
if prompt_wav_path is not None and not os.path.exists(prompt_wav_path):
|
||||
raise FileNotFoundError(f"prompt_wav_path does not exist: {prompt_wav_path}")
|
||||
|
||||
prompt_cache = None
|
||||
if prompt_wav_path is not None and prompt_text is not None:
|
||||
prompt_cache = self.tts_model.build_prompt_cache(
|
||||
prompt_text=prompt_text,
|
||||
prompt_wav_path=prompt_wav_path,
|
||||
)
|
||||
|
||||
gen_kw = dict(
|
||||
target_text=" ".join(text.split()),
|
||||
prompt_cache=prompt_cache,
|
||||
min_len=min_len,
|
||||
max_len=max_len,
|
||||
inference_timesteps=inference_timesteps,
|
||||
cfg_value=cfg_value,
|
||||
retry_badcase=retry_badcase,
|
||||
retry_badcase_max_times=retry_badcase_max_times,
|
||||
retry_badcase_ratio_threshold=retry_badcase_ratio_threshold,
|
||||
streaming_prefix_len=streaming_prefix_len,
|
||||
)
|
||||
stream_entry = getattr(self.tts_model, "generate_latents_with_prompt_cache_streaming", None)
|
||||
if stream_entry is not None:
|
||||
gen = stream_entry(**gen_kw)
|
||||
else:
|
||||
fallback_stream_entry = getattr(self.tts_model, "generate_with_prompt_cache_streaming", None)
|
||||
if fallback_stream_entry is not None:
|
||||
gen = fallback_stream_entry(**gen_kw, latents_only=True)
|
||||
else:
|
||||
gen = self.tts_model._generate_with_prompt_cache(streaming=True, latents_only=True, **gen_kw)
|
||||
|
||||
iterator = iter(gen)
|
||||
previous = next(iterator, None)
|
||||
while previous is not None:
|
||||
current = next(iterator, None)
|
||||
_, _target_tok, chunk_latent = previous
|
||||
if not isinstance(chunk_latent, torch.Tensor):
|
||||
chunk_latent = torch.as_tensor(chunk_latent)
|
||||
yield chunk_latent, current is None
|
||||
previous = current
|
||||
|
||||
|
||||
class _DirectVoxCPMAudioVAE:
|
||||
def __init__(self, audio_vae: nn.Module, *, patch_size: int = 2):
|
||||
self.audio_vae = audio_vae
|
||||
self.sample_rate = int(getattr(audio_vae, "sample_rate", 24000))
|
||||
self.latent_dim = int(getattr(audio_vae, "latent_dim", 64))
|
||||
self.patch_size = int(patch_size)
|
||||
self._chunk_size = int(getattr(audio_vae, "chunk_size", 1))
|
||||
self._stream_audio_patch_samples = max(1, self.patch_size * self._chunk_size)
|
||||
|
||||
def _prepare_latents_for_decode(self, latent_audio_feat: Any) -> torch.Tensor:
|
||||
latents = latent_audio_feat
|
||||
if not isinstance(latents, torch.Tensor):
|
||||
latents = torch.tensor(latents, dtype=torch.float32)
|
||||
latents = latents.detach().to(torch.float32)
|
||||
|
||||
if latents.ndim == 3:
|
||||
if latents.shape[-1] == self.latent_dim:
|
||||
latents = rearrange(latents, "t p d -> 1 d (t p)")
|
||||
elif latents.shape[1] == self.latent_dim:
|
||||
latents = latents.contiguous()
|
||||
else:
|
||||
raise ValueError(f"Unsupported latent_audio_feat shape: {tuple(latents.shape)}")
|
||||
elif latents.ndim == 2:
|
||||
if latents.shape[0] == self.latent_dim:
|
||||
latents = latents.unsqueeze(0)
|
||||
elif latents.shape[1] == self.latent_dim:
|
||||
latents = rearrange(latents, "t d -> 1 d t")
|
||||
else:
|
||||
raise ValueError(f"Unsupported latent_audio_feat shape: {tuple(latents.shape)}")
|
||||
else:
|
||||
raise ValueError(f"Unsupported latent_audio_feat ndim: {latents.ndim}")
|
||||
|
||||
return latents
|
||||
|
||||
@torch.no_grad()
|
||||
def decode(self, latent_audio_feat: Any, *, trim_streaming_patch: bool = False) -> torch.Tensor:
|
||||
latents = self._prepare_latents_for_decode(latent_audio_feat)
|
||||
device = next(self.audio_vae.parameters()).device
|
||||
raw = self.audio_vae.decode(latents.to(device=device, dtype=torch.float32))
|
||||
if isinstance(raw, dict):
|
||||
audio = raw.get("audio")
|
||||
if audio is None:
|
||||
audio = next(v for v in raw.values() if isinstance(v, torch.Tensor))
|
||||
else:
|
||||
audio = raw
|
||||
if audio.dim() == 3:
|
||||
stream = audio.squeeze(1)
|
||||
elif audio.dim() == 2:
|
||||
stream = audio
|
||||
else:
|
||||
stream = audio.reshape(audio.shape[0], -1)
|
||||
if trim_streaming_patch:
|
||||
stream = stream[..., -self._stream_audio_patch_samples :]
|
||||
return stream.reshape(-1).detach().cpu().to(torch.float32)
|
||||
@@ -31,6 +31,7 @@ from vllm.model_executor.models.utils import (
|
||||
maybe_prefix,
|
||||
)
|
||||
from vllm.multimodal.audio import AudioResampler
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
@@ -456,7 +457,6 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
|
||||
self._scaffold_graphs: dict[int, _CapturedGraph] = {}
|
||||
self._residual_graphs: dict[int, _CapturedGraph] = {}
|
||||
self._max_cached_graphs = self._max_batch_size
|
||||
self._cuda_graph_pool: tuple | None = None
|
||||
self._cuda_graph_warmup_steps = 0
|
||||
self._cuda_graph_warmup_threshold = 3
|
||||
|
||||
@@ -661,11 +661,6 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
|
||||
tts = self.tts
|
||||
return tts.stop_head(tts.stop_actn(tts.stop_proj(lm_h)))
|
||||
|
||||
def _get_cuda_graph_pool(self) -> tuple:
|
||||
if self._cuda_graph_pool is None:
|
||||
self._cuda_graph_pool = torch.cuda.graph_pool_handle()
|
||||
return self._cuda_graph_pool
|
||||
|
||||
@staticmethod
|
||||
def _nullify_volatile_metadata(ctx: Any) -> Any:
|
||||
"""Set ``scheduler_metadata`` to None on all attention layers.
|
||||
@@ -698,7 +693,6 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
|
||||
hidden_size = self.config.hidden_size
|
||||
dtype = self._side_dtype
|
||||
dev = torch.device(self._device)
|
||||
pool = self._get_cuda_graph_pool()
|
||||
|
||||
model.precompute_fused_qkv()
|
||||
|
||||
@@ -721,7 +715,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
|
||||
for _ in range(3):
|
||||
_ = model(**call_kwargs)
|
||||
|
||||
with torch.cuda.graph(g.graph, pool=pool):
|
||||
with torch.cuda.graph(g.graph, pool=current_platform.get_global_graph_pool()):
|
||||
g.output = model(**call_kwargs)
|
||||
|
||||
logger.info("CUDA Graph captured for %s (batch_size=%d)", label, batch_size)
|
||||
|
||||
Reference in new issue
Block a user