src/models/SoulX-LiveAct/kokoro/custom_stft.py
7,746 bytes · 198 lines · capsule://quake0day/[email protected]
raw on github
from attr import attr
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
class CustomSTFT(nn.Module):
"""
STFT/iSTFT without unfold/complex ops, using conv1d and conv_transpose1d.
- forward STFT => Real-part conv1d + Imag-part conv1d
- inverse STFT => Real-part conv_transpose1d + Imag-part conv_transpose1d + sum
- avoids F.unfold, so easier to export to ONNX
- uses replicate or constant padding for 'center=True' to approximate 'reflect'
(reflect is not supported for dynamic shapes in ONNX)
"""
def __init__(
self,
filter_length=800,
hop_length=200,
win_length=800,
window="hann",
center=True,
pad_mode="replicate", # or 'constant'
):
super().__init__()
self.filter_length = filter_length
self.hop_length = hop_length
self.win_length = win_length
self.n_fft = filter_length
self.center = center
self.pad_mode = pad_mode
# Number of frequency bins for real-valued STFT with onesided=True
self.freq_bins = self.n_fft // 2 + 1
# Build window
assert window == 'hann', window
window_tensor = torch.hann_window(win_length, periodic=True, dtype=torch.float32)
if self.win_length < self.n_fft:
# Zero-pad up to n_fft
extra = self.n_fft - self.win_length
window_tensor = F.pad(window_tensor, (0, extra))
elif self.win_length > self.n_fft:
window_tensor = window_tensor[: self.n_fft]
self.register_buffer("window", window_tensor)
# Precompute forward DFT (real, imag)
# PyTorch stft uses e^{-j 2 pi k n / N} => real=cos(...), imag=-sin(...)
n = np.arange(self.n_fft)
k = np.arange(self.freq_bins)
angle = 2 * np.pi * np.outer(k, n) / self.n_fft # shape (freq_bins, n_fft)
dft_real = np.cos(angle)
dft_imag = -np.sin(angle) # note negative sign
# Combine window and dft => shape (freq_bins, filter_length)
# We'll make 2 conv weight tensors of shape (freq_bins, 1, filter_length).
forward_window = window_tensor.numpy() # shape (n_fft,)
forward_real = dft_real * forward_window # (freq_bins, n_fft)
forward_imag = dft_imag * forward_window
# Convert to PyTorch
forward_real_torch = torch.from_numpy(forward_real).float()
forward_imag_torch = torch.from_numpy(forward_imag).float()
# Register as Conv1d weight => (out_channels, in_channels, kernel_size)
# out_channels = freq_bins, in_channels=1, kernel_size=n_fft
self.register_buffer(
"weight_forward_real", forward_real_torch.unsqueeze(1)
)
self.register_buffer(
"weight_forward_imag", forward_imag_torch.unsqueeze(1)
)
# Precompute inverse DFT
# Real iFFT formula => scale = 1/n_fft, doubling for bins 1..freq_bins-2 if n_fft even, etc.
# For simplicity, we won't do the "DC/nyquist not doubled" approach here.
# If you want perfect real iSTFT, you can add that logic.
# This version just yields good approximate reconstruction with Hann + typical overlap.
inv_scale = 1.0 / self.n_fft
n = np.arange(self.n_fft)
angle_t = 2 * np.pi * np.outer(n, k) / self.n_fft # shape (n_fft, freq_bins)
idft_cos = np.cos(angle_t).T # => (freq_bins, n_fft)
idft_sin = np.sin(angle_t).T # => (freq_bins, n_fft)
# Multiply by window again for typical overlap-add
# We also incorporate the scale factor 1/n_fft
inv_window = window_tensor.numpy() * inv_scale
backward_real = idft_cos * inv_window # (freq_bins, n_fft)
backward_imag = idft_sin * inv_window
# We'll implement iSTFT as real+imag conv_transpose with stride=hop.
self.register_buffer(
"weight_backward_real", torch.from_numpy(backward_real).float().unsqueeze(1)
)
self.register_buffer(
"weight_backward_imag", torch.from_numpy(backward_imag).float().unsqueeze(1)
)
def transform(self, waveform: torch.Tensor):
"""
Forward STFT => returns magnitude, phase
Output shape => (batch, freq_bins, frames)
"""
# waveform shape => (B, T). conv1d expects (B, 1, T).
# Optional center pad
if self.center:
pad_len = self.n_fft // 2
waveform = F.pad(waveform, (pad_len, pad_len), mode=self.pad_mode)
x = waveform.unsqueeze(1) # => (B, 1, T)
# Convolution to get real part => shape (B, freq_bins, frames)
real_out = F.conv1d(
x,
self.weight_forward_real,
bias=None,
stride=self.hop_length,
padding=0,
)
# Imag part
imag_out = F.conv1d(
x,
self.weight_forward_imag,
bias=None,
stride=self.hop_length,
padding=0,
)
# magnitude, phase
magnitude = torch.sqrt(real_out**2 + imag_out**2 + 1e-14)
phase = torch.atan2(imag_out, real_out)
# Handle the case where imag_out is 0 and real_out is negative to correct ONNX atan2 to match PyTorch
# In this case, PyTorch returns pi, ONNX returns -pi
correction_mask = (imag_out == 0) & (real_out < 0)
phase[correction_mask] = torch.pi
return magnitude, phase
def inverse(self, magnitude: torch.Tensor, phase: torch.Tensor, length=None):
"""
Inverse STFT => returns waveform shape (B, T).
"""
# magnitude, phase => (B, freq_bins, frames)
# Re-create real/imag => shape (B, freq_bins, frames)
real_part = magnitude * torch.cos(phase)
imag_part = magnitude * torch.sin(phase)
# conv_transpose wants shape (B, freq_bins, frames). We'll treat "frames" as time dimension
# so we do (B, freq_bins, frames) => (B, freq_bins, frames)
# But PyTorch conv_transpose1d expects (B, in_channels, input_length)
real_part = real_part # (B, freq_bins, frames)
imag_part = imag_part
# real iSTFT => convolve with "backward_real", "backward_imag", and sum
# We'll do 2 conv_transpose calls, each giving (B, 1, time),
# then add them => (B, 1, time).
real_rec = F.conv_transpose1d(
real_part,
self.weight_backward_real, # shape (freq_bins, 1, filter_length)
bias=None,
stride=self.hop_length,
padding=0,
)
imag_rec = F.conv_transpose1d(
imag_part,
self.weight_backward_imag,
bias=None,
stride=self.hop_length,
padding=0,
)
# sum => (B, 1, time)
waveform = real_rec - imag_rec # typical real iFFT has minus for imaginary part
# If we used "center=True" in forward, we should remove pad
if self.center:
pad_len = self.n_fft // 2
# Because of transposed convolution, total length might have extra samples
# We remove `pad_len` from start & end if possible
waveform = waveform[..., pad_len:-pad_len]
# If a specific length is desired, clamp
if length is not None:
waveform = waveform[..., :length]
# shape => (B, T)
return waveform
def forward(self, x: torch.Tensor):
"""
Full STFT -> iSTFT pass: returns time-domain reconstruction.
Same interface as your original code.
"""
mag, phase = self.transform(x)
return self.inverse(mag, phase, length=x.shape[-1])