toolbox/torchaudio/models/dfnet2/modeling_dfnet2.py

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

社区开发者（如 Rikorose）提供了基于 Torch 的流式推理实现
https://github.com/grazder/DeepFilterNet/tree/1097015d53ced78fb234e7d7071a5dd4446e3952/torchDF

此文件试图实现一个支持流式推理的 dfnet

"""
import os
import math
from collections import defaultdict
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.dfnet2.configuration_dfnet2 import DfNet2Config
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
from toolbox.torchaudio.modules.utils.ema import ErbEMA, SpecEMA


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.Module):
    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 kernel_size: Union[int, Iterable[int]],
                 fstride: int = 1,
                 dilation: int = 1,
                 pad_f_dim: bool = True,
                 bias: bool = True,
                 separable: bool = False,
                 norm_layer: str = "batch_norm_2d",
                 activation_layer: str = "relu",
                 ):
        super(CausalConv2d, self).__init__()
        kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else tuple(kernel_size)

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

        # for last 2 dim, pad (left, right, top, bottom).
        self.lookback = kernel_size[0] - 1
        if self.lookback > 0:
            self.tpad = nn.ConstantPad2d(padding=(0, 0, self.lookback, 0), value=0.0)
        else:
            self.tpad = nn.Identity()

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

        self.conv = 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:
            self.convp = nn.Conv2d(
                out_channels,
                out_channels,
                kernel_size=1,
                bias=False,
            )
        else:
            self.convp = nn.Identity()

        if norm_layer is not None:
            norm_layer = norm_layer_dict[norm_layer]
            self.norm = norm_layer(out_channels)
        else:
            self.norm = nn.Identity()

        if activation_layer is not None:
            activation_layer = activation_layer_dict[activation_layer]
            self.activation = activation_layer()
        else:
            self.activation = nn.Identity()

    def forward(self, inputs: torch.Tensor, cache: Tuple[torch.Tensor, torch.Tensor] = None):
        """
        :param inputs: shape: [b, c, t, f]
        :param cache: shape: [b, c, lookback, f];
        :return:
        """
        x = inputs

        if cache is None:
            x = self.tpad(x)
        else:
            x = torch.concat(tensors=[cache, x], dim=2)

        new_cache = None
        if self.lookback > 0:
            new_cache = x[:, :, -self.lookback:, :]

        x = self.conv(x)

        x = self.convp(x)
        x = self.norm(x)
        x = self.activation(x)

        return x, new_cache


class CausalConvTranspose2dErrorCase(nn.Module):
    """
    错误的缓存方法。
    """
    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 kernel_size: Union[int, Iterable[int]],
                 fstride: int = 1,
                 dilation: int = 1,
                 pad_f_dim: bool = True,
                 bias: bool = True,
                 separable: bool = False,
                 norm_layer: str = "batch_norm_2d",
                 activation_layer: str = "relu",
                 ):
        super(CausalConvTranspose2dErrorCase, self).__init__()

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

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

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

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

        self.convt = nn.ConvTranspose2d(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            padding=(0, fpad),
            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:
            self.convp = nn.Conv2d(
                out_channels,
                out_channels,
                kernel_size=1,
                bias=False,
            )
        else:
            self.convp = nn.Identity()

        if norm_layer is not None:
            norm_layer = norm_layer_dict[norm_layer]
            self.norm = norm_layer(out_channels)
        else:
            self.norm = nn.Identity()

        if activation_layer is not None:
            activation_layer = activation_layer_dict[activation_layer]
            self.activation = activation_layer()
        else:
            self.activation = nn.Identity()

    def forward(self, inputs: torch.Tensor, cache: torch.Tensor = None):
        """
        :param inputs: shape: [b, c, t, f]
        :param cache: shape: [b, c, lookback, f];
        :return:
        """
        x = inputs

        # x shape: [b, c, t, f]
        x = self.convt(x)
        # x shape: [b, c, t+lookback, f]

        new_cache = None
        if self.lookback > 0:
            if cache is not None:
                x = torch.concat(tensors=[
                    x[:, :, :self.lookback, :] + cache,
                    x[:, :, self.lookback:, :]
                ], dim=2)

            x = x[:, :, :-self.lookback, :]
            new_cache = x[:, :, -self.lookback:, :]

        x = self.convp(x)
        x = self.norm(x)
        x = self.activation(x)

        return x, new_cache


class CausalConvTranspose2d(nn.Module):
    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 kernel_size: Union[int, Iterable[int]],
                 fstride: int = 1,
                 dilation: int = 1,
                 pad_f_dim: bool = True,
                 bias: bool = True,
                 separable: bool = False,
                 norm_layer: str = "batch_norm_2d",
                 activation_layer: str = "relu",
                 ):
        super(CausalConvTranspose2d, self).__init__()

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

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

        # for last 2 dim, pad (left, right, top, bottom).
        self.lookback = kernel_size[0] - 1
        if self.lookback > 0:
            self.tpad = nn.ConstantPad2d(padding=(0, 0, self.lookback, 0), value=0.0)
        else:
            self.tpad = nn.Identity()

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

        self.convt = 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:
            self.convp = nn.Conv2d(
                out_channels,
                out_channels,
                kernel_size=1,
                bias=False,
            )
        else:
            self.convp = nn.Identity()

        if norm_layer is not None:
            norm_layer = norm_layer_dict[norm_layer]
            self.norm = norm_layer(out_channels)
        else:
            self.norm = nn.Identity()

        if activation_layer is not None:
            activation_layer = activation_layer_dict[activation_layer]
            self.activation = activation_layer()
        else:
            self.activation = nn.Identity()

    def forward(self, inputs: torch.Tensor, cache: torch.Tensor = None):
        """
        :param inputs: shape: [b, c, t, f]
        :param cache: shape: [b, c, lookback, f];
        :return:
        """
        x = inputs

        # x shape: [b, c, t, f]
        x = self.convt(x)
        # x shape: [b, c, t+lookback, f]

        if cache is None:
            x = self.tpad(x)
        else:
            x = torch.concat(tensors=[cache, x], dim=2)

        new_cache = None
        if self.lookback > 0:
            new_cache = x[:, :, -self.lookback:, :]

        x = self.convp(x)
        x = self.norm(x)
        x = self.activation(x)

        return x, new_cache


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, hx: torch.Tensor = None) -> Tuple[torch.Tensor, torch.Tensor]:
        # inputs: shape: [b, t, h]
        x = self.linear_in.forward(inputs)

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

        x = self.linear_out(x)

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

        return x, hx


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: DfNet2Config):
        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,
        )
        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,
        )
        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,
        )
        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,
        )

        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,
                cache_dict: dict = None,
                ):
        if cache_dict is None:
            cache_dict = defaultdict(lambda: None)
        cache0 = cache_dict["cache0"]
        cache1 = cache_dict["cache1"]
        cache2 = cache_dict["cache2"]
        cache3 = cache_dict["cache3"]
        cache4 = cache_dict["cache4"]
        cache5 = cache_dict["cache5"]
        cache6 = cache_dict["cache6"]

        # feat_erb shape: (b, 1, t, erb_bins)
        e0, new_cache0 = self.spec_conv0.forward(feat_erb, cache=cache0)
        e1, new_cache1 = self.spec_conv1.forward(e0, cache=cache1)
        e2, new_cache2 = self.spec_conv2.forward(e1, cache=cache2)
        e3, new_cache3 = self.spec_conv3.forward(e2, cache=cache3)
        # 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, new_cache4 = self.df_conv0.forward(feat_spec, cache=cache4)
        c1, new_cache5 = self.df_conv1.forward(c0, cache=cache5)
        # 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, new_cache6 = self.emb_gru.forward(emb, hx=cache6)

        # 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]

        new_cache_dict = {
            "cache0": new_cache0,
            "cache1": new_cache1,
            "cache2": new_cache2,
            "cache3": new_cache3,
            "cache4": new_cache4,
            "cache5": new_cache5,
            "cache6": new_cache6,
        }
        return e0, e1, e2, e3, emb, c0, lsnr, new_cache_dict


class ErbDecoder(nn.Module):
    def __init__(self, config: DfNet2Config):
        super(ErbDecoder, 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,
        )
        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,
        )
        self.conv2p = CausalConv2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=1,
            bias=False,
            separable=True,
            fstride=1,
        )
        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,
        )
        self.conv1p = CausalConv2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=1,
            bias=False,
            separable=True,
            fstride=1,
        )
        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,
        )
        self.conv0p = CausalConv2d(
            in_channels=config.conv_channels,
            out_channels=config.conv_channels,
            kernel_size=1,
            bias=False,
            separable=True,
            fstride=1,
        )
        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,
        )

    def forward(self, emb, e3, e2, e1, e0, cache_dict: dict = None) -> torch.Tensor:
        if cache_dict is None:
            cache_dict = defaultdict(lambda: None)
        cache0 = cache_dict["cache0"]
        cache1 = cache_dict["cache1"]
        cache2 = cache_dict["cache2"]
        cache3 = cache_dict["cache3"]
        cache4 = cache_dict["cache4"]

        # Estimates erb mask
        b, _, t, f8 = e3.shape

        # emb shape: [batch_size, time_steps, (freq_dim // 4) * conv_channels]
        emb, new_cache0 = self.emb_gru.forward(emb, hx=cache0)
        # emb shape: [batch_size, conv_channels, time_steps, freq_dim // 4]
        emb = emb.view(b, t, f8, -1).permute(0, 3, 1, 2)

        e3, new_cache1 = self.convt3.forward(self.conv3p(e3)[0] + emb, cache=cache1)
        # e3 shape: [batch_size, conv_channels, time_steps, freq_dim // 4]
        e2, new_cache2 = self.convt2.forward(self.conv2p(e2)[0] + e3, cache=cache2)
        # e2 shape: [batch_size, conv_channels, time_steps, freq_dim // 2]
        e1, new_cache3 = self.convt1.forward(self.conv1p(e1)[0] + e2, cache=cache3)
        # e1 shape: [batch_size, conv_channels, time_steps, freq_dim]
        mask, new_cache4 = self.conv0_out.forward(self.conv0p(e0)[0] + e1, cache=cache4)
        # mask shape: [batch_size, 1, time_steps, freq_dim]

        new_cache_dict = {
            "cache0": new_cache0,
            "cache1": new_cache1,
            "cache2": new_cache2,
            "cache3": new_cache3,
            "cache4": new_cache4,
        }
        return mask, new_cache_dict


class DfDecoder(nn.Module):
    def __init__(self, config: DfNet2Config):
        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, cache_dict: dict = None) -> torch.Tensor:
        if cache_dict is None:
            cache_dict = defaultdict(lambda: None)
        cache0 = cache_dict["cache0"]
        cache1 = cache_dict["cache1"]

        # emb shape: [batch_size, time_steps, df_bins // 4 * channels]
        b, t, _ = emb.shape
        df_coefs, new_cache0 = self.df_gru.forward(emb, hx=cache0)
        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, new_cache1 = self.df_convp.forward(c0, cache=cache1)
        # 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]

        new_cache_dict = {
            "cache0": new_cache0,
            "cache1": new_cache1,
        }
        return df_coefs, new_cache_dict


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.lookahead = lookahead

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

    def forward(self, *args, **kwargs):
        raise AssertionError("use `forward_offline` or `forward_online` stead.")

    def spec_unfold_offline(self, spec: torch.Tensor) -> torch.Tensor:
        """
        Pads and unfolds the spectrogram according to frame_size.
        :param spec: shape: [b, c, t, f], dtype: torch.complex64
        :return: shape: [b, c, t, f, df_order]
        """
        if self.df_order <= 1:
            return spec.unsqueeze(-1)

        # spec shape: [b, 1, t, f], dtype: torch.complex64
        spec = self.pad(spec)
        # spec_pad shape: [b, 1, t+df_order-1, f], dtype: torch.complex64
        spec_unfold = spec.unfold(dimension=2, size=self.df_order, step=1)
        # spec_unfold shape: [b, 1, t, f, df_order], dtype: torch.complex64
        return spec_unfold

    def forward_offline(self,
                        spec: torch.Tensor,
                        coefs: torch.Tensor,
                        ):
        # spec shape: [b, 1, t, spec_bins, 2]
        spec_c = torch.view_as_complex(spec.contiguous())
        # spec_c shape: [b, 1, t, spec_bins]
        spec_u = self.spec_unfold_offline(spec_c)
        # spec_u shape: [b, 1, t, spec_bins, df_order]
        spec_f = spec_u.narrow(dim=-2, start=0, length=self.df_bins)
        # spec_f shape: [b, 1, t, df_bins, df_order]

        # coefs shape: [b, df_order, t, df_bins, 2]
        coefs = torch.view_as_complex(coefs.contiguous())
        # coefs shape: [b, df_order, t, df_bins]
        coefs = coefs.unsqueeze(dim=1)
        # coefs shape: [b, 1, df_order, t, df_bins]

        spec_f = self.df_offline(spec_f, coefs)
        # spec_f shape: [b, 1, t, df_bins]

        spec_f = torch.view_as_real(spec_f)
        # spec_f shape: [b, 1, t, df_bins, 2]
        return spec_f

    def df_offline(self, spec: torch.Tensor, coefs: torch.Tensor):
        """
        Deep filter implementation using `torch.einsum`. Requires unfolded spectrogram.
        :param spec: [b, 1, t, df_bins, df_order] complex.
        :param coefs: [b, 1, df_order, t, df_bins] complex.
        :return: [b, 1, t, df_bins] complex.
        """
        spec_f = torch.einsum("...tfn,...ntf->...tf", spec, coefs)
        return spec_f

    def spec_unfold_online(self, spec: torch.Tensor, cache_spec: torch.Tensor = None):
        """
        Pads and unfolds the spectrogram according to frame_size.
        :param spec: shape: [b, c, t, f], dtype: torch.complex64
        :param cache_spec: shape: [b, c, df_order-1, f], dtype: torch.complex64
        :return: shape: [b, c, t, f, df_order]
        """
        if self.df_order <= 1:
            return spec.unsqueeze(-1)

        if cache_spec is None:
            b, c, _, f = spec.shape
            cache_spec = spec.new_zeros(size=(b, c, self.df_order-1, f))
        spec_pad = torch.concat(tensors=[
            cache_spec, spec
        ], dim=2)
        new_cache_spec = spec_pad[:, :, -(self.df_order-1):, :]

        # spec_pad shape: [b, 1, t+df_order-1, f], dtype: torch.complex64
        spec_unfold = spec_pad.unfold(dimension=2, size=self.df_order, step=1)
        # spec_unfold shape: [b, 1, t, f, df_order], dtype: torch.complex64
        return spec_unfold, new_cache_spec

    def forward_online(self,
                       spec: torch.Tensor,
                       coefs: torch.Tensor,
                       cache_dict: dict = None,
                       ):
        if cache_dict is None:
            cache_dict = defaultdict(lambda: None)
        cache0 = cache_dict["cache0"]
        cache1 = cache_dict["cache1"]

        # spec shape: [b, 1, t, spec_bins, 2]
        spec_c = torch.view_as_complex(spec.contiguous())
        # spec_c shape: [b, 1, t, spec_bins]
        spec_u, new_cache0 = self.spec_unfold_online(spec_c, cache_spec=cache0)
        # spec_u shape: [b, 1, t, spec_bins, df_order]
        spec_f = spec_u.narrow(dim=-2, start=0, length=self.df_bins)
        # spec_f shape: [b, 1, t, df_bins, df_order]

        # coefs shape: [b, df_order, t, df_bins, 2]
        coefs = torch.view_as_complex(coefs.contiguous())
        # coefs shape: [b, df_order, t, df_bins]
        coefs = coefs.unsqueeze(dim=1)
        # coefs shape: [b, 1, df_order, t, df_bins]

        spec_f, new_cache1 = self.df_online(spec_f, coefs, cache_coefs=cache1)
        # spec_f shape: [b, 1, t, df_bins]

        spec_f = torch.view_as_real(spec_f)
        # spec_f shape: [b, 1, t, df_bins, 2]

        new_cache_dict = {
            "cache0": new_cache0,
            "cache1": new_cache1,
        }
        return spec_f, new_cache_dict

    def df_online(self, spec: torch.Tensor, coefs: torch.Tensor, cache_coefs: torch.Tensor = None) -> torch.Tensor:
        """
        Deep filter implementation using `torch.einsum`. Requires unfolded spectrogram.
        :param spec: [b, 1, 1, df_bins, df_order] complex.
        :param coefs: [b, 1, df_order, 1, df_bins] complex.
        :param cache_coefs: [b, 1, df_order, lookahead, df_bins] complex.
        :return: [b, 1, 1, df_bins] complex.
        """

        if cache_coefs is None:
            b, c, _, _, f = coefs.shape
            cache_coefs = coefs.new_zeros(size=(b, c, self.df_order, self.lookahead, f))
        coefs_pad = torch.concat(tensors=[
            cache_coefs, coefs
        ], dim=3)

        # coefs_pad shape: [b, 1, df_order, 1+lookahead, df_bins], torch.complex64.
        coefs = coefs_pad[:, :, :, :-self.lookahead, :]
        # coefs shape: [b, 1, df_order, 1, df_bins], torch.complex64.
        new_cache_coefs = coefs_pad[:, :, :, -self.lookahead:, :]
        # new_cache_coefs shape: [b, 1, df_order, lookahead, df_bins], torch.complex64.
        spec_f = torch.einsum("...tfn,...ntf->...tf", spec, coefs)
        return spec_f, new_cache_coefs


class DfNet2(nn.Module):
    def __init__(self, config: DfNet2Config):
        super(DfNet2, 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.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.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.erb_ema = ErbEMA(
            sample_rate=config.sample_rate,
            hop_size=config.hop_size,
            erb_bins=config.erb_bins,
        )
        self.spec_ema = SpecEMA(
            sample_rate=config.sample_rate,
            hop_size=config.hop_size,
            df_bins=config.df_bins,
        )

        self.encoder = Encoder(config)
        self.erb_decoder = ErbDecoder(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]

        spec = spec.detach()
        feat_erb = feat_erb.detach()
        feat_spec = feat_spec.detach()
        return spec, feat_erb, feat_spec

    def feature_norm(self, feat_erb, feat_spec, cache_dict: dict = None):
        if cache_dict is None:
            cache_dict = defaultdict(lambda: None)
        cache0 = cache_dict["cache0"]
        cache1 = cache_dict["cache1"]

        feat_erb, new_cache0 = self.erb_ema.norm(feat_erb, state=cache0)
        feat_spec, new_cache1 = self.spec_ema.norm(feat_spec, state=cache1)

        new_cache_dict = {
            "cache0": new_cache0,
            "cache1": new_cache1,
        }

        feat_erb = feat_erb.detach()
        feat_spec = feat_spec.detach()
        return feat_erb, feat_spec, new_cache_dict

    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)
        if self.config.use_ema_norm:
            feat_erb, feat_spec, _ = self.feature_norm(feat_erb, feat_spec)

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

        mask, _ = self.erb_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_f = self.df_op.forward_offline(spec_, df_coefs)
        # spec_f shape: [b, 1, t, df_bins, 2], torch.float32

        spec_e = torch.concat(tensors=[
            spec_f, spec_m[..., self.df_decoder.df_bins:, :]
        ], dim=3)

        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.view_as_complex(spec_e.contiguous())
        # 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 forward_chunk_by_chunk(self,
                               noisy: torch.Tensor,
                               ):
        noisy = self.signal_prepare(noisy)
        b, _, _ = noisy.shape
        noisy = torch.concat(tensors=[
            noisy, noisy.new_zeros(size=(b, 1, (self.config.df_lookahead+1)*self.hop_size))
        ], dim=2)
        b, _, num_samples = noisy.shape

        t = (num_samples - self.win_size) // self.hop_size + 1

        cache_dict0 = None
        cache_dict1 = None
        cache_dict2 = None
        cache_dict3 = None
        cache_dict4 = None
        cache_dict5 = None
        cache_dict6 = None

        waveform_list = list()
        for i in range(int(t)):
            begin = i * self.hop_size
            end = begin + self.win_size
            sub_noisy = noisy[:, :, begin: end]

            spec, feat_erb, feat_spec = self.feature_prepare(sub_noisy)
            # spec shape: [b, 1, t, f, 2]
            # feat_erb shape: [b, 1, t, erb_bins]
            # feat_spec shape: [b, 2, t, df_bins]
            if self.config.use_ema_norm:
                feat_erb, feat_spec, cache_dict0 = self.feature_norm(feat_erb, feat_spec, cache_dict=cache_dict0)

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

            mask, cache_dict2 = self.erb_decoder.forward(emb, e3, e2, e1, e0, cache_dict=cache_dict2)
            # mask shape: [b, 1, t, erb_bins]
            mask = self.erb_bands.erb_scale_inv(mask)
            # mask shape: [b, 1, t, f]

            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, cache_dict3 = self.df_decoder.forward(emb, c0, cache_dict=cache_dict3)
            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_f, cache_dict4 = self.df_op.forward_online(spec_, df_coefs, cache_dict=cache_dict4)
            # spec_f shape: [b, 1, t, df_bins, 2], torch.float32

            spec_e, cache_dict5 = self.spec_e_m_combine_online(spec_f, spec_m, cache_dict=cache_dict5)

            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.view_as_complex(spec_e.contiguous())
            # 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, cache_dict6 = self.istft.forward_chunk(est_spec, cache_dict=cache_dict6)
            # est_wav shape: [b, 1, hop_size]

            waveform_list.append(est_wav)

        waveform = torch.concat(tensors=waveform_list, dim=-1)
        # waveform shape: [b, 1, n]
        return waveform

    def spec_e_m_combine_online(self, spec_f: torch.Tensor, spec_m: torch.Tensor, cache_dict: dict = None):
        """
        :param spec_f: shape: [b, 1, t, df_bins, 2], torch.float32
        :param spec_m: shape: [b, 1, t, spec_bins, 2]
        :param cache_dict:
        :return:
        """
        if cache_dict is None:
            cache_dict = defaultdict(lambda: None)
        cache_spec_m = cache_dict["cache_spec_m"]

        if cache_spec_m is None:
            b, c, t, f, _ = spec_m.shape
            cache_spec_m = spec_m.new_zeros(size=(b, c, self.config.df_lookahead, f, 2))
            # cache0 shape: [b, 1, lookahead, f, 2]
        spec_m_cat = torch.concat(tensors=[
            cache_spec_m, spec_m,
        ], dim=2)

        spec_m = spec_m_cat[:, :, :-self.config.df_lookahead, :, :]
        new_cache_spec_m = spec_m_cat[:, :, -self.config.df_lookahead:, :, :]

        spec_e = torch.concat(tensors=[
            spec_f, spec_m[..., self.df_decoder.df_bins:, :]
        ], dim=3)

        new_cache_dict = {
            "cache_spec_m": new_cache_spec_m,
        }
        return spec_e, new_cache_dict

    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 DfNet2PretrainedModel(DfNet2):
    def __init__(self,
                 config: DfNet2Config,
                 ):
        super(DfNet2PretrainedModel, self).__init__(
            config=config,
        )

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        config = DfNet2Config.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():
    import time
    # torch.set_num_threads(1)

    config = DfNet2Config()
    model = DfNet2PretrainedModel(config=config)
    model.eval()

    num_samples = 16000
    noisy = torch.randn(size=(1, num_samples), dtype=torch.float32)
    duration = num_samples / config.sample_rate

    begin = time.time()
    est_spec, est_wav, est_mask, lsnr = model.forward(noisy)
    time_cost = time.time() - begin
    print(f"time_cost: {time_cost:.4f}, audio_duration: {duration:.4f}, fpr: {time_cost / duration:.4f}")

    # 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}")

    waveform = est_wav
    print(f"waveform.shape: {waveform.shape}, waveform.dtype: {waveform.dtype}")
    print(waveform[:, :, 300: 302])
    print(waveform[:, :, 15680: 15682])
    print(waveform[:, :, 15760: 15762])
    print(waveform[:, :, 15840: 15842])

    begin = time.time()
    waveform = model.forward_chunk_by_chunk(noisy)
    time_cost = time.time() - begin
    print(f"time_cost: {time_cost:.4f}, audio_duration: {duration:.4f}, fpr: {time_cost / duration:.4f}")

    waveform = waveform[:, :, (config.df_lookahead*config.hop_size):]
    print(f"waveform.shape: {waveform.shape}, waveform.dtype: {waveform.dtype}")
    print(waveform[:, :, 300: 302])
    print(waveform[:, :, 15680: 15682])
    print(waveform[:, :, 15760: 15762])
    print(waveform[:, :, 15840: 15842])

    return


if __name__ == "__main__":
    main()