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 models.av_mossformer2_tse.visual_frontend import Visual_encoder 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.av_conv = nn.Conv1d(args.network_audio.encoder_out_nchannels+args.network_reference.emb_size, args.network_audio.encoder_out_nchannels, 1, bias=True) def forward(self, x, visual): """ 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) visual = F.interpolate(visual, (D), mode='linear') x = torch.cat((x, visual),1) x = self.av_conv(x) 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() # 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 class av_mossformer2(nn.Module): def __init__(self, args): super(av_mossformer2, self).__init__() args.causal=0 self.sep_network = Mossformer(args) self.ref_encoder = Visual_encoder(args) def forward(self, mixture, ref): ref = self.ref_encoder(ref) return self.sep_network(mixture, ref) class AV_MossFormer2_TSE_16K(nn.Module): """MossFormer2 model wrapper for outside calling""" def __init__(self, args): super(AV_MossFormer2_TSE_16K, self).__init__() self.model = av_mossformer2(args) def forward(self, x): outputs = self.model(x) return outputs