models/av_mossformer2_tse/av_mossformer_tmp.py

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 VisualFrontend

EPS = 1e-8

class avMossformer(nn.Module):
    def __init__(self, args):
        super(avMossformer, 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
        # visual
        stacks = []
        for x in range(5):
            stacks +=[VisualConv1D(V=256, H=512)]
        self.visual_conv = nn.Sequential(*stacks)
        self.v_ds = nn.Conv1d(512, 256, 1, bias=False)
        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 = visual.transpose(1,2)
        visual = self.v_ds(visual)
        visual = self.visual_conv(visual)
        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().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


class av_mossformer_tmp(nn.Module):
    def __init__(self, args):
        super(av_mossformer_tmp, self).__init__()
        args.causal=0
        self.sep_network = avMossformer(args)
        self.v_front_end = VisualFrontend(args)

    def forward(self, mixture, ref):
        ref = self.v_front_end(ref.unsqueeze(1)).transpose(1,2)
        return self.sep_network(mixture, ref)



class VisualConv1D(nn.Module):
    def __init__(self, V=256, H=512):
        super(VisualConv1D, self).__init__()
        relu_0 = nn.ReLU()
        norm_0 = GlobalLayerNorm(V)
        conv1x1 = nn.Conv1d(V, H, 1, bias=False)
        relu = nn.ReLU()
        norm_1 = GlobalLayerNorm(H)
        dsconv = nn.Conv1d(H, H, 3, stride=1, padding=1,dilation=1, groups=H, bias=False)
        prelu = nn.PReLU()
        norm_2 = GlobalLayerNorm(H)
        pw_conv = nn.Conv1d(H, V, 1, bias=False)
        self.net = nn.Sequential(relu_0, norm_0, conv1x1, relu, norm_1 ,dsconv, prelu, norm_2, pw_conv)

    def forward(self, x):
        out = self.net(x)
        return out + x

class GlobalLayerNorm(nn.Module):
    """Global Layer Normalization (gLN)"""
    def __init__(self, channel_size):
        super(GlobalLayerNorm, self).__init__()
        self.gamma = nn.Parameter(torch.Tensor(1, channel_size, 1))  # [1, N, 1]
        self.beta = nn.Parameter(torch.Tensor(1, channel_size,1 ))  # [1, N, 1]
        self.reset_parameters()

    def reset_parameters(self):
        self.gamma.data.fill_(1)
        self.beta.data.zero_()

    def forward(self, y):
        """
        Args:
            y: [M, N, K], M is batch size, N is channel size, K is length
        Returns:
            gLN_y: [M, N, K]
        """
        # TODO: in torch 1.0, torch.mean() support dim list
        mean = y.mean(dim=1, keepdim=True).mean(dim=2, keepdim=True) #[M, 1, 1]
        var = (torch.pow(y-mean, 2)).mean(dim=1, keepdim=True).mean(dim=2, keepdim=True)
        gLN_y = self.gamma * (y - mean) / torch.pow(var + EPS, 0.5) + self.beta
        return gLN_y