toolbox/torchaudio/models/dfnet/modeling_dfnet.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
DeepFilterNet 的原生实现不直接支持流式推理

社区开发者（如 Rikorose）提供了基于 Torch 的流式推理实现
https://github.com/grazder/DeepFilterNet/tree/1097015d53ced78fb234e7d7071a5dd4446e3952/torchDF
"""
import os
import math
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F

from toolbox.torchaudio.configuration_utils import CONFIG_FILE
from toolbox.torchaudio.models.dfnet.configuration_dfnet import DfNetConfig
from toolbox.torchaudio.modules.conv_stft import ConvSTFT, ConviSTFT
from toolbox.torchaudio.modules.local_snr_target import LocalSnrTarget
from toolbox.torchaudio.modules.freq_bands.erb_bands import ErbBands


MODEL_FILE = "model.pt"


norm_layer_dict = {
    "batch_norm_2d": torch.nn.BatchNorm2d
}


activation_layer_dict = {
    "relu": torch.nn.ReLU,
    "identity": torch.nn.Identity,
    "sigmoid": torch.nn.Sigmoid,
}


class CausalConv2d(nn.Sequential):
    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 kernel_size: Union[int, Iterable[int]],
                 fstride: int = 1,
                 dilation: int = 1,
                 fpad: bool = True,
                 bias: bool = True,
                 separable: bool = False,
                 norm_layer: str = "batch_norm_2d",
                 activation_layer: str = "relu",
                 lookahead: int = 0
                 ):
        """
        Causal Conv2d by delaying the signal for any lookahead.

        Expected input format: [batch_size, channels, time_steps, spec_dim]

        :param in_channels:
        :param out_channels:
        :param kernel_size:
        :param fstride:
        :param dilation:
        :param fpad:
        """
        super(CausalConv2d, self).__init__()
        kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else tuple(kernel_size)

        if fpad:
            fpad_ = kernel_size[1] // 2 + dilation - 1
        else:
            fpad_ = 0

        # for last 2 dim, pad (left, right, top, bottom).
        pad = (0, 0, kernel_size[0] - 1 - lookahead, lookahead)

        layers = list()
        if any(x > 0 for x in pad):
            layers.append(nn.ConstantPad2d(pad, 0.0))

        groups = math.gcd(in_channels, out_channels) if separable else 1
        if groups == 1:
            separable = False
        if max(kernel_size) == 1:
            separable = False

        layers.append(
            nn.Conv2d(
                in_channels,
                out_channels,
                kernel_size=kernel_size,
                padding=(0, fpad_),
                stride=(1, fstride),  # stride over time is always 1
                dilation=(1, dilation),  # dilation over time is always 1
                groups=groups,
                bias=bias,
            )
        )

        if separable:
            layers.append(
                nn.Conv2d(
                    out_channels,
                    out_channels,
                    kernel_size=1,
                    bias=False,
                )
            )

        if norm_layer is not None:
            norm_layer = norm_layer_dict[norm_layer]
            layers.append(norm_layer(out_channels))

        if activation_layer is not None:
            activation_layer = activation_layer_dict[activation_layer]
            layers.append(activation_layer())

        super().__init__(*layers)

    def forward(self, inputs):
        for module in self:
            inputs = module(inputs)
        return inputs


class CausalConvTranspose2d(nn.Sequential):
    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 kernel_size: Union[int, Iterable[int]],
                 fstride: int = 1,
                 dilation: int = 1,
                 fpad: bool = True,
                 bias: bool = True,
                 separable: bool = False,
                 norm_layer: str = "batch_norm_2d",
                 activation_layer: str = "relu",
                 lookahead: int = 0
                 ):
        """
        Causal ConvTranspose2d.

        Expected input format: [batch_size, channels, time_steps, spec_dim]
        """
        super(CausalConvTranspose2d, self).__init__()

        kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size

        if fpad:
            fpad_ = kernel_size[1] // 2
        else:
            fpad_ = 0

        # for last 2 dim, pad (left, right, top, bottom).
        pad = (0, 0, kernel_size[0] - 1 - lookahead, lookahead)

        layers = []
        if any(x > 0 for x in pad):
            layers.append(nn.ConstantPad2d(pad, 0.0))

        groups = math.gcd(in_channels, out_channels) if separable else 1
        if groups == 1:
            separable = False

        layers.append(
            nn.ConvTranspose2d(
                in_channels,
                out_channels,
                kernel_size=kernel_size,
                padding=(kernel_size[0] - 1, fpad_ + dilation - 1),
                output_padding=(0, fpad_),
                stride=(1, fstride),  # stride over time is always 1
                dilation=(1, dilation),  # dilation over time is always 1
                groups=groups,
                bias=bias,
            )
        )

        if separable:
            layers.append(
                nn.Conv2d(
                    out_channels,
                    out_channels,
                    kernel_size=1,
                    bias=False,
                )
            )

        if norm_layer is not None:
            norm_layer = norm_layer_dict[norm_layer]
            layers.append(norm_layer(out_channels))

        if activation_layer is not None:
            activation_layer = activation_layer_dict[activation_layer]
            layers.append(activation_layer())

        super().__init__(*layers)


class GroupedLinear(nn.Module):

    def __init__(self, input_size: int, hidden_size: int, groups: int = 1):
        super().__init__()
        # self.weight: Tensor
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.groups = groups
        assert input_size % groups == 0, f"Input size {input_size} not divisible by {groups}"
        assert hidden_size % groups == 0, f"Hidden size {hidden_size} not divisible by {groups}"
        self.ws = input_size // groups
        self.register_parameter(
            "weight",
            torch.nn.Parameter(
                torch.zeros(groups, input_size // groups, hidden_size // groups), requires_grad=True
            ),
        )
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))  # type: ignore

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [..., I]
        b, t, f = x.shape
        if f != self.input_size:
            raise AssertionError

        # new_shape = list(x.shape)[:-1] + [self.groups, self.ws]
        new_shape = (b, t, self.groups, self.ws)
        x = x.view(new_shape)
        # The better way, but not supported by torchscript
        # x = x.unflatten(-1, (self.groups, self.ws))  # [..., G, I/G]
        x = torch.einsum("btgi,gih->btgh", x, self.weight)  # [..., G, H/G]
        x = x.flatten(2, 3)
        # x: [b, t, h]
        return x

    def __repr__(self):
        cls = self.__class__.__name__
        return f"{cls}(input_size: {self.input_size}, hidden_size: {self.hidden_size}, groups: {self.groups})"


class SqueezedGRU_S(nn.Module):
    """
    SGE net: Video object detection with squeezed GRU and information entropy map
    https://arxiv.org/abs/2106.07224
    """

    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        output_size: Optional[int] = None,
        num_layers: int = 1,
        linear_groups: int = 8,
        batch_first: bool = True,
        skip_op: str = "none",
        activation_layer: str = "identity",
    ):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size

        self.linear_in = nn.Sequential(
            GroupedLinear(
                input_size=input_size,
                hidden_size=hidden_size,
                groups=linear_groups,
            ),
            activation_layer_dict[activation_layer](),
        )

        # gru skip operator
        self.gru_skip_op = None

        if skip_op == "none":
            self.gru_skip_op = None
        elif skip_op == "identity":
            if not input_size != output_size:
                raise AssertionError("Dimensions do not match")
            self.gru_skip_op = nn.Identity()
        elif skip_op == "grouped_linear":
            self.gru_skip_op = GroupedLinear(
                input_size=hidden_size,
                hidden_size=hidden_size,
                groups=linear_groups,
            )
        else:
            raise NotImplementedError()

        self.gru = nn.GRU(
            input_size=hidden_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=batch_first,
            bidirectional=False,
        )

        if output_size is not None:
            self.linear_out = nn.Sequential(
                GroupedLinear(
                    input_size=hidden_size,
                    hidden_size=output_size,
                    groups=linear_groups,
                ),
                activation_layer_dict[activation_layer](),
            )
        else:
            self.linear_out = nn.Identity()

    def forward(self, inputs: torch.Tensor, h=None) -> Tuple[torch.Tensor, torch.Tensor]:
        # inputs: shape: [b, t, h]
        x = self.linear_in.forward(inputs)

        x, h = self.gru.forward(x, h)

        x = self.linear_out(x)

        if self.gru_skip_op is not None:
            x = x + self.gru_skip_op(inputs)

        return x, h


class Add(nn.Module):
    def forward(self, a, b):
        return a + b


class Concat(nn.Module):
    def forward(self, a, b):
        return torch.cat((a, b), dim=-1)


class Encoder(nn.Module):
    def __init__(self, config: DfNetConfig):
        super(Encoder, self).__init__()
        self.embedding_input_size = config.conv_channels * config.erb_bins // 4
        self.embedding_output_size = config.conv_channels * config.erb_bins // 4
        self.embedding_hidden_size = config.embedding_hidden_size

        self.spec_conv0 = CausalConv2d(
            in_channels=1,
            out_channels=config.conv_channels,
            kernel_size=config.conv_kernel_size_input,
            bias=False,
            separable=True,
            fstride=1,
            lookahead=config.conv_lookahead,
        )
        self.spec_conv1 = CausalConv2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=config.conv_kernel_size_inner,
            bias=False,
            separable=True,
            fstride=2,
            lookahead=config.conv_lookahead,
        )
        self.spec_conv2 = CausalConv2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=config.conv_kernel_size_inner,
            bias=False,
            separable=True,
            fstride=2,
            lookahead=config.conv_lookahead,
        )
        self.spec_conv3 = CausalConv2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=config.conv_kernel_size_inner,
            bias=False,
            separable=True,
            fstride=1,
            lookahead=config.conv_lookahead,
        )

        self.df_conv0 = CausalConv2d(
            in_channels=2,
            out_channels=config.conv_channels,
            kernel_size=config.conv_kernel_size_input,
            bias=False,
            separable=True,
            fstride=1,
        )
        self.df_conv1 = CausalConv2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=config.conv_kernel_size_inner,
            bias=False,
            separable=True,
            fstride=2,
        )
        self.df_fc_emb = nn.Sequential(
            GroupedLinear(
                config.conv_channels * config.df_bins // 2,
                self.embedding_input_size,
                groups=config.encoder_linear_groups
            ),
            nn.ReLU(inplace=True)
        )

        if config.encoder_combine_op == "concat":
            self.embedding_input_size *= 2
            self.combine = Concat()
        else:
            self.combine = Add()

        # emb_gru
        if config.spec_bins % 8 != 0:
            raise AssertionError("spec_bins should be divisible by 8")

        self.emb_gru = SqueezedGRU_S(
            self.embedding_input_size,
            self.embedding_hidden_size,
            output_size=self.embedding_output_size,
            num_layers=1,
            batch_first=True,
            skip_op=config.encoder_emb_skip_op,
            linear_groups=config.encoder_emb_linear_groups,
            activation_layer="relu",
        )

        # lsnr
        self.lsnr_fc = nn.Sequential(
            nn.Linear(self.embedding_output_size, 1),
            nn.Sigmoid()
        )
        self.lsnr_scale = config.max_local_snr - config.min_local_snr
        self.lsnr_offset = config.min_local_snr

    def forward(self,
                feat_erb: torch.Tensor,
                feat_spec: torch.Tensor,
                hidden_state: torch.Tensor = None,
                ):
        # feat_erb shape: (b, 1, t, erb_bins)
        e0 = self.spec_conv0.forward(feat_erb)
        e1 = self.spec_conv1.forward(e0)
        e2 = self.spec_conv2.forward(e1)
        e3 = self.spec_conv3.forward(e2)
        # e0 shape: [b, c, t, erb_bins]
        # e1 shape: [b, c, t, erb_bins // 2]
        # e2 shape: [b, c, t, erb_bins // 4]
        # e3 shape: [b, c, t, erb_bins // 4]
        # e3 shape: [b, 64, t, 32/4=8]

        # feat_spec, shape: (b, 2, t, df_bins)
        c0 = self.df_conv0(feat_spec)
        c1 = self.df_conv1(c0)
        # c0 shape: [b, c, t, df_bins]
        # c1 shape: [b, c, t, df_bins // 2]
        # c1 shape: [b, 64, t, 96/2=48]

        cemb = c1.permute(0, 2, 3, 1)
        # cemb shape: [b, t, df_bins // 2, c]
        cemb = cemb.flatten(2)
        # cemb shape: [b, t, df_bins // 2 * c]
        # cemb shape: [b, t, 96/2*64=3072]
        cemb = self.df_fc_emb.forward(cemb)
        # cemb shape: [b, t, erb_bins // 4 * c]
        # cemb shape: [b, t, 32/4*64=512]

        # e3 shape: [b, c, t, erb_bins // 4]
        emb = e3.permute(0, 2, 3, 1)
        # emb shape: [b, t, erb_bins // 4, c]
        emb = emb.flatten(2)
        # emb shape: [b, t, erb_bins // 4 * c]
        # emb shape: [b, t, 32/4*64=512]

        emb = self.combine(emb, cemb)
        # if concat; emb shape: [b, t, spec_bins // 4 * c * 2]
        # if add; emb shape: [b, t, spec_bins // 4 * c]

        emb, h = self.emb_gru.forward(emb, hidden_state)

        # emb shape: [b, t, spec_dim // 4 * c]
        # h shape: [b, 1, spec_dim]

        lsnr = self.lsnr_fc(emb) * self.lsnr_scale + self.lsnr_offset
        # lsnr shape: [b, t, 1]

        return e0, e1, e2, e3, emb, c0, lsnr, h


class Decoder(nn.Module):
    """ErbDecoder"""
    def __init__(self, config: DfNetConfig):
        super(Decoder, self).__init__()

        if config.spec_bins % 8 != 0:
            raise AssertionError("spec_bins should be divisible by 8")

        self.emb_in_dim = config.conv_channels * config.erb_bins // 4
        self.emb_out_dim = config.conv_channels * config.erb_bins // 4
        self.emb_hidden_dim = config.decoder_emb_hidden_size

        self.emb_gru = SqueezedGRU_S(
            self.emb_in_dim,
            self.emb_hidden_dim,
            output_size=self.emb_out_dim,
            num_layers=config.decoder_emb_num_layers - 1,
            batch_first=True,
            skip_op=config.decoder_emb_skip_op,
            linear_groups=config.decoder_emb_linear_groups,
            activation_layer="relu",
        )
        self.conv3p = CausalConv2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=1,
            bias=False,
            separable=True,
            fstride=1,
            lookahead=config.conv_lookahead,
        )
        self.convt3 = CausalConv2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=config.conv_kernel_size_inner,
            bias=False,
            separable=True,
            fstride=1,
            lookahead=config.conv_lookahead,
        )
        self.conv2p = CausalConv2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=1,
            bias=False,
            separable=True,
            fstride=1,
            lookahead=config.conv_lookahead,
        )
        self.convt2 = CausalConvTranspose2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=config.convt_kernel_size_inner,
            bias=False,
            separable=True,
            fstride=2,
            lookahead=config.conv_lookahead,
        )
        self.conv1p = CausalConv2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=1,
            bias=False,
            separable=True,
            fstride=1,
            lookahead=config.conv_lookahead,
        )
        self.convt1 = CausalConvTranspose2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=config.convt_kernel_size_inner,
            bias=False,
            separable=True,
            fstride=2,
            lookahead=config.conv_lookahead,
        )
        self.conv0p = CausalConv2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=1,
            bias=False,
            separable=True,
            fstride=1,
            lookahead=config.conv_lookahead,
        )
        self.conv0_out = CausalConv2d(
            in_channels=config.conv_channels,
            out_channels=1,
            kernel_size=config.conv_kernel_size_inner,
            activation_layer="sigmoid",
            bias=False,
            separable=True,
            fstride=1,
            lookahead=config.conv_lookahead,
        )

    def forward(self, emb, e3, e2, e1, e0) -> torch.Tensor:
        # Estimates erb mask
        b, _, t, f8 = e3.shape

        # emb shape: [batch_size, time_steps, (freq_dim // 4) * conv_channels]
        emb, _ = self.emb_gru.forward(emb)
        # emb shape: [batch_size, conv_channels, time_steps, freq_dim // 4]
        emb = emb.view(b, t, f8, -1).permute(0, 3, 1, 2)
        e3 = self.convt3(self.conv3p(e3) + emb)
        # e3 shape: [batch_size, conv_channels, time_steps, freq_dim // 4]
        e2 = self.convt2(self.conv2p(e2) + e3)
        # e2 shape: [batch_size, conv_channels, time_steps, freq_dim // 2]
        e1 = self.convt1(self.conv1p(e1) + e2)
        # e1 shape: [batch_size, conv_channels, time_steps, freq_dim]
        mask = self.conv0_out(self.conv0p(e0) + e1)
        # mask shape: [batch_size, 1, time_steps, freq_dim]
        return mask


class DfDecoder(nn.Module):
    def __init__(self, config: DfNetConfig):
        super(DfDecoder, self).__init__()

        self.embedding_input_size = config.conv_channels * config.erb_bins // 4
        self.df_decoder_hidden_size = config.df_decoder_hidden_size
        self.df_num_layers = config.df_num_layers

        self.df_order = config.df_order

        self.df_bins = config.df_bins
        self.df_out_ch = config.df_order * 2

        self.df_convp = CausalConv2d(
            config.conv_channels,
            self.df_out_ch,
            fstride=1,
            kernel_size=(config.df_pathway_kernel_size_t, 1),
            separable=True,
            bias=False,
        )
        self.df_gru = SqueezedGRU_S(
            self.embedding_input_size,
            self.df_decoder_hidden_size,
            num_layers=self.df_num_layers,
            batch_first=True,
            skip_op="none",
            activation_layer="relu",
        )

        if config.df_gru_skip == "none":
            self.df_skip = None
        elif config.df_gru_skip == "identity":
            if config.embedding_hidden_size != config.df_decoder_hidden_size:
                raise AssertionError("Dimensions do not match")
            self.df_skip = nn.Identity()
        elif config.df_gru_skip == "grouped_linear":
            self.df_skip = GroupedLinear(
                self.embedding_input_size,
                self.df_decoder_hidden_size,
                groups=config.df_decoder_linear_groups
            )
        else:
            raise NotImplementedError()

        self.df_out: nn.Module
        out_dim = self.df_bins * self.df_out_ch

        self.df_out = nn.Sequential(
            GroupedLinear(
                input_size=self.df_decoder_hidden_size,
                hidden_size=out_dim,
                groups=config.df_decoder_linear_groups,
                # groups = self.df_bins // 5,
        ),
            nn.Tanh()
        )
        self.df_fc_a = nn.Sequential(
            nn.Linear(self.df_decoder_hidden_size, 1),
            nn.Sigmoid()
        )

    def forward(self, emb: torch.Tensor, c0: torch.Tensor) -> torch.Tensor:
        # emb shape: [batch_size, time_steps, df_bins // 4 * channels]
        b, t, _ = emb.shape
        df_coefs, _ = self.df_gru(emb)
        if self.df_skip is not None:
            df_coefs = df_coefs + self.df_skip(emb)
        # df_coefs shape: [batch_size, time_steps, df_decoder_hidden_size]

        # c0 shape: [batch_size, channels, time_steps, df_bins]
        c0 = self.df_convp(c0)
        # c0 shape: [batch_size, df_order * 2, time_steps, df_bins]
        c0 = c0.permute(0, 2, 3, 1)
        # c0 shape: [batch_size, time_steps, df_bins, df_order * 2]

        df_coefs = self.df_out(df_coefs)  # [B, T, F*O*2], O: df_order
        # df_coefs shape: [batch_size, time_steps, df_bins * df_order * 2]
        df_coefs = df_coefs.view(b, t, self.df_bins, self.df_out_ch)
        # df_coefs shape: [batch_size, time_steps, df_bins, df_order * 2]
        df_coefs = df_coefs + c0
        # df_coefs shape: [batch_size, time_steps, df_bins, df_order * 2]
        return df_coefs


class DfOutputReshapeMF(nn.Module):
    """Coefficients output reshape for multiframe/MultiFrameModule

    Requires input of shape B, C, T, F, 2.
    """

    def __init__(self, df_order: int, df_bins: int):
        super().__init__()
        self.df_order = df_order
        self.df_bins = df_bins

    def forward(self, coefs: torch.Tensor) -> torch.Tensor:
        # [B, T, F, O*2] -> [B, O, T, F, 2]
        new_shape = list(coefs.shape)
        new_shape[-1] = -1
        new_shape.append(2)
        coefs = coefs.view(new_shape)
        coefs = coefs.permute(0, 3, 1, 2, 4)
        return coefs


class Mask(nn.Module):
    def __init__(self, use_post_filter: bool = False, eps: float = 1e-12):
        super().__init__()
        self.use_post_filter = use_post_filter
        self.eps = eps

    def post_filter(self, mask: torch.Tensor, beta: float = 0.02) -> torch.Tensor:
        """
        Post-Filter

        A Perceptually-Motivated Approach for Low-Complexity, Real-Time Enhancement of Fullband Speech.
        https://arxiv.org/abs/2008.04259

        :param mask: Real valued mask, typically of shape [B, C, T, F].
        :param beta: Global gain factor.
        :return:
        """
        mask_sin = mask * torch.sin(np.pi * mask / 2)
        mask_pf = (1 + beta) * mask / (1 + beta * mask.div(mask_sin.clamp_min(self.eps)).pow(2))
        return mask_pf

    def forward(self, spec: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        # spec shape: [b, 1, t, spec_bins, 2]

        if not self.training and self.use_post_filter:
            mask = self.post_filter(mask)

        # mask shape: [b, 1, t, spec_bins]
        mask = mask.unsqueeze(4)
        # mask shape: [b, 1, t, spec_bins, 1]
        return spec * mask


class DeepFiltering(nn.Module):
    def __init__(self,
                 df_bins: int,
                 df_order: int,
                 lookahead: int = 0,
                 ):
        super(DeepFiltering, self).__init__()
        self.df_bins = df_bins
        self.df_order = df_order
        self.need_unfold = df_order > 1
        self.lookahead = lookahead

        self.pad = nn.ConstantPad2d((0, 0, df_order - 1 - lookahead, lookahead), 0.0)

    def spec_unfold(self, spec: torch.Tensor):
        """
        Pads and unfolds the spectrogram according to frame_size.
        :param spec: complex Tensor, Spectrogram of shape [B, C, T, F].
        :return: Tensor, Unfolded spectrogram of shape [B, C, T, F, N], where N: frame_size.
        """
        if self.need_unfold:
            # spec shape: [batch_size, spec_bins, time_steps]
            spec_pad = self.pad(spec)
            # spec_pad shape: [batch_size, 1, time_steps_pad, spec_bins]
            spec_unfold = spec_pad.unfold(2, self.df_order, 1)
            # spec_unfold shape: [batch_size, 1, time_steps, spec_bins, df_order]
            return spec_unfold
        else:
            return spec.unsqueeze(-1)

    def forward(self,
                spec: torch.Tensor,
                coefs: torch.Tensor,
                ):
        # spec shape: [batch_size, 1, time_steps, spec_bins, 2]
        spec_u = self.spec_unfold(torch.view_as_complex(spec.contiguous()))
        # spec_u shape: [batch_size, 1, time_steps, spec_bins, df_order]

        # coefs shape: [batch_size, df_order, time_steps, df_bins, 2]
        coefs = torch.view_as_complex(coefs.contiguous())
        # coefs shape: [batch_size, df_order, time_steps, df_bins]
        spec_f = spec_u.narrow(-2, 0, self.df_bins)
        # spec_f shape: [batch_size, 1, time_steps, df_bins, df_order]

        coefs = coefs.view(coefs.shape[0], -1, self.df_order, *coefs.shape[2:])
        # coefs shape: [batch_size, 1, df_order, time_steps, df_bins]

        spec_f = self.df(spec_f, coefs)
        # spec_f shape: [batch_size, 1, time_steps, df_bins]

        if self.training:
            spec = spec.clone()
        spec[..., :self.df_bins, :] = torch.view_as_real(spec_f)
        # spec shape: [batch_size, 1, time_steps, spec_bins, 2]
        return spec

    @staticmethod
    def df(spec: torch.Tensor, coefs: torch.Tensor) -> torch.Tensor:
        """
        Deep filter implementation using `torch.einsum`. Requires unfolded spectrogram.
        :param spec: (complex Tensor). Spectrogram of shape [B, C, T, F, N].
        :param coefs: (complex Tensor). Coefficients of shape [B, C, N, T, F].
        :return: (complex Tensor). Spectrogram of shape [B, C, T, F].
        """
        return torch.einsum("...tfn,...ntf->...tf", spec, coefs)


class DfNet(nn.Module):
    """
    我感觉这个模型没办法实现完全一致的流式推理。
    """
    def __init__(self, config: DfNetConfig):
        super(DfNet, self).__init__()
        self.config = config
        self.eps = 1e-12

        self.freq_bins = self.config.nfft // 2 + 1

        self.nfft = config.nfft
        self.win_size = config.win_size
        self.hop_size = config.hop_size
        self.win_type = config.win_type

        self.erb_bands = ErbBands(
            sample_rate=config.sample_rate,
            nfft=config.nfft,
            erb_bins=config.erb_bins,
            min_freq_bins_for_erb=config.min_freq_bins_for_erb,
        )

        self.stft = ConvSTFT(
            nfft=config.nfft,
            win_size=config.win_size,
            hop_size=config.hop_size,
            win_type=config.win_type,
            power=None,
            requires_grad=False
        )
        self.istft = ConviSTFT(
            nfft=config.nfft,
            win_size=config.win_size,
            hop_size=config.hop_size,
            win_type=config.win_type,
            requires_grad=False
        )

        self.encoder = Encoder(config)
        self.decoder = Decoder(config)

        self.df_decoder = DfDecoder(config)
        self.df_out_transform = DfOutputReshapeMF(config.df_order, config.df_bins)
        self.df_op = DeepFiltering(
            df_bins=config.df_bins,
            df_order=config.df_order,
            lookahead=config.df_lookahead,
        )

        self.mask = Mask(use_post_filter=config.use_post_filter)

        self.lsnr_fn = LocalSnrTarget(
            sample_rate=config.sample_rate,
            nfft=config.nfft,
            win_size=config.win_size,
            hop_size=config.hop_size,
            n_frame=config.n_frame,
            min_local_snr=config.min_local_snr,
            max_local_snr=config.max_local_snr,
            db=True,
        )

    def signal_prepare(self, signal: torch.Tensor) -> torch.Tensor:
        if signal.dim() == 2:
            signal = torch.unsqueeze(signal, dim=1)
        _, _, n_samples = signal.shape
        remainder = (n_samples - self.win_size) % self.hop_size
        if remainder > 0:
            n_samples_pad = self.hop_size - remainder
            signal = F.pad(signal, pad=(0, n_samples_pad), mode="constant", value=0)
        return signal

    def feature_prepare(self, signal: torch.Tensor):
        # noisy shape: [b, num_samples_pad]
        spec_cmp = self.stft.forward(signal)
        # spec_complex shape: [b, f, t], torch.complex64
        spec_cmp = torch.transpose(spec_cmp, dim0=1, dim1=2)
        # spec_complex shape: [b, t, f], torch.complex64
        spec_cmp_real = torch.view_as_real(spec_cmp)
        # spec_cmp_real shape: [b, t, f, 2]
        spec_mag = torch.abs(spec_cmp)
        spec_pow = torch.square(spec_mag)
        # shape: [b, t, f]

        spec = torch.unsqueeze(spec_cmp_real, dim=1)
        # spec shape: [b, 1, t, f, 2]

        feat_erb = self.erb_bands.erb_scale(spec_pow, db=True)
        # feat_erb shape: [b, t, erb_bins]
        feat_erb = torch.unsqueeze(feat_erb, dim=1)
        # feat_erb shape: [b, 1, t, erb_bins]

        feat_spec = spec_cmp_real.permute(0, 3, 1, 2)
        # feat_spec shape: [b, 2, t, f]
        feat_spec = feat_spec[..., :self.df_decoder.df_bins]
        # feat_spec shape: [b, 2, t, df_bins]

        return spec, feat_erb, feat_spec

    def forward(self,
                noisy: torch.Tensor,
                ):
        """
        :param noisy:
        :return:
        est_spec: shape: [b, 257*2, t]
        est_wav:  shape: [b, num_samples]
        est_mask: shape: [b, 257, t]
        lsnr:     shape: [b, 1, t]
        """
        n_samples = noisy.shape[-1]
        noisy = self.signal_prepare(noisy)

        spec, feat_erb, feat_spec = self.feature_prepare(noisy)

        e0, e1, e2, e3, emb, c0, lsnr, h = self.encoder.forward(feat_erb, feat_spec)

        mask = self.decoder.forward(emb, e3, e2, e1, e0)
        # mask shape: [b, 1, t, erb_bins]
        mask = self.erb_bands.erb_scale_inv(mask)
        # mask shape: [b, 1, t, f]
        if torch.any(mask > 1) or torch.any(mask < 0):
            raise AssertionError

        spec_m = self.mask.forward(spec, mask)
        # spec_m shape: [b, 1, t, f, 2]
        spec_m = spec_m[:, :, :, :self.config.spec_bins, :]
        # spec_m shape: [b, 1, t, spec_bins, 2]

        # lsnr shape: [b, t, 1]
        lsnr = torch.transpose(lsnr, dim0=2, dim1=1)
        # lsnr shape: [b, 1, t]

        df_coefs = self.df_decoder.forward(emb, c0)
        df_coefs = self.df_out_transform(df_coefs)
        # df_coefs shape: [b, df_order, t, df_bins, 2]

        spec_ = spec[:, :, :, :self.config.spec_bins, :]
        # spec shape: [b, 1, t, spec_bins, 2]
        spec_e = self.df_op.forward(spec_, df_coefs)
        # spec_e shape: [b, 1, t, spec_bins, 2]

        spec_e[..., self.df_decoder.df_bins:, :] = spec_m[..., self.df_decoder.df_bins:, :]

        spec_e = torch.squeeze(spec_e, dim=1)
        spec_e = spec_e.permute(0, 2, 1, 3)
        # spec_e shape: [b, spec_bins, t, 2]

        # spec_e shape: [b, spec_bins, t, 2]
        est_spec = torch.complex(real=spec_e[..., 0], imag=spec_e[..., 1])
        # est_spec shape: [b, spec_bins, t], torch.complex64
        est_spec = torch.concat(tensors=[est_spec, est_spec[:, -1:, :]], dim=1)
        # est_spec shape: [b, f, t], torch.complex64

        est_wav = self.istft.forward(est_spec)
        est_wav = est_wav[:, :, :n_samples]
        # est_wav shape: [b, 1, n_samples]

        est_mask = torch.squeeze(mask, dim=1)
        est_mask = est_mask.permute(0, 2, 1)
        # est_mask shape: [b, f, t]

        return est_spec, est_wav, est_mask, lsnr

    def mask_loss_fn(self, est_mask: torch.Tensor, clean: torch.Tensor, noisy: torch.Tensor):
        """
        :param est_mask: torch.Tensor, shape: [b, 257, t]
        :param clean:
        :param noisy:
        :return:
        """
        if noisy.shape != clean.shape:
            raise AssertionError("Input signals must have the same shape")
        noise = noisy - clean

        clean = self.signal_prepare(clean)
        noise = self.signal_prepare(noise)

        stft_clean = self.stft.forward(clean)
        mag_clean = torch.abs(stft_clean)

        stft_noise = self.stft.forward(noise)
        mag_noise = torch.abs(stft_noise)

        gth_irm_mask = (mag_clean / (mag_clean + mag_noise + self.eps)).clamp(0, 1)

        loss = F.l1_loss(gth_irm_mask, est_mask, reduction="mean")

        return loss

    def lsnr_loss_fn(self, lsnr: torch.Tensor, clean: torch.Tensor, noisy: torch.Tensor):
        if noisy.shape != clean.shape:
            raise AssertionError("Input signals must have the same shape")
        noise = noisy - clean

        clean = self.signal_prepare(clean)
        noise = self.signal_prepare(noise)

        stft_clean = self.stft.forward(clean)
        stft_noise = self.stft.forward(noise)
        # shape: [b, f, t]
        stft_clean = torch.transpose(stft_clean, dim0=1, dim1=2)
        stft_noise = torch.transpose(stft_noise, dim0=1, dim1=2)
        # shape: [b, t, f]
        stft_clean = torch.unsqueeze(stft_clean, dim=1)
        stft_noise = torch.unsqueeze(stft_noise, dim=1)
        # shape: [b, 1, t, f]

        # lsnr shape: [b, 1, t]
        lsnr = lsnr.squeeze(1)
        # lsnr shape: [b, t]

        lsnr_gth = self.lsnr_fn.forward(stft_clean, stft_noise)
        # lsnr_gth shape: [b, t]

        loss = F.mse_loss(lsnr, lsnr_gth)
        return loss


class DfNetPretrainedModel(DfNet):
    def __init__(self,
                 config: DfNetConfig,
                 ):
        super(DfNetPretrainedModel, self).__init__(
            config=config,
        )

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        config = DfNetConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)

        model = cls(config)

        if os.path.isdir(pretrained_model_name_or_path):
            ckpt_file = os.path.join(pretrained_model_name_or_path, MODEL_FILE)
        else:
            ckpt_file = pretrained_model_name_or_path

        with open(ckpt_file, "rb") as f:
            state_dict = torch.load(f, map_location="cpu", weights_only=True)
        model.load_state_dict(state_dict, strict=True)
        return model

    def save_pretrained(self,
                        save_directory: Union[str, os.PathLike],
                        state_dict: Optional[dict] = None,
                        ):

        model = self

        if state_dict is None:
            state_dict = model.state_dict()

        os.makedirs(save_directory, exist_ok=True)

        # save state dict
        model_file = os.path.join(save_directory, MODEL_FILE)
        torch.save(state_dict, model_file)

        # save config
        config_file = os.path.join(save_directory, CONFIG_FILE)
        self.config.to_yaml_file(config_file)
        return save_directory


def main():

    config = DfNetConfig()
    model = DfNetPretrainedModel(config=config)

    noisy = torch.randn(size=(1, 16000), dtype=torch.float32)

    est_spec, est_wav, est_mask, lsnr = model.forward(noisy)
    print(f"est_spec.shape: {est_spec.shape}")
    print(f"est_wav.shape: {est_wav.shape}")
    print(f"est_mask.shape: {est_mask.shape}")
    print(f"lsnr.shape: {lsnr.shape}")

    return


if __name__ == "__main__":
    main()