1

config/EEYD_large.yaml

@@ -1,22 +0,0 @@
-#!/bin/bash 
-mode: 'inference'
-use_cuda: 1 # 1 for True, 0 for False
-num_gpu: 1
-sampling_rate: 16000
-network: "EEYD_base"  # network type
-checkpoint_dir: "checkpoints/EEYD_base"
-
-# decode parameters
-one_time_decode_length: 10 # maximum segment length for one-pass decoding (seconds), longer audio (>5s) will use segmented decoding
-decode_window: 10 # one-pass decoding length
-
-
-# network settings
-network_reference:
-  cue: text           # 
-  emb_size: 512       # resnet18:256
-  text_layers: 3
-  text_network: t5  # default t5, or clap
-network_audio:
-  backbone: mrx
-

extract_everything.py

@@ -25,6 +25,7 @@ class main(nn.Module):
         parser.add_argument('--sampling-rate', dest='sampling_rate', type=int, default=16000, help='Sampling rate')
         parser.add_argument('--one-time-decode-length', dest='one_time_decode_length', type=int, default=60, help='Max segment length for one-pass decoding')
         parser.add_argument('--decode-window', dest='decode_window', type=int, default=1, help='Decoding chunk size')
+        parser.add_argument('--output_residual',  type=int, default=0)
 
         # Parse arguments from the config file
         self.args = parser.parse_args(['--config', self.config_path])

models/mossformer2/mossformer/utils/Transformer.py

@@ -1,460 +0,0 @@
-"""Transformer implementaion for Mossformer2 
-
-Authors
-* Shengkui Zhao 2024
-* Jia Qi Yip 2024
-"""
-
-import math
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch import einsum
-
-from typing import Optional
-import numpy as np
-
-# from ..utils.flash_pytorch_fsmn import FLASHTransformer_DualA_FSMN
-# from sb.nnet.normalization import LayerNorm
-# from speechbrain.lobes.models.layer_norm import CLayerNorm, GLayerNorm, GlobLayerNorm, ILayerNorm
-# from speechbrain.lobes.models.fsmn import UniDeepFsmn, UniDeepFsmn_dilated
-# from speechbrain.lobes.models.conv_module import ConvModule
-from einops import rearrange
-from rotary_embedding_torch import RotaryEmbedding
-
-from ..utils.fsmn import UniDeepFsmn, UniDeepFsmn_dilated
-from ..utils.normalization import LayerNorm, CLayerNorm, ScaleNorm
-from ..utils.conv_module import ConvModule
-
-def exists(val):
-    return val is not None
-
-def padding_to_multiple_of(n, mult):
-    remainder = n % mult
-    if remainder == 0:
-        return 0
-    return mult - remainder
-
-def default(val, d):
-    return val if exists(val) else d
-
-class FFConvM(nn.Module):
-    def __init__(
-        self,
-        dim_in,
-        dim_out,
-        norm_klass = nn.LayerNorm,
-        dropout = 0.1
-    ):
-        super().__init__()
-        self.mdl = nn.Sequential(
-            norm_klass(dim_in),
-            nn.Linear(dim_in, dim_out),
-            nn.SiLU(),
-            ConvModule(dim_out),
-            nn.Dropout(dropout)
-        )
-    def forward(
-        self,
-        x,
-    ):
-        output = self.mdl(x)
-        return output
-
-class Gated_FSMN_dilated(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        lorder,
-        hidden_size
-    ):
-        super().__init__()
-        self.to_u = FFConvM(
-            dim_in = in_channels,
-            dim_out = hidden_size,
-            norm_klass = nn.LayerNorm,
-            dropout = 0.1,
-            )
-        self.to_v = FFConvM(
-            dim_in = in_channels,
-            dim_out = hidden_size,
-            norm_klass = nn.LayerNorm,
-            dropout = 0.1,
-            )
-        self.fsmn = UniDeepFsmn_dilated(in_channels, out_channels, lorder, hidden_size)
-
-    def forward(
-        self,
-        x,
-    ):
-        input = x
-        x_u = self.to_u(x)
-        x_v = self.to_v(x) 
-        x_u = self.fsmn(x_u)
-        x = x_v * x_u + input               
-        return x
-
-class Gated_FSMN_Block_Dilated(nn.Module):
-    """1-D convolutional block."""
-
-    def __init__(self,
-                 dim,
-                 inner_channels = 256,
-                 group_size = 256, #384, #128, #256,
-                 #query_key_dim = 128, #256, #128,
-                 #expansion_factor = 4.,
-                 #causal = False,
-                 #dropout = 0.1,
-                 norm_type = 'scalenorm',
-                 #shift_tokens = True,
-                 #rotary_pos_emb = None,
-                 ):
-        super(Gated_FSMN_Block_Dilated, self).__init__()
-        if norm_type == 'scalenorm':
-            norm_klass = ScaleNorm
-        elif norm_type == 'layernorm':
-            norm_klass = nn.LayerNorm
-
-        self.group_size = group_size
-
-        # rotary_pos_emb = RotaryEmbedding(dim = min(32, query_key_dim))
-        self.conv1 = nn.Sequential(
-            nn.Conv1d(dim, inner_channels, kernel_size=1),
-            nn.PReLU(),
-        )
-        self.norm1 = CLayerNorm(inner_channels)
-        #block dilated without gating
-        #self.gated_fsmn = UniDeepFsmn_dilated(inner_channels, inner_channels, 20, inner_channels)
-        #block dilated with gating
-        self.gated_fsmn = Gated_FSMN_dilated(inner_channels, inner_channels, lorder=20, hidden_size=inner_channels)
-        self.norm2 = CLayerNorm(inner_channels)
-        self.conv2 = nn.Conv1d(inner_channels, dim, kernel_size=1)
-
-    def forward(self, input):
-        conv1 = self.conv1(input.transpose(2,1))
-        norm1 = self.norm1(conv1)
-        seq_out = self.gated_fsmn(norm1.transpose(2,1))
-        norm2 = self.norm2(seq_out.transpose(2,1))
-        conv2 = self.conv2(norm2)
-        return conv2.transpose(2,1) + input
-
-class OffsetScale(nn.Module):
-    def __init__(self, dim, heads = 1):
-        super().__init__()
-        self.gamma = nn.Parameter(torch.ones(heads, dim))
-        self.beta = nn.Parameter(torch.zeros(heads, dim))
-        nn.init.normal_(self.gamma, std = 0.02)
-
-    def forward(self, x):
-        out = einsum('... d, h d -> ... h d', x, self.gamma) + self.beta
-        return out.unbind(dim = -2)
-        
-class FLASH_ShareA_FFConvM(nn.Module):
-    def __init__(
-        self,
-        *,
-        dim,
-        group_size = 256,
-        query_key_dim = 128,
-        expansion_factor = 1.,
-        causal = False,
-        dropout = 0.1,
-        rotary_pos_emb = None,
-        norm_klass = nn.LayerNorm,
-        shift_tokens = True
-    ):
-        super().__init__()
-        hidden_dim = int(dim * expansion_factor)        
-        self.group_size = group_size
-        self.causal = causal
-        self.shift_tokens = shift_tokens
-
-        # positional embeddings
-        self.rotary_pos_emb = rotary_pos_emb
-        # norm
-        self.dropout = nn.Dropout(dropout)
-        #self.move = MultiHeadEMA(embed_dim=dim, ndim=4, bidirectional=False, truncation=None)
-        # projections
-        
-        self.to_hidden = FFConvM(
-            dim_in = dim,
-            dim_out = hidden_dim,
-            norm_klass = norm_klass,
-            dropout = dropout,
-            )
-        self.to_qk = FFConvM(
-            dim_in = dim,
-            dim_out = query_key_dim,
-            norm_klass = norm_klass,
-            dropout = dropout,
-            )
-
-        self.qk_offset_scale = OffsetScale(query_key_dim, heads = 4)
-
-        self.to_out = FFConvM(
-            dim_in = dim*2,
-            dim_out = dim,
-            norm_klass = norm_klass,
-            dropout = dropout,
-            )
-        
-        self.gateActivate=nn.Sigmoid() #exp3
-
-    def forward(
-        self,
-        x,
-        *,
-        mask = None
-    ):
-
-        """
-        b - batch
-        n - sequence length (within groups)
-        g - group dimension
-        d - feature dimension (keys)
-        e - feature dimension (values)
-        i - sequence dimension (source)
-        j - sequence dimension (target)
-        """
-
-        #b, n, device, g = x.shape[0], x.shape[-2], x.device, self.group_size
-
-        # prenorm
-        #x = self.fsmn(x)
-        normed_x = x #self.norm(x)
-
-        # do token shift - a great, costless trick from an independent AI researcher in Shenzhen
-        residual = x
-
-        if self.shift_tokens:
-            x_shift, x_pass = normed_x.chunk(2, dim = -1)
-            x_shift = F.pad(x_shift, (0, 0, 1, -1), value = 0.)
-            normed_x = torch.cat((x_shift, x_pass), dim = -1)
-
-        # initial projections
-
-        v, u = self.to_hidden(normed_x).chunk(2, dim = -1)
-        qk = self.to_qk(normed_x)
-        #print('normed_x: {}'.format(normed_x.shape)) 
-
-        # offset and scale
-        quad_q, lin_q, quad_k, lin_k = self.qk_offset_scale(qk)
-        #print('q {}, k {}, v {}'.format(quad_q.shape, quad_k.shape, v.shape))
-        att_v, att_u = self.cal_attention(x, quad_q, lin_q, quad_k, lin_k, v, u)
-
-        #exp5: self.gateActivate=nn.SiLU()
-        out = (att_u*v ) * self.gateActivate(att_v*u)
-        
-        x = x + self.to_out(out)
-        #x = x + self.conv_module(x)
-        return x
-
-    def cal_attention(self, x, quad_q, lin_q, quad_k, lin_k, v, u, mask = None):
-        b, n, device, g = x.shape[0], x.shape[-2], x.device, self.group_size
-
-        if exists(mask):
-            lin_mask = rearrange(mask, '... -> ... 1')
-            lin_k = lin_k.masked_fill(~lin_mask, 0.)
-
-        # rotate queries and keys
-
-        if exists(self.rotary_pos_emb):
-            quad_q, lin_q, quad_k, lin_k = map(self.rotary_pos_emb.rotate_queries_or_keys, (quad_q, lin_q, quad_k, lin_k))
-
-        # padding for groups
-
-        padding = padding_to_multiple_of(n, g)
-
-        if padding > 0:
-            quad_q, quad_k, lin_q, lin_k, v, u = map(lambda t: F.pad(t, (0, 0, 0, padding), value = 0.), (quad_q, quad_k, lin_q, lin_k, v, u))
-
-            mask = default(mask, torch.ones((b, n), device = device, dtype = torch.bool))
-            mask = F.pad(mask, (0, padding), value = False)
-
-        # group along sequence
-
-        quad_q, quad_k, lin_q, lin_k, v, u = map(lambda t: rearrange(t, 'b (g n) d -> b g n d', n = self.group_size), (quad_q, quad_k, lin_q, lin_k, v, u))
-
-        if exists(mask):
-            mask = rearrange(mask, 'b (g j) -> b g 1 j', j = g)
-
-        # calculate quadratic attention output
-
-        sim = einsum('... i d, ... j d -> ... i j', quad_q, quad_k) / g
-
-        ###eddy REMOVE this part can solve infinite loss prob!!!!!!!!!!!!!
-        #sim = sim + self.rel_pos_bias(sim)
-
-        attn = F.relu(sim) ** 2
-        #attn = F.relu(sim)
-        attn = self.dropout(attn)
-
-        if exists(mask):
-            attn = attn.masked_fill(~mask, 0.)
-
-        if self.causal:
-            causal_mask = torch.ones((g, g), dtype = torch.bool, device = device).triu(1)
-            attn = attn.masked_fill(causal_mask, 0.)
-
-        quad_out_v = einsum('... i j, ... j d -> ... i d', attn, v)
-        quad_out_u = einsum('... i j, ... j d -> ... i d', attn, u)
-
-        # calculate linear attention output
-
-        if self.causal:
-            lin_kv = einsum('b g n d, b g n e -> b g d e', lin_k, v) / g
-            # exclusive cumulative sum along group dimension
-            lin_kv = lin_kv.cumsum(dim = 1)
-            lin_kv = F.pad(lin_kv, (0, 0, 0, 0, 1, -1), value = 0.)
-            lin_out_v = einsum('b g d e, b g n d -> b g n e', lin_kv, lin_q)
-
-            lin_ku = einsum('b g n d, b g n e -> b g d e', lin_k, u) / g
-            # exclusive cumulative sum along group dimension
-            lin_ku = lin_ku.cumsum(dim = 1)
-            lin_ku = F.pad(lin_ku, (0, 0, 0, 0, 1, -1), value = 0.)
-            lin_out_u = einsum('b g d e, b g n d -> b g n e', lin_ku, lin_q)
-        else:
-            lin_kv = einsum('b g n d, b g n e -> b d e', lin_k, v) / n
-            lin_out_v = einsum('b g n d, b d e -> b g n e', lin_q, lin_kv)
-
-            lin_ku = einsum('b g n d, b g n e -> b d e', lin_k, u) / n
-            lin_out_u = einsum('b g n d, b d e -> b g n e', lin_q, lin_ku)
-
-        # fold back groups into full sequence, and excise out padding
-        '''
-        quad_attn_out_v, lin_attn_out_v = map(lambda t: rearrange(t, 'b g n d -> b (g n) d')[:, :n], (quad_out_v, lin_out_v))
-        quad_attn_out_u, lin_attn_out_u = map(lambda t: rearrange(t, 'b g n d -> b (g n) d')[:, :n], (quad_out_u, lin_out_u))
-        return quad_attn_out_v+lin_attn_out_v, quad_attn_out_u+lin_attn_out_u
-        '''
-        return map(lambda t: rearrange(t, 'b g n d -> b (g n) d')[:, :n], (quad_out_v+lin_out_v, quad_out_u+lin_out_u))
-
-class FLASHTransformer_DualA_FSMN(nn.Module):
-    def __init__(
-        self,
-        *,
-        dim,
-        depth,
-        group_size = 256, #384, #128, #256,
-        query_key_dim = 128, #256, #128,
-        expansion_factor = 4.,
-        causal = False,
-        attn_dropout = 0.1,
-        norm_type = 'scalenorm',
-        shift_tokens = True
-    ):
-        super().__init__()
-        assert norm_type in ('scalenorm', 'layernorm'), 'norm_type must be one of scalenorm or layernorm'
-
-        if norm_type == 'scalenorm':
-            norm_klass = ScaleNorm
-        elif norm_type == 'layernorm':
-            norm_klass = nn.LayerNorm
-
-        self.group_size = group_size
-
-        rotary_pos_emb = RotaryEmbedding(dim = min(32, query_key_dim))
-        # max rotary embedding dimensions of 32, partial Rotary embeddings, from Wang et al - GPT-J
-        #self.fsmn = nn.ModuleList([Gated_FSMN(dim, dim, lorder=20, hidden_size=dim) for _ in range(depth)])
-        #self.fsmn = nn.ModuleList([Gated_FSMN_Block(dim) for _ in range(depth)])
-        self.fsmn = nn.ModuleList([Gated_FSMN_Block_Dilated(dim) for _ in range(depth)])
-        self.layers = nn.ModuleList([FLASH_ShareA_FFConvM(dim = dim, group_size = group_size, query_key_dim = query_key_dim, expansion_factor = expansion_factor, causal = causal, dropout = attn_dropout, rotary_pos_emb = rotary_pos_emb, norm_klass = norm_klass, shift_tokens = shift_tokens) for _ in range(depth)])
-
-    def _build_repeats(self, in_channels, out_channels, lorder, hidden_size, repeats=1):
-        repeats = [
-            UniDeepFsmn(in_channels, out_channels, lorder, hidden_size)
-            for i in range(repeats)
-        ]
-        return nn.Sequential(*repeats)
-
-    def forward(
-        self,
-        x,
-        *,
-        mask = None
-    ):
-        ii = 0
-        for flash in self.layers:
-            #x_residual = x
-            x = flash(x, mask = mask)
-            x = self.fsmn[ii](x)
-            #x = x + x_residual
-            ii = ii + 1
-        return x
-
-class TransformerEncoder_FLASH_DualA_FSMN(nn.Module):
-    """This class implements the transformer encoder.
-
-    Arguments
-    ---------
-    num_layers : int
-        Number of transformer layers to include.
-    nhead : int
-        Number of attention heads.
-    d_ffn : int
-        Hidden size of self-attention Feed Forward layer.
-    d_model : int
-        The dimension of the input embedding.
-    kdim : int
-        Dimension for key (Optional).
-    vdim : int
-        Dimension for value (Optional).
-    dropout : float
-        Dropout for the encoder (Optional).
-    input_module: torch class
-        The module to process the source input feature to expected
-        feature dimension (Optional).
-
-    Example
-    -------
-    >>> import torch
-    >>> x = torch.rand((8, 60, 512))
-    >>> net = TransformerEncoder(1, 8, 512, d_model=512)
-    >>> output, _ = net(x)
-    >>> output.shape
-    torch.Size([8, 60, 512])
-    """
-    def __init__(
-        self,
-        num_layers,
-        nhead,
-        d_ffn,
-        input_shape=None,
-        d_model=None,
-        kdim=None,
-        vdim=None,
-        dropout=0.0,
-        activation=nn.ReLU,
-        normalize_before=False,
-        causal=False,
-        attention_type="regularMHA",
-    ):
-
-        super().__init__()
-
-        self.flashT = FLASHTransformer_DualA_FSMN(dim=d_model, depth=num_layers)
-        self.norm = LayerNorm(d_model, eps=1e-6)
-
-    def forward(
-        self,
-        src,
-        src_mask: Optional[torch.Tensor] = None,
-        src_key_padding_mask: Optional[torch.Tensor] = None,
-        pos_embs: Optional[torch.Tensor] = None,
-    ):
-        """
-        Arguments
-        ----------
-        src : tensor
-            The sequence to the encoder layer (required).
-        src_mask : tensor
-            The mask for the src sequence (optional).
-        src_key_padding_mask : tensor
-            The mask for the src keys per batch (optional).
-        """
-        output = self.flashT(src)
-        #summary(self.flashT, [(src.size())])
-        output = self.norm(output)
-        #summary(self.norm, [(output.size())])
-
-        return output

models/mossformer2/mossformer/utils/conv_module.py

@@ -1,87 +0,0 @@
-import torch
-import torch.nn as nn
-from torch import Tensor
-import torch.nn.init as init
-import torch.nn.functional as F
-
-class Transpose(nn.Module):
-    """ Wrapper class of torch.transpose() for Sequential module. """
-    def __init__(self, shape: tuple):
-        super(Transpose, self).__init__()
-        self.shape = shape
-
-    def forward(self, x: Tensor) -> Tensor:
-        return x.transpose(*self.shape)
-
-class DepthwiseConv1d(nn.Module):
-    """
-    When groups == in_channels and out_channels == K * in_channels, where K is a positive integer,
-    this operation is termed in literature as depthwise convolution.
-    Args:
-        in_channels (int): Number of channels in the input
-        out_channels (int): Number of channels produced by the convolution
-        kernel_size (int or tuple): Size of the convolving kernel
-        stride (int, optional): Stride of the convolution. Default: 1
-        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
-        bias (bool, optional): If True, adds a learnable bias to the output. Default: True
-    Inputs: inputs
-        - **inputs** (batch, in_channels, time): Tensor containing input vector
-    Returns: outputs
-        - **outputs** (batch, out_channels, time): Tensor produces by depthwise 1-D convolution.
-    """
-    def __init__(
-            self,
-            in_channels: int,
-            out_channels: int,
-            kernel_size: int,
-            stride: int = 1,
-            padding: int = 0,
-            bias: bool = False,
-    ) -> None:
-        super(DepthwiseConv1d, self).__init__()
-        assert out_channels % in_channels == 0, "out_channels should be constant multiple of in_channels"
-        self.conv = nn.Conv1d(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            kernel_size=kernel_size,
-            groups=in_channels,
-            stride=stride,
-            padding=padding,
-            bias=bias,
-        )
-
-    def forward(self, inputs: Tensor) -> Tensor:
-        return self.conv(inputs)
-
-class ConvModule(nn.Module):
-    """
-    Conformer convolution module starts with a pointwise convolution and a gated linear unit (GLU).
-    This is followed by a single 1-D depthwise convolution layer. Batchnorm is  deployed just after the convolution
-    to aid training deep models.
-    Args:
-        in_channels (int): Number of channels in the input
-        kernel_size (int or tuple, optional): Size of the convolving kernel Default: 31
-        dropout_p (float, optional): probability of dropout
-    Inputs: inputs
-        inputs (batch, time, dim): Tensor contains input sequences
-    Outputs: outputs
-        outputs (batch, time, dim): Tensor produces by conformer convolution module.
-    """
-    def __init__(
-            self,
-            in_channels: int,
-            kernel_size: int = 17,
-            expansion_factor: int = 2,
-            dropout_p: float = 0.1,
-    ) -> None:
-        super(ConvModule, self).__init__()
-        assert (kernel_size - 1) % 2 == 0, "kernel_size should be a odd number for 'SAME' padding"
-        assert expansion_factor == 2, "Currently, Only Supports expansion_factor 2"
-
-        self.sequential = nn.Sequential(
-            Transpose(shape=(1, 2)),
-            DepthwiseConv1d(in_channels, in_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2),
-        )
-
-    def forward(self, inputs: Tensor) -> Tensor:
-        return inputs + self.sequential(inputs).transpose(1, 2)

models/mossformer2/mossformer/utils/fsmn.py

@@ -1,108 +0,0 @@
-import torch.nn as nn
-import torch.nn.functional as F
-import torch as th
-from torch.nn.parameter import Parameter
-import numpy as np
-import os
-
-class UniDeepFsmn(nn.Module):
-
-    def __init__(self, input_dim, output_dim, lorder=None, hidden_size=None):
-        super(UniDeepFsmn, self).__init__()
-
-        self.input_dim = input_dim
-        self.output_dim = output_dim
-
-        if lorder is None:
-            return
-
-        self.lorder = lorder
-        self.hidden_size = hidden_size
-
-        self.linear = nn.Linear(input_dim, hidden_size)
-
-        self.project = nn.Linear(hidden_size, output_dim, bias=False)
-
-        self.conv1 = nn.Conv2d(output_dim, output_dim, [lorder+lorder-1, 1], [1, 1], groups=output_dim, bias=False)
-
-    def forward(self, input):
-
-        f1 = F.relu(self.linear(input))
-
-        p1 = self.project(f1)
-
-        x = th.unsqueeze(p1, 1)
-
-        x_per = x.permute(0, 3, 2, 1)
-
-        y = F.pad(x_per, [0, 0, self.lorder - 1, self.lorder - 1])
-
-        out = x_per + self.conv1(y)
-
-        out1 = out.permute(0, 3, 2, 1)
-
-        return input + out1.squeeze()
-
-class DilatedDenseNet(nn.Module):
-    def __init__(self, depth=4, lorder=20, in_channels=64):
-        super(DilatedDenseNet, self).__init__()
-        self.depth = depth
-        self.in_channels = in_channels
-        self.pad = nn.ConstantPad2d((1, 1, 1, 0), value=0.)
-        self.twidth = lorder*2-1
-        self.kernel_size = (self.twidth, 1)
-        for i in range(self.depth):
-            dil = 2 ** i
-            pad_length = lorder + (dil - 1) * (lorder - 1) - 1
-            setattr(self, 'pad{}'.format(i + 1), nn.ConstantPad2d((0, 0, pad_length, pad_length), value=0.))
-            setattr(self, 'conv{}'.format(i + 1),
-                    nn.Conv2d(self.in_channels*(i+1), self.in_channels, kernel_size=self.kernel_size,
-                              dilation=(dil, 1), groups=self.in_channels, bias=False))
-            setattr(self, 'norm{}'.format(i + 1), nn.InstanceNorm2d(in_channels, affine=True))
-            setattr(self, 'prelu{}'.format(i + 1), nn.PReLU(self.in_channels))
-
-    def forward(self, x):
-        skip = x
-        for i in range(self.depth):
-            out = getattr(self, 'pad{}'.format(i + 1))(skip)
-            out = getattr(self, 'conv{}'.format(i + 1))(out)
-            out = getattr(self, 'norm{}'.format(i + 1))(out)
-            out = getattr(self, 'prelu{}'.format(i + 1))(out)            
-            skip = th.cat([out, skip], dim=1)
-        return out
-
-class UniDeepFsmn_dilated(nn.Module):
-
-    def __init__(self, input_dim, output_dim, lorder=None, hidden_size=None):
-        super(UniDeepFsmn_dilated, self).__init__()
-
-        self.input_dim = input_dim
-        self.output_dim = output_dim
-
-        if lorder is None:
-            return
-
-        self.lorder = lorder
-        self.hidden_size = hidden_size
-
-        self.linear = nn.Linear(input_dim, hidden_size)
-
-        self.project = nn.Linear(hidden_size, output_dim, bias=False)
-
-        self.conv = DilatedDenseNet(depth=2, lorder=lorder, in_channels=output_dim)
-
-    def forward(self, input):
-
-        f1 = F.relu(self.linear(input))
-
-        p1 = self.project(f1)
-
-        x = th.unsqueeze(p1, 1)
-
-        x_per = x.permute(0, 3, 2, 1)
-
-        out = self.conv(x_per)
-
-        out1 = out.permute(0, 3, 2, 1)
-
-        return input + out1.squeeze()

models/mossformer2/mossformer/utils/normalization.py

@@ -1,94 +0,0 @@
-import torch
-import torch.nn as nn
-
-class LayerNorm(nn.Module):
-    """
-    This code came from sb.nnet.normalization
-    # from sb.nnet.normalization import LayerNorm
-
-
-    Applies layer normalization to the input tensor.
-
-    Arguments
-    ---------
-    input_shape : tuple
-        The expected shape of the input.
-    eps : float
-        This value is added to std deviation estimation to improve the numerical
-        stability.
-    elementwise_affine : bool
-        If True, this module has learnable per-element affine parameters
-        initialized to ones (for weights) and zeros (for biases).
-
-    Example
-    -------
-    >>> input = torch.randn(100, 101, 128)
-    >>> norm = LayerNorm(input_shape=input.shape)
-    >>> output = norm(input)
-    >>> output.shape
-    torch.Size([100, 101, 128])
-    """
-
-    def __init__(
-        self,
-        input_size=None,
-        input_shape=None,
-        eps=1e-05,
-        elementwise_affine=True,
-    ):
-        super().__init__()
-        self.eps = eps
-        self.elementwise_affine = elementwise_affine
-
-        if input_shape is not None:
-            input_size = input_shape[2:]
-
-        self.norm = torch.nn.LayerNorm(
-            input_size,
-            eps=self.eps,
-            elementwise_affine=self.elementwise_affine,
-        )
-
-    def forward(self, x):
-        """Returns the normalized input tensor.
-
-        Arguments
-        ---------
-        x : torch.Tensor (batch, time, channels)
-            input to normalize. 3d or 4d tensors are expected.
-        """
-        return self.norm(x)
-
-class CLayerNorm(nn.LayerNorm):
-    """Channel-wise layer normalization."""
-
-    def __init__(self, *args, **kwargs):
-        super(CLayerNorm, self).__init__(*args, **kwargs)
-
-    def forward(self, sample):
-        """Forward function.
-
-        Args:
-            sample: [batch_size, channels, length]
-        """
-        if sample.dim() != 3:
-            raise RuntimeError('{} only accept 3-D tensor as input'.format(
-                self.__name__))
-        # [N, C, T] -> [N, T, C]
-        sample = torch.transpose(sample, 1, 2)
-        # LayerNorm
-        sample = super().forward(sample)
-        # [N, T, C] -> [N, C, T]
-        sample = torch.transpose(sample, 1, 2)
-        return sample
-
-class ScaleNorm(nn.Module):
-    def __init__(self, dim, eps = 1e-5):
-        super().__init__()
-        self.scale = dim ** -0.5
-        self.eps = eps
-        self.g = nn.Parameter(torch.ones(1))
-
-    def forward(self, x):
-        norm = torch.norm(x, dim = -1, keepdim = True) * self.scale
-        return x / norm.clamp(min = self.eps) * self.g

models/mossformer2/mossformer/utils/one_path_flash_fsmn.py

@@ -1,800 +0,0 @@
-import copy
-import math
-import torch
-
-import torch.nn as nn
-import torch.nn.functional as F
-
-from torch import einsum
-from ..utils.Transformer import TransformerEncoder_FLASH_DualA_FSMN
-
-EPS = 1e-8
-
-class ScaledSinuEmbedding(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-        self.scale = nn.Parameter(torch.ones(1,))
-        inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
-        self.register_buffer('inv_freq', inv_freq)
-
-    def forward(self, x):
-        n, device = x.shape[1], x.device
-        t = torch.arange(n, device = device).type_as(self.inv_freq)
-        sinu = einsum('i , j -> i j', t, self.inv_freq)
-        emb = torch.cat((sinu.sin(), sinu.cos()), dim = -1)
-        return emb * self.scale
-
-class Linear(torch.nn.Module):
-    """Computes a linear transformation y = wx + b.
-
-    Arguments
-    ---------
-    n_neurons : int
-        It is the number of output neurons (i.e, the dimensionality of the
-        output).
-    input_shape: tuple
-        It is the shape of the input tensor.
-    input_size: int
-        Size of the input tensor.
-    bias : bool
-        If True, the additive bias b is adopted.
-    combine_dims : bool
-        If True and the input is 4D, combine 3rd and 4th dimensions of input.
-
-    Example
-    -------
-    >>> inputs = torch.rand(10, 50, 40)
-    >>> lin_t = Linear(input_shape=(10, 50, 40), n_neurons=100)
-    >>> output = lin_t(inputs)
-    >>> output.shape
-    torch.Size([10, 50, 100])
-    """
-
-    def __init__(
-        self,
-        n_neurons,
-        input_shape=None,
-        input_size=None,
-        bias=True,
-        combine_dims=False,
-    ):
-        super().__init__()
-        self.combine_dims = combine_dims
-
-        if input_shape is None and input_size is None:
-            raise ValueError("Expected one of input_shape or input_size")
-
-        if input_size is None:
-            input_size = input_shape[-1]
-            if len(input_shape) == 4 and self.combine_dims:
-                input_size = input_shape[2] * input_shape[3]
-
-        # Weights are initialized following pytorch approach
-        self.w = nn.Linear(input_size, n_neurons, bias=bias)
-
-    def forward(self, x):
-        """Returns the linear transformation of input tensor.
-
-        Arguments
-        ---------
-        x : torch.Tensor
-            Input to transform linearly.
-        """
-        if x.ndim == 4 and self.combine_dims:
-            x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3])
-
-        wx = self.w(x)
-
-        return wx
-
-class GlobalLayerNorm(nn.Module):
-    """Calculate Global Layer Normalization.
-
-    Arguments
-    ---------
-       dim : (int or list or torch.Size)
-           Input shape from an expected input of size.
-       eps : float
-           A value added to the denominator for numerical stability.
-       elementwise_affine : bool
-          A boolean value that when set to True,
-          this module has learnable per-element affine parameters
-          initialized to ones (for weights) and zeros (for biases).
-
-    Example
-    -------
-    >>> x = torch.randn(5, 10, 20)
-    >>> GLN = GlobalLayerNorm(10, 3)
-    >>> x_norm = GLN(x)
-    """
-
-    def __init__(self, dim, shape, eps=1e-8, elementwise_affine=True):
-        super(GlobalLayerNorm, self).__init__()
-        self.dim = dim
-        self.eps = eps
-        self.elementwise_affine = elementwise_affine
-
-        if self.elementwise_affine:
-            if shape == 3:
-                self.weight = nn.Parameter(torch.ones(self.dim, 1))
-                self.bias = nn.Parameter(torch.zeros(self.dim, 1))
-            if shape == 4:
-                self.weight = nn.Parameter(torch.ones(self.dim, 1, 1))
-                self.bias = nn.Parameter(torch.zeros(self.dim, 1, 1))
-        else:
-            self.register_parameter("weight", None)
-            self.register_parameter("bias", None)
-
-    def forward(self, x):
-        """Returns the normalized tensor.
-
-        Arguments
-        ---------
-        x : torch.Tensor
-            Tensor of size [N, C, K, S] or [N, C, L].
-        """
-        # x = N x C x K x S or N x C x L
-        # N x 1 x 1
-        # cln: mean,var N x 1 x K x S
-        # gln: mean,var N x 1 x 1
-        if x.dim() == 3:
-            mean = torch.mean(x, (1, 2), keepdim=True)
-            var = torch.mean((x - mean) ** 2, (1, 2), keepdim=True)
-            if self.elementwise_affine:
-                x = (
-                    self.weight * (x - mean) / torch.sqrt(var + self.eps)
-                    + self.bias
-                )
-            else:
-                x = (x - mean) / torch.sqrt(var + self.eps)
-
-        if x.dim() == 4:
-            mean = torch.mean(x, (1, 2, 3), keepdim=True)
-            var = torch.mean((x - mean) ** 2, (1, 2, 3), keepdim=True)
-            if self.elementwise_affine:
-                x = (
-                    self.weight * (x - mean) / torch.sqrt(var + self.eps)
-                    + self.bias
-                )
-            else:
-                x = (x - mean) / torch.sqrt(var + self.eps)
-        return x
-
-
-class CumulativeLayerNorm(nn.LayerNorm):
-    """Calculate Cumulative Layer Normalization.
-
-       Arguments
-       ---------
-       dim : int
-        Dimension that you want to normalize.
-       elementwise_affine : True
-        Learnable per-element affine parameters.
-
-    Example
-    -------
-    >>> x = torch.randn(5, 10, 20)
-    >>> CLN = CumulativeLayerNorm(10)
-    >>> x_norm = CLN(x)
-    """
-
-    def __init__(self, dim, elementwise_affine=True):
-        super(CumulativeLayerNorm, self).__init__(
-            dim, elementwise_affine=elementwise_affine, eps=1e-8
-        )
-
-    def forward(self, x):
-        """Returns the normalized tensor.
-
-        Arguments
-        ---------
-        x : torch.Tensor
-            Tensor size [N, C, K, S] or [N, C, L]
-        """
-        # x: N x C x K x S or N x C x L
-        # N x K x S x C
-        if x.dim() == 4:
-            x = x.permute(0, 2, 3, 1).contiguous()
-            # N x K x S x C == only channel norm
-            x = super().forward(x)
-            # N x C x K x S
-            x = x.permute(0, 3, 1, 2).contiguous()
-        if x.dim() == 3:
-            x = torch.transpose(x, 1, 2)
-            # N x L x C == only channel norm
-            x = super().forward(x)
-            # N x C x L
-            x = torch.transpose(x, 1, 2)
-        return x
-
-
-def select_norm(norm, dim, shape):
-    """Just a wrapper to select the normalization type.
-    """
-
-    if norm == "gln":
-        return GlobalLayerNorm(dim, shape, elementwise_affine=True)
-    if norm == "cln":
-        return CumulativeLayerNorm(dim, elementwise_affine=True)
-    if norm == "ln":
-        return nn.GroupNorm(1, dim, eps=1e-8)
-    else:
-        return nn.BatchNorm1d(dim)
-
-class Encoder(nn.Module):
-    """Convolutional Encoder Layer.
-
-    Arguments
-    ---------
-    kernel_size : int
-        Length of filters.
-    in_channels : int
-        Number of  input channels.
-    out_channels : int
-        Number of output channels.
-
-    Example
-    -------
-    >>> x = torch.randn(2, 1000)
-    >>> encoder = Encoder(kernel_size=4, out_channels=64)
-    >>> h = encoder(x)
-    >>> h.shape
-    torch.Size([2, 64, 499])
-    """
-
-    def __init__(self, kernel_size=2, out_channels=64, in_channels=1):
-        super(Encoder, self).__init__()
-        self.conv1d = nn.Conv1d(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            kernel_size=kernel_size,
-            stride=kernel_size // 2,
-            groups=1,
-            bias=False,
-        )
-        self.in_channels = in_channels
-
-    def forward(self, x):
-        """Return the encoded output.
-
-        Arguments
-        ---------
-        x : torch.Tensor
-            Input tensor with dimensionality [B, L].
-        Return
-        ------
-        x : torch.Tensor
-            Encoded tensor with dimensionality [B, N, T_out].
-
-        where B = Batchsize
-              L = Number of timepoints
-              N = Number of filters
-              T_out = Number of timepoints at the output of the encoder
-        """
-        # B x L -> B x 1 x L
-        if self.in_channels == 1:
-            x = torch.unsqueeze(x, dim=1)
-        # B x 1 x L -> B x N x T_out
-        x = self.conv1d(x)
-        x = F.relu(x)
-
-        return x
-        
-class Decoder(nn.ConvTranspose1d):
-    """A decoder layer that consists of ConvTranspose1d.
-
-    Arguments
-    ---------
-    kernel_size : int
-        Length of filters.
-    in_channels : int
-        Number of  input channels.
-    out_channels : int
-        Number of output channels.
-
-
-    Example
-    ---------
-    >>> x = torch.randn(2, 100, 1000)
-    >>> decoder = Decoder(kernel_size=4, in_channels=100, out_channels=1)
-    >>> h = decoder(x)
-    >>> h.shape
-    torch.Size([2, 1003])
-    """
-
-    def __init__(self, *args, **kwargs):
-        super(Decoder, self).__init__(*args, **kwargs)
-
-    def forward(self, x):
-        """Return the decoded output.
-
-        Arguments
-        ---------
-        x : torch.Tensor
-            Input tensor with dimensionality [B, N, L].
-                where, B = Batchsize,
-                       N = number of filters
-                       L = time points
-        """
-
-        if x.dim() not in [2, 3]:
-            raise RuntimeError(
-                "{} accept 3/4D tensor as input".format(self.__name__)
-            )
-        x = super().forward(x if x.dim() == 3 else torch.unsqueeze(x, 1))
-
-        if torch.squeeze(x).dim() == 1:
-            x = torch.squeeze(x, dim=1)
-        else:
-            x = torch.squeeze(x)
-        return x
-
-class SBFLASHBlock_DualA(nn.Module):
-    """A wrapper for the SpeechBrain implementation of the transformer encoder.
-
-    Arguments
-    ---------
-    num_layers : int
-        Number of layers.
-    d_model : int
-        Dimensionality of the representation.
-    nhead : int
-        Number of attention heads.
-    d_ffn : int
-        Dimensionality of positional feed forward.
-    input_shape : tuple
-        Shape of input.
-    kdim : int
-        Dimension of the key (Optional).
-    vdim : int
-        Dimension of the value (Optional).
-    dropout : float
-        Dropout rate.
-    activation : str
-        Activation function.
-    use_positional_encoding : bool
-        If true we use a positional encoding.
-    norm_before: bool
-        Use normalization before transformations.
-
-    Example
-    ---------
-    >>> x = torch.randn(10, 100, 64)
-    >>> block = SBTransformerBlock(1, 64, 8)
-    >>> x = block(x)
-    >>> x.shape
-    torch.Size([10, 100, 64])
-    """
-
-    def __init__(
-        self,
-        num_layers,
-        d_model,
-        nhead,
-        d_ffn=2048,
-        input_shape=None,
-        kdim=None,
-        vdim=None,
-        dropout=0.1,
-        activation="relu",
-        use_positional_encoding=False,
-        norm_before=False,
-        attention_type="regularMHA",
-    ):
-
-        super(SBFLASHBlock_DualA, self).__init__()
-        self.use_positional_encoding = use_positional_encoding
-
-        if activation == "relu":
-            activation = nn.ReLU
-        elif activation == "gelu":
-            activation = nn.GELU
-        else:
-            raise ValueError("unknown activation")
-
-
-        self.mdl = TransformerEncoder_FLASH_DualA_FSMN(
-            num_layers=num_layers,
-            nhead=nhead,
-            d_ffn=d_ffn,
-            input_shape=input_shape,
-            d_model=d_model,
-            kdim=kdim,
-            vdim=vdim,
-            dropout=dropout,
-            activation=activation,
-            normalize_before=norm_before,
-            attention_type=attention_type,
-        )
-
-    def forward(self, x):
-        """Returns the transformed output.
-
-        Arguments
-        ---------
-        x : torch.Tensor
-            Tensor shape [B, L, N],
-            where, B = Batchsize,
-                   L = time points
-                   N = number of filters
-
-        """
-        output = self.mdl(x)
-
-        return output
-
-
-def _get_activation_fn(activation):
-    """Just a wrapper to get the activation functions.
-    """
-
-    if activation == "relu":
-        return F.relu
-    elif activation == "gelu":
-        return F.gelu
-
-
-class Dual_Computation_Block(nn.Module):
-    """Computation block for dual-path processing.
-
-    Arguments
-    ---------
-    intra_mdl : torch.nn.module
-        Model to process within the chunks.
-     inter_mdl : torch.nn.module
-        Model to process across the chunks.
-     out_channels : int
-        Dimensionality of inter/intra model.
-     norm : str
-        Normalization type.
-     skip_around_intra : bool
-        Skip connection around the intra layer.
-     linear_layer_after_inter_intra : bool
-        Linear layer or not after inter or intra.
-
-    Example
-    ---------
-        >>> intra_block = SBTransformerBlock(1, 64, 8)
-        >>> inter_block = SBTransformerBlock(1, 64, 8)
-        >>> dual_comp_block = Dual_Computation_Block(intra_block, inter_block, 64)
-        >>> x = torch.randn(10, 64, 100, 10)
-        >>> x = dual_comp_block(x)
-        >>> x.shape
-        torch.Size([10, 64, 100, 10])
-    """
-
-    def __init__(
-        self,
-        intra_mdl,
-        out_channels,
-        norm="ln",
-        skip_around_intra=True,
-        linear_layer_after_inter_intra=True,
-    ):
-        super(Dual_Computation_Block, self).__init__()
-
-        self.intra_mdl = intra_mdl
-        self.skip_around_intra = skip_around_intra
-        self.linear_layer_after_inter_intra = linear_layer_after_inter_intra
-
-        # Norm
-        self.norm = norm
-        if norm is not None:
-            self.intra_norm = select_norm(norm, out_channels, 3)
-
-        # Linear
-        if linear_layer_after_inter_intra:
-            self.intra_linear = Linear(
-                    out_channels, input_size=out_channels
-            )
-
-    def forward(self, x):
-        """Returns the output tensor.
-
-        Arguments
-        ---------
-        x : torch.Tensor
-            Input tensor of dimension [B, N, K, S].
-
-
-        Return
-        ---------
-        out: torch.Tensor
-            Output tensor of dimension [B, N, K, S].
-            where, B = Batchsize,
-               N = number of filters
-               K = time points in each chunk
-               S = the number of chunks
-        """
-        B, N, S = x.shape
-        # intra RNN
-        # [B, S, N]
-        intra = x.permute(0, 2, 1).contiguous() #.view(B, S, N)
-
-        intra = self.intra_mdl(intra)
-
-        # [B, S, N]
-        if self.linear_layer_after_inter_intra:
-            intra = self.intra_linear(intra)
-
-        # [B, N, S]
-        intra = intra.permute(0, 2, 1).contiguous()
-        if self.norm is not None:
-            intra = self.intra_norm(intra)
-
-        # [B, N, S]
-        if self.skip_around_intra:
-            intra = intra + x
-
-        # inter RNN
-        # [B, S, N]
-        '''
-        inter = intra.permute(0, 2, 1).contiguous() #.view(B, S, N)
-        # [BK, S, H]
-        inter = self.inter_mdl(inter)
-
-        # [BK, S, N]
-        if self.linear_layer_after_inter_intra:
-            inter = self.inter_linear(inter)
-
-        # [B, N, S]
-        inter = inter.permute(0, 2, 1).contiguous()
-        if self.norm is not None:
-            inter = self.inter_norm(inter)
-        # [B, N, K, S]
-        out = inter + intra
-        '''
-        out = intra
-        return out
-
-
-class Dual_Path_Model(nn.Module):
-    """The dual path model which is the basis for dualpathrnn, sepformer, dptnet.
-
-    Arguments
-    ---------
-    in_channels : int
-        Number of channels at the output of the encoder.
-    out_channels : int
-        Number of channels that would be inputted to the intra and inter blocks.
-    intra_model : torch.nn.module
-        Model to process within the chunks.
-    inter_model : torch.nn.module
-        model to process across the chunks,
-    num_layers : int
-        Number of layers of Dual Computation Block.
-    norm : str
-        Normalization type.
-    K : int
-        Chunk length.
-    num_spks : int
-        Number of sources (speakers).
-    skip_around_intra : bool
-        Skip connection around intra.
-    linear_layer_after_inter_intra : bool
-        Linear layer after inter and intra.
-    use_global_pos_enc : bool
-        Global positional encodings.
-    max_length : int
-        Maximum sequence length.
-
-    Example
-    ---------
-    >>> intra_block = SBTransformerBlock(1, 64, 8)
-    >>> inter_block = SBTransformerBlock(1, 64, 8)
-    >>> dual_path_model = Dual_Path_Model(64, 64, intra_block, inter_block, num_spks=2)
-    >>> x = torch.randn(10, 64, 2000)
-    >>> x = dual_path_model(x)
-    >>> x.shape
-    torch.Size([2, 10, 64, 2000])
-    """
-
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        intra_model,
-        #inter_model,
-        num_layers=1,
-        norm="ln",
-        K=200,
-        num_spks=2,
-        skip_around_intra=True,
-        linear_layer_after_inter_intra=True,
-        use_global_pos_enc=True,
-        max_length=20000,
-    ):
-        super(Dual_Path_Model, self).__init__()
-        self.K = K
-        self.num_spks = num_spks
-        self.num_layers = num_layers
-        # self.norm = select_norm(norm, in_channels, 3)
-        # self.conv1d_encoder = nn.Conv1d(in_channels, out_channels, 1, bias=False)
-        self.use_global_pos_enc = use_global_pos_enc
-
-        if self.use_global_pos_enc:
-            self.pos_enc = ScaledSinuEmbedding(out_channels)
-
-        self.dual_mdl = nn.ModuleList([])
-        for i in range(num_layers):
-            self.dual_mdl.append(
-                copy.deepcopy(
-                    Dual_Computation_Block(
-                        intra_model,
-                        #inter_model,
-                        out_channels,
-                        norm,
-                        skip_around_intra=skip_around_intra,
-                        linear_layer_after_inter_intra=linear_layer_after_inter_intra,
-                    )
-                )
-            )
-
-        self.conv1d_out = nn.Conv1d(
-            out_channels, out_channels * num_spks, kernel_size=1
-        )
-        self.conv1_decoder = nn.Conv1d(out_channels, in_channels, 1, bias=False)
-        self.prelu = nn.PReLU()
-        self.activation = nn.ReLU()
-        # gated output layer
-        self.output = nn.Sequential(
-            nn.Conv1d(out_channels, out_channels, 1), nn.Tanh()
-        )
-        self.output_gate = nn.Sequential(
-            nn.Conv1d(out_channels, out_channels, 1), nn.Sigmoid()
-        )
-
-    def forward(self, x):
-        """Returns the output tensor.
-
-        Arguments
-        ---------
-        x : torch.Tensor
-            Input tensor of dimension [B, N, L].
-
-        Returns
-        -------
-        out : torch.Tensor
-            Output tensor of dimension [spks, B, N, L]
-            where, spks = Number of speakers
-               B = Batchsize,
-               N = number of filters
-               L = the number of time points
-        """
-
-        # before each line we indicate the shape after executing the line
-
-        # # [B, N, L]
-        # x = self.norm(x)
-
-        # # [B, N, L]
-        # x = self.conv1d_encoder(x)
-         
-        if self.use_global_pos_enc:
-            base = x
-            x = x.transpose(1, -1)
-            emb = self.pos_enc(x)
-            emb = emb.transpose(0, -1) 
-            x = base + emb
-        
-        # [B, N, S]
-        for i in range(self.num_layers):
-            x = self.dual_mdl[i](x)
-        x = self.prelu(x)
-
-        # [B, N*spks, K, S]
-        x = self.conv1d_out(x)
-        B, _, S = x.shape
-
-        # [B*spks, N, K, S]
-        x = x.view(B * self.num_spks, -1, S)
-
-        # [B*spks, N, L]
-        x = self.output(x) * self.output_gate(x)
-
-        # [B*spks, N, L]
-        x = self.conv1_decoder(x)
-
-        # [B, spks, N, L]
-        _, N, L = x.shape
-        x = x.view(B, self.num_spks, N, L)
-        x = self.activation(x)
-
-        # [spks, B, N, L]
-        x = x.transpose(0, 1)
-
-        return x
-
-    def _padding(self, input, K):
-        """Padding the audio times.
-
-        Arguments
-        ---------
-        K : int
-            Chunks of length.
-        P : int
-            Hop size.
-        input : torch.Tensor
-            Tensor of size [B, N, L].
-            where, B = Batchsize,
-                   N = number of filters
-                   L = time points
-        """
-        B, N, L = input.shape
-        P = K // 2
-        gap = K - (P + L % K) % K
-        if gap > 0:
-            pad = torch.Tensor(torch.zeros(B, N, gap)).type(input.type())
-            input = torch.cat([input, pad], dim=2)
-
-        _pad = torch.Tensor(torch.zeros(B, N, P)).type(input.type())
-        input = torch.cat([_pad, input, _pad], dim=2)
-
-        return input, gap
-
-    def _Segmentation(self, input, K):
-        """The segmentation stage splits
-
-        Arguments
-        ---------
-        K : int
-            Length of the chunks.
-        input : torch.Tensor
-            Tensor with dim [B, N, L].
-
-        Return
-        -------
-        output : torch.tensor
-            Tensor with dim [B, N, K, S].
-            where, B = Batchsize,
-               N = number of filters
-               K = time points in each chunk
-               S = the number of chunks
-               L = the number of time points
-        """
-        B, N, L = input.shape
-        P = K // 2
-        input, gap = self._padding(input, K)
-        # [B, N, K, S]
-        input1 = input[:, :, :-P].contiguous().view(B, N, -1, K)
-        input2 = input[:, :, P:].contiguous().view(B, N, -1, K)
-        input = (
-            torch.cat([input1, input2], dim=3).view(B, N, -1, K).transpose(2, 3)
-        )
-
-        return input.contiguous(), gap
-
-    def _over_add(self, input, gap):
-        """Merge the sequence with the overlap-and-add method.
-
-        Arguments
-        ---------
-        input : torch.tensor
-            Tensor with dim [B, N, K, S].
-        gap : int
-            Padding length.
-
-        Return
-        -------
-        output : torch.tensor
-            Tensor with dim [B, N, L].
-            where, B = Batchsize,
-               N = number of filters
-               K = time points in each chunk
-               S = the number of chunks
-               L = the number of time points
-
-        """
-        B, N, K, S = input.shape
-        P = K // 2
-        # [B, N, S, K]
-        input = input.transpose(2, 3).contiguous().view(B, N, -1, K * 2)
-
-        input1 = input[:, :, :, :K].contiguous().view(B, N, -1)[:, :, P:]
-        input2 = input[:, :, :, K:].contiguous().view(B, N, -1)[:, :, :-P]
-        input = input1 + input2
-        # [B, N, L]
-        if gap > 0:
-            input = input[:, :, :-gap]
-
-        return input

models/mossformer2/mossformer2.py

@@ -1,216 +0,0 @@
-import os
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torchaudio
-
-import math
-
-from .mossformer.utils.one_path_flash_fsmn import Dual_Path_Model, SBFLASHBlock_DualA
-from torch.nn import TransformerEncoder, TransformerEncoderLayer
-
-EPS = 1e-8
-
-class Mossformer(nn.Module):
-    def __init__(self, args):
-        super(Mossformer, self).__init__()
-        
-        N, L, = args.network_audio.encoder_out_nchannels, args.network_audio.encoder_kernel_size
-
-        self.encoder = Encoder(L, N)
-        self.separator = Separator(args)
-        self.decoder = Decoder(args, N, L)
-
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_normal_(p)
-
-    def forward(self, mixture, visual):
-        """
-        Args:
-            mixture: [M, T], M is batch size, T is #samples
-        Returns:
-            est_source: [M, C, T]
-        """
-        mixture_w = self.encoder(mixture)
-        est_mask = self.separator(mixture_w, visual)
-        est_source = self.decoder(mixture_w, est_mask)
-
-        # T changed after conv1d in encoder, fix it here
-        T_origin = mixture.size(-1)
-        T_conv = est_source.size(-1)
-        est_source = F.pad(est_source, (0, T_origin - T_conv))
-        return est_source
-
-class Encoder(nn.Module):
-    def __init__(self, L, N):
-        super(Encoder, self).__init__()
-        self.L, self.N = L, N
-        self.conv1d_U = nn.Conv1d(1, N, kernel_size=L, stride=L // 2, bias=False)
-
-    def forward(self, mixture):
-        """
-        Args:
-            mixture: [M, T], M is batch size, T is #samples
-        Returns:
-            mixture_w: [M, N, K], where K = (T-L)/(L/2)+1 = 2T/L-1
-        """
-        mixture = torch.unsqueeze(mixture, 1)  # [M, 1, T]
-        mixture_w = F.relu(self.conv1d_U(mixture))  # [M, N, K]
-        return mixture_w
-
-
-class Decoder(nn.Module):
-    def __init__(self, args, N, L):
-        super(Decoder, self).__init__()
-        self.N, self.L, self.args = N, L, args
-        self.basis_signals = nn.Linear(N, L, bias=False)
-
-    def forward(self, mixture_w, est_mask):
-        """
-        Args:
-            mixture_w: [M, N, K]
-            est_mask: [M, C, N, K]
-        Returns:
-            est_source: [M, C, T]
-        """
-        est_source = mixture_w * est_mask
-        est_source = torch.transpose(est_source, 2, 1) # [M,  K, N]
-        est_source = self.basis_signals(est_source)  # [M,  K, L]
-        est_source = overlap_and_add(est_source, self.L//2) # M x C x T
-        return est_source
-
-
-
-
-class Separator(nn.Module):
-    def __init__(self, args):
-        super(Separator, self).__init__()
-
-        self.layer_norm = nn.GroupNorm(1, args.network_audio.encoder_out_nchannels, eps=1e-8)
-        self.bottleneck_conv1x1 = nn.Conv1d(args.network_audio.encoder_out_nchannels, args.network_audio.encoder_out_nchannels, 1, bias=False)
-
-        # mossformer 2
-        intra_model = SBFLASHBlock_DualA(
-            num_layers=args.network_audio.intra_numlayers,
-            d_model=args.network_audio.encoder_out_nchannels,
-            nhead=args.network_audio.intra_nhead,
-            d_ffn=args.network_audio.intra_dffn,
-            dropout=args.network_audio.intra_dropout,
-            use_positional_encoding=args.network_audio.intra_use_positional,
-            norm_before=args.network_audio.intra_norm_before
-        )
-
-        self.masknet = Dual_Path_Model(
-            in_channels=args.network_audio.encoder_out_nchannels,
-            out_channels=args.network_audio.encoder_out_nchannels,
-            intra_model=intra_model,
-            num_layers=args.network_audio.masknet_numlayers,
-            norm=args.network_audio.masknet_norm,
-            K=args.network_audio.masknet_chunksize,
-            num_spks=args.network_audio.masknet_numspks,
-            skip_around_intra=args.network_audio.masknet_extraskipconnection,
-            linear_layer_after_inter_intra=args.network_audio.masknet_useextralinearlayer
-        )
-
-        # reference
-        self.args = args
-        if self.args.network_reference.cue == 'text':
-            self.ref_ds = nn.Linear(768, args.network_reference.emb_size)
-            encoder_layers = TransformerEncoderLayer(d_model=args.network_reference.emb_size, nhead=2, dim_feedforward=args.network_reference.emb_size*2, batch_first=True)
-            self.text_net = TransformerEncoder(encoder_layers, num_layers=args.network_reference.text_layers)
-            self.summarize = nn.LSTM(args.network_reference.emb_size, args.network_reference.emb_size, num_layers=1, batch_first=True)
-            self.fusion = nn.Linear(512+args.network_reference.emb_size, 512)
-        elif self.args.network_reference.cue == 'audio':
-            self.ref_ds = nn.Linear(768, args.network_reference.emb_size)
-            encoder_layers = TransformerEncoderLayer(d_model=args.network_reference.emb_size, nhead=2, dim_feedforward=args.network_reference.emb_size*2, batch_first=True)
-            self.audio_net = TransformerEncoder(encoder_layers, num_layers=args.network_reference.text_layers)
-            self.summarize = nn.LSTM(args.network_reference.emb_size, args.network_reference.emb_size, num_layers=1, batch_first=True)
-            self.fusion = nn.Linear(512+args.network_reference.emb_size, 512)
-
-    def forward(self, x, ref):
-        """
-        Keep this API same with TasNet
-        Args:
-            mixture_w: [M, N, K], M is batch size
-        returns:
-            est_mask: [M, C, N, K]
-        """
-        M, N, D = x.size()
-
-        x = self.layer_norm(x)
-        x = self.bottleneck_conv1x1(x)
-
-        cross_0 = x.transpose(1,2)
-        if self.args.network_reference.cue == 'text':
-            text_embedding, text_attention_mask, text_len = ref
-            text_embedding = self.ref_ds(text_embedding)
-            text_attention_mask = (text_attention_mask==0)
-            text_embedding = self.text_net(text_embedding, src_key_padding_mask=text_attention_mask)
-            text_embedding, _ = self.summarize(text_embedding)
-            text_len = text_len-1
-            batch_indices = torch.arange(text_embedding.size(0))
-            text = text_embedding[batch_indices, text_len]
-
-            text = torch.repeat_interleave(text.unsqueeze(1), repeats=cross_0.shape[1], dim=1)
-            cross_1 = torch.cat((cross_0, text),2)
-            cross_1  = self.fusion(cross_1)
-
-        elif self.args.network_reference.cue == 'audio':
-            audio = self.ref_ds(ref)
-            audio = self.audio_net(audio)
-            audio, _ = self.summarize(audio)
-            audio = audio[:,-1,:]
-            
-            audio = torch.repeat_interleave(audio.unsqueeze(1), repeats=cross_0.shape[1], dim=1)
-            cross_1 = torch.cat((cross_0, audio),2)
-            cross_1  = self.fusion(cross_1)
-        x = cross_1.transpose(1,2)
-
-
-        x = self.masknet(x)
-
-        x = x.squeeze(0)
-
-        return x
-
-
-
-def overlap_and_add(signal, frame_step):
-    """Reconstructs a signal from a framed representation.
-
-    Adds potentially overlapping frames of a signal with shape
-    `[..., frames, frame_length]`, offsetting subsequent frames by `frame_step`.
-    The resulting tensor has shape `[..., output_size]` where
-
-        output_size = (frames - 1) * frame_step + frame_length
-
-    Args:
-        signal: A [..., frames, frame_length] Tensor. All dimensions may be unknown, and rank must be at least 2.
-        frame_step: An integer denoting overlap offsets. Must be less than or equal to frame_length.
-
-    Returns:
-        A Tensor with shape [..., output_size] containing the overlap-added frames of signal's inner-most two dimensions.
-        output_size = (frames - 1) * frame_step + frame_length
-
-    Based on https://github.com/tensorflow/tensorflow/blob/r1.12/tensorflow/contrib/signal/python/ops/reconstruction_ops.py
-    """
-    outer_dimensions = signal.size()[:-2]
-    frames, frame_length = signal.size()[-2:]
-
-    subframe_length = math.gcd(frame_length, frame_step)  # gcd=Greatest Common Divisor
-    subframe_step = frame_step // subframe_length
-    subframes_per_frame = frame_length // subframe_length
-    output_size = frame_step * (frames - 1) + frame_length
-    output_subframes = output_size // subframe_length
-
-    subframe_signal = signal.view(*outer_dimensions, -1, subframe_length)
-
-    frame = torch.arange(0, output_subframes).unfold(0, subframes_per_frame, subframe_step)
-    frame = signal.new_tensor(frame).long().cuda()  # signal may in GPU or CPU
-    frame = frame.contiguous().view(-1)
-
-    result = signal.new_zeros(*outer_dimensions, output_subframes, subframe_length)
-    result.index_add_(-2, frame, subframe_signal)
-    result = result.view(*outer_dimensions, -1)
-    return result