Add CLAP & HiFiGAN

CLAP/__init__.py

@@ -0,0 +1,5 @@
+import os
+import sys
+dir_path = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(dir_path)
+from .clap_module import CLAP_Module

CLAP/clap_model.py

@@ -0,0 +1,165 @@
+import numpy as np
+from pathlib import Path
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from transformers import RobertaModel
+from .htsat import create_htsat_model
+
+BASE_DIR = Path(__file__).resolve().parent
+
+class MLPLayers(nn.Module):
+    def __init__(self, units=[512, 512, 512], nonlin=nn.ReLU(), dropout=0.1):
+        super(MLPLayers, self).__init__()
+        self.nonlin = nonlin
+        self.dropout = dropout
+
+        sequence = []
+        for u0, u1 in zip(units[:-1], units[1:]):
+            sequence.append(nn.Linear(u0, u1))
+            sequence.append(self.nonlin)
+            sequence.append(nn.Dropout(self.dropout))
+        sequence = sequence[:-2]
+
+        self.sequential = nn.Sequential(*sequence)
+
+    def forward(self, X):
+        X = self.sequential(X)
+        return X
+
+
+# Audio Config Class
+@dataclass
+class CLAPAudioCfp:
+    model_type: str = "PANN"
+    model_name: str = "Cnn14"
+    sample_rate: int = 48000
+    # Param
+    audio_length: int = 1024
+    window_size: int = 1024
+    hop_size: int = 1024
+    fmin: int = 50
+    fmax: int = 14000
+    class_num: int = 527
+    mel_bins: int = 64
+    clip_samples: int = 480000
+
+
+@dataclass
+class CLAPTextCfg:
+    context_length: int
+    vocab_size: int
+    width: int
+    heads: int
+    layers: int
+    model_type: str
+
+
+class CLAP(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        audio_cfg: CLAPAudioCfp,
+        text_cfg: CLAPTextCfg,
+    ):
+        super().__init__()
+
+        audio_cfg = CLAPAudioCfp(**audio_cfg)
+        text_cfg = CLAPTextCfg(**text_cfg)
+        self.context_length = text_cfg.context_length
+
+        mlp_act_layer = nn.ReLU()
+
+        # audio branch
+        self.audio_branch = create_htsat_model(audio_cfg)
+
+        # audio branch parameters
+        self.audio_transform = MLPLayers(units=[512,
+                                                512,
+                                                512], dropout=0.1)
+
+        self.audio_projection = nn.Sequential(
+                nn.Linear(embed_dim, 512),
+                mlp_act_layer,
+                nn.Linear(512, 512)
+            )
+        
+        # text branch
+        self.text_branch = RobertaModel.from_pretrained('roberta-base')
+            
+        # text branch parameters
+        self.text_transform = MLPLayers(units=[512,
+                                               512,
+                                               512], dropout=0.1)
+        self.text_projection = nn.Sequential(
+            nn.Linear(768, 512),
+            mlp_act_layer,
+            nn.Linear(512, 512)
+        )
+
+        self.logit_scale_a = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        self.logit_scale_t = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        self.register_buffer("attn_mask", self.build_attention_mask(), persistent=False)
+
+        self.init_text_branch_parameters()
+
+    def init_text_branch_parameters(self):
+        nn.init.constant_(self.logit_scale_a, np.log(1 / 0.07))
+        nn.init.constant_(self.logit_scale_t, np.log(1 / 0.07))
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+    
+    def get_word_embedding(self, data):
+        device = next(self.parameters()).device
+        for k in data:
+            data[k] = data[k].to(device)
+        
+        word_embeds = self.text_branch.embeddings.word_embeddings(
+            data['input_ids'].to(device=device, non_blocking=True)
+        )
+
+        return word_embeds
+
+    def get_text_embedding(self, data, normalize=False):
+        
+        device = next(self.parameters()).device
+        for k in data:
+            data[k] = data[k].to(device)
+            
+        x = self.text_branch(
+            input_ids=data["input_ids"].to(device=device, non_blocking=True),
+            attention_mask=data["attention_mask"].to(
+                device=device, non_blocking=True
+            ),
+        )["pooler_output"]
+        text_embeds = self.text_projection(x)
+        
+        if normalize:
+            text_embeds = F.normalize(text_embeds, dim=-1)
+        
+        return text_embeds
+
+    def get_audio_embedding(self, data, normalize=False):
+        
+        device = next(self.parameters()).device
+        input_dict = {}
+        keys = data[0].keys()
+        for k in keys:
+            input_dict[k] = torch.cat([d[k].unsqueeze(0) for d in data], dim=0).to(device)
+        
+        audio_embeds = self.audio_branch(input_dict, mixup_lambda=None, device=device)["embedding"]
+        audio_embeds = self.audio_projection(audio_embeds)
+        
+        if normalize:
+            audio_embeds = F.normalize(audio_embeds, dim=-1)
+        return audio_embeds
+

CLAP/clap_module.py

@@ -0,0 +1,141 @@
+"""
+Contrastive Language-Audio Pretraining Model from LAION
+--------------------------------------------------------
+Paper: https://arxiv.org/abs/2211.06687
+Authors (equal contributions): Ke Chen, Yusong Wu, Tianyu Zhang, Yuchen Hui
+Support: LAION
+"""
+import os
+import json
+import torch
+import librosa
+import torchaudio
+import transformers
+import numpy as np
+from pathlib import Path
+from packaging import version
+
+from .data import get_audio_features
+from .data import int16_to_float32, float32_to_int16
+from .clap_model import CLAP
+
+from transformers import RobertaTokenizer
+import wget
+
+BASE_DIR = Path(__file__).resolve().parent
+
+class CLAP_Module(torch.nn.Module):
+    def __init__(self, amodel='HTSAT-tiny', tmodel='roberta') -> None:
+        super(CLAP_Module, self).__init__()
+        
+        config_path = os.path.join(BASE_DIR, 'model_configs', f'{amodel}.json')
+        with open(config_path, "r") as f:
+            model_cfg = json.load(f)
+        
+        self.tokenize = RobertaTokenizer.from_pretrained("roberta-base")
+                
+        model_cfg["text_cfg"]["model_type"] = tmodel
+        model = CLAP(**model_cfg)
+        
+        self.model = model
+        self.model_cfg = model_cfg
+
+    def tokenizer(self, text):
+        result = self.tokenize(
+            text,
+            padding="max_length",
+            truncation=True,
+            max_length=77,
+            return_tensors="pt",
+        )
+        return result
+
+    def load_ckpt(self, ckpt_folder_path, ckpt_name):
+        ckpt_path = os.path.join(ckpt_folder_path, ckpt_name)
+        
+        if os.path.exists(ckpt_path):
+            print(f'Load checkpoint from {ckpt_path}')
+        else:
+            download_link = 'https://huggingface.co/lukewys/laion_clap/resolve/main/'
+            print(f'Download checkpoint from {download_link + ckpt_name}.')
+            ckpt_path = wget.download(download_link + ckpt_name, ckpt_folder_path)
+            print('Download completed!')
+        print()
+        
+        checkpoint = torch.load(ckpt_path, map_location='cpu', weights_only=False)
+        
+        if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
+            state_dict = checkpoint["state_dict"]
+        else:
+            state_dict = checkpoint
+            
+        if next(iter(state_dict.items()))[0].startswith("module"):
+            state_dict = {k[7:]: v for k, v in state_dict.items()}
+             
+        if version.parse(transformers.__version__) >= version.parse("4.31.0"): 
+            del state_dict["text_branch.embeddings.position_ids"]
+    
+        self.model.load_state_dict(state_dict)
+    
+    def get_audio_embedding(self, x, sr=16000, normalize=False, use_tensor=True):
+        self.model.eval()
+        if isinstance(x, str):
+            x = [x]
+
+        audio_input = []
+        for audio_waveform in x:
+            
+            if isinstance(audio_waveform, str):
+                # load the waveform of the shape (T,), should resample to 48000
+                audio_waveform, _ = librosa.load(audio_waveform, sr=48000)
+            elif sr != 48000:
+                audio_waveform = torchaudio.functional.resample(audio_waveform, orig_freq=sr, new_freq=48000)                
+                
+            if isinstance(audio_waveform, torch.Tensor):
+                audio_waveform = audio_waveform.numpy()
+                
+            # quantize
+            audio_waveform = int16_to_float32(float32_to_int16(audio_waveform))
+            audio_waveform = torch.from_numpy(audio_waveform).float()
+
+            temp_dict = {}
+            temp_dict = get_audio_features(
+                temp_dict, audio_waveform, 480000, 
+                data_truncating='rand_trunc', 
+                data_filling='repeatpad',
+                audio_cfg=self.model_cfg['audio_cfg'],
+                require_grad=audio_waveform.requires_grad
+            )
+            
+            audio_input.append(temp_dict)
+            
+        audio_embed = self.model.get_audio_embedding(audio_input, normalize)
+        
+        if not use_tensor:
+            audio_embed = audio_embed.detach().cpu().numpy()
+
+        return audio_embed
+
+    def get_text_embedding(self, x, normalize=False, use_tensor=True):
+        self.model.eval()
+        if isinstance(x, str):
+            x = [x]
+            
+        token_data = self.tokenizer(x)
+        sequence_lengths = (torch.ne(token_data['attention_mask'], 0).sum(-1) - 1)
+        setence_embeds = self.model.get_text_embedding(token_data, normalize)
+        word_embeds = self.model.get_word_embedding(token_data)
+        
+        if not use_tensor:
+            setence_embeds = setence_embeds.detach().cpu().numpy()
+            word_embeds = word_embeds.detach().cpu().numpy()
+            
+        return setence_embeds, word_embeds, sequence_lengths
+        
+    def get_clap_score(self, text, audio, sr=16000):
+        setence_embeds, word_embeds, sequence_lengths = self.get_text_embedding(text, normalize=True)
+        audio_embeds = self.get_audio_embedding(audio, sr=16000, normalize=True)
+        
+        clap_score = torch.nn.functional.cosine_similarity(setence_embeds, audio_embeds, dim=-1)
+
+        return clap_score

CLAP/data.py

@@ -0,0 +1,137 @@
+import torch
+import torchaudio
+import torchvision
+import numpy as np
+from contextlib import suppress
+import torch.nn.functional as F
+
+def int16_to_float32(x):
+    return (x / 32767.0).astype(np.float32)
+
+def float32_to_int16(x):
+    x = np.clip(x, a_min=-1., a_max=1.)
+    return (x * 32767.).astype(np.int16)
+
+def get_mel(audio_data, audio_cfg):
+    # mel shape: (n_mels, T)
+    mel_tf = torchaudio.transforms.MelSpectrogram(
+        sample_rate=audio_cfg['sample_rate'],
+        n_fft=audio_cfg['window_size'],
+        win_length=audio_cfg['window_size'],
+        hop_length=audio_cfg['hop_size'],
+        center=True,
+        pad_mode="reflect",
+        power=2.0,
+        norm=None,
+        onesided=True,
+        n_mels=audio_cfg['mel_bins'],
+        f_min=audio_cfg['fmin'],
+        f_max=audio_cfg['fmax']
+    ).to(audio_data.device)
+    
+    mel = mel_tf(audio_data)
+    
+    # we use log mel spectrogram as input
+    mel = torchaudio.transforms.AmplitudeToDB(top_db=None)(mel)
+    return mel.T  # (T, n_mels)
+
+def get_audio_features(sample, audio_data, max_len, data_truncating, data_filling, audio_cfg, require_grad=False):
+    """
+    Calculate and add audio features to sample.
+    Sample: a dict containing all the data of current sample.
+    audio_data: a tensor of shape (T) containing audio data.
+    max_len: the maximum length of audio data.
+    data_truncating: the method of truncating data.
+    data_filling: the method of filling data.
+    audio_cfg: a dict containing audio configuration. Comes from model_cfg['audio_cfg'].
+    require_grad: whether to require gradient for audio data.
+        This is useful when we want to apply gradient-based classifier-guidance.
+    """
+    grad_fn = suppress if require_grad else torch.no_grad
+    with grad_fn():
+        if len(audio_data) > max_len:
+            if data_truncating == "rand_trunc":
+                longer = torch.tensor([True])
+            elif data_truncating == "fusion":
+                # fusion
+                mel = get_mel(audio_data, audio_cfg)
+                # split to three parts
+                chunk_frames = max_len // audio_cfg['hop_size'] + 1  # the +1 related to how the spectrogram is computed
+                total_frames = mel.shape[0]
+                if chunk_frames == total_frames:
+                    # there is a corner case where the audio length is
+                    # larger than max_len but smaller than max_len+hop_size.
+                    # In this case, we just use the whole audio.
+                    mel_fusion = torch.stack([mel, mel, mel, mel], dim=0)
+                    sample["mel_fusion"] = mel_fusion
+                    longer = torch.tensor([False])
+                else:
+                    ranges = np.array_split(list(range(0, total_frames - chunk_frames + 1)), 3)
+                    
+                    if len(ranges[1]) == 0:
+                        # if the audio is too short, we just use the first chunk
+                        ranges[1] = [0]
+                    if len(ranges[2]) == 0:
+                        # if the audio is too short, we just use the first chunk
+                        ranges[2] = [0]
+                    # randomly choose index for each part
+                    idx_front = np.random.choice(ranges[0])
+                    idx_middle = np.random.choice(ranges[1])
+                    idx_back = np.random.choice(ranges[2])
+                    # select mel
+                    mel_chunk_front = mel[idx_front:idx_front + chunk_frames, :]
+                    mel_chunk_middle = mel[idx_middle:idx_middle + chunk_frames, :]
+                    mel_chunk_back = mel[idx_back:idx_back + chunk_frames, :]
+
+                    # shrink the mel
+                    mel_shrink = torchvision.transforms.Resize(size=[chunk_frames, audio_cfg['mel_bins']])(mel[None])[0]
+
+                    # stack
+                    mel_fusion = torch.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back], dim=0)
+                    sample["mel_fusion"] = mel_fusion
+                    longer = torch.tensor([True])
+            else:
+                raise NotImplementedError(
+                    f"data_truncating {data_truncating} not implemented"
+                )
+            # random crop to max_len (for compatibility)
+            overflow = len(audio_data) - max_len
+            idx = np.random.randint(0, overflow + 1)
+            audio_data = audio_data[idx: idx + max_len]
+
+        else:  # padding if too short
+            if len(audio_data) < max_len:  # do nothing if equal
+                if data_filling == "repeatpad":
+                    n_repeat = int(max_len / len(audio_data))
+                    audio_data = audio_data.repeat(n_repeat)
+                    
+                    audio_data = F.pad(
+                        audio_data,
+                        (0, max_len - len(audio_data)),
+                        mode="constant",
+                        value=0,
+                    )
+                elif data_filling == "pad":
+                    audio_data = F.pad(
+                        audio_data,
+                        (0, max_len - len(audio_data)),
+                        mode="constant",
+                        value=0,
+                    )
+                elif data_filling == "repeat":
+                    n_repeat = int(max_len / len(audio_data))
+                    audio_data = audio_data.repeat(n_repeat + 1)[:max_len]
+                else:
+                    raise NotImplementedError(
+                        f"data_filling {data_filling} not implemented"
+                    )
+            if data_truncating == 'fusion':
+                mel = get_mel(audio_data, audio_cfg)
+                mel_fusion = torch.stack([mel, mel, mel, mel], dim=0)
+                sample["mel_fusion"] = mel_fusion
+            longer = torch.tensor([False])
+
+    sample["longer"] = longer
+    sample["waveform"] = audio_data
+
+    return sample

CLAP/htsat.py

@@ -0,0 +1,894 @@
+# Ke Chen
+# knutchen@ucsd.edu
+# HTS-AT: A HIERARCHICAL TOKEN-SEMANTIC AUDIO TRANSFORMER FOR SOUND CLASSIFICATION AND DETECTION
+# Some layers designed on the model
+# below codes are based and referred from https://github.com/microsoft/Swin-Transformer
+# Swin Transformer for Computer Vision: https://arxiv.org/pdf/2103.14030.pdf
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from itertools import repeat
+import collections.abc
+import math
+import warnings
+
+from torch.nn.init import _calculate_fan_in_and_fan_out
+import torch.utils.checkpoint as checkpoint
+
+import random
+
+from torchlibrosa.stft import Spectrogram, LogmelFilterBank
+from torchlibrosa.augmentation import SpecAugmentation
+
+from itertools import repeat
+
+def interpolate(x, ratio):
+    """Interpolate data in time domain. This is used to compensate the
+    resolution reduction in downsampling of a CNN.
+
+    Args:
+      x: (batch_size, time_steps, classes_num)
+      ratio: int, ratio to interpolate
+    Returns:
+      upsampled: (batch_size, time_steps * ratio, classes_num)
+    """
+    (batch_size, time_steps, classes_num) = x.shape
+    upsampled = x[:, :, None, :].repeat(1, 1, ratio, 1)
+    upsampled = upsampled.reshape(batch_size, time_steps * ratio, classes_num)
+    return upsampled
+
+def do_mixup(x, mixup_lambda):
+    """
+    Args:
+      x: (batch_size , ...)
+      mixup_lambda: (batch_size,)
+    Returns:
+      out: (batch_size, ...)
+    """
+    out = (
+            x.transpose(0, -1) * mixup_lambda
+            + torch.flip(x, dims=[0]).transpose(0, -1) * (1 - mixup_lambda)
+    ).transpose(0, -1)
+    return out
+
+# from PyTorch internals
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable):
+            return x
+        return tuple(repeat(x, n))
+    return parse
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = _ntuple
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+class PatchEmbed(nn.Module):
+    """ 2D Image to Patch Embedding
+    """
+    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.Conv2d(in_chans, embed_dim, 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.proj(x)
+    
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+class Mlp(nn.Module):
+    """ MLP as used in Vision Transformer, MLP-Mixer and related networks
+    """
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    r"""Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+def variance_scaling_(tensor, scale=1.0, mode='fan_in', distribution='normal'):
+    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
+    if mode == 'fan_in':
+        denom = fan_in
+    elif mode == 'fan_out':
+        denom = fan_out
+    elif mode == 'fan_avg':
+        denom = (fan_in + fan_out) / 2
+
+    variance = scale / denom
+
+    if distribution == "truncated_normal":
+        # constant is stddev of standard normal truncated to (-2, 2)
+        trunc_normal_(tensor, std=math.sqrt(variance) / .87962566103423978)
+    elif distribution == "normal":
+        tensor.normal_(std=math.sqrt(variance))
+    elif distribution == "uniform":
+        bound = math.sqrt(3 * variance)
+        tensor.uniform_(-bound, bound)
+    else:
+        raise ValueError(f"invalid distribution {distribution}")
+
+
+def lecun_normal_(tensor):
+    variance_scaling_(tensor, mode='fan_in', distribution='truncated_normal')
+
+def window_partition(x, window_size):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows
+
+
+def window_reverse(windows, window_size, H, W):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+    Returns:
+        x: (B, H, W, C)
+    """
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+
+class WindowAttention(nn.Module):
+    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+    Args:
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        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
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
+
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size[0])
+        coords_w = torch.arange(self.window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        trunc_normal_(self.relative_position_bias_table, std=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        """
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x, attn
+
+    def extra_repr(self):
+        return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'
+
+
+# We use the model based on Swintransformer Block, therefore we can use the swin-transformer pretrained model
+class SwinTransformerBlock(nn.Module):
+    r""" Swin Transformer Block.
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resulotion.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        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, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, norm_before_mlp='ln'):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        self.norm_before_mlp = norm_before_mlp
+        if min(self.input_resolution) <= self.window_size:
+            # if window size is larger than input resolution, we don't partition windows
+            self.shift_size = 0
+            self.window_size = min(self.input_resolution)
+        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        if self.norm_before_mlp == 'ln':
+            self.norm2 = nn.LayerNorm(dim)
+        elif self.norm_before_mlp == 'bn':
+            self.norm2 = lambda x: nn.BatchNorm1d(dim)(x.transpose(1, 2)).transpose(1, 2)
+        else:
+            raise NotImplementedError
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if self.shift_size > 0:
+            # calculate attention mask for SW-MSA
+            H, W = self.input_resolution
+            img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
+            h_slices = (slice(0, -self.window_size),
+                        slice(-self.window_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            w_slices = (slice(0, -self.window_size),
+                        slice(-self.window_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            cnt = 0
+            for h in h_slices:
+                for w in w_slices:
+                    img_mask[:, h, w, :] = cnt
+                    cnt += 1
+
+            mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+            mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+        else:
+            attn_mask = None
+
+        self.register_buffer("attn_mask", attn_mask)
+
+    def forward(self, x):
+        # pdb.set_trace()
+        H, W = self.input_resolution
+        # print("H: ", H)
+        # print("W: ", W)
+        # pdb.set_trace()
+        B, L, C = x.shape
+        # assert L == H * W, "input feature has wrong size"
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+        else:
+            shifted_x = x
+
+        # partition windows
+        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
+
+        # W-MSA/SW-MSA
+        attn_windows, attn = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+        shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            x = shifted_x
+        x = x.view(B, H * W, C)
+
+        # FFN
+        x = shortcut + self.drop_path(x)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x, attn
+
+    def extra_repr(self):
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
+               f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
+
+
+
+class PatchMerging(nn.Module):
+    r""" Patch Merging Layer.
+    Args:
+        input_resolution (tuple[int]): Resolution of input feature.
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, x):
+        """
+        x: B, H*W, C
+        """
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."
+
+        x = x.view(B, H, W, C)
+
+        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
+        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
+        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
+        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
+        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
+        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
+
+        x = self.norm(x)
+        x = self.reduction(x)
+
+        return x
+
+    def extra_repr(self):
+        return f"input_resolution={self.input_resolution}, dim={self.dim}"
+
+
+class BasicLayer(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, dim=dim, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        attns = []
+        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 self.downsample is not None:
+            x = self.downsample(x)
+        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}"
+
+
+# The Core of HTSAT
+class HTSAT_Swin_Transformer(nn.Module):
+    r"""HTSAT 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
+        config (module): The configuration Module from config.py
+    """
+
+    def __init__(self, spec_size=256, patch_size=4, patch_stride=(4,4), 
+                in_chans=1, num_classes=527,
+                 embed_dim=96, depths=[2, 2, 6, 2], 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', config = None, **kwargs):
+        super(HTSAT_Swin_Transformer, self).__init__()
+
+        self.config = config
+        self.spec_size = spec_size 
+        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.config.mel_bins
+        window = 'hann'
+        center = True
+        pad_mode = 'reflect'
+        ref = 1.0
+        amin = 1e-10
+        top_db = None
+        self.interpolate_ratio = 32     # Downsampled ratio
+        # Spectrogram extractor
+        self.spectrogram_extractor = Spectrogram(n_fft=config.window_size, hop_length=config.hop_size, 
+            win_length=config.window_size, window=window, center=center, pad_mode=pad_mode, 
+            freeze_parameters=True)
+        # Logmel feature extractor
+        self.logmel_extractor = LogmelFilterBank(sr=config.sample_rate, n_fft=config.window_size, 
+            n_mels=config.mel_bins, fmin=config.fmin, fmax=config.fmax, ref=ref, amin=amin, top_db=top_db, 
+            freeze_parameters=True)
+        # Spec augmenter
+        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2, 
+            freq_drop_width=8, freq_stripes_num=2) # 2 2
+        self.bn0 = nn.BatchNorm2d(self.config.mel_bins)
+
+
+        # split spctrogram into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            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,
+            )
+
+        num_patches = self.patch_embed.num_patches
+        patches_resolution = self.patch_embed.grid_size
+        self.patches_resolution = patches_resolution
+
+        # absolute position embedding
+        if self.ape:
+            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, self.embed_dim))
+            trunc_normal_(self.absolute_pos_embed, std=.02)
+
+        self.pos_drop = nn.Dropout(p=self.drop_rate)
+
+        # 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()
+        for i_layer in range(self.num_layers):
+            layer = BasicLayer(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=PatchMerging 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.norm = self.norm_layer(self.num_features)
+        self.avgpool = nn.AdaptiveAvgPool1d(1)
+        self.maxpool = nn.AdaptiveMaxPool1d(1)
+        
+        SF = self.spec_size // (2 ** (len(self.depths) - 1)) // self.patch_stride[0] // self.freq_ratio
+        self.tscam_conv = nn.Conv2d(
+            in_channels = self.num_features,
+            out_channels = self.num_classes,
+            kernel_size = (SF,3),
+            padding = (0,1)
+        )
+        self.head = nn.Linear(num_classes, num_classes)
+
+        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_features(self, x):
+        # A deprecated optimization for using a hierarchical output from different blocks
+
+        frames_num = x.shape[2]
+        x = self.patch_embed(x)
+        if self.ape:
+            x = x + self.absolute_pos_embed
+        x = self.pos_drop(x)
+        for i, layer in enumerate(self.layers):
+            x, attn = layer(x)
+        # for x
+        x = self.norm(x)
+        B, N, C = x.shape
+        SF = frames_num // (2 ** (len(self.depths) - 1)) // self.patch_stride[0]
+        ST = frames_num // (2 ** (len(self.depths) - 1)) // self.patch_stride[1]
+        x = x.permute(0,2,1).contiguous().reshape(B, C, SF, ST)
+        B, C, F, T = x.shape
+        # group 2D CNN
+        c_freq_bin = F // self.freq_ratio
+        x = x.reshape(B, C, F // c_freq_bin, c_freq_bin, T)
+        x = x.permute(0,1,3,2,4).contiguous().reshape(B, C, c_freq_bin, -1)
+        # get latent_output
+        fine_grained_latent_output = torch.mean(x, dim = 2)
+        fine_grained_latent_output = interpolate(fine_grained_latent_output.permute(0,2,1).contiguous(), 8 * self.patch_stride[1]) 
+        
+        latent_output = self.avgpool(torch.flatten(x,2))
+        latent_output = torch.flatten(latent_output, 1)
+
+        # display the attention map, if needed
+
+        x = self.tscam_conv(x)
+        x = torch.flatten(x, 2) # B, C, T
+ 
+        fpx = interpolate(torch.sigmoid(x).permute(0,2,1).contiguous(), 8 * self.patch_stride[1]) 
+            
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+
+        output_dict = {
+            'framewise_output': fpx, # already sigmoided
+            'clipwise_output': torch.sigmoid(x),
+            'fine_grained_embedding': fine_grained_latent_output,
+            'embedding': latent_output
+        }
+
+        return output_dict
+
+    def crop_wav(self, x, crop_size, spe_pos = None):
+        time_steps = x.shape[2]
+        tx = torch.zeros(x.shape[0], x.shape[1], crop_size, x.shape[3]).to(x.device)
+        for i in range(len(x)):
+            if spe_pos is None:
+                crop_pos = random.randint(0, time_steps - crop_size - 1)
+            else:
+                crop_pos = spe_pos
+            tx[i][0] = x[i, 0, crop_pos:crop_pos + crop_size,:]
+        return tx
+
+    # Reshape the wavform to a img size, if you want to use the pretrained swin transformer model
+    def reshape_wav2img(self, x):
+        B, C, T, F = x.shape
+        target_T = int(self.spec_size * self.freq_ratio)
+        target_F = self.spec_size // self.freq_ratio
+        assert T <= target_T and F <= target_F, "the wav size should less than or equal to the swin input size"
+        # to avoid bicubic zero error
+        if T < target_T:
+            x = nn.functional.interpolate(x, (target_T, x.shape[3]), mode="bicubic", align_corners=True)
+        if F < target_F:
+            x = nn.functional.interpolate(x, (x.shape[2], target_F), mode="bicubic", align_corners=True)
+        x = x.permute(0,1,3,2).contiguous()
+        x = x.reshape(x.shape[0], x.shape[1], x.shape[2], self.freq_ratio, x.shape[3] // self.freq_ratio)
+        # print(x.shape)
+        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])
+        return x
+    
+    # Repeat the wavform to a img size, if you want to use the pretrained swin transformer model
+    def repeat_wat2img(self, x, cur_pos):
+        B, C, T, F = x.shape
+        target_T = int(self.spec_size * self.freq_ratio)
+        target_F = self.spec_size // self.freq_ratio
+        assert T <= target_T and F <= target_F, "the wav size should less than or equal to the swin input size"
+        # to avoid bicubic zero error
+        if T < target_T:
+            x = nn.functional.interpolate(x, (target_T, x.shape[3]), mode="bicubic", align_corners=True)
+        if F < target_F:
+            x = nn.functional.interpolate(x, (x.shape[2], target_F), mode="bicubic", align_corners=True)  
+        x = x.permute(0,1,3,2).contiguous() # B C F T
+        x = x[:,:,:,cur_pos:cur_pos + self.spec_size]
+        x = x.repeat(repeats = (1,1,4,1))
+        return x
+
+    def forward(self, x: torch.Tensor, mixup_lambda = None, infer_mode = False, device=None):# out_feat_keys: List[str] = None):
+            
+        x = x["waveform"].to(device=device, non_blocking=True)
+        x = self.spectrogram_extractor(x)   # (batch_size, 1, time_steps, freq_bins)
+        x = self.logmel_extractor(x)    # (batch_size, 1, time_steps, mel_bins)
+        x = x.transpose(1, 3)
+        x = self.bn0(x)
+        x = x.transpose(1, 3)
+        if self.training:
+            x = self.spec_augmenter(x)
+
+        if self.training and mixup_lambda is not None:
+            x = do_mixup(x, mixup_lambda)
+            
+        x = self.reshape_wav2img(x)
+        output_dict = self.forward_features(x)
+
+        return output_dict
+
+def create_htsat_model(audio_cfg):
+    try:
+        assert audio_cfg.model_name in ["tiny", "base", "large"], "model name for HTS-AT is wrong!"
+        if audio_cfg.model_name == "tiny":
+            model = HTSAT_Swin_Transformer(
+                spec_size=256,
+                patch_size=4,
+                patch_stride=(4,4),
+                num_classes=audio_cfg.class_num,
+                embed_dim=96,
+                depths=[2,2,6,2],
+                num_heads=[4,8,16,32],
+                window_size=8,
+                config = audio_cfg
+            )
+        elif audio_cfg.model_name == "base":
+            model = HTSAT_Swin_Transformer(
+                spec_size=256,
+                patch_size=4,
+                patch_stride=(4,4),
+                num_classes=audio_cfg.class_num,
+                embed_dim=128,
+                depths=[2,2,12,2],
+                num_heads=[4,8,16,32],
+                window_size=8,
+                config = audio_cfg
+            )
+        elif audio_cfg.model_name == "large":
+            model = HTSAT_Swin_Transformer(
+                spec_size=256,
+                patch_size=4,
+                patch_stride=(4,4),
+                num_classes=audio_cfg.class_num,
+                embed_dim=256,
+                depths=[2,2,12,2],
+                num_heads=[4,8,16,32],
+                window_size=8,
+                config = audio_cfg
+            )
+        return model
+    except:
+        raise RuntimeError(f'Import Model for {audio_cfg.model_name} not found, or the audio cfg parameters are not enough.')
+        

CLAP/model_configs/HTSAT-base.json

@@ -0,0 +1,23 @@
+{
+    "embed_dim": 1024,
+    "audio_cfg": {
+        "audio_length": 1024,
+        "clip_samples": 480000,
+        "mel_bins": 64,
+        "sample_rate": 48000,
+        "window_size": 1024,
+        "hop_size": 480,
+        "fmin": 50,
+        "fmax": 14000,
+        "class_num": 527,
+        "model_type": "HTSAT",
+        "model_name": "base"
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

CLAP/model_configs/HTSAT-tiny.json

@@ -0,0 +1,23 @@
+{
+    "embed_dim": 768,
+    "audio_cfg": {
+        "audio_length": 1024,
+        "clip_samples": 480000,
+        "mel_bins": 64,
+        "sample_rate": 48000,
+        "window_size": 1024,
+        "hop_size": 480,
+        "fmin": 50,
+        "fmax": 14000,
+        "class_num": 527,
+        "model_type": "HTSAT",
+        "model_name": "tiny"
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

HiFiGAN/hifigan_model.py

@@ -0,0 +1,177 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+LRELU_SLOPE = 0.1
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs):
+        super(AttrDict, self).__init__(*args, **kwargs)
+        self.__dict__ = self
+
+class ResBlock(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock, self).__init__()
+        self.h = h
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=get_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
+        self.convs2.apply(init_weights)
+
+    def forward(self, x):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class Generator(torch.nn.Module):
+    def __init__(self, h):
+        super(Generator, self).__init__()
+        self.h = h
+        self.num_kernels = len(h.resblock_kernel_sizes)
+        self.num_upsamples = len(h.upsample_rates)
+        self.conv_pre = weight_norm(
+            Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)
+        )
+        resblock = ResBlock
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        h.upsample_initial_channel // (2**i),
+                        h.upsample_initial_channel // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = h.upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)
+            ):
+                self.resblocks.append(resblock(h, ch, k, d))
+
+        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
+        self.ups.apply(init_weights)
+        self.conv_post.apply(init_weights)
+
+    def forward(self, x):
+        x = self.conv_pre(x)
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+        remove_weight_norm(self.conv_pre)
+        remove_weight_norm(self.conv_post)

HiFiGAN/inference.py

@@ -0,0 +1,32 @@
+import os
+import json
+import torch
+
+from .hifigan_model import AttrDict, Generator
+
+def torch_version_orig_mod_remove(state_dict):
+    new_state_dict = {}
+    new_state_dict["generator"] = {}
+    for key in state_dict["generator"].keys():
+        if "_orig_mod." in key:
+            new_state_dict["generator"][key.replace("_orig_mod.", "")] = state_dict[
+                "generator"
+            ][key]
+        else:
+            new_state_dict["generator"][key] = state_dict["generator"][key]
+    return new_state_dict
+
+def get_vocoder(sr, ckpt_path):
+
+    with open(os.path.join(ckpt_path, "hifigan_16k_64bins.json"), "r") as f:
+        config = json.load(f)
+    config = AttrDict(config)
+    vocoder = Generator(config)
+
+    ckpt = torch.load(os.path.join(ckpt_path, "hifigan_16k_64bins.ckpt"), map_location='cpu')
+    ckpt = torch_version_orig_mod_remove(ckpt)
+    vocoder.load_state_dict(ckpt["generator"])
+    vocoder.eval()
+    vocoder.remove_weight_norm()
+    
+    return vocoder

generator.py

@@ -0,0 +1,439 @@
+import torch
+from torch import nn
+from torch.nn import init
+from torch.nn import functional as F
+
+class SpectralNorm:
+    def __init__(self, name):
+        self.name = name
+
+    def compute_weight(self, module):
+        weight = getattr(module, self.name + '_orig')
+        u = getattr(module, self.name + '_u')
+        size = weight.size()
+        weight_mat = weight.contiguous().view(size[0], -1)
+        with torch.no_grad():
+            v = weight_mat.t() @ u
+            v = v / v.norm()
+            u = weight_mat @ v
+            u = u / u.norm()
+        sigma = u @ weight_mat @ v
+        weight_sn = weight / sigma
+
+        return weight_sn, u, sigma
+
+    @staticmethod
+    def apply(module, name):
+        fn = SpectralNorm(name)
+
+        weight = getattr(module, name)
+        del module._parameters[name]
+        module.register_parameter(name + '_orig', weight)
+        input_size = weight.size(0)
+        u = weight.new_empty(input_size).normal_()
+        module.register_buffer(name, weight)
+        module.register_buffer(name + '_u', u)
+        module.register_buffer(name + '_sv', torch.ones(1).squeeze())
+
+        module.register_forward_pre_hook(fn)
+
+        return fn
+
+    def __call__(self, module, input):
+        weight_sn, u, sigma = self.compute_weight(module)
+        setattr(module, self.name, weight_sn)
+        setattr(module, self.name + '_u', u)
+        setattr(module, self.name + '_sv', sigma)
+
+def spectral_norm(module, name='weight'):
+    SpectralNorm.apply(module, name)
+    return module
+
+def spectral_init(module, gain=1):
+    init.xavier_uniform_(module.weight, gain)
+    if module.bias is not None:
+        module.bias.data.zero_()
+    return spectral_norm(module)
+
+class ConditionalNorm(nn.Module):
+    def __init__(self, in_channel, condition_dim):
+        super().__init__()
+
+        self.bn = nn.BatchNorm2d(in_channel, affine=False)
+        self.linear1 = nn.Linear(condition_dim, in_channel)
+        self.linear2 = nn.Linear(condition_dim, in_channel)
+
+    def forward(self, input, condition):
+        out = self.bn(input)
+        gamma, beta = self.linear1(condition), self.linear2(condition)
+        gamma = gamma.unsqueeze(2).unsqueeze(3)
+        beta = beta.unsqueeze(2).unsqueeze(3)
+        out = gamma * out + beta
+
+        return out
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channel, out_channel, kernel_size=[3, 3],
+                 padding=1, stride=1, condition_dim=None, bn=True,
+                 activation=F.relu, upsample=True, downsample=False):
+        super().__init__()
+
+        gain = 2 ** 0.5
+
+        self.conv1 = spectral_init(nn.Conv2d(in_channel, out_channel,
+                                             kernel_size, stride, padding,
+                                             bias=False if bn else True),
+                                   gain=gain)
+        self.conv2 = spectral_init(nn.Conv2d(out_channel, out_channel,
+                                             kernel_size, stride, padding,
+                                             bias=False if bn else True),
+                                   gain=gain)
+
+        self.skip_proj = False
+        if in_channel != out_channel or upsample or downsample:
+            self.conv_skip = spectral_init(nn.Conv2d(in_channel, out_channel,
+                                                     1, 1, 0))
+            self.skip_proj = True
+
+        self.upsample = upsample
+        self.downsample = downsample
+        self.activation = activation
+        self.bn = bn
+
+        if bn:
+            self.norm1 = ConditionalNorm(in_channel, condition_dim)
+            self.norm2 = ConditionalNorm(out_channel, condition_dim)
+
+    def forward(self, input, condition=None, condition1=None):
+        out = input
+
+        if self.bn:
+            out = self.norm1(out, condition)
+        out = self.activation(out)
+        if self.upsample:
+            out = F.interpolate(out, scale_factor=2, mode='nearest')
+        out = self.conv1(out)
+        if self.bn:
+            out = self.norm2(out, condition)
+        out = self.activation(out)
+        out = self.conv2(out)
+
+        if self.downsample:
+            out = F.avg_pool2d(out, 2)
+
+        if self.skip_proj:
+            skip = input
+            if self.upsample:
+                skip = F.interpolate(skip, scale_factor=2, mode='nearest')
+            skip = self.conv_skip(skip)
+            if self.downsample:
+                skip = F.avg_pool2d(skip, 2)
+        else:
+            skip = input
+
+        return out + skip
+    
+class SelfAttention(nn.Module):
+    def __init__(self, in_channel, embed_dim, gain=2 ** 0.5):
+        super().__init__()
+
+        self.query = spectral_init(nn.Conv1d(in_channel, embed_dim, 1),
+                                   gain=gain)
+        self.key = spectral_init(nn.Conv1d(in_channel, embed_dim, 1),
+                                 gain=gain)
+        self.value = spectral_init(nn.Conv1d(in_channel, in_channel, 1),
+                                   gain=gain)
+
+        self.gamma = nn.Parameter(torch.tensor(0.0))
+
+    def forward(self, input): # [bsz, channel, freq, time]
+        shape = input.shape
+        flatten = input.view(shape[0], shape[1], -1) # [bsz, channel, freq*time]
+        query = self.query(flatten).permute(0, 2, 1)
+        key = self.key(flatten)
+        value = self.value(flatten)
+        query_key = torch.bmm(query, key) # [bsz, freq*time, freq*time]
+        attention_map = F.softmax(query_key, 1)
+        out = torch.bmm(value, attention_map)
+        out = out.view(*shape)
+        out = self.gamma * out + input
+
+        return (out, attention_map)
+
+class CrossAttention(nn.Module):
+    def __init__(self, in_channel, cond_channel, embed_dim, gain=2 ** 0.5):
+        super().__init__()
+
+        self.key = spectral_init(nn.Conv1d(cond_channel, embed_dim, 1),
+                                   gain=gain)
+        self.value = spectral_init(nn.Conv1d(cond_channel, in_channel, 1),
+                                   gain=gain)
+        self.query = spectral_init(nn.Conv1d(in_channel, embed_dim, 1),
+                                   gain=gain)
+
+        self.gamma = nn.Parameter(torch.tensor(0.0))
+
+    def forward(self, input, condition, sequence_lengths=None):
+        # input : mel [bsz, channel, freq, time] or sentence [bsz, channel]
+        # condition : sentence [bsz, channel] or word [bsz, word_num, channel]
+        input_shape = input.shape
+        if len(input.shape) == 4: # mel [bsz, channel, freq, time]
+            batch_size, c, w, h = input.shape
+            num = w * h
+            x = input.reshape([batch_size, c, num]) #[bsz, channel, input_num]
+        elif len(input.shape) == 2: # sentence [bsz, channel]
+            batch_size, c = input.shape
+            num = 1
+            x = input.unsqueeze(2) # [bsz, channel, input_num]
+
+        if len(condition.shape) == 2: # sentence [bsz, channel]
+            condition = condition.unsqueeze(2) # [bsz, channel, cond_num]
+        else: # word [bsz, word_num, channel]
+            condition = condition.permute(0, 2, 1) # [bsz, channel, cond_num]
+
+        query = self.query(x).permute(0, 2, 1) # [bsz, input_num, channel]
+        key = self.key(condition) # [bsz, channel, cond_num]
+        value = self.value(condition).permute(0, 2, 1) # [bsz, cond_num, channel]
+        attention_map = torch.bmm(query, key) # [bsz, input_num, cond_num]
+
+        if sequence_lengths is not None: # condition is word embedding
+            total_len = condition.shape[2]
+
+            mask = torch.tile(torch.arange(total_len), [batch_size, num, 1]).to(condition.device)
+            for i in range(batch_size):
+                sequence_lengths_i = sequence_lengths[i]
+                mask[i,:,:] = mask[i,:,:] >= sequence_lengths_i.item()
+            attention_map = attention_map + mask * (-1e9)
+        
+        attention_map = F.softmax(attention_map, dim=-1) # [bsz, input_num, cond_num]
+        out = torch.bmm(attention_map, value).permute(0, 2, 1) # [bsz, input_num, channel]
+        out = out.permute(0, 2, 1).reshape(input_shape).squeeze()
+        out = self.gamma * out + input
+
+        return out, attention_map
+
+class Spec_Attention(nn.Module):
+    def __init__(self, in_channel, cond_channel=None, embed_dim=64, gain=2 ** 0.5):
+        super().__init__()
+        if cond_channel is None:
+            cond_channel = in_channel
+
+        self.f_query = spectral_init(nn.Conv1d(in_channel, embed_dim, 1),
+                                   gain=gain)
+        self.t_key = spectral_init(nn.Conv1d(cond_channel, embed_dim, 1),
+                                   gain=gain)
+
+        self.t_query = spectral_init(nn.Conv1d(in_channel, embed_dim, 1),
+                                   gain=gain)
+        self.f_key = spectral_init(nn.Conv1d(cond_channel, embed_dim, 1),
+                                   gain=gain)
+
+        self.value = spectral_init(nn.Conv1d(cond_channel, in_channel, 1),
+                                   gain=gain)
+
+        self.gamma = nn.Parameter(torch.tensor(0.0))
+
+    def forward(self, input, condition=None, sequence_lengths=None):
+        # input : mel [bsz, channel, freq, time]
+        # condition : sentence [bsz, channel] or word [bsz, word_num, channel]
+        
+        batch_size, c, f, t = input.shape
+
+        freq_embedding = input.mean(dim=3) # [bsz, channel, freq]
+        time_embedding = input.mean(dim=2) # [bsz, channel, time]
+
+        if condition is not None:
+            if len(condition.shape) == 2: # sentence [bsz, channel]
+                condition = condition.unsqueeze(2) # [bsz, channel, 1]
+            else: # word [bsz, word_num, channel]
+                condition = condition.permute(0, 2, 1) # [bsz, channel, cond_num]
+            t_condition = condition
+            f_condition = condition
+        else:
+            t_condition = time_embedding
+            f_condition = freq_embedding
+
+        f_query = self.f_query(freq_embedding).permute(0, 2, 1) # [bsz, freq, channel]
+        t_key = self.t_key(t_condition) # [bsz, channel, time] or [bsz, channel, cond_num]
+        freq_cond_map = torch.bmm(f_query, t_key) # [bsz, freq, time] or [bsz, freq, cond_num]
+
+        t_query = self.t_query(time_embedding).permute(0, 2, 1) # [bsz, time, channel]
+        f_key = self.f_key(f_condition) # [bsz, channel, freq] or [bsz, channel, cond_num]
+        time_cond_map = torch.bmm(t_query, f_key) # [bsz, time, freq] or [bsz, time, cond_num]
+
+        if sequence_lengths is not None: # condition is word embedding
+            total_len = condition.shape[2]
+
+            mask = torch.arange(total_len, device=condition.device)[None, None, :]
+            mask = mask >= sequence_lengths[:, None, None]
+
+            freq_cond_map = freq_cond_map + mask * (-1e9)
+            time_cond_map = time_cond_map + mask * (-1e9)
+
+        freq_cond_map = F.softmax(freq_cond_map, dim=-1) # [bsz, freq, time] or [bsz, freq, cond_num]
+        time_cond_map = F.softmax(time_cond_map, dim=-1) # [bsz, time, freq] or [bsz, time, cond_num]
+        
+        if condition is None:
+            freq_time_embedding = input.reshape([batch_size, c, f*t]) # [bsz, channel, freq*time]
+            weight_map = torch.add(freq_cond_map, time_cond_map.permute(0, 2, 1)).reshape([batch_size, f*t]).unsqueeze(-1) # [bsz, freq*time, 1]
+            value = self.value(freq_time_embedding).permute(0, 2, 1) # [bsz, freq*time, channel]
+            out = torch.mul(value, weight_map).permute(0, 2, 1).reshape(batch_size, c, f, t) # [bsz, channel, freq, time]
+        else:
+            freq_cond_map = torch.tile(freq_cond_map.unsqueeze(2), [1, 1, t, 1]) # [bsz, freq, time, cond_num]
+            time_cond_map = torch.tile(time_cond_map.unsqueeze(1), [1, f, 1, 1]) # [bsz, freq, time, cond_num]
+            weight_map = torch.add(freq_cond_map, time_cond_map).reshape([batch_size, f*t, -1]) # [bsz, freq*time, cond_num]
+            value = self.value(condition).permute(0, 2, 1) # [bsz, cond_num, channel]
+            out = torch.bmm(weight_map, value).permute(0, 2, 1).reshape(batch_size, c, f, t) # [bsz, channel, freq, time]
+
+        out = self.gamma * out + input
+
+        return out, weight_map
+
+class Multi_Triple_Attention(nn.Module):
+    def __init__(self, in_channel, sentence_embed_dim=768, word_embed_dim=768, embed_dim=64, reverse=False, gain=2 ** 0.5, n_heads=2, attention_list="self,word,sentence", spec_attention=False):
+        super().__init__()
+        self.reverse = reverse
+        self.n_heads = n_heads
+        self.attention_list = attention_list.split(",")
+        
+        if "self" in self.attention_list:
+            if spec_attention:
+                self.self_attention_modules = nn.ModuleList([Spec_Attention(in_channel, embed_dim=embed_dim) for _ in range(self.n_heads)])
+            else:
+                self.self_attention_modules = nn.ModuleList([SelfAttention(in_channel, embed_dim=embed_dim) for _ in range(self.n_heads)])
+        
+        if "word" in self.attention_list:
+            if spec_attention:
+                self.cross_attention_for_word_modules = nn.ModuleList([Spec_Attention(in_channel, cond_channel=word_embed_dim, embed_dim=embed_dim) for _ in range(self.n_heads)])
+            else:
+                self.cross_attention_for_word_modules = nn.ModuleList([CrossAttention(in_channel, cond_channel=word_embed_dim, embed_dim=embed_dim) for _ in range(self.n_heads)])
+        
+        if "sentence" in self.attention_list:
+            if spec_attention:
+                self.cross_attention_for_sent_modules = nn.ModuleList([Spec_Attention(in_channel, cond_channel=sentence_embed_dim, embed_dim=embed_dim) for _ in range(self.n_heads)])
+            else:
+                self.cross_attention_for_sent_modules = nn.ModuleList([CrossAttention(in_channel, cond_channel=sentence_embed_dim, embed_dim=embed_dim) for _ in range(self.n_heads)])
+        
+        self.gamma = [nn.Parameter(torch.tensor(0.0)) for _ in range(self.n_heads)]
+
+        self.conv_for_attention = spectral_init(nn.Conv1d(in_channel * len(self.attention_list), in_channel, 1), gain=gain)
+            
+        self.out = spectral_init(nn.Conv1d(in_channel * self.n_heads, in_channel, 1), gain=gain)
+
+    def forward(self, input, sentence_embedding, word_embedding, sequence_lengths):
+        batch_size, c, f, t = input.shape
+        x = input
+
+        result = []
+        for head in range(self.n_heads):
+            out_list = []
+            if "self" in self.attention_list:
+                x_self, attention_map = self.self_attention_modules[head](x)
+                out_list.append(x_self)
+            if "word" in self.attention_list:
+                x_word, attention_map = self.cross_attention_for_word_modules[head](x, word_embedding, sequence_lengths)
+                out_list.append(x_word)
+            if "sentence" in self.attention_list:
+                x_sent, attention_map = self.cross_attention_for_sent_modules[head](x, sentence_embedding)
+                out_list.append(x_sent)
+            out = torch.cat(out_list, dim=1)
+            out = self.conv_for_attention(out.reshape([batch_size, c*len(out_list), f*t])).reshape([batch_size, c, f, t])
+            out = self.gamma[head] * out + x
+            result.append(out)
+        x = torch.cat(result, dim=1)
+        x = self.out(x.reshape([batch_size, c * self.n_heads, f*t])).reshape([batch_size, c, f, t])
+
+        x = input + x
+
+        return x
+
+class Generator(nn.Module):
+    def __init__(self, model_config=None):
+        super().__init__()
+
+        if model_config is None:
+            model_config = {
+                "noise_dim":128,
+                "g_chaneel":128,
+                "n_heads":10,
+                "sentence_embed_dim":512,
+                "word_embed_dim":768,
+                "attention_list":["self,word,sentence", "word,sentence", "sentence"],
+                "spec_attention":True,
+            }
+
+        self.noise_dim = model_config['noise_dim']
+        self.channel = model_config['g_chaneel']
+        self.n_heads = model_config['n_heads']
+        
+        self.sentence_embed_dim = model_config['sentence_embed_dim']
+        self.word_embed_dim = model_config['word_embed_dim']
+
+        self.attention_list = model_config['attention_list']
+        self.spec_attention = model_config['spec_attention']
+
+        channel_list = [self.channel, self.channel, self.channel//2, self.channel//2, self.channel//4, self.channel//4, self.channel//4, self.channel//8, self.channel//8]
+
+        self.lin_code = spectral_init(nn.Linear(self.noise_dim, channel_list[0] * 2 * 32))
+
+        self.conv1 = ConvBlock(channel_list[0], channel_list[1], condition_dim=self.sentence_embed_dim)
+        self.conv2 = ConvBlock(channel_list[1], channel_list[2], condition_dim=self.sentence_embed_dim)
+        self.multi_triple_attention_1 = Multi_Triple_Attention(channel_list[2],
+                                                   sentence_embed_dim=self.sentence_embed_dim,
+                                                   word_embed_dim=self.word_embed_dim,
+                                                   embed_dim=channel_list[2],
+                                                   reverse=False,
+                                                   n_heads=self.n_heads,
+                                                   attention_list=self.attention_list[0],
+                                                   spec_attention=self.spec_attention)
+
+        self.conv3 = ConvBlock(channel_list[2], channel_list[3], condition_dim=self.sentence_embed_dim)
+        self.conv4 = ConvBlock(channel_list[3], channel_list[4], condition_dim=self.sentence_embed_dim, upsample=False)
+        self.multi_triple_attention_2 = Multi_Triple_Attention(channel_list[4],
+                                                   sentence_embed_dim=self.sentence_embed_dim,
+                                                   word_embed_dim=self.word_embed_dim,
+                                                   embed_dim=channel_list[4],
+                                                   reverse=False,
+                                                   n_heads=self.n_heads,
+                                                   attention_list=self.attention_list[1],
+                                                   spec_attention=self.spec_attention)
+
+        self.conv5 = ConvBlock(channel_list[4], channel_list[5], condition_dim=self.sentence_embed_dim)
+        self.conv6 = ConvBlock(channel_list[5], channel_list[6], condition_dim=self.sentence_embed_dim, upsample=False)
+        self.multi_triple_attention_3 = Multi_Triple_Attention(channel_list[6],
+                                                   sentence_embed_dim=self.sentence_embed_dim,
+                                                   word_embed_dim=self.word_embed_dim,
+                                                   embed_dim=channel_list[6],
+                                                   reverse=False,
+                                                   n_heads=self.n_heads,
+                                                   attention_list=self.attention_list[2],
+                                                   spec_attention=self.spec_attention)
+
+        self.conv7 = ConvBlock(channel_list[6], channel_list[7], condition_dim=self.sentence_embed_dim)
+        self.bn = nn.BatchNorm2d(channel_list[8])
+        self.colorize = spectral_init(nn.Conv1d(channel_list[8], 1, 1))
+
+    def forward(self, z, sentence_embedding, word_embedding, sequence_lengths):
+        batch_size = z.shape[0]
+
+        x = self.lin_code(z)
+        x = x.view(-1, self.channel, 2, 32) # [bsz, c, 2, 32]
+
+        x = self.conv1(x, sentence_embedding) # [bsz, c, 4, 64]
+        x = self.conv2(x, sentence_embedding) # [bsz, c, 8, 128]
+        x = self.multi_triple_attention_1(x, sentence_embedding, word_embedding, sequence_lengths) # [bsz, c, 8, 128]
+
+        x = self.conv3(x, sentence_embedding) # [bsz, c, 16, 256]
+        x = self.conv4(x, sentence_embedding) # [bsz, c, 16, 256]
+        x = self.multi_triple_attention_2(x, sentence_embedding, word_embedding, sequence_lengths) # [bsz, c, 16, 256]
+
+        x = self.conv5(x, sentence_embedding) # [bsz, c, 32, 512]
+        x = self.conv6(x, sentence_embedding) # [bsz, c, 32, 512]
+        x = self.multi_triple_attention_3(x, sentence_embedding, word_embedding, sequence_lengths) # [bsz, c, 32, 512]
+
+        x = self.conv7(x, sentence_embedding) # [bsz, c, 64, 1024]
+        x = self.bn(x) # [bsz, c // 8, 64, 1024]
+        x = F.relu(x)
+        x = self.colorize(x.reshape([batch_size, -1, 64*1024])).reshape([batch_size, 1, 64, 1024]) # [bsz, 1, 64, 1024]
+
+        return x