[Refactor] Unify Snake/SnakeBeta and alias-free activation into common modules (#3886)

Signed-off-by: BeatSeat <wendavid552@gmail.com>

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

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

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.
 

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

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(

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(