toolbox/torchaudio/modules/conv_stft.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
https://github.com/modelscope/modelscope/blob/master/modelscope/models/audio/ans/conv_stft.py
"""
from collections import defaultdict
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy.signal import get_window


def init_kernels(nfft: int, win_size: int, hop_size: int, win_type: str = None, inverse=False):
    if win_type == "None" or win_type is None:
        window = np.ones(win_size)
    else:
        window = get_window(win_type, win_size, fftbins=True)**0.5

    fourier_basis = np.fft.rfft(np.eye(nfft))[:win_size]
    real_kernel = np.real(fourier_basis)
    image_kernel = np.imag(fourier_basis)
    kernel = np.concatenate([real_kernel, image_kernel], 1).T

    if inverse:
        kernel = np.linalg.pinv(kernel).T

    kernel = kernel * window
    kernel = kernel[:, None, :]
    result = (
        torch.from_numpy(kernel.astype(np.float32)),
        torch.from_numpy(window[None, :, None].astype(np.float32))
    )
    return result


class ConvSTFT(nn.Module):

    def __init__(self,
                 nfft: int,
                 win_size: int,
                 hop_size: int,
                 win_type: str = "hamming",
                 power: int = None,
                 requires_grad: bool = False):
        super(ConvSTFT, self).__init__()

        if nfft is None:
            self.nfft = int(2**np.ceil(np.log2(win_size)))
        else:
            self.nfft = nfft

        kernel, _ = init_kernels(self.nfft, win_size, hop_size, win_type)
        self.weight = nn.Parameter(kernel, requires_grad=requires_grad)

        self.win_size = win_size
        self.hop_size = hop_size

        self.stride = hop_size
        self.dim = self.nfft
        self.power = power

    def forward(self, waveform: torch.Tensor):
        if waveform.dim() == 2:
            waveform = torch.unsqueeze(waveform, 1)

        matrix = F.conv1d(waveform, self.weight, stride=self.stride)
        dim = self.dim // 2 + 1
        real = matrix[:, :dim, :]
        imag = matrix[:, dim:, :]
        spec = torch.complex(real, imag)
        # spec shape: [b, f, t], torch.complex64

        if self.power is None:
            return spec
        elif self.power == 1:
            mags = torch.sqrt(real**2 + imag**2)
            # phase = torch.atan2(imag, real)
            return mags
        elif self.power == 2:
            power = real**2 + imag**2
            return power
        else:
            raise AssertionError


class ConviSTFT(nn.Module):

    def __init__(self,
                 win_size: int,
                 hop_size: int,
                 nfft: int = None,
                 win_type: str = "hamming",
                 requires_grad: bool = False):
        super(ConviSTFT, self).__init__()
        if nfft is None:
            self.nfft = int(2**np.ceil(np.log2(win_size)))
        else:
            self.nfft = nfft

        kernel, window = init_kernels(self.nfft, win_size, hop_size, win_type, inverse=True)
        self.weight = nn.Parameter(kernel, requires_grad=requires_grad)
        # weight shape: [f*2, 1, nfft]
        # f = nfft // 2 + 1

        self.win_size = win_size
        self.hop_size = hop_size
        self.win_type = win_type

        self.stride = hop_size
        self.dim = self.nfft

        self.register_buffer("window", window)
        self.register_buffer("enframe", torch.eye(win_size)[:, None, :])
        # window shape: [1, nfft, 1]
        # enframe shape: [nfft, 1, nfft]

    def forward(self,
                spec: torch.Tensor):
        """
        self.weight shape: [f*2, 1, win_size]
        self.window shape: [1, win_size, 1]
        self.enframe shape: [win_size, 1, win_size]

        :param spec: torch.Tensor, shape: [b, f, t, 2]
        :return:
        """
        spec = torch.view_as_real(spec)
        # spec shape: [b, f, t, 2]
        matrix = torch.concat(tensors=[spec[..., 0], spec[..., 1]], dim=1)
        # matrix shape: [b, f*2, t]

        waveform = F.conv_transpose1d(matrix, self.weight, stride=self.stride)
        # waveform shape: [b, 1, num_samples]

        # this is from torch-stft: https://github.com/pseeth/torch-stft
        t = self.window.repeat(1, 1, matrix.size(-1))**2
        # t shape: [1, win_size, t]
        coff = F.conv_transpose1d(t, self.enframe, stride=self.stride)
        # coff shape: [1, 1, num_samples]
        waveform = waveform / (coff + 1e-8)
        # waveform = waveform / coff
        return waveform

    @torch.no_grad()
    def forward_chunk(self,
                      spec: torch.Tensor,
                      cache_dict: dict = None
                      ):
        """
        :param spec: shape: [b, f, t]
        :param cache_dict: dict,
        waveform_cache shape: [b, 1, win_size - hop_size]
        coff_cache shape: [b, 1, win_size - hop_size]
        :return:
        """
        if cache_dict is None:
            cache_dict = defaultdict(lambda: None)
        waveform_cache = cache_dict["waveform_cache"]
        coff_cache = cache_dict["coff_cache"]

        spec = torch.view_as_real(spec)
        matrix = torch.concat(tensors=[spec[..., 0], spec[..., 1]], dim=1)

        waveform_current = F.conv_transpose1d(matrix, self.weight, stride=self.stride)

        t = self.window.repeat(1, 1, matrix.size(-1))**2
        coff_current = F.conv_transpose1d(t, self.enframe, stride=self.stride)

        overlap_size = self.win_size - self.hop_size

        if waveform_cache is not None:
            waveform_current[:, :, :overlap_size] += waveform_cache
        waveform_output = waveform_current[:, :, :self.hop_size]
        new_waveform_cache = waveform_current[:, :, self.hop_size:]

        if coff_cache is not None:
            coff_current[:, :, :overlap_size] += coff_cache
        coff_output = coff_current[:, :, :self.hop_size]
        new_coff_cache = coff_current[:, :, self.hop_size:]

        waveform_output = waveform_output / (coff_output + 1e-8)

        new_cache_dict = {
            "waveform_cache": new_waveform_cache,
            "coff_cache": new_coff_cache,
        }
        return waveform_output, new_cache_dict


def main():
    nfft = 512
    win_size = 512
    hop_size = 256

    stft = ConvSTFT(nfft=nfft, win_size=win_size, hop_size=hop_size, power=None)
    istft = ConviSTFT(nfft=nfft, win_size=win_size, hop_size=hop_size)

    mixture = torch.rand(size=(1, 16000), dtype=torch.float32)
    b, num_samples = mixture.shape
    t = (num_samples - win_size) / hop_size + 1

    spec = stft.forward(mixture)
    b, f, t = spec.shape

    # 如果 spec 是由 stft 变换得来的，以下两种 waveform 还原方法就是一致的，否则还原出的 waveform 会有差异。
    # spec = spec + 0.01 * torch.randn(size=(1, nfft//2+1, t), dtype=torch.float32)
    print(f"spec.shape: {spec.shape}, spec.dtype: {spec.dtype}")

    waveform = istft.forward(spec)
    # shape: [batch_size, channels, num_samples]
    print(f"waveform.shape: {waveform.shape}, waveform.dtype: {waveform.dtype}")
    print(waveform[:, :, 300: 302])

    waveform = torch.zeros(size=(b, 1, num_samples), dtype=torch.float32)
    for i in range(int(t)):
        begin = i * hop_size
        end = begin + win_size
        sub_spec = spec[:, :, i:i+1]
        sub_waveform = istft.forward(sub_spec)
        # (b, 1, win_size)
        waveform[:, :, begin:end] = sub_waveform
    print(f"waveform.shape: {waveform.shape}, waveform.dtype: {waveform.dtype}")
    print(waveform[:, :, 300: 302])

    return


def main2():
    nfft = 512
    win_size = 512
    hop_size = 256

    stft = ConvSTFT(nfft=nfft, win_size=win_size, hop_size=hop_size, power=None)
    istft = ConviSTFT(nfft=nfft, win_size=win_size, hop_size=hop_size)

    mixture = torch.rand(size=(1, 16128), dtype=torch.float32)
    b, num_samples = mixture.shape

    spec = stft.forward(mixture)
    b, f, t = spec.shape

    # 如果 spec 是由 stft 变换得来的，以下两种 waveform 还原方法就是一致的，否则还原出的 waveform 会有差异。
    spec = spec + 0.01 * torch.randn(size=(1, nfft//2+1, t), dtype=torch.float32)
    print(f"spec.shape: {spec.shape}, spec.dtype: {spec.dtype}")

    waveform = istft.forward(spec)
    # shape: [batch_size, channels, num_samples]
    print(f"waveform.shape: {waveform.shape}, waveform.dtype: {waveform.dtype}")
    print(waveform[:, :, 300: 302])

    cache_dict = None
    waveform = torch.zeros(size=(b, 1, num_samples), dtype=torch.float32)
    for i in range(int(t)):
        sub_spec = spec[:, :, i:i+1]
        begin = i * hop_size

        end = begin + win_size - hop_size
        sub_waveform, cache_dict = istft.forward_chunk(sub_spec, cache_dict=cache_dict)
        # end = begin + win_size
        # sub_waveform = istft.forward(sub_spec)

        waveform[:, :, begin:end] = sub_waveform
    print(f"waveform.shape: {waveform.shape}, waveform.dtype: {waveform.dtype}")
    print(waveform[:, :, 300: 302])

    return


if __name__ == "__main__":
    main2()