reverting changes

app.py

@@ -3,7 +3,6 @@ import gradio as gr
 import models as MOD
 import process_data as PD
 from transformers import pipeline
-from manipulate_model.utils import get_config_and_model, infere
 
 model_master = {
     "SSL-AASIST (Trained on ASV-Spoof5)": {"eer_threshold": 3.3330237865448,
@@ -23,8 +22,6 @@ model.load_state_dict(torch.load("ssl_aasist_epoch_7.pth", map_location="cpu"))
 model.eval()
 loaded_model = "SSL-AASIST (Trained on ASV-Spoof5)"
 
-manpulate_config, manipulate_model = get_config_and_model()
-
 def process(file, type):
     global model
     global loaded_model
@@ -98,32 +95,7 @@ transcribe_proc = gr.Interface(
     allow_flagging="never"
 )
 
-#############################################################################################
-#For manipulation detection interface
-
-def detect_manipulation(inputs):
-    global manipulate_model
-    global manpulate_config
-    out = infere(manipulate_model, inputs, manpulate_config)
-    out = out.tolist()
-    return str(out)
-
-manipulate_proc = gr.Interface(
-    fn = detect_manipulation,
-    inputs=[
-    gr.Audio(type="filepath", label="Speech file (<30s)", max_length=30, sources=["microphone", "upload"], show_download_button=True)
-    ],
-    outputs=[
-        gr.Text(label="Predicted manipulations", info="Manipulation detection is performed automatically."),
-    ],
-    title="Find the manipulated segments",
-    description=(
-        "Automatactic manipulation detection service. Upload a audio file."
-    ),
-    allow_flagging="never"
-)
-
 with demo:
-    gr.TabbedInterface([file_proc, transcribe_proc, manipulate_proc], ["Analyze Audio File", "Transcribe Audio File", "Manipulation Detection"])
+    gr.TabbedInterface([file_proc, transcribe_proc], ["Analyze Audio File", "Transcribe Audio File"])
 demo.queue(max_size=10)
 demo.launch(share=True)

app_backup.py

@@ -0,0 +1,129 @@
+import torch
+import gradio as gr
+import models as MOD
+import process_data as PD
+from transformers import pipeline
+from manipulate_model.utils import get_config_and_model, infere
+
+model_master = {
+    "SSL-AASIST (Trained on ASV-Spoof5)": {"eer_threshold": 3.3330237865448,
+                                           "data_process_func": "process_ssl_assist_input",
+                                           "note": "This model is trained only on ASVSpoof 2024 training data.",
+                                           "model_class": "Model",
+                                           "model_checkpoint": "ssl_aasist_epoch_7.pth"},
+    "AASIST": {"eer_threshold": 1.8018419742584229,
+               "data_process_func": "process_assist_input",
+               "note": "This model is trained on ASVSpoof 2024 training data.",
+               "model_class":"AASIST_Model",
+               "model_checkpoint": "orig_aasist_epoch_1.pth"}
+}
+
+model = MOD.Model(None, "cpu")
+model.load_state_dict(torch.load("ssl_aasist_epoch_7.pth", map_location="cpu"))
+model.eval()
+loaded_model = "SSL-AASIST (Trained on ASV-Spoof5)"
+
+manpulate_config, manipulate_model = get_config_and_model()
+
+def process(file, type):
+    global model
+    global loaded_model
+    inp = getattr(PD, model_master[type]["data_process_func"])(file)
+    if not loaded_model == type:
+        model = getattr(MOD, model_master[type]["model_class"])(None, "cpu")
+        model.load_state_dict(torch.load(model_master[type]["model_checkpoint"], map_location="cpu"))
+        model.eval()
+        loaded_model = type
+
+    op = model(inp).detach().squeeze()[1].item()
+
+    response_text = "Decision score: {} \nDecision threshold: {} \nNotes: 1. Any score below threshold is indicative of fake. \n2. {} ".format(
+        str(op), str(model_master[type]["eer_threshold"]), model_master[type]["note"])
+    return response_text
+
+
+demo = gr.Blocks()
+file_proc = gr.Interface(
+    fn=process,
+    inputs=[
+        gr.Audio(sources=["upload"], label="Audio file", type="filepath"),
+        gr.Radio(["SSL-AASIST (Trained on ASV-Spoof5)", "AASIST"], label="Select Model", type="value"),
+    ],
+    outputs="text",
+    title="Find the Fake: Analyze 'Real' or 'Fake'.",
+    description=(
+        "Analyze fake or real with a click of a button. Upload a .wav or .flac file."
+    ),
+    examples=[
+        ["./bonafide.flac", "SSL-AASIST (Trained on ASV-Spoof5)"],
+        ["./fake.flac", "SSL-AASIST (Trained on ASV-Spoof5)"],
+        ["./bonafide.flac", "AASIST"],
+        ["./fake.flac", "AASIST"],
+    ],
+    cache_examples=True,
+    allow_flagging="never",
+)
+#####################################################################################
+# For ASR interface
+pipe = pipeline(
+    task="automatic-speech-recognition",
+    model="openai/whisper-large-v3",
+    chunk_length_s=30,
+    device="cpu",
+)
+
+def transcribe(inputs):
+    if inputs is None:
+        raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
+
+    op = pipe(inputs, batch_size=8, generate_kwargs={"task": "transcribe"}, return_timestamps=False, return_language=True)
+    lang = op["chunks"][0]["language"]
+    text = op["text"]
+
+    return  lang, text
+
+transcribe_proc = gr.Interface(
+    fn = transcribe,
+    inputs = [
+        gr.Audio(type="filepath", label="Speech file (<30s)", max_length=30, sources=["microphone", "upload"], show_download_button=True)
+    ],
+    outputs=[
+        gr.Text(label="Predicted Language", info="Language identification is performed automatically."),
+        gr.Text(label="Predicted transcription", info="Best hypothesis."),
+    ],
+    title="Transcribe Anything.",
+    description=(
+        "Automatactic language identification and transcription service by Whisper Large V3. Upload a .wav or .flac file."
+    ),
+    allow_flagging="never"
+)
+
+#############################################################################################
+#For manipulation detection interface
+
+def detect_manipulation(inputs):
+    global manipulate_model
+    global manpulate_config
+    out = infere(manipulate_model, inputs, manpulate_config)
+    out = out.tolist()
+    return str(out)
+
+manipulate_proc = gr.Interface(
+    fn = detect_manipulation,
+    inputs=[
+    gr.Audio(type="filepath", label="Speech file (<30s)", max_length=30, sources=["microphone", "upload"], show_download_button=True)
+    ],
+    outputs=[
+        gr.Text(label="Predicted manipulations", info="Manipulation detection is performed automatically."),
+    ],
+    title="Find the manipulated segments",
+    description=(
+        "Automatactic manipulation detection service. Upload a audio file."
+    ),
+    allow_flagging="never"
+)
+
+with demo:
+    gr.TabbedInterface([file_proc, transcribe_proc, manipulate_proc], ["Analyze Audio File", "Transcribe Audio File", "Manipulation Detection"])
+demo.queue(max_size=10)
+demo.launch(share=True)

manipulate_model/decoder/aasist/aasist.py

@@ -1,617 +0,0 @@
-import random
-from typing import Union
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch import Tensor
-
-# import fairseq
-
-
-___author__ = "Hemlata Tak"
-__email__ = "[email protected]"
-
-############################
-## FOR fine-tuned SSL MODEL
-############################
-
-
-# class SSLModel(nn.Module):
-#     def __init__(self, device):
-#         super(SSLModel, self).__init__()
-
-#         cp_path = "xlsr2_300m.pt"  # Change the pre-trained XLSR model path.
-#         model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task(
-#             [cp_path]
-#         )
-#         self.model = model[0]
-#         self.device = device
-#         self.out_dim = 1024
-#         return
-
-#     def extract_feat(self, input_data):
-
-#         # put the model to GPU if it not there
-#         if (
-#             next(self.model.parameters()).device != input_data.device
-#             or next(self.model.parameters()).dtype != input_data.dtype
-#         ):
-#             self.model.to(input_data.device, dtype=input_data.dtype)
-#             self.model.train()
-
-#         if True:
-#             # input should be in shape (batch, length)
-#             if input_data.ndim == 3:
-#                 input_tmp = input_data[:, :, 0]
-#             else:
-#                 input_tmp = input_data
-
-#             # [batch, length, dim]
-#             emb = self.model(input_tmp, mask=False, features_only=True)["x"]
-#         return emb
-
-
-# ---------AASIST back-end------------------------#
-""" Jee-weon Jung, Hee-Soo Heo, Hemlata Tak, Hye-jin Shim, Joon Son Chung, Bong-Jin Lee, Ha-Jin Yu and Nicholas Evans. 
-    AASIST: Audio Anti-Spoofing Using Integrated Spectro-Temporal Graph Attention Networks. 
-    In Proc. ICASSP 2022, pp: 6367--6371."""
-
-
-class GraphAttentionLayer(nn.Module):
-    def __init__(self, in_dim, out_dim, **kwargs):
-        super().__init__()
-
-        # attention map
-        self.att_proj = nn.Linear(in_dim, out_dim)
-        self.att_weight = self._init_new_params(out_dim, 1)
-
-        # project
-        self.proj_with_att = nn.Linear(in_dim, out_dim)
-        self.proj_without_att = nn.Linear(in_dim, out_dim)
-
-        # batch norm
-        self.bn = nn.BatchNorm1d(out_dim)
-
-        # dropout for inputs
-        self.input_drop = nn.Dropout(p=0.2)
-
-        # activate
-        self.act = nn.SELU(inplace=True)
-
-        # temperature
-        self.temp = 1.0
-        if "temperature" in kwargs:
-            self.temp = kwargs["temperature"]
-
-    def forward(self, x):
-        """
-        x   :(#bs, #node, #dim)
-        """
-        # apply input dropout
-        x = self.input_drop(x)
-
-        # derive attention map
-        att_map = self._derive_att_map(x)
-
-        # projection
-        x = self._project(x, att_map)
-
-        # apply batch norm
-        x = self._apply_BN(x)
-        x = self.act(x)
-        return x
-
-    def _pairwise_mul_nodes(self, x):
-        """
-        Calculates pairwise multiplication of nodes.
-        - for attention map
-        x           :(#bs, #node, #dim)
-        out_shape   :(#bs, #node, #node, #dim)
-        """
-
-        nb_nodes = x.size(1)
-        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
-        x_mirror = x.transpose(1, 2)
-
-        return x * x_mirror
-
-    def _derive_att_map(self, x):
-        """
-        x           :(#bs, #node, #dim)
-        out_shape   :(#bs, #node, #node, 1)
-        """
-        att_map = self._pairwise_mul_nodes(x)
-        # size: (#bs, #node, #node, #dim_out)
-        att_map = torch.tanh(self.att_proj(att_map))
-        # size: (#bs, #node, #node, 1)
-        att_map = torch.matmul(att_map, self.att_weight)
-
-        # apply temperature
-        att_map = att_map / self.temp
-
-        att_map = F.softmax(att_map, dim=-2)
-
-        return att_map
-
-    def _project(self, x, att_map):
-        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
-        x2 = self.proj_without_att(x)
-
-        return x1 + x2
-
-    def _apply_BN(self, x):
-        org_size = x.size()
-        x = x.view(-1, org_size[-1])
-        x = self.bn(x)
-        x = x.view(org_size)
-
-        return x
-
-    def _init_new_params(self, *size):
-        out = nn.Parameter(torch.FloatTensor(*size))
-        nn.init.xavier_normal_(out)
-        return out
-
-
-class HtrgGraphAttentionLayer(nn.Module):
-    def __init__(self, in_dim, out_dim, **kwargs):
-        super().__init__()
-
-        self.proj_type1 = nn.Linear(in_dim, in_dim)
-        self.proj_type2 = nn.Linear(in_dim, in_dim)
-
-        # attention map
-        self.att_proj = nn.Linear(in_dim, out_dim)
-        self.att_projM = nn.Linear(in_dim, out_dim)
-
-        self.att_weight11 = self._init_new_params(out_dim, 1)
-        self.att_weight22 = self._init_new_params(out_dim, 1)
-        self.att_weight12 = self._init_new_params(out_dim, 1)
-        self.att_weightM = self._init_new_params(out_dim, 1)
-
-        # project
-        self.proj_with_att = nn.Linear(in_dim, out_dim)
-        self.proj_without_att = nn.Linear(in_dim, out_dim)
-
-        self.proj_with_attM = nn.Linear(in_dim, out_dim)
-        self.proj_without_attM = nn.Linear(in_dim, out_dim)
-
-        # batch norm
-        self.bn = nn.BatchNorm1d(out_dim)
-
-        # dropout for inputs
-        self.input_drop = nn.Dropout(p=0.2)
-
-        # activate
-        self.act = nn.SELU(inplace=True)
-
-        # temperature
-        self.temp = 1.0
-        if "temperature" in kwargs:
-            self.temp = kwargs["temperature"]
-
-    def forward(self, x1, x2, master=None):
-        """
-        x1  :(#bs, #node, #dim)
-        x2  :(#bs, #node, #dim)
-        """
-        # print('x1',x1.shape)
-        # print('x2',x2.shape)
-        num_type1 = x1.size(1)
-        num_type2 = x2.size(1)
-        # print('num_type1',num_type1)
-        # print('num_type2',num_type2)
-        x1 = self.proj_type1(x1)
-        # print('proj_type1',x1.shape)
-        x2 = self.proj_type2(x2)
-        # print('proj_type2',x2.shape)
-        x = torch.cat([x1, x2], dim=1)
-        # print('Concat x1 and x2',x.shape)
-
-        if master is None:
-            master = torch.mean(x, dim=1, keepdim=True)
-            # print('master',master.shape)
-        # apply input dropout
-        x = self.input_drop(x)
-
-        # derive attention map
-        att_map = self._derive_att_map(x, num_type1, num_type2)
-        # print('master',master.shape)
-        # directional edge for master node
-        master = self._update_master(x, master)
-        # print('master',master.shape)
-        # projection
-        x = self._project(x, att_map)
-        # print('proj x',x.shape)
-        # apply batch norm
-        x = self._apply_BN(x)
-        x = self.act(x)
-
-        x1 = x.narrow(1, 0, num_type1)
-        # print('x1',x1.shape)
-        x2 = x.narrow(1, num_type1, num_type2)
-        # print('x2',x2.shape)
-        return x1, x2, master
-
-    def _update_master(self, x, master):
-
-        att_map = self._derive_att_map_master(x, master)
-        master = self._project_master(x, master, att_map)
-
-        return master
-
-    def _pairwise_mul_nodes(self, x):
-        """
-        Calculates pairwise multiplication of nodes.
-        - for attention map
-        x           :(#bs, #node, #dim)
-        out_shape   :(#bs, #node, #node, #dim)
-        """
-
-        nb_nodes = x.size(1)
-        x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1)
-        x_mirror = x.transpose(1, 2)
-
-        return x * x_mirror
-
-    def _derive_att_map_master(self, x, master):
-        """
-        x           :(#bs, #node, #dim)
-        out_shape   :(#bs, #node, #node, 1)
-        """
-        att_map = x * master
-        att_map = torch.tanh(self.att_projM(att_map))
-
-        att_map = torch.matmul(att_map, self.att_weightM)
-
-        # apply temperature
-        att_map = att_map / self.temp
-
-        att_map = F.softmax(att_map, dim=-2)
-
-        return att_map
-
-    def _derive_att_map(self, x, num_type1, num_type2):
-        """
-        x           :(#bs, #node, #dim)
-        out_shape   :(#bs, #node, #node, 1)
-        """
-        att_map = self._pairwise_mul_nodes(x)
-        # size: (#bs, #node, #node, #dim_out)
-        att_map = torch.tanh(self.att_proj(att_map))
-        # size: (#bs, #node, #node, 1)
-
-        att_board = torch.zeros_like(att_map[:, :, :, 0]).unsqueeze(-1)
-
-        att_board[:, :num_type1, :num_type1, :] = torch.matmul(
-            att_map[:, :num_type1, :num_type1, :], self.att_weight11
-        )
-        att_board[:, num_type1:, num_type1:, :] = torch.matmul(
-            att_map[:, num_type1:, num_type1:, :], self.att_weight22
-        )
-        att_board[:, :num_type1, num_type1:, :] = torch.matmul(
-            att_map[:, :num_type1, num_type1:, :], self.att_weight12
-        )
-        att_board[:, num_type1:, :num_type1, :] = torch.matmul(
-            att_map[:, num_type1:, :num_type1, :], self.att_weight12
-        )
-
-        att_map = att_board
-
-        # apply temperature
-        att_map = att_map / self.temp
-
-        att_map = F.softmax(att_map, dim=-2)
-
-        return att_map
-
-    def _project(self, x, att_map):
-        x1 = self.proj_with_att(torch.matmul(att_map.squeeze(-1), x))
-        x2 = self.proj_without_att(x)
-
-        return x1 + x2
-
-    def _project_master(self, x, master, att_map):
-
-        x1 = self.proj_with_attM(torch.matmul(att_map.squeeze(-1).unsqueeze(1), x))
-        x2 = self.proj_without_attM(master)
-
-        return x1 + x2
-
-    def _apply_BN(self, x):
-        org_size = x.size()
-        x = x.view(-1, org_size[-1])
-        x = self.bn(x)
-        x = x.view(org_size)
-
-        return x
-
-    def _init_new_params(self, *size):
-        out = nn.Parameter(torch.FloatTensor(*size))
-        nn.init.xavier_normal_(out)
-        return out
-
-
-class GraphPool(nn.Module):
-    def __init__(self, k: float, in_dim: int, p: Union[float, int]):
-        super().__init__()
-        self.k = k
-        self.sigmoid = nn.Sigmoid()
-        self.proj = nn.Linear(in_dim, 1)
-        self.drop = nn.Dropout(p=p) if p > 0 else nn.Identity()
-        self.in_dim = in_dim
-
-    def forward(self, h):
-        Z = self.drop(h)
-        weights = self.proj(Z)
-        scores = self.sigmoid(weights)
-        new_h = self.top_k_graph(scores, h, self.k)
-
-        return new_h
-
-    def top_k_graph(self, scores, h, k):
-        """
-        args
-        =====
-        scores: attention-based weights (#bs, #node, 1)
-        h: graph data (#bs, #node, #dim)
-        k: ratio of remaining nodes, (float)
-        returns
-        =====
-        h: graph pool applied data (#bs, #node', #dim)
-        """
-        _, n_nodes, n_feat = h.size()
-        n_nodes = max(int(n_nodes * k), 1)
-        _, idx = torch.topk(scores, n_nodes, dim=1)
-        idx = idx.expand(-1, -1, n_feat)
-
-        h = h * scores
-        h = torch.gather(h, 1, idx)
-
-        return h
-
-
-class Residual_block(nn.Module):
-    def __init__(self, nb_filts, first=False):
-        super().__init__()
-        self.first = first
-
-        if not self.first:
-            self.bn1 = nn.BatchNorm2d(num_features=nb_filts[0])
-        self.conv1 = nn.Conv2d(
-            in_channels=nb_filts[0],
-            out_channels=nb_filts[1],
-            kernel_size=(2, 3),
-            padding=(1, 1),
-            stride=1,
-        )
-        self.selu = nn.SELU(inplace=True)
-
-        self.bn2 = nn.BatchNorm2d(num_features=nb_filts[1])
-        self.conv2 = nn.Conv2d(
-            in_channels=nb_filts[1],
-            out_channels=nb_filts[1],
-            kernel_size=(2, 3),
-            padding=(0, 1),
-            stride=1,
-        )
-
-        if nb_filts[0] != nb_filts[1]:
-            self.downsample = True
-            self.conv_downsample = nn.Conv2d(
-                in_channels=nb_filts[0],
-                out_channels=nb_filts[1],
-                padding=(0, 1),
-                kernel_size=(1, 3),
-                stride=1,
-            )
-
-        else:
-            self.downsample = False
-
-    def forward(self, x):
-        identity = x
-        if not self.first:
-            out = self.bn1(x)
-            out = self.selu(out)
-        else:
-            out = x
-
-        # print('out',out.shape)
-        out = self.conv1(x)
-
-        # print('aft conv1 out',out.shape)
-        out = self.bn2(out)
-        out = self.selu(out)
-        # print('out',out.shape)
-        out = self.conv2(out)
-        # print('conv2 out',out.shape)
-
-        if self.downsample:
-            identity = self.conv_downsample(identity)
-
-        out += identity
-        # out = self.mp(out)
-        return out
-
-
-class AASIST(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-
-        # AASIST parameters
-        filts = [128, [1, 32], [32, 32], [32, 64], [64, 64]]
-        gat_dims = [64, 32]
-        pool_ratios = [0.5, 0.5, 0.5, 0.5]
-        temperatures = [2.0, 2.0, 100.0, 100.0]
-
-        ####
-        # create network wav2vec 2.0
-        ####
-        # self.ssl_model = SSLModel(self.device)
-        self.LL = nn.Linear(self.config.model.decoder.encoding_dim, 128)
-
-        self.first_bn = nn.BatchNorm2d(num_features=1)
-        self.first_bn1 = nn.BatchNorm2d(num_features=64)
-        self.drop = nn.Dropout(0.5, inplace=True)
-        self.drop_way = nn.Dropout(0.2, inplace=True)
-        self.selu = nn.SELU(inplace=True)
-
-        # RawNet2 encoder
-        self.encoder = nn.Sequential(
-            nn.Sequential(Residual_block(nb_filts=filts[1], first=True)),
-            nn.Sequential(Residual_block(nb_filts=filts[2])),
-            nn.Sequential(Residual_block(nb_filts=filts[3])),
-            nn.Sequential(Residual_block(nb_filts=filts[4])),
-            nn.Sequential(Residual_block(nb_filts=filts[4])),
-            nn.Sequential(Residual_block(nb_filts=filts[4])),
-        )
-
-        self.attention = nn.Sequential(
-            nn.Conv2d(64, 128, kernel_size=(1, 1)),
-            nn.SELU(inplace=True),
-            nn.BatchNorm2d(128),
-            nn.Conv2d(128, 64, kernel_size=(1, 1)),
-        )
-        # position encoding
-        self.pos_S = nn.Parameter(torch.randn(1, 42, filts[-1][-1]))
-
-        self.master1 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))
-        self.master2 = nn.Parameter(torch.randn(1, 1, gat_dims[0]))
-
-        # Graph module
-        self.GAT_layer_S = GraphAttentionLayer(
-            filts[-1][-1], gat_dims[0], temperature=temperatures[0]
-        )
-        self.GAT_layer_T = GraphAttentionLayer(
-            filts[-1][-1], gat_dims[0], temperature=temperatures[1]
-        )
-        # HS-GAL layer
-        self.HtrgGAT_layer_ST11 = HtrgGraphAttentionLayer(
-            gat_dims[0], gat_dims[1], temperature=temperatures[2]
-        )
-        self.HtrgGAT_layer_ST12 = HtrgGraphAttentionLayer(
-            gat_dims[1], gat_dims[1], temperature=temperatures[2]
-        )
-        self.HtrgGAT_layer_ST21 = HtrgGraphAttentionLayer(
-            gat_dims[0], gat_dims[1], temperature=temperatures[2]
-        )
-        self.HtrgGAT_layer_ST22 = HtrgGraphAttentionLayer(
-            gat_dims[1], gat_dims[1], temperature=temperatures[2]
-        )
-
-        # Graph pooling layers
-        self.pool_S = GraphPool(pool_ratios[0], gat_dims[0], 0.3)
-        self.pool_T = GraphPool(pool_ratios[1], gat_dims[0], 0.3)
-        self.pool_hS1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
-        self.pool_hT1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
-
-        self.pool_hS2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
-        self.pool_hT2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3)
-
-        if self.config.model.task == "audio-video":
-            self.out_layer = nn.Linear(5 * gat_dims[1], 4)
-        else:
-            self.out_layer = nn.Linear(5 * gat_dims[1], 2)
-
-    def forward(self, x):
-        # -------pre-trained Wav2vec model fine tunning ------------------------##
-        # x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1))
-        x = self.LL(x)  # (bs,frame_number,feat_out_dim)
-
-        # post-processing on front-end features
-        x = x.transpose(1, 2)  # (bs,feat_out_dim,frame_number)
-        x = x.unsqueeze(dim=1)  # add channel
-        x = F.max_pool2d(x, (3, 3))
-        x = self.first_bn(x)
-        x = self.selu(x)
-
-        # RawNet2-based encoder
-        x = self.encoder(x)
-        x = self.first_bn1(x)
-        x = self.selu(x)
-
-        w = self.attention(x)
-
-        # ------------SA for spectral feature-------------#
-        w1 = F.softmax(w, dim=-1)
-        m = torch.sum(x * w1, dim=-1)
-        e_S = m.transpose(1, 2) + self.pos_S
-
-        # graph module layer
-        gat_S = self.GAT_layer_S(e_S)
-        out_S = self.pool_S(gat_S)  # (#bs, #node, #dim)
-
-        # ------------SA for temporal feature-------------#
-        w2 = F.softmax(w, dim=-2)
-        m1 = torch.sum(x * w2, dim=-2)
-
-        e_T = m1.transpose(1, 2)
-
-        # graph module layer
-        gat_T = self.GAT_layer_T(e_T)
-        out_T = self.pool_T(gat_T)
-
-        # learnable master node
-        master1 = self.master1.expand(x.size(0), -1, -1)
-        master2 = self.master2.expand(x.size(0), -1, -1)
-
-        # inference 1
-        out_T1, out_S1, master1 = self.HtrgGAT_layer_ST11(
-            out_T, out_S, master=self.master1
-        )
-
-        out_S1 = self.pool_hS1(out_S1)
-        out_T1 = self.pool_hT1(out_T1)
-
-        out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST12(
-            out_T1, out_S1, master=master1
-        )
-        out_T1 = out_T1 + out_T_aug
-        out_S1 = out_S1 + out_S_aug
-        master1 = master1 + master_aug
-
-        # inference 2
-        out_T2, out_S2, master2 = self.HtrgGAT_layer_ST21(
-            out_T, out_S, master=self.master2
-        )
-        out_S2 = self.pool_hS2(out_S2)
-        out_T2 = self.pool_hT2(out_T2)
-
-        out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST22(
-            out_T2, out_S2, master=master2
-        )
-        out_T2 = out_T2 + out_T_aug
-        out_S2 = out_S2 + out_S_aug
-        master2 = master2 + master_aug
-
-        out_T1 = self.drop_way(out_T1)
-        out_T2 = self.drop_way(out_T2)
-        out_S1 = self.drop_way(out_S1)
-        out_S2 = self.drop_way(out_S2)
-        master1 = self.drop_way(master1)
-        master2 = self.drop_way(master2)
-
-        out_T = torch.max(out_T1, out_T2)
-        out_S = torch.max(out_S1, out_S2)
-        master = torch.max(master1, master2)
-
-        # Readout operation
-        T_max, _ = torch.max(torch.abs(out_T), dim=1)
-        T_avg = torch.mean(out_T, dim=1)
-
-        S_max, _ = torch.max(torch.abs(out_S), dim=1)
-        S_avg = torch.mean(out_S, dim=1)
-
-        last_hidden = torch.cat([T_max, T_avg, S_max, S_avg, master.squeeze(1)], dim=1)
-
-        last_hidden = self.drop(last_hidden)
-        output = self.out_layer(last_hidden)
-
-        if self.config.model.task == "audio-video":
-            output = output.view(-1, 2, 2)
-
-        return output

manipulate_model/decoder/decoder.py

@@ -1,20 +0,0 @@
-import torch
-import torch.nn as nn
-
-
-class Decoder(nn.Module):
-    def __init__(self, config):
-        super(Decoder, self).__init__()
-        self.config = config
-
-        self.decoder = None
-
-        if config.model.decoder.name.lower() == "aasist":
-            from manipulate_model.decoder.aasist.aasist import AASIST
-
-            self.decoder = AASIST(config)
-        else:
-            raise ValueError("Invalid decoder name")
-
-    def forward(self, x):
-        return self.decoder(x)

manipulate_model/demo-model/audio/config.yaml

@@ -1,44 +0,0 @@
-model:
-  task: audio
-  encoder:
-    name: wavlm
-    version: base
-    pretrained: true
-    pretrained_path: manipulate_model/encoder_checkpoints/wavlm/WavLM-Base+.pt
-    output_layer: 3
-  encoder_freeze: false
-  decoder:
-    name: aasist
-    version: default
-    output_size: 2
-  online_encoding: true
-data:
-  name: av1m
-  train_parts: all
-  val_parts: all
-  test_parts: all
-  train_size: -1
-  val_size: -1
-  test_size: -1
-  shape:
-  - 3
-  - 224
-  - 224
-  sr: 16000
-  fps: 25
-  center_transition: true
-  window_size: 4
-  sliding_window: false
-train:
-  num_workers: 16
-  batch_size: 64
-  num_epochs: 15
-  optimizer: adam
-  scheduler: step
-  lr: 0.0001
-  step_size: 1
-  gamma: 0.1
-  loss: bce
-  log_interval: 100
-  shuffle: true
-debug: false

manipulate_model/demo-model/audio/weights.pt

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

manipulate_model/encoder/encoder.py

@@ -1,40 +0,0 @@
-import torch
-import torch.nn as nn
-
-from torch.nn.functional import pad
-from collections import OrderedDict
-
-
-class Encoder(nn.Module):
-    def __init__(self, config):
-        super(Encoder, self).__init__()
-        self.config = config
-
-        self.encoder = None
-        self.succeeding_layers = None
-
-        # AUDIO
-        if self.config.model.task == "audio":
-            if self.config.model.encoder.name.lower() == "wavlm":
-                from manipulate_model.encoder.wavlm.WavLM import WavLM, WavLMConfig
-
-                ckpt = torch.load(
-                    config.model.encoder.pretrained_path, map_location="cpu"
-                )
-                cfg = WavLMConfig(ckpt["cfg"])
-                self.encoder = WavLM(cfg)
-
-
-    def forward(self, x):
-        if self.config.model.encoder.name.lower() == "wavlm":
-            return self.encoder(x, output_layer=self.config.model.encoder.output_layer)
-        elif self.config.model.encoder.name.lower() == "videomamba":
-            return self.encoder(x)
-
-        return self.encoder(x)
-
-    def get_encoding_dim(self):
-        return self.encoder.get_encoding_dim()
-
-    def get_temporal_dim(self):
-        return self.encoder.get_temporal_dim(window_size=self.config.data.window_size)

manipulate_model/encoder/wavlm/WavLM.py

@@ -1,795 +0,0 @@
-# --------------------------------------------------------
-# WavLM: Large-Scale Self-Supervised  Pre-training  for Full Stack Speech Processing (https://arxiv.org/abs/2110.13900.pdf)
-# Github source: https://github.com/microsoft/unilm/tree/master/wavlm
-# Copyright (c) 2021 Microsoft
-# Licensed under The MIT License [see LICENSE for details]
-# Based on fairseq code bases
-# https://github.com/pytorch/fairseq
-# --------------------------------------------------------
-
-import math
-import logging
-from typing import List, Optional, Tuple
-
-import numpy as np
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.nn import LayerNorm
-from manipulate_model.encoder.wavlm.modules import (
-    Fp32GroupNorm,
-    Fp32LayerNorm,
-    GradMultiply,
-    MultiheadAttention,
-    SamePad,
-    init_bert_params,
-    get_activation_fn,
-    TransposeLast,
-    GLU_Linear,
-)
-
-logger = logging.getLogger(__name__)
-
-
-def compute_mask_indices(
-    shape: Tuple[int, int],
-    padding_mask: Optional[torch.Tensor],
-    mask_prob: float,
-    mask_length: int,
-    mask_type: str = "static",
-    mask_other: float = 0.0,
-    min_masks: int = 0,
-    no_overlap: bool = False,
-    min_space: int = 0,
-) -> np.ndarray:
-    """
-    Computes random mask spans for a given shape
-
-    Args:
-        shape: the the shape for which to compute masks.
-            should be of size 2 where first element is batch size and 2nd is timesteps
-        padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements
-        mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by
-            number of timesteps divided by length of mask span to mask approximately this percentage of all elements.
-            however due to overlaps, the actual number will be smaller (unless no_overlap is True)
-        mask_type: how to compute mask lengths
-            static = fixed size
-            uniform = sample from uniform distribution [mask_other, mask_length*2]
-            normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element
-            poisson = sample from possion distribution with lambda = mask length
-        min_masks: minimum number of masked spans
-        no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping
-        min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans
-    """
-
-    bsz, all_sz = shape
-    mask = np.full((bsz, all_sz), False)
-
-    all_num_mask = int(
-        # add a random number for probabilistic rounding
-        mask_prob * all_sz / float(mask_length)
-        + np.random.rand()
-    )
-
-    all_num_mask = max(min_masks, all_num_mask)
-
-    mask_idcs = []
-    for i in range(bsz):
-        if padding_mask is not None:
-            sz = all_sz - padding_mask[i].long().sum().item()
-            num_mask = int(
-                # add a random number for probabilistic rounding
-                mask_prob * sz / float(mask_length)
-                + np.random.rand()
-            )
-            num_mask = max(min_masks, num_mask)
-        else:
-            sz = all_sz
-            num_mask = all_num_mask
-
-        if mask_type == "static":
-            lengths = np.full(num_mask, mask_length)
-        elif mask_type == "uniform":
-            lengths = np.random.randint(mask_other, mask_length * 2 + 1, size=num_mask)
-        elif mask_type == "normal":
-            lengths = np.random.normal(mask_length, mask_other, size=num_mask)
-            lengths = [max(1, int(round(x))) for x in lengths]
-        elif mask_type == "poisson":
-            lengths = np.random.poisson(mask_length, size=num_mask)
-            lengths = [int(round(x)) for x in lengths]
-        else:
-            raise Exception("unknown mask selection " + mask_type)
-
-        if sum(lengths) == 0:
-            lengths[0] = min(mask_length, sz - 1)
-
-        if no_overlap:
-            mask_idc = []
-
-            def arrange(s, e, length, keep_length):
-                span_start = np.random.randint(s, e - length)
-                mask_idc.extend(span_start + i for i in range(length))
-
-                new_parts = []
-                if span_start - s - min_space >= keep_length:
-                    new_parts.append((s, span_start - min_space + 1))
-                if e - span_start - keep_length - min_space > keep_length:
-                    new_parts.append((span_start + length + min_space, e))
-                return new_parts
-
-            parts = [(0, sz)]
-            min_length = min(lengths)
-            for length in sorted(lengths, reverse=True):
-                lens = np.fromiter(
-                    (e - s if e - s >= length + min_space else 0 for s, e in parts),
-                    np.int,
-                )
-                l_sum = np.sum(lens)
-                if l_sum == 0:
-                    break
-                probs = lens / np.sum(lens)
-                c = np.random.choice(len(parts), p=probs)
-                s, e = parts.pop(c)
-                parts.extend(arrange(s, e, length, min_length))
-            mask_idc = np.asarray(mask_idc)
-        else:
-            min_len = min(lengths)
-            if sz - min_len <= num_mask:
-                min_len = sz - num_mask - 1
-
-            mask_idc = np.random.choice(sz - min_len, num_mask, replace=False)
-
-            mask_idc = np.asarray(
-                [
-                    mask_idc[j] + offset
-                    for j in range(len(mask_idc))
-                    for offset in range(lengths[j])
-                ]
-            )
-
-        mask_idcs.append(np.unique(mask_idc[mask_idc < sz]))
-
-    min_len = min([len(m) for m in mask_idcs])
-    for i, mask_idc in enumerate(mask_idcs):
-        if len(mask_idc) > min_len:
-            mask_idc = np.random.choice(mask_idc, min_len, replace=False)
-        mask[i, mask_idc] = True
-
-    return mask
-
-
-class WavLMConfig:
-    def __init__(self, cfg=None):
-        self.extractor_mode: str = (
-            "default"  # mode for feature extractor. default has a single group norm with d groups in the first conv block, whereas layer_norm has layer norms in every block (meant to use with normalize=True)
-        )
-        self.encoder_layers: int = 12  # num encoder layers in the transformer
-
-        self.encoder_embed_dim: int = 768  # encoder embedding dimension
-        self.encoder_ffn_embed_dim: int = 3072  # encoder embedding dimension for FFN
-        self.encoder_attention_heads: int = 12  # num encoder attention heads
-        self.activation_fn: str = "gelu"  # activation function to use
-
-        self.layer_norm_first: bool = False  # apply layernorm first in the transformer
-        self.conv_feature_layers: str = (
-            "[(512,10,5)] + [(512,3,2)] * 4 + [(512,2,2)] * 2"  # string describing convolutional feature extraction layers in form of a python list that contains [(dim, kernel_size, stride), ...]
-        )
-        self.conv_bias: bool = False  # include bias in conv encoder
-        self.feature_grad_mult: float = (
-            1.0  # multiply feature extractor var grads by this
-        )
-
-        self.normalize: bool = (
-            False  # normalize input to have 0 mean and unit variance during training
-        )
-
-        # dropouts
-        self.dropout: float = 0.1  # dropout probability for the transformer
-        self.attention_dropout: float = 0.1  # dropout probability for attention weights
-        self.activation_dropout: float = (
-            0.0  # dropout probability after activation in FFN
-        )
-        self.encoder_layerdrop: float = (
-            0.0  # probability of dropping a tarnsformer layer
-        )
-        self.dropout_input: float = (
-            0.0  # dropout to apply to the input (after feat extr)
-        )
-        self.dropout_features: float = (
-            0.0  # dropout to apply to the features (after feat extr)
-        )
-
-        # masking
-        self.mask_length: int = 10  # mask length
-        self.mask_prob: float = 0.65  # probability of replacing a token with mask
-        self.mask_selection: str = "static"  # how to choose mask length
-        self.mask_other: float = (
-            0  # secondary mask argument (used for more complex distributions), see help in compute_mask_indicesh
-        )
-        self.no_mask_overlap: bool = False  # whether to allow masks to overlap
-        self.mask_min_space: int = (
-            1  # min space between spans (if no overlap is enabled)
-        )
-
-        # channel masking
-        self.mask_channel_length: int = 10  # length of the mask for features (channels)
-        self.mask_channel_prob: float = 0.0  # probability of replacing a feature with 0
-        self.mask_channel_selection: str = (
-            "static"  # how to choose mask length for channel masking
-        )
-        self.mask_channel_other: float = (
-            0  # secondary mask argument (used for more complex distributions), see help in compute_mask_indices
-        )
-        self.no_mask_channel_overlap: bool = (
-            False  # whether to allow channel masks to overlap
-        )
-        self.mask_channel_min_space: int = (
-            1  # min space between spans (if no overlap is enabled)
-        )
-
-        # positional embeddings
-        self.conv_pos: int = (
-            128  # number of filters for convolutional positional embeddings
-        )
-        self.conv_pos_groups: int = (
-            16  # number of groups for convolutional positional embedding
-        )
-
-        # relative position embedding
-        self.relative_position_embedding: bool = (
-            False  # apply relative position embedding
-        )
-        self.num_buckets: int = 320  # number of buckets for relative position embedding
-        self.max_distance: int = (
-            1280  # maximum distance for relative position embedding
-        )
-        self.gru_rel_pos: bool = False  # apply gated relative position embedding
-
-        if cfg is not None:
-            self.update(cfg)
-
-    def update(self, cfg: dict):
-        self.__dict__.update(cfg)
-
-
-class WavLM(nn.Module):
-    def __init__(
-        self,
-        cfg: WavLMConfig,
-    ) -> None:
-        super().__init__()
-        logger.info(f"WavLM Config: {cfg.__dict__}")
-
-        self.cfg = cfg
-        feature_enc_layers = eval(cfg.conv_feature_layers)
-        self.embed = feature_enc_layers[-1][0]
-
-        self.feature_extractor = ConvFeatureExtractionModel(
-            conv_layers=feature_enc_layers,
-            dropout=0.0,
-            mode=cfg.extractor_mode,
-            conv_bias=cfg.conv_bias,
-        )
-
-        self.post_extract_proj = (
-            nn.Linear(self.embed, cfg.encoder_embed_dim)
-            if self.embed != cfg.encoder_embed_dim
-            else None
-        )
-
-        self.mask_prob = cfg.mask_prob
-        self.mask_selection = cfg.mask_selection
-        self.mask_other = cfg.mask_other
-        self.mask_length = cfg.mask_length
-        self.no_mask_overlap = cfg.no_mask_overlap
-        self.mask_min_space = cfg.mask_min_space
-
-        self.mask_channel_prob = cfg.mask_channel_prob
-        self.mask_channel_selection = cfg.mask_channel_selection
-        self.mask_channel_other = cfg.mask_channel_other
-        self.mask_channel_length = cfg.mask_channel_length
-        self.no_mask_channel_overlap = cfg.no_mask_channel_overlap
-        self.mask_channel_min_space = cfg.mask_channel_min_space
-
-        self.dropout_input = nn.Dropout(cfg.dropout_input)
-        self.dropout_features = nn.Dropout(cfg.dropout_features)
-
-        self.feature_grad_mult = cfg.feature_grad_mult
-
-        self.mask_emb = nn.Parameter(
-            torch.FloatTensor(cfg.encoder_embed_dim).uniform_()
-        )
-
-        self.encoder = TransformerEncoder(cfg)
-        self.layer_norm = LayerNorm(self.embed)
-
-    def apply_mask(self, x, padding_mask):
-        B, T, C = x.shape
-        if self.mask_prob > 0:
-            mask_indices = compute_mask_indices(
-                (B, T),
-                padding_mask,
-                self.mask_prob,
-                self.mask_length,
-                self.mask_selection,
-                self.mask_other,
-                min_masks=2,
-                no_overlap=self.no_mask_overlap,
-                min_space=self.mask_min_space,
-            )
-            mask_indices = torch.from_numpy(mask_indices).to(x.device)
-            x[mask_indices] = self.mask_emb
-        else:
-            mask_indices = None
-
-        if self.mask_channel_prob > 0:
-            mask_channel_indices = compute_mask_indices(
-                (B, C),
-                None,
-                self.mask_channel_prob,
-                self.mask_channel_length,
-                self.mask_channel_selection,
-                self.mask_channel_other,
-                no_overlap=self.no_mask_channel_overlap,
-                min_space=self.mask_channel_min_space,
-            )
-            mask_channel_indices = (
-                torch.from_numpy(mask_channel_indices)
-                .to(x.device)
-                .unsqueeze(1)
-                .expand(-1, T, -1)
-            )
-            x[mask_channel_indices] = 0
-
-        return x, mask_indices
-
-    def forward_padding_mask(
-        self,
-        features: torch.Tensor,
-        padding_mask: torch.Tensor,
-    ) -> torch.Tensor:
-        extra = padding_mask.size(1) % features.size(1)
-        if extra > 0:
-            padding_mask = padding_mask[:, :-extra]
-        padding_mask = padding_mask.view(padding_mask.size(0), features.size(1), -1)
-        padding_mask = padding_mask.all(-1)
-        return padding_mask
-
-    def extract_features(
-        self,
-        source: torch.Tensor,
-        padding_mask: Optional[torch.Tensor] = None,
-        mask: bool = False,
-        ret_conv: bool = False,
-        output_layer: Optional[int] = None,
-        ret_layer_results: bool = False,
-    ):
-
-        if self.feature_grad_mult > 0:
-            features = self.feature_extractor(source)
-            if self.feature_grad_mult != 1.0:
-                features = GradMultiply.apply(features, self.feature_grad_mult)
-        else:
-            with torch.no_grad():
-                features = self.feature_extractor(source)
-
-        features = features.transpose(1, 2)
-        features = self.layer_norm(features)
-
-        if padding_mask is not None:
-            padding_mask = self.forward_padding_mask(features, padding_mask)
-
-        if self.post_extract_proj is not None:
-            features = self.post_extract_proj(features)
-
-        features = self.dropout_input(features)
-
-        if mask:
-            x, mask_indices = self.apply_mask(features, padding_mask)
-        else:
-            x = features
-
-        # feature: (B, T, D), float
-        # target: (B, T), long
-        # x: (B, T, D), float
-        # padding_mask: (B, T), bool
-        # mask_indices: (B, T), bool
-        x, layer_results = self.encoder(
-            x,
-            padding_mask=padding_mask,
-            layer=None if output_layer is None else output_layer - 1,
-        )
-
-        res = {
-            "x": x,
-            "padding_mask": padding_mask,
-            "features": features,
-            "layer_results": layer_results,
-        }
-
-        feature = res["features"] if ret_conv else res["x"]
-        if ret_layer_results:
-            feature = (feature, res["layer_results"])
-        return feature, res["padding_mask"]
-
-    def forward(self, x, output_layer=None):
-        return self.extract_features(x, output_layer=output_layer)[0]
-
-    def get_encoding_dim(self):
-        return self.cfg.encoder_embed_dim
-
-    def get_temporal_dim(self, window_size):
-        return 2 * window_size - 1
-
-
-class ConvFeatureExtractionModel(nn.Module):
-    def __init__(
-        self,
-        conv_layers: List[Tuple[int, int, int]],
-        dropout: float = 0.0,
-        mode: str = "default",
-        conv_bias: bool = False,
-        conv_type: str = "default",
-    ):
-        super().__init__()
-
-        assert mode in {"default", "layer_norm"}
-
-        def block(
-            n_in,
-            n_out,
-            k,
-            stride,
-            is_layer_norm=False,
-            is_group_norm=False,
-            conv_bias=False,
-        ):
-            def make_conv():
-                conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias)
-                nn.init.kaiming_normal_(conv.weight)
-                return conv
-
-            assert (
-                is_layer_norm and is_group_norm
-            ) == False, "layer norm and group norm are exclusive"
-
-            if is_layer_norm:
-                return nn.Sequential(
-                    make_conv(),
-                    nn.Dropout(p=dropout),
-                    nn.Sequential(
-                        TransposeLast(),
-                        Fp32LayerNorm(dim, elementwise_affine=True),
-                        TransposeLast(),
-                    ),
-                    nn.GELU(),
-                )
-            elif is_group_norm:
-                return nn.Sequential(
-                    make_conv(),
-                    nn.Dropout(p=dropout),
-                    Fp32GroupNorm(dim, dim, affine=True),
-                    nn.GELU(),
-                )
-            else:
-                return nn.Sequential(make_conv(), nn.Dropout(p=dropout), nn.GELU())
-
-        self.conv_type = conv_type
-        if self.conv_type == "default":
-            in_d = 1
-            self.conv_layers = nn.ModuleList()
-            for i, cl in enumerate(conv_layers):
-                assert len(cl) == 3, "invalid conv definition: " + str(cl)
-                (dim, k, stride) = cl
-
-                self.conv_layers.append(
-                    block(
-                        in_d,
-                        dim,
-                        k,
-                        stride,
-                        is_layer_norm=mode == "layer_norm",
-                        is_group_norm=mode == "default" and i == 0,
-                        conv_bias=conv_bias,
-                    )
-                )
-                in_d = dim
-        elif self.conv_type == "conv2d":
-            in_d = 1
-            self.conv_layers = nn.ModuleList()
-            for i, cl in enumerate(conv_layers):
-                assert len(cl) == 3
-                (dim, k, stride) = cl
-
-                self.conv_layers.append(torch.nn.Conv2d(in_d, dim, k, stride))
-                self.conv_layers.append(torch.nn.ReLU())
-                in_d = dim
-        elif self.conv_type == "custom":
-            in_d = 1
-            idim = 80
-            self.conv_layers = nn.ModuleList()
-            for i, cl in enumerate(conv_layers):
-                assert len(cl) == 3
-                (dim, k, stride) = cl
-                self.conv_layers.append(
-                    torch.nn.Conv2d(in_d, dim, k, stride, padding=1)
-                )
-                self.conv_layers.append(torch.nn.LayerNorm([dim, idim]))
-                self.conv_layers.append(torch.nn.ReLU())
-                in_d = dim
-                if (i + 1) % 2 == 0:
-                    self.conv_layers.append(
-                        torch.nn.MaxPool2d(2, stride=2, ceil_mode=True)
-                    )
-                    idim = int(math.ceil(idim / 2))
-        else:
-            pass
-
-    def forward(self, x, mask=None):
-
-        # BxT -> BxCxT
-        x = x.unsqueeze(1)
-        if self.conv_type == "custom":
-            for conv in self.conv_layers:
-                if isinstance(conv, nn.LayerNorm):
-                    x = x.transpose(1, 2)
-                    x = conv(x).transpose(1, 2)
-                else:
-                    x = conv(x)
-            x = x.transpose(2, 3).contiguous()
-            x = x.view(x.size(0), -1, x.size(-1))
-        else:
-            for conv in self.conv_layers:
-                x = conv(x)
-            if self.conv_type == "conv2d":
-                b, c, t, f = x.size()
-                x = x.transpose(2, 3).contiguous().view(b, c * f, t)
-        return x
-
-
-class TransformerEncoder(nn.Module):
-    def __init__(self, args):
-        super().__init__()
-
-        self.dropout = args.dropout
-        self.embedding_dim = args.encoder_embed_dim
-
-        self.pos_conv = nn.Conv1d(
-            self.embedding_dim,
-            self.embedding_dim,
-            kernel_size=args.conv_pos,
-            padding=args.conv_pos // 2,
-            groups=args.conv_pos_groups,
-        )
-        dropout = 0
-        std = math.sqrt((4 * (1.0 - dropout)) / (args.conv_pos * self.embedding_dim))
-        nn.init.normal_(self.pos_conv.weight, mean=0, std=std)
-        nn.init.constant_(self.pos_conv.bias, 0)
-
-        self.pos_conv = nn.utils.weight_norm(self.pos_conv, name="weight", dim=2)
-        self.pos_conv = nn.Sequential(self.pos_conv, SamePad(args.conv_pos), nn.GELU())
-
-        if hasattr(args, "relative_position_embedding"):
-            self.relative_position_embedding = args.relative_position_embedding
-            self.num_buckets = args.num_buckets
-            self.max_distance = args.max_distance
-        else:
-            self.relative_position_embedding = False
-            self.num_buckets = 0
-            self.max_distance = 0
-
-        self.layers = nn.ModuleList(
-            [
-                TransformerSentenceEncoderLayer(
-                    embedding_dim=self.embedding_dim,
-                    ffn_embedding_dim=args.encoder_ffn_embed_dim,
-                    num_attention_heads=args.encoder_attention_heads,
-                    dropout=self.dropout,
-                    attention_dropout=args.attention_dropout,
-                    activation_dropout=args.activation_dropout,
-                    activation_fn=args.activation_fn,
-                    layer_norm_first=args.layer_norm_first,
-                    has_relative_attention_bias=(
-                        self.relative_position_embedding and i == 0
-                    ),
-                    num_buckets=self.num_buckets,
-                    max_distance=self.max_distance,
-                    gru_rel_pos=args.gru_rel_pos,
-                )
-                for i in range(args.encoder_layers)
-            ]
-        )
-
-        self.layer_norm_first = args.layer_norm_first
-        self.layer_norm = LayerNorm(self.embedding_dim)
-        self.layerdrop = args.encoder_layerdrop
-
-        self.apply(init_bert_params)
-
-    def forward(self, x, padding_mask=None, streaming_mask=None, layer=None):
-        x, layer_results = self.extract_features(x, padding_mask, streaming_mask, layer)
-
-        if self.layer_norm_first and layer is None:
-            x = self.layer_norm(x)
-
-        return x, layer_results
-
-    def extract_features(
-        self, x, padding_mask=None, streaming_mask=None, tgt_layer=None
-    ):
-
-        if padding_mask is not None:
-            x[padding_mask] = 0
-
-        x_conv = self.pos_conv(x.transpose(1, 2))
-        x_conv = x_conv.transpose(1, 2)
-        x = x + x_conv
-
-        if not self.layer_norm_first:
-            x = self.layer_norm(x)
-
-        x = F.dropout(x, p=self.dropout, training=self.training)
-
-        # B x T x C -> T x B x C
-        x = x.transpose(0, 1)
-
-        layer_results = []
-        z = None
-        if tgt_layer is not None:
-            layer_results.append((x, z))
-        r = None
-        pos_bias = None
-        for i, layer in enumerate(self.layers):
-            dropout_probability = np.random.random()
-            if not self.training or (dropout_probability > self.layerdrop):
-                x, z, pos_bias = layer(
-                    x,
-                    self_attn_padding_mask=padding_mask,
-                    need_weights=False,
-                    self_attn_mask=streaming_mask,
-                    pos_bias=pos_bias,
-                )
-            if tgt_layer is not None:
-                layer_results.append((x, z))
-            if i == tgt_layer:
-                r = x
-                break
-
-        if r is not None:
-            x = r
-
-        # T x B x C -> B x T x C
-        x = x.transpose(0, 1)
-
-        return x, layer_results
-
-
-class TransformerSentenceEncoderLayer(nn.Module):
-    """
-    Implements a Transformer Encoder Layer used in BERT/XLM style pre-trained
-    models.
-    """
-
-    def __init__(
-        self,
-        embedding_dim: float = 768,
-        ffn_embedding_dim: float = 3072,
-        num_attention_heads: float = 8,
-        dropout: float = 0.1,
-        attention_dropout: float = 0.1,
-        activation_dropout: float = 0.1,
-        activation_fn: str = "relu",
-        layer_norm_first: bool = False,
-        has_relative_attention_bias: bool = False,
-        num_buckets: int = 0,
-        max_distance: int = 0,
-        rescale_init: bool = False,
-        gru_rel_pos: bool = False,
-    ) -> None:
-
-        super().__init__()
-        # Initialize parameters
-        self.embedding_dim = embedding_dim
-        self.dropout = dropout
-        self.activation_dropout = activation_dropout
-
-        # Initialize blocks
-        self.activation_name = activation_fn
-        self.activation_fn = get_activation_fn(activation_fn)
-        self.self_attn = MultiheadAttention(
-            self.embedding_dim,
-            num_attention_heads,
-            dropout=attention_dropout,
-            self_attention=True,
-            has_relative_attention_bias=has_relative_attention_bias,
-            num_buckets=num_buckets,
-            max_distance=max_distance,
-            rescale_init=rescale_init,
-            gru_rel_pos=gru_rel_pos,
-        )
-
-        self.dropout1 = nn.Dropout(dropout)
-        self.dropout2 = nn.Dropout(self.activation_dropout)
-        self.dropout3 = nn.Dropout(dropout)
-
-        self.layer_norm_first = layer_norm_first
-
-        # layer norm associated with the self attention layer
-        self.self_attn_layer_norm = LayerNorm(self.embedding_dim)
-
-        if self.activation_name == "glu":
-            self.fc1 = GLU_Linear(self.embedding_dim, ffn_embedding_dim, "swish")
-        else:
-            self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
-        self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)
-
-        # layer norm associated with the position wise feed-forward NN
-        self.final_layer_norm = LayerNorm(self.embedding_dim)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        self_attn_mask: torch.Tensor = None,
-        self_attn_padding_mask: torch.Tensor = None,
-        need_weights: bool = False,
-        pos_bias=None,
-    ):
-        """
-        LayerNorm is applied either before or after the self-attention/ffn
-        modules similar to the original Transformer imlementation.
-        """
-        residual = x
-
-        if self.layer_norm_first:
-            x = self.self_attn_layer_norm(x)
-            x, attn, pos_bias = self.self_attn(
-                query=x,
-                key=x,
-                value=x,
-                key_padding_mask=self_attn_padding_mask,
-                need_weights=False,
-                attn_mask=self_attn_mask,
-                position_bias=pos_bias,
-            )
-            x = self.dropout1(x)
-            x = residual + x
-
-            residual = x
-            x = self.final_layer_norm(x)
-            if self.activation_name == "glu":
-                x = self.fc1(x)
-            else:
-                x = self.activation_fn(self.fc1(x))
-            x = self.dropout2(x)
-            x = self.fc2(x)
-            x = self.dropout3(x)
-            x = residual + x
-        else:
-            x, attn, pos_bias = self.self_attn(
-                query=x,
-                key=x,
-                value=x,
-                key_padding_mask=self_attn_padding_mask,
-                need_weights=need_weights,
-                attn_mask=self_attn_mask,
-                position_bias=pos_bias,
-            )
-
-            x = self.dropout1(x)
-            x = residual + x
-
-            x = self.self_attn_layer_norm(x)
-
-            residual = x
-            if self.activation_name == "glu":
-                x = self.fc1(x)
-            else:
-                x = self.activation_fn(self.fc1(x))
-            x = self.dropout2(x)
-            x = self.fc2(x)
-            x = self.dropout3(x)
-            x = residual + x
-            x = self.final_layer_norm(x)
-
-        return x, attn, pos_bias

manipulate_model/encoder/wavlm/modules.py

@@ -1,827 +0,0 @@
-# --------------------------------------------------------
-# WavLM: Large-Scale Self-Supervised  Pre-training  for Full Stack Speech Processing (https://arxiv.org/abs/2110.13900.pdf)
-# Github source: https://github.com/microsoft/unilm/tree/master/wavlm
-# Copyright (c) 2021 Microsoft
-# Licensed under The MIT License [see LICENSE for details]
-# Based on fairseq code bases
-# https://github.com/pytorch/fairseq
-# --------------------------------------------------------
-
-import math
-import warnings
-from typing import Dict, Optional, Tuple
-import torch
-from torch import Tensor, nn
-from torch.nn import Parameter
-import torch.nn.functional as F
-
-
-class TransposeLast(nn.Module):
-    def __init__(self, deconstruct_idx=None):
-        super().__init__()
-        self.deconstruct_idx = deconstruct_idx
-
-    def forward(self, x):
-        if self.deconstruct_idx is not None:
-            x = x[self.deconstruct_idx]
-        return x.transpose(-2, -1)
-
-
-class Fp32LayerNorm(nn.LayerNorm):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    def forward(self, input):
-        output = F.layer_norm(
-            input.float(),
-            self.normalized_shape,
-            self.weight.float() if self.weight is not None else None,
-            self.bias.float() if self.bias is not None else None,
-            self.eps,
-        )
-        return output.type_as(input)
-
-
-class Fp32GroupNorm(nn.GroupNorm):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    def forward(self, input):
-        output = F.group_norm(
-            input.float(),
-            self.num_groups,
-            self.weight.float() if self.weight is not None else None,
-            self.bias.float() if self.bias is not None else None,
-            self.eps,
-        )
-        return output.type_as(input)
-
-
-class GradMultiply(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, x, scale):
-        ctx.scale = scale
-        res = x.new(x)
-        return res
-
-    @staticmethod
-    def backward(ctx, grad):
-        return grad * ctx.scale, None
-
-
-class SamePad(nn.Module):
-    def __init__(self, kernel_size, causal=False):
-        super().__init__()
-        if causal:
-            self.remove = kernel_size - 1
-        else:
-            self.remove = 1 if kernel_size % 2 == 0 else 0
-
-    def forward(self, x):
-        if self.remove > 0:
-            x = x[:, :, : -self.remove]
-        return x
-
-
-class Swish(nn.Module):
-    """Swish function
-    """
-
-    def __init__(self):
-        """Construct an MultiHeadedAttention object."""
-        super(Swish, self).__init__()
-        self.act = torch.nn.Sigmoid()
-
-    def forward(self, x):
-        return x * self.act(x)
-
-
-class GLU_Linear(nn.Module):
-    def __init__(self, input_dim, output_dim, glu_type="sigmoid", bias_in_glu=True):
-        super(GLU_Linear, self).__init__()
-
-        self.glu_type = glu_type
-        self.output_dim = output_dim
-
-        if glu_type == "sigmoid":
-            self.glu_act = torch.nn.Sigmoid()
-        elif glu_type == "swish":
-            self.glu_act = Swish()
-        elif glu_type == "relu":
-            self.glu_act = torch.nn.ReLU()
-        elif glu_type == "gelu":
-            self.glu_act = torch.nn.GELU()
-
-        if bias_in_glu:
-            self.linear = nn.Linear(input_dim, output_dim * 2, True)
-        else:
-            self.linear = nn.Linear(input_dim, output_dim * 2, False)
-
-    def forward(self, x):
-        # to be consistent with GLU_Linear, we assume the input always has the #channel (#dim) in the last dimension of the tensor, so need to switch the dimension first for 1D-Conv case
-        x = self.linear(x)
-
-        if self.glu_type == "bilinear":
-            x = (x[:, :, 0:self.output_dim] * x[:, :, self.output_dim:self.output_dim * 2])
-        else:
-            x = (x[:, :, 0:self.output_dim] * self.glu_act(x[:, :, self.output_dim:self.output_dim * 2]))
-
-        return x
-
-
-def gelu_accurate(x):
-    if not hasattr(gelu_accurate, "_a"):
-        gelu_accurate._a = math.sqrt(2 / math.pi)
-    return (
-        0.5 * x * (1 + torch.tanh(gelu_accurate._a * (x + 0.044715 * torch.pow(x, 3))))
-    )
-
-
-def gelu(x: torch.Tensor) -> torch.Tensor:
-    return torch.nn.functional.gelu(x.float()).type_as(x)
-
-
-def get_activation_fn(activation: str):
-    """Returns the activation function corresponding to `activation`"""
-
-    if activation == "relu":
-        return F.relu
-    elif activation == "gelu":
-        return gelu
-    elif activation == "gelu_fast":
-        warnings.warn(
-            "--activation-fn=gelu_fast has been renamed to gelu_accurate"
-        )
-        return gelu_accurate
-    elif activation == "gelu_accurate":
-        return gelu_accurate
-    elif activation == "tanh":
-        return torch.tanh
-    elif activation == "linear":
-        return lambda x: x
-    elif activation == "glu":
-        return lambda x: x
-    else:
-        raise RuntimeError("--activation-fn {} not supported".format(activation))
-
-
-def init_bert_params(module):
-    """
-    Initialize the weights specific to the BERT Model.
-    This overrides the default initializations depending on the specified arguments.
-        1. If normal_init_linear_weights is set then weights of linear
-           layer will be initialized using the normal distribution and
-           bais will be set to the specified value.
-        2. If normal_init_embed_weights is set then weights of embedding
-           layer will be initialized using the normal distribution.
-        3. If normal_init_proj_weights is set then weights of
-           in_project_weight for MultiHeadAttention initialized using
-           the normal distribution (to be validated).
-    """
-
-    def normal_(data):
-        # with FSDP, module params will be on CUDA, so we cast them back to CPU
-        # so that the RNG is consistent with and without FSDP
-        data.copy_(
-            data.cpu().normal_(mean=0.0, std=0.02).to(data.device)
-        )
-
-    if isinstance(module, nn.Linear):
-        normal_(module.weight.data)
-        if module.bias is not None:
-            module.bias.data.zero_()
-    if isinstance(module, nn.Embedding):
-        normal_(module.weight.data)
-        if module.padding_idx is not None:
-            module.weight.data[module.padding_idx].zero_()
-    if isinstance(module, MultiheadAttention):
-        normal_(module.q_proj.weight.data)
-        normal_(module.k_proj.weight.data)
-        normal_(module.v_proj.weight.data)
-
-
-def quant_noise(module, p, block_size):
-    """
-    Wraps modules and applies quantization noise to the weights for
-    subsequent quantization with Iterative Product Quantization as
-    described in "Training with Quantization Noise for Extreme Model Compression"
-
-    Args:
-        - module: nn.Module
-        - p: amount of Quantization Noise
-        - block_size: size of the blocks for subsequent quantization with iPQ
-
-    Remarks:
-        - Module weights must have the right sizes wrt the block size
-        - Only Linear, Embedding and Conv2d modules are supported for the moment
-        - For more detail on how to quantize by blocks with convolutional weights,
-          see "And the Bit Goes Down: Revisiting the Quantization of Neural Networks"
-        - We implement the simplest form of noise here as stated in the paper
-          which consists in randomly dropping blocks
-    """
-
-    # if no quantization noise, don't register hook
-    if p <= 0:
-        return module
-
-    # supported modules
-    assert isinstance(module, (nn.Linear, nn.Embedding, nn.Conv2d))
-
-    # test whether module.weight has the right sizes wrt block_size
-    is_conv = module.weight.ndim == 4
-
-    # 2D matrix
-    if not is_conv:
-        assert (
-            module.weight.size(1) % block_size == 0
-        ), "Input features must be a multiple of block sizes"
-
-    # 4D matrix
-    else:
-        # 1x1 convolutions
-        if module.kernel_size == (1, 1):
-            assert (
-                module.in_channels % block_size == 0
-            ), "Input channels must be a multiple of block sizes"
-        # regular convolutions
-        else:
-            k = module.kernel_size[0] * module.kernel_size[1]
-            assert k % block_size == 0, "Kernel size must be a multiple of block size"
-
-    def _forward_pre_hook(mod, input):
-        # no noise for evaluation
-        if mod.training:
-            if not is_conv:
-                # gather weight and sizes
-                weight = mod.weight
-                in_features = weight.size(1)
-                out_features = weight.size(0)
-
-                # split weight matrix into blocks and randomly drop selected blocks
-                mask = torch.zeros(
-                    in_features // block_size * out_features, device=weight.device
-                )
-                mask.bernoulli_(p)
-                mask = mask.repeat_interleave(block_size, -1).view(-1, in_features)
-
-            else:
-                # gather weight and sizes
-                weight = mod.weight
-                in_channels = mod.in_channels
-                out_channels = mod.out_channels
-
-                # split weight matrix into blocks and randomly drop selected blocks
-                if mod.kernel_size == (1, 1):
-                    mask = torch.zeros(
-                        int(in_channels // block_size * out_channels),
-                        device=weight.device,
-                    )
-                    mask.bernoulli_(p)
-                    mask = mask.repeat_interleave(block_size, -1).view(-1, in_channels)
-                else:
-                    mask = torch.zeros(
-                        weight.size(0), weight.size(1), device=weight.device
-                    )
-                    mask.bernoulli_(p)
-                    mask = (
-                        mask.unsqueeze(2)
-                        .unsqueeze(3)
-                        .repeat(1, 1, mod.kernel_size[0], mod.kernel_size[1])
-                    )
-
-            # scale weights and apply mask
-            mask = mask.to(
-                torch.bool
-            )  # x.bool() is not currently supported in TorchScript
-            s = 1 / (1 - p)
-            mod.weight.data = s * weight.masked_fill(mask, 0)
-
-    module.register_forward_pre_hook(_forward_pre_hook)
-    return module
-
-
-class MultiheadAttention(nn.Module):
-    """Multi-headed attention.
-
-    See "Attention Is All You Need" for more details.
-    """
-
-    def __init__(
-            self,
-            embed_dim,
-            num_heads,
-            kdim=None,
-            vdim=None,
-            dropout=0.0,
-            bias=True,
-            add_bias_kv=False,
-            add_zero_attn=False,
-            self_attention=False,
-            encoder_decoder_attention=False,
-            q_noise=0.0,
-            qn_block_size=8,
-            has_relative_attention_bias=False,
-            num_buckets=32,
-            max_distance=128,
-            gru_rel_pos=False,
-            rescale_init=False,
-    ):
-        super().__init__()
-        self.embed_dim = embed_dim
-        self.kdim = kdim if kdim is not None else embed_dim
-        self.vdim = vdim if vdim is not None else embed_dim
-        self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
-
-        self.num_heads = num_heads
-        self.dropout_module = nn.Dropout(dropout)
-
-        self.has_relative_attention_bias = has_relative_attention_bias
-        self.num_buckets = num_buckets
-        self.max_distance = max_distance
-        if self.has_relative_attention_bias:
-            self.relative_attention_bias = nn.Embedding(num_buckets, num_heads)
-
-        self.head_dim = embed_dim // num_heads
-        self.q_head_dim = self.head_dim
-        self.k_head_dim = self.head_dim
-        assert (
-                self.head_dim * num_heads == self.embed_dim
-        ), "embed_dim must be divisible by num_heads"
-        self.scaling = self.head_dim ** -0.5
-
-        self.self_attention = self_attention
-        self.encoder_decoder_attention = encoder_decoder_attention
-
-        assert not self.self_attention or self.qkv_same_dim, (
-            "Self-attention requires query, key and " "value to be of the same size"
-        )
-
-        k_bias = True
-        if rescale_init:
-            k_bias = False
-
-        k_embed_dim = embed_dim
-        q_embed_dim = embed_dim
-
-        self.k_proj = quant_noise(
-            nn.Linear(self.kdim, k_embed_dim, bias=k_bias), q_noise, qn_block_size
-        )
-        self.v_proj = quant_noise(
-            nn.Linear(self.vdim, embed_dim, bias=bias), q_noise, qn_block_size
-        )
-        self.q_proj = quant_noise(
-            nn.Linear(embed_dim, q_embed_dim, bias=bias), q_noise, qn_block_size
-        )
-
-        self.out_proj = quant_noise(
-            nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size
-        )
-
-        if add_bias_kv:
-            self.bias_k = Parameter(torch.Tensor(1, 1, embed_dim))
-            self.bias_v = Parameter(torch.Tensor(1, 1, embed_dim))
-        else:
-            self.bias_k = self.bias_v = None
-
-        self.add_zero_attn = add_zero_attn
-
-        self.gru_rel_pos = gru_rel_pos
-        if self.gru_rel_pos:
-            self.grep_linear = nn.Linear(self.q_head_dim, 8)
-            self.grep_a = nn.Parameter(torch.ones(1, num_heads, 1, 1))
-
-        self.reset_parameters()
-
-    def reset_parameters(self):
-        if self.qkv_same_dim:
-            # Empirically observed the convergence to be much better with
-            # the scaled initialization
-            nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
-            nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
-            nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
-        else:
-            nn.init.xavier_uniform_(self.k_proj.weight)
-            nn.init.xavier_uniform_(self.v_proj.weight)
-            nn.init.xavier_uniform_(self.q_proj.weight)
-
-        nn.init.xavier_uniform_(self.out_proj.weight)
-        if self.out_proj.bias is not None:
-            nn.init.constant_(self.out_proj.bias, 0.0)
-        if self.bias_k is not None:
-            nn.init.xavier_normal_(self.bias_k)
-        if self.bias_v is not None:
-            nn.init.xavier_normal_(self.bias_v)
-        if self.has_relative_attention_bias:
-            nn.init.xavier_normal_(self.relative_attention_bias.weight)
-
-    def _relative_positions_bucket(self, relative_positions, bidirectional=True):
-        num_buckets = self.num_buckets
-        max_distance = self.max_distance
-        relative_buckets = 0
-
-        if bidirectional:
-            num_buckets = num_buckets // 2
-            relative_buckets += (relative_positions > 0).to(torch.long) * num_buckets
-            relative_positions = torch.abs(relative_positions)
-        else:
-            relative_positions = -torch.min(relative_positions, torch.zeros_like(relative_positions))
-
-        max_exact = num_buckets // 2
-        is_small = relative_positions < max_exact
-
-        relative_postion_if_large = max_exact + (
-                torch.log(relative_positions.float() / max_exact)
-                / math.log(max_distance / max_exact)
-                * (num_buckets - max_exact)
-        ).to(torch.long)
-        relative_postion_if_large = torch.min(
-            relative_postion_if_large, torch.full_like(relative_postion_if_large, num_buckets - 1)
-        )
-
-        relative_buckets += torch.where(is_small, relative_positions, relative_postion_if_large)
-        return relative_buckets
-
-    def compute_bias(self, query_length, key_length):
-        context_position = torch.arange(query_length, dtype=torch.long)[:, None]
-        memory_position = torch.arange(key_length, dtype=torch.long)[None, :]
-        relative_position = memory_position - context_position
-        relative_position_bucket = self._relative_positions_bucket(
-            relative_position,
-            bidirectional=True
-        )
-        relative_position_bucket = relative_position_bucket.to(self.relative_attention_bias.weight.device)
-        values = self.relative_attention_bias(relative_position_bucket)
-        values = values.permute([2, 0, 1])
-        return values
-
-    def forward(
-            self,
-            query,
-            key: Optional[Tensor],
-            value: Optional[Tensor],
-            key_padding_mask: Optional[Tensor] = None,
-            incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
-            need_weights: bool = True,
-            static_kv: bool = False,
-            attn_mask: Optional[Tensor] = None,
-            before_softmax: bool = False,
-            need_head_weights: bool = False,
-            position_bias: Optional[Tensor] = None
-    ) -> Tuple[Tensor, Optional[Tensor], Optional[Tensor]]:
-        """Input shape: Time x Batch x Channel
-
-        Args:
-            key_padding_mask (ByteTensor, optional): mask to exclude
-                keys that are pads, of shape `(batch, src_len)`, where
-                padding elements are indicated by 1s.
-            need_weights (bool, optional): return the attention weights,
-                averaged over heads (default: False).
-            attn_mask (ByteTensor, optional): typically used to
-                implement causal attention, where the mask prevents the
-                attention from looking forward in time (default: None).
-            before_softmax (bool, optional): return the raw attention
-                weights and values before the attention softmax.
-            need_head_weights (bool, optional): return the attention
-                weights for each head. Implies *need_weights*. Default:
-                return the average attention weights over all heads.
-        """
-        if need_head_weights:
-            need_weights = True
-
-        is_tpu = query.device.type == "xla"
-
-        tgt_len, bsz, embed_dim = query.size()
-        src_len = tgt_len
-        assert embed_dim == self.embed_dim
-        assert list(query.size()) == [tgt_len, bsz, embed_dim]
-        if key is not None:
-            src_len, key_bsz, _ = key.size()
-            if not torch.jit.is_scripting():
-                assert key_bsz == bsz
-                assert value is not None
-                assert src_len, bsz == value.shape[:2]
-
-        if self.has_relative_attention_bias and position_bias is None:
-            position_bias = self.compute_bias(tgt_len, src_len)
-            position_bias = position_bias.unsqueeze(0).repeat(bsz, 1, 1, 1).view(bsz * self.num_heads, tgt_len, src_len)
-
-        if (
-                not is_tpu  # don't use PyTorch version on TPUs
-                and incremental_state is None
-                and not static_kv
-                # A workaround for quantization to work. Otherwise JIT compilation
-                # treats bias in linear module as method.
-                and not torch.jit.is_scripting()
-                and self.q_head_dim == self.head_dim
-        ):
-            assert key is not None and value is not None
-            assert attn_mask is None
-
-            attn_mask_rel_pos = None
-            if position_bias is not None:
-                attn_mask_rel_pos = position_bias
-                if self.gru_rel_pos:
-                    query_layer = query.transpose(0, 1)
-                    new_x_shape = query_layer.size()[:-1] + (self.num_heads, -1)
-                    query_layer = query_layer.view(*new_x_shape)
-                    query_layer = query_layer.permute(0, 2, 1, 3)
-                    _B, _H, _L, __ = query_layer.size()
-
-                    gate_a, gate_b = torch.sigmoid(self.grep_linear(query_layer).view(
-                        _B, _H, _L, 2, 4).sum(-1, keepdim=False)).chunk(2, dim=-1)
-                    gate_a_1 = gate_a * (gate_b * self.grep_a - 1.0) + 2.0
-                    attn_mask_rel_pos = gate_a_1.view(bsz * self.num_heads, -1, 1) * position_bias
-
-                attn_mask_rel_pos = attn_mask_rel_pos.view((-1, tgt_len, tgt_len))
-            k_proj_bias = self.k_proj.bias
-            if k_proj_bias is None:
-                k_proj_bias = torch.zeros_like(self.q_proj.bias)
-
-            x, attn = F.multi_head_attention_forward(
-                query,
-                key,
-                value,
-                self.embed_dim,
-                self.num_heads,
-                torch.empty([0]),
-                torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias)),
-                self.bias_k,
-                self.bias_v,
-                self.add_zero_attn,
-                self.dropout_module.p,
-                self.out_proj.weight,
-                self.out_proj.bias,
-                self.training,
-                # self.training or self.dropout_module.apply_during_inference,
-                key_padding_mask,
-                need_weights,
-                attn_mask_rel_pos,
-                use_separate_proj_weight=True,
-                q_proj_weight=self.q_proj.weight,
-                k_proj_weight=self.k_proj.weight,
-                v_proj_weight=self.v_proj.weight,
-            )
-            return x, attn, position_bias
-
-        if incremental_state is not None:
-            saved_state = self._get_input_buffer(incremental_state)
-            if saved_state is not None and "prev_key" in saved_state:
-                # previous time steps are cached - no need to recompute
-                # key and value if they are static
-                if static_kv:
-                    assert self.encoder_decoder_attention and not self.self_attention
-                    key = value = None
-        else:
-            saved_state = None
-
-        if self.self_attention:
-            q = self.q_proj(query)
-            k = self.k_proj(query)
-            v = self.v_proj(query)
-        elif self.encoder_decoder_attention:
-            # encoder-decoder attention
-            q = self.q_proj(query)
-            if key is None:
-                assert value is None
-                k = v = None
-            else:
-                k = self.k_proj(key)
-                v = self.v_proj(key)
-
-        else:
-            assert key is not None and value is not None
-            q = self.q_proj(query)
-            k = self.k_proj(key)
-            v = self.v_proj(value)
-        q *= self.scaling
-
-        if self.bias_k is not None:
-            assert self.bias_v is not None
-            k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
-            v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
-            if attn_mask is not None:
-                attn_mask = torch.cat(
-                    [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1
-                )
-            if key_padding_mask is not None:
-                key_padding_mask = torch.cat(
-                    [
-                        key_padding_mask,
-                        key_padding_mask.new_zeros(key_padding_mask.size(0), 1),
-                    ],
-                    dim=1,
-                )
-
-        q = (
-            q.contiguous()
-                .view(tgt_len, bsz * self.num_heads, self.q_head_dim)
-                .transpose(0, 1)
-        )
-        if k is not None:
-            k = (
-                k.contiguous()
-                    .view(-1, bsz * self.num_heads, self.k_head_dim)
-                    .transpose(0, 1)
-            )
-        if v is not None:
-            v = (
-                v.contiguous()
-                    .view(-1, bsz * self.num_heads, self.head_dim)
-                    .transpose(0, 1)
-            )
-
-        if saved_state is not None:
-            # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)
-            if "prev_key" in saved_state:
-                _prev_key = saved_state["prev_key"]
-                assert _prev_key is not None
-                prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)
-                if static_kv:
-                    k = prev_key
-                else:
-                    assert k is not None
-                    k = torch.cat([prev_key, k], dim=1)
-                src_len = k.size(1)
-            if "prev_value" in saved_state:
-                _prev_value = saved_state["prev_value"]
-                assert _prev_value is not None
-                prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)
-                if static_kv:
-                    v = prev_value
-                else:
-                    assert v is not None
-                    v = torch.cat([prev_value, v], dim=1)
-            prev_key_padding_mask: Optional[Tensor] = None
-            if "prev_key_padding_mask" in saved_state:
-                prev_key_padding_mask = saved_state["prev_key_padding_mask"]
-            assert k is not None and v is not None
-            key_padding_mask = MultiheadAttention._append_prev_key_padding_mask(
-                key_padding_mask=key_padding_mask,
-                prev_key_padding_mask=prev_key_padding_mask,
-                batch_size=bsz,
-                src_len=k.size(1),
-                static_kv=static_kv,
-            )
-
-            saved_state["prev_key"] = k.view(bsz, self.num_heads, -1, self.head_dim)
-            saved_state["prev_value"] = v.view(bsz, self.num_heads, -1, self.head_dim)
-            saved_state["prev_key_padding_mask"] = key_padding_mask
-            # In this branch incremental_state is never None
-            assert incremental_state is not None
-            incremental_state = self._set_input_buffer(incremental_state, saved_state)
-        assert k is not None
-        assert k.size(1) == src_len
-
-        # This is part of a workaround to get around fork/join parallelism
-        # not supporting Optional types.
-        if key_padding_mask is not None and key_padding_mask.dim() == 0:
-            key_padding_mask = None
-
-        if key_padding_mask is not None:
-            assert key_padding_mask.size(0) == bsz
-            assert key_padding_mask.size(1) == src_len
-
-        if self.add_zero_attn:
-            assert v is not None
-            src_len += 1
-            k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1)
-            v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1)
-            if attn_mask is not None:
-                attn_mask = torch.cat(
-                    [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1
-                )
-            if key_padding_mask is not None:
-                key_padding_mask = torch.cat(
-                    [
-                        key_padding_mask,
-                        torch.zeros(key_padding_mask.size(0), 1).type_as(
-                            key_padding_mask
-                        ),
-                    ],
-                    dim=1,
-                )
-
-        attn_weights = torch.bmm(q, k.transpose(1, 2))
-        attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)
-
-        assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]
-
-        if attn_mask is not None:
-            attn_mask = attn_mask.unsqueeze(0)
-            attn_weights += attn_mask
-
-        if key_padding_mask is not None:
-            # don't attend to padding symbols
-            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
-            if not is_tpu:
-                attn_weights = attn_weights.masked_fill(
-                    key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool),
-                    float("-inf"),
-                )
-            else:
-                attn_weights = attn_weights.transpose(0, 2)
-                attn_weights = attn_weights.masked_fill(key_padding_mask, float("-inf"))
-                attn_weights = attn_weights.transpose(0, 2)
-            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-
-        if before_softmax:
-            return attn_weights, v, position_bias
-
-        if position_bias is not None:
-            if self.gru_rel_pos == 1:
-                query_layer = q.view(bsz, self.num_heads, tgt_len, self.q_head_dim)
-                _B, _H, _L, __ = query_layer.size()
-                gate_a, gate_b = torch.sigmoid(self.grep_linear(query_layer).view(
-                    _B, _H, _L, 2, 4).sum(-1, keepdim=False)).chunk(2, dim=-1)
-                gate_a_1 = gate_a * (gate_b * self.grep_a - 1.0) + 2.0
-                position_bias = gate_a_1.view(bsz * self.num_heads, -1, 1) * position_bias
-
-            position_bias = position_bias.view(attn_weights.size())
-
-            attn_weights = attn_weights + position_bias
-
-        attn_weights_float = F.softmax(
-            attn_weights, dim=-1
-        )
-        attn_weights = attn_weights_float.type_as(attn_weights)
-        attn_probs = self.dropout_module(attn_weights)
-
-        assert v is not None
-        attn = torch.bmm(attn_probs, v)
-        assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
-        attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
-        attn = self.out_proj(attn)
-        attn_weights: Optional[Tensor] = None
-        if need_weights:
-            attn_weights = attn_weights_float.view(
-                bsz, self.num_heads, tgt_len, src_len
-            ).transpose(1, 0)
-            if not need_head_weights:
-                # average attention weights over heads
-                attn_weights = attn_weights.mean(dim=0)
-
-        return attn, attn_weights, position_bias
-
-    @staticmethod
-    def _append_prev_key_padding_mask(
-            key_padding_mask: Optional[Tensor],
-            prev_key_padding_mask: Optional[Tensor],
-            batch_size: int,
-            src_len: int,
-            static_kv: bool,
-    ) -> Optional[Tensor]:
-        # saved key padding masks have shape (bsz, seq_len)
-        if prev_key_padding_mask is not None and static_kv:
-            new_key_padding_mask = prev_key_padding_mask
-        elif prev_key_padding_mask is not None and key_padding_mask is not None:
-            new_key_padding_mask = torch.cat(
-                [prev_key_padding_mask.float(), key_padding_mask.float()], dim=1
-            )
-        # During incremental decoding, as the padding token enters and
-        # leaves the frame, there will be a time when prev or current
-        # is None
-        elif prev_key_padding_mask is not None:
-            if src_len > prev_key_padding_mask.size(1):
-                filler = torch.zeros(
-                    (batch_size, src_len - prev_key_padding_mask.size(1)),
-                    device=prev_key_padding_mask.device,
-                )
-                new_key_padding_mask = torch.cat(
-                    [prev_key_padding_mask.float(), filler.float()], dim=1
-                )
-            else:
-                new_key_padding_mask = prev_key_padding_mask.float()
-        elif key_padding_mask is not None:
-            if src_len > key_padding_mask.size(1):
-                filler = torch.zeros(
-                    (batch_size, src_len - key_padding_mask.size(1)),
-                    device=key_padding_mask.device,
-                )
-                new_key_padding_mask = torch.cat(
-                    [filler.float(), key_padding_mask.float()], dim=1
-                )
-            else:
-                new_key_padding_mask = key_padding_mask.float()
-        else:
-            new_key_padding_mask = prev_key_padding_mask
-        return new_key_padding_mask
-
-    def _get_input_buffer(
-            self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]]
-    ) -> Dict[str, Optional[Tensor]]:
-        result = self.get_incremental_state(incremental_state, "attn_state")
-        if result is not None:
-            return result
-        else:
-            empty_result: Dict[str, Optional[Tensor]] = {}
-            return empty_result
-
-    def _set_input_buffer(
-            self,
-            incremental_state: Dict[str, Dict[str, Optional[Tensor]]],
-            buffer: Dict[str, Optional[Tensor]],
-    ):
-        return self.set_incremental_state(incremental_state, "attn_state", buffer)
-
-    def apply_sparse_mask(self, attn_weights, tgt_len: int, src_len: int, bsz: int):
-        return attn_weights

manipulate_model/encoder_checkpoints/wavlm/WavLM-Base+.pt

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

manipulate_model/model.py

@@ -1,25 +0,0 @@
-import torch
-import torch.nn as nn
-
-from manipulate_model.encoder.encoder import Encoder
-from manipulate_model.decoder.decoder import Decoder
-
-
-class Model(nn.Module):
-    def __init__(self, config):
-        super(Model, self).__init__()
-        self.config = config
-
-        self.encoder = Encoder(self.config)
-        self.config.model.decoder.temporal_dim = self.encoder.get_temporal_dim()
-        self.config.model.decoder.encoding_dim = self.encoder.get_encoding_dim()
-        self.decoder = Decoder(self.config)
-
-    def forward(self, x):
-        if self.config.model.encoder_freeze:
-            with torch.no_grad():
-                x = self.encoder(x)
-        else:
-            x = self.encoder(x)
-        x = self.decoder(x)
-        return x

manipulate_model/utils.py

@@ -1,138 +0,0 @@
-import os
-import torch
-import torchaudio
-import numpy as np
-from omegaconf import OmegaConf
-from torchvision.io import read_video
-from torch.nn.functional import pad, normalize, softmax
-
-from manipulate_model.model import Model
-
-
-
-def get_config_and_model(model_root="manipulate_model/demo-model/audio"):
-    config_path = os.path.join(model_root, "config.yaml")
-    config = OmegaConf.load(config_path)
-    if isinstance(config.model.encoder, str):
-        config.model.encoder = OmegaConf.load(config.model.encoder)
-    if isinstance(config.model.decoder, str):
-        config.model.decoder = OmegaConf.load(config.model.decoder)
-
-    model = Model(config)
-    weights = torch.load(os.path.join(model_root, "weights.pt"))
-    model.load_state_dict(weights["model_state_dict"])
-
-    return config, model
-
-
-def load_audio(file_path, config):
-    # Load audio
-    # Parameters
-    # ----------
-    # file_path : str
-    #     Path to audio file
-    # Returns
-    # -------
-    # torch.Tensor
-
-    audio = None
-
-    if file_path.endswith(".wav") or file_path.endswith(".flac"):
-        audio, sample_rate = torchaudio.load(file_path)
-    elif file_path.endswith(".mp3"):
-        pass
-    elif file_path.endswith(".mp4"):
-        _, audio, _ = read_video(file_path)
-
-    return preprocess_audio(audio, config)
-
-
-def preprocess_audio(audio, config, step_size=1):
-    # Preprocess audio
-    # Parameters
-    # ----------
-    # audio : torch.Tensor
-    #     Audio signal
-    # config : OmegaConf
-    #     Configuration object
-    # Returns
-    # -------
-    # torch.Tensor : Normalized audio signal
-
-    window_size = config.data.window_size
-    sr = config.data.sr
-    fps = config.data.fps
-
-    audio_len = audio.shape[1]
-    step_size = step_size * (sr // fps)
-    window_size = window_size * (sr // fps)
-    audio = pad(audio, (window_size, window_size), "constant", 0)
-
-    sliced_audio = []
-
-    for i in range(0, audio_len + window_size, step_size):
-        audio_slice = audio[:, i : i + window_size]
-
-        if audio_slice.shape[1] < window_size:
-            audio_slice = pad(
-                audio_slice, (0, window_size - audio_slice.shape[1]), "constant", 0
-            )
-
-        audio_slice = normalize(audio_slice, dim=1)
-        sliced_audio.append(audio_slice)
-
-    sliced_audio = torch.stack(sliced_audio).squeeze()
-
-    return sliced_audio
-
-
-def infere(model, x, config, device="cpu", bs=8):
-    print(x)
-    model.eval()
-
-    x = load_audio(x, config)
-
-    # Inference (x is a stack of windows)
-    frame_predictions = []
-
-    with torch.no_grad():
-        n_iter = x.shape[0]
-
-        for i in range(0, n_iter, bs):
-            input_batch = x[i: i + bs]
-            input_batch = input_batch.to(device)
-
-            output = softmax(model(input_batch), dim=1)
-            frame_predictions.append(output.cpu().numpy())
-
-        frame_predictions = np.concatenate(frame_predictions, axis=0)[:,0]
-
-
-    return frame_predictions
-
-def convert_frame_predictions_to_timestamps(frame_predictions, fps, window_size):
-    # Convert frame predictions to timestamps
-    # Parameters
-    # ----------
-    # frame_predictions : np.ndarray
-    #     Frame predictions
-    # fps : int
-    #     Frames per second
-    # Returns
-    # -------
-    # np.ndarray : Timestamps
-
-    frame_predictions = (
-        frame_predictions[
-            int(window_size / 2) : -int(window_size / 2), 0
-        ]  # removes the padding, does not consider step size as of now
-        .round()
-        .astype(int)
-    )
-    timestamps = []
-
-    for i, frame_prediction in enumerate(frame_predictions):
-        if frame_prediction == 1:
-            timestamps.append(i / fps)
-
-    return timestamps

requirements.txt

@@ -1,6 +1,5 @@
 gradio==3.50.2
 torch==2.3.0
-torchaudio ==2.3.0
 fairseq @ git+https://github.com/facebookresearch/fairseq.git
 librosa==0.10.1
 numpy==1.24.4