diff --git a/vllm_omni/model_executor/models/common/alias_free_activation.py b/vllm_omni/model_executor/models/common/alias_free_activation.py new file mode 100644 index 000000000..adb7508af --- /dev/null +++ b/vllm_omni/model_executor/models/common/alias_free_activation.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Alias-free activation for BigVGAN-style speech decoders. + +Provides anti-aliased activation (upsample → activation → downsample) with +Kaiser-windowed sinc filters. Includes NPU bf16 workarounds for platforms +where F.pad(mode='replicate') does not support bfloat16. + +Used by: Qwen2.5-Omni, Qwen3-TTS v1, CoVo-Audio, and future BigVGAN vocoders. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.logger import init_logger + +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# Helper utilities +# --------------------------------------------------------------------------- + + +def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: + """Generate a 1-D Kaiser-windowed sinc low-pass filter. + + Args: + cutoff: Normalized cutoff frequency (0 to 0.5). + half_width: Transition bandwidth. + kernel_size: Number of filter taps. + + Returns: + Filter tensor of shape ``(1, 1, kernel_size)``. + """ + is_even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + + delta_f = 4 * half_width + attenuation = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + + if attenuation > 50.0: + beta = 0.1102 * (attenuation - 8.7) + elif attenuation >= 21.0: + beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0) + else: + beta = 0.0 + + # torch.kaiser_window is not supported on NPU / XPU — compute on CPU. + if current_omni_platform.is_npu(): + kaiser_window = torch.kaiser_window( + kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu" + ).to("npu") + elif current_omni_platform.is_xpu(): + kaiser_window = torch.kaiser_window( + kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu" + ).to("xpu") + else: + kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32) + + if is_even: + time_indices = torch.arange(-half_size, half_size) + 0.5 + else: + time_indices = torch.arange(kernel_size) - half_size + + if cutoff == 0: + return torch.zeros((1, 1, kernel_size), dtype=torch.float32) + + sinc_filter = torch.sinc(2 * cutoff * time_indices) + normalized_filter = 2 * cutoff * kaiser_window * sinc_filter + normalized_filter /= normalized_filter.sum() + + return normalized_filter.view(1, 1, kernel_size) + + +def replication_pad_1d(hidden_states: torch.Tensor, pad_left: int, pad_right: int) -> torch.Tensor: + """Manual replicate padding for NPU bf16 compatibility. + + ``F.pad(mode='replicate')`` does not support bfloat16 on NPU (and older + CUDA). This expand+cat fallback produces identical results. + + TODO: Remove when upstream PyTorch fixes replication_pad1d for bf16. + """ + if pad_left == 0 and pad_right == 0: + return hidden_states + + segments = [] + if pad_left > 0: + left = hidden_states[..., :1].expand(*hidden_states.shape[:-1], pad_left) + segments.append(left) + + segments.append(hidden_states) + + if pad_right > 0: + right = hidden_states[..., -1:].expand(*hidden_states.shape[:-1], pad_right) + segments.append(right) + + return torch.cat(segments, dim=-1) + + +# --------------------------------------------------------------------------- +# Up / Down sampling with Kaiser sinc filters +# --------------------------------------------------------------------------- + + +class UpSample1d(nn.Module): + """Kaiser-sinc interpolation upsample (default: 2x, 12-tap filter).""" + + def __init__(self, ratio: int = 2, kernel_size: int | None = None): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.stride = ratio + self.pad = self.kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 + + filt = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size) + self.register_buffer("filter", filt, persistent=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + channels = hidden_states.shape[1] + if current_omni_platform.is_npu(): + input_dtype = hidden_states.dtype + hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad, self.pad) + hidden_states = self.ratio * F.conv_transpose1d( + hidden_states, + self.filter.expand(channels, -1, -1), + stride=self.stride, + groups=channels, + ).to(input_dtype) + else: + input_dtype = hidden_states.dtype + hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(self.filter.dtype) + hidden_states = self.ratio * F.conv_transpose1d( + hidden_states, + self.filter.expand(channels, -1, -1), + stride=self.stride, + groups=channels, + ).to(input_dtype) + hidden_states = hidden_states[..., self.pad_left : -self.pad_right] + return hidden_states + + +class DownSample1d(nn.Module): + """Kaiser-sinc anti-aliasing downsample (default: 2x, 12-tap filter).""" + + def __init__(self, ratio: int = 2, kernel_size: int | None = None): + super().__init__() + if kernel_size is None: + kernel_size = int(6 * ratio // 2) * 2 + + cutoff = 0.5 / ratio + half_width = 0.6 / ratio + + if cutoff < 0.0: + raise ValueError("Minimum cutoff must be larger than zero.") + if cutoff > 0.5: + raise ValueError("A cutoff above 0.5 does not make sense.") + + self.even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(self.even) + self.pad_right = kernel_size // 2 + self.stride = ratio + + filt = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) + self.register_buffer("filter", filt, persistent=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + channels = hidden_states.shape[1] + if current_omni_platform.is_npu(): + input_dtype = hidden_states.dtype + hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad_left, self.pad_right) + out = F.conv1d( + hidden_states, + self.filter.to(device=hidden_states.device, dtype=hidden_states.dtype).expand(channels, -1, -1), + stride=self.stride, + groups=channels, + ).to(input_dtype) + else: + input_dtype = hidden_states.dtype + hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate").to( + self.filter.dtype + ) + out = F.conv1d( + hidden_states, + self.filter.expand(channels, -1, -1), + stride=self.stride, + groups=channels, + ).to(input_dtype) + return out + + +# --------------------------------------------------------------------------- +# Alias-free activation: upsample → activation → downsample +# --------------------------------------------------------------------------- + + +class AliasFreeActivation1d(nn.Module): + """Anti-aliased activation following BigVGAN (ICLR 2023). + + Upsamples the signal, applies a pointwise activation (typically SnakeBeta), + then downsamples with a low-pass filter to suppress aliasing from the + high-frequency components introduced by the activation function. + + The ``activation`` should be an ``nn.Module`` instance — typically + ``SnakeBeta`` from ``common.snake_activation``, which already dispatches + between Triton and eager PyTorch internally. + + Args: + activation: Pointwise activation module (e.g. ``SnakeBeta(channels)``). + up_ratio: Upsample factor before activation. + down_ratio: Downsample factor after activation. + up_kernel_size: Kaiser sinc filter taps for upsampling. + down_kernel_size: Kaiser sinc filter taps for downsampling. + """ + + def __init__( + self, + activation: nn.Module, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + ): + super().__init__() + if not callable(activation): + raise TypeError("Activation function must be callable") + self.act = activation + self.upsample = UpSample1d(up_ratio, up_kernel_size) + self.downsample = DownSample1d(down_ratio, down_kernel_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.upsample(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.downsample(hidden_states) + return hidden_states diff --git a/vllm_omni/model_executor/models/common/snake_activation.py b/vllm_omni/model_executor/models/common/snake_activation.py index 35bb3384a..f50e27674 100644 --- a/vllm_omni/model_executor/models/common/snake_activation.py +++ b/vllm_omni/model_executor/models/common/snake_activation.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Copyright 2025 The Qwen team. -"""Shared SnakeBeta activation for speech decoders (Qwen3-TTS, Qwen3-Omni Code2Wav).""" +"""Shared Snake/SnakeBeta activations for speech decoders. + +Used by: Qwen3-TTS, Qwen3-Omni Code2Wav, Qwen2.5-Omni, CoVo-Audio, CosyVoice3. +""" import torch from torch import nn @@ -70,12 +73,17 @@ class SnakeBeta(nn.Module): SnakeBeta._triton_kernel = _kernel return True - def __init__(self, in_features, alpha=1.0): + def __init__(self, in_features, alpha=1.0, alpha_logscale=True): super().__init__() self.in_features = in_features + self.alpha_logscale = alpha_logscale - self.alpha = Parameter(torch.zeros(in_features) * alpha) - self.beta = Parameter(torch.zeros(in_features) * alpha) + if alpha_logscale: + self.alpha = Parameter(torch.zeros(in_features) * alpha) + self.beta = Parameter(torch.zeros(in_features) * alpha) + else: + self.alpha = Parameter(torch.ones(in_features) * alpha) + self.beta = Parameter(torch.ones(in_features) * alpha) self.no_div_by_zero = 0.000000001 @@ -86,8 +94,12 @@ class SnakeBeta(nn.Module): def precompute_exp_cache(self): """Materialize exp(alpha) and 1/(exp(beta)+eps) as frozen buffers.""" with torch.no_grad(): - self._exp_alpha = torch.exp(self.alpha).contiguous() - self._inv_beta = (1.0 / (torch.exp(self.beta) + self.no_div_by_zero)).contiguous() + if self.alpha_logscale: + self._exp_alpha = torch.exp(self.alpha).contiguous() + self._inv_beta = (1.0 / (torch.exp(self.beta) + self.no_div_by_zero)).contiguous() + else: + self._exp_alpha = self.alpha.contiguous() + self._inv_beta = (1.0 / (self.beta + self.no_div_by_zero)).contiguous() @property def _cached(self): @@ -104,12 +116,10 @@ class SnakeBeta(nn.Module): return self._eager_forward(hidden_states) def _eager_forward(self, hidden_states): - if self._cached: - exp_alpha = self._exp_alpha.unsqueeze(0).unsqueeze(-1) - inv_beta = self._inv_beta.unsqueeze(0).unsqueeze(-1) - else: - exp_alpha = torch.exp(self.alpha).unsqueeze(0).unsqueeze(-1) - inv_beta = (1.0 / (torch.exp(self.beta) + self.no_div_by_zero)).unsqueeze(0).unsqueeze(-1) + if not self._cached: + self.precompute_exp_cache() + exp_alpha = self._exp_alpha.unsqueeze(0).unsqueeze(-1) + inv_beta = self._inv_beta.unsqueeze(0).unsqueeze(-1) hidden_states = hidden_states + inv_beta * torch.pow(torch.sin(hidden_states * exp_alpha), 2) return hidden_states @@ -134,3 +144,33 @@ class SnakeBeta(nn.Module): block_t=block_t, ) return out + + +class Snake(SnakeBeta): + """Original Snake activation with a single parameter: x + 1/α * sin²(αx). + + Unlike SnakeBeta which has separate alpha (frequency) and beta (magnitude) + parameters, Snake uses alpha for both. Only ``alpha`` appears in the + state_dict — ``beta`` is absent, keeping checkpoint compatibility with + CosyVoice3's HiFi-GAN. + + The Triton kernel and precomputed-cache path from SnakeBeta are reused; + ``precompute_exp_cache`` derives ``_inv_beta`` from ``alpha`` so the + forward path is identical. + """ + + def __init__(self, in_features, alpha=1.0, alpha_logscale=False): + # Initialise SnakeBeta — creates both alpha and beta Parameters. + super().__init__(in_features, alpha=alpha, alpha_logscale=alpha_logscale) + # Drop beta as a Parameter so it won't appear in state_dict. + del self.beta + + def precompute_exp_cache(self): + """Derive both exp_alpha and inv_beta from the single alpha parameter.""" + with torch.no_grad(): + if self.alpha_logscale: + self._exp_alpha = torch.exp(self.alpha).contiguous() + self._inv_beta = (1.0 / (torch.exp(self.alpha) + self.no_div_by_zero)).contiguous() + else: + self._exp_alpha = self.alpha.contiguous() + self._inv_beta = (1.0 / (self.alpha + self.no_div_by_zero)).contiguous() diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py index bd31443cc..63434d57f 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -11,7 +11,6 @@ import torch import torch.nn as nn import torch.nn.functional as F from scipy.signal import get_window -from torch import pow, sin from torch.nn import Conv1d, ConvTranspose1d from torch.nn.utils import remove_weight_norm @@ -20,64 +19,8 @@ try: except ImportError: from torch.nn.utils import weight_norm from torch.distributions.uniform import Uniform -from torch.nn import Parameter - - -# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license. -# LICENSE is in incl_licenses directory. -class Snake(nn.Module): - """ - Implementation of a sine-based periodic activation function - Shape: - - Input: (B, C, T) - - Output: (B, C, T), same shape as the input - Parameters: - - alpha - trainable parameter - References: - - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: - https://arxiv.org/abs/2006.08195 - Examples: - >>> a1 = snake(256) - >>> x = torch.randn(256) - >>> x = a1(x) - """ - - def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): - """ - Initialization. - INPUT: - - in_features: shape of the input - - alpha: trainable parameter - alpha is initialized to 1 by default, higher values = higher-frequency. - alpha will be trained along with the rest of your model. - """ - super().__init__() - self.in_features = in_features - - # initialize alpha - self.alpha_logscale = alpha_logscale - if self.alpha_logscale: # log scale alphas initialized to zeros - self.alpha = Parameter(torch.zeros(in_features) * alpha) - else: # linear scale alphas initialized to ones - self.alpha = Parameter(torch.ones(in_features) * alpha) - - self.alpha.requires_grad = alpha_trainable - - self.no_div_by_zero = 0.000000001 - - def forward(self, x): - """ - Forward pass of the function. - Applies the function to the input elementwise. - Snake ∶= x + 1/a * sin^2 (xa) - """ - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] - if self.alpha_logscale: - alpha = torch.exp(alpha) - x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2) - - return x +from vllm_omni.model_executor.models.common.snake_activation import Snake """hifigan based generator implementation. diff --git a/vllm_omni/model_executor/models/covo_audio/token2wav.py b/vllm_omni/model_executor/models/covo_audio/token2wav.py index f6bed125c..4a6ce4b7e 100644 --- a/vllm_omni/model_executor/models/covo_audio/token2wav.py +++ b/vllm_omni/model_executor/models/covo_audio/token2wav.py @@ -10,10 +10,11 @@ import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange -from torch import pow, sin -from torch.nn import Parameter from torch.nn.utils import weight_norm +from vllm_omni.model_executor.models.common.alias_free_activation import AliasFreeActivation1d +from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta + # --------------- Utilities --------------- @@ -93,41 +94,7 @@ def apply_rotary_pos_emb(pos, t): return t * pos.cos() + rotate_half(t) * pos.sin() -if "sinc" in dir(torch): - sinc = torch.sinc -else: - - def sinc(x: torch.Tensor): - return torch.where( - x == 0, torch.tensor(1.0, device=x.device, dtype=x.dtype), torch.sin(math.pi * x) / math.pi / x - ) - - -def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): - even = kernel_size % 2 == 0 - half_size = kernel_size // 2 - - delta_f = 4 * half_width - A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 - if A > 50.0: - beta = 0.1102 * (A - 8.7) - elif A >= 21.0: - beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0) - else: - beta = 0.0 - window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) - - if even: - time = torch.arange(-half_size, half_size) + 0.5 - else: - time = torch.arange(kernel_size) - half_size - if cutoff == 0: - filter_ = torch.zeros_like(time) - else: - filter_ = 2 * cutoff * window * sinc(2 * cutoff * time) - filter_ /= filter_.sum() - - return filter_.view(1, 1, kernel_size) +# sinc, kaiser_sinc_filter1d removed — now in common.alias_free_activation # --------------- Basic Layers --------------- @@ -320,116 +287,14 @@ class ConvPositionEmbed(nn.Module): # --------------- Anti-aliasing (BigVGAN) --------------- -class LowPassFilter1d(nn.Module): - def __init__( - self, - cutoff=0.5, - half_width=0.6, - stride: int = 1, - padding: bool = True, - padding_mode: str = "replicate", - kernel_size: int = 12, - ): - super().__init__() - if cutoff < -0.0: - raise ValueError("Minimum cutoff must be larger than zero.") - if cutoff > 0.5: - raise ValueError("A cutoff above 0.5 does not make sense.") - self.kernel_size = kernel_size - self.even = kernel_size % 2 == 0 - self.pad_left = kernel_size // 2 - int(self.even) - self.pad_right = kernel_size // 2 - self.stride = stride - self.padding = padding - self.padding_mode = padding_mode - filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) - self.register_buffer("filter", filter) - - def forward(self, x): - _, C, _ = x.shape - if self.padding: - x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) - out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) - return out - - -class UpSample1d(nn.Module): - def __init__(self, ratio=2, kernel_size=None): - super().__init__() - self.ratio = ratio - self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size - self.stride = ratio - self.pad = self.kernel_size // ratio - 1 - self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 - self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 - filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size) - self.register_buffer("filter", filter) - - def forward(self, x): - _, C, _ = x.shape - x = F.pad(x, (self.pad, self.pad), mode="replicate") - x = self.ratio * F.conv_transpose1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) - x = x[..., self.pad_left : -self.pad_right] - return x - - -class DownSample1d(nn.Module): - def __init__(self, ratio=2, kernel_size=None): - super().__init__() - self.ratio = ratio - self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size - self.lowpass = LowPassFilter1d( - cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=self.kernel_size - ) - - def forward(self, x): - return self.lowpass(x) - - -class Activation1d(nn.Module): - def __init__( - self, activation, up_ratio: int = 2, down_ratio: int = 2, up_kernel_size: int = 12, down_kernel_size: int = 12 - ): - super().__init__() - self.up_ratio = up_ratio - self.down_ratio = down_ratio - self.act = activation - self.upsample = UpSample1d(up_ratio, up_kernel_size) - self.downsample = DownSample1d(down_ratio, down_kernel_size) - - def forward(self, x): - x = self.upsample(x) - x = self.act(x) - x = self.downsample(x) - return x +# LowPassFilter1d, UpSample1d, DownSample1d, Activation1d removed — +# now provided by common.alias_free_activation.AliasFreeActivation1d # --------------- Activations --------------- -class SnakeBeta(nn.Module): - def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): - super().__init__() - self.in_features = in_features - self.alpha_logscale = alpha_logscale - if self.alpha_logscale: - self.alpha = Parameter(torch.zeros(in_features) * alpha) - self.beta = Parameter(torch.zeros(in_features) * alpha) - else: - self.alpha = Parameter(torch.ones(in_features) * alpha) - self.beta = Parameter(torch.ones(in_features) * alpha) - self.alpha.requires_grad = alpha_trainable - self.beta.requires_grad = alpha_trainable - self.no_div_by_zero = 0.000000001 - - def forward(self, x): - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) - beta = self.beta.unsqueeze(0).unsqueeze(-1) - if self.alpha_logscale: - alpha = torch.exp(alpha) - beta = torch.exp(beta) - x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2) - return x +# SnakeBeta removed — now in common.snake_activation (with alpha_logscale support) # --------------- Attention --------------- @@ -828,7 +693,7 @@ class AMPBlock1(torch.nn.Module): self.activations = nn.ModuleList( [ - Activation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale)) + AliasFreeActivation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale)) for _ in range(self.num_layers) ] ) @@ -882,7 +747,7 @@ class BigVGANFlowVAE(nn.Module): self.resblocks.append(AMPBlock1(h, ch, k, d, causal=causal)) activation_post = SnakeBeta(ch, alpha_logscale=h.snake_logscale) - self.activation_post = Activation1d(activation=activation_post) + self.activation_post = AliasFreeActivation1d(activation=activation_post) self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, causal=causal)) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py index 34c88bad9..49d3ac114 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py @@ -2,14 +2,12 @@ # Start Token2Wav # ############################ -import math from collections.abc import Iterable import numpy as np import torch import torch.nn as nn import torch.nn.functional as F -from torch.nn import Parameter from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import ( Qwen2_5OmniBigVGANConfig, @@ -35,8 +33,9 @@ from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding +from vllm_omni.model_executor.models.common.alias_free_activation import AliasFreeActivation1d +from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import cap_and_align_mel_length, resolve_max_mel_frames -from vllm_omni.platforms import current_omni_platform # Provide a no-op auto_docstring decorator to satisfy annotations if missing @@ -634,238 +633,9 @@ class DiTDecoderLayer(nn.Module): return hidden_states -class SnakeBeta(nn.Module): - """ - A modified Snake function which uses separate parameters for the magnitude - of the periodic components - Shape: - - Input: (B, C, T) - - Output: (B, C, T), same shape as the input - Parameters: - - alpha - trainable parameter that controls frequency - - beta - trainable parameter that controls magnitude - References: - - This activation function is a modified version based on this paper - by Liu Ziyin, Tilman Hartwig, Masahito Ueda: - https://huggingface.co/papers/2006.08195 - """ - - def __init__(self, in_features, alpha=1.0): - super().__init__() - self.in_features = in_features - - # initialize alpha - self.alpha = Parameter(torch.zeros(in_features) * alpha) - self.beta = Parameter(torch.zeros(in_features) * alpha) - - self.no_div_by_zero = 0.000000001 - - def forward(self, hidden_states): - """ - Forward pass of the function. - Applies the function to the input elementwise. - SnakeBeta ∶= x + 1/b * sin^2 (xa) - """ - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] - beta = self.beta.unsqueeze(0).unsqueeze(-1) - alpha = torch.exp(alpha) - beta = torch.exp(beta) - hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow( - torch.sin(hidden_states * alpha), 2 - ) - - return hidden_states - - -def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: - """Generates a 1D Kaiser-windowed sinc filter. - - Args: - cutoff (float): Normalized cutoff frequency (0 to 0.5). - half_width (float): Transition bandwidth. - kernel_size (int): Number of filter taps. - - Returns: - torch.Tensor: A tensor of shape (1, 1, kernel_size) representing the filter. - """ - is_even = kernel_size % 2 == 0 - half_size = kernel_size // 2 - - # Compute Kaiser window parameters - delta_f = 4 * half_width - attenuation = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 - - if attenuation > 50.0: - beta = 0.1102 * (attenuation - 8.7) - elif attenuation >= 21.0: - beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0) - else: - beta = 0.0 - - # TODO: When torch.kaiser_window supports NPU, remove the device="cpu" argument - if current_omni_platform.is_npu(): - kaiser_window = torch.kaiser_window( - kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu" - ).to("npu") - elif current_omni_platform.is_xpu(): - kaiser_window = torch.kaiser_window( - kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu" - ).to("xpu") - else: - kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32) - - # Compute time indices - if is_even: - time_indices = torch.arange(-half_size, half_size) + 0.5 - else: - time_indices = torch.arange(kernel_size) - half_size - - # Compute sinc filter - if cutoff == 0: - return torch.zeros((1, 1, kernel_size), dtype=torch.float32) - - sinc_filter = torch.sinc(2 * cutoff * time_indices) - normalized_filter = 2 * cutoff * kaiser_window * sinc_filter - - # Normalize to ensure sum = 1 (avoid leakage of constant component) - normalized_filter /= normalized_filter.sum() - - return normalized_filter.view(1, 1, kernel_size) - - -def replication_pad_1d(hidden_states: torch.Tensor, pad_left: int, pad_right: int) -> torch.Tensor: - """ - Manual replicate padding to avoid replication_pad1d kernel limits on NPU. - TODO: remove when F.pad supports replicate mode on NPU. - """ - # NOTE: a immature implementation for running in NPU. Need to discuss. - if pad_left == 0 and pad_right == 0: - return hidden_states - - segments = [] - if pad_left > 0: - left = hidden_states[..., :1].expand(*hidden_states.shape[:-1], pad_left) - segments.append(left) - - segments.append(hidden_states) - - if pad_right > 0: - right = hidden_states[..., -1:].expand(*hidden_states.shape[:-1], pad_right) - segments.append(right) - - return torch.cat(segments, dim=-1) - - -class UpSample1d(nn.Module): - def __init__(self, ratio=2, kernel_size=None): - super().__init__() - self.ratio = ratio - self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size - self.stride = ratio - self.pad = self.kernel_size // ratio - 1 - self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 - self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 - - filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size) - self.register_buffer("filter", filter, persistent=False) - - def forward(self, hidden_states): - channels = hidden_states.shape[1] - if current_omni_platform.is_npu(): - # TODO: When F.pad supports replicate mode on NPU, remove this branch - input_dtype = hidden_states.dtype - # F.pad in NPU doesn't support BF16 when mode is replicate. - # To ensure the accuracy, manually pad the input tensor. - hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad, self.pad) - filter_convert_dtype = self.filter.to(hidden_states.dtype) - hidden_states = self.ratio * F.conv_transpose1d( - hidden_states, - filter_convert_dtype.expand(channels, -1, -1), - stride=self.stride, - groups=channels, - ).to(input_dtype) - else: - hidden_states_dtype = hidden_states.dtype - hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(self.filter.dtype) - hidden_states = self.ratio * F.conv_transpose1d( - hidden_states, - self.filter.expand(channels, -1, -1), - stride=self.stride, - groups=channels, - ).to(hidden_states_dtype) - hidden_states = hidden_states[..., self.pad_left : -self.pad_right] - - return hidden_states - - -class DownSample1d(nn.Module): - def __init__(self, ratio=2, kernel_size=None): - super().__init__() - cutoff = 0.5 / ratio - half_width = 0.6 / ratio - - if cutoff < 0.0: - raise ValueError("Minimum cutoff must be larger than zero.") - if cutoff > 0.5: - raise ValueError("A cutoff above 0.5 does not make sense.") - - self.even = kernel_size % 2 == 0 - self.pad_left = kernel_size // 2 - int(self.even) - self.pad_right = kernel_size // 2 - self.stride = ratio - filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) - self.register_buffer("filter", filter, persistent=False) - - def forward(self, hidden_states): - channels = hidden_states.shape[1] - if current_omni_platform.is_npu(): - input_dtype = hidden_states.dtype - # F.pad in NPU doesn't support BF16 when mode is replicate. - # To ensure the accuracy, manually pad the input tensor. - hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad_left, self.pad_right) - filter_on_device = self.filter.to(device=hidden_states.device, dtype=hidden_states.dtype) - out = F.conv1d( - hidden_states, - filter_on_device.expand(channels, -1, -1), - stride=self.stride, - groups=channels, - ).to(input_dtype) - else: - hidden_states_dtype = hidden_states.dtype - hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate").to( - self.filter.dtype - ) - out = F.conv1d( - hidden_states, - self.filter.expand(channels, -1, -1), - stride=self.stride, - groups=channels, - ).to(hidden_states_dtype) - return out - - -class TorchActivation1d(nn.Module): - def __init__( - self, - activation, - up_ratio: int = 2, - down_ratio: int = 2, - up_kernel_size: int = 12, - down_kernel_size: int = 12, - ): - super().__init__() - if not callable(activation): - raise TypeError("Activation function must be callable") - self.act = activation - self.upsample = UpSample1d(up_ratio, up_kernel_size) - self.downsample = DownSample1d(down_ratio, down_kernel_size) - - def forward(self, hidden_states): - hidden_states = self.upsample(hidden_states) - hidden_states = self.act(hidden_states) - hidden_states = self.downsample(hidden_states) - - return hidden_states +# SnakeBeta, UpSample1d, DownSample1d, AliasFreeActivation1d (formerly +# TorchActivation1d) are now provided by the shared ``common`` modules. +# See: vllm_omni/model_executor/models/common/{snake_activation,alias_free_activation}.py class AMPBlock(torch.nn.Module): @@ -938,7 +708,7 @@ class AMPBlock(torch.nn.Module): self.num_layers = len(self.convs1) + len(self.convs2) # total number of conv layers self.activations = nn.ModuleList( - [TorchActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)] + [AliasFreeActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)] ) def _get_padding(self, kernel_size, dilation=1): @@ -1002,7 +772,7 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel): ] ) - self.activation_post = TorchActivation1d( + self.activation_post = AliasFreeActivation1d( activation=SnakeBeta(config.upsample_initial_channel // (2**self.num_upsample_layers)) ) self.conv_post = nn.Conv1d( diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py index 7bc0030bc..6ddbedd9b 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py @@ -13,13 +13,11 @@ # limitations under the License. """PyTorch Qwen3TTSTokenizerV1 model.""" -import math from dataclasses import dataclass import numpy as np import torch from torch import nn -from torch.nn import Parameter from torch.nn import functional as F from torch.nn.utils.rnn import pad_sequence from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel @@ -27,6 +25,8 @@ from transformers.utils import ModelOutput, auto_docstring, logging from transformers.utils.hub import cached_file from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding +from vllm_omni.model_executor.models.common.alias_free_activation import AliasFreeActivation1d +from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta from .configuration_qwen3_tts_tokenizer_v1 import ( Qwen3TTSTokenizerV1Config, @@ -663,166 +663,9 @@ class DiTDecoderLayer(nn.Module): return hidden_states -class SnakeBeta(nn.Module): - """ - A modified Snake function which uses separate parameters for the magnitude of the periodic components - Shape: - - Input: (B, C, T) - - Output: (B, C, T), same shape as the input - Parameters: - - alpha - trainable parameter that controls frequency - - beta - trainable parameter that controls magnitude - References: - - This activation function is a modified version based on this paper - by Liu Ziyin, Tilman Hartwig, Masahito Ueda: - https://huggingface.co/papers/2006.08195 - """ - - def __init__(self, in_features, alpha=1.0): - super().__init__() - self.in_features = in_features - - # initialize alpha - self.alpha = Parameter(torch.zeros(in_features) * alpha) - self.beta = Parameter(torch.zeros(in_features) * alpha) - - self.no_div_by_zero = 0.000000001 - - def forward(self, hidden_states): - """ - Forward pass of the function. - Applies the function to the input elementwise. - SnakeBeta ∶= x + 1/b * sin^2 (xa) - """ - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] - beta = self.beta.unsqueeze(0).unsqueeze(-1) - alpha = torch.exp(alpha) - beta = torch.exp(beta) - hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow( - torch.sin(hidden_states * alpha), 2 - ) - - return hidden_states - - -def kaiser_sinc_filter1d(cutoff, half_width, kernel_size) -> torch.Tensor: - """Generates a 1D Kaiser-windowed sinc filter. - - Args: - cutoff (float): Normalized cutoff frequency (0 to 0.5). - half_width (float): Transition bandwidth. - kernel_size (int): Number of filter taps. - - Returns: - torch.Tensor: A tensor of shape (1, 1, kernel_size) representing the filter. - """ - is_even = kernel_size % 2 == 0 - half_size = kernel_size // 2 - - # Compute Kaiser window parameters - delta_f = 4 * half_width - attenuation = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 - - if attenuation > 50.0: - beta = 0.1102 * (attenuation - 8.7) - elif attenuation >= 21.0: - beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0) - else: - beta = 0.0 - - kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32) - - # Compute time indices - if is_even: - time_indices = torch.arange(-half_size, half_size) + 0.5 - else: - time_indices = torch.arange(kernel_size) - half_size - - # Compute sinc filter - if cutoff == 0: - return torch.zeros((1, 1, kernel_size), dtype=torch.float32) # Ensures correct shape - - sinc_filter = torch.sinc(2 * cutoff * time_indices) - normalized_filter = 2 * cutoff * kaiser_window * sinc_filter - - # Normalize to ensure sum = 1 (avoid leakage of constant component) - normalized_filter /= normalized_filter.sum() - - return normalized_filter.view(1, 1, kernel_size) - - -class UpSample1d(nn.Module): - def __init__(self, ratio=2, kernel_size=None): - super().__init__() - self.ratio = ratio - self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size - self.stride = ratio - self.pad = self.kernel_size // ratio - 1 - self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 - self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 - - filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size) - self.register_buffer("filter", filter, persistent=False) - - def forward(self, hidden_states): - channels = hidden_states.shape[1] - - hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate") - hidden_states = self.ratio * F.conv_transpose1d( - hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels - ) - hidden_states = hidden_states[..., self.pad_left : -self.pad_right] - - return hidden_states - - -class DownSample1d(nn.Module): - def __init__(self, ratio=2, kernel_size=None): - super().__init__() - cutoff = 0.5 / ratio - half_width = 0.6 / ratio - - if cutoff < 0.0: - raise ValueError("Minimum cutoff must be larger than zero.") - if cutoff > 0.5: - raise ValueError("A cutoff above 0.5 does not make sense.") - - self.even = kernel_size % 2 == 0 - self.pad_left = kernel_size // 2 - int(self.even) - self.pad_right = kernel_size // 2 - self.stride = ratio - filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) - self.register_buffer("filter", filter, persistent=False) - - def forward(self, hidden_states): - channels = hidden_states.shape[1] - hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate") - out = F.conv1d(hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels) - return out - - -class TorchActivation1d(nn.Module): - def __init__( - self, - activation, - up_ratio: int = 2, - down_ratio: int = 2, - up_kernel_size: int = 12, - down_kernel_size: int = 12, - ): - super().__init__() - if not callable(activation): - raise TypeError("Activation function must be callable") - self.act = activation - self.upsample = UpSample1d(up_ratio, up_kernel_size) - self.downsample = DownSample1d(down_ratio, down_kernel_size) - - def forward(self, hidden_states): - hidden_states = self.upsample(hidden_states) - hidden_states = self.act(hidden_states) - hidden_states = self.downsample(hidden_states) - - return hidden_states +# SnakeBeta, UpSample1d, DownSample1d, AliasFreeActivation1d (formerly +# TorchActivation1d) are now provided by the shared ``common`` modules. +# See: vllm_omni/model_executor/models/common/{snake_activation,alias_free_activation}.py class CausalConv1d(nn.Conv1d): @@ -929,7 +772,7 @@ class AMPBlock(torch.nn.Module): self.num_layers = len(self.convs1) + len(self.convs2) # total number of conv layers self.activations = nn.ModuleList( - [TorchActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)] + [AliasFreeActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)] ) if causal_type == "2": @@ -940,7 +783,7 @@ class AMPBlock(torch.nn.Module): stride=1, padding=self._get_padding(kernel_size, 1), ) - self.pre_act = TorchActivation1d(activation=SnakeBeta(channels)) + self.pre_act = AliasFreeActivation1d(activation=SnakeBeta(channels)) else: self.pre_conv = nn.Identity() self.pre_act = nn.Identity() @@ -1002,7 +845,7 @@ class Qwen3TTSTokenizerV1DecoderBigVGANModel(Qwen3TTSTokenizerV1DecoderPreTraine ] ) - self.activation_post = TorchActivation1d( + self.activation_post = AliasFreeActivation1d( activation=SnakeBeta(config.upsample_initial_channel // (2**self.num_upsample_layers)) ) self.conv_post = nn.Conv1d(