Upload model files

app.py

@@ -1,7 +1,45 @@
 import gradio as gr
+import torch
+import numpy as np
+from modules.models import *
+from util import get_prompt_template
+from PIL import Image
+
 
 def greet(name):
     return "Hello " + name + "!!"
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+
+def main():
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+    # Get model
+    model_conf_file = f'./config/model/ACL_ViT16.yaml'
+    model = ACL(model_conf_file, device)
+    model.train(False)
+    model.load('./pretrain/Param_best.pth')
+
+    # Get placeholder text
+    prompt_template, text_pos_at_prompt, prompt_length = get_prompt_template()
+
+    # Input pre processing
+
+    # Inference
+    placeholder_tokens = model.get_placeholder_token(prompt_template.replace('{}', ''))
+    # audio_driven_embedding = model.encode_audio(audios.to(model.device), placeholder_tokens, text_pos_at_prompt,
+    #                                             prompt_length)
+
+    # Localization result
+    # out_dict = model(images.to(model.device), audio_driven_embedding, 352)
+    # seg = out_dict['heatmap'][j:j + 1]
+    # seg_image = ((1 - seg.squeeze().detach().cpu().numpy()) * 255).astype(np.uint8)
+    # seg_image = Image.fromarray(seg_image)
+    heatmap_image = cv2.applyColorMap(np.array(seg_image), cv2.COLORMAP_JET)
+    # overlaid_image = cv2.addWeighted(np.array(original_image), 0.5, heatmap_image, 0.5, 0)
+
+
+if __name__ == "__main__":
+    iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+    iface.launch()
+
+    main()

config/model/ACL_ViT16.yaml

@@ -0,0 +1,89 @@
+model:
+  clip: ViT16
+  vision_backbone: null
+  audio_backbone: BEATs
+  audio_proj: FGA512
+
+pretrain:
+  vision_backbone: null
+  audio_backbone: ./pretrain/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt2.pt
+  audio_proj: null
+
+fga_conf:
+  FGA:
+    input_size: 768
+    output_size: 768
+
+  FGA512:
+    input_size: 768
+    output_size: 512
+
+clip_conf:
+  RN50:
+    name: RN50
+    vision:
+      image_resolution: 224
+      vision_layers: [3, 4, 6, 3]
+      vision_width: 64
+      heads: 8
+      vision_patch_size: null
+    text:
+      transformer_layers: 12
+      transformer_width: 512
+      transformer_heads: 8
+      vocab_size: 49408
+      context_length: 77
+    embedding_dim: 1024
+
+  ViT16:
+    name: ViT-B/16
+    vision:
+      image_resolution: 224
+      vision_layers: 12
+      vision_width: 768
+      heads: 12
+      vision_patch_size: 16
+    text:
+      transformer_layers: 12
+      transformer_width: 512
+      transformer_heads: 8
+      vocab_size: 49408
+      context_length: 77
+    embedding_dim: 512
+
+  ViT14:
+    name: ViT-L/14
+    vision:
+      image_resolution: 224
+      vision_layers: 24
+      vision_width: 1024
+      heads: 16
+      vision_patch_size: 14
+    text:
+      transformer_layers: 12
+      transformer_width: 768
+      transformer_heads: 12
+      vocab_size: 49408
+      context_length: 77
+    embedding_dim: 768
+
+vision_backbone_conf:
+  maskclip_plus_rn50_512:
+    name: maskclip_plus_rn50_512
+    image_resolution: 512
+    vision_layers: [ 3, 4, 6, 3 ]
+    vision_width: 2048
+    aspp:
+      dilations: [ 6, 12, 18, 24 ]
+      in_channels: 2048
+      channels: 512
+
+  maskclip_plus_rn101_512:
+    name: maskclip_plus_rn101_512
+    image_resolution: 512
+    vision_layers: [ 3, 4, 23, 3 ]
+    vision_width: 2048
+    aspp:
+      dilations: [ 6, 12, 18, 24 ]
+      in_channels: 2048
+      channels: 1024

config/train/Exp_ACL_v1.yaml

@@ -0,0 +1,45 @@
+model: ACL
+
+common:
+  train_data: vggss
+  epoch: 20
+  batch_size: 8
+  input_resolution: 352
+  num_workers: 4
+  seed: 0
+  loss:
+    - acl_i
+    - acl_f
+    - area_reg
+  loss_w:
+    - 1
+    - 1
+    - 1
+
+optimizer: Adam
+scheduler: null
+amp: True
+
+optim_conf:
+  Adam:
+    module_path: torch.optim
+    module_name: Adam
+    lr: 0.0001
+    weight_decay: 0.0001
+
+  AdamW:
+    module_path: torch.optim
+    module_name: AdamW
+    lr: 0.001
+
+  SGDR:
+    module_path: torch.optim
+    module_name: SGD
+    lr: 0.5
+    weight_decay: 0.00001
+
+sched_conf:
+  Cosine:
+    module_path: torch.optim.lr_scheduler
+    module_name: CosineAnnealingLR
+    eta_ratio: 0.0

modules/AudioToken/AudioToken.py

@@ -0,0 +1,100 @@
+import torch
+from diffusers.loaders import AttnProcsLayers
+
+from modules.BEATs.BEATs import BEATs, BEATsConfig
+from modules.AudioToken.embedder import FGAEmbedder
+from diffusers import AutoencoderKL, UNet2DConditionModel
+from diffusers.models.attention_processor import LoRAAttnProcessor
+
+
+class AudioTokenWrapper(torch.nn.Module):
+    """Simple wrapper module for Stable Diffusion that holds all the models together"""
+
+    def __init__(
+        self,
+        args,
+        accelerator,
+    ):
+
+        super().__init__()
+        # Load scheduler and models
+        from modules.clip_text_model.modeling_clip import CLIPTextModel
+        self.text_encoder = CLIPTextModel.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
+        )
+        self.unet = UNet2DConditionModel.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
+        )
+        self.vae = AutoencoderKL.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision
+        )
+
+        checkpoint = torch.load(
+            'models/BEATs/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt2.pt')
+        cfg = BEATsConfig(checkpoint['cfg'])
+        self.aud_encoder = BEATs(cfg)
+        self.aud_encoder.load_state_dict(checkpoint['model'])
+        self.aud_encoder.predictor = None
+        input_size = 768 * 3
+
+        if args.pretrained_model_name_or_path == "CompVis/stable-diffusion-v1-4":
+            self.embedder = FGAEmbedder(input_size=input_size, output_size=768)
+
+        else:
+            self.embedder = FGAEmbedder(input_size=input_size, output_size=1024)
+
+        self.vae.eval()
+        self.unet.eval()
+        self.text_encoder.eval()
+        self.aud_encoder.eval()
+
+        if 'lora' in args and args.lora:
+            # Set correct lora layers
+            lora_attn_procs = {}
+            for name in self.unet.attn_processors.keys():
+                cross_attention_dim = None if name.endswith(
+                    "attn1.processor") else self.unet.config.cross_attention_dim
+                if name.startswith("mid_block"):
+                    hidden_size = self.unet.config.block_out_channels[-1]
+                elif name.startswith("up_blocks"):
+                    block_id = int(name[len("up_blocks.")])
+                    hidden_size = list(reversed(self.unet.config.block_out_channels))[block_id]
+                elif name.startswith("down_blocks"):
+                    block_id = int(name[len("down_blocks.")])
+                    hidden_size = self.unet.config.block_out_channels[block_id]
+
+                lora_attn_procs[name] = LoRAAttnProcessor(hidden_size=hidden_size,
+                                                          cross_attention_dim=cross_attention_dim)
+
+            self.unet.set_attn_processor(lora_attn_procs)
+            self.lora_layers = AttnProcsLayers(self.unet.attn_processors)
+
+        if args.data_set == 'train':
+
+            # Freeze vae, unet, text_enc and aud_encoder
+            self.vae.requires_grad_(False)
+            self.unet.requires_grad_(False)
+            self.text_encoder.requires_grad_(False)
+            self.aud_encoder.requires_grad_(False)
+            self.embedder.requires_grad_(True)
+            self.embedder.train()
+
+            if 'lora' in args and args.lora:
+                self.unet.train()
+
+        if args.data_set == 'test':
+
+            from transformers import CLIPTextModel
+            self.text_encoder = CLIPTextModel.from_pretrained(
+                args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
+            )
+
+            self.embedder.eval()
+            embedder_learned_embeds = args.learned_embeds
+            self.embedder.load_state_dict(torch.load(embedder_learned_embeds, map_location=accelerator.device))
+
+            if 'lora' in args and args.lora:
+                self.lora_layers.eval()
+                lora_layers_learned_embeds = args.lora_learned_embeds
+                self.lora_layers.load_state_dict(torch.load(lora_layers_learned_embeds, map_location=accelerator.device))
+                self.unet.load_attn_procs(lora_layers_learned_embeds)

modules/AudioToken/embedder.py

@@ -0,0 +1,17 @@
+import torch.nn as nn
+from modules.FGA.atten import Atten
+
+class FGAEmbedder(nn.Module):
+    def __init__(self, input_size=768*3, output_size=768):
+        super(FGAEmbedder, self).__init__()
+        self.fc1 = nn.Linear(input_size, input_size)
+        self.fc2 = nn.Linear(input_size, output_size)
+        self.gelu = nn.GELU()
+        self.fga = Atten(util_e=[output_size], pairwise_flag=False)
+
+    def forward(self, audio_embs):
+        audio_embs = self.fc1(audio_embs)
+        audio_embs = self.gelu(audio_embs)
+        audio_embs = self.fc2(audio_embs)
+        attend = self.fga([audio_embs])[0]
+        return attend

modules/BEATs/BEATs.py

@@ -0,0 +1,180 @@
+# --------------------------------------------------------
+# BEATs: Audio Pre-Training with Acoustic Tokenizers (https://arxiv.org/abs/2212.09058)
+# Github source: https://github.com/microsoft/unilm/tree/master/beats
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Based on fairseq code bases
+# https://github.com/pytorch/fairseq
+# --------------------------------------------------------
+
+
+import torch
+import torch.nn as nn
+from torch.nn import LayerNorm
+import torchaudio.compliance.kaldi as ta_kaldi
+from torch.cuda.amp import autocast
+
+from modules.BEATs.backbone import (
+    TransformerEncoder,
+)
+
+import logging
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+class BEATsConfig:
+    def __init__(self, cfg=None):
+        self.input_patch_size: int = -1  # path size of patch embedding
+        self.embed_dim: int = 512  # patch embedding dimension
+        self.conv_bias: bool = False  # include bias in conv encoder
+
+        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_wise_gradient_decay_ratio: float = 1.0  # ratio for layer-wise gradient decay
+        self.layer_norm_first: bool = False  # apply layernorm first in the transformer
+        self.deep_norm: bool = False  # apply deep_norm first in the transformer
+
+        # 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)
+
+        # 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
+
+        # label predictor
+        self.finetuned_model: bool = False  # whether the model is a fine-tuned model.
+        self.predictor_dropout: float = 0.1  # dropout probability for the predictor
+        self.predictor_class: int = 527  # target class number for the predictor
+
+        if cfg is not None:
+            self.update(cfg)
+
+    def update(self, cfg: dict):
+        self.__dict__.update(cfg)
+
+
+class BEATs(nn.Module):
+    def __init__(
+            self,
+            cfg: BEATsConfig,
+    ) -> None:
+        super().__init__()
+        logger.info(f"BEATs Config: {cfg.__dict__}")
+
+        self.cfg = cfg
+
+        self.embed = cfg.embed_dim
+        self.post_extract_proj = (
+            nn.Linear(self.embed, cfg.encoder_embed_dim)
+            if self.embed != cfg.encoder_embed_dim
+            else None
+        )
+
+        self.input_patch_size = cfg.input_patch_size
+        self.patch_embedding = nn.Conv2d(1, self.embed, kernel_size=self.input_patch_size, stride=self.input_patch_size,
+                                         bias=cfg.conv_bias)
+
+        self.dropout_input = nn.Dropout(cfg.dropout_input)
+
+        assert not cfg.deep_norm or not cfg.layer_norm_first
+        self.encoder = TransformerEncoder(cfg)
+        self.layer_norm = LayerNorm(self.embed)
+
+        if cfg.finetuned_model:
+            self.predictor_dropout = nn.Dropout(cfg.predictor_dropout)
+            self.predictor = nn.Linear(cfg.encoder_embed_dim, cfg.predictor_class)
+        else:
+            self.predictor = None
+
+    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
+
+    @autocast(enabled=False)
+    def preprocess(
+            self,
+            source: torch.Tensor,
+            fbank_mean: float = 15.41663,
+            fbank_std: float = 6.55582,
+    ) -> torch.Tensor:
+        fbanks = []
+        for waveform in source:
+            waveform = waveform.unsqueeze(0) * 2 ** 15
+            fbank = ta_kaldi.fbank(waveform, num_mel_bins=128, sample_frequency=16000, frame_length=25, frame_shift=10)
+            fbanks.append(fbank)
+        fbank = torch.stack(fbanks, dim=0)
+        fbank = (fbank - fbank_mean) / (2 * fbank_std)
+        return fbank
+
+    def extract_features(
+            self,
+            source: torch.Tensor,
+            padding_mask: Optional[torch.Tensor] = None,
+            fbank_mean: float = 15.41663,
+            fbank_std: float = 6.55582,
+    ):
+        fbank = self.preprocess(source, fbank_mean=fbank_mean, fbank_std=fbank_std)
+        if padding_mask is not None:
+            padding_mask = self.forward_padding_mask(fbank, padding_mask)
+        # ToDo Aug here
+        fbank = fbank.unsqueeze(1)
+        features = self.patch_embedding(fbank)
+        features = features.reshape(features.shape[0], features.shape[1], -1)
+        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)
+
+        x = self.dropout_input(features)
+
+        x, layers_sum, layers = self.encoder(
+            x,
+            padding_mask=padding_mask,
+        )
+
+        if self.predictor is not None:
+            x = self.predictor_dropout(x)
+            logits = self.predictor(x)
+
+            if padding_mask is not None and padding_mask.any():
+                logits[padding_mask] = 0
+                logits = logits.sum(dim=1)
+                logits = logits / (~padding_mask).sum(dim=1).unsqueeze(-1).expand_as(logits)
+            else:
+                logits = logits.mean(dim=1)
+
+            lprobs = torch.sigmoid(logits)
+
+            return lprobs, padding_mask
+        else:
+            return x, layers_sum, layers

modules/BEATs/Tokenizers.py

@@ -0,0 +1,172 @@
+# --------------------------------------------------------
+# BEATs: Audio Pre-Training with Acoustic Tokenizers (https://arxiv.org/abs/2212.09058)
+# Github source: https://github.com/microsoft/unilm/tree/master/beats
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Based on fairseq code bases
+# https://github.com/pytorch/fairseq
+# --------------------------------------------------------
+
+
+import torch
+import torch.nn as nn
+from torch.nn import LayerNorm
+import torchaudio.compliance.kaldi as ta_kaldi
+
+from modules.BEATs.backbone import (
+    TransformerEncoder,
+)
+from modules.BEATs.quantizer import (
+    NormEMAVectorQuantizer,
+)
+
+import logging
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+class TokenizersConfig:
+    def __init__(self, cfg=None):
+        self.input_patch_size: int = -1  # path size of patch embedding
+        self.embed_dim: int = 512  # patch embedding dimension
+        self.conv_bias: bool = False  # include bias in conv encoder
+
+        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.deep_norm: bool = False  # apply deep_norm first in the transformer
+
+        # 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)
+
+        # 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
+
+        # quantizer
+        self.quant_n: int = 1024 # codebook number in quantizer
+        self.quant_dim: int = 256    # codebook dimension in quantizer
+
+        if cfg is not None:
+            self.update(cfg)
+
+    def update(self, cfg: dict):
+        self.__dict__.update(cfg)
+
+
+class Tokenizers(nn.Module):
+    def __init__(
+            self,
+            cfg: TokenizersConfig,
+    ) -> None:
+        super().__init__()
+        logger.info(f"Tokenizers Config: {cfg.__dict__}")
+
+        self.cfg = cfg
+
+        self.embed = cfg.embed_dim
+        self.post_extract_proj = (
+            nn.Linear(self.embed, cfg.encoder_embed_dim)
+            if self.embed != cfg.encoder_embed_dim
+            else None
+        )
+
+        self.input_patch_size = cfg.input_patch_size
+        self.patch_embedding = nn.Conv2d(1, self.embed, kernel_size=self.input_patch_size, stride=self.input_patch_size,
+                                         bias=cfg.conv_bias)
+
+        self.dropout_input = nn.Dropout(cfg.dropout_input)
+
+        assert not cfg.deep_norm or not cfg.layer_norm_first
+        self.encoder = TransformerEncoder(cfg)
+        self.layer_norm = LayerNorm(self.embed)
+
+        self.quantize = NormEMAVectorQuantizer(
+            n_embed=cfg.quant_n, embedding_dim=cfg.quant_dim, beta=1.0, kmeans_init=True, decay=0.99,
+        )
+        self.quant_n = cfg.quant_n
+        self.quantize_layer = nn.Sequential(
+            nn.Linear(cfg.encoder_embed_dim, cfg.encoder_embed_dim),
+            nn.Tanh(),
+            nn.Linear(cfg.encoder_embed_dim, cfg.quant_dim)  # for quantize
+        )
+
+    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 preprocess(
+            self,
+            source: torch.Tensor,
+            fbank_mean: float = 15.41663,
+            fbank_std: float = 6.55582,
+    ) -> torch.Tensor:
+        fbanks = []
+        for waveform in source:
+            waveform = waveform.unsqueeze(0) * 2 ** 15
+            fbank = ta_kaldi.fbank(waveform, num_mel_bins=128, sample_frequency=16000, frame_length=25, frame_shift=10)
+            fbanks.append(fbank)
+        fbank = torch.stack(fbanks, dim=0)
+        fbank = (fbank - fbank_mean) / (2 * fbank_std)
+        return fbank
+
+    def extract_labels(
+            self,
+            source: torch.Tensor,
+            padding_mask: Optional[torch.Tensor] = None,
+            fbank_mean: float = 15.41663,
+            fbank_std: float = 6.55582,
+    ):
+        fbank = self.preprocess(source, fbank_mean=fbank_mean, fbank_std=fbank_std)
+
+        if padding_mask is not None:
+            padding_mask = self.forward_padding_mask(fbank, padding_mask)
+
+        fbank = fbank.unsqueeze(1)
+        features = self.patch_embedding(fbank)
+        features = features.reshape(features.shape[0], features.shape[1], -1)
+        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)
+
+        x = self.dropout_input(features)
+
+        x, layer_results = self.encoder(
+            x,
+            padding_mask=padding_mask,
+        )
+
+        quantize_input = self.quantize_layer(x)
+        quantize_feature, embed_loss, embed_ind = self.quantize(quantize_input)
+
+        return embed_ind

modules/BEATs/backbone.py

@@ -0,0 +1,788 @@
+# --------------------------------------------------------
+# BEATs: Audio Pre-Training with Acoustic Tokenizers (https://arxiv.org/abs/2212.09058)
+# Github source: https://github.com/microsoft/unilm/tree/master/beats
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Based on fairseq code bases
+# https://github.com/pytorch/fairseq
+# --------------------------------------------------------
+
+import math
+import numpy as np
+from typing import Dict, Optional, Tuple
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+from torch.nn import LayerNorm, Parameter
+from modules.BEATs.modules import (
+    GradMultiply,
+    SamePad,
+    get_activation_fn,
+    GLU_Linear,
+    quant_noise,
+)
+
+
+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,
+                    deep_norm=args.deep_norm,
+                    has_relative_attention_bias=self.relative_position_embedding,
+                    num_buckets=self.num_buckets,
+                    max_distance=self.max_distance,
+                    gru_rel_pos=args.gru_rel_pos,
+                    encoder_layers=args.encoder_layers,
+                )
+                for i in range(args.encoder_layers)
+            ]
+        )
+        if self.relative_position_embedding:
+            for i in range(1, args.encoder_layers):
+                del self.layers[i].self_attn.relative_attention_bias
+                self.layers[i].self_attn.relative_attention_bias = self.layers[0].self_attn.relative_attention_bias
+
+        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)
+
+        if args.deep_norm:
+            deep_norm_beta = math.pow(8 * args.encoder_layers, -1 / 4)
+            for i in range(args.encoder_layers):
+                nn.init.xavier_normal_(self.layers[i].self_attn.k_proj.weight, gain=1)
+                nn.init.xavier_normal_(self.layers[i].self_attn.v_proj.weight, gain=deep_norm_beta)
+                nn.init.xavier_normal_(self.layers[i].self_attn.q_proj.weight, gain=1)
+                nn.init.xavier_normal_(self.layers[i].self_attn.out_proj.weight, gain=deep_norm_beta)
+                nn.init.xavier_normal_(self.layers[i].fc1.weight, gain=deep_norm_beta)
+                nn.init.xavier_normal_(self.layers[i].fc2.weight, gain=deep_norm_beta)
+
+        self.layer_wise_gradient_decay_ratio = getattr(args, "layer_wise_gradient_decay_ratio", 1)
+
+    def forward(self, x, padding_mask=None, layer=None):
+        x, layers_sum, layers = self.extract_features(x, padding_mask, layer)
+
+        if self.layer_norm_first and layer is None:
+            x = self.layer_norm(x)
+
+        return x, layers_sum, layers
+
+    def extract_features(self, x, padding_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_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)
+        layers = []
+
+        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):
+            if self.layer_wise_gradient_decay_ratio != 1.0:
+                x = GradMultiply.apply(x, self.layer_wise_gradient_decay_ratio)
+            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, pos_bias=pos_bias)
+            if tgt_layer is not None:
+                layer_results.append((x, z))
+            if i == tgt_layer:
+                r = x
+                break
+            if i in [3, 7, 11]:
+                layers.append(x.transpose(0, 1))
+
+        if r is not None:
+            x = r
+
+        # T x B x C -> B x T x C
+        x = x.transpose(0, 1)
+        layers_cat = torch.cat(layers, dim=2)
+        # layers = layers[0] + layers[1] + layers[2]
+
+        return x, layers_cat, layers
+
+
+class TransformerSentenceEncoderLayer(nn.Module):
+    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,
+            deep_norm: 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,
+            encoder_layers: int = 0,
+    ) -> None:
+
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.dropout = dropout
+        self.activation_dropout = activation_dropout
+
+        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
+
+        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)
+
+        self.final_layer_norm = LayerNorm(self.embedding_dim)
+
+        self.deep_norm = deep_norm
+        if self.deep_norm:
+            self.deep_norm_alpha = math.pow(2 * encoder_layers, 1 / 4)
+        else:
+            self.deep_norm_alpha = 1
+
+    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
+    ):
+        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 * self.deep_norm_alpha + 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 * self.deep_norm_alpha + x
+            x = self.final_layer_norm(x)
+
+        return x, attn, pos_bias
+
+
+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 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
+        alpha = 32
+        q *= 1 / alpha
+
+        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 = (attn_weights - attn_weights.max(dim=-1, keepdim=True)[0]) * alpha
+        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:
+            attn_mask_rel_pos = position_bias
+            if self.gru_rel_pos == 1:
+                query_layer = q.view(bsz, self.num_heads, tgt_len, self.q_head_dim) * alpha / self.scaling
+                _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, tgt_len, 1) * position_bias
+
+            attn_mask_rel_pos = attn_mask_rel_pos.view(attn_weights.size())
+
+            attn_weights = attn_weights + attn_mask_rel_pos
+
+        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
+
+
+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)

modules/BEATs/modules.py

@@ -0,0 +1,218 @@
+# --------------------------------------------------------
+# BEATs: Audio Pre-Training with Acoustic Tokenizers (https://arxiv.org/abs/2212.09058)
+# Github source: https://github.com/microsoft/unilm/tree/master/beats
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Based on fairseq code bases
+# https://github.com/pytorch/fairseq
+# --------------------------------------------------------
+
+import math
+import warnings
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+
+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):
+    def __init__(self):
+        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 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

modules/BEATs/quantizer.py

@@ -0,0 +1,215 @@
+# --------------------------------------------------------
+# BEATs: Audio Pre-Training with Acoustic Tokenizers (https://arxiv.org/abs/2212.09058)
+# Github source: https://github.com/microsoft/unilm/tree/master/beats
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Based on VQGAN code bases
+# https://github.com/CompVis/taming-transformers
+# --------------------------------------------------------'
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as distributed
+
+try:
+    from einops import rearrange, repeat
+except ImportError:
+    pass
+
+
+def l2norm(t):
+    return F.normalize(t, p=2, dim=-1)
+
+
+def ema_inplace(moving_avg, new, decay):
+    moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
+
+
+def sample_vectors(samples, num):
+    num_samples, device = samples.shape[0], samples.device
+
+    if num_samples >= num:
+        indices = torch.randperm(num_samples, device=device)[:num]
+    else:
+        indices = torch.randint(0, num_samples, (num,), device=device)
+
+    return samples[indices]
+
+
+def kmeans(samples, num_clusters, num_iters=10, use_cosine_sim=False):
+    dim, dtype, device = samples.shape[-1], samples.dtype, samples.device
+
+    means = sample_vectors(samples, num_clusters)
+
+    for _ in range(num_iters):
+        if use_cosine_sim:
+            dists = samples @ means.t()
+        else:
+            diffs = rearrange(samples, 'n d -> n () d') \
+                    - rearrange(means, 'c d -> () c d')
+            dists = -(diffs ** 2).sum(dim=-1)
+
+        buckets = dists.max(dim=-1).indices
+        bins = torch.bincount(buckets, minlength=num_clusters)
+        zero_mask = bins == 0
+        bins_min_clamped = bins.masked_fill(zero_mask, 1)
+
+        new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)
+        new_means.scatter_add_(0, repeat(buckets, 'n -> n d', d=dim), samples)
+        new_means = new_means / bins_min_clamped[..., None]
+
+        if use_cosine_sim:
+            new_means = l2norm(new_means)
+
+        means = torch.where(zero_mask[..., None], means, new_means)
+
+    return means, bins
+
+
+class EmbeddingEMA(nn.Module):
+    def __init__(self, num_tokens, codebook_dim, decay=0.99, eps=1e-5, kmeans_init=True, codebook_init_path=''):
+        super().__init__()
+        self.num_tokens = num_tokens
+        self.codebook_dim = codebook_dim
+        self.decay = decay
+        self.eps = eps
+        if codebook_init_path == '':
+            if not kmeans_init:
+                weight = torch.randn(num_tokens, codebook_dim)
+                weight = l2norm(weight)
+            else:
+                weight = torch.zeros(num_tokens, codebook_dim)
+            self.register_buffer('initted', torch.Tensor([not kmeans_init]))
+        else:
+            print(f"load init codebook weight from {codebook_init_path}")
+            codebook_ckpt_weight = torch.load(codebook_init_path, map_location='cpu')
+            weight = codebook_ckpt_weight.clone()
+            self.register_buffer('initted', torch.Tensor([True]))
+
+        self.weight = nn.Parameter(weight, requires_grad=False)
+        self.cluster_size = nn.Parameter(torch.zeros(num_tokens), requires_grad=False)
+        self.embed_avg = nn.Parameter(weight.clone(), requires_grad=False)
+        # self.register_buffer('initted', torch.Tensor([not kmeans_init]))
+        self.update = True
+
+    @torch.jit.ignore
+    def init_embed_(self, data):
+        if self.initted:
+            return
+        print("Performing Kemans init for codebook")
+        embed, cluster_size = kmeans(data, self.num_tokens, 10, use_cosine_sim=True)
+        self.weight.data.copy_(embed)
+        self.cluster_size.data.copy_(cluster_size)
+        self.initted.data.copy_(torch.Tensor([True]))
+
+    def forward(self, embed_id):
+        return F.embedding(embed_id, self.weight)
+
+    def cluster_size_ema_update(self, new_cluster_size):
+        self.cluster_size.data.mul_(self.decay).add_(new_cluster_size, alpha=1 - self.decay)
+
+    def embed_avg_ema_update(self, new_embed_avg):
+        self.embed_avg.data.mul_(self.decay).add_(new_embed_avg, alpha=1 - self.decay)
+
+    def weight_update(self, num_tokens):
+        n = self.cluster_size.sum()
+        smoothed_cluster_size = (
+                (self.cluster_size + self.eps) / (n + num_tokens * self.eps) * n
+        )
+        # normalize embedding average with smoothed cluster size
+        embed_normalized = self.embed_avg / smoothed_cluster_size.unsqueeze(1)
+        # embed_normalized = l2norm(self.embed_avg / smoothed_cluster_size.unsqueeze(1))
+        self.weight.data.copy_(embed_normalized)
+
+
+def norm_ema_inplace(moving_avg, new, decay):
+    moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
+    moving_avg.data.copy_(l2norm(moving_avg.data))
+
+
+class NormEMAVectorQuantizer(nn.Module):
+    def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5,
+                 statistic_code_usage=True, kmeans_init=False, codebook_init_path=''):
+        super().__init__()
+        self.codebook_dim = embedding_dim
+        self.num_tokens = n_embed
+        self.beta = beta
+        self.decay = decay
+
+        # learnable = True if orthogonal_reg_weight > 0 else False
+        self.embedding = EmbeddingEMA(self.num_tokens, self.codebook_dim, decay, eps, kmeans_init, codebook_init_path)
+
+        self.statistic_code_usage = statistic_code_usage
+        if statistic_code_usage:
+            self.register_buffer('cluster_size', torch.zeros(n_embed))
+        if distributed.is_available() and distributed.is_initialized():
+            print("ddp is enable, so use ddp_reduce to sync the statistic_code_usage for each gpu!")
+            self.all_reduce_fn = distributed.all_reduce
+        else:
+            self.all_reduce_fn = nn.Identity()
+
+    def reset_cluster_size(self, device):
+        if self.statistic_code_usage:
+            self.register_buffer('cluster_size', torch.zeros(self.num_tokens))
+            self.cluster_size = self.cluster_size.to(device)
+
+    def forward(self, z):
+        # reshape z -> (batch, height, width, channel) and flatten
+        # z, 'b c h w -> b h w c'
+        # z = rearrange(z, 'b c h w -> b h w c')
+        # z = z.transpose(1, 2)
+        z = l2norm(z)
+        z_flattened = z.reshape(-1, self.codebook_dim)
+
+        self.embedding.init_embed_(z_flattened)
+
+        d = z_flattened.pow(2).sum(dim=1, keepdim=True) + \
+            self.embedding.weight.pow(2).sum(dim=1) - 2 * \
+            torch.einsum('bd,nd->bn', z_flattened, self.embedding.weight)  # 'n d -> d n'
+
+        encoding_indices = torch.argmin(d, dim=1)
+
+        z_q = self.embedding(encoding_indices).view(z.shape)
+
+        encodings = F.one_hot(encoding_indices, self.num_tokens).type(z.dtype)
+
+        if not self.training:
+            with torch.no_grad():
+                cluster_size = encodings.sum(0)
+                self.all_reduce_fn(cluster_size)
+                ema_inplace(self.cluster_size, cluster_size, self.decay)
+
+        if self.training and self.embedding.update:
+            # EMA cluster size
+
+            bins = encodings.sum(0)
+            self.all_reduce_fn(bins)
+
+            # self.embedding.cluster_size_ema_update(bins)
+            ema_inplace(self.cluster_size, bins, self.decay)
+
+            zero_mask = (bins == 0)
+            bins = bins.masked_fill(zero_mask, 1.)
+
+            embed_sum = z_flattened.t() @ encodings
+            self.all_reduce_fn(embed_sum)
+
+            embed_normalized = (embed_sum / bins.unsqueeze(0)).t()
+            embed_normalized = l2norm(embed_normalized)
+
+            embed_normalized = torch.where(zero_mask[..., None], self.embedding.weight,
+                                           embed_normalized)
+            norm_ema_inplace(self.embedding.weight, embed_normalized, self.decay)
+
+        # compute loss for embedding
+        loss = self.beta * F.mse_loss(z_q.detach(), z)
+
+        # preserve gradients
+        z_q = z + (z_q - z).detach()
+
+        # reshape back to match original input shape
+        # z_q, 'b h w c -> b c h w'
+        # z_q = rearrange(z_q, 'b h w c -> b c h w')
+        # z_q = z_q.transpose(1, 2)
+        return z_q, loss, encoding_indices

modules/CLIPSeg/clipseg_for_audio.py

@@ -0,0 +1,182 @@
+import transformers
+import torch
+import torch.nn.functional as F
+from torch import nn
+from typing import List, Tuple, Union, Optional
+import numpy as np
+from transformers.models.clipseg.modeling_clipseg import _expand_mask
+
+
+class CLIPSeg(transformers.CLIPSegForImageSegmentation):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def encode_text(self, text: torch.Tensor) -> torch.Tensor:
+        """
+        Encode textual input and return the text embeddings.
+
+        Args:
+            text (torch.Tensor): Input text tensor.
+
+        Returns:
+            torch.Tensor: Text embeddings.
+        """
+        tokens = text
+        if text.ndim == 3:
+            tokens = torch.squeeze(text, dim=1)
+        non_zero_index = torch.nonzero(tokens.sum(axis=0) == 0)[0]
+        input_ids = tokens[:, :non_zero_index]
+        attention_mask = (input_ids > 0).to(tokens.dtype)
+        input_ids += torch.max(input_ids) * (1 - attention_mask)
+        conditional_embeddings = self.clip.get_text_features(input_ids, attention_mask=attention_mask,
+                                                             position_ids=None)
+
+        return conditional_embeddings
+
+    def similarity(self, image: torch.Tensor, embeddings: List[torch.Tensor]) -> torch.Tensor:
+        """
+        Calculate the similarity score between an image and a list of embeddings.
+
+        Args:
+            image (torch.Tensor): Input image tensor of shape (B, C, H, W).
+            embeddings (List[torch.Tensor]): List of N embedding tensors of shape (dim,).
+
+        Returns:
+            torch.Tensor: Similarity scores of shape (B, N) for each batch.
+        """
+        B, c, h, w = image.shape
+        if (h, w) != (352, 352):
+            vision_outputs = self.clip.vision_model(pixel_values=F.interpolate(image, 352, mode='bicubic'),
+                                                    output_attentions=False,
+                                                    output_hidden_states=False,
+                                                    return_dict=False)
+            img_embedding = self.clip.visual_projection(vision_outputs[1])
+        else:
+            vision_outputs = self.clip.vision_model(pixel_values=image,
+                                                    output_attentions=False,
+                                                    output_hidden_states=False,
+                                                    return_dict=False)
+            img_embedding = self.clip.visual_projection(vision_outputs[1])
+
+        paired_embedding = torch.cat(embeddings, dim=0)
+        paired_embedding = paired_embedding.repeat(B, 1)  # Batch-wise replication of embeddings
+        paired_embedding = paired_embedding.view(B, -1, img_embedding.size(-1))
+
+        result = torch.matmul(F.normalize(paired_embedding, dim=-1), F.normalize(img_embedding, dim=-1).unsqueeze(-1))
+        result = result.squeeze(-1).view(B, -1)
+        return F.softmax(result, dim=-1)
+
+    def encode_audio(self, placeholder_token: torch.Tensor, audio_token: torch.Tensor, pos: int,
+                     length: int) -> torch.Tensor:
+        """
+        Encode audio token into the audio-driven embeddings. (Audio-Driven Embedder)
+
+        Args:
+            placeholder_token (torch.Tensor): Placeholder text token tensor.
+            audio_token (torch.Tensor): Audio token tensor.
+            pos (int): Position index for audio token.
+            length (int): Length of the input token.
+
+        Returns:
+            torch.Tensor: Audio-driven embeddings.
+
+        Reference:
+            "Can CLIP Help Sound Source Localization?" WACV 2024
+            - https://arxiv.org/abs/2311.04066
+        """
+        tokens = placeholder_token
+        if placeholder_token.ndim == 3:
+            tokens = torch.squeeze(placeholder_token, dim=1)
+
+        inputs_embeds = self.clip.text_model.embeddings.token_embedding(tokens).type(
+            self.dtype)  # [batch_size, n_ctx, d_model]
+        inputs_embeds = torch.cat((inputs_embeds[:, :pos, :], audio_token, inputs_embeds[:, pos:, :]),
+                                  dim=1)  # Inject Audio token
+        inputs_embeds = inputs_embeds[:, :length, :]
+
+        bsz, seq_len, _ = inputs_embeds.shape
+        attention_mask = torch.ones((bsz, seq_len)).to(placeholder_token.device)
+        position_ids = torch.arange(length).unsqueeze(0).to(placeholder_token.device)
+
+        position_embeddings = self.clip.text_model.embeddings.position_embedding(position_ids)
+        hidden_states = inputs_embeds + position_embeddings
+
+        bsz, seq_len, _ = inputs_embeds.shape
+        # CLIPSeg's text model uses causal mask, prepare it here.
+        # https://github.com/openai/CLIPSeg/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clipseg/model.py#L324
+        causal_attention_mask = self.clip.text_model._build_causal_attention_mask(bsz, seq_len, hidden_states.dtype).to(
+            hidden_states.device
+        )
+        # expand attention_mask
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            attention_mask = _expand_mask(attention_mask, hidden_states.dtype)
+
+        encoder_outputs = self.clip.text_model.encoder(
+            inputs_embeds=hidden_states,
+            attention_mask=attention_mask,
+            causal_attention_mask=causal_attention_mask,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+        )
+
+        last_hidden_state = encoder_outputs[0]
+        last_hidden_state = self.clip.text_model.final_layer_norm(last_hidden_state)
+
+        # text_embeds.shape = [batch_size, sequence_length, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14
+        pooled_output = last_hidden_state[:, -1, :]
+        audio_driven_embeddings = self.clip.text_projection(pooled_output)
+        return audio_driven_embeddings
+
+    def get_pixels(self, image: torch.Tensor) -> torch.Tensor:
+        """
+        Extract spatial features (pixel-level) from the CLIP image encoder.
+
+        Args:
+            image (torch.Tensor): Input image tensor.
+
+        Returns:
+            torch.Tensor: Spatial visual features (pixel-level).
+        """
+        vision_outputs = self.clip.vision_model(pixel_values=image,
+                                                output_attentions=None,
+                                                output_hidden_states=True,
+                                                return_dict=True)
+        last_layer = self.clip.vision_model.encoder.layers[-1]
+
+        hidden_states = vision_outputs.hidden_states[-2]
+        residual = hidden_states
+
+        hidden_states = last_layer.layer_norm1(hidden_states)
+
+        bsz, tgt_len, embed_dim = hidden_states.size()
+
+        # get query proj
+        # query_states = last_layer.self_attn.q_proj(hidden_states) * last_layer.self_attn.scale
+        # key_states = last_layer.self_attn.k_proj(hidden_states)
+        value_states = last_layer.self_attn.v_proj(hidden_states)
+
+        value_states = last_layer.self_attn.out_proj(value_states)
+
+        value_states += residual
+
+        residual = value_states
+        value_states = last_layer.layer_norm2(value_states)
+        value_states = last_layer.mlp(value_states)
+        value_states += residual
+
+        value_states = self.clip.vision_model.post_layernorm(value_states)
+        output = self.clip.visual_projection(value_states)
+
+        width = int(np.sqrt(tgt_len - 1))
+        output = output[:, 1:]
+        if output.ndim == 2:
+            output = output.unsqueeze(0)
+
+        output = output.permute(0, 2, 1)
+        output = output.reshape(bsz, self.clip.visual_projection.out_features, width, width)
+
+        return output

modules/FGA/atten.py

@@ -0,0 +1,303 @@
+#!/usr/bin/env python
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+from itertools import product, permutations, combinations_with_replacement, chain
+
+
+class Unary(nn.Module):
+    def __init__(self, embed_size):
+        """
+            Captures local entity information
+        :param embed_size:  the embedding dimension
+        """
+        super(Unary, self).__init__()
+        self.embed = nn.Conv1d(embed_size, embed_size, 1)
+        self.feature_reduce = nn.Conv1d(embed_size, 1, 1)
+
+    def forward(self, X):
+        X = X.transpose(1, 2)
+
+        X_embed = self.embed(X)
+
+        X_nl_embed = F.dropout(F.relu(X_embed), training=self.training)
+        X_poten = self.feature_reduce(X_nl_embed)
+        return X_poten.squeeze(1)
+
+
+class Pairwise(nn.Module):
+    def __init__(self, embed_x_size, x_spatial_dim=None, embed_y_size=None, y_spatial_dim=None):
+        """
+            Captures interaction between utilities or entities of the same utility
+        :param embed_x_size: the embedding dimension of the first utility
+        :param x_spatial_dim: the spatial dimension of the first utility for batch norm and weighted marginalization
+        :param embed_y_size: the embedding dimension of the second utility (none for self-interactions)
+        :param y_spatial_dim: the spatial dimension of the second utility for batch norm and weighted marginalization
+        """
+
+        super(Pairwise, self).__init__()
+        embed_y_size = embed_y_size if y_spatial_dim is not None else embed_x_size
+        self.y_spatial_dim = y_spatial_dim if y_spatial_dim is not None else x_spatial_dim
+
+        self.embed_size = max(embed_x_size, embed_y_size)
+        self.x_spatial_dim = x_spatial_dim
+
+        self.embed_X = nn.Conv1d(embed_x_size, self.embed_size, 1)
+        self.embed_Y = nn.Conv1d(embed_y_size, self.embed_size, 1)
+        if x_spatial_dim is not None:
+            self.normalize_S = nn.BatchNorm1d(self.x_spatial_dim * self.y_spatial_dim)
+
+            self.margin_X = nn.Conv1d(self.y_spatial_dim, 1, 1)
+            self.margin_Y = nn.Conv1d(self.x_spatial_dim, 1, 1)
+
+    def forward(self, X, Y=None):
+
+        X_t = X.transpose(1, 2)
+        Y_t = Y.transpose(1, 2) if Y is not None else X_t
+
+
+        X_embed = self.embed_X(X_t)
+        Y_embed = self.embed_Y(Y_t)
+
+        X_norm = F.normalize(X_embed)
+        Y_norm = F.normalize(Y_embed)
+
+        S = X_norm.transpose(1, 2).bmm(Y_norm)
+        if self.x_spatial_dim is not None:
+            S = self.normalize_S(S.view(-1, self.x_spatial_dim * self.y_spatial_dim)) \
+                .view(-1, self.x_spatial_dim, self.y_spatial_dim)
+
+            X_poten = self.margin_X(S.transpose(1, 2)).transpose(1, 2).squeeze(2)
+            Y_poten = self.margin_Y(S).transpose(1, 2).squeeze(2)
+        else:
+            X_poten = S.mean(dim=2, keepdim=False)
+            Y_poten = S.mean(dim=1, keepdim=False)
+
+        if Y is None:
+            return X_poten
+        else:
+            return X_poten, Y_poten
+
+
+class Atten(nn.Module):
+    def __init__(self, util_e, sharing_factor_weights=[], prior_flag=False,
+                 sizes=[], size_force=False, pairwise_flag=True,
+                 unary_flag=True, self_flag=True):
+        """
+            The class performs an attention on a given list of utilities representation.
+        :param util_e: the embedding dimensions
+        :param sharing_factor_weights:  To share weights, provide a dict of tuples:
+         {idx: (num_utils, connected utils)
+         Note, for efficiency, the shared utils (i.e., history, are connected to ans
+          and question only.
+         TODO: connections between shared utils
+        :param prior_flag: is prior factor provided
+        :param sizes: the spatial simension (used for batch-norm and weighted marginalization)
+        :param size_force: force spatial size with adaptive avg pooling.
+        :param pairwise_flag: use pairwise interaction between utilities
+        :param unary_flag: use local information
+        :param self_flag: use self interactions between utilitie's entities
+        """
+        super(Atten, self).__init__()
+        self.util_e = util_e
+
+        self.prior_flag = prior_flag
+
+        self.n_utils = len(util_e)
+
+        self.spatial_pool = nn.ModuleDict()
+
+        self.un_models = nn.ModuleList()
+
+        self.self_flag = self_flag
+        self.pairwise_flag = pairwise_flag
+        self.unary_flag = unary_flag
+        self.size_force = size_force
+
+        if len(sizes) == 0:
+            sizes = [None for _ in util_e]
+
+        self.sharing_factor_weights = sharing_factor_weights
+
+        #force the provided size
+        for idx, e_dim in enumerate(util_e):
+            self.un_models.append(Unary(e_dim))
+            if self.size_force:
+                self.spatial_pool[str(idx)] = nn.AdaptiveAvgPool1d(sizes[idx])
+
+        #Pairwise
+        self.pp_models = nn.ModuleDict()
+        for ((idx1, e_dim_1), (idx2, e_dim_2)) \
+                in combinations_with_replacement(enumerate(util_e), 2):
+            # self
+            if self.self_flag and idx1 == idx2:
+                self.pp_models[str(idx1)] = Pairwise(e_dim_1, sizes[idx1])
+            else:
+                if pairwise_flag:
+                    if idx1 in self.sharing_factor_weights:
+                        # not connected
+                        if idx2 not in self.sharing_factor_weights[idx1][1]:
+                           continue
+                    if idx2 in self.sharing_factor_weights:
+                        # not connected
+                        if idx1 not in self.sharing_factor_weights[idx2][1]:
+                            continue
+                    self.pp_models[str((idx1, idx2))] = Pairwise(e_dim_1, sizes[idx1], e_dim_2, sizes[idx2])
+
+        # Handle reduce potentials (with scalars)
+        self.reduce_potentials = nn.ModuleList()
+
+        self.num_of_potentials = dict()
+
+        self.default_num_of_potentials = 0
+
+        if self.self_flag:
+            self.default_num_of_potentials += 1
+        if self.unary_flag:
+            self.default_num_of_potentials += 1
+        if self.prior_flag:
+            self.default_num_of_potentials += 1
+        for idx in range(self.n_utils):
+            self.num_of_potentials[idx] = self.default_num_of_potentials
+
+        '''
+         All other utilities
+        '''
+        if pairwise_flag:
+            for idx, (num_utils, connected_utils) in sharing_factor_weights:
+                for c_u in connected_utils:
+                    self.num_of_potentials[c_u] += num_utils
+                    self.num_of_potentials[idx] += 1
+            for k in self.num_of_potentials:
+                if k not in self.sharing_factor_weights:
+                    self.num_of_potentials[k] += (self.n_utils - 1) \
+                                                 - len(sharing_factor_weights)
+
+        for idx in range(self.n_utils):
+            self.reduce_potentials.append(nn.Conv1d(self.num_of_potentials[idx],
+                                                    1, 1, bias=False))
+
+    def forward(self, utils, priors=None):
+        assert self.n_utils == len(utils)
+        assert (priors is None and not self.prior_flag) \
+               or (priors is not None
+                   and self.prior_flag
+                   and len(priors) == self.n_utils)
+        b_size = utils[0].size(0)
+        util_factors = dict()
+        attention = list()
+
+        #Force size, constant size is used for pairwise batch normalization
+        if self.size_force:
+            for i, (num_utils, _) in self.sharing_factor_weights.items():
+                if str(i) not in self.spatial_pool.keys():
+                    continue
+                else:
+                    high_util = utils[i]
+                    high_util = high_util.view(num_utils * b_size, high_util.size(2), high_util.size(3))
+                    high_util = high_util.transpose(1, 2)
+                    utils[i] = self.spatial_pool[str(i)](high_util).transpose(1, 2)
+
+            for i in range(self.n_utils):
+                if i in self.sharing_factor_weights \
+                        or str(i) not in self.spatial_pool.keys():
+                    continue
+                utils[i] = utils[i].transpose(1, 2)
+                utils[i] = self.spatial_pool[str(i)](utils[i]).transpose(1, 2)
+                if self.prior_flag and priors[i] is not None:
+                    priors[i] = self.spatial_pool[str(i)](priors[i].unsqueeze(1)).squeeze(1)
+
+        # handle Shared weights
+        for i, (num_utils, connected_list) in self.sharing_factor_weights:
+            if self.unary_flag:
+                util_factors.setdefault(i, []).append(self.un_models[i](utils[i]))
+
+            if self.self_flag:
+                util_factors.setdefault(i, []).append(self.pp_models[str(i)](utils[i]))
+
+            if self.pairwise_flag:
+                for j in connected_list:
+                    other_util = utils[j]
+                    expanded_util = other_util.unsqueeze(1).expand(b_size,
+                                                                   num_utils,
+                                                                   other_util.size(1),
+                                                                   other_util.size(2)).contiguous().view(
+                        b_size * num_utils,
+                        other_util.size(1),
+                        other_util.size(2))
+
+                    if i < j:
+                        factor_ij, factor_ji = self.pp_models[str((i, j))](utils[i], expanded_util)
+                    else:
+                        factor_ji, factor_ij = self.pp_models[str((j, i))](expanded_util, utils[i])
+                    util_factors[i].append(factor_ij)
+                    util_factors.setdefault(j, []).append(factor_ji.view(b_size, num_utils, factor_ji.size(1)))
+
+        # handle local factors
+        for i in range(self.n_utils):
+            if i in self.sharing_factor_weights:
+                continue
+            if self.unary_flag:
+                util_factors.setdefault(i, []).append(self.un_models[i](utils[i]))
+            if self.self_flag:
+                util_factors.setdefault(i, []).append(self.pp_models[str(i)](utils[i]))
+
+        # joint
+        if self.pairwise_flag:
+            for (i, j) in combinations_with_replacement(range(self.n_utils), 2):
+                if i in self.sharing_factor_weights \
+                        or j in self.sharing_factor_weights:
+                    continue
+                if i == j:
+                    continue
+                else:
+                    factor_ij, factor_ji = self.pp_models[str((i, j))](utils[i], utils[j])
+                    util_factors.setdefault(i, []).append(factor_ij)
+                    util_factors.setdefault(j, []).append(factor_ji)
+
+        # perform attention
+        for i in range(self.n_utils):
+            if self.prior_flag:
+                prior = priors[i] \
+                    if priors[i] is not None \
+                    else torch.zeros_like(util_factors[i][0], requires_grad=False).cuda()
+
+                util_factors[i].append(prior)
+
+            util_factors[i] = torch.cat([p if len(p.size()) == 3 else p.unsqueeze(1)
+                                       for p in util_factors[i]], dim=1)
+            util_factors[i] = self.reduce_potentials[i](util_factors[i]).squeeze(1)
+            util_factors[i] = F.softmax(util_factors[i], dim=1).unsqueeze(2)
+            attention.append(torch.bmm(utils[i].transpose(1, 2), util_factors[i]).squeeze(2))
+
+        return attention
+
+
+class NaiveAttention(nn.Module):
+    def __init__(self):
+        """
+            Used for ablation analysis - removing attention.
+        """
+        super(NaiveAttention, self).__init__()
+
+    def forward(self, utils, priors):
+        atten = []
+        spatial_atten = []
+        for u, p in zip(utils, priors):
+            if type(u) is tuple:
+                u = u[1]
+                num_elements = u.shape[0]
+                if p is not None:
+                    u = u.view(-1, u.shape[-2], u.shape[-1])
+                    p = p.view(-1, p.shape[-2], p.shape[-1])
+                    spatial_atten.append(
+                        torch.bmm(p.transpose(1, 2), u).squeeze(2).view(num_elements, -1, u.shape[-2], u.shape[-1]))
+                else:
+                    spatial_atten.append(u.mean(2))
+                continue
+            if p is not None:
+                atten.append(torch.bmm(u.transpose(1, 2), p.unsqueeze(2)).squeeze(2))
+            else:
+                atten.append(u.mean(1))
+        return atten, spatial_atten

modules/FGA/fga_model.py

@@ -0,0 +1,219 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from modules.FGA.atten import Atten
+
+
+class FGA(nn.Module):
+    def __init__(self, vocab_size, word_embed_dim, hidden_ques_dim, hidden_ans_dim,
+                 hidden_hist_dim, hidden_cap_dim, hidden_img_dim):
+        '''
+        Factor Graph Attention
+        :param vocab_size: vocabulary size
+        :param word_embed_dim
+        :param hidden_ques_dim:
+        :param hidden_ans_dim:
+        :param hidden_hist_dim:
+        :param img_features_dim:
+        '''
+        super(FGA, self).__init__()
+
+        print("Init FGA with vocab size %s, word embed %s, hidden ques %s, hidden ans %s,"
+              " hidden hist %s, hidden cap %s, hidden img %s" % (vocab_size, word_embed_dim,
+                                                                hidden_ques_dim,
+                                                                hidden_ans_dim,
+                                                                hidden_hist_dim,
+                                                                hidden_cap_dim,
+                                                                hidden_img_dim))
+        self.hidden_ques_dim = hidden_ques_dim
+        self.hidden_ans_dim = hidden_ans_dim
+        self.hidden_cap_dim = hidden_cap_dim
+        self.hidden_img_dim = hidden_img_dim
+        self.hidden_hist_dim = hidden_hist_dim
+
+        # Vocab of History LSTMs is one more as we are keeping a stop id (the last id)
+        self.word_embedddings = nn.Embedding(vocab_size+1+1, word_embed_dim, padding_idx=0)
+
+        self.lstm_ques = nn.LSTM(word_embed_dim, self.hidden_ques_dim, batch_first=True)
+        self.lstm_ans = nn.LSTM(word_embed_dim, self.hidden_ans_dim, batch_first=True)
+
+        self.lstm_hist_ques = nn.LSTM(word_embed_dim, self.hidden_hist_dim, batch_first=True)
+        self.lstm_hist_ans = nn.LSTM(word_embed_dim, self.hidden_hist_dim, batch_first=True)
+
+        self.lstm_hist_cap = nn.LSTM(word_embed_dim, self.hidden_cap_dim, batch_first=True)
+
+
+        self.qahistnet = nn.Sequential(
+            nn.Linear(self.hidden_hist_dim*2, self.hidden_hist_dim),
+            nn.ReLU(inplace=True)
+        )
+
+        self.concat_dim = self.hidden_ques_dim + self.hidden_ans_dim + \
+                          self.hidden_ans_dim + self.hidden_img_dim + \
+                          self.hidden_cap_dim + self.hidden_hist_dim*9
+
+        self.simnet = nn.Sequential(
+            nn.Linear(self.concat_dim, (self.concat_dim)//2, bias=False),
+            nn.BatchNorm1d((self.concat_dim) // 2),
+            nn.ReLU(inplace=True),
+            nn.Linear((self.concat_dim)//2, (self.concat_dim)//4, bias=False),
+            nn.BatchNorm1d((self.concat_dim) // 4),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.5),
+            nn.Linear((self.concat_dim)//4, 1)
+        )
+
+        # To share weights, provide list of tuples: (idx, list of connected utils)
+        # Note, for efficiency, the shared utils (i.e., history, are connected to ans and question only.
+        # connecting shared factors is not supported (!)
+        sharing_factor_weights = {4: (9, [0, 1]),
+                                  5: (9, [0, 1])}
+
+        self.mul_atten = Atten(util_e=[self.hidden_ans_dim, # Answer modal
+                                       self.hidden_ques_dim, # Question modal
+                                       self.hidden_cap_dim, # Caption modal
+                                       self.hidden_img_dim, # Image modal
+                                       self.hidden_hist_dim, # Question-history modal
+                                       self.hidden_hist_dim # Answer-history modal
+                                       ],
+                               sharing_factor_weights=sharing_factor_weights,
+                               sizes=[100, # 100 Answers
+                                      21, # Question length
+                                      41, # Caption length
+                                      37, # 36 Image regions
+                                      21, # History-Question length
+                                      21 #  History-Answer length
+                                      ] # The spatial dim used for pairwise normalization (use force for adaptive)
+                               , prior_flag=True,
+                               pairwise_flag=True)
+
+
+
+    def forward(self, input_ques, input_ans, input_hist_ques, input_hist_ans, input_hist_cap,
+                input_ques_length, input_ans_length, input_cap_length, i_e):
+        """
+
+        :param input_ques:
+        :param input_ans:
+        :param input_hist_ques:
+        :param input_hist_ans:
+        :param input_hist_cap:
+        :param input_ques_length:
+        :param input_ans_length:
+        :param input_cap_length:
+        :param i_e:
+        :return:
+        """
+
+
+        n_options = input_ans.size()[1]
+        batch_size = input_ques.size()[0]
+
+
+
+        nqa_per_dial, nwords_per_qa = input_hist_ques.size()[1], input_hist_ques.size()[2]
+        nwords_per_cap = input_hist_cap.size()[1]
+        max_length_input_ans = input_ans.size()[-1]
+
+        assert batch_size == input_hist_ques.size()[0] == input_hist_ans.size()[0] == input_ques.size()[0] == \
+               input_ans.size()[0] == input_hist_cap.size()[0]
+        assert nqa_per_dial == input_hist_ques.size()[1] == input_hist_ans.size()[1]
+        assert nwords_per_qa == input_hist_ques.size()[2] == input_hist_ans.size()[2]
+
+        q_we = self.word_embedddings(input_ques)
+        a_we = self.word_embedddings(input_ans.view(-1, max_length_input_ans))
+        hq_we = self.word_embedddings(input_hist_ques.view(-1, nwords_per_qa))
+        ha_we = self.word_embedddings(input_hist_ans.view(-1, nwords_per_qa))
+        c_we = self.word_embedddings(input_hist_cap.view(-1, nwords_per_cap))
+
+
+
+        '''
+            q_we = batch x 20 x embed_ques_dim
+            a_we = 100*batch x 20 x embed_ans_dim
+            hq_we = batch*nqa_per_dial, nwords_per_qa, embed_hist_dim
+            ha_we = batch*nqa_per_dial, nwords_per_qa, embed_hist_dim
+            c_we = batch*ncap_per_dial, nwords_per_cap, embed_hist_dim
+        '''
+        self.lstm_ques.flatten_parameters()
+        self.lstm_ans.flatten_parameters()
+        self.lstm_hist_ques.flatten_parameters()
+        self.lstm_hist_ans.flatten_parameters()
+        self.lstm_hist_cap.flatten_parameters()
+
+
+        i_feat = i_e
+
+        q_seq, self.hidden_ques = self.lstm_ques(q_we)
+        a_seq, self.hidden_ans = self.lstm_ans(a_we)
+        hq_seq, self.hidden_hist_ques = self.lstm_hist_ques(hq_we)
+        ha_seq, self.hidden_hist_ans = self.lstm_hist_ans(ha_we)
+        cap_seq, self.hidden_cap = self.lstm_hist_cap(c_we)
+
+
+        '''
+            length is used for attention prior
+        '''
+        q_len = input_ques_length.data - 1
+        c_len = input_cap_length.data.view(-1) - 1
+
+
+        ans_index = torch.arange(0, n_options * batch_size).long().cuda()
+        ans_len = input_ans_length.data.view(-1) - 1
+        ans_seq = a_seq[ans_index, ans_len, :]
+        ans_seq = ans_seq.view(batch_size, n_options, self.hidden_ans_dim)
+
+        batch_index = torch.arange(0, batch_size).long().cuda()
+        q_prior = torch.zeros(batch_size, q_seq.size(1)).cuda()
+        q_prior[batch_index, q_len] = 100
+        c_prior = torch.zeros(batch_size, cap_seq.size(1)).cuda()
+        c_prior[batch_index, c_len] = 100
+        ans_prior = torch.ones(batch_size, ans_seq.size(1)).cuda()
+        img_prior = torch.ones(batch_size, i_feat.size(1)).cuda()
+
+        (ans_atten, ques_atten, cap_atten, img_atten, hq_atten, ha_atten) = \
+            self.mul_atten([ans_seq, q_seq, cap_seq, i_feat, hq_seq, ha_seq],
+                           priors=[ans_prior, q_prior, c_prior, img_prior, None, None])
+
+        '''
+            expand to answers based
+        '''
+        ques_atten = torch.unsqueeze(ques_atten, 1).expand(batch_size,
+                                                           n_options,
+                                                           self.hidden_ques_dim)
+        cap_atten = torch.unsqueeze(cap_atten, 1).expand(batch_size,
+                                                         n_options,
+                                                         self.hidden_cap_dim)
+        img_atten = torch.unsqueeze(img_atten, 1).expand(batch_size, n_options,
+                                                         self.hidden_img_dim)
+        ans_atten = torch.unsqueeze(ans_atten, 1).expand(batch_size, n_options,
+                                                         self.hidden_ans_dim)
+
+
+        '''
+            combine history
+        '''
+
+        input_qahistnet = torch.cat((hq_atten, ha_atten), 1)
+        # input_qahistnet: (nqa_per_dial*batch x 2*hidden_hist_dim)
+        output_qahistnet = self.qahistnet(input_qahistnet)
+        # output_qahistnet: (nqa_per_dial*batch x hidden_hist_dim)
+        output_qahistnet = output_qahistnet.view(batch_size,
+                                                 nqa_per_dial * self.hidden_hist_dim)
+        # output_qahistnet: (batch x nqa_per_dial*hidden_hist_dim)
+        output_qahistnet = torch.unsqueeze(output_qahistnet, 1)\
+            .expand(batch_size,
+                    n_options,
+                    nqa_per_dial * self.hidden_hist_dim)
+
+        input_qa = torch.cat((ans_seq, ques_atten, ans_atten, img_atten,
+                              output_qahistnet, cap_atten), 2)  # Concatenate last dimension
+
+        input_qa = input_qa.view(batch_size * n_options, self.concat_dim)
+
+        out_scores = self.simnet(input_qa)
+
+        out_scores = out_scores.squeeze(dim=1)
+        out_scores = out_scores.view(batch_size, n_options)
+
+        return out_scores

modules/arg_utils.py

@@ -0,0 +1,42 @@
+import argparse
+from typing import List, Optional, Union, Tuple
+
+
+def int_or_int_list_or_none(value: Optional[Union[int, str]]) -> List[Optional[int]]:
+    """
+    Parse an input value into a list of integers or a single integer, or None.
+
+    Args:
+        value (Optional[Union[int, str]]): The input value to parse.
+
+    Returns:
+        List[Optional[int]]: A list containing either a single integer, a list of integers,
+                             or a single None value.
+
+    Raises:
+        argparse.ArgumentTypeError: If the input value cannot be parsed into the specified formats.
+    """
+    if value in ['None', 'null']:
+        return [None]
+    try:
+        # If the value contains commas, parse it as a comma-separated list of integers
+        if ',' in value:
+            return [int(x) for x in value.split(',')]
+        # If it's a single integer, pack it into a list
+        else:
+            return [int(value)]
+    except ValueError:
+        raise argparse.ArgumentTypeError("Invalid format. Use an integer, a comma-separated list of integers, or None.")
+
+
+def int_or_float(value):
+    if '.' in value:
+        try:
+            return float(value)
+        except ValueError:
+            raise argparse.ArgumentTypeError("Quality level must be an integer or a float")
+    else:
+        try:
+            return int(value)
+        except ValueError:
+            raise argparse.ArgumentTypeError("Quality level must be an integer or a float")

modules/mask_utils.py

@@ -0,0 +1,144 @@
+import torch
+import torch.nn as nn
+
+
+def gumbel_sigmoid(logits: torch.Tensor, tau: float = 1, hard: bool = False):
+    """Samples from the Gumbel-Sigmoid distribution and optionally discretizes.
+    References:
+        - https://github.com/yandexdataschool/gumbel_dpg/blob/master/gumbel.py
+        - https://pytorch.org/docs/stable/_modules/torch/nn/functional.html#gumbel_softmax
+    Note:
+        X - Y ~ Logistic(0,1) s.t. X, Y ~ Gumbel(0, 1).
+        That is, we can implement gumbel_sigmoid using Logistic distribution.
+    """
+    logistic = torch.rand_like(logits)
+    logistic = logistic.div_(1. - logistic).log_()  # ~Logistic(0,1)
+
+    gumbels = (logits + logistic) / tau  # ~Logistic(logits, tau)
+    y_soft = gumbels.sigmoid_()
+
+    if hard:
+        # Straight through.
+        y_hard = y_soft.gt(0.5).type(y_soft.dtype)
+        # gt_ break gradient flow
+        #  y_hard = y_soft.gt_(0.5)  # gt_() maintain dtype, different to gt()
+        ret = y_hard - y_soft.detach() + y_soft
+    else:
+        # Reparametrization trick.
+        ret = y_soft
+
+    return ret
+
+
+class Sim2Mask(nn.Module):
+    def __init__(self, init_w: float = 1.0, init_b: float = 0.0, gumbel_tau: float = 1.0, learnable: bool = True):
+        """
+        Sim2Mask module for generating binary masks.
+
+        Args:
+            init_w (float): Initial value for weight.
+            init_b (float): Initial value for bias.
+            gumbel_tau (float): Gumbel-Softmax temperature.
+            learnable (bool): If True, weight and bias are learnable parameters.
+
+        Reference:
+            "Learning to Generate Text-grounded Mask for Open-world Semantic Segmentation from Only Image-Text Pairs" CVPR 2023
+            - https://github.com/kakaobrain/tcl
+            - https://arxiv.org/abs/2212.00785
+        """
+        super().__init__()
+        self.init_w = init_w
+        self.init_b = init_b
+        self.gumbel_tau = gumbel_tau
+        self.learnable = learnable
+
+        assert not ((init_w is None) ^ (init_b is None))
+        if learnable:
+            self.w = nn.Parameter(torch.full([], float(init_w)))
+            self.b = nn.Parameter(torch.full([], float(init_b)))
+        else:
+            self.w = init_w
+            self.b = init_b
+
+    def forward(self, x, deterministic=False):
+        logits = x * self.w + self.b
+
+        soft_mask = torch.sigmoid(logits)
+        if deterministic:
+            hard_mask = soft_mask.gt(0.5).type(logits.dtype)
+        else:
+            hard_mask = gumbel_sigmoid(logits, hard=True, tau=self.gumbel_tau)
+
+        return hard_mask, soft_mask
+
+    def extra_repr(self):
+        return f'init_w={self.init_w}, init_b={self.init_b}, learnable={self.learnable}, gumbel_tau={self.gumbel_tau}'
+
+
+def norm_img_tensor(tensor: torch.Tensor) -> torch.Tensor:
+    """
+    Normalize image tensor to the range [0, 1].
+
+    Args:
+        tensor (torch.Tensor): Input image tensor.
+
+    Returns:
+        torch.Tensor: Normalized image tensor.
+    """
+    vmin = tensor.amin((2, 3), keepdims=True) - 1e-7
+    vmax = tensor.amax((2, 3), keepdims=True) + 1e-7
+    tensor = (tensor - vmin) / (vmax - vmin)
+    return tensor
+
+
+class ImageMasker(Sim2Mask):
+    def forward(self, x: torch.Tensor, infer: bool = False) -> torch.Tensor:
+        """
+        Forward pass for generating image-level binary masks.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+            infer (bool): True for only inference stage.
+
+        Returns:
+            torch.Tensor: Binary mask.
+
+        Reference:
+            "Can CLIP Help Sound Source Localization?" WACV 2024
+            - https://arxiv.org/abs/2311.04066
+        """
+        if self.training or not infer:
+            output = super().forward(x, False)[0]
+        else:
+            output = torch.sigmoid(x + self.b / self.w)
+        return output
+
+
+class FeatureMasker(nn.Module):
+    def __init__(self, thr: float = 0.5, tau: float = 0.07):
+        """
+        Masker module for generating feature-level masks.
+
+        Args:
+            thr (float): Threshold for generating the mask.
+            tau (float): Temperature for the sigmoid function.
+
+        Reference:
+            "Can CLIP Help Sound Source Localization?" WACV 2024
+            - https://arxiv.org/abs/2311.04066
+        """
+        super().__init__()
+        self.thr = thr
+        self.tau = tau
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass for generating feature-level masks
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            torch.Tensor: Generated mask.
+        """
+        return torch.sigmoid((norm_img_tensor(x) - self.thr) / self.tau)

modules/models.py

@@ -0,0 +1,294 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+import yaml
+import argparse
+
+from modules.BEATs.BEATs import BEATs, BEATsConfig
+from modules.AudioToken.embedder import FGAEmbedder
+from modules.CLIPSeg.clipseg_for_audio import CLIPSeg
+from modules.mask_utils import ImageMasker, FeatureMasker
+from transformers import AutoTokenizer
+
+
+class ACL(nn.Module):
+    def __init__(self, conf_file: str, device: str):
+        """
+        Audio-Grounded Contrastive Learning (ACL) model.
+
+        Args:
+            conf_file (str): Path to the configuration file.
+            device (str): Device to move the model to.
+        """
+        super(ACL, self).__init__()
+
+        # Get configuration
+        with open(conf_file) as f:
+            config = yaml.load(f, Loader=yaml.FullLoader)
+            self.args = argparse.Namespace()
+            self.args.model = argparse.Namespace(**config['model'])
+            self.args.clip_embedding_dim = config['clip_conf'][self.args.model.clip]['embedding_dim']
+            self.args.clip_name = config['clip_conf'][self.args.model.clip]['name']
+            self.pretrain = argparse.Namespace(**config['pretrain'])
+            self.args.audio_proj = argparse.Namespace(**config['fga_conf'][self.args.model.audio_proj])
+
+        # Init audio encoder
+        checkpoint = torch.load(self.pretrain.audio_backbone)
+        cfg = BEATsConfig(checkpoint['cfg'])
+        self.audio_backbone = BEATs(cfg)
+
+        # Text Tokenizer for placeholder prompt
+        self.tokenizer = AutoTokenizer.from_pretrained("CIDAS/clipseg-rd64-refined")
+
+        # Init audio projection layer
+        self.audio_proj = FGAEmbedder(input_size=self.args.audio_proj.input_size * 3,
+                                      output_size=self.args.audio_proj.output_size)
+
+        # Init audio-visual grounder (Grounder: CLIPSeg)
+        self.av_grounder = CLIPSeg.from_pretrained("CIDAS/clipseg-rd64-refined")
+
+        # Init maskers
+        self.masker_i = ImageMasker(10.0, 14.0, 1.0)
+        self.masker_f = FeatureMasker(0.5, 0.07)
+
+        # Load weights
+        self.audio_backbone.load_state_dict(checkpoint['model'])
+        self.audio_backbone.predictor = None
+
+        if self.pretrain.audio_proj is not None:
+            self.audio_proj.load_state_dict(torch.load(self.pretrain.audio_embedder))
+
+        # Set device
+        self.device = device
+        self.audio_backbone.to(device=self.device)
+        self.av_grounder.to(device=self.device)
+        self.audio_proj.to(device=self.device)
+        self.masker_i.to(self.device)
+        self.masker_f.to(self.device)
+
+    def get_placeholder_token(self, prompt_text: str):
+        """
+        Get placeholder token from prompt text
+
+        Args:
+            prompt_text (str): prompt text without '{}'
+
+        Returns:
+            CLIPTokenizerFast result with prompt text
+        """
+        placeholder_token = self.tokenizer(prompt_text, return_tensors="pt").data['input_ids']
+        placeholder_token = F.pad(placeholder_token, (0, 77 - placeholder_token.shape[-1])).to(self.device)
+        return placeholder_token
+
+    def train(self, bool: bool = True):
+        """
+        Set the module in training mode.
+
+        Args:
+            bool (bool): If True, set the module in training mode.
+        """
+        super().train(bool)
+        self.av_grounder.requires_grad_(False)
+        self.audio_backbone.requires_grad_(False)
+
+    def encode_audio(self, audio: torch.Tensor, placeholder_token: torch.Tensor, pos: int,
+                     prompt_size: int) -> torch.Tensor:
+        """
+        Encode audio input into audio-driven embedding (Audio-Driven Embedder)
+
+        Args:
+            audio (torch.Tensor): Input audio tensor.
+            placeholder_token (torch.Tensor): Placeholder token for CLIP Text encoder.
+            pos (int): Position of audio token.
+            prompt_size (int): Size of the placeholder prompt.
+
+        Returns:
+            torch.Tensor: Audio-driven embeddings.
+        """
+        audio_feat = self.audio_backbone.extract_features(audio)[1]
+        audio_token_emb = self.audio_proj(audio_feat).unsqueeze(1)
+        audio_driven_embedding = self.av_grounder.encode_audio(placeholder_token, audio_token_emb, pos,
+                                                               prompt_size + audio_token_emb.shape[1])
+
+        return audio_driven_embedding
+
+    def encode_vision(self, image: torch.Tensor) -> torch.Tensor:
+        """
+        Encode visual input and generate visual embeddings.
+
+        Args:
+            image (torch.Tensor): Input image tensor.
+
+        Returns:
+            torch.Tensor: Visual embeddings.
+        """
+        vision_outputs = self.av_grounder.clip.vision_model(pixel_values=image,
+                                                            output_attentions=None,
+                                                            output_hidden_states=True,
+                                                            return_dict=True)
+        pooled_output = self.av_grounder.clip.visual_projection(vision_outputs[1])
+
+        return pooled_output
+
+    def forward_decoder(self, image: torch.Tensor, embedding: torch.Tensor, resolution: int = 224) -> torch.Tensor:
+        """
+        Forward pass of audio-visual grounder
+
+        Args:
+            image (torch.Tensor): Input image tensor.
+            embedding (torch.Tensor): Condition embedding tensor for grounder.
+            resolution (int): Resolution of the output.
+            ignore_indices (list): List of indices to ignore.
+
+        Returns:
+            torch.Tensor: Logits from the decoder.
+        """
+        # step 1: forward the query images through the frozen CLIP vision encoder
+        vision_outputs = self.av_grounder.clip.vision_model(pixel_values=image,
+                                                            output_attentions=None,
+                                                            output_hidden_states=True,
+                                                            return_dict=True)
+
+        hidden_states = vision_outputs.hidden_states
+        # we add +1 here as the hidden states also include the initial embeddings
+        activations = [hidden_states[i + 1] for i in self.av_grounder.extract_layers]
+
+        # step 2: compute conditional embeddings, either from text, images or an own provided embedding
+        # Audio injected embedding from input argument
+
+        # step 3: forward both the pooled output and the activations through the lightweight decoder to predict masks
+        decoder_outputs = self.av_grounder.decoder(
+            activations,
+            embedding,
+            output_attentions=None,
+            output_hidden_states=None,
+            return_dict=True,
+        )
+        logits = decoder_outputs.logits
+
+        if logits.ndim == 2:
+            logits = logits.unsqueeze(0).unsqueeze(1)
+        else:
+            logits = logits.unsqueeze(1)
+
+        B, c, h, w = image.shape
+        if (h, w) != (resolution, resolution):
+            logits = F.interpolate(logits, resolution, mode='bicubic')
+
+        return logits
+
+    def forward_module(self, image: torch.Tensor, embedding: torch.Tensor, resolution: int = 224,
+                       force_comb: bool = False) -> torch.Tensor:
+        """
+        Forward pass through the module.
+
+        Args:
+            image (torch.Tensor): Input image tensor.
+            embedding (torch.Tensor): Condition embedding tensor for grounder.
+            resolution (int): Resolution of the output tensor.
+            force_comb (bool): If True, force to get logits with all combination audio and image.
+
+        Returns:
+            torch.Tensor: Logits from the decoder.
+        """
+        # N image, 1 embedding case -> [B_i, h, w]
+        if embedding.shape[0] != image.shape[0] and embedding.shape[0] == 1:
+            embeddings = embedding.repeat(image.shape[0], 1)
+            logits = self.forward_decoder(image, embeddings, resolution)
+
+        # N image, M embedding case -> [B_i, B_e, h, w]
+        elif embedding.shape[0] != image.shape[0] and embedding.shape[0] != 1 and image.shape[0] != 1 or force_comb:
+            logit_list = []
+            for i in range(embedding.shape[0]):
+                embeddings = embedding[i].unsqueeze(0).repeat(image.shape[0], 1)
+                logit_list.append(self.forward_decoder(image, embeddings, resolution))
+            logits = torch.cat(logit_list, dim=1)
+
+        # N image, N embedding or 1 image, N embedding -> [B_e, h, w]
+        else:
+            logits = self.forward_decoder(image, embedding, resolution)
+
+        return logits
+
+    def encode_masked_vision(self, image: torch.Tensor, embedding: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, float, float]:
+        """
+        Encode masked visual feature both image-level and feature-level.
+
+        Args:
+            image (torch.Tensor): Input image tensor.
+            embedding (torch.Tensor): Condition embedding tensor for grounder.
+
+        Returns:
+            tuple[torch.Tensor, torch.Tensor, float, float]: Feature masked embeddings, masked image embeddings, positive area, negative area.
+        """
+        B, c, h, w = image.shape
+        maskclip_feat = self.av_grounder.get_pixels(image)  # v^D: [B, c, h, w]
+        clipseg_mask = self.forward_module(image, embedding, h, force_comb=True)  # M^G: [B, B, H, W]
+
+        # Area
+        area_matrix = self.masker_i(clipseg_mask).mean((2, 3))
+        positive_area = area_matrix.diagonal().mean()
+        negative_area = area_matrix.mean() - positive_area / B
+
+        # Feature level masker
+        feature_mask = F.interpolate(self.masker_f(clipseg_mask), maskclip_feat.shape[2])
+
+        # Image level masker
+        ind = torch.arange(B).to(image.device)
+        image_mask = self.masker_i(clipseg_mask[ind, ind].unsqueeze(1))  # Positive pair only
+        feature_masked_emb = torch.einsum('bchw,bnhw->bnc', maskclip_feat, feature_mask) / (feature_mask.sum() + 1e-6)
+
+        # step 1: forward the query images through the frozen CLIP vision encoder
+        masked_vision_outputs = self.av_grounder.clip.vision_model(pixel_values=image * image_mask,
+                                                                   output_attentions=None,
+                                                                   output_hidden_states=True,
+                                                                   return_dict=True)
+        masked_image_emb = self.av_grounder.clip.visual_projection(masked_vision_outputs[1])
+
+        return feature_masked_emb, masked_image_emb, positive_area, negative_area
+
+    def forward(self, image: torch.Tensor, embedding: torch.Tensor, resolution: int = 224) -> dict:
+        """
+        Forward pass of ACL model.
+
+        Args:
+            image (torch.Tensor): Input image tensor.
+            embedding (torch.Tensor): Condition embedding tensor for grounder.
+            resolution (int): Resolution of the output tensor.
+
+        Returns:
+            dict: Output dictionary containing relevant tensors.
+        """
+        if self.training:
+            # seg_logit = self.forward_module(image, embedding, resolution)
+            v_f, v_i, p_area, n_area = self.encode_masked_vision(image, embedding)
+            out_dict = {'v_f': v_f, 'v_i': v_i, 'p_area': p_area, 'n_area': n_area}
+
+        else:
+            seg_logit = self.forward_module(image, embedding, resolution)
+            heatmap = self.masker_i(seg_logit, infer=True)
+            out_dict = {'heatmap': heatmap}
+
+        return out_dict
+
+    def save(self, model_dir: str):
+        """
+        Save model parameters to a file. (Only trainable parts)
+
+        Args:
+            model_dir (str): Directory to save the model.
+        """
+        ckp = {'audio_proj': self.audio_proj.state_dict(), 'masker_i': self.masker_i.state_dict()}
+        torch.save(ckp, model_dir)
+
+    def load(self, model_dir: str):
+        """
+        Load model parameters from a file. (Only trainable parts)
+
+        Args:
+            model_dir (str): Directory to load the model from.
+        """
+        ckp = torch.load(model_dir)
+        self.audio_proj.load_state_dict(ckp['audio_proj'])
+        self.masker_i.load_state_dict(ckp['masker_i'])