Upload 7 files

app.py

@@ -0,0 +1,69 @@
+import gradio as gr
+from model.CLAPSep import CLAPSep
+import torch
+import librosa
+import numpy as np
+
+
+model_config = {"lan_embed_dim": 1024,
+    "depths": [1, 1, 1, 1],
+    "embed_dim": 128,
+    "encoder_embed_dim": 128,
+    "phase": False,
+    "spec_factor": 8,
+    "d_attn": 640,
+    "n_masker_layer": 3,
+    "conv": False}
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+CLAP_path = "model/music_audioset_epoch_15_esc_90.14.pt"
+
+
+model = CLAPSep(model_config, CLAP_path).to(DEVICE)
+ckpt = torch.load('model/best_model.ckpt', map_location=DEVICE)
+model.load_state_dict(ckpt, strict=False)
+model.eval()
+
+
+
+def inference(audio_file_path: str, text_p: str, text_n: str):
+    print(f"Separate audio from [{audio_file_path}] with textual query p: [{text_p}] and n: [{text_n}]")
+
+    mixture, _ = librosa.load(audio_file_path, sr=32000)
+
+    pad = (320000 - (len(mixture) % 320000))if len(mixture) % 320000 != 0 else 0
+
+    mixture =torch.tensor(np.pad(mixture,(0,pad)))
+
+    mixture_chunks = torch.chunk(mixture, dim=0, chunks=len(mixture)//320000)
+    sep_segments = []
+    for chunk in mixture_chunks:
+        with torch.no_grad():
+            sep_segments.append(model.inference_from_data(chunk.unsqueeze(0), [text_p], [text_n]))
+
+    sep_segment = torch.concat(sep_segments, dim=1)
+
+    return 32000, sep_segment.squeeze().numpy()
+
+
+with gr.Blocks(title="CLAPSep") as demo:
+    with gr.Row():
+        with gr.Column():
+            input_audio = gr.Audio(label="Mixture", type="filepath")
+            text_p = gr.Textbox(label="Positive Query")
+            text_n = gr.Textbox(label="Negative Query")
+        with gr.Column():
+            with gr.Column():
+                output_audio = gr.Audio(label="Separation Result", scale=10)
+                button = gr.Button(
+                    "Separate",
+                    variant="primary",
+                    scale=2,
+                    size="lg",
+                    interactive=True,
+                )
+                button.click(
+                    fn=inference, inputs=[input_audio, text_p, text_n], outputs=[output_audio]
+                )
+
+
+demo.queue().launch(share=True)

model/CLAPSep.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+'''
+@Project ：Waveformer-main 
+@File    ：CLAPSep.py
+@IDE     ：PyCharm 
+@Author  ：Aisaka/Hao Ma @SDU
+@Date    ：2024/2/28 下午1:12 
+'''
+
+import torch
+from torch import nn
+import torchaudio
+import laion_clap
+from .CLAPSep_decoder import HTSAT_Decoder
+import copy
+import loralib as lora
+from torchlibrosa import ISTFT, STFT
+from torchlibrosa.stft import magphase
+import librosa
+
+def set_module(model, submodule_key, module):
+    tokens = submodule_key.split('.')
+    sub_tokens = tokens[:-1]
+    cur_mod = model
+    for s in sub_tokens:
+        cur_mod = getattr(cur_mod, s)
+    setattr(cur_mod, tokens[-1], module)
+
+
+def process_model(model, rank):
+    for n, module in model.named_modules():
+        if 'WindowAttention' in str(type(module)):
+            for n_, layer in module.named_modules():
+                if isinstance(layer, torch.nn.Linear):
+                    lora_layer = lora.Linear(layer.in_features, layer.out_features, r=rank,
+                                             bias=hasattr(layer, 'bias'), merge_weights=True)
+                    lora_layer.weight = layer.weight
+                    if hasattr(layer, 'bias'):
+                        lora_layer.bias = layer.bias
+                    set_module(model, n+'.'+n_, lora_layer)
+    return model
+
+
+class CLAPSep(nn.Module):
+    def __init__(self, model_config, CLAP_path, use_lora=True, rank=16, nfft=1024):
+        super().__init__()
+        self.resampler = torchaudio.transforms.Resample(32000, 48000)
+        self.clap_model = laion_clap.CLAP_Module(enable_fusion=False, amodel='HTSAT-base', device='cpu')
+        self.clap_model.load_ckpt(CLAP_path)
+        for p in self.clap_model.parameters():
+            p.requires_grad = False
+        self.audio_branch = copy.deepcopy(self.clap_model.model.audio_branch)
+        if use_lora:
+            process_model(self.audio_branch, rank)
+        self.decoder_model = HTSAT_Decoder(**model_config)
+        self.stft = STFT(n_fft=nfft, hop_length=320,
+                         win_length=nfft, window='hann', center=True, pad_mode='reflect',
+                         freeze_parameters=True)
+        self.istft = ISTFT(n_fft=nfft, hop_length=320,
+                           win_length=nfft, window='hann', center=True, pad_mode='reflect',
+                           freeze_parameters=True)
+        self.features = self.install_forward_hooks()
+
+    def wav_reconstruct(self, mask, mag_x, cos_x, sin_x, length):
+        mag_y = torch.nn.functional.relu_(mag_x * mask)
+        cos_y = cos_x
+        sin_y = sin_x
+        pred = self.istft(mag_y * cos_y, mag_y * sin_y, length=length)
+        return pred
+
+    def inference_from_data(self, mixed, pos_prompt, neg_prompt):
+        self.eval()
+        real, imag = self.stft(mixed)
+        mag, cos, sin = magphase(real, imag)
+        self.features.append(mag)
+        with torch.no_grad():
+            embed_pos, embed_neg = torch.chunk(self.clap_model.get_text_embedding(pos_prompt + neg_prompt,
+                                                                                  use_tensor=True), dim=0, chunks=2)
+            embed_pos = torch.zeros_like(embed_pos) if pos_prompt == '' else embed_pos
+            embed_neg = torch.zeros_like(embed_neg) if neg_prompt == '' else embed_neg
+            embed = torch.concat([embed_pos, embed_neg], dim=-1)
+            self.audio_branch({"waveform": self.resampler(mixed)})
+            mask = self.decoder_model(hidden_state=self.features[-1], skip_features=self.features[:-1], embed=embed)
+            pred = self.wav_reconstruct(mask, mag, cos, sin, length=mixed.size(-1))
+        del self.features[:]
+        return pred
+
+    def install_forward_hooks(self):
+        features = []
+
+        def get_features_list(_, __, output):
+            features.append(output)
+
+        def get_features_list_basic_layer(_, __, output):
+            features.append(output[0])
+
+        def spectrogram_padding(_, __, out):
+            return torch.nn.functional.pad(out, (0, 0, 0, 1024 - out.size(2)))
+
+        self.audio_branch.spectrogram_extractor.register_forward_hook(spectrogram_padding)
+        self.audio_branch.patch_embed.register_forward_hook(get_features_list)
+        for module in self.audio_branch.layers:
+            module.register_forward_hook(get_features_list_basic_layer)
+        return features
+
+if __name__ == '__main__':
+    model_config = {"lan_embed_dim": 1024,
+    "depths": [1, 1, 1, 1],
+    "embed_dim": 128,
+    "encoder_embed_dim": 128,
+    "phase": False,
+    "spec_factor": 8,
+    "d_attn": 640,
+    "n_masker_layer": 3,
+    "conv": False}
+    CLAP_path = "./music_audioset_epoch_15_esc_90.14.pt"
+
+    model = CLAPSep(model_config, CLAP_path)
+    ckpt = torch.load('best_model.ckpt', map_location='cpu')
+    model.load_state_dict(ckpt, strict=False)
+    model.eval()
+    audio, fs = librosa.load("./510_25.221254348754883_mixture.wav", sr=32000)
+    pred = model.inference_from_data(torch.tensor(audio).unsqueeze(0), pos_prompt=[''], neg_prompt=['A vehicle engine revving then powering down.'])
+    import soundfile as sf
+    sf.write('./pred.wav', pred.squeeze().numpy(), 32000)

model/CLAPSep_decoder.py

@@ -0,0 +1,605 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+'''
+@Project ：Waveformer-main
+@File    ：CLAPsep_decoder.py
+@IDE     ：PyCharm
+@Author  ：Aisaka/Hao Ma @SDU
+@Date    ：2023/10/31 下午8:34
+'''
+
+from laion_clap.clap_module.htsat import *
+from einops import rearrange
+import numpy as np
+
+class Transpose(nn.Module):
+
+    def __init__(self, dim0, dim1):
+        super(Transpose, self).__init__()
+        self.dim0 = dim0
+        self.dim1 = dim1
+
+    def forward(self, x):
+        return x.transpose(self.dim0, self.dim1)
+
+
+class Swish(nn.Module):
+
+    def __init__(self):
+        super(Swish, self).__init__()
+
+    def forward(self, x):
+        return x * x.sigmoid()
+
+
+class Glu(nn.Module):
+
+    def __init__(self, dim):
+        super(Glu, self).__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        x_in, x_gate = x.chunk(2, dim=self.dim)
+        return x_in * x_gate.sigmoid()
+
+
+class FiLM(nn.Module):
+    def __init__(self, dim_in=1024, hidden_dim=768):
+        super(FiLM, self).__init__()
+        self.beta = nn.Linear(dim_in, hidden_dim)
+        self.gamma = nn.Linear(dim_in, hidden_dim)
+
+    def forward(self, hidden_state, embed):
+        embed = embed.unsqueeze(1)
+        return self.gamma(embed) * hidden_state + self.beta(embed)
+
+
+class SkipTrans(nn.Module):
+    def __init__(self, in_features, out_features, embed_dim=512, film=True):
+        super(SkipTrans, self).__init__()
+        self.film = film
+        if film:
+            self.skip_conv = FiLM(embed_dim, out_features)
+        self.feature_proj = nn.Linear(in_features, out_features)
+        self.norm = nn.LayerNorm(out_features)
+
+    def forward(self, skip, embed, x=None):
+        out = self.feature_proj(skip)
+        if self.film:
+            out = self.skip_conv(out, embed)
+        return self.norm(out) if x is None else self.norm(out + x)
+
+class Conv1d(nn.Conv1d):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride = 1,
+        padding = "same",
+        dilation = 1,
+        groups = 1,
+        bias = True
+    ):
+        super(Conv1d, self).__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=0,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            padding_mode="zeros")
+
+        # Assert
+        assert padding in ["valid", "same", "causal"]
+
+        # Padding
+        if padding == "valid":
+            self.pre_padding = None
+        elif padding == "same":
+            self.pre_padding = nn.ConstantPad1d(padding=((kernel_size - 1) // 2, (kernel_size - 1) // 2), value=0)
+        elif padding == "causal":
+            self.pre_padding = nn.ConstantPad1d(padding=(kernel_size - 1, 0), value=0)
+
+        # Variational Noise
+        self.noise = None
+        self.vn_std = None
+
+    def init_vn(self, vn_std):
+
+        # Variational Noise
+        self.vn_std = vn_std
+
+    def sample_synaptic_noise(self, distributed):
+
+        # Sample Noise
+        self.noise = torch.normal(mean=0.0, std=1.0, size=self.weight.size(), device=self.weight.device, dtype=self.weight.dtype)
+
+        # Broadcast Noise
+        if distributed:
+            torch.distributed.broadcast(self.noise, 0)
+
+    def forward(self, input):
+
+        # Weight
+        weight = self.weight
+
+        # Add Noise
+        if self.noise is not None and self.training:
+            weight = weight + self.vn_std * self.noise
+
+        # Padding
+        if self.pre_padding is not None:
+            input = self.pre_padding(input)
+
+        # Apply Weight
+        return F.conv1d(input, weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+
+
+class ConvolutionModule(nn.Module):
+    """Conformer Convolution Module
+
+    Args:
+        dim_model: input feature dimension
+        dim_expand: output feature dimension
+        kernel_size: 1D depthwise convolution kernel size
+        Pdrop: residual dropout probability
+        stride: 1D depthwise convolution stride
+        padding: "valid", "same" or "causal"
+
+    Input: (batch size, input length, dim_model)
+    Output: (batch size, output length, dim_expand)
+
+    """
+
+    def __init__(self, dim_model, dim_expand, kernel_size, Pdrop, stride, padding):
+        super(ConvolutionModule, self).__init__()
+
+        # Layers
+        self.layers = nn.Sequential(
+            nn.LayerNorm(dim_model, eps=1e-6),
+            Transpose(1, 2),
+            Conv1d(dim_model, 2 * dim_expand, kernel_size=1),
+            Glu(dim=1),
+            Conv1d(dim_expand, dim_expand, kernel_size, stride=stride, padding=padding, groups=dim_expand),
+            nn.BatchNorm1d(dim_expand),
+            Swish(),
+            Conv1d(dim_expand, dim_expand, kernel_size=1),
+            Transpose(1, 2),
+            nn.Dropout(p=Pdrop)
+        )
+        self.ln = nn.LayerNorm(dim_expand)
+
+    def forward(self, x):
+        return self.ln(self.layers(x)+x)
+
+
+class BasicLayerDec(nn.Module):
+    """ A basic Swin Transformer layer for one stage.
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Local window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,
+                 norm_before_mlp='ln'):
+
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            SwinTransformerBlock(dim=dim, input_resolution=input_resolution,
+                                 num_heads=num_heads, window_size=window_size,
+                                 shift_size=0 if (i % 2 == 0) else window_size // 2,
+                                 mlp_ratio=mlp_ratio,
+                                 qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                 drop=drop, attn_drop=attn_drop,
+                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                                 norm_layer=norm_layer, norm_before_mlp=norm_before_mlp)
+            for i in range(depth)])
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample((input_resolution[0]//2, input_resolution[1]//2), dim=dim * 2, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        attns = []
+        if self.downsample is not None:
+            x = self.downsample(x)
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x, attn = blk(x)
+                if not self.training:
+                    attns.append(attn.unsqueeze(0))
+        if not self.training:
+            attn = torch.cat(attns, dim = 0)
+            attn = torch.mean(attn, dim = 0)
+        return x, attn
+
+    def extra_repr(self):
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
+
+
+class PatchExpand(nn.Module):
+    def __init__(self, input_resolution, dim, dim_scale=2, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.expand = nn.Linear(dim, 2 * dim, bias=False) if dim_scale == 2 else nn.Identity()
+        self.norm = norm_layer(dim // dim_scale)
+
+    def forward(self, x):
+        """
+        x: B, H*W, C
+        """
+        H, W = self.input_resolution
+        x = self.expand(x)
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        x = x.view(B, H, W, C)
+        # This is the original implementation in SwinUnet
+        # x = rearrange(x, 'b h w (p1 p2 c)-> b (h p1) (w p2) c', p1=2, p2=2, c=C // 4)
+
+        # here is our implementation
+        # can reverse patch-emerging in Swin-Transformer encoder, seems helpful
+        x0, x2, x1, x3 = x.chunk(4, dim=-1)
+        x = torch.stack((x0, x1, x2, x3), dim=-1)
+        x = torch.chunk(x, C // 4, dim=-2)
+        x = torch.concat(x, dim=-1).squeeze(-2)
+        x = rearrange(x, 'b h w c -> b c h w')
+        x = torch.nn.functional.pixel_shuffle(x, 2)
+        x = rearrange(x, 'b c h w -> b h w c')
+        x = x.view(B, -1, C // 4)
+        x = self.norm(x)
+
+        return x
+
+
+class InversePatchEmbed(nn.Module):
+    """
+    Patch Embedding to 2D Image.
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True,
+                 patch_stride=16):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patch_stride = to_2tuple(patch_stride)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patch_stride = patch_stride
+        self.grid_size = (img_size[0] // patch_stride[0], img_size[1] // patch_stride[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        padding = ((patch_size[0] - patch_stride[0]) // 2, (patch_size[1] - patch_stride[1]) // 2)
+
+        self.proj = nn.ConvTranspose2d(embed_dim, in_chans, kernel_size=patch_size, stride=patch_stride, padding=padding)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        # B, C, H, W = x.shape
+        # assert H == self.img_size[0] and W == self.img_size[1], \
+        #     f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.norm(x)
+        if self.flatten:
+            # x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+            x = x.transpose(1, 2).unflatten(2, self.grid_size).contiguous()  # BNC -> BCHW
+        x = self.proj(x)
+
+        return x
+
+
+class HTSAT_Decoder(nn.Module):
+    r"""HTSAT_decoder based on the Swin Transformer
+    Args:
+        spec_size (int | tuple(int)): Input Spectrogram size. Default 256
+        patch_size (int | tuple(int)): Patch size. Default: 4
+        path_stride (iot | tuple(int)): Patch Stride for Frequency and Time Axis. Default: 4
+        in_chans (int): Number of input image channels. Default: 1 (mono)
+        num_classes (int): Number of classes for classification head. Default: 527
+        embed_dim (int): Patch embedding dimension. Default: 96
+        depths (tuple(int)): Depth of each HTSAT-Swin Transformer layer.
+        num_heads (tuple(int)): Number of attention heads in different layers.
+        window_size (int): Window size. Default: 8
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop_rate (float): Dropout rate. Default: 0
+        attn_drop_rate (float): Attention dropout rate. Default: 0
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
+    """
+
+    def __init__(self, lan_embed_dim=512, spec_size=256, patch_size=4, patch_stride=(4, 4),
+                 in_chans=1, num_classes=527,
+                 embed_dim=48, depths=[1, 1, 1, 1], num_heads=[4, 8, 16, 32],
+                 window_size=8, mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
+                 norm_layer=nn.LayerNorm,
+                 ape=False, patch_norm=True,
+                 use_checkpoint=False, norm_before_mlp='ln', encoder_embed_dim=96, phase=False,
+                 spec_factor=8, d_attn=640, n_masker_layer=4, conv=False):
+        super(HTSAT_Decoder, self).__init__()
+        self.mel_bins = 64
+        self.spec_size = spec_size
+        self.phase = phase
+        self.patch_stride = patch_stride
+        self.patch_size = patch_size
+        self.window_size = window_size
+        self.embed_dim = embed_dim
+        self.depths = depths
+        self.ape = ape
+        self.in_chans = in_chans
+        self.num_classes = num_classes
+        self.num_heads = num_heads
+        self.num_layers = len(self.depths)
+        self.num_features = int(self.embed_dim * 2 ** (self.num_layers - 1))
+
+        self.drop_rate = drop_rate
+        self.attn_drop_rate = attn_drop_rate
+        self.drop_path_rate = drop_path_rate
+
+        self.qkv_bias = qkv_bias
+        self.qk_scale = None
+
+        self.patch_norm = patch_norm
+        self.norm_layer = norm_layer if self.patch_norm else None
+        self.norm_before_mlp = norm_before_mlp
+        self.mlp_ratio = mlp_ratio
+
+        self.use_checkpoint = use_checkpoint
+
+        #  process mel-spec ; used only once
+        self.freq_ratio = self.spec_size // self.mel_bins
+
+
+        # split spctrogram into non-overlapping patches
+        self.inverse_patch_embed = InversePatchEmbed(
+            img_size=self.spec_size, patch_size=self.patch_size, in_chans=self.in_chans,
+            embed_dim=self.embed_dim, norm_layer=self.norm_layer, patch_stride=patch_stride)
+
+        patches_resolution = self.inverse_patch_embed.grid_size
+        self.patches_resolution = patches_resolution
+
+
+        # stochastic depth
+        dpr = [x.item() for x in
+               torch.linspace(0, self.drop_path_rate, sum(self.depths))]  # stochastic depth decay rule
+
+        # build layers
+        self.layers = nn.ModuleList()
+        self.skip = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = BasicLayerDec(dim=int(self.embed_dim * 2 ** i_layer),
+                               input_resolution=(patches_resolution[0] // (2 ** i_layer),
+                                                 patches_resolution[1] // (2 ** i_layer)),
+                               depth=self.depths[i_layer],
+                               num_heads=self.num_heads[i_layer],
+                               window_size=self.window_size,
+                               mlp_ratio=self.mlp_ratio,
+                               qkv_bias=self.qkv_bias, qk_scale=self.qk_scale,
+                               drop=self.drop_rate, attn_drop=self.attn_drop_rate,
+                               drop_path=dpr[sum(self.depths[:i_layer]):sum(self.depths[:i_layer + 1])],
+                               norm_layer=self.norm_layer,
+                               downsample=PatchExpand if (i_layer < self.num_layers - 1) else None,
+                               use_checkpoint=use_checkpoint,
+                               norm_before_mlp=self.norm_before_mlp)
+            self.layers.append(layer)
+            self.skip.append(
+                SkipTrans(embed_dim=lan_embed_dim, in_features=int(encoder_embed_dim * 2 ** i_layer), out_features=int(self.embed_dim * 2 ** i_layer)),
+            )
+        self.layers = self.layers[::-1]
+        self.skip = self.skip[::-1]
+        # self.skip.append(
+        #     SkipTrans(embed_dim=lan_embed_dim, in_features=self.mel_bins, out_features=self.mel_bins),
+        # )
+
+        d_spec = self.mel_bins * spec_factor + 1
+
+        self.spec_norm = nn.BatchNorm2d(d_spec, momentum=0.01)
+        self.conv = conv
+        if not conv:
+            encoder_layer = nn.TransformerEncoderLayer(d_model=d_attn, nhead=8,
+                                                       dim_feedforward=int(d_attn * self.mlp_ratio),
+                                                       batch_first=True, dropout=0)
+            transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_masker_layer)
+
+            self.mask_net = nn.Sequential(
+                nn.Linear(self.mel_bins + d_spec, d_attn),
+                nn.LayerNorm(d_attn),
+                transformer_encoder,
+                nn.Linear(d_attn, d_spec)
+            )
+        else:
+            self.mask_net = nn.Sequential(
+                nn.Linear(self.mel_bins + d_spec, d_spec),
+                nn.LayerNorm(d_spec),
+                *[ConvolutionModule(dim_model=d_spec, dim_expand=d_spec, kernel_size=9, padding='same',
+                                    Pdrop=0, stride=1) for i in range(n_masker_layer)]
+            )
+        if self.phase:
+            self.phase_net = nn.Sequential(
+                nn.Linear(self.mel_bins + d_spec, d_spec * 2),
+                nn.LayerNorm(d_spec * 2),
+                *[ConvolutionModule(dim_model=d_spec * 2, dim_expand=d_spec * 2, kernel_size=9, padding='same',
+                                    Pdrop=0, stride=1) for i in range(n_masker_layer)]
+            )
+
+        self.film = SkipTrans(embed_dim=lan_embed_dim, in_features=encoder_embed_dim * 8, out_features=self.num_features)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    # @torch.jit.ignore
+    # def no_weight_decay(self):
+    #     return {'absolute_pos_embed'}
+    #
+    # @torch.jit.ignore
+    # def no_weight_decay_keywords(self):
+    #     return {'relative_position_bias_table'}
+
+    def forward(self, hidden_state, skip_features, embed):
+        skip_features = skip_features[::-1]
+        # hidden_state = torch.randn(hidden_state.shape).type_as(hidden_state)
+
+        spec = skip_features[-1]
+
+        h = self.film(hidden_state, embed)
+
+        for i, (layer, f, skip) in enumerate(zip(self.layers, skip_features, self.skip)):
+            h = layer(h)[0]
+            h = skip(skip=f, embed=embed, x=h)
+
+        h = self.reshape_img2wav(self.inverse_patch_embed(h)).squeeze(1)
+
+        h = h[:, :spec.size(2), :]
+
+        spec = spec.transpose(1, 3)
+
+        spec = self.spec_norm(spec).transpose(1, 3).squeeze(1)
+
+        h = torch.concat([spec, h], dim=-1)
+
+        mask = self.mask_net(h).unsqueeze(1)
+
+        if self.phase:
+            mask_r, mask_i = torch.chunk(self.phase_net(h).unsqueeze(1), chunks=2, dim=-1)
+            return torch.sigmoid(mask), torch.tanh(mask_r), torch.tanh(mask_i)
+        else:
+            return torch.sigmoid(mask)
+
+    def reshape_img2wav(self, x):
+        # (B, 1, 256, 256)
+        x = x.reshape(x.shape[0], x.shape[1], self.freq_ratio, x.shape[2]//self.freq_ratio, x.shape[3]) # (B, 1, 4, 64, 256)
+        x = x.permute(0, 1, 3, 2, 4).contiguous()
+        x = x.reshape(x.shape[0], x.shape[1], x.shape[2], x.shape[3] * x.shape[4])
+        x = x.permute(0, 1, 3, 2).contiguous()
+        return x
+
+
+# if __name__ == "__main__":
+#     import torch
+#     from msclap import CLAP
+#     import os
+#     import torchaudio
+#     import torchaudio.transforms as T
+#     import numpy as np
+#     import random
+#     from torchlibrosa import Spectrogram, LogmelFilterBank
+#     clap_model = CLAP(model_fp="/home/user/202212661/clapsep/Waveformer-main/checkpoint_path/CLAP_weights_2023.pth",
+#                       version='2023', use_cuda=True)
+#     text_data = [
+#         "Acoustic_guitar", "Applause", "Bark", "Bass_drum", "Burping_or_eructation",
+#         "Bus", "Cello", "Chime", "Clarinet", "Computer_keyboard",
+#         "Cough", "Cowbell", "Double_bass", "Drawer_open_or_close", "Electric_piano",
+#         "Fart", "Finger_snapping", "Fireworks", "Flute", "Glockenspiel",
+#         "Gong", "Gunshot_or_gunfire", "Harmonica", "Hi-hat", "Keys_jangling",
+#         "Knock", "Laughter", "Meow", "Microwave_oven", "Oboe",
+#         "Saxophone", "Scissors", "Shatter", "Snare_drum", "Squeak",
+#         "Tambourine", "Tearing", "Telephone", "Trumpet", "Violin_or_fiddle",
+#         "Writing"]
+#     # Extract text embeddings
+#     text_embeddings = clap_model.get_text_embeddings(text_data)
+#     path = "/home/user/202212661/clapsep/Waveformer-main/data/FSDSoundScapes/FSDKaggle2018/train/Tearing/2232ce13.wav"
+#     # Extract audio embeddings
+#     audio_embeddings_ = clap_model.get_audio_embeddings([path])
+#
+#     window = 'hann'
+#     center = True
+#     pad_mode = 'reflect'
+#     ref = 1.0
+#     amin = 1e-10
+#     top_db = None
+#
+#     spectrogram_extractor = Spectrogram(n_fft=512, hop_length=160,
+#                                         win_length=512, window=window, center=center, pad_mode=pad_mode,
+#                                         freeze_parameters=True).cuda()
+#     # Logmel feature extractor
+#     logmel_extractor = LogmelFilterBank(sr=16000, n_fft=512,
+#                                         n_mels=64, fmin=0, fmax=8000, ref=ref, amin=amin,
+#                                         top_db=top_db,
+#                                         freeze_parameters=True).cuda()
+#
+#     clap_model.clap.audio_encoder.base.htsat.spectrogram_extractor = spectrogram_extractor
+#     clap_model.clap.audio_encoder.base.htsat.logmel_extractor = logmel_extractor
+#
+#     features = []
+#
+#
+#     def get_features_list(module, input, output):
+#         features.append(output)
+#
+#
+#     def get_features_list_basic_layer(module, input, output):
+#         features.append(output[0])
+#
+#
+#     clap_model.clap.audio_encoder.base.htsat.patch_embed.register_forward_hook(get_features_list)
+#     for module in clap_model.clap.audio_encoder.base.htsat.layers:
+#         module.register_forward_hook(get_features_list_basic_layer)
+#
+#     audio_time_series, sample_rate = torchaudio.load(path)
+#     resample_rate = 16000
+#     if resample_rate != sample_rate:
+#         resampler = T.Resample(sample_rate, resample_rate)
+#         audio_time_series = resampler(audio_time_series)
+#
+#     sample_rate = resample_rate
+#     audio_duration = 10
+#     audio_time_series = audio_time_series.reshape(-1)
+#     if audio_duration * sample_rate >= audio_time_series.shape[0]:
+#         repeat_factor = int(np.ceil((audio_duration * sample_rate) /
+#                                     audio_time_series.shape[0]))
+#         # Repeat audio_time_series by repeat_factor to match audio_duration
+#         audio_time_series = audio_time_series.repeat(repeat_factor)
+#         # remove excess part of audio_time_series
+#         audio_time_series = audio_time_series[0:audio_duration * sample_rate]
+#     else:
+#         # audio_time_series is longer than predefined audio duration,
+#         # so audio_time_series is trimmed
+#         start_index = random.randrange(
+#             audio_time_series.shape[0] - audio_duration * sample_rate)
+#         audio_time_series = audio_time_series[start_index:start_index +
+#                                                           audio_duration * sample_rate]
+#

model/best_model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6fcc8dbcd7174af86266cf16b4105eced0802352762f81dbfefdce29af3dba04
+size 177818986

model/music_audioset_epoch_15_esc_90.14.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fae3e9c087f2909c28a09dc31c8dfcdacbc42ba44c70e972b58c1bd1caf6dedd
+size 2352471003

requirements.txt

@@ -0,0 +1,8 @@
+torch
+librosa
+torchaudio
+torchlibrosa
+numpy
+einops
+loralib
+laion-clap