Upload folder using huggingface_hub

README.md

@@ -0,0 +1,38 @@
+# Synchformer Hugging Face Model
+
+This repository contains a Synchformer model for audio-visual synchronization. The model predicts the offset between audio and video tracks.
+
+## Usage
+
+```python
+from transformers import AutoModel
+import torch
+
+# Load the model
+model = AutoModel.from_pretrained("AmrMKayid/synchformer-hf")
+model.to("cuda" if torch.cuda.is_available() else "cpu")
+
+# Predict offset for a video
+results = model.predict_offset(
+    "path/to/your/video.mp4",
+    offset_sec=0.0,  # Ground truth offset (if known)
+    v_start_i_sec=0.0  # Start time in seconds for video
+)
+
+# Print results
+print("\nPrediction Results:")
+for pred in results["predictions"]:
+    print(f'p={pred["probability"]:.4f}, "{pred["offset_sec"]:.2f}" (class {pred["class_idx"]})')
+```
+
+## Model Details
+
+This model is based on the Synchformer architecture, which uses a transformer to predict the offset between audio and video tracks.
+
+## Requirements
+
+- torch
+- torchaudio
+- torchvision
+- omegaconf
+- ffmpeg (for video processing)

__init__.py

@@ -0,0 +1,8 @@
+from .synchformer_config import SynchformerConfig
+from .synchformer_model import SynchformerModel
+from .modeling_synchformer import *
+
+__all__ = [
+    "SynchformerConfig",
+    "SynchformerModel",
+] 

cfg.yaml

@@ -0,0 +1,282 @@
+action: train_avsync_model
+model:
+  target: model.sync_model.Synchformer
+  params:
+    afeat_extractor:
+      is_trainable: false
+      target: model.modules.feat_extractors.audio.ast.AST
+      params:
+        ckpt_path: /scratch/project_462000293/vladimir/logs/sync/avclip_models/23-12-22T16-13-38/checkpoints/epoch_e28.pt
+        extract_features: true
+        max_spec_t: 66
+        factorize_freq_time: true
+        agg_freq_module: TransformerEncoderLayer
+        agg_time_module: torch.nn.Identity
+        add_global_repr: false
+    vfeat_extractor:
+      is_trainable: false
+      target: model.modules.feat_extractors.visual.motionformer.MotionFormer
+      params:
+        ckpt_path: /scratch/project_462000293/vladimir/logs/sync/avclip_models/23-12-22T16-13-38/checkpoints/epoch_e28.pt
+        extract_features: true
+        factorize_space_time: true
+        agg_space_module: TransformerEncoderLayer
+        agg_time_module: torch.nn.Identity
+        add_global_repr: false
+    aproj:
+      target: torch.nn.Linear
+      params:
+        in_features: 768
+        out_features: 768
+    vproj:
+      target: torch.nn.Linear
+      params:
+        in_features: 768
+        out_features: 768
+    transformer:
+      target: model.sync_model.GlobalTransformer
+      params:
+        n_layer: 3
+        n_head: 8
+        n_embd: 768
+        tok_pdrop: 0.0
+        embd_pdrop: 0.1
+        resid_pdrop: 0.1
+        attn_pdrop: 0.1
+        pos_emb_cfg:
+          target: model.modules.transformer.RandInitPositionalEncoding
+          params:
+            block_shape:
+            - 198
+            n_embd: 768
+        off_head_cfg:
+          target: torch.nn.Linear
+          params:
+            in_features: 768
+            out_features: 21
+training:
+  base_learning_rate: 2.0e-06
+  base_batch_size: 16
+  num_workers: 7
+  num_epochs: 10000
+  patience: 50
+  to_max_metric: true
+  metric_name: accuracy_1
+  early_stop_phase: valid
+  use_half_precision: true
+  seed: 1337
+  compile: false
+  skip_test: false
+  run_test_only: false
+  resume: false
+  finetune: false
+  dist_backend: nccl
+  max_clip_norm: 1
+  lr_scheduler:
+    name: constant_with_warmup
+    warmup: 1000
+  optimizer:
+    name: adam
+    betas:
+    - 0.9
+    - 0.999
+    momentum: 0.9
+    weight_decay: 0
+  local_rank: 0
+  global_rank: 0
+  world_size: 32
+data:
+  offset_type: grid
+  num_off_cls: 21
+  prob_oos: null
+  max_off_sec: 2
+  crop_len_sec: 5
+  step_size_seg: 0.5
+  vids_path: /scratch/project_462000293/vladimir/data/audioset/h264_video_25fps_256side_16000hz_aac/
+  size_before_crop: 256
+  input_size: 224
+  segment_size_vframes: 16
+  vfps: 25
+  afps: 16000
+  n_segments: 14
+  do_offset: true
+  p_color_jitter: 0.0
+  p_gray_scale: 0.0
+  sometimes_upscale_p: 0.0
+  is_spatial_crop_random: true
+  is_temporal_crop_random: true
+  audio_jitter_sec: 0.05
+  p_horizontal_flip: 0.5
+  p_audio_aug: 0.0
+  dataset:
+    target: dataset.audioset.AudioSet
+    params:
+      load_fixed_offsets_on:
+      - valid
+      - test
+      vis_load_backend: read_video
+      size_ratio: null
+transform_sequence_train:
+- target: dataset.transforms.EqualifyFromRight
+  params:
+    clip_max_len_sec: 10
+- target: dataset.transforms.RGBSpatialCropSometimesUpscale
+  params:
+    sometimes_p: 0.0
+    smaller_input_size: 192
+    target_input_size: 224
+    is_random: true
+- target: dataset.transforms.TemporalCropAndOffset
+  params:
+    crop_len_sec: 5
+    max_off_sec: 2
+    max_wiggle_sec: 0.05
+    do_offset: true
+    offset_type: grid
+    prob_oos: null
+    grid_size: 21
+    segment_size_vframes: 16
+    n_segments: 14
+    step_size_seg: 0.5
+    vfps: 25
+- target: dataset.transforms.RandomApplyColorDistortion
+  params:
+    p_color_jitter: 0.0
+    s: 1.0
+    p_gray_scale: 0.0
+- target: dataset.transforms.RandomHorizontalFlip
+  params:
+    p: 0.5
+- target: dataset.transforms.AudioRandomReverb
+  params:
+    p: 0.0
+- target: dataset.transforms.AudioRandomVolume
+  params:
+    p: 0.0
+    gain: 2.0
+    gain_type: amplitude
+- target: dataset.transforms.AudioRandomPitchShift
+  params:
+    p: 0.0
+    shift: 1000
+- target: dataset.transforms.AudioRandomLowpassFilter
+  params:
+    p: 0.0
+    cutoff_freq: 100
+- target: dataset.transforms.AudioRandomGaussNoise
+  params:
+    p: 0.0
+    amplitude: 0.01
+- target: dataset.transforms.GenerateMultipleSegments
+  params:
+    segment_size_vframes: 16
+    n_segments: 14
+    is_start_random: true
+    step_size_seg: 0.5
+- target: dataset.transforms.RGBToHalfToZeroOne
+- target: dataset.transforms.RGBNormalize
+  params:
+    mean:
+    - 0.5
+    - 0.5
+    - 0.5
+    std:
+    - 0.5
+    - 0.5
+    - 0.5
+- target: dataset.transforms.AudioMelSpectrogram
+  params:
+    sample_rate: 16000
+    win_length: 400
+    hop_length: 160
+    n_fft: 1024
+    n_mels: 128
+- target: dataset.transforms.AudioLog
+- target: dataset.transforms.PadOrTruncate
+  params:
+    max_spec_t: 66
+- target: dataset.transforms.AudioNormalizeAST
+  params:
+    mean: -4.2677393
+    std: 4.5689974
+- target: dataset.transforms.PermuteStreams
+  params:
+    einops_order_audio: S F T -> S 1 F T
+    einops_order_rgb: S T C H W -> S T C H W
+transform_sequence_test:
+- target: dataset.transforms.EqualifyFromRight
+- target: dataset.transforms.RGBSpatialCrop
+  params:
+    input_size: 224
+    is_random: false
+- target: dataset.transforms.TemporalCropAndOffset
+  params:
+    crop_len_sec: 5
+    max_off_sec: 2
+    max_wiggle_sec: 0.0
+    do_offset: true
+    grid_size: 21
+    offset_type: grid
+    prob_oos: null
+    segment_size_vframes: 16
+    n_segments: 14
+    step_size_seg: 0.5
+    vfps: 25
+- target: dataset.transforms.GenerateMultipleSegments
+  params:
+    segment_size_vframes: 16
+    n_segments: 14
+    is_start_random: false
+    step_size_seg: 0.5
+- target: dataset.transforms.RGBToHalfToZeroOne
+- target: dataset.transforms.RGBNormalize
+  params:
+    mean:
+    - 0.5
+    - 0.5
+    - 0.5
+    std:
+    - 0.5
+    - 0.5
+    - 0.5
+- target: dataset.transforms.AudioMelSpectrogram
+  params:
+    sample_rate: 16000
+    win_length: 400
+    hop_length: 160
+    n_fft: 1024
+    n_mels: 128
+- target: dataset.transforms.AudioLog
+- target: dataset.transforms.PadOrTruncate
+  params:
+    max_spec_t: 66
+- target: dataset.transforms.AudioNormalizeAST
+  params:
+    mean: -4.2677393
+    std: 4.5689974
+- target: dataset.transforms.PermuteStreams
+  params:
+    einops_order_audio: S F T -> S 1 F T
+    einops_order_rgb: S T C H W -> S T C H W
+logging:
+  logdir: /scratch/project_462000293/vladimir/logs/sync/sync_models/
+  log_code_state: true
+  log_frequency: 20
+  patterns_to_ignore:
+  - logs
+  - .git
+  - __pycache__
+  - data
+  - '*.pt'
+  - sbatch_logs
+  - '*.mp4'
+  - '*.wav'
+  - '*.jpg'
+  - '*.gif'
+  - misc*
+  vis_segment_sim: true
+  log_max_items: 500000
+  use_wandb: true
+start_time: 24-01-04T16-39-21
+config: ./configs/sync.yaml
+ckpt_path: /scratch/project_462000293/vladimir/logs/sync/sync_models/24-01-04T16-39-21/24-01-04T16-39-21.pt

config.json

@@ -0,0 +1,12 @@
+{
+  "afeat_extractor_config": {},
+  "afps": 16000,
+  "in_size": 256,
+  "max_off_sec": 2,
+  "model_type": "synchformer",
+  "transformer_config": {},
+  "transformers_version": "4.45.2",
+  "use_half_precision": true,
+  "vfeat_extractor_config": {},
+  "vfps": 25
+}

model.pt

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

model_files.zip

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

modeling_synchformer.py

@@ -0,0 +1,10 @@
+from transformers.models.auto.modeling_auto import _BaseAutoModelClass
+from transformers.models.auto.configuration_auto import _LazyAutoMapping
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.auto import AutoModel
+
+from synchformer_config import SynchformerConfig
+from synchformer_model import SynchformerModel
+
+# Register the model with the transformers library
+AutoModel.register(SynchformerConfig, SynchformerModel) 

module_loader.py

@@ -0,0 +1,32 @@
+
+import os
+import sys
+import importlib.util
+
+def setup_modules():
+    # Get the directory where this script is located
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    modules_dir = os.path.join(current_dir, "modules")
+    
+    # Add modules directory to path
+    if modules_dir not in sys.path:
+        sys.path.insert(0, modules_dir)
+    
+    # Import required modules
+    required_modules = ["dataset", "model", "scripts", "utils"]
+    for module_name in required_modules:
+        module_path = os.path.join(modules_dir, module_name)
+        if os.path.exists(module_path):
+            # Check if module is already imported
+            if module_name not in sys.modules:
+                # Import the module
+                spec = importlib.util.spec_from_file_location(
+                    module_name, 
+                    os.path.join(module_path, "__init__.py")
+                )
+                if spec:
+                    module = importlib.util.module_from_spec(spec)
+                    sys.modules[module_name] = module
+                    spec.loader.exec_module(module)
+    
+    return True

modules/dataset/audioset.py

@@ -0,0 +1,216 @@
+import csv
+import logging
+import random
+import sys
+from glob import glob
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, '.')  # nopep8
+from dataset.dataset_utils import (get_fixed_offsets, get_video_and_audio)
+
+
+class AudioSet(torch.utils.data.Dataset):
+
+    def __init__(self,
+                 split,
+                 vids_dir,
+                 transforms=None,
+                 to_filter_bad_examples=True,
+                 splits_path='./data',
+                 meta_path='./data/audioset.csv',
+                 seed=1337,
+                 load_fixed_offsets_on=['valid' 'test'],
+                 vis_load_backend='read_video',
+                 size_ratio=None,
+                 attr_annot_path=None,
+                 max_attr_per_vid=None):
+        super().__init__()
+        self.max_clip_len_sec = None
+        self.split = split
+        self.vids_dir = Path(vids_dir)
+        self.transforms = transforms
+        self.to_filter_bad_examples = to_filter_bad_examples
+        self.splits_path = Path(splits_path)
+        self.meta_path = Path(meta_path)
+        self.seed = seed
+        self.load_fixed_offsets_on = [] if load_fixed_offsets_on is None else load_fixed_offsets_on
+        self.vis_load_backend = vis_load_backend
+        self.size_ratio = size_ratio
+
+        self.split2short = {'train': 'unbalanced', 'valid': 'balanced', 'test': 'eval'}
+        short2long = {'unbalanced': 'unbalanced_train_segments',
+                      'balanced': 'balanced_train_segments',
+                      'eval': 'eval_segments'}
+
+        # read meta
+        split_meta = []
+        for shortdir_vid, start, end, targets, phase in csv.reader(open(meta_path), quotechar='"'):
+            if shortdir_vid.startswith(self.split2short[split]):
+                # shortdir_vid 'unbalanced/NFap9qgsI_s' -> 'unbalanced_train_segments/NFap9qgsI_s'
+                shortdir, vid = shortdir_vid.split('/')
+                longdir_vid = '/'.join([short2long[shortdir], vid])
+                split_meta.append([longdir_vid, float(start), float(end), targets, phase])
+
+        # filter "bad" examples
+        if to_filter_bad_examples:
+            split_meta = self.filter_bad_examples(split_meta)
+
+        # label maps
+        self.label2target = {l: int(t) for t, _, l in csv.reader(open(self.splits_path / 'audioset_labels.csv'))}
+        self.target2label = {t: l for l, t in self.label2target.items()}
+        self.video2target = {key: list(map(int, targets.split(','))) for key, _, _, targets, _ in split_meta}
+
+        clip_paths = [self.vids_dir / f'{k}_{int(s*1000)}_{int(e*1000)}.mp4' for k, s, e, t, p in split_meta]
+        clip_paths = sorted(clip_paths)
+
+        # loading the fixed offsets. COMMENT THIS IF YOU DON'T HAVE A FILE YET
+        if transforms is not None and split in load_fixed_offsets_on:
+            logging.info(f'Using fixed offset for {split}')
+            self.vid2offset_params = get_fixed_offsets(transforms, split, splits_path, 'audioset')
+
+        self.dataset = clip_paths
+        if size_ratio is not None and 0.0 < size_ratio < 1.0:
+            cut_off = int(len(self.dataset) * size_ratio)
+            random.seed(seed)
+            random.shuffle(self.dataset)
+            self.dataset = self.dataset[:cut_off]
+
+        logging.info(f'{split} has {len(self.dataset)} items')
+
+    def filter_bad_examples(self, audioset_meta):
+        bad = set()
+        base_path = Path('./data/filtered_examples_audioset')
+        files = sorted(glob(str(base_path / '*.txt')))
+        lists = [open(p).read().splitlines() for p in files]
+        logging.info(f'Filtering for {files}')
+        for s in lists:
+            bad = bad.union(s)
+        # the ugly string converts '---g-f_I2yQ', '1' into `---g-f_I2yQ_1000_11000`
+        audioset_meta = [r for r in audioset_meta if f'{r[0]}_{int(r[1]*1000)}_{int(r[2]*1000)}' not in bad]
+        return audioset_meta
+
+    def __getitem__(self, index):
+        path = self.dataset[index]
+        rgb, audio, meta = self.load_media(path)
+        item = self.make_datapoint(path, rgb, audio, meta)
+        if self.transforms is not None:
+            item = self.transforms(item)
+        return item
+
+    def make_datapoint(self, path, rgb, audio, meta):
+        # (Tv, 3, H, W) in [0, 225], (Ta, C) in [-1, 1]
+        # TODO: since audioset is annotated by tagging, targets have multiple elemenets: default collate fails
+        # targets = self.video2target[f'{Path(path).parent.stem}/{Path(path).stem[:11]}']
+        item = {
+            'video': rgb,
+            'audio': audio,
+            'meta': meta,
+            'path': str(path),
+            # 'targets': {'audioset_target': [targets], 'audioset_label': [self.target2label[t] for t in targets]},
+            'targets': {},
+            'split': self.split,
+        }
+
+        # loading the fixed offsets. COMMENT THIS IF YOU DON'T HAVE A FILE YET
+        if self.transforms is not None and self.split in self.load_fixed_offsets_on:
+            key = f'{self.split2short[self.split]}/{Path(path).stem}'
+            item['targets']['offset_sec'] = self.vid2offset_params[key]['offset_sec']
+            item['targets']['v_start_i_sec'] = self.vid2offset_params[key]['v_start_i_sec']
+
+        return item
+
+    def load_media(self, path):
+        rgb, audio, meta = get_video_and_audio(path, get_meta=True, end_sec=self.max_clip_len_sec)
+        return rgb, audio, meta
+
+    def __len__(self):
+        return len(self.dataset)
+
+class AudioSetBalanced737k(AudioSet):
+
+    def __init__(self, split, vids_dir, transforms=None, to_filter_bad_examples=True, splits_path='./data',
+                 # here
+                 meta_path='./data/audioset_balanced_737k.csv',
+                 seed=1337, load_fixed_offsets_on=['valid', 'test'], vis_load_backend='read_video', size_ratio=None,
+                 attr_annot_path=None, max_attr_per_vid=None):
+        super().__init__(split, vids_dir, transforms, to_filter_bad_examples, splits_path, meta_path,
+                         seed, load_fixed_offsets_on, vis_load_backend, size_ratio)
+
+class AudioSetBalanced540k(AudioSet):
+    ''' MBT's balanced 500k (from unbalanced part) + 20k from balaced part + 20k from eval part '''
+
+    def __init__(self, split, vids_dir, transforms=None, to_filter_bad_examples=True, splits_path='./data',
+                 # here
+                 meta_path='./data/audioset_balanced_540k.csv',
+                 seed=1337, load_fixed_offsets_on=['valid', 'test'], vis_load_backend='read_video', size_ratio=None,
+                 attr_annot_path=None, max_attr_per_vid=None):
+        super().__init__(split, vids_dir, transforms, to_filter_bad_examples, splits_path, meta_path,
+                         seed, load_fixed_offsets_on, vis_load_backend, size_ratio)
+
+
+if __name__ == '__main__':
+    from omegaconf import OmegaConf
+    from scripts.train_utils import get_transforms
+    from utils.utils import cfg_sanity_check_and_patch
+    cfg = OmegaConf.load('./configs/sparse_sync.yaml')
+    vis_load_backend = 'read_video'
+
+    transforms = get_transforms(cfg)
+
+    # vids_path = 'PLACEHOLDER'
+    vids_path = '/scratch/project_2000936/vladimir/data/audioset/h264_video_25fps_256side_16000hz_aac'
+    load_fixed_offsets_on = []
+
+    cfg.data.dataset.params.size_ratio = 0.1
+
+    cfg_sanity_check_and_patch(cfg)
+
+    datasets = {
+        'train': AudioSet('train', vids_path, transforms['train'], vis_load_backend=vis_load_backend,
+                          to_filter_bad_examples=True, size_ratio=cfg.data.dataset.params.size_ratio,
+                          load_fixed_offsets_on=load_fixed_offsets_on),
+        'valid': AudioSet('valid', vids_path, transforms['test'], vis_load_backend=vis_load_backend,
+                          to_filter_bad_examples=True, load_fixed_offsets_on=load_fixed_offsets_on),
+        'test': AudioSet('test', vids_path, transforms['test'], vis_load_backend=vis_load_backend,
+                         to_filter_bad_examples=True, load_fixed_offsets_on=load_fixed_offsets_on),
+    }
+    for phase in ['train', 'valid', 'test']:
+        print(phase, len(datasets[phase]))
+
+    print(datasets['train'][0]['audio'].shape, datasets['train'][0]['video'].shape)
+    print(datasets['train'][0]['meta'])
+    print(datasets['valid'][0]['audio'].shape, datasets['valid'][0]['video'].shape)
+    print(datasets['valid'][0]['meta'])
+    print(datasets['test'][0]['audio'].shape, datasets['test'][0]['video'].shape)
+    print(datasets['test'][0]['meta'])
+
+    for i in range(300, 1000):
+        datasets['train'][i]['path']
+        print(datasets['train'][0]['audio'].shape, datasets['train'][0]['video'].shape)
+        print(datasets['train'][0]['meta'])
+
+    datasets = {
+        'train': AudioSetBalanced737k('train', vids_path, transforms['train'], vis_load_backend=vis_load_backend,
+                          to_filter_bad_examples=True, size_ratio=cfg.data.dataset.params.size_ratio,
+                          load_fixed_offsets_on=load_fixed_offsets_on),
+        'valid': AudioSetBalanced737k('valid', vids_path, transforms['test'], vis_load_backend=vis_load_backend,
+                          to_filter_bad_examples=True, load_fixed_offsets_on=load_fixed_offsets_on),
+        'test': AudioSetBalanced737k('test', vids_path, transforms['test'], vis_load_backend=vis_load_backend,
+                         to_filter_bad_examples=True, load_fixed_offsets_on=load_fixed_offsets_on),
+    }
+    for phase in ['train', 'valid', 'test']:
+        print(phase, len(datasets[phase]))
+
+    print(datasets['train'][0]['audio'].shape, datasets['train'][0]['video'].shape)
+    print(datasets['train'][0]['meta'])
+    print(datasets['valid'][0]['audio'].shape, datasets['valid'][0]['video'].shape)
+    print(datasets['valid'][0]['meta'])
+    print(datasets['test'][0]['audio'].shape, datasets['test'][0]['video'].shape)
+    print(datasets['test'][0]['meta'])
+
+    for i in range(300, 1000):
+        datasets['train'][i]['path']
+        print(datasets['train'][0]['audio'].shape, datasets['train'][0]['video'].shape)
+        print(datasets['train'][0]['meta'])

modules/dataset/dataset_utils.py

@@ -0,0 +1,112 @@
+import csv
+import os
+import random
+from pathlib import Path
+from glob import glob
+import shutil
+import logging
+
+import torchaudio
+import torchvision
+
+from utils.utils import get_fixed_off_fname
+
+
+def get_fixed_offsets(transforms, split, splits_path, dataset_name):
+    '''dataset_name: `vggsound` or `lrs3`'''
+    logging.info(f'Using fixed offset for {split}')
+    vid2offset_params = {}
+    fixed_offset_fname = get_fixed_off_fname(transforms, split)
+    if fixed_offset_fname is None:
+        raise ValueError('Cant find fixed offsets for given params. Perhaps you need to make it first?')
+    fixed_offset_path = os.path.join(splits_path, f'fixed_offsets_{dataset_name}', fixed_offset_fname)
+    fixed_offset_paths = sorted(glob(fixed_offset_path.replace(split, '*')))
+    assert len(fixed_offset_paths) > 0, f'Perhaps: {fixed_offset_path} does not exist. Make fixed offsets'
+
+    for fix_off_path in fixed_offset_paths:
+        reader = csv.reader(open(fix_off_path))
+        # k700_2020 has no header, and also `vstart` comes before `offset_sec`
+        if dataset_name == 'k700_2020':
+            header = ['path', 'vstart_sec', 'offset_sec', 'oos_target']
+        else:
+            header = next(reader)
+        for line in reader:
+            data = dict()
+            for f, value in zip(header, line):
+                if f == 'path':
+                    v = value
+                elif f == 'offset_sec':
+                    data[f] = float(value)
+                elif f in ['vstart_sec', 'v_start_sec']:
+                    f = 'v_start_i_sec'
+                    data[f] = float(value)
+                elif f == 'oos_target':
+                    data[f] = int(value)
+                else:
+                    data[f] = value
+            # assert v not in vid2offset_params, 'otherwise, offs from other splits will override each other'
+
+            # even if we have multiple splits (val=test), we want to make sure that the offsets are the same
+            if v in vid2offset_params:
+                assert all([vid2offset_params[v][k] == data[k] for k in data]), f'{v} isnt unique and vary'
+
+            vid2offset_params[v] = data
+    return vid2offset_params
+
+
+def maybe_cache_file(path: os.PathLike):
+    '''Motivation: if every job reads from a shared disk it`ll get very slow, consider an image can
+    be 2MB, then with batch size 32, 16 workers in dataloader you`re already requesting 1GB!! -
+    imagine this for all users and all jobs simultaneously.'''
+    # checking if we are on cluster, not on a local machine
+    if 'LOCAL_SCRATCH' in os.environ:
+        cache_dir = os.environ.get('LOCAL_SCRATCH')
+        # a bit ugly but we need not just fname to be appended to `cache_dir` but parent folders,
+        # otherwise the same fnames in multiple folders will create a bug (the same input for multiple paths)
+        cache_path = os.path.join(cache_dir, Path(path).relative_to('/'))
+        if not os.path.exists(cache_path):
+            os.makedirs(Path(cache_path).parent, exist_ok=True)
+            shutil.copyfile(path, cache_path)
+        return cache_path
+    else:
+        return path
+
+
+def get_video_and_audio(path, get_meta=False, start_sec=0, end_sec=None):
+    orig_path = path
+    path = maybe_cache_file(path)
+    # (Tv, 3, H, W) [0, 255, uint8]; (Ca, Ta)
+    rgb, audio, meta = torchvision.io.read_video(str(path), start_sec, end_sec, 'sec', output_format='TCHW')
+    assert meta['video_fps'], f'No video fps for {orig_path}'
+    # (Ta) <- (Ca, Ta)
+    audio = audio.mean(dim=0)
+    # FIXME: this is legacy format of `meta` as it used to be loaded by VideoReader.
+    meta = {'video': {'fps': [meta['video_fps']]}, 'audio': {'framerate': [meta['audio_fps']]}, }
+    return rgb, audio, meta
+
+
+def get_audio_stream(path, get_meta=False):
+    '''Used only in feature extractor training'''
+    path = str(Path(path).with_suffix('.wav'))
+    path = maybe_cache_file(path)
+    waveform, _ = torchaudio.load(path)
+    waveform = waveform.mean(dim=0)
+    if get_meta:
+        info = torchaudio.info(path)
+        duration = info.num_frames / info.sample_rate
+        meta = {'audio': {'duration': [duration], 'framerate': [info.sample_rate]}}
+        return waveform, meta
+    else:
+        return waveform
+
+def subsample_dataset(dataset: list, size_ratio: float, shuffle: bool = False):
+    if size_ratio is not None and 0.0 < size_ratio < 1.0:
+        logging.info(f'Subsampling dataset to {size_ratio}')
+        # shuffling is important only during subsampling (sometimes paths are sorted by class)
+        if shuffle:
+            random.shuffle(dataset)
+        cut_off = int(len(dataset) * size_ratio)
+        # making sure that we have at least one example
+        dataset = dataset[:max(1, cut_off)]
+        logging.info(f'Subsampled dataset to {size_ratio} (size: {len(dataset)})')
+    return dataset

modules/dataset/lrs.py

@@ -0,0 +1,192 @@
+import csv
+import logging
+import os
+import random
+import sys
+from glob import glob
+from pathlib import Path
+
+import torch
+
+
+sys.path.insert(0, '.')  # nopep8
+from dataset.dataset_utils import (get_fixed_offsets, get_video_and_audio, subsample_dataset)
+
+
+class LRS3(torch.utils.data.Dataset):
+
+    def __init__(self,
+                 split,
+                 vids_dir,
+                 transforms=None,
+                 splits_path='./data',
+                 seed=1337,
+                 load_fixed_offsets_on=['valid', 'test'],
+                 vis_load_backend='VideoReader',
+                 size_ratio=None,
+                 attr_annot_path=None,
+                 max_attr_per_vid=None,
+                 to_filter_bad_examples=True,):
+        super().__init__()
+        self.max_clip_len_sec = 11
+        logging.info(f'During IO, the length of clips is limited to {self.max_clip_len_sec} sec')
+        self.split = split
+        self.vids_dir = vids_dir
+        self.transforms = transforms
+        self.splits_path = splits_path
+        self.seed = seed
+        self.load_fixed_offsets_on = [] if load_fixed_offsets_on is None else load_fixed_offsets_on
+        self.vis_load_backend = vis_load_backend
+        self.size_ratio = size_ratio
+
+        split_clip_ids_path = os.path.join(splits_path, f'lrs3_{split}.txt')
+        if not os.path.exists(split_clip_ids_path):
+            vid_folder = Path(vids_dir) / 'pretrain'
+            clip_paths = sorted(vid_folder.rglob('*/*.mp4'))
+            if to_filter_bad_examples:
+                clip_paths = self.filter_bad_examples(clip_paths)
+            self.make_split_files(clip_paths)
+
+        # read the ids from a split
+        split_clip_ids = sorted(open(split_clip_ids_path).read().splitlines())
+
+        # make paths from the ids
+        clip_paths = [os.path.join(vids_dir, v + '.mp4') for v in split_clip_ids]
+
+        if split in self.load_fixed_offsets_on:
+            logging.info(f'Using fixed offset for {split}')
+            self.vid2offset_params = get_fixed_offsets(transforms, split, splits_path, 'lrs3')
+
+        self.dataset = clip_paths
+        self.dataset = subsample_dataset(self.dataset, size_ratio, shuffle=split == 'train')
+
+        logging.info(f'{split} has {len(self.dataset)} items')
+
+    def __getitem__(self, index):
+        path = self.dataset[index]
+        rgb, audio, meta = get_video_and_audio(path, get_meta=True, end_sec=self.max_clip_len_sec)
+
+        # (Tv, 3, H, W) in [0, 225], (Ta, C) in [-1, 1]
+        item = {'video': rgb, 'audio': audio, 'meta': meta, 'path': path, 'targets': {}, 'split': self.split}
+
+        # loading fixed offsets so we could evaluate on the same data each time (valid and test)
+        if self.split in self.load_fixed_offsets_on:
+            unique_id = path.replace(f'{self.vids_dir}/', '').replace(self.vids_dir, '').replace('.mp4', '')
+            offset_params = self.vid2offset_params[unique_id]
+            item['targets']['offset_sec'] = offset_params['offset_sec']
+            item['targets']['v_start_i_sec'] = offset_params['v_start_i_sec']
+            if 'oos_target' in offset_params:
+                item['targets']['offset_target'] = {
+                    'oos': offset_params['oos_target'], 'offset': item['targets']['offset_sec'],
+                }
+
+        if self.transforms is not None:
+            item = self.transforms(item)
+
+        return item
+
+    def filter_bad_examples(self, paths):
+        bad = set()
+        base_path = Path('./data/filtered_examples_lrs')
+        lists = [open(p).read().splitlines() for p in sorted(glob(str(base_path / '*.txt')))]
+        for s in lists:
+            bad = bad.union(s)
+        logging.info(f'Number of clips before filtering: {len(paths)}')
+        video_ids = [str(i).replace(self.vids_dir, '') for i in paths]
+        video_ids = [str(i).replace(f'{self.vids_dir}/', '') for i in video_ids]
+        paths = sorted([r for r in video_ids if r not in bad])
+        logging.info(f'Number of clips after filtering: {len(paths)}')
+        return paths
+
+    def make_split_files(self, paths):
+        logging.warning(f'The split files do not exist @ {self.splits_path}. Calculating the new ones.')
+
+        # will be splitting using videos, not clips to prevent train-test intersection
+        all_vids = sorted(list(set([Path(p).parent.name for p in paths])))
+        random.Random(self.seed).shuffle(all_vids)
+
+        # 0.1: splits are 8:1:1
+        hold_out_ratio = 0.1
+        hold_out_size = int(len(all_vids) * hold_out_ratio)
+        test_vids, train_valid_vids = all_vids[:hold_out_size], all_vids[hold_out_size:]
+        valid_vids, train_vids = train_valid_vids[:hold_out_size], train_valid_vids[hold_out_size:]
+
+        # making files
+        for phase, vids in zip(['train', 'valid', 'test'], [train_vids, valid_vids, test_vids]):
+            with open(os.path.join(self.splits_path, f'lrs3_{phase}.txt'), 'w') as wfile:
+                for path in paths:
+                    vid_name = Path(path).parent.name
+                    # just in the case I forgot the trailing '/' in the path
+                    unique_id = path.replace(f'{self.vids_dir}/', '').replace(self.vids_dir, '') \
+                                    .replace('.mp4', '')
+                    if vid_name in vids:
+                        wfile.write(unique_id + '\n')
+
+    def __len__(self):
+        return len(self.dataset)
+
+class LongerLRS3(LRS3):
+    '''This class is different to the parent in the extra filtering it does. If the parent was
+    making the splits with filtering for shorter than 9 second, this class filters for shorter than 9.5 sec.
+    by applying extra filtering.
+    '''
+
+    def __init__(self,
+                 split,
+                 vids_dir,
+                 transforms=None,
+                 splits_path='./data',
+                 seed=1337,
+                 load_fixed_offsets_on=['valid', 'test'],
+                 vis_load_backend='VideoReader',
+                 size_ratio=None,
+                 attr_annot_path=None,
+                 max_attr_per_vid=None,
+                 to_filter_bad_examples=True,):
+        # size_ratio is not used here as we are doing it this class (avoiding double subsampling)
+        super().__init__(split, vids_dir, transforms, splits_path, seed, load_fixed_offsets_on,
+                         vis_load_backend, None, attr_annot_path, max_attr_per_vid,
+                         to_filter_bad_examples)
+        # does extra filtering
+        if to_filter_bad_examples:
+            self.dataset = self.filter_bad_examples(self.dataset)
+        self.dataset = subsample_dataset(self.dataset, size_ratio, shuffle=split == 'train')
+        logging.info(f'{split} has {len(self.dataset)} items')
+
+    def filter_bad_examples(self, paths):
+        bad = set()
+        base_path = Path('./data/filtered_examples_lrs_extra')
+        lists = [open(p).read().splitlines() for p in sorted(glob(str(base_path / '*.txt')))]
+        for s in lists:
+            bad = bad.union(s)
+        logging.info(f'Number of clips before filtering: {len(paths)}')
+        video_ids = [str(i).replace(self.vids_dir, '').replace(f'{self.vids_dir}/', '') for i in paths]
+        paths = sorted([os.path.join(self.vids_dir, r) for r in video_ids if r not in bad])
+        logging.info(f'Number of clips after filtering: {len(paths)}')
+        return paths
+
+
+if __name__ == '__main__':
+    from time import time
+    from omegaconf import OmegaConf
+    import sys
+    sys.path.insert(0, '.')  # nopep8
+    from scripts.train_utils import get_transforms
+    from utils.utils import cfg_sanity_check_and_patch
+    cfg = OmegaConf.load('./configs/sparse_sync.yaml')
+    cfg.data.vids_path = '/scratch/local/hdd/vi/data/lrs3/h264_uncropped_25fps_256side_16000hz_aac/'
+    cfg.data.dataset.params.load_fixed_offsets_on = ['valid', 'test']
+
+    cfg_sanity_check_and_patch(cfg)
+    transforms = get_transforms(cfg)
+
+    datasets = {
+        'train': LRS3('train', cfg.data.vids_path, transforms['train'], load_fixed_offsets_on=[]),
+        'valid': LRS3('valid', cfg.data.vids_path, transforms['test'], load_fixed_offsets_on=[]),
+        'test': LRS3('test', cfg.data.vids_path, transforms['test'], load_fixed_offsets_on=[]),
+    }
+    for phase in ['train', 'valid', 'test']:
+        print(phase, len(datasets[phase]))
+
+    print(datasets['train'][1]['audio'].shape, datasets['train'][1]['video'].shape)
+    print(datasets['valid'][1]['audio'].shape, datasets['valid'][1]['video'].shape)

modules/dataset/transforms.py

@@ -0,0 +1,1074 @@
+import logging
+import math
+import random
+from typing import Tuple
+import torch
+import torchvision
+import torchaudio
+import numpy as np
+import einops
+
+
+def sec2frames(sec, fps):
+    return int(sec * fps)
+
+def frames2sec(frames, fps):
+    return frames / fps
+
+
+class EqualifyFromRight(torch.nn.Module):
+
+    def __init__(self, clip_max_len_sec=10):
+        '''
+        Takes the dataset item and makes sure more streams are of an equal size in terms of fps.
+        It, however, assumes that the signal is synched and trims the ending parts ('from the right').
+        '''
+        super().__init__()
+        self.clip_max_len_sec = clip_max_len_sec
+
+    def forward(self, item):
+        '''
+        `item`: {'video': (Tv, C, H, W), 'audio': (Ta,),
+                 'meta': {
+                     'audio': {'framerate': [float], 'duration': [float]}
+                     'video': {'fps': [float], 'duration': [float]}}
+        '''
+        a_fps = item['meta']['audio']['framerate'][0]
+        v_fps = item['meta']['video']['fps'][0]
+
+        Ta = item['audio'].shape[0]
+        Tv, C, H, W = item['video'].shape
+
+        a_len_secs = Ta / a_fps
+        v_len_secs = Tv / v_fps
+        min_len = min(self.clip_max_len_sec, a_len_secs, v_len_secs)
+
+        a_frames_per_v_frame = a_fps // v_fps
+        v_len_frames = int(v_fps * min_len)
+        a_len_frames = int(a_frames_per_v_frame * v_len_frames)
+        # print(a_len_frames, v_len_frames)
+
+        assert a_len_frames <= Ta and v_len_frames <= Tv
+
+        item['audio'] = item['audio'][:a_len_frames]
+        item['video'] = item['video'][:v_len_frames, :, :, :]
+
+        return item
+
+
+class RGBSpatialCrop(torch.nn.Module):
+
+    def __init__(self, input_size, is_random):
+        super().__init__()
+        assert input_size is not None, f'smaller_input_size is `{input_size}`'
+        if isinstance(input_size, int):
+            input_size = (input_size, input_size)
+        self.input_size = input_size
+        self.is_random = is_random
+
+    @staticmethod
+    def get_random_crop_sides(vid, output_size):
+        '''Slice parameters for random crop'''
+        h, w = vid.shape[-2:]
+        th, tw = output_size
+        if w == tw and h == th:
+            return 0, 0, h, w
+        i = random.randint(0, h - th)
+        j = random.randint(0, w - tw)
+        return i, j, th, tw
+
+    @staticmethod
+    def get_center_crop_sides(vid, output_size):
+        '''Slice parameters for center crop'''
+        h, w = vid.shape[-2:]
+        th, tw = output_size
+
+        i = int(round((h - th) / 2.))
+        j = int(round((w - tw) / 2.))
+        return i, j, th, tw
+
+    def forward(self, item):
+        # (Tv, C, H, W)
+        vid = item['video']
+        if self.is_random:
+            i, j, h, w = self.get_random_crop_sides(vid, self.input_size)
+        else:
+            i, j, h, w = self.get_center_crop_sides(vid, self.input_size)
+        item['video'] = vid[..., i:(i + h), j:(j + w)]
+        return item
+
+class Resize(torchvision.transforms.Resize):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, item):
+        item['video'] = super().forward(item['video'])
+        return item
+
+
+class RGBSpatialCropSometimesUpscale(torch.nn.Module):
+    '''This (randomly) crops the input video and with prob `sometimes_p` this crop is smaller but upscaled
+    to `target_input_size`'''
+
+    def __init__(self, sometimes_p, target_input_size, is_random, smaller_input_size=None):
+        super().__init__()
+        self.sometimes_p = sometimes_p
+        self.do_sometimes_upscale = sometimes_p is not None and sometimes_p > 0
+
+        self.crop_only = RGBSpatialCrop(target_input_size, is_random)
+
+        if self.do_sometimes_upscale:
+            self.crop_further_and_upscale = torchvision.transforms.Compose([
+                RGBSpatialCrop(smaller_input_size, is_random),
+                Resize(target_input_size, antialias=None),
+            ])
+
+    def forward(self, item):
+        assert len(item['video'].shape) == 4, \
+            f"{item['video'].shape}: if it is applied after GenerateMultipleClips," \
+            "augs should be applied to each clip separately, not to the whole video array. " \
+            "Otherwise, ignore this warning (comment it)."
+        if self.do_sometimes_upscale and self.sometimes_p > torch.rand(1):
+            return self.crop_further_and_upscale(item)
+        else:
+            return self.crop_only(item)
+
+
+class RandomApplyColorDistortion(torch.nn.Module):
+
+    def __init__(self, p_gray_scale=0., p_color_jitter=0., s=1.) -> None:
+        super().__init__()
+        self.p_gray_scale = p_gray_scale
+        self.p_color_jitter = p_color_jitter
+        self.s = s
+        assert 0 <= self.p_color_jitter <= 1 and 0 <= self.p_gray_scale <= 1, (p_color_jitter, p_gray_scale)
+        # SimCLR params
+        color_jitter = torchvision.transforms.ColorJitter(0.8*s, 0.8*s, 0.8*s, 0.2*s)
+        rand_color_jitter = torchvision.transforms.RandomApply([color_jitter], p_color_jitter)
+        rand_gray = torchvision.transforms.RandomGrayscale(p_gray_scale)
+        self.transforms = torchvision.transforms.Compose([rand_color_jitter, rand_gray])
+
+    def apply_to_single_clip(self, clip):
+        return self.transforms(clip)
+
+    def apply_to_each_clip(self, clips):
+        for i, clip in enumerate(clips):
+            clips[i] = self.apply_to_single_clip(clip)
+        return clips
+
+    def forward(self, item):
+        has_batch_dim = len(item['video'].shape) == 5
+        if has_batch_dim:
+            fn = self.apply_to_each_clip
+        else:
+            fn = self.apply_to_single_clip
+        item['video'] = fn(item['video'])
+        return item
+
+
+class ApplyColorJitterFrameWise(torch.nn.Module):
+
+    def __init__(self, s=1.) -> None:
+        super().__init__()
+        self.s = s
+        # SimCLR params
+        self.transform = torchvision.transforms.ColorJitter(0.8*s, 0.8*s, 0.8*s, 0.2*s)
+
+    def apply_to_single_clip(self, clip):
+        for i, frame in enumerate(clip):
+            clip[i] = self.transform(frame)
+        return clip
+
+    def apply_to_each_clip(self, clips):
+        for i, clip in enumerate(clips):
+            clips[i] = self.apply_to_single_clip(clip)
+        return clips
+
+    def forward(self, item):
+        has_batch_dim = len(item['video'].shape) == 5
+        if has_batch_dim:
+            fn = self.apply_to_each_clip
+        else:
+            fn = self.apply_to_single_clip
+        item['video'] = fn(item['video'])
+        return item
+
+
+class RandomHorizontalFlip(torchvision.transforms.RandomHorizontalFlip):
+
+    def __init__(self, p=0.5):
+        super().__init__(p)
+
+    def apply_to_single_clip(self, clip):
+        return super().forward(clip)
+
+    def apply_to_each_clip(self, clips):
+        for i, clip in enumerate(clips):
+            clips[i] = self.apply_to_single_clip(clip)
+        return clips
+
+    def forward(self, item):
+        has_batch_dim = len(item['video'].shape) == 5
+        if has_batch_dim:
+            fn = self.apply_to_each_clip
+        else:
+            fn = self.apply_to_single_clip
+        item['video'] = fn(item['video'])
+        return item
+
+
+def make_class_grid(leftmost_val, rightmost_val, grid_size, add_extreme_offset: bool = False,
+                    seg_size_vframes: int = None, nseg: int = None, step_size_seg: float = None,
+                    vfps: float = None):
+    assert grid_size >= 3, f'grid_size: {grid_size} doesnot make sense. If =2 -> (-1,1); =1 -> (-1); =0 -> ()'
+    grid = torch.from_numpy(np.linspace(leftmost_val, rightmost_val, grid_size)).float()
+    if add_extreme_offset:
+        assert all([seg_size_vframes, nseg, step_size_seg]), f'{seg_size_vframes} {nseg} {step_size_seg}'
+        seg_size_sec = seg_size_vframes / vfps
+        trim_size_in_seg = nseg - (1 - step_size_seg) * (nseg - 1)
+        extreme_value = trim_size_in_seg * seg_size_sec
+        grid = torch.cat([grid, torch.tensor([extreme_value])])  # adding extreme offset to the class grid
+    return grid
+
+
+def quantize_offset(grid: torch.Tensor, off_sec: float) -> Tuple[float, int]:
+    '''Takes in the offset in seconds and snaps it onto the closest grid element.
+    Returns the grid value and its index.'''
+    closest_grid_el = (grid - off_sec).abs().argmin()
+    return grid[closest_grid_el], closest_grid_el
+
+def apply_a_jitter(a_start_i, a_len_frames, a_crop_len_frames, a_fps, max_a_jitter_sec):
+    max_a_start_i = a_len_frames - a_crop_len_frames
+    max_a_jitter_i = sec2frames(max_a_jitter_sec, a_fps)
+    max_a_jitter_i_left = min(a_start_i, max_a_jitter_i)
+    max_a_jitter_i_right = min(max_a_start_i - a_start_i, max_a_jitter_i)
+    # jitter is U[left, right]
+    a_jitter_i = random.randint(-max_a_jitter_i_left, max_a_jitter_i_right)
+    # apply jitter
+    a_start_i = a_start_i + a_jitter_i
+    # making sure that any value from `a_start_i + U[left, right]` will be inside of [0, len-crop] region
+    assert 0 <= a_start_i <= max_a_start_i, f'{a_jitter_i} {max_a_jitter_i_left} {max_a_jitter_i_right} {max_a_start_i}'
+    return a_start_i, a_jitter_i
+
+
+class TemporalCropAndOffset(torch.nn.Module):
+
+    def __init__(self, crop_len_sec: float, max_off_sec: float, offset_type='grid', do_offset: bool = True,
+                 grid_size: int = None, max_wiggle_sec: float = None, add_doubt_cls: bool = False,
+                 segment_size_vframes: int = None, n_segments: int = None, step_size_seg: float = None,
+                 vfps: float = None, prob_oos: float = None):
+        super().__init__()
+        self.crop_len_sec = crop_len_sec
+        self.do_offset = do_offset
+        self.grid_size = grid_size
+        self.offset_type = offset_type
+        self.max_off_sec = max_off_sec
+        self.max_a_jitter_sec = max_wiggle_sec
+        if do_offset:
+            if offset_type == 'grid':
+                self.class_grid = make_class_grid(-max_off_sec, max_off_sec, grid_size, add_doubt_cls,
+                                                  segment_size_vframes, n_segments, step_size_seg, vfps)
+                logging.info(f'Offsets class grid: {self.class_grid}')
+                if self.max_a_jitter_sec is not None:
+                    assert (max_wiggle_sec-1e-6) <= ((self.class_grid[1] - self.class_grid[0]) / 2), f'{self.class_grid}'
+            elif offset_type == 'uniform':
+                self.off_dist = torch.distributions.uniform.Uniform(-max_off_sec, max_off_sec)
+                logging.info(f'Offset uniform distribution: {self.off_dist}')
+            elif offset_type == 'uniform_binary':
+                self.itu_t_range = (-0.125, 0.045)
+                self.prob_oos = prob_oos
+                self.ins_dist = torch.distributions.uniform.Uniform(self.itu_t_range[0], self.itu_t_range[1])
+                self.off_dist = torch.distributions.uniform.Uniform(-max_off_sec, max_off_sec)
+            else:
+                raise NotImplementedError(f'Unknown offset type: {offset_type}')
+
+    def forward(self, item):
+        vid = item['video']
+        aud = item['audio']
+        v_len_frames, C, H, W = vid.shape
+        a_len_frames = aud.shape[0]
+
+        v_fps = int(item['meta']['video']['fps'][0])
+        a_fps = int(item['meta']['audio']['framerate'][0])
+
+        v_crop_len_frames = sec2frames(self.crop_len_sec, v_fps)
+        a_crop_len_frames = sec2frames(self.crop_len_sec, a_fps)
+
+        if self.do_offset:
+            # trying to get the offset parameters (for instance during valid and test we have fixed offsets)
+            offset_sec = item['targets'].get('offset_sec', None)
+            v_start_i_sec = item['targets'].get('v_start_i_sec', None)
+            if 'offset_target' in item['targets']:
+                is_oos = item['targets']['offset_target'].get('oos', None)
+            # train-time
+            if offset_sec is None and v_start_i_sec is None:
+                # aud starts `offset_sec` earlier than it should; aud has what will be shown after offset_sec
+                if self.offset_type == 'grid':
+                    offset_sec = random.choice(self.class_grid.tolist())
+                elif self.offset_type == 'uniform':
+                    offset_sec = self.off_dist.sample().item()
+                elif self.offset_type == 'uniform_binary':
+                    # in-sync: Uniform(-0.125, 0.045)
+                    # out-of-sync: Uniform(-5.5, 5.5) and resampled until not in Uniform(-0.125, 0.045)
+                    # first, we sample if the offset is out-of-sync with prob_oss
+                    is_oos = (torch.rand(1) < self.prob_oos).item()
+                    if is_oos:
+                        # second, we sample the offset itself (if in in-sync range, trying again)
+                        offset_sec = self.off_dist.sample().item()
+                        while self.itu_t_range[0] <= offset_sec <= self.itu_t_range[1]:
+                            offset_sec = self.off_dist.sample().item()
+                    else:
+                        offset_sec = self.ins_dist.sample().item()
+                offset_sec = round(offset_sec, 2)
+                v_start_max_sec = frames2sec(v_len_frames - v_crop_len_frames, v_fps)
+                assert v_start_max_sec > 0, f'{v_len_frames} {v_crop_len_frames} {v_fps} @ {item["path"]}'
+                # `v_start_sec` IS NOT rounded to the fps grid
+                v_start_sec = random.uniform(max(0, -offset_sec), min(v_start_max_sec, v_start_max_sec-offset_sec))
+                assert 0 <= v_start_sec <= v_start_max_sec, f'{v_start_sec} {v_start_max_sec} {item["path"]}'
+                v_start_i = sec2frames(v_start_sec, v_fps)
+                # `v_start_i_sec` IS rounded to the fps grid
+                v_start_i_sec = frames2sec(v_start_i, v_fps)
+            else:
+                offset_sec = round(offset_sec, 2)
+                v_start_i = sec2frames(v_start_i_sec, v_fps)
+            v_end_i = v_start_i + v_crop_len_frames
+            # `a_start_i` depends on the rounded value `v_start_i_sec`, otherwise
+            # (v_start_sec) we have ±0.1 jittering
+            a_start_i = sec2frames(v_start_i_sec + offset_sec, a_fps)
+        else:
+            offset_sec = 0.0
+            is_random_crop = item['split'] == 'train'
+            v_start_i, v_end_i = self.get_crop_idx(v_len_frames, v_crop_len_frames, is_random=is_random_crop)
+            v_start_i_sec = frames2sec(v_start_i, v_fps)
+            a_start_i = sec2frames(v_start_i_sec, a_fps)
+
+        # sometimes due to the rounding error e.g. v_start_sec = 1.505 but sec2frames(1.505, 25) = 1.48
+        # given offset is -1.5, the a_start_i will be a small negative value. (likely a_fps * 1/v_fps * 0.5)
+        if a_start_i < 0:
+            how_much_out = a_start_i
+            logging.info(f'a_start_i is negative ({how_much_out}) at {item["path"]}')
+            if abs(how_much_out) <= a_fps / v_fps:
+                logging.info('fixing it')
+                a_start_i += abs(how_much_out)
+            else:
+                raise Exception(f'{how_much_out} {item["path"]}')
+
+        if self.max_a_jitter_sec is not None and self.max_a_jitter_sec > 0:
+            a_start_i, a_jitter_i = apply_a_jitter(a_start_i, a_len_frames, a_crop_len_frames, a_fps,
+                                                   self.max_a_jitter_sec)
+            item['meta']['a_jitter_i'] = a_jitter_i
+
+        a_end_i = a_start_i + a_crop_len_frames
+
+        assert v_start_i < v_end_i and a_start_i < a_end_i
+        assert aud.shape[0] >= a_end_i, f'{aud.shape} {a_end_i} {item["path"]}'
+        assert vid.shape[0] >= v_end_i, f'{vid.shape} {v_end_i} {item["path"]}'
+
+        vid, aud = vid[v_start_i:v_end_i, :, :, :], aud[a_start_i:a_end_i]
+
+        item['video'] = vid
+        item['audio'] = aud
+
+        assert item['video'].shape[0] == v_fps * self.crop_len_sec, f'{item["video"].shape} {item["path"]}'
+        assert item['audio'].shape[0] == a_fps * self.crop_len_sec, f'{item["audio"].shape} {item["path"]}'
+
+        # caching parameters
+        if self.do_offset:
+            if self.offset_type == 'grid':
+                offset_label, offset_target = quantize_offset(self.class_grid, offset_sec)
+            elif self.offset_type == 'uniform':
+                offset_label, offset_target = offset_sec, offset_sec
+            elif self.offset_type == 'uniform_binary':
+                offset_label, offset_target = offset_sec, {'oos': is_oos, 'offset': offset_sec}
+            item['targets']['offset_sec'] = offset_sec
+            item['targets']['v_start_i_sec'] = v_start_i_sec
+            item['targets']['offset_label'] = offset_label
+            # assert 'offset_target' not in item['targets'], f'{item["targets"]}. What passed it there?'
+            item['targets']['offset_target'] = offset_target
+
+        return item
+
+    def get_crop_idx(self, len_frames: int, crop_len_frames: int, is_random=True):
+        if len_frames == crop_len_frames:
+            return 0, len_frames
+        if is_random:
+            left_i = random.randint(0, len_frames - crop_len_frames)
+        else:
+            left_i = int(round((len_frames - crop_len_frames) / 2.))
+        return left_i, left_i+crop_len_frames
+
+
+class GenerateMultipleSegments(torch.nn.Module):
+    '''
+    Given an item with video and audio, generates a batch of `n_segments` segments
+    of length `segment_size_vframes` (if None, the max number of segments will be made).
+    If `is_start_random` is True, the starting position of the 1st segment will be random but respecting
+    n_segments.
+    `audio_jitter_sec` is the amount of audio offset in seconds.
+    '''
+
+    def __init__(self, segment_size_vframes: int, n_segments: int = None, is_start_random: bool = False,
+                 audio_jitter_sec: float = 0., step_size_seg: float = 1):
+        super().__init__()
+        self.segment_size_vframes = segment_size_vframes
+        self.n_segments = n_segments
+        self.is_start_random = is_start_random
+        self.audio_jitter_sec = audio_jitter_sec
+        self.step_size_seg = step_size_seg
+        logging.info(f'Segment step size: {self.step_size_seg}')
+
+    def forward(self, item):
+        v_len_frames, C, H, W = item['video'].shape
+        a_len_frames = item['audio'].shape[0]
+
+        v_fps = int(item['meta']['video']['fps'][0])
+        a_fps = int(item['meta']['audio']['framerate'][0])
+
+        ## Determining the number of segments
+        # segment size
+        segment_size_vframes = self.segment_size_vframes
+        segment_size_aframes = sec2frames(frames2sec(self.segment_size_vframes, v_fps), a_fps)
+        # step size (stride)
+        stride_vframes = int(self.step_size_seg * segment_size_vframes)
+        stride_aframes = int(self.step_size_seg * segment_size_aframes)
+        # calculating the number of segments. (W - F + 2P) / S + 1
+        n_segments_max_v = math.floor((v_len_frames - segment_size_vframes) / stride_vframes) + 1
+        n_segments_max_a = math.floor((a_len_frames - segment_size_aframes) / stride_aframes) + 1
+        # making sure audio and video can accommodate the same number of segments
+        n_segments_max = min(n_segments_max_v, n_segments_max_a)
+        n_segments = n_segments_max if self.n_segments is None else self.n_segments
+
+        assert n_segments <= n_segments_max, \
+            f'cant make {n_segments} segs of len {self.segment_size_vframes} in a vid ' \
+            f'of len {v_len_frames} for {item["path"]}'
+
+        # (n_segments, 2) each
+        v_ranges, a_ranges = self.get_sequential_seg_ranges(v_len_frames, a_len_frames, v_fps, a_fps,
+                                                            n_segments, segment_size_aframes)
+
+        # segmenting original streams (n_segments, segment_size_frames, C, H, W)
+        item['video'] = torch.stack([item['video'][s:e] for s, e in v_ranges], dim=0)
+        item['audio'] = torch.stack([item['audio'][s:e] for s, e in a_ranges], dim=0)
+        return item
+
+    def get_sequential_seg_ranges(self, v_len_frames, a_len_frames, v_fps, a_fps, n_seg, seg_size_aframes):
+        # if is_start_random is True, the starting position of the 1st segment will
+        # be random but respecting n_segments like so: "-CCCCCCCC---" (maybe with fixed overlap),
+        # else the segments are taken from the middle of the video respecting n_segments: "--CCCCCCCC--"
+
+        seg_size_vframes = self.segment_size_vframes  # for brevity
+
+        # calculating the step size in frames
+        step_size_vframes = int(self.step_size_seg * seg_size_vframes)
+        step_size_aframes = int(self.step_size_seg * seg_size_aframes)
+
+        # calculating the length of the sequence of segments (and in frames)
+        seg_seq_len = n_seg * self.step_size_seg + (1 - self.step_size_seg)
+        vframes_seg_seq_len = int(seg_seq_len * seg_size_vframes)
+        aframes_seg_seq_len = int(seg_seq_len * seg_size_aframes)
+
+        # doing temporal crop
+        max_v_start_i = v_len_frames - vframes_seg_seq_len
+        if self.is_start_random:
+            v_start_i = random.randint(0, max_v_start_i)
+        else:
+            v_start_i = max_v_start_i // 2
+        a_start_i = sec2frames(frames2sec(v_start_i, v_fps), a_fps)  # vid frames -> seconds -> aud frames
+
+        # make segments starts
+        v_start_seg_i = torch.tensor([v_start_i + i * step_size_vframes for i in range(n_seg)]).int()
+        a_start_seg_i = torch.tensor([a_start_i + i * step_size_aframes for i in range(n_seg)]).int()
+
+        # apply jitter to audio
+        if self.audio_jitter_sec > 0:
+            jitter_aframes = sec2frames(self.audio_jitter_sec, a_fps)
+            # making sure after applying jitter, the audio is still within the audio boundaries
+            jitter_aframes = min(jitter_aframes, a_start_i, a_len_frames-a_start_i-aframes_seg_seq_len)
+            a_start_seg_i += random.randint(-jitter_aframes, jitter_aframes)  # applying jitter to segments
+
+        # make segment ends
+        v_ends_seg_i = v_start_seg_i + seg_size_vframes
+        a_ends_seg_i = a_start_seg_i + seg_size_aframes  # using the adjusted a_start_seg_i (with jitter)
+
+        # make ranges
+        v_ranges = torch.stack([v_start_seg_i, v_ends_seg_i], dim=1)
+        a_ranges = torch.stack([a_start_seg_i, a_ends_seg_i], dim=1)
+        assert (a_ranges >= 0).all() and (a_ranges <= a_len_frames).all(), f'{a_ranges} out of {a_len_frames}'
+        assert (v_ranges <= v_len_frames).all(), f'{v_ranges} out of {v_len_frames}'
+        return v_ranges, a_ranges
+
+
+class TemporalCropAndOffsetForSyncabilityTraining(torch.nn.Module):
+
+    def __init__(self, max_off_sec: float, do_offset: bool = True,
+                 grid_size: int = None, max_wiggle_sec: float = None,
+                 segment_size_vframes: int = None, n_segments: int = None, step_size_seg: float = None,
+                 vfps: float = None):
+        super().__init__()
+        seg_size_sec = segment_size_vframes / vfps
+        trim_size_in_seg = n_segments - (1 - step_size_seg) * (n_segments - 1)
+        self.crop_len_sec = round(trim_size_in_seg * seg_size_sec, 2)
+        logging.info(f'Crop len: {self.crop_len_sec}')
+        self.do_offset = do_offset
+        self.grid_size = grid_size
+        self.max_off_sec = max_off_sec
+        self.max_a_jitter_sec = max_wiggle_sec
+        self.segment_size_vframes = segment_size_vframes
+        self.n_segments = n_segments
+        self.step_size_seg = step_size_seg
+        self.prob_syncable = 0.5
+        if do_offset:
+            self.class_grid = make_class_grid(-max_off_sec, max_off_sec, grid_size)
+            logging.info(f'Offset class grid: {self.class_grid}')
+            if self.max_a_jitter_sec is not None:
+                assert (max_wiggle_sec-1e-6) <= ((self.class_grid[1] - self.class_grid[0]) / 2), f'{self.class_grid}'
+
+    def forward(self, item):
+        vid = item['video']
+        aud = item['audio']
+        v_len_frames, C, H, W = vid.shape
+        a_len_frames = aud.shape[0]
+
+        v_fps = int(item['meta']['video']['fps'][0])
+        a_fps = int(item['meta']['audio']['framerate'][0])
+
+        v_crop_len_frames = sec2frames(self.crop_len_sec, v_fps)
+        a_crop_len_frames = sec2frames(self.crop_len_sec, a_fps)
+
+        if self.do_offset:
+            # trying to get the offset parameters (for instance during valid and test we have fixed offsets)
+            offset_sec = item['targets'].get('offset_sec', None)
+            v_start_i_sec = item['targets'].get('v_start_i_sec', None)
+            # train-time
+            if offset_sec is None and v_start_i_sec is None:
+
+                # for the syncability training, we want to have a syncable or non-syncable offset with 50% prob
+                offset_is_syncable = random.random() < self.prob_syncable  # 1=syncable, 0=non-syncable
+                if offset_is_syncable:
+                    offset_sec = random.choice(self.class_grid.tolist())
+                else:
+                    offset_sec = random.choice([-self.crop_len_sec, self.crop_len_sec])  # either - or + offset
+                # aud starts `offset_sec` earlier than it should; aud has what will be shown after offset_sec
+
+                offset_sec = round(offset_sec, 2)
+                v_start_max_sec = frames2sec(v_len_frames - v_crop_len_frames, v_fps)
+                assert v_start_max_sec > 0, f'{v_len_frames} {v_crop_len_frames} {v_fps} @ {item["path"]}'
+                # `v_start_sec` IS NOT rounded to the fps grid
+                v_start_sec = random.uniform(max(0, -offset_sec), min(v_start_max_sec, v_start_max_sec-offset_sec))
+                assert 0 <= v_start_sec <= v_start_max_sec, f'{v_start_sec} {v_start_max_sec} {item["path"]}'
+                v_start_i = sec2frames(v_start_sec, v_fps)
+                v_end_i = v_start_i + v_crop_len_frames
+                # `v_start_i_sec` IS rounded to the fps grid
+                v_start_i_sec = frames2sec(v_start_i, v_fps)
+                # `a_start_i` depends on the rounded value `v_start_i_sec`, otherwise
+                # (v_start_sec) we have ±0.1 jittering
+                a_start_i = sec2frames(v_start_i_sec + offset_sec, a_fps)
+                if self.max_a_jitter_sec is not None and self.max_a_jitter_sec > 0:
+                    a_start_i, a_jitter_i = apply_a_jitter(a_start_i, a_len_frames, a_crop_len_frames, a_fps,
+                                                           self.max_a_jitter_sec)
+                    item['meta']['a_jitter_i'] = a_jitter_i
+                a_end_i = a_start_i + a_crop_len_frames
+            else:
+                offset_sec = round(offset_sec, 2)
+                v_start_i = sec2frames(v_start_i_sec, v_fps)
+                a_start_i = sec2frames(v_start_i_sec + offset_sec, a_fps)
+                v_end_i = v_start_i + v_crop_len_frames
+                a_end_i = a_start_i + a_crop_len_frames
+        else:
+            offset_sec = 0.0
+            is_random_crop = item['split'] == 'train'
+            v_start_i, v_end_i = self.get_crop_idx(v_len_frames, v_crop_len_frames, is_random=is_random_crop)
+            v_start_i_sec = frames2sec(v_start_i, v_fps)
+            a_start_i = sec2frames(v_start_i_sec, a_fps)
+            if self.max_a_jitter_sec is not None and self.max_a_jitter_sec > 0:
+                a_start_i, a_jitter_i = apply_a_jitter(a_start_i, a_len_frames, a_crop_len_frames, a_fps,
+                                                       self.max_a_jitter_sec)
+                item['meta']['a_jitter_i'] = a_jitter_i
+            a_end_i = a_start_i + a_crop_len_frames
+
+        # sometimes due to the rounding error e.g. v_start_sec = 1.505 but sec2frames(1.505, 25) = 1.48
+        # given offset is -1.5, the a_start_i will be a small negative value. (likely a_fps * 1/v_fps * 0.5)
+        if a_start_i < 0:
+            how_much_out = a_start_i
+            logging.info(f'a_start_i is negative ({how_much_out}) at {item["path"]}')
+            if abs(how_much_out) <= a_fps / v_fps:
+                logging.info('fixing it')
+                a_start_i += abs(how_much_out)
+                a_end_i += abs(how_much_out)
+            else:
+                raise Exception(f'{how_much_out} {item["path"]}')
+
+        assert v_start_i < v_end_i and a_start_i < a_end_i
+        assert aud.shape[0] >= a_end_i, f'{aud.shape} {a_end_i} {item["path"]}'
+        assert vid.shape[0] >= v_end_i, f'{vid.shape} {v_end_i} {item["path"]}'
+
+        vid, aud = vid[v_start_i:v_end_i, :, :, :], aud[a_start_i:a_end_i]
+
+        item['video'] = vid
+        item['audio'] = aud
+
+        assert item['video'].shape[0] == int(v_fps*self.crop_len_sec), f'{item["video"].shape} {item["path"]}'
+        assert item['audio'].shape[0] == int(a_fps*self.crop_len_sec), f'{item["audio"].shape} {item["path"]}'
+
+        # caching parameters
+        if self.do_offset:
+            # NOTE: this is useless for the extreme offsetting
+            offset_label, offset_target = quantize_offset(self.class_grid, offset_sec)
+            item['targets']['offset_sec'] = offset_sec
+            item['targets']['offset_label'] = offset_label
+            # assert 'offset_target' not in item['targets'], f'{item["targets"]}. What passed it there?'
+            item['targets']['offset_target'] = offset_target
+            item['targets']['v_start_i_sec'] = v_start_i_sec
+            item['targets']['sync_target'] = int(offset_is_syncable)
+
+        return item
+
+    def get_crop_idx(self, len_frames: int, crop_len_frames: int, is_random=True):
+        if len_frames == crop_len_frames:
+            return 0, len_frames
+        if is_random:
+            left_i = random.randint(0, len_frames - crop_len_frames)
+        else:
+            left_i = int(round((len_frames - crop_len_frames) / 2.))
+        return left_i, left_i+crop_len_frames
+
+
+class RGBToFloatToZeroOne(torch.nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, item):
+        item['video'] = item['video'].to(torch.float32).div(255.)
+        return item
+
+
+class RGBToHalfToZeroOne(torch.nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, item):
+        item['video'] = item['video'].half().div(255.)
+        return item
+
+
+class RGBNormalize(torchvision.transforms.Normalize):
+    '''The same as the torchvision`s but with different interface for the dict.
+    This should work for any shape (..., C, H, W)'''
+
+    def __init__(self, mean, std, inplace=False):
+        super().__init__(mean, std, inplace)
+        logging.info(f'RGBNormalize: mean={mean}, std={std}')
+
+    def forward(self, item):
+        item['video'] = super().forward(item['video'])
+        item['meta']['video']['norm_stats'] = {'mean': torch.as_tensor(self.mean),
+                                               'std': torch.as_tensor(self.std)}
+        return item
+
+
+class AudioRandomVolume(torch.nn.Module):
+
+    def __init__(self, p: float, **kwargs):
+        super().__init__()
+        transform = torchaudio.transforms.Vol(**kwargs)
+        self.transform = torchvision.transforms.RandomApply([transform], p)
+
+    def apply_to_single_clip(self, clip):
+        return self.transform(clip)
+
+    def apply_to_each_clip(self, clips):
+        for i, clip in enumerate(clips):
+            clips[i] = self.apply_to_single_clip(clip)
+        return clips
+
+    def forward(self, item):
+        has_batch_dim = len(item['audio'].shape) == 2
+        if has_batch_dim:
+            fn = self.apply_to_each_clip
+        else:
+            fn = self.apply_to_single_clip
+        item['audio'] = fn(item['audio'])
+        return item
+
+
+class AudioRandomLowpassFilter(torch.nn.Module):
+
+    def __init__(self, p: float, cutoff_freq: float, Q: float = 0.707):
+        super().__init__()
+        self.p = p
+        self.cutoff_freq = cutoff_freq
+        self.Q = Q
+
+    def apply_to_single_clip(self, clip, sr):
+        if self.p > torch.rand(1):
+            return torchaudio.functional.lowpass_biquad(clip, sr, self.cutoff_freq, self.Q)
+        else:
+            return clip
+
+    def apply_to_each_clip(self, clips, sr):
+        for i, clip in enumerate(clips):
+            clips[i] = self.apply_to_single_clip(clip, sr)
+        return clips
+
+    def forward(self, item):
+        has_batch_dim = len(item['audio'].shape) == 2
+        sr = int(item['meta']['audio']['framerate'][0])
+        if has_batch_dim:
+            fn = self.apply_to_each_clip
+        else:
+            fn = self.apply_to_single_clip
+        item['audio'] = fn(item['audio'], sr)
+        return item
+
+
+class AudioRandomPitchShift(torch.nn.Module):
+
+    def __init__(self, p: float, shift: int) -> None:
+        super().__init__()
+        self.p = p
+        self.shift = shift
+
+    def apply_to_single_clip(self, wave, sr):
+        if self.p > torch.rand(1):
+            effects = [['pitch', f'{self.shift}'], ['rate', f'{sr}']]
+            wave = wave.unsqueeze(0)
+            wave, _ = torchaudio.sox_effects.apply_effects_tensor(wave, sr, effects)
+            wave = wave.squeeze(0)
+        return wave
+
+    def apply_to_each_clip(self, waves, sr):
+        for i, wave in enumerate(waves):
+            waves[i] = self.apply_to_single_clip(wave, sr)
+        return waves
+
+    def forward(self, item):
+        has_batch_dim = len(item['audio'].shape) == 2
+        sr = int(item['meta']['audio']['framerate'][0])
+        if has_batch_dim:
+            fn = self.apply_to_each_clip
+        else:
+            fn = self.apply_to_single_clip
+        item['audio'] = fn(item['audio'], sr)
+        return item
+
+
+class AudioRandomReverb(torch.nn.Module):
+
+    def __init__(self, p: float) -> None:
+        super().__init__()
+        self.p = p
+        self.effects = [['reverb', '-w']]
+
+    def apply_to_single_clip(self, wave, fps):
+        if self.p > torch.rand(1):
+            wave = wave.unsqueeze(0)
+            wave, _ = torchaudio.sox_effects.apply_effects_tensor(wave, fps, self.effects)
+            wave = wave.mean(dim=0)
+        return wave
+
+    def apply_to_each_clip(self, waves, fps):
+        for i, wave in enumerate(waves):
+            waves[i] = self.apply_to_single_clip(wave, fps)
+        return waves
+
+    def forward(self, item):
+        has_batch_dim = len(item['audio'].shape) == 2
+        sr = int(item['meta']['audio']['framerate'][0])
+        if has_batch_dim:
+            fn = self.apply_to_each_clip
+        else:
+            fn = self.apply_to_single_clip
+        item['audio'] = fn(item['audio'], sr)
+        return item
+
+class AudioRandomGaussNoise(torch.nn.Module):
+
+    def __init__(self, p: float, amplitude=0.01) -> None:
+        super().__init__()
+        self.p = p
+        self.amplitude = amplitude
+
+    def apply_to_single_clip(self, wave):
+        if self.p > torch.rand(1):
+            noise = torch.randn_like(wave, dtype=wave.dtype)
+            wave = wave + self.amplitude * noise
+        return wave
+
+    def apply_to_each_clip(self, waves):
+        for i, wave in enumerate(waves):
+            waves[i] = self.apply_to_single_clip(wave)
+        return waves
+
+    def forward(self, item):
+        has_batch_dim = len(item['audio'].shape) == 2
+        if has_batch_dim:
+            fn = self.apply_to_each_clip
+        else:
+            fn = self.apply_to_single_clip
+        item['audio'] = fn(item['audio'])
+        return item
+
+
+class AudioMelSpectrogram(torch.nn.Module):
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.spec = torchaudio.transforms.MelSpectrogram(**kwargs)
+
+    def forward(self, item):
+        item['audio'] = self.spec(item['audio'])  # safe for batched input
+        return item
+
+
+class AudioLog(torch.nn.Module):
+
+    def __init__(self, eps=1e-6) -> None:
+        super().__init__()
+        self.eps = eps
+
+    def forward(self, item):
+        item['audio'] = torch.log(item['audio'] + self.eps)
+        return item
+
+class PadOrTruncate(torch.nn.Module):
+
+    def __init__(self, max_spec_t: int, pad_mode: str = 'constant', pad_value: float = 0.0):
+        super().__init__()
+        self.max_spec_t = max_spec_t
+        self.pad_mode = pad_mode
+        self.pad_value = pad_value
+
+    def forward(self, item):
+        item['audio'] = self.pad_or_truncate(item['audio'])
+        return item
+
+    def pad_or_truncate(self, audio):
+        difference = self.max_spec_t - audio.shape[-1]  # safe for batched input
+        # pad or truncate, depending on difference
+        if difference > 0:
+            # pad the last dim (time) -> (..., n_mels, 0+time+difference)  # safe for batched input
+            pad_dims = (0, difference)
+            audio = torch.nn.functional.pad(audio, pad_dims, self.pad_mode, self.pad_value)
+        elif difference < 0:
+            logging.warning(f'Truncating spec ({audio.shape}) to max_spec_t ({self.max_spec_t}).')
+            audio = audio[..., :self.max_spec_t]  # safe for batched input
+        return audio
+
+
+class AudioNormalizeAST(torch.nn.Module):
+    '''Normalization is done with two specified mean and std (half)'''
+    def __init__(self, mean: float, std: float) -> None:
+        super().__init__()
+        self.mean = mean
+        self.std = std
+
+    def forward(self, item):
+        item['audio'] = (item['audio'] - self.mean) / (2 * self.std)
+        item['meta']['audio']['norm_stats'] = {'mean': self.mean, 'std': self.std}
+        return item
+
+
+class PermuteStreams(torch.nn.Module):
+
+    def __init__(self, einops_order_audio: str, einops_order_rgb: str) -> None:
+        ''' For example:
+                einops_order_audio: "S F T -> S T F"
+                einops_order_rgb: "S T C H W -> S C T H W"'''
+        super().__init__()
+        self.einops_order_audio = einops_order_audio
+        self.einops_order_rgb = einops_order_rgb
+
+    def forward(self, item):
+        if self.einops_order_audio is not None:
+            item['audio'] = einops.rearrange(item['audio'], self.einops_order_audio).contiguous()
+        if self.einops_order_rgb is not None:
+            item['video'] = einops.rearrange(item['video'], self.einops_order_rgb).contiguous()
+        return item
+
+
+class ResampleAudio(torch.nn.Module):
+
+    def __init__(self, new_fps: int):
+        super().__init__()
+        self.new_fps = new_fps
+
+    def forward(self, item):
+        orig_fps = int(item['meta']['audio']['framerate'][0])
+        item['meta']['audio']['orig_shape'] = item['audio'].shape
+        if orig_fps != self.new_fps:
+            item['audio'] = torchaudio.functional.resample(item['audio'], orig_fps, self.new_fps)
+            item['meta']['audio']['framerate'][0] = self.new_fps
+        return item
+
+class ResampleRGB(torch.nn.Module):
+
+    def __init__(self, new_fps: int) -> None:
+        super().__init__()
+        self.new_fps = new_fps
+
+    def forward(self, item):
+        orig_fps = float(item['meta']['video']['fps'][0])
+        item['meta']['video']['orig_shape'] = item['video'].shape
+        if orig_fps != self.new_fps:
+            duration_sec = item['video'].shape[0] / orig_fps
+            indices = torch.arange(0, orig_fps * duration_sec - 1e-9, orig_fps / self.new_fps)
+            # basically, rounding
+            indices = indices.to(dtype=torch.long)
+            item['video'] = item['video'][indices]
+            item['meta']['video']['fps'][0] = self.new_fps
+        return item
+
+class ResizeAndLetterboxPad(torch.nn.Module):
+    '''Adapted from WACV24 Amazon`s challenge'''
+
+    def __init__(self, new_h, new_w):
+        super().__init__()
+        self.new_h = new_h
+        self.new_w = new_w
+        self.aspect_ratio = new_w / new_h
+
+    def forward(self, item):
+        item['video'] = self.resize_and_pad(item['video'])
+        return item
+
+    def resize_and_pad(self, rgb: torch.Tensor):
+        _, _, height, width = rgb.shape
+        current_aspect_ratio = width / height
+        if current_aspect_ratio > self.aspect_ratio:
+            scaled_height = round(self.new_w / current_aspect_ratio)
+            rgb = torchvision.transforms.functional.resize(rgb, (scaled_height, self.new_w), antialias=None)
+            top = (self.new_h - scaled_height) // 2
+            bottom = self.new_h - (scaled_height + top)
+            rgb = torch.nn.ConstantPad2d((0, 0, top, bottom), 0)(rgb)
+        elif current_aspect_ratio < self.aspect_ratio:
+            scaled_width = round(self.new_h*current_aspect_ratio)
+            rgb = torchvision.transforms.functional.resize(rgb, (self.new_h, scaled_width), antialias=None)
+            left = (self.new_w - scaled_width) // 2
+            right = self.new_w - (scaled_width + left)
+            rgb = torch.nn.ConstantPad2d((left, right, 0, 0), 0)(rgb)
+        return rgb
+
+
+class ResampleResizeLetterboxPad(torch.nn.Module):
+
+    def __init__(self, afps, vfps, new_h, new_w) -> None:
+        super().__init__()
+        self.transforms = torchvision.transforms.Compose([
+            ResampleAudio(new_fps=afps),
+            ResampleRGB(new_fps=vfps),
+            ResizeAndLetterboxPad(new_h=new_h, new_w=new_w)
+        ])
+
+    def forward(self, x: dict) -> dict:
+        return self.transforms(x)
+
+class DoNothing(torch.nn.Module):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__()
+
+    def forward(self, x: dict) -> dict:
+        return x
+
+
+if __name__ == '__main__':
+    grid = make_class_grid(-1, 1, 21)
+    grid = make_class_grid(-2, 2, 41)
+    print('grid:', grid)
+    print('value quantization:', quantize_offset(grid, 0.06))
+    v_fps = 25.0
+    duration = 10.0
+
+    input = {
+        'video': torch.randint(0, 256, (int(duration * v_fps), 3, 720//2, 1280//2), dtype=torch.uint8),
+        'audio': torch.arange(221184-1).float(),
+        'targets': {},
+        'meta': {
+            'video': {'duration': [duration], 'fps': [v_fps]},
+            'audio': {'duration': [duration], 'framerate': [22050.0]},
+            'subtitles': {'duration': []},
+            'cc': {'duration': []},
+        },
+        'path': '/home/nvme/data/vggsound/video/-5cWCaoEDlE_261000_271000.mp4',
+        'split': 'train',
+    }
+
+    print(input['audio'].shape, input['video'].shape)
+
+    fn = EqualifyFromRight(clip_max_len_sec=10)
+    input = fn(input)
+    print(input['audio'].shape, input['video'].shape)
+
+    fn = RGBSpatialCrop((224, 224), is_random=True)
+    # fn = RGBSpatialCrop((112, 112), is_random=True)
+    input = fn(input)
+    print(input['audio'].shape, input['video'].shape, input['meta']['audio'])
+
+    fn = Resize((224, 224))
+    input = fn(input)
+    print(input['audio'].shape, input['video'].shape, input['meta']['audio'])
+
+    fn = GenerateMultipleSegments(segment_size_vframes=16, n_segments=14,
+                                  is_start_random=False, audio_jitter_sec=0.05, step_size_seg=0.5)
+    input = fn(input)
+    print(input['audio'].shape, input['video'].shape, input['meta']['audio'])
+
+    fn = RandomApplyColorDistortion(p_gray_scale=0.5, p_color_jitter=0.5, s=1.0)
+    input = fn(input)
+    print(input['audio'].shape, input['video'].shape, input['meta']['audio'])
+
+    fn = RGBToFloatToZeroOne()
+    input = fn(input)
+    print(input['audio'].shape, input['video'].shape, input['meta']['audio'])
+    print(input['meta'])
+
+    fn = RGBNormalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    input = fn(input)
+    print(input['audio'].shape, input['video'].shape, input['meta']['audio'])
+    print(input['video'].mean(dim=(0, 2, 3)))
+    print(input['meta'])
+
+    fn = AudioRandomReverb(p=1.0)
+    input = fn(input)
+
+    fn = AudioRandomVolume(p=1.0, gain=2.0, gain_type='amplitude')
+    input = fn(input)
+    print(input['audio'].shape, input['video'].shape, input['meta']['audio'])
+
+    fn = AudioRandomPitchShift(p=1.0, shift=1000)
+    input = fn(input)
+    print(input['audio'].shape, input['video'].shape, input['meta']['audio'])
+
+    fn = AudioRandomLowpassFilter(p=1.0, cutoff_freq=100)
+    input = fn(input)
+    print(input['audio'].shape, input['video'].shape, input['meta']['audio'])
+
+    fn = AudioRandomGaussNoise(p=1.0, amplitude=0.01)
+    input = fn(input)
+    print(input['audio'].shape, input['video'].shape, input['meta']['audio'])
+
+    fn = AudioLog()
+    input = fn(input)
+    print(input['audio'].shape, input['video'].shape, input['meta']['audio'])
+
+    # audio only
+    input = {
+        'audio': torch.arange(221184).float(),
+        'meta': {
+            'video': {'duration': [10.0], 'fps': [10.0]},
+            'audio': {'duration': [11.0], 'framerate': [22050.0]},
+            'subtitles': {'duration': []},
+            'cc': {'duration': []}
+        },
+        'path': '/home/nvme/data/vggsound/video/-5cWCaoEDlE_261000_271000.mp4'
+    }
+
+    print(input['audio'].shape)
+
+    fn = AudioLog()
+    input = fn(input)
+    print(input['audio'].shape, input['meta']['audio'])
+    print(input['meta'])
+    print(input['audio'].min(), input['audio'].max())

modules/dataset/vggsound.py

@@ -0,0 +1,394 @@
+import csv
+import logging
+import os
+import random
+import sys
+from collections import Counter
+from glob import glob
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, '.')  # nopep8
+from dataset.dataset_utils import get_fixed_offsets, get_video_and_audio, subsample_dataset
+
+
+class VGGSound(torch.utils.data.Dataset):
+
+    def __init__(self,
+                 split,
+                 vids_dir,
+                 transforms=None,
+                 to_filter_bad_examples=True,
+                 splits_path='./data',
+                 meta_path='./data/vggsound.csv',
+                 seed=1337,
+                 load_fixed_offsets_on=['valid', 'test'],
+                 vis_load_backend='read_video',
+                 size_ratio=None,
+                 attr_annot_path=None,
+                 max_attr_per_vid=None):
+        super().__init__()
+        self.max_clip_len_sec = None
+        self.split = split
+        self.vids_dir = vids_dir
+        self.transforms = transforms
+        self.to_filter_bad_examples = to_filter_bad_examples
+        self.splits_path = splits_path
+        self.meta_path = meta_path
+        self.seed = seed
+        self.load_fixed_offsets_on = [] if load_fixed_offsets_on is None else load_fixed_offsets_on
+        self.vis_load_backend = vis_load_backend
+        self.size_ratio = size_ratio
+
+        vggsound_meta = list(csv.reader(open(meta_path), quotechar='"'))
+
+        # filter "bad" examples
+        if to_filter_bad_examples:
+            vggsound_meta = self.filter_bad_examples(vggsound_meta)
+
+        unique_classes = sorted(list(set(row[2] for row in vggsound_meta)))
+        self.label2target = {label: target for target, label in enumerate(unique_classes)}
+        self.target2label = {target: label for label, target in self.label2target.items()}
+        self.video2target = {row[0]: self.label2target[row[2]] for row in vggsound_meta}
+
+        split_clip_ids_path = os.path.join(splits_path, f'vggsound_{split}.txt')
+        if not os.path.exists(split_clip_ids_path):
+            self.make_split_files()
+        # the ugly string converts ['AdfsGsfII2yQ', '1'] into `AdfsGsfII2yQ_1000_11000`
+        meta_available = set([f'{r[0]}_{int(r[1])*1000}_{(int(r[1])+10)*1000}' for r in vggsound_meta])
+        within_split = set(open(split_clip_ids_path).read().splitlines())
+        clip_paths = [os.path.join(vids_dir, v + '.mp4') for v in meta_available.intersection(within_split)]
+        clip_paths = sorted(clip_paths)
+
+        if split in self.load_fixed_offsets_on:
+            logging.info(f'Using fixed offset for {split}')
+            self.vid2offset_params = get_fixed_offsets(transforms, split, splits_path, 'vggsound')
+
+        # making sure that all classes have at least one example
+        counter = Counter([self.video2target[Path(p).stem[:11]] for p in clip_paths])
+        assert all(counter[c] > 0 for c in self.target2label.keys()), \
+            f'Some classes have 0 count: {dict(counter)}'
+
+        self.dataset = clip_paths
+        self.dataset = subsample_dataset(self.dataset, size_ratio, shuffle=split == 'train')
+
+    def filter_bad_examples(self, vggsound_meta):
+        bad = set()
+        base_path = Path('./data/filtered_examples_vggsound')
+        lists = [open(p).read().splitlines() for p in sorted(glob(str(base_path / '*.txt')))]
+        for s in lists:
+            bad = bad.union(s)
+        # the ugly string converts '---g-f_I2yQ', '1' into `---g-f_I2yQ_1000_11000`
+        vggsound_meta = [r for r in vggsound_meta if f'{r[0]}_{int(r[1])*1000}_{(int(r[1])+10)*1000}' not in bad]
+        return vggsound_meta
+
+    def __getitem__(self, index):
+        path = self.dataset[index]
+        rgb, audio, meta = self.load_media(path)
+        item = self.make_datapoint(path, rgb, audio, meta)
+        if self.transforms is not None:
+            item = self.transforms(item)
+        return item
+
+    def make_datapoint(self, path, rgb, audio, meta):
+        # (Tv, 3, H, W) in [0, 225], (Ta, C) in [-1, 1]
+        target = self.video2target[Path(path).stem[:11]]
+        item = {
+            'video': rgb,
+            'audio': audio,
+            'meta': meta,
+            'path': str(path),
+            'targets': {'vggsound_target': target, 'vggsound_label': self.target2label[target]},
+            'split': self.split,
+        }
+
+        if self.split in self.load_fixed_offsets_on:
+            unique_id = Path(path).stem
+            offset_params = self.vid2offset_params[unique_id]
+            item['targets']['offset_sec'] = self.vid2offset_params[unique_id]['offset_sec']
+            item['targets']['v_start_i_sec'] = self.vid2offset_params[unique_id]['v_start_i_sec']
+            if 'oos_target' in offset_params:
+                item['targets']['offset_target'] = {
+                    'oos': offset_params['oos_target'], 'offset': item['targets']['offset_sec'],
+                }
+
+        return item
+
+    def load_media(self, path):
+        rgb, audio, meta = get_video_and_audio(path, get_meta=True, end_sec=self.max_clip_len_sec)
+        return rgb, audio, meta
+
+    def make_split_files(self):
+        if self.to_filter_bad_examples:
+            logging.warning('`to_filter_bad_examples` is True. `make_split_files` expects otherwise')
+
+        logging.info(f'The split files do not exist @ {self.splits_path}. Calculating the new ones.')
+        # The downloaded videos (some went missing on YouTube and no longer available)
+        available_vid_paths = sorted(glob(os.path.join(self.vids_dir, '*.mp4')))
+        logging.info(f'The number of clips available after download: {len(available_vid_paths)}')
+
+        # original (full) train and test sets
+        vggsound_meta = list(csv.reader(open(self.meta_path), quotechar='"'))
+        train_vids = {row[0] for row in vggsound_meta if row[3] == 'train'}
+        test_vids = {row[0] for row in vggsound_meta if row[3] == 'test'}
+
+        # # the cleaned test set
+        # vggsound_meta_test_v2 = list(csv.reader(open(self.meta_path_clean_test), quotechar='"'))
+        logging.info(f'The number of videos in vggsound train set: {len(train_vids)}')
+        logging.info(f'The number of videos in vggsound test set: {len(test_vids)}')
+
+        # class counts in test set. We would like to have the same distribution in valid
+        unique_classes = sorted(list(set(row[2] for row in vggsound_meta)))
+        label2target = {label: target for target, label in enumerate(unique_classes)}
+        video2target = {row[0]: label2target[row[2]] for row in vggsound_meta}
+        test_vid_classes = [video2target[vid] for vid in test_vids]
+        test_target2count = Counter(test_vid_classes)
+
+        # now given the counts from test set, sample the same count for validation and the rest leave in train
+        train_vids_wo_valid, valid_vids = set(), set()
+        for target, label in enumerate(label2target.keys()):
+            class_train_vids = [vid for vid in sorted(list(train_vids)) if video2target[vid] == target]
+            random.Random(self.seed).shuffle(class_train_vids)
+            count = test_target2count[target]
+            valid_vids.update(class_train_vids[:count])
+            train_vids_wo_valid.update(class_train_vids[count:])
+
+        # make file with a list of available test videos (each video should contain timestamps as well)
+        train_i = valid_i = test_i = 0
+        with open(os.path.join(self.splits_path, 'vggsound_train.txt'), 'w') as train_file, \
+             open(os.path.join(self.splits_path, 'vggsound_valid.txt'), 'w') as valid_file, \
+             open(os.path.join(self.splits_path, 'vggsound_test.txt'), 'w') as test_file:
+            #  open(os.path.join(self.splits_path, 'vggsound_test_v2.txt'), 'w') as test_file_v2:
+            for path in available_vid_paths:
+                path = path.replace('.mp4', '')
+                vid_name = Path(path).name
+                # 'zyTX_1BXKDE_16000_26000'[:11] -> 'zyTX_1BXKDE'
+                if vid_name[:11] in train_vids_wo_valid:
+                    train_file.write(vid_name + '\n')
+                    train_i += 1
+                elif vid_name[:11] in valid_vids:
+                    valid_file.write(vid_name + '\n')
+                    valid_i += 1
+                elif vid_name[:11] in test_vids:
+                    test_file.write(vid_name + '\n')
+                    test_i += 1
+                # else:
+                #     raise Exception(f'Clip {vid_name} is neither in train, valid nor test. Strange.')
+
+        logging.info(f'Put {train_i} clips to the train set and saved it to ./data/vggsound_train.txt')
+        logging.info(f'Put {valid_i} clips to the valid set and saved it to ./data/vggsound_valid.txt')
+        logging.info(f'Put {test_i} clips to the test set and saved it to ./data/vggsound_test.txt')
+
+    def __len__(self):
+        return len(self.dataset)
+
+class VGGSoundSparse(VGGSound):
+    '''
+        The same as VGGSound, except the list of videos is filtered for sparse sounds (sparse_meta_path)
+    '''
+
+    def __init__(self, split, vids_dir, transforms=None, to_filter_bad_examples=True,
+                 splits_path='./data', meta_path='./data/vggsound.csv',
+                 sparse_meta_path='./data/sparse_classes.csv', seed=1337, load_fixed_offsets_on=['valid', 'test'],
+                 vis_load_backend='read_video', size_ratio=None, attr_annot_path=None, max_attr_per_vid=None):
+        super().__init__(split, vids_dir, transforms, to_filter_bad_examples, splits_path, meta_path, seed,
+                         load_fixed_offsets_on, vis_load_backend, size_ratio)
+        self.sparse_meta_path = sparse_meta_path
+        sparse_meta = list(csv.reader(open(sparse_meta_path), quotechar='"', delimiter='\t'))
+        sparse_classes = set([row[0] for row in sparse_meta if row[1] == 'y'])
+        label2new_target = {label: target for target, label in enumerate(sorted(list(sparse_classes)))}
+        new_target2label = {target: label for label, target in label2new_target.items()}
+
+        sparse_dataset = []
+        video2new_target = {}
+        for path in self.dataset:
+            vid_id = Path(path).stem[:11]
+            vid_target = self.video2target[vid_id]
+            vid_label = self.target2label[vid_target]
+            if vid_label in sparse_classes:
+                sparse_dataset.append(path)
+                video2new_target[vid_id] = label2new_target[vid_label]
+
+        self.dataset = sparse_dataset
+        logging.debug(f'Filtered VGGSound dataset to sparse classes from {Path(sparse_meta_path).name}.')
+
+        # redefining the label <-> target variable
+        self.label2target = label2new_target
+        self.target2label = new_target2label
+        self.video2target = video2new_target
+        logging.debug('Redefined label <-> target mapping to sparse classes.')
+
+        counter = Counter([self.video2target[Path(p).stem[:11]] for p in self.dataset])
+        assert len(self.dataset) < 1000 or all(counter[c] > 0 for c in self.target2label.keys()), \
+            f'Some classes have 0 count: {dict(counter)}'
+
+
+class VGGSoundSparsePicked(VGGSoundSparse):
+
+    def __init__(self, split, vids_dir, transforms=None, to_filter_bad_examples=True,
+                 splits_path='./data', meta_path='./data/vggsound.csv',
+                 sparse_meta_path='./data/picked_sparse_classes.csv', seed=1337,
+                 load_fixed_offsets_on=['valid', 'test'], vis_load_backend='read_video', size_ratio=None,
+                 attr_annot_path=None, max_attr_per_vid=None):
+        super().__init__(split, vids_dir, transforms, to_filter_bad_examples, splits_path,
+                         meta_path, sparse_meta_path, seed, load_fixed_offsets_on, vis_load_backend,
+                         size_ratio)
+
+
+class VGGSoundSparsePickedCleanTest(VGGSoundSparse):
+
+    def __init__(self, split, vids_dir, transforms=None, to_filter_bad_examples=True,
+                 splits_path='./data', meta_path='./data/vggsound.csv',
+                 sparse_meta_path='./data/picked_sparse_classes.csv', seed=1337,
+                 load_fixed_offsets_on=['valid', 'test'], vis_load_backend='read_video', size_ratio=None,
+                 attr_annot_path=None, max_attr_per_vid=None):
+        super().__init__(split, vids_dir, transforms, to_filter_bad_examples, splits_path,
+                         meta_path, sparse_meta_path, seed, load_fixed_offsets_on, vis_load_backend,
+                         size_ratio)
+
+    def filter_bad_examples(self, vggsound_meta):
+        bad = set()
+        base_path1 = Path('./data/filtered_examples_vggsound')
+        base_path2 = Path('./data/filtered_examples_vggsound_extra')
+        lists1 = [open(p).read().splitlines() for p in sorted(glob(str(base_path1 / '*.txt')))]
+        lists2 = [open(p).read().splitlines() for p in sorted(glob(str(base_path2 / '*.txt')))]
+        lists = lists1 + lists2
+        for s in lists:
+            bad = bad.union(s)
+        # the ugly string converts '---g-f_I2yQ', '1' into `---g-f_I2yQ_1000_11000`
+        vggsound_meta = [r for r in vggsound_meta if f'{r[0]}_{int(r[1])*1000}_{(int(r[1])+10)*1000}' not in bad]
+        return vggsound_meta
+
+
+class VGGSoundSparsePickedCleanTestFixedOffsets(VGGSoundSparse):
+    '''This dataset only operates on manually annotated fixed offsets. Meant for evaluation purpose.'''
+
+    def __init__(self, split, vids_dir, transforms=None, to_filter_bad_examples=True,
+                 splits_path='./data', meta_path='./data/vggsound.csv',
+                 sparse_meta_path='./data/picked_sparse_classes.csv', seed=1337,
+                 load_fixed_offsets_on=['valid', 'test'], vis_load_backend='read_video', size_ratio=None,
+                 attr_annot_path=None, max_attr_per_vid=None):
+        super().__init__(split, vids_dir, transforms, to_filter_bad_examples, splits_path,
+                         meta_path, sparse_meta_path, seed, load_fixed_offsets_on, vis_load_backend,
+                         size_ratio)
+        # redefine the dataset to only use fixed offsets
+        fix_off_path = './data/vggsound_sparse_clean_fixed_offsets.csv'
+        self.vid2offset_params = {}
+        # lines have the format: dataset_name,video_id,vstart_sec,offset_sec,is_sync
+        reader = csv.reader(open(fix_off_path))
+        next(reader)  # skipping the header
+        for _, vi, st, of, sy in reader:
+            # FIXME: if there are more than one fixed offset per video, we will need to change the logic here
+            assert vi not in self.vid2offset_params, 'offsets from other splits will override each other'
+            if sy == '1':
+                self.vid2offset_params[vi] = {'offset_sec': float(of), 'v_start_i_sec': float(st)}
+        logging.debug(f'Loaded {len(self.vid2offset_params)} fixed offsets from {Path(fix_off_path).name}.')
+        # since we care only about the videos in the fixed offset split, we filter the dataset
+        self.dataset = [p for p in self.dataset if Path(p).stem in self.vid2offset_params]
+        logging.debug(f'Filtered VGGSoundSparse dataset to fixed offsets from {Path(fix_off_path).name}.')
+
+
+class LongerVGGSound(VGGSound):
+    '''This class is different to the parent in the extra filtering it does. If the parent was
+    making the splits with filtering for shorter than 9 second, this class filters for shorter than 9.5 sec.
+    by applying extra filtering.
+    '''
+
+    def __init__(self,
+                 split,
+                 vids_dir,
+                 transforms=None,
+                 to_filter_bad_examples=True,
+                 splits_path='./data',
+                 meta_path='./data/vggsound.csv',
+                 seed=1337,
+                 load_fixed_offsets_on=['valid', 'test'],
+                 vis_load_backend='read_video',
+                 size_ratio=None,
+                 attr_annot_path=None,
+                 max_attr_per_vid=None):
+        # size_ratio=None and fixed_offsets_on=[] to avoid double subsampling and loading fixed offsets
+        super().__init__(split, vids_dir, transforms, to_filter_bad_examples, splits_path, meta_path, seed,
+                         [], vis_load_backend, None, attr_annot_path, max_attr_per_vid)
+        # redefining the load_fixed_offsets_on because the parent class does not load them (see above)
+        self.load_fixed_offsets_on = load_fixed_offsets_on
+        # doing the extra filtering for longer than 9.5 sec
+        if to_filter_bad_examples:
+            filt_ex_path = Path('./data/filtered_examples_vggsound_shorter/less_than_9.5s.txt')
+            bad = set([p for p in open(filt_ex_path).read().splitlines()])
+            self.dataset = [p for p in self.dataset if Path(p).stem not in bad]
+
+        # longer vid clips require new fixed offsets, loading them here again (VGGSound loads them in init)
+        if split in self.load_fixed_offsets_on:
+            logging.info(f'Using fixed offset for {split}')
+            self.vid2offset_params = get_fixed_offsets(transforms, split, splits_path, 'vggsound')
+
+        self.dataset = subsample_dataset(self.dataset, size_ratio, shuffle=split == 'train')
+        logging.info(f'{split} has {len(self.dataset)} items')
+
+
+
+if __name__ == '__main__':
+    from omegaconf import OmegaConf
+    from scripts.train_utils import get_transforms
+    from utils.utils import cfg_sanity_check_and_patch
+    cfg = OmegaConf.load('./configs/transformer.yaml')
+    vis_load_backend = 'read_video'
+
+    # set logging for info level
+    logging.basicConfig(level=logging.INFO)
+
+    transforms = get_transforms(cfg)
+
+    vids_path = 'PLACEHOLDER'
+
+    # cfg.data.dataset.params.size_ratio = 0.1
+
+    cfg_sanity_check_and_patch(cfg)
+
+    datasets = {
+        'train': VGGSound('train', vids_path, transforms['train'], vis_load_backend=vis_load_backend,
+                          size_ratio=cfg.data.dataset.params.size_ratio),
+        'valid': VGGSound('valid', vids_path, transforms['test'], vis_load_backend=vis_load_backend),
+        'test': VGGSound('test', vids_path, transforms['test'], vis_load_backend=vis_load_backend),
+    }
+    for phase in ['train', 'valid', 'test']:
+        print(phase, len(datasets[phase]))
+
+    print(datasets['train'][0]['audio'].shape, datasets['train'][0]['video'].shape)
+    print(datasets['train'][0]['meta'])
+    print(datasets['valid'][0]['audio'].shape, datasets['valid'][0]['video'].shape)
+    print(datasets['valid'][0]['meta'])
+    print(datasets['test'][0]['audio'].shape, datasets['test'][0]['video'].shape)
+    print(datasets['test'][0]['meta'])
+
+    # for i in range(300, 1000):
+    #     datasets['train'][i]['path']
+    #     print(datasets['train'][0]['audio'].shape, datasets['train'][0]['video'].shape)
+    #     print(datasets['train'][0]['meta'])
+
+    datasets = {
+        'train': VGGSoundSparse('train', vids_path, transforms['train']),
+        'valid': VGGSoundSparse('valid', vids_path, transforms['test']),
+        'test': VGGSoundSparse('test', vids_path, transforms['test']),
+    }
+    for phase in ['train', 'valid', 'test']:
+        print(phase, len(datasets[phase]))
+
+    print(datasets['train'][0]['audio'].shape, datasets['train'][0]['video'].shape)
+    print(datasets['train'][0]['meta'])
+    print(datasets['valid'][0]['audio'].shape, datasets['valid'][0]['video'].shape)
+    print(datasets['valid'][0]['meta'])
+    print(datasets['test'][0]['audio'].shape, datasets['test'][0]['video'].shape)
+    print(datasets['test'][0]['meta'])
+
+    datasets = {
+        'test': VGGSoundSparsePickedCleanTestFixedOffsets('test', vids_path, transforms['test']),
+    }
+    for phase in ['test']:
+        print(phase, len(datasets[phase]))
+
+    for i in range(len(datasets['test'])):
+        datasets['test'][i]['path']
+        print(datasets['test'][i]['audio'].shape, datasets['test'][i]['video'].shape)

modules/model/__init__.py

@@ -0,0 +1 @@
+# Auto-generated __init__.py

modules/model/modules/bridges.py

@@ -0,0 +1,178 @@
+import torch
+from torch.nn import functional as F
+
+
+class BridgeBase(torch.nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.bridge = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        try:
+            x = self.bridge(x)
+            return x
+        except TypeError as e:
+            raise TypeError('The class cant be called on its own. Please, use a class that inherits it', e)
+
+
+class ConvBridgeBase(BridgeBase):
+
+    def __init__(self, block, **kwargs) -> None:
+        super().__init__()
+        self.bridge = torch.nn.Sequential(
+            block(**kwargs),
+            torch.nn.GELU(),
+        )
+
+
+class AvgPoolBridgeBase(BridgeBase):
+
+    def __init__(self, block, **kwargs) -> None:
+        super().__init__()
+        # a list [1, 2, 3] specified in omegaconfig, has type ListConf which is not accepted by pytorch
+        # interestingly, conv layers don't care
+        for k in ['kernel_size', 'stride']:
+            kwargs[k] = list(kwargs[k])
+        self.bridge = block(**kwargs)
+
+
+class ConvBridgeAudio(ConvBridgeBase):
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(block=torch.nn.Conv2d, **kwargs)
+
+
+class ConvBridgeVisual(ConvBridgeBase):
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(block=torch.nn.Conv3d, **kwargs)
+
+
+class AvgPoolBridgeVisual(AvgPoolBridgeBase):
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(block=torch.nn.AvgPool3d, **kwargs)
+
+
+class AvgPoolBridgeAudio(AvgPoolBridgeBase):
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(block=torch.nn.AvgPool2d, **kwargs)
+
+
+class DoNothingBridge(BridgeBase):
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.bridge = torch.nn.Identity(**kwargs)
+
+
+class AppendZerosToHidden(BridgeBase):
+
+    def __init__(self, target_hidden_size, dim) -> None:
+        super().__init__()
+        self.target_hidden_size = target_hidden_size
+        self.dim = dim
+
+    def forward(self, x):
+        d_res = self.target_hidden_size - x.shape[self.dim]
+        # going to insert the new dimension into the x.shape output
+        shape_target = list(x.shape[:self.dim]) + [d_res] + list(x.shape[self.dim+1:])
+        # creating the zeros to append to x
+        zeros = torch.zeros(shape_target).to(x)
+        x = torch.cat([x, zeros], self.dim)
+        return x
+
+
+class SpatialpoolConvTemporalpool(torch.nn.Module):
+    '''Similar to S3D but with slightly different kernel for F.avg_pool3d.
+    To be used in AVCLIP in visual branch'''
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.conv = torch.nn.Conv3d(**kwargs)
+
+    def forward(self, x: torch.Tensor):
+        B, t, d, h, w = x.shape
+        x = x.permute(0, 2, 1, 3, 4)  # (B, d, t, h, w)
+        # pool as in S3D but without temporal pooling (2-->1, h, w)
+        x = F.avg_pool3d(x, (1, h, w), stride=1)  # (B, d, t, 1, 1)
+        x = self.conv(x)  # (B, D, t, 1, 1)
+        x = x.view(B, self.conv.out_channels, t)  # squeeze the spatial dimensions
+        x = x.mean(dim=-1)  # temporal pooling
+        return x  # (B, d)
+
+
+class FrequencypoolConvTemporalpool(torch.nn.Module):
+    '''Similar to the visual branch of S3D, which is a stack of spatial pool, conv, temporal pool blocks.
+    Instead, this is a stack of frequency pool, conv, temporal pooling blocks.
+    To be used in AVCLIP in audio branch'''
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.conv = torch.nn.Conv2d(**kwargs)
+
+    def forward(self, x: torch.Tensor):
+        B, d, f, t = x.shape
+        # frequency pooling (f-->1)
+        x = F.avg_pool2d(x, (f, 1), stride=1)  # (B, d, 1, t)
+        x = self.conv(x)  # (B, D, 1, t)
+        x = x.view(B, self.conv.out_channels, t)  # squeeze the frequency dimension
+        x = x.mean(dim=-1)  # temporal pooling
+        return x  # (B, d)
+
+
+if __name__ == '__main__':
+    v = torch.rand(2, 50, 512, 7, 7)
+    a = torch.rand(2, 512, 9, 27)
+
+    in_channels = 512
+    out_channels = 512
+    kernel_size_v = [1, 7, 7]
+    kernel_size_a = [9, 1]
+    stride = 1
+    bias = True
+
+    conv_bridge_a = ConvBridgeAudio(in_channels=in_channels, out_channels=out_channels,
+                                    kernel_size=kernel_size_a, stride=stride, bias=bias)
+    conv_bridge_v = ConvBridgeVisual(in_channels=in_channels, out_channels=out_channels,
+                                     kernel_size=kernel_size_v, stride=stride, bias=bias)
+    avg_bridge_a = AvgPoolBridgeAudio(kernel_size=kernel_size_a)
+    avg_bridge_v = AvgPoolBridgeVisual(kernel_size=kernel_size_v)
+    i_bridge_a = DoNothingBridge(some_arg=123)
+    i_bridge_v = DoNothingBridge(some_arg=123)
+    h_bridge_a = AppendZerosToHidden(target_hidden_size=1024, dim=1)
+    h_bridge_v = AppendZerosToHidden(target_hidden_size=1024, dim=1)
+
+    print('v', v.shape)
+    print('conv_v(v)', conv_bridge_v(v.permute(0, 2, 1, 3, 4)).permute(0, 2, 1, 3, 4).shape)
+    print()
+
+    print('a', a.shape)
+    print('conv_a(a)', conv_bridge_a(a).shape)
+    print()
+
+    print('v', v.shape)
+    print('avg3d(v)', avg_bridge_v(v.permute(0, 2, 1, 3, 4)).permute(0, 2, 1, 3, 4).shape)
+    print()
+
+    print('a', a.shape)
+    print('avg2d(a)', avg_bridge_a(a).shape)
+    print()
+
+    print('v', v.shape)
+    print('i(v)', i_bridge_v(v.permute(0, 2, 1, 3, 4)).permute(0, 2, 1, 3, 4).shape)
+    print()
+
+    print('a', a.shape)
+    print('i(a)', i_bridge_a(a).shape)
+    print()
+
+    print('v', v.shape)
+    print('h(v)', h_bridge_v(v.permute(0, 2, 1, 3, 4)).permute(0, 2, 1, 3, 4).shape)
+    print()
+
+    print('a', a.shape)
+    print('h(a)', h_bridge_a(a).shape)
+    print()

modules/model/modules/feat_extractors/audio/ast.py

@@ -0,0 +1,279 @@
+import logging
+import torch
+from torch import nn
+# importing modified version of AST
+from model.modules.feat_extractors.audio.hf_src.modeling_ast import ASTForAudioClassification, ASTConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+
+from model.modules.feat_extractors.visual.motionformer import (AveragePooling, BaseEncoderLayer,
+                                                               TemporalTransformerEncoderLayer)
+from utils.utils import check_if_file_exists_else_download
+
+
+class AST(torch.nn.Module):
+    def __init__(self,
+                 extract_features: bool = False,
+                 ckpt_path: str = None,
+                 feat_type: str = None,
+                 max_spec_t: int = None,
+                 factorize_freq_time: bool = None,
+                 agg_freq_module: str = None,
+                 agg_time_module: str = None,
+                 add_global_repr: bool = True,
+                 agg_segments_module: str = None,
+                 max_segments: int = None,
+                 ) -> None:
+        '''
+            extract_features: if True, then the model will return the features instead of head's output
+            ckpt_path: is not a path to a ckpt file, but a name of a model from the HuggingFace model hub.
+            feat_type: if extract_features is True, this parameter specifies the type of features to return
+            max_spec_t: if specified, then the model (pos emb) will be patched to support this length of spec
+            factorize_freq_time: if True, then the model will use a factorized freq/time aggregation
+            agg_freq_module: if specified, then the model will use this module for freq aggregation
+            agg_time_module: if specified, then the model will use this module for time aggregation
+            add_global_repr: if True, adds a global representation to the features (aggregation on segments)
+            agg_segments_module: if specified, then the model will use this module for segments aggregation
+            max_segments: if specified, the initialization of PE in the global agg module will use this value.
+                          This should correspond to the max number of segments per video (if None, 16 is used)
+        '''
+        super().__init__()
+        self.extract_features = extract_features
+        self.ckpt_path = ckpt_path
+        self.max_spec_t = max_spec_t
+        self.max_segments = max_segments
+
+        # depending on whether the feat extractor was pre-trained contrastively or not, we need to
+        # load the state dict differently.
+
+        # if ckpt is specified, then load the model from the HuggingFace model hub, otherwise init a new model
+        if ckpt_path == 'MIT/ast-finetuned-audioset-10-10-0.4593':
+            revision = 'c1c0c66'  # fixing the revision for compatibility (V4.27.4)
+            self.config = ASTConfig.from_pretrained(ckpt_path, revision=revision)
+            full_model = ASTForAudioClassification.from_pretrained(ckpt_path, revision=revision)
+            logging.info(f'Loaded AST from {ckpt_path}')
+        else:
+            self.config = ASTConfig()
+            self.config.num_labels = 527  # 2 by default, audioset has 527 labels
+            full_model = ASTForAudioClassification(self.config)
+            logging.info('Initialized AST from scratch with the AST AudioSet config')
+
+        was_pt_on_avclip = ckpt_path is not None and ckpt_path.endswith('.pt')
+
+        # feature extractor
+        self.ast = full_model.audio_spectrogram_transformer
+
+        if self.extract_features:
+            # assign `feat_type` (use default if not specified)
+            self.feat_type = 'last_hidden_state' if feat_type is None else feat_type
+            # define adapters if needed
+            self.factorize_freq_time = factorize_freq_time
+            # avoiding code duplication (used only if agg_*_module is TransformerEncoderLayer)
+            transf_enc_layer_kwargs = dict(
+                d_model=self.config.hidden_size, nhead=self.config.num_attention_heads,
+                dim_feedforward=self.config.intermediate_size, activation=nn.GELU(), batch_first=True,
+                dropout=self.config.attention_probs_dropout_prob, layer_norm_eps=1e-6, norm_first=True,
+            )
+            if factorize_freq_time:
+                self.feat_type = 'last_hidden_state'  # this feat_type supports factorization
+                # frequency aggreration
+                if agg_freq_module == 'TransformerEncoderLayer':
+                    self.freq_attn_agg = FrequencyTransformerEncoderLayer(**transf_enc_layer_kwargs)
+                elif agg_freq_module == 'AveragePooling':
+                    self.freq_attn_agg = AveragePooling(avg_pattern='BS D f t -> BS D t',
+                                                        then_permute_pattern='BS D t -> BS t D')
+                # time aggreration
+                if agg_time_module == 'TransformerEncoderLayer':
+                    self.temp_attn_agg = TemporalTransformerEncoderLayer(**transf_enc_layer_kwargs)
+                elif agg_time_module == 'AveragePooling':
+                    self.temp_attn_agg = AveragePooling(avg_pattern='BS t D -> BS D')
+                elif 'Identity' in agg_time_module:
+                    self.temp_attn_agg = nn.Identity()
+            # define a global aggregation layer (aggregarate over segments)
+            self.add_global_repr = add_global_repr
+            if add_global_repr:
+                if agg_segments_module == 'TransformerEncoderLayer':
+                    # we can reuse the same layer as for temporal factorization (B, dim_to_agg, D) -> (B, D)
+                    # we need to add pos emb (PE) because previously we added the same PE for each segment
+                    pos_max_len = max_segments if max_segments is not None else 16  # 16 = 10sec//0.64sec + 1
+                    self.global_attn_agg = TemporalTransformerEncoderLayer(
+                        add_pos_emb=True, pos_emb_drop=self.config.hidden_dropout_prob,
+                        pos_max_len=pos_max_len, **transf_enc_layer_kwargs
+                    )
+                elif agg_segments_module == 'AveragePooling':
+                    self.global_attn_agg = AveragePooling(avg_pattern='B S D -> B D')
+        else:
+            self.classifier = full_model.classifier
+
+        # AST.device fails with AttributeError. This is a workaround
+        self.device = full_model.device
+
+        # pre-trained on 12*101+2=1214 tokens, but we have less (e.g. 12*6+2=74)
+        self.patch_position_emb()
+
+        if was_pt_on_avclip:
+            # we need to filter out the state_dict of the AVCLIP model (has both A and V extractors)
+            # and keep only the state_dict of the feat extractor
+            check_if_file_exists_else_download(self.ckpt_path)
+            ckpt = torch.load(ckpt_path, map_location='cpu')
+            ckpt_weights = dict()
+            for k, v in ckpt['state_dict'].items():
+                if k.startswith(('module.a_encoder.', 'a_encoder.')):
+                    k = k.replace('module.', '').replace('a_encoder.', '')
+                    ckpt_weights[k] = v
+            _load_status = self.load_state_dict(ckpt_weights, strict=False)
+            if len(_load_status.missing_keys) > 0 or len(_load_status.unexpected_keys) > 0:
+                logging.warning(f'Loading exact afeat_extractor ckpt from {self.ckpt_path} failed. \n' \
+                                f'Missing keys ({len(_load_status.missing_keys)}): ' \
+                                f'{_load_status.missing_keys}, \n' \
+                                f'Unexpected keys ({len(_load_status.unexpected_keys)}): ' \
+                                f'{_load_status.unexpected_keys} \n' \
+                                f'temp_attn_agg are expected to be missing if ckpt was pt contrastively.')
+            else:
+                logging.info(f'Loading afeat_extractor ckpt from {self.ckpt_path} succeeded.')
+
+        # print the number of parameters
+        logging.info(f'AST: {sum(p.numel() for p in self.parameters() if p.requires_grad):,}')
+
+    def forward(self, x: torch.Tensor, for_loop: bool = False, cont_mask: torch.Tensor = None,
+                **ast_kwargs) -> torch.Tensor:
+        '''
+            x: (B, S, T, F) where S is number of segments, F is number of (mel) frequency bins,
+            ast_kwargs: additional arguments for the AST model
+            cont_mask: (B, S, T, F) where 0s are the values to be masked out
+            if `for_loop=True`, we use a for loop to extract features for each segment separately.
+            if `for_loop=False`, we extract features for all segments at once.
+                Using the for loop is slower but more memory efficient, while using all segments at once
+                is faster but more memory inefficient.
+                Using for loop allows to control the memory footprint by varying the number of videos in a
+                batch (batch size) rather than the number of segments in a video.
+        '''
+        B, S, T, F = x.shape
+
+        if for_loop:
+            assert cont_mask is None, 'cont_mask is not supported with for_loop=True'
+            orig_shape_s = (B, 1, T, F)
+            # NOTE: since x is (B, S, T, F), and forward_segments expects (BS, T, F).
+            # (B, S, T, F)[:, s] is (B, T, F) or (BS, T, F) if S=1.
+            x = torch.cat(
+                [self.forward_segments(x[:, s], orig_shape_s, **ast_kwargs).unsqueeze(1) for s in range(S)],
+                dim=1)
+        else:
+            orig_shape = (B, S, T, F)
+            x = x.view(B * S, T, F)
+            if cont_mask is not None:
+                cont_mask = cont_mask.reshape(B * S, T, F)
+            # AST expects a tensor of shape (B*S, T, F).
+            x = self.forward_segments(x, orig_shape=orig_shape, cont_mask=cont_mask, **ast_kwargs)
+            # unpack the segments (using rest dimensions to support different shapes e.g. (BS, D) or (BS, t, D))
+            x = x.view(B, S, *x.shape[1:])
+        # x now is of shape (B, S, D) or (B, S, t, D) if `self.temp_attn_agg` is `Identity`
+
+        global_x = None
+        if self.extract_features and self.add_global_repr:  # lazy execution, throws AttributeError
+            assert len(x.shape) == 3, f'Local representation should be (B, S, D) {x.shape}'
+            global_x = self.global_attn_agg(x)  # (B, D)
+
+        return x, global_x  # x is (B, S, ...), global_x is (B, D) or None
+
+    def forward_segments(self, x, orig_shape: tuple, cont_mask: torch.Tensor = None, **ast_kwargs):
+        '''x is (BS, T, F), where S is the number of segments; cont_mask is (BS, T, F): 0s to be masked out'''
+        # 'pooler_output': (B, D); or 'last_hidden_state: (B, T, D) where T is [CLS, DISTILL, <tokens>]
+        # x_mask is (B, T) where 0s are the values to be masked out
+        x, x_mask = self.ast(x, cont_mask=cont_mask, **ast_kwargs)
+
+        if self.extract_features:
+            x = self.get_features_by_type(x)
+            if self.factorize_freq_time:
+                x = self.restore_freq_temp_dims(x, orig_shape)  # (BS, D, f, t) <- (B*S, T, D)
+                if cont_mask is not None:
+                    # duplicating the mask for the latent dimension (D) to be compatible with the next func
+                    x_mask = x_mask.unsqueeze(-1).expand(-1, -1, self.config.hidden_size)
+                    x_mask = self.restore_freq_temp_dims(x_mask, orig_shape)  # (BS, D, f, t) <- (B*S, T, D)
+                    # again removing the latent
+                    x_mask = x_mask[:, 0, :, :]
+                else:
+                    x_mask = None
+                x = self.freq_attn_agg(x, x_mask)  # (BS, t, D)
+                x = self.temp_attn_agg(x)  # (BS, D) or (BS, t, D) if self.temp_attn_agg is Identity
+        else:
+            x = x['pooler_output']
+            x = self.classifier(x)
+        return x
+
+    def get_features_by_type(self, x: BaseModelOutputWithPooling) -> torch.Tensor:
+        if self.feat_type == 'pooler_output':
+            return x['pooler_output']  # (B, D)
+        elif self.feat_type == 'CLS':
+            return x['last_hidden_state'][:, 0, :]  # (B, D)
+        elif self.feat_type == 'last_hidden_state':
+            return x['last_hidden_state']  # (B, 2+T, D)
+        elif self.feat_type == 'last_hidden_state_no_AUX':
+            return x['last_hidden_state'][:, 2:, :]  # (B, T, D) removing CLS and distill tokens
+        else:
+            raise ValueError(f'Unknown feature type: {self.feat_type}')
+
+    def restore_freq_temp_dims(self, feats, orig_shape: tuple):
+        '''
+            feats are of shape (B*S, T, D)
+                where T = 2 + f * t (if feat_type == 'last_hidden_state')
+                where T =     f * t (if feat_type == 'last_hidden_state_no_AUX')
+            Our goal is to make them of shape (B*S, f, t, D) where f and t are dimensions after patching.
+            From `self.ast.embeddings.patch_embeddings`, it follows that we could reshape feats:
+                `feats.transpose(1, 2).view(B*S, D, f, t)`
+
+            (Similar function is defined in for RGB features in `motionformer.py`)
+        '''
+        B, S, T, F = orig_shape
+        D = self.config.hidden_size
+
+        # num patches in each dimension
+        f, t = self.ast.embeddings.get_shape(self.config)
+
+        if self.feat_type == 'last_hidden_state':
+            feats = feats[:, 2:, :]  # removing CLS and distill tokens
+
+        feats = feats.permute(0, 2, 1)  # (B*S, D, T)
+        feats = feats.view(B * S, D, f, t)  # (B*S, D, f, t)
+
+        return feats
+
+    def patch_position_emb(self):
+        if self.max_spec_t is not None:
+            self.config.max_length = self.max_spec_t
+        f, t = self.ast.embeddings.get_shape(self.config)
+        shortened = self.ast.embeddings.position_embeddings[:, :f*t+2].clone()  # +2 for CLS and distill tokens
+        self.ast.embeddings.position_embeddings = torch.nn.Parameter(shortened).to(self.device)
+
+    def to(self, device):
+        '''AST.device fails with AttributeError. This is a workaround. '''
+        self.device = torch.device(device)
+        return super().to(device)
+
+
+class FrequencyTransformerEncoderLayer(BaseEncoderLayer):
+    ''' This layer is used to aggregate the features along the frequency axis.
+    It follows the same logic as spatio-temporal aggregation in visual feature extractor.
+    Thus, it is recommended to check the definition of `BaseEncoderLayer` in `motionformer.py` '''
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor = None) -> torch.Tensor:
+        ''' x: (B*S, D, f, t); if specified x_mask (B*S, f, t), 0s are the values to be masked out '''
+        BS, D, f, t = x.shape
+
+        # time as a batch dimension
+        x = x.permute(0, 3, 2, 1)  # (B*S, t, f, D)
+        x = x.reshape(BS * t, f, D)  # .view() fails with non-contiguous memory
+        # similar to mask
+        if x_mask is not None:
+            x_mask = x_mask.permute(0, 2, 1)  # (B*S, t, f)
+            x_mask = x_mask.reshape(BS * t, f)
+
+        # apply encoder layer (BaseEncoderLayer.forward) - it will add CLS token and output its representation
+        x = super().forward(x=x, x_mask=x_mask)  # (B*S*t, D)
+
+        # reshape back to (B*S, t, D)
+        x = x.view(BS, t, D)
+
+        return x  # (B*S, t, D)

modules/model/modules/feat_extractors/audio/hf_src/modeling_ast.py

@@ -0,0 +1,662 @@
+# coding=utf-8
+# Copyright 2022 MIT and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Modified by v-iashin to support token masking
+
+""" PyTorch Audio Spectrogram Transformer (AST) model."""
+
+import math
+from typing import Dict, List, Optional, Set, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, SequenceClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
+from transformers.models.audio_spectrogram_transformer.modeling_audio_spectrogram_transformer import ASTConfig
+from transformers.utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
+
+
+logger = logging.get_logger(__name__)
+
+# General docstring
+_CONFIG_FOR_DOC = "ASTConfig"
+
+# Base docstring
+_CHECKPOINT_FOR_DOC = "MIT/ast-finetuned-audioset-10-10-0.4593"
+_EXPECTED_OUTPUT_SHAPE = [1, 1214, 768]
+
+# Audio classification docstring
+_SEQ_CLASS_CHECKPOINT = "MIT/ast-finetuned-audioset-10-10-0.4593"
+_SEQ_CLASS_EXPECTED_OUTPUT = "'Speech'"
+_SEQ_CLASS_EXPECTED_LOSS = 0.17
+
+
+AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "MIT/ast-finetuned-audioset-10-10-0.4593",
+    # See all Audio Spectrogram Transformer models at https://huggingface.co/models?filter=ast
+]
+
+
+class ASTEmbeddings(nn.Module):
+    """
+    Construct the CLS token, position and patch embeddings.
+    """
+
+    def __init__(self, config: ASTConfig) -> None:
+        super().__init__()
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
+        self.distillation_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
+        self.patch_embeddings = ASTPatchEmbeddings(config)
+
+        frequency_out_dimension, time_out_dimension = self.get_shape(config)
+        num_patches = frequency_out_dimension * time_out_dimension
+        self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 2, config.hidden_size))
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.config = config
+
+    def get_shape(self, config):
+        # see Karpathy's cs231n blog on how to calculate the output dimensions
+        # https://cs231n.github.io/convolutional-networks/#conv
+        frequency_out_dimension = (config.num_mel_bins - config.patch_size) // config.frequency_stride + 1
+        time_out_dimension = (config.max_length - config.patch_size) // config.time_stride + 1
+
+        return frequency_out_dimension, time_out_dimension
+
+    def forward(self, input_values: torch.Tensor) -> torch.Tensor:
+        batch_size = input_values.shape[0]
+        embeddings = self.patch_embeddings(input_values)
+
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+        distillation_tokens = self.distillation_token.expand(batch_size, -1, -1)
+        embeddings = torch.cat((cls_tokens, distillation_tokens, embeddings), dim=1)
+        embeddings = embeddings + self.position_embeddings
+        embeddings = self.dropout(embeddings)
+
+        return embeddings
+
+
+class ASTPatchEmbeddings(nn.Module):
+    """
+    This class turns `input_values` into the initial `hidden_states` (patch embeddings) of shape `(batch_size,
+    seq_length, hidden_size)` to be consumed by a Transformer.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+
+        patch_size = config.patch_size
+        frequency_stride = config.frequency_stride
+        time_stride = config.time_stride
+
+        self.projection = nn.Conv2d(
+            1, config.hidden_size, kernel_size=(patch_size, patch_size), stride=(frequency_stride, time_stride)
+        )
+
+    def forward(self, input_values: torch.Tensor) -> torch.Tensor:
+        input_values = input_values.unsqueeze(1)
+        input_values = input_values.transpose(2, 3)
+        embeddings = self.projection(input_values).flatten(2).transpose(1, 2)
+        return embeddings
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTSelfAttention with ViT->AST
+class ASTSelfAttention(nn.Module):
+    def __init__(self, config: ASTConfig) -> None:
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+            raise ValueError(
+                f"The hidden size {config.hidden_size,} is not a multiple of the number of attention "
+                f"heads {config.num_attention_heads}."
+            )
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
+        self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
+        self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self, hidden_states, tok_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False
+    ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
+        mixed_query_layer = self.query(hidden_states)
+
+        key_layer = self.transpose_for_scores(self.key(hidden_states))
+        value_layer = self.transpose_for_scores(self.value(hidden_states))
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+
+        # apply masking if provided, tok_mask is (BS, N): 1s - keep; attention_scores is (BS, H, N, N)
+        if tok_mask is not None:
+            BS, N = tok_mask.shape
+            attention_scores = attention_scores.masked_fill(tok_mask.view(BS, 1, 1, N) == 0, float('-inf'))
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        attention_probs = self.dropout(attention_probs)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attention_probs = attention_probs * head_mask
+
+        context_layer = torch.matmul(attention_probs, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(new_context_layer_shape)
+
+        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+
+        return outputs
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTSelfOutput with ViT->AST
+class ASTSelfOutput(nn.Module):
+    """
+    The residual connection is defined in ASTLayer instead of here (as is the case with other models), due to the
+    layernorm applied before each block.
+    """
+
+    def __init__(self, config: ASTConfig) -> None:
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+
+        return hidden_states
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTAttention with ViT->AST
+class ASTAttention(nn.Module):
+    def __init__(self, config: ASTConfig) -> None:
+        super().__init__()
+        self.attention = ASTSelfAttention(config)
+        self.output = ASTSelfOutput(config)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads: Set[int]) -> None:
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads
+        )
+
+        # Prune linear layers
+        self.attention.query = prune_linear_layer(self.attention.query, index)
+        self.attention.key = prune_linear_layer(self.attention.key, index)
+        self.attention.value = prune_linear_layer(self.attention.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)
+        self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        tok_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
+        self_outputs = self.attention(hidden_states, tok_mask, head_mask, output_attentions)
+
+        attention_output = self.output(self_outputs[0], hidden_states)
+
+        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
+        return outputs
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTIntermediate with ViT->AST
+class ASTIntermediate(nn.Module):
+    def __init__(self, config: ASTConfig) -> None:
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+
+        return hidden_states
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTOutput with ViT->AST
+class ASTOutput(nn.Module):
+    def __init__(self, config: ASTConfig) -> None:
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+
+        hidden_states = hidden_states + input_tensor
+
+        return hidden_states
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTLayer with ViT->AST
+class ASTLayer(nn.Module):
+    """This corresponds to the Block class in the timm implementation."""
+
+    def __init__(self, config: ASTConfig) -> None:
+        super().__init__()
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        self.attention = ASTAttention(config)
+        self.intermediate = ASTIntermediate(config)
+        self.output = ASTOutput(config)
+        self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        tok_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
+        self_attention_outputs = self.attention(
+            self.layernorm_before(hidden_states),  # in AST, layernorm is applied before self-attention
+            tok_mask,
+            head_mask,
+            output_attentions=output_attentions,
+        )
+        attention_output = self_attention_outputs[0]
+        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights
+
+        # first residual connection
+        hidden_states = attention_output + hidden_states
+
+        # in AST, layernorm is also applied after self-attention
+        layer_output = self.layernorm_after(hidden_states)
+        layer_output = self.intermediate(layer_output)
+
+        # second residual connection is done here
+        layer_output = self.output(layer_output, hidden_states)
+
+        outputs = (layer_output,) + outputs
+
+        return outputs
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTEncoder with ViT->AST
+class ASTEncoder(nn.Module):
+    def __init__(self, config: ASTConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList([ASTLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        tok_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+    ) -> Union[tuple, BaseModelOutput]:
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+
+        for i, layer_module in enumerate(self.layer):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer_module),
+                    hidden_states,
+                    tok_mask,
+                    layer_head_mask,
+                )
+            else:
+                layer_outputs = layer_module(hidden_states, tok_mask, layer_head_mask, output_attentions)
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+        )
+
+
+class ASTPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = ASTConfig
+    base_model_prefix = "audio_spectrogram_transformer"
+    main_input_name = "input_values"
+    supports_gradient_checkpointing = True
+
+    # Copied from transformers.models.deit.modeling_deit.DeiTPreTrainedModel._init_weights
+    def _init_weights(self, module: Union[nn.Linear, nn.Conv2d, nn.LayerNorm]) -> None:
+        """Initialize the weights"""
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            # Upcast the input in `fp32` and cast it back to desired `dtype` to avoid
+            # `trunc_normal_cpu` not implemented in `half` issues
+            module.weight.data = nn.init.trunc_normal_(
+                module.weight.data.to(torch.float32), mean=0.0, std=self.config.initializer_range
+            ).to(module.weight.dtype)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    # Copied from transformers.models.vit.modeling_vit.ViTPreTrainedModel._set_gradient_checkpointing with ViT->AST
+    def _set_gradient_checkpointing(self, module: ASTEncoder, value: bool = False) -> None:
+        if isinstance(module, ASTEncoder):
+            module.gradient_checkpointing = value
+
+
+AUDIO_SPECTROGRAM_TRANSFORMER_START_DOCSTRING = r"""
+    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
+    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
+    behavior.
+
+    Parameters:
+        config ([`ASTConfig`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+AUDIO_SPECTROGRAM_TRANSFORMER_INPUTS_DOCSTRING = r"""
+    Args:
+        input_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Pixel values. Pixel values can be obtained using [`AutoFeatureExtractor`]. See
+            [`ASTFeatureExtractor.__call__`] for details.
+
+        head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
+            Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare AST Model transformer outputting raw hidden-states without any specific head on top.",
+    AUDIO_SPECTROGRAM_TRANSFORMER_START_DOCSTRING,
+)
+class ASTModel(ASTPreTrainedModel):
+    def __init__(self, config: ASTConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = ASTEmbeddings(config)
+        self.encoder = ASTEncoder(config)
+
+        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> ASTPatchEmbeddings:
+        return self.embeddings.patch_embeddings
+
+    def _prune_heads(self, heads_to_prune: Dict[int, List[int]]) -> None:
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    @add_start_docstrings_to_model_forward(AUDIO_SPECTROGRAM_TRANSFORMER_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=BaseModelOutputWithPooling,
+        config_class=_CONFIG_FOR_DOC,
+        modality="audio",
+        expected_output=_EXPECTED_OUTPUT_SHAPE,
+    )
+    def forward(
+        self,
+        input_values: Optional[torch.Tensor] = None,
+        cont_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_values is None:
+            raise ValueError("You have to specify input_values")
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+        embedding_output = self.embeddings(input_values)
+
+        # transforms the mask that has spectrogram dims to the token masking which is obtained after patching.
+        # Due to the ovelap in patching, getting the token mask from spectrogram mask is not straightforward,
+        # because one 16x16 content patch is encoded in two tokens if stride is <16. So, to get the mask for
+        # tokens I will apply the patching func (self.embeddings) to the tensor with infinities at the masked
+        # content position. For infs, the patching fn will return nans, which I'll use to get the token mask.
+        if cont_mask is not None:
+            indicator = torch.ones_like(input_values).to(input_values.dtype)
+            # replace content mask (0s) with infs
+            indicator[~cont_mask] = torch.inf
+            # apply patching; now nans are where the content mask was
+            with torch.no_grad():
+                indicator = self.embeddings(indicator)  # BS, N, D
+            # replace nans with 0s; these are the tokens that correspond to the masked content
+            tok_mask = ~torch.isnan(indicator)
+            # since all values in the D-dimension (latent) will also be nans, we can just use the first el
+            tok_mask = tok_mask[:, :, 0]   # (BS, 2+num_patches) -- 2 is from CLS and DISTIL tokens
+        else:
+            tok_mask = None
+
+        encoder_outputs = self.encoder(
+            embedding_output,
+            tok_mask=tok_mask,
+            head_mask=head_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = encoder_outputs[0]
+        sequence_output = self.layernorm(sequence_output)
+
+        pooled_output = (sequence_output[:, 0] + sequence_output[:, 1]) / 2
+
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        ), tok_mask
+
+
+class ASTMLPHead(nn.Module):
+    def __init__(self, config: ASTConfig):
+        super().__init__()
+        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dense = nn.Linear(
+            config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
+
+    def forward(self, hidden_state):
+        hidden_state = self.layernorm(hidden_state)
+        hidden_state = self.dense(hidden_state)
+        return hidden_state
+
+
+@add_start_docstrings(
+    """
+    Audio Spectrogram Transformer model with an audio classification head on top (a linear layer on top of the pooled
+    output) e.g. for datasets like AudioSet, Speech Commands v2.
+    """,
+    AUDIO_SPECTROGRAM_TRANSFORMER_START_DOCSTRING,
+)
+class ASTForAudioClassification(ASTPreTrainedModel):
+    def __init__(self, config: ASTConfig) -> None:
+        super().__init__(config)
+
+        self.num_labels = config.num_labels
+        self.audio_spectrogram_transformer = ASTModel(config)
+
+        # Classifier head
+        self.classifier = ASTMLPHead(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(AUDIO_SPECTROGRAM_TRANSFORMER_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_SEQ_CLASS_CHECKPOINT,
+        output_type=SequenceClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+        modality="audio",
+        expected_output=_SEQ_CLASS_EXPECTED_OUTPUT,
+        expected_loss=_SEQ_CLASS_EXPECTED_LOSS,
+    )
+    def forward(
+        self,
+        input_values: Optional[torch.Tensor] = None,
+        cont_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[tuple, SequenceClassifierOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the audio classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.audio_spectrogram_transformer(
+            input_values,
+            cont_mask=cont_mask,
+            head_mask=head_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        pooled_output = outputs[1]
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

modules/model/modules/feat_extractors/audio/resnet.py

@@ -0,0 +1,249 @@
+import sys
+from pathlib import Path
+import logging
+import einops
+
+import torch
+from torchvision.models.resnet import BasicBlock, Bottleneck, ResNet
+
+sys.path.append('.')  # nopep8
+
+from utils.utils import check_if_file_exists_else_download
+from model.modules.feat_extractors.audio.ast import FrequencyTransformerEncoderLayer
+from model.modules.feat_extractors.visual.motionformer import AveragePooling, TemporalTransformerEncoderLayer
+
+
+class ResNetAudio(ResNet):
+
+    def __init__(self, arch_name, num_classes, extract_features, ckpt_path=None, **kwargs):
+
+        if arch_name == 'resnet18':
+            block = BasicBlock
+            layers = [2, 2, 2, 2]
+        elif arch_name == 'resnet34':
+            block = BasicBlock
+            layers = [3, 4, 6, 3]
+        elif arch_name == 'resnet50':
+            block = Bottleneck
+            layers = [3, 4, 6, 3]
+        elif arch_name == 'resnet101':
+            block = Bottleneck
+            layers = [3, 4, 23, 3]
+        elif arch_name == 'resnet152':
+            block = Bottleneck
+            layers = [3, 8, 36, 3]
+        else:
+            raise NotImplementedError
+
+        super().__init__(block, layers, num_classes, **kwargs)
+
+        # replacing the old conv1 to the new one (RGB - 3; spectrogram - 1)
+        conv1 = self.conv1
+        self.conv1 = torch.nn.Conv2d(1, conv1.out_channels, conv1.kernel_size,
+                                     conv1.stride, conv1.padding, bias=conv1.bias)
+        self.extract_features = extract_features
+        self.embed_dim = self.fc.in_features
+
+        # load the ckpt
+        load_state_dict_resnet(self, ckpt_path, prefix='afeat_extractor.')
+
+    def _forward_impl(self, x):
+        # See note [TorchScript super()]
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        if self.extract_features:
+            return x
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return x
+
+    def forward(self, x):
+        return super().forward(x)
+
+
+class ResNet18AudioFeatures(ResNetAudio):
+
+    # ckpt_path should default to None, otherwise when no pre-training is desired it will throw an error
+    def __init__(self,
+                 extract_features: bool = False,
+                 ckpt_path: str = None,
+                 feat_type: str = None,
+                 max_spec_t: int = None,
+                 factorize_freq_time: bool = None,
+                 agg_freq_module: str = None,
+                 agg_time_module: str = None,
+                 add_global_repr: bool = True,
+                 agg_segments_module: str = None,
+                 max_segments: int = None,
+                 ) -> None:
+        super().__init__(arch_name='resnet18', num_classes=308, extract_features=extract_features,
+                         ckpt_path=ckpt_path)
+        assert extract_features, 'Not implemented otherwise'
+        self.extract_features = extract_features
+        self.feat_type = feat_type
+        self.max_spec_t = max_spec_t
+        # similar to s3d
+        self.nhead = 8
+        self.mlp_ratio = 4
+        self.drop_rate = 0.0
+
+        if ckpt_path is not None:
+            ckpt = torch.load(ckpt_path, map_location=torch.device('cpu'))
+            was_pt_on_vgs_cls = 'ResNetAudio-' in Path(ckpt_path).stem
+            if was_pt_on_vgs_cls:
+                self.load_state_dict(ckpt['model'], strict=True)
+
+        # saving some memory
+        if extract_features:
+            self.avgpool = torch.nn.Identity()
+            self.fc = torch.nn.Identity()
+
+        # define adapters if needed
+        self.factorize_freq_time = factorize_freq_time
+        # avoiding code duplication (used only if agg_*_module is TransformerEncoderLayer)
+        transf_enc_layer_kwargs = dict(
+            d_model=self.embed_dim, nhead=self.nhead, dim_feedforward=self.mlp_ratio*self.embed_dim,
+            activation=torch.nn.GELU(), batch_first=True, dropout=self.drop_rate, layer_norm_eps=1e-6,
+            norm_first=True,
+        )
+        if factorize_freq_time:
+            self.feat_type = 'last_hidden_state'  # this feat_type supports factorization
+            # frequency aggreration
+            if agg_freq_module == 'TransformerEncoderLayer':
+                self.freq_attn_agg = FrequencyTransformerEncoderLayer(**transf_enc_layer_kwargs)
+            elif agg_freq_module == 'AveragePooling':
+                self.freq_attn_agg = AveragePooling(avg_pattern='BS D f t -> BS D t',
+                                                    then_permute_pattern='BS D t -> BS t D')
+            # time aggreration
+            if agg_time_module == 'TransformerEncoderLayer':
+                self.temp_attn_agg = TemporalTransformerEncoderLayer(**transf_enc_layer_kwargs)
+            elif agg_time_module == 'AveragePooling':
+                self.temp_attn_agg = AveragePooling(avg_pattern='BS t D -> BS D')
+            elif 'Identity' in agg_time_module:
+                self.temp_attn_agg = torch.nn.Identity()
+        # define a global aggregation layer (aggregarate over segments)
+        self.add_global_repr = add_global_repr
+        if add_global_repr:
+            if agg_segments_module == 'TransformerEncoderLayer':
+                # we can reuse the same layer as for temporal factorization (B, dim_to_agg, D) -> (B, D)
+                # we need to add pos emb (PE) because previously we added the same PE for each segment
+                pos_max_len = max_segments if max_segments is not None else 16  # 16 = 10sec//0.64sec + 1
+                self.global_attn_agg = TemporalTransformerEncoderLayer(
+                    add_pos_emb=True, pos_emb_drop=self.drop_rate,
+                    pos_max_len=pos_max_len, **transf_enc_layer_kwargs
+                )
+            elif agg_segments_module == 'AveragePooling':
+                self.global_attn_agg = AveragePooling(avg_pattern='B S D -> B D')
+
+        # do not keep fc to save memory
+        self.fc = torch.nn.Identity()
+
+        if ckpt_path is not None:
+            ckpt = ckpt['state_dict']
+            was_pt_on_avclip = any('a_encoder.' in k[0] or 'v_encoder.' in k[0] for k in ckpt.items())
+            assert was_pt_on_vgs_cls is False, f'Unexpected ckpt: {ckpt_path}'
+            if was_pt_on_avclip:
+                ckpt_weights = dict()
+                for k, v in ckpt.items():
+                    if k.startswith(('module.a_encoder.', 'a_encoder.')):
+                        k = k.replace('module.', '').replace('a_encoder.', '')
+                        ckpt_weights[k] = v
+                _load_status = self.load_state_dict(ckpt_weights, strict=False)
+                if len(_load_status.missing_keys) > 0 or len(_load_status.unexpected_keys) > 0:
+                    logging.warning(f'Loading exact ckpt from {ckpt_path} failed. \n' \
+                                    f'Missing keys ({len(_load_status.missing_keys)}): ' \
+                                    f'{_load_status.missing_keys}, \n' \
+                                    f'Unexpected keys ({len(_load_status.unexpected_keys)}): ' \
+                                    f'{_load_status.unexpected_keys} \n' \
+                                    f'freq_attn_agg are expected to be unexpected if ckpt was pt contrastively '\
+                                    f'as well as fc could be missing because we use features, not a classifier.')
+                else:
+                    logging.info(f'Loading ResNet ckpt from {ckpt_path} succeeded.')
+
+        # print the number of parameters
+        logging.info(f'afeat_extractor: {sum(p.numel() for p in self.parameters() if p.requires_grad):,}')
+
+    def forward(self, x: torch.Tensor, for_loop: bool = False, cont_mask: torch.Tensor = None):
+        assert for_loop is False and cont_mask is None, 'Not implemented'
+        B, S, T, F = x.shape
+
+        # (BS, D) <- (B, S, T, D)
+        x = self.forward_segments(x)
+
+        # unpack the segments (using rest dimensions to support different shapes e.g. (BS, D) or (BS, t, D))
+        x = x.view(B, S, *x.shape[1:])
+        # x now is of shape (B, S, D) or (B, S, t, D) if `self.temp_attn_agg` is `Identity`
+
+        global_x = None
+        if self.extract_features and self.add_global_repr:  # lazy execution, throws AttributeError
+            assert len(x.shape) == 3, f'Local representation should be (B, S, D) {x.shape}'
+            global_x = self.global_attn_agg(x)  # (B, D)
+
+        return x, global_x  # x is (B, S, ...), global_x is (B, D) or None
+
+    def forward_segments(self, x):
+        x = einops.rearrange(x, 'B S T F -> (B S) 1 F T')
+        # (BS, D, f, t) <- (BS, 1, F, T)
+        x = super().forward(x)
+
+        if self.extract_features:
+            if self.factorize_freq_time:
+                x = self.freq_attn_agg(x)  # (BS, t, D)
+                x = self.temp_attn_agg(x)  # (BS, D) or (BS, t, D) if self.temp_attn_agg is Identity
+        return x
+
+
+def load_state_dict_resnet(model, ckpt_path, prefix):
+    if ckpt_path is not None:
+        check_if_file_exists_else_download(ckpt_path)
+        ckpt = torch.load(ckpt_path, map_location=torch.device('cpu'))
+        ckpt = ckpt.get('model', ckpt.get('state_dict', ckpt))
+        # we need to filter out the state_dict of the AVCLIP model (has both A and V extractors)
+        # and keep only the state_dict of the feat extractor
+        # FIXME: this is a bit hacky, but it works
+        was_pt_on_avclip = any('a_encoder.' in k[0] or 'v_encoder.' in k[0] for k in ckpt.items())
+        if not was_pt_on_avclip:
+            model.load_state_dict(ckpt)
+            logging.info(f'Loading ResNet ckpt from {ckpt_path} succeeded.')
+        # if was_pt_on_avclip:
+        #     ckpt_weights = dict()
+        #     for k, v in ckpt.items():
+        #         if k.startswith(('module.a_encoder.', 'a_encoder.')):
+        #             k = k.replace('module.', '').replace('a_encoder.', '')
+        #             ckpt_weights[k] = v
+        #     _load_status = model.load_state_dict(ckpt_weights, strict=False)
+        #     if len(_load_status.missing_keys) > 0 or len(_load_status.unexpected_keys) > 0:
+        #         logging.warning(f'Loading exact ckpt from {ckpt_path} failed. \n' \
+        #                         f'Missing keys ({len(_load_status.missing_keys)}): ' \
+        #                         f'{_load_status.missing_keys}, \n' \
+        #                         f'Unexpected keys ({len(_load_status.unexpected_keys)}): ' \
+        #                         f'{_load_status.unexpected_keys} \n' \
+        #                         f'freq_attn_agg are expected to be unexpected if ckpt was pt contrastively '\
+        #                         f'as well as fc could be missing because we use features, not a classifier.')
+        #     else:
+        #         logging.info(f'Loading ResNet ckpt from {ckpt_path} succeeded.')
+        # else:
+        #     model.load_state_dict(ckpt)
+        #     logging.info(f'Loading ResNet ckpt from {ckpt_path} succeeded.')
+
+
+if __name__ == '__main__':
+    B = 2
+    ckpt_path = './model/modules/feat_extractors/audio/22-06-24T08-10-33/ResNetAudio-22-06-24T08-10-33.pt'
+    afeat_extractor = ResNet18AudioFeatures(ckpt_path=ckpt_path)
+    # x = torch.rand(B, 1, 257, 1551)
+    # x = torch.rand(B, 1, 257, 1379)
+    x = torch.rand(B, 1, 128, 66)
+    x = afeat_extractor(x)
+    print(x.shape)

modules/model/modules/feat_extractors/train_clip_src/__init__.py

@@ -0,0 +1,3 @@
+import sys
+sys.path.append('model/modules/feat_extractors')
+sys.path.append('model/modules/feat_extractors/train_clip_src')

modules/model/modules/feat_extractors/train_clip_src/open_clip/__init__.py

@@ -0,0 +1,13 @@
+from .coca_model import CoCa
+from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from .factory import create_model, create_model_from_pretrained, get_tokenizer, create_loss
+from .factory import list_models, add_model_config, get_model_config, load_checkpoint
+from .loss import ClipLoss, DistillClipLoss, CoCaLoss
+from .model import CLIP, CustomTextCLIP, CLIPTextCfg, CLIPVisionCfg, \
+    convert_weights_to_lp, convert_weights_to_fp16, trace_model, get_cast_dtype
+from .openai import load_openai_model, list_openai_models
+from .pretrained import list_pretrained, list_pretrained_models_by_tag, list_pretrained_tags_by_model, \
+    get_pretrained_url, download_pretrained_from_url, is_pretrained_cfg, get_pretrained_cfg, download_pretrained
+from .push_to_hf_hub import push_pretrained_to_hf_hub, push_to_hf_hub
+from .tokenizer import SimpleTokenizer, tokenize, decode
+from .transform import image_transform, AugmentationCfg

modules/model/modules/feat_extractors/train_clip_src/open_clip/bpe_simple_vocab_16e6.txt.gz

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

modules/model/modules/feat_extractors/train_clip_src/open_clip/coca_model.py

@@ -0,0 +1,458 @@
+from typing import Optional
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+import numpy as np
+from dataclasses import dataclass
+
+from .transformer import (
+    LayerNormFp32,
+    LayerNorm,
+    QuickGELU,
+    MultimodalTransformer,
+)
+from .model import CLIPTextCfg, CLIPVisionCfg, _build_vision_tower, _build_text_tower
+
+try:
+    from transformers import (
+        BeamSearchScorer,
+        LogitsProcessorList,
+        TopPLogitsWarper,
+        TopKLogitsWarper,
+        RepetitionPenaltyLogitsProcessor,
+        MinLengthLogitsProcessor,
+        MaxLengthCriteria,
+        StoppingCriteriaList
+    )
+
+    GENERATION_TYPES = {
+        "top_k": TopKLogitsWarper,
+        "top_p": TopPLogitsWarper,
+        "beam_search": "beam_search"
+    }
+    _has_transformers = True
+except ImportError as e:
+    GENERATION_TYPES = {
+        "top_k": None,
+        "top_p": None,
+        "beam_search": "beam_search"
+    }
+    _has_transformers = False
+
+
+@dataclass
+class MultimodalCfg(CLIPTextCfg):
+    mlp_ratio: int = 4
+    dim_head: int = 64
+    heads: int = 8
+    n_queries: int = 256
+    attn_pooler_heads: int = 8
+
+
+def _build_text_decoder_tower(
+        embed_dim,
+        multimodal_cfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    multimodal_cfg = MultimodalCfg(**multimodal_cfg) if isinstance(multimodal_cfg, dict) else multimodal_cfg
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+    norm_layer = (
+        LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+    )
+
+    decoder = MultimodalTransformer(
+        context_length=multimodal_cfg.context_length,
+        width=multimodal_cfg.width,
+        heads=multimodal_cfg.heads,
+        layers=multimodal_cfg.layers,
+        ls_init_value=multimodal_cfg.ls_init_value,
+        output_dim=embed_dim,
+        act_layer=act_layer,
+        norm_layer=norm_layer,
+    )
+
+    return decoder
+
+
+class CoCa(nn.Module):
+    def __init__(
+            self,
+            embed_dim,
+            multimodal_cfg: MultimodalCfg,
+            text_cfg: CLIPTextCfg,
+            vision_cfg: CLIPVisionCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            pad_id: int = 0,
+    ):
+        super().__init__()
+        multimodal_cfg = MultimodalCfg(**multimodal_cfg) if isinstance(multimodal_cfg, dict) else multimodal_cfg
+        text_cfg = CLIPTextCfg(**text_cfg) if isinstance(text_cfg, dict) else text_cfg
+        vision_cfg = CLIPVisionCfg(**vision_cfg) if isinstance(vision_cfg, dict) else vision_cfg
+
+        self.text = _build_text_tower(
+            embed_dim=embed_dim,
+            text_cfg=text_cfg,
+            quick_gelu=quick_gelu,
+            cast_dtype=cast_dtype,
+        )
+
+        vocab_size = (
+            text_cfg.vocab_size  # for hf models
+            if hasattr(text_cfg, "hf_model_name") and text_cfg.hf_model_name is not None
+            else text_cfg.vocab_size
+        )
+
+        self.visual = _build_vision_tower(
+            embed_dim=embed_dim,
+            vision_cfg=vision_cfg,
+            quick_gelu=quick_gelu,
+            cast_dtype=cast_dtype,
+        )
+
+        self.text_decoder = _build_text_decoder_tower(
+            vocab_size,
+            multimodal_cfg=multimodal_cfg,
+            quick_gelu=quick_gelu,
+            cast_dtype=cast_dtype,
+        )
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        self.pad_id = pad_id
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.text.set_grad_checkpointing(enable)
+        self.text_decoder.set_grad_checkpointing(enable)
+
+    def _encode_image(self, images, normalize=True):
+        image_latent, tokens_embs = self.visual(images)
+        image_latent = F.normalize(image_latent, dim=-1) if normalize else image_latent
+        return image_latent, tokens_embs
+
+    def _encode_text(self, text, normalize=True, embed_cls=True):
+        text = text[:, :-1] if embed_cls else text # make space for CLS token
+        text_latent, token_emb = self.text(text)
+        text_latent = F.normalize(text_latent, dim=-1) if normalize else text_latent
+        return text_latent, token_emb
+
+    def encode_image(self, images, normalize=True):
+        image_latent, _ = self._encode_image(images, normalize=normalize)
+        return image_latent
+
+    def encode_text(self, text, normalize=True, embed_cls=True):
+        text_latent, _ = self._encode_text(text, normalize=normalize, embed_cls=embed_cls)
+        return text_latent
+
+    def forward(self, image, text, embed_cls=True, image_latent=None, image_embs=None):
+        text_latent, token_embs = self._encode_text(text, embed_cls=embed_cls)
+        if image_latent is None or image_embs is None:
+            image_latent, image_embs = self._encode_image(image)
+
+        # TODO: add assertion to avoid bugs?
+        labels = text[:, -token_embs.shape[1]:]
+
+        logits = self.text_decoder(image_embs, token_embs)
+        return {
+            "image_features": image_latent,
+            "text_features": text_latent,
+            "logits": logits,
+            "labels": labels,
+            "logit_scale": self.logit_scale.exp()
+        }
+
+    def generate(
+        self,
+        image,
+        text=None,
+        seq_len=30,
+        max_seq_len=77,
+        temperature=1.,
+        generation_type="beam_search",
+        top_p=0.1,  # keep tokens in the 1 - top_p quantile
+        top_k=1,  # keeps the top_k most probable tokens
+        pad_token_id=None,
+        eos_token_id=None,
+        sot_token_id=None,
+        num_beams=6,
+        num_beam_groups=3,
+        min_seq_len=5,
+        stopping_criteria=None,
+        repetition_penalty=1.0,
+        fixed_output_length=False # if True output.shape == (batch_size, seq_len)
+    ):
+        # taking many ideas and components from HuggingFace GenerationMixin
+        # https://huggingface.co/docs/transformers/main/en/main_classes/text_generation
+        assert _has_transformers, "Please install transformers for generate functionality. `pip install transformers`."
+        assert seq_len > min_seq_len, "seq_len must be larger than min_seq_len"
+
+        with torch.no_grad():
+            sot_token_id = 49406 if sot_token_id is None else sot_token_id
+            eos_token_id = 49407 if eos_token_id is None else eos_token_id
+            pad_token_id = self.pad_id if pad_token_id is None else pad_token_id
+            logit_processor = LogitsProcessorList(
+                [
+                    MinLengthLogitsProcessor(min_seq_len, eos_token_id),
+                    RepetitionPenaltyLogitsProcessor(repetition_penalty),
+                ]
+            )
+
+            if stopping_criteria is None:
+                stopping_criteria = [MaxLengthCriteria(max_length=seq_len)]
+
+            stopping_criteria = StoppingCriteriaList(
+                stopping_criteria
+            )
+
+            device = image.device
+
+            if generation_type == "beam_search":
+                output = self._generate_beamsearch(
+                    image_inputs = image,
+                    pad_token_id=pad_token_id,
+                    eos_token_id=eos_token_id,
+                    sot_token_id=sot_token_id,
+                    num_beams=num_beams,
+                    num_beam_groups=num_beam_groups,
+                    min_seq_len=min_seq_len,
+                    stopping_criteria=stopping_criteria,
+                    logit_processor=logit_processor,
+                )
+                if fixed_output_length and output.shape[1] < seq_len:
+                    return torch.cat(
+                        (output, torch.ones(output.shape[0], seq_len-output.shape[1], device=device, dtype=output.dtype) * self.pad_id),
+                        dim=1
+                    )
+                return output
+
+            elif generation_type == "top_p":
+                logit_warper = GENERATION_TYPES[generation_type](top_p)
+            elif generation_type == "top_k":
+                logit_warper = GENERATION_TYPES[generation_type](top_k)
+            else:
+                raise ValueError(
+                    f"generation_type has to be one of "
+                    f"{'| ' + ' | '.join(list(GENERATION_TYPES.keys())) + ' |'}."
+                )
+
+            image_latent, image_embs = self._encode_image(image)
+
+            if text is None:
+                text = torch.ones((image.shape[0], 1), device=device, dtype=torch.long) * sot_token_id
+
+            was_training = self.training
+            num_dims = len(text.shape)
+
+            if num_dims == 1:
+                text = text[None, :]
+
+            cur_len = text.shape[1]
+            self.eval()
+            out = text
+
+            while True:
+                x = out[:, -max_seq_len:]
+                cur_len = x.shape[1]
+                logits = self(image, x, image_latent=image_latent, image_embs=image_embs, embed_cls=False)["logits"][:, -1]
+                mask = (out[:, -1] == eos_token_id) | (out[:, -1] == pad_token_id)
+                sample = torch.ones((out.shape[0], 1), device=device, dtype=torch.long) * pad_token_id
+
+                if mask.all():
+                    if not fixed_output_length:
+                        break
+                else:
+                    logits = logits[~mask, :]
+                    filtered_logits = logit_processor(x[~mask, :], logits)
+                    filtered_logits = logit_warper(x[~mask, :], filtered_logits)
+                    probs = F.softmax(filtered_logits / temperature, dim=-1)
+
+                    if (cur_len + 1 == seq_len):
+                        sample[~mask, :] = torch.ones((sum(~mask), 1), device=device, dtype=torch.long) * eos_token_id
+                    else:
+                        sample[~mask, :] = torch.multinomial(probs, 1)
+
+                out = torch.cat((out, sample), dim=-1)
+
+                cur_len += 1
+
+                if stopping_criteria(out, None):
+                    break
+
+            if num_dims == 1:
+                out = out.squeeze(0)
+
+            self.train(was_training)
+            return out
+
+    def _generate_beamsearch(
+            self,
+            image_inputs,
+            pad_token_id=None,
+            eos_token_id=None,
+            sot_token_id=None,
+            num_beams=6,
+            num_beam_groups=3,
+            min_seq_len=5,
+            stopping_criteria=None,
+            logit_processor=None,
+            logit_warper=None,
+    ):
+        device = image_inputs.device
+        batch_size = image_inputs.shape[0]
+        image_inputs = torch.repeat_interleave(image_inputs, num_beams, dim=0)
+        image_latent, image_embs = self._encode_image(image_inputs)
+
+        input_ids = torch.ones((batch_size * num_beams, 1), device=device, dtype=torch.long)
+        input_ids = input_ids * sot_token_id
+        beam_scorer = BeamSearchScorer(
+            batch_size=batch_size,
+            num_beams=num_beams,
+            device=device,
+            num_beam_groups=num_beam_groups,
+        )
+        # instantiate logits processors
+        logits_processor = (
+            LogitsProcessorList([MinLengthLogitsProcessor(min_seq_len, eos_token_id=eos_token_id)])
+            if logit_processor is None
+            else logit_processor
+        )
+
+        batch_size = len(beam_scorer._beam_hyps)
+        num_beams = beam_scorer.num_beams
+        num_beam_groups = beam_scorer.num_beam_groups
+        num_sub_beams = num_beams // num_beam_groups
+        batch_beam_size, cur_len = input_ids.shape
+        beam_indices = None
+
+        if num_beams * batch_size != batch_beam_size:
+            raise ValueError(
+                f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
+            )
+
+        beam_scores = torch.full((batch_size, num_beams), -1e9, dtype=torch.float, device=device)
+        # initialise score of first beam of each group with 0 and the rest with 1e-9. This ensures that the beams in
+        # the same group don't produce same tokens everytime.
+        beam_scores[:, ::num_sub_beams] = 0
+        beam_scores = beam_scores.view((batch_size * num_beams,))
+
+        while True:
+
+            # predicted tokens in cur_len step
+            current_tokens = torch.zeros(batch_size * num_beams, dtype=input_ids.dtype, device=device)
+
+            # indices which will form the beams in the next time step
+            reordering_indices = torch.zeros(batch_size * num_beams, dtype=torch.long, device=device)
+
+            # do one decoder step on all beams of all sentences in batch
+            model_inputs = prepare_inputs_for_generation(input_ids=input_ids, image_inputs=image_inputs)
+            outputs = self(
+                model_inputs['images'],
+                model_inputs['text'],
+                embed_cls=False,
+                image_latent=image_latent,
+                image_embs=image_embs
+            )
+
+            for beam_group_idx in range(num_beam_groups):
+                group_start_idx = beam_group_idx * num_sub_beams
+                group_end_idx = min(group_start_idx + num_sub_beams, num_beams)
+                group_size = group_end_idx - group_start_idx
+
+                # indices of beams of current group among all sentences in batch
+                batch_group_indices = []
+
+                for batch_idx in range(batch_size):
+                    batch_group_indices.extend(
+                        [batch_idx * num_beams + idx for idx in range(group_start_idx, group_end_idx)]
+                    )
+                group_input_ids = input_ids[batch_group_indices]
+
+                # select outputs of beams of currentg group only
+                next_token_logits = outputs['logits'][batch_group_indices, -1, :]
+                vocab_size = next_token_logits.shape[-1]
+
+                next_token_scores_processed = logits_processor(
+                    group_input_ids, next_token_logits, current_tokens=current_tokens, beam_group_idx=beam_group_idx
+                )
+                next_token_scores = next_token_scores_processed + beam_scores[batch_group_indices].unsqueeze(-1)
+                next_token_scores = next_token_scores.expand_as(next_token_scores_processed)
+
+                # reshape for beam search
+                next_token_scores = next_token_scores.view(batch_size, group_size * vocab_size)
+
+                next_token_scores, next_tokens = torch.topk(
+                    next_token_scores, 2 * group_size, dim=1, largest=True, sorted=True
+                )
+
+                next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor")
+                next_tokens = next_tokens % vocab_size
+
+                # stateless
+                process_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None
+                beam_outputs = beam_scorer.process(
+                    group_input_ids,
+                    next_token_scores,
+                    next_tokens,
+                    next_indices,
+                    pad_token_id=pad_token_id,
+                    eos_token_id=eos_token_id,
+                    beam_indices=process_beam_indices,
+                )
+                beam_scores[batch_group_indices] = beam_outputs["next_beam_scores"]
+                beam_next_tokens = beam_outputs["next_beam_tokens"]
+                beam_idx = beam_outputs["next_beam_indices"]
+
+                input_ids[batch_group_indices] = group_input_ids[beam_idx]
+                group_input_ids = torch.cat([group_input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)
+                current_tokens[batch_group_indices] = group_input_ids[:, -1]
+
+                # (beam_idx // group_size) -> batch_idx
+                # (beam_idx % group_size) -> offset of idx inside the group
+                reordering_indices[batch_group_indices] = (
+                    num_beams * torch.div(beam_idx, group_size, rounding_mode="floor") + group_start_idx + (beam_idx % group_size)
+                )
+
+            input_ids = torch.cat([input_ids, current_tokens.unsqueeze(-1)], dim=-1)
+
+            # increase cur_len
+            cur_len = cur_len + 1
+            if beam_scorer.is_done or stopping_criteria(input_ids, None):
+                break
+
+        final_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None
+        sequence_outputs = beam_scorer.finalize(
+            input_ids,
+            beam_scores,
+            next_tokens,
+            next_indices,
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            max_length=stopping_criteria.max_length,
+            beam_indices=final_beam_indices,
+        )
+        return sequence_outputs['sequences']
+
+
+def prepare_inputs_for_generation(input_ids, image_inputs, past=None, **kwargs):
+    if past:
+        input_ids = input_ids[:, -1].unsqueeze(-1)
+
+    attention_mask = kwargs.get("attention_mask", None)
+    position_ids = kwargs.get("position_ids", None)
+
+    if attention_mask is not None and position_ids is None:
+        # create position_ids on the fly for batch generation
+        position_ids = attention_mask.long().cumsum(-1) - 1
+        position_ids.masked_fill_(attention_mask == 0, 1)
+    else:
+        position_ids = None
+    return {
+        "text": input_ids,
+        "images": image_inputs,
+        "past_key_values": past,
+        "position_ids": position_ids,
+        "attention_mask": attention_mask,
+    }

modules/model/modules/feat_extractors/train_clip_src/open_clip/constants.py

@@ -0,0 +1,2 @@
+OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
+OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)

modules/model/modules/feat_extractors/train_clip_src/open_clip/factory.py

@@ -0,0 +1,193 @@
+import json
+import re
+from copy import deepcopy
+from pathlib import Path
+from typing import Optional, Tuple, Union
+
+import torch
+
+from utils.utils import instantiate_from_config
+
+from .model import convert_to_custom_text_state_dict, resize_pos_embed
+from .loss import AVCLIPLoss, ClipLoss, DistillClipLoss, CoCaLoss, MultilevelAVCLIPLoss
+from .transform import image_transform
+from .tokenizer import HFTokenizer, tokenize
+
+
+HF_HUB_PREFIX = 'hf-hub:'
+_MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
+_MODEL_CONFIGS = {}  # directory (model_name: config) of model architecture configs
+
+
+def _natural_key(string_):
+    return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]
+
+
+def _rescan_model_configs():
+    global _MODEL_CONFIGS
+
+    config_ext = ('.json',)
+    config_files = []
+    for config_path in _MODEL_CONFIG_PATHS:
+        if config_path.is_file() and config_path.suffix in config_ext:
+            config_files.append(config_path)
+        elif config_path.is_dir():
+            for ext in config_ext:
+                config_files.extend(config_path.glob(f'*{ext}'))
+
+    for cf in config_files:
+        with open(cf, 'r') as f:
+            model_cfg = json.load(f)
+            if all(a in model_cfg for a in ('embed_dim', 'vision_cfg', 'text_cfg')):
+                _MODEL_CONFIGS[cf.stem] = model_cfg
+
+    _MODEL_CONFIGS = {k: v for k, v in sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0]))}
+
+
+_rescan_model_configs()  # initial populate of model config registry
+
+
+def list_models():
+    """ enumerate available model architectures based on config files """
+    return list(_MODEL_CONFIGS.keys())
+
+
+def add_model_config(path):
+    """ add model config path or file and update registry """
+    if not isinstance(path, Path):
+        path = Path(path)
+    _MODEL_CONFIG_PATHS.append(path)
+    _rescan_model_configs()
+
+
+def get_model_config(model_name):
+    if model_name in _MODEL_CONFIGS:
+        return deepcopy(_MODEL_CONFIGS[model_name])
+    else:
+        return None
+
+
+def get_tokenizer(model_name):
+    if model_name.startswith(HF_HUB_PREFIX):
+        tokenizer = HFTokenizer(model_name[len(HF_HUB_PREFIX):])
+    else:
+        config = get_model_config(model_name)
+        tokenizer = HFTokenizer(
+            config['text_cfg']['hf_tokenizer_name']) if 'hf_tokenizer_name' in config['text_cfg'] else tokenize
+    return tokenizer
+
+
+def load_state_dict(checkpoint_path: str, map_location='cpu'):
+    checkpoint = torch.load(checkpoint_path, map_location=map_location)
+    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
+        state_dict = checkpoint['state_dict']
+    else:
+        state_dict = checkpoint
+    if next(iter(state_dict.items()))[0].startswith('module'):
+        state_dict = {k[7:]: v for k, v in state_dict.items()}
+    return state_dict
+
+
+def load_checkpoint(model, checkpoint_path, strict=True):
+    state_dict = load_state_dict(checkpoint_path)
+    # detect old format and make compatible with new format
+    if 'positional_embedding' in state_dict and not hasattr(model, 'positional_embedding'):
+        state_dict = convert_to_custom_text_state_dict(state_dict)
+    resize_pos_embed(state_dict, model)
+    incompatible_keys = model.load_state_dict(state_dict, strict=strict)
+    return incompatible_keys
+
+
+def create_model(cfg, device: Union[str, torch.device] = 'cpu'):
+    if isinstance(device, str):
+        device = torch.device(device)
+    model = instantiate_from_config(cfg.model)
+    model.to(device=device)
+    return model
+
+
+def create_loss(args):
+    if args.distill:
+        return DistillClipLoss(
+            local_loss=args.local_loss,
+            gather_with_grad=args.gather_with_grad,
+            cache_labels=True,
+            rank=args.rank,
+            world_size=args.world_size,
+        )
+    elif 'coca' in args.model.target.split('.')[-1].lower():
+        return CoCaLoss(
+            caption_loss_weight=args.coca_caption_loss_weight,
+            clip_loss_weight=args.coca_contrastive_loss_weight,
+            local_loss=args.local_loss,
+            gather_with_grad=args.gather_with_grad,
+            cache_labels=True,
+            rank=args.rank,
+            world_size=args.world_size,
+        )
+    elif 'multilevel' in args.model.target.split('.')[-1].lower():
+        return MultilevelAVCLIPLoss(
+            local_loss=args.local_loss,
+            gather_with_grad=args.gather_with_grad,
+            cache_labels=True,
+            rank=args.rank,
+            world_size=args.world_size,
+        )
+    elif 'avclip' in args.model.target.split('.')[-1].lower():
+        return AVCLIPLoss(
+            local_loss=args.local_loss,
+            gather_with_grad=args.gather_with_grad,
+            cache_labels=True,
+            rank=args.rank,
+            world_size=args.world_size,
+        )
+
+    return ClipLoss(
+        local_loss=args.local_loss,
+        gather_with_grad=args.gather_with_grad,
+        cache_labels=True,
+        rank=args.rank,
+        world_size=args.world_size,
+    )
+
+
+def create_model_from_pretrained(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        return_transform: bool = True,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        cache_dir: Optional[str] = None,
+):
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_text=force_custom_text,
+        force_image_size=force_image_size,
+        cache_dir=cache_dir,
+        require_pretrained=True,
+    )
+
+    if not return_transform:
+        return model
+
+    image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+    image_std = image_std or getattr(model.visual, 'image_std', None)
+    preprocess = image_transform(
+        model.visual.image_size,
+        is_train=False,
+        mean=image_mean,
+        std=image_std,
+    )
+
+    return model, preprocess

modules/model/modules/feat_extractors/train_clip_src/open_clip/hf_configs.py

@@ -0,0 +1,45 @@
+# HF architecture dict:
+arch_dict = {
+    # https://huggingface.co/docs/transformers/model_doc/roberta#roberta
+    "roberta": {
+        "config_names": {
+            "context_length": "max_position_embeddings",
+            "vocab_size": "vocab_size",
+            "width": "hidden_size",
+            "heads": "num_attention_heads",
+            "layers": "num_hidden_layers",
+            "layer_attr": "layer",
+            "token_embeddings_attr": "embeddings"
+        },
+        "pooler": "mean_pooler",
+    },
+    # https://huggingface.co/docs/transformers/model_doc/xlm-roberta#transformers.XLMRobertaConfig
+    "xlm-roberta": {
+        "config_names": {
+            "context_length": "max_position_embeddings",
+            "vocab_size": "vocab_size",
+            "width": "hidden_size",
+            "heads": "num_attention_heads",
+            "layers": "num_hidden_layers",
+            "layer_attr": "layer",
+            "token_embeddings_attr": "embeddings"
+        },
+        "pooler": "mean_pooler",
+    },
+    # https://huggingface.co/docs/transformers/model_doc/mt5#mt5
+    "mt5": {
+        "config_names": {
+            # unlimited seqlen
+            # https://github.com/google-research/text-to-text-transfer-transformer/issues/273
+            # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/t5/modeling_t5.py#L374
+            "context_length": "",
+            "vocab_size": "vocab_size",
+            "width": "d_model",
+            "heads": "num_heads",
+            "layers": "num_layers",
+            "layer_attr": "block",
+            "token_embeddings_attr": "embed_tokens"
+        },
+        "pooler": "mean_pooler",
+    },
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/hf_model.py

@@ -0,0 +1,176 @@
+""" huggingface model adapter
+
+Wraps HuggingFace transformers (https://github.com/huggingface/transformers) models for use as a text tower in CLIP model.
+"""
+
+import re
+
+import torch
+import torch.nn as nn
+from torch import TensorType
+
+try:
+    import transformers
+    from transformers import AutoModel, AutoTokenizer, AutoConfig, PretrainedConfig
+    from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, \
+        BaseModelOutputWithPoolingAndCrossAttentions
+except ImportError as e:
+    transformers = None
+
+
+    class BaseModelOutput:
+        pass
+
+
+    class PretrainedConfig:
+        pass
+
+from .hf_configs import arch_dict
+
+
+# utils
+def _camel2snake(s):
+    return re.sub(r'(?<!^)(?=[A-Z])', '_', s).lower()
+
+
+# TODO: ?last - for gpt-like models
+_POOLERS = {}
+
+
+def register_pooler(cls):
+    """Decorator registering pooler class"""
+    _POOLERS[_camel2snake(cls.__name__)] = cls
+    return cls
+
+
+@register_pooler
+class MeanPooler(nn.Module):
+    """Mean pooling"""
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        masked_output = x.last_hidden_state * attention_mask.unsqueeze(-1)
+        return masked_output.sum(dim=1) / attention_mask.sum(-1, keepdim=True)
+
+
+@register_pooler
+class MaxPooler(nn.Module):
+    """Max pooling"""
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        masked_output = x.last_hidden_state.masked_fill(attention_mask.unsqueeze(-1), -torch.inf)
+        return masked_output.max(1).values
+
+
+@register_pooler
+class ClsPooler(nn.Module):
+    """CLS token pooling"""
+
+    def __init__(self, use_pooler_output=True):
+        super().__init__()
+        self.cls_token_position = 0
+        self.use_pooler_output = use_pooler_output
+
+    def forward(self, x: BaseModelOutput, attention_mask: TensorType):
+        if (self.use_pooler_output and
+            isinstance(x, (BaseModelOutputWithPooling, BaseModelOutputWithPoolingAndCrossAttentions)) and
+            (x.pooler_output is not None)
+        ):
+            return x.pooler_output
+
+        return x.last_hidden_state[:, self.cls_token_position, :]
+
+
+class HFTextEncoder(nn.Module):
+    """HuggingFace model adapter"""
+    output_tokens: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            model_name_or_path: str,
+            output_dim: int,
+            config: PretrainedConfig = None,
+            pooler_type: str = None,
+            proj: str = None,
+            pretrained: bool = True,
+            output_tokens: bool = False,
+    ):
+        super().__init__()
+        self.output_tokens = output_tokens
+        self.output_dim = output_dim
+
+        # TODO: find better way to get this information
+        uses_transformer_pooler = (pooler_type == "cls_pooler")
+
+        if transformers is None:
+            raise RuntimeError("Please `pip install transformers` to use pre-trained HuggingFace models")
+        if config is None:
+            self.config = AutoConfig.from_pretrained(model_name_or_path)
+            create_func, model_args = (AutoModel.from_pretrained, model_name_or_path) if pretrained else (
+                AutoModel.from_config, self.config)
+            # TODO: do all model configs have this attribute? PretrainedConfig does so yes??
+            if hasattr(self.config, "is_encoder_decoder") and self.config.is_encoder_decoder:
+                self.transformer = create_func(model_args)
+                self.transformer = self.transformer.encoder
+            else:
+                self.transformer = create_func(model_args, add_pooling_layer=uses_transformer_pooler)
+        else:
+            self.config = config
+            self.transformer = AutoModel.from_config(config)
+        if pooler_type is None:  # get default arch pooler
+            pooler_type = (arch_dict[self.config.model_type]["pooler"])
+        
+        self.pooler = _POOLERS[pooler_type]()
+
+        d_model = getattr(self.config, arch_dict[self.config.model_type]["config_names"]["width"])
+        if (d_model == output_dim) and (proj is None):  # do we always need a proj?
+            self.proj = nn.Identity()
+        elif proj == 'linear':
+            self.proj = nn.Linear(d_model, output_dim, bias=False)
+        elif proj == 'mlp':
+            hidden_size = (d_model + output_dim) // 2
+            self.proj = nn.Sequential(
+                nn.Linear(d_model, hidden_size, bias=False),
+                nn.GELU(),
+                nn.Linear(hidden_size, output_dim, bias=False),
+            )
+
+    def forward(self, x: TensorType):
+        attn_mask = (x != self.config.pad_token_id).long()
+        out = self.transformer(input_ids=x, attention_mask=attn_mask)
+        pooled_out = self.pooler(out, attn_mask)
+        projected = self.proj(pooled_out)
+
+        seq_len = out.last_hidden_state.shape[1]
+        tokens = (
+            out.last_hidden_state[:, torch.arange(seq_len) != self.pooler.cls_token_position, :] 
+            if type(self.pooler) == ClsPooler 
+            else out.last_hidden_state
+        )
+        
+        if self.output_tokens:
+            return projected, tokens
+        return projected
+
+    def lock(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
+        if not unlocked_layers:  # full freezing
+            for n, p in self.transformer.named_parameters():
+                p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
+            return
+
+        encoder = self.transformer.encoder if hasattr(self.transformer, 'encoder') else self.transformer
+        layer_list = getattr(encoder, arch_dict[self.config.model_type]["config_names"]["layer_attr"])
+        print(f"Unlocking {unlocked_layers}/{len(layer_list) + 1} layers of hf model")
+        embeddings = getattr(
+            self.transformer, arch_dict[self.config.model_type]["config_names"]["token_embeddings_attr"])
+        modules = [embeddings, *layer_list][:-unlocked_layers]
+        # freeze layers
+        for module in modules:
+            for n, p in module.named_parameters():
+                p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.transformer.gradient_checkpointing_enable()
+
+    def init_parameters(self):
+        pass

modules/model/modules/feat_extractors/train_clip_src/open_clip/loss.py

@@ -0,0 +1,229 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+try:
+    import torch.distributed.nn
+    from torch import distributed as dist
+
+    has_distributed = True
+except ImportError:
+    has_distributed = False
+
+
+def gather_features(
+        image_features,
+        text_features,
+        local_loss=False,
+        gather_with_grad=False,
+        rank=0,
+        world_size=1,
+):
+    assert has_distributed, 'torch.distributed did not import correctly, please use a PyTorch version with support.'
+    # We gather tensors from all gpus
+    if gather_with_grad:
+        all_image_features = torch.cat(torch.distributed.nn.all_gather(image_features), dim=0)
+        all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features), dim=0)
+    else:
+        gathered_image_features = [torch.zeros_like(image_features) for _ in range(world_size)]
+        gathered_text_features = [torch.zeros_like(text_features) for _ in range(world_size)]
+        dist.all_gather(gathered_image_features, image_features)
+        dist.all_gather(gathered_text_features, text_features)
+        if not local_loss:
+            # ensure grads for local rank when all_* features don't have a gradient
+            gathered_image_features[rank] = image_features
+            gathered_text_features[rank] = text_features
+        all_image_features = torch.cat(gathered_image_features, dim=0)
+        all_text_features = torch.cat(gathered_text_features, dim=0)
+
+    return all_image_features, all_text_features
+
+
+class ClipLoss(nn.Module):
+
+    def __init__(
+            self,
+            local_loss=False,
+            gather_with_grad=False,
+            cache_labels=False,
+            rank=0,
+            world_size=1,
+    ):
+        super().__init__()
+        self.local_loss = local_loss
+        self.gather_with_grad = gather_with_grad
+        self.cache_labels = cache_labels
+        self.rank = rank
+        self.world_size = world_size
+
+        # cache state
+        self.prev_num_logits = 0
+        self.labels = {}
+
+    def get_ground_truth(self, device, num_logits) -> torch.Tensor:
+        # calculated ground-truth and cache if enabled
+        if self.prev_num_logits != num_logits or device not in self.labels:
+            labels = torch.arange(num_logits, device=device, dtype=torch.long)
+            if self.world_size > 1 and self.local_loss:
+                labels = labels + num_logits * self.rank
+            if self.cache_labels:
+                self.labels[device] = labels
+                self.prev_num_logits = num_logits
+        else:
+            labels = self.labels[device]
+        return labels
+
+    def get_logits(self, image_features, text_features, logit_scale):
+        if self.world_size > 1:
+            all_image_features, all_text_features = gather_features(
+                image_features, text_features,
+                self.local_loss, self.gather_with_grad, self.rank, self.world_size)
+
+            if self.local_loss:
+                logits_per_image = logit_scale * image_features @ all_text_features.T
+                logits_per_text = logit_scale * text_features @ all_image_features.T
+            else:
+                logits_per_image = logit_scale * all_image_features @ all_text_features.T
+                logits_per_text = logits_per_image.T
+        else:
+            logits_per_image = logit_scale * image_features @ text_features.T
+            logits_per_text = logit_scale * text_features @ image_features.T
+
+        return logits_per_image, logits_per_text
+
+    def forward(self, image_features, text_features, logit_scale, output_dict=False):
+        device = image_features.device
+        logits_per_image, logits_per_text = self.get_logits(image_features, text_features, logit_scale)
+
+        labels = self.get_ground_truth(device, logits_per_image.shape[0])
+
+        total_loss = (
+            F.cross_entropy(logits_per_image, labels) + F.cross_entropy(logits_per_text, labels)
+        ) / 2
+
+        return {"contrastive_loss": total_loss} if output_dict else total_loss
+
+
+class CoCaLoss(ClipLoss):
+    def __init__(
+            self,
+            caption_loss_weight,
+            clip_loss_weight,
+            pad_id=0,  # pad_token for open_clip custom tokenizer
+            local_loss=False,
+            gather_with_grad=False,
+            cache_labels=False,
+            rank=0,
+            world_size=1,
+    ):
+        super().__init__(
+            local_loss=local_loss,
+            gather_with_grad=gather_with_grad,
+            cache_labels=cache_labels,
+            rank=rank,
+            world_size=world_size,
+        )
+
+        self.clip_loss_weight = clip_loss_weight
+        self.caption_loss_weight = caption_loss_weight
+        self.caption_loss = nn.CrossEntropyLoss(ignore_index=pad_id)
+
+    def forward(self, image_features, text_features, logits, labels, logit_scale, output_dict=False):
+        clip_loss = super().forward(image_features, text_features, logit_scale)
+        clip_loss = self.clip_loss_weight * clip_loss
+
+        caption_loss = self.caption_loss(
+            logits.permute(0, 2, 1),
+            labels,
+        )
+        caption_loss = caption_loss * self.caption_loss_weight
+
+        if output_dict:
+            return {"contrastive_loss": clip_loss, "caption_loss": caption_loss}
+
+        return clip_loss, caption_loss
+
+
+class DistillClipLoss(ClipLoss):
+
+    def dist_loss(self, teacher_logits, student_logits):
+        return -(teacher_logits.softmax(dim=1) * student_logits.log_softmax(dim=1)).sum(dim=1).mean(dim=0)
+
+    def forward(
+            self,
+            image_features,
+            text_features,
+            logit_scale,
+            dist_image_features,
+            dist_text_features,
+            dist_logit_scale,
+            output_dict=False,
+    ):
+        logits_per_image, logits_per_text = \
+            self.get_logits(image_features, text_features, logit_scale)
+
+        dist_logits_per_image, dist_logits_per_text = \
+            self.get_logits(dist_image_features, dist_text_features, dist_logit_scale)
+
+        labels = self.get_ground_truth(image_features.device, logits_per_image.shape[0])
+
+        contrastive_loss = (
+            F.cross_entropy(logits_per_image, labels) +
+            F.cross_entropy(logits_per_text, labels)
+        ) / 2
+
+        distill_loss = (
+            self.dist_loss(dist_logits_per_image, logits_per_image) +
+            self.dist_loss(dist_logits_per_text, logits_per_text)
+        ) / 2
+
+        if output_dict:
+            return {"contrastive_loss": contrastive_loss, "distill_loss": distill_loss}
+
+        return contrastive_loss, distill_loss
+
+
+class AVCLIPLoss(ClipLoss):
+    '''This loss is resembles the CLIP loss, but it simply renames the variables'''
+
+    def __init__(self, local_loss=False, gather_with_grad=False, cache_labels=False, rank=0, world_size=1):
+        super().__init__(local_loss, gather_with_grad, cache_labels, rank, world_size)
+
+    def forward(self, rgb_features, audio_features, logit_scale, output_dict=False):
+        return super().forward(rgb_features, audio_features, logit_scale, output_dict)
+
+
+class MultilevelAVCLIPLoss(nn.Module):
+
+    def __init__(self, local_loss=False, gather_with_grad=False, cache_labels=False, rank=0, world_size=1):
+        super().__init__()
+        self.segment_avclip_loss = AVCLIPLoss(
+            local_loss, gather_with_grad, cache_labels, rank, world_size)
+        self.global_avclip_loss = AVCLIPLoss(
+            local_loss, gather_with_grad, cache_labels, rank, world_size)
+
+    def forward(self, rgb_features, audio_features, logit_scales, output_dict=False):
+        segment_rgb, global_rgb = rgb_features
+        segment_audio, global_audio = audio_features
+        segment_logit_scale, global_logit_scale = logit_scales
+
+        assert segment_rgb.shape[0] == segment_audio.shape[0], f'{segment_rgb.shape} != {segment_audio.shape}'
+        B, S, D = segment_rgb.shape
+
+        # (B*S, D) <- (B, S, D)
+        segment_loss = self.segment_avclip_loss(segment_rgb.view(B*S, D), segment_audio.view(B*S, D),
+                                                segment_logit_scale, output_dict)
+        global_loss = None
+        if global_rgb is not None:
+            global_loss = self.global_avclip_loss(global_rgb, global_audio, global_logit_scale, output_dict)
+
+        if output_dict:
+            losses = {'segment_contrastive_loss': segment_loss['contrastive_loss']}
+            if global_rgb is not None:
+                losses['global_contrastive_loss'] = global_loss['contrastive_loss']
+            return losses
+        else:
+            if global_loss is None:
+                return segment_loss
+            else:
+                return segment_loss, global_loss

modules/model/modules/feat_extractors/train_clip_src/open_clip/model.py

@@ -0,0 +1,883 @@
+""" CLIP Model
+
+Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+from dataclasses import dataclass
+import logging
+import math
+from typing import Optional, Tuple, Union
+
+import numpy as np
+from omegaconf import OmegaConf
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from utils.utils import instantiate_from_config
+
+from .hf_model import HFTextEncoder
+from .modified_resnet import ModifiedResNet
+from .timm_model import TimmModel
+from .transformer import LayerNormFp32, LayerNorm, QuickGELU, Attention, VisionTransformer, TextTransformer
+from .utils import to_2tuple
+
+
+@dataclass
+class CLIPVisionCfg:
+    layers: Union[Tuple[int, int, int, int], int] = 12
+    width: int = 768
+    head_width: int = 64
+    mlp_ratio: float = 4.0
+    patch_size: int = 16
+    image_size: Union[Tuple[int, int], int] = 224
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    patch_dropout: float = 0.  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+    input_patchnorm: bool = False # whether to use dual patchnorm - would only apply the input layernorm on each patch, as post-layernorm already exist in original clip vit design
+    global_average_pool: bool = False  # whether to global average pool the last embedding layer, instead of using CLS token (https://arxiv.org/abs/2205.01580)
+    attentional_pool: bool = False # whether to use attentional pooler in the last embedding layer
+    n_queries: int = 256 # n_queries for attentional pooler
+    attn_pooler_heads: int = 8 # n heads for attentional_pooling
+    timm_model_name: str = None  # a valid model name overrides layers, width, patch_size
+    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
+    timm_pool: str = 'avg'  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+    timm_proj: str = 'linear'  # linear projection for timm model output ('linear', 'mlp', '')
+    timm_proj_bias: bool = False  # enable bias final projection
+    timm_drop: float = 0.  # head dropout
+    timm_drop_path: Optional[float] = None  # backbone stochastic depth
+    output_tokens: bool = False
+
+
+@dataclass
+class CLIPTextCfg:
+    context_length: int = 77
+    vocab_size: int = 49408
+    width: int = 512
+    heads: int = 8
+    layers: int = 12
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    hf_model_name: str = None
+    hf_tokenizer_name: str = None
+    hf_model_pretrained: bool = True
+    proj: str = 'mlp'
+    pooler_type: str = 'mean_pooler'
+    embed_cls: bool = False
+    pad_id: int = 0
+    output_tokens: bool = False
+
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == 'bf16':
+        cast_dtype = torch.bfloat16
+    elif precision == 'fp16':
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def _build_vision_tower(
+        embed_dim: int,
+        vision_cfg: CLIPVisionCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None
+):
+    if isinstance(vision_cfg, dict):
+        vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+    # OpenAI models are pretrained w/ QuickGELU but native nn.GELU is both faster and more
+    # memory efficient in recent PyTorch releases (>= 1.10).
+    # NOTE: timm models always use native GELU regardless of quick_gelu flag.
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+
+    if vision_cfg.timm_model_name:
+        visual = TimmModel(
+            vision_cfg.timm_model_name,
+            pretrained=vision_cfg.timm_model_pretrained,
+            pool=vision_cfg.timm_pool,
+            proj=vision_cfg.timm_proj,
+            proj_bias=vision_cfg.timm_proj_bias,
+            drop=vision_cfg.timm_drop,
+            drop_path=vision_cfg.timm_drop_path,
+            embed_dim=embed_dim,
+            image_size=vision_cfg.image_size,
+        )
+        act_layer = nn.GELU  # so that text transformer doesn't use QuickGELU w/ timm models
+    elif isinstance(vision_cfg.layers, (tuple, list)):
+        vision_heads = vision_cfg.width * 32 // vision_cfg.head_width
+        visual = ModifiedResNet(
+            layers=vision_cfg.layers,
+            output_dim=embed_dim,
+            heads=vision_heads,
+            image_size=vision_cfg.image_size,
+            width=vision_cfg.width,
+        )
+    else:
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+        visual = VisionTransformer(
+            image_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            width=vision_cfg.width,
+            layers=vision_cfg.layers,
+            heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            ls_init_value=vision_cfg.ls_init_value,
+            patch_dropout=vision_cfg.patch_dropout,
+            input_patchnorm=vision_cfg.input_patchnorm,
+            global_average_pool=vision_cfg.global_average_pool,
+            attentional_pool=vision_cfg.attentional_pool,
+            n_queries=vision_cfg.n_queries,
+            attn_pooler_heads=vision_cfg.attn_pooler_heads,
+            output_tokens=vision_cfg.output_tokens,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+
+    return visual
+
+
+def _build_text_tower(
+        embed_dim: int,
+        text_cfg: CLIPTextCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    if isinstance(text_cfg, dict):
+        text_cfg = CLIPTextCfg(**text_cfg)
+
+    if text_cfg.hf_model_name:
+        text = HFTextEncoder(
+            text_cfg.hf_model_name,
+            output_dim=embed_dim,
+            proj=text_cfg.proj,
+            pooler_type=text_cfg.pooler_type,
+            pretrained=text_cfg.hf_model_pretrained,
+            output_tokens=text_cfg.output_tokens,
+        )
+    else:
+        act_layer = QuickGELU if quick_gelu else nn.GELU
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+
+        text = TextTransformer(
+            context_length=text_cfg.context_length,
+            vocab_size=text_cfg.vocab_size,
+            width=text_cfg.width,
+            heads=text_cfg.heads,
+            layers=text_cfg.layers,
+            ls_init_value=text_cfg.ls_init_value,
+            output_dim=embed_dim,
+            embed_cls=text_cfg.embed_cls,
+            output_tokens=text_cfg.output_tokens,
+            pad_id=text_cfg.pad_id,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+    return text
+
+
+class CLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+
+        text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.transformer = text.transformer
+        self.vocab_size = text.vocab_size
+        self.token_embedding = text.token_embedding
+        self.positional_embedding = text.positional_embedding
+        self.ln_final = text.ln_final
+        self.text_projection = text.text_projection
+        self.register_buffer('attn_mask', text.attn_mask, persistent=False)
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.transformer.grad_checkpointing = enable
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        cast_dtype = self.transformer.get_cast_dtype()
+
+        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.to(cast_dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x, attn_mask=self.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+        return F.normalize(x, dim=-1) if normalize else x
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image, normalize=True)
+        text_features = self.encode_text(text, normalize=True)
+        if self.output_dict:
+            return {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return image_features, text_features, self.logit_scale.exp()
+
+
+class CustomTextCLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+        self.text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
+        self.text.lock(unlocked_layers, freeze_layer_norm)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.text.set_grad_checkpointing(enable)
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        features = self.text(text)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image, normalize=True)
+        text_features = self.encode_text(text, normalize=True)
+        if self.output_dict:
+            return {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return image_features, text_features, self.logit_scale.exp()
+
+
+def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
+    """Convert applicable model parameters to low-precision (bf16 or fp16)"""
+
+    def _convert_weights(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.to(dtype)
+            if l.bias is not None:
+                l.bias.data = l.bias.data.to(dtype)
+
+        if isinstance(l, (nn.MultiheadAttention, Attention)):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.to(dtype)
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.to(dtype)
+
+    model.apply(_convert_weights)
+
+
+convert_weights_to_fp16 = convert_weights_to_lp  # backwards compat
+
+
+# used to maintain checkpoint compatibility
+def convert_to_custom_text_state_dict(state_dict: dict):
+    if 'text_projection' in state_dict:
+        # old format state_dict, move text tower -> .text
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            if any(k.startswith(p) for p in (
+                'text_projection',
+                'positional_embedding',
+                'token_embedding',
+                'transformer',
+                'ln_final',
+            )):
+                k = 'text.' + k
+            new_state_dict[k] = v
+        return new_state_dict
+    return state_dict
+
+
+def build_model_from_openai_state_dict(
+        state_dict: dict,
+        quick_gelu=True,
+        cast_dtype=torch.float16,
+):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len(
+            [k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_size = vision_patch_size * grid_size
+    else:
+        counts: list = [
+            len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_size = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    vision_cfg = CLIPVisionCfg(
+        layers=vision_layers,
+        width=vision_width,
+        patch_size=vision_patch_size,
+        image_size=image_size,
+    )
+    text_cfg = CLIPTextCfg(
+        context_length=context_length,
+        vocab_size=vocab_size,
+        width=transformer_width,
+        heads=transformer_heads,
+        layers=transformer_layers,
+    )
+    model = CLIP(
+        embed_dim,
+        vision_cfg=vision_cfg,
+        text_cfg=text_cfg,
+        quick_gelu=quick_gelu,  # OpenAI models were trained with QuickGELU
+        cast_dtype=cast_dtype,
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        state_dict.pop(key, None)
+
+    convert_weights_to_fp16(model)  # OpenAI state dicts are partially converted to float16
+    model.load_state_dict(state_dict)
+    return model.eval()
+
+
+def trace_model(model, batch_size=256, device=torch.device('cpu')):
+    model.eval()
+    image_size = model.visual.image_size
+    example_images = torch.ones((batch_size, 3, image_size, image_size), device=device)
+    example_text = torch.zeros((batch_size, model.context_length), dtype=torch.int, device=device)
+    model = torch.jit.trace_module(
+        model,
+        inputs=dict(
+            forward=(example_images, example_text),
+            encode_text=(example_text,),
+            encode_image=(example_images,)
+        ))
+    model.visual.image_size = image_size
+    return model
+
+
+def resize_pos_embed(state_dict, model, interpolation: str = 'bicubic', antialias: bool = True):
+    # Rescale the grid of position embeddings when loading from state_dict
+    old_pos_embed = state_dict.get('visual.positional_embedding', None)
+    if old_pos_embed is None or not hasattr(model.visual, 'grid_size'):
+        return
+    grid_size = to_2tuple(model.visual.grid_size)
+    extra_tokens = 1  # FIXME detect different token configs (ie no class token, or more)
+    new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
+    if new_seq_len == old_pos_embed.shape[0]:
+        return
+
+    if extra_tokens:
+        pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
+    else:
+        pos_emb_tok, pos_emb_img = None, old_pos_embed
+    old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
+
+    logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size)
+    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
+    pos_emb_img = F.interpolate(
+        pos_emb_img,
+        size=grid_size,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
+    if pos_emb_tok is not None:
+        new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
+    else:
+        new_pos_embed = pos_emb_img
+    state_dict['visual.positional_embedding'] = new_pos_embed
+
+class AVCLIP(nn.Module):
+
+    def __init__(self, n_embd: int, afeat_extractor: OmegaConf, vfeat_extractor: OmegaConf,
+                 aproj: OmegaConf, vproj: OmegaConf, init_scale: float = 0.07, clamp_scale_min: float = 0.001,
+                 clamp_scale_max: float = 0.5, gather_for_loss: bool = False):
+        super().__init__()
+        self.output_dict = True
+        self.n_embd = n_embd
+
+        # loading audio and rgb towers
+        self.v_encoder = instantiate_from_config(vfeat_extractor)
+        self.a_encoder = instantiate_from_config(afeat_extractor)
+        # loading audio and rgb towers and projection layers to account for different feature dimensions
+        self.aproj = instantiate_from_config(aproj)
+        self.vproj = instantiate_from_config(vproj)
+
+        self.clamp_scale_min, self.clamp_scale_max = clamp_scale_min, clamp_scale_max
+        self.init_scale = init_scale
+        self.logit_scale = nn.Parameter(torch.ones([]) * self.init_scale)  # NOTE: exp(1/OpenCLIP)
+
+        self.gather_for_loss = gather_for_loss
+
+        # self.ln_final = text.ln_final  # perhaps only useful for transformer towers
+        # self.register_buffer('attn_mask', text.attn_mask, persistent=False)
+
+    def forward(self, vis: torch.Tensor, aud: torch.Tensor, alpha: float = 0.0, for_loop: bool = False,
+                world_size=1):
+        '''
+        Args:
+            vis (torch.Tensor): RGB frames (B, S, C, Tv, H, W)
+            aud (torch.Tensor): audio spectrograms (B, S, Ta, F)
+            alpha (float): linear interpolation coefficient for pseudo-targets (with 0.0 targets are 1-hot)
+            for_loop (bool): whether to process each segment in a for loop or all at once
+                             (speed-memory tradeoff)
+        Returns:
+            rgb_features (tuple(torch.Tensor)): local (B, S, D) and global (B, D) or None RGB features
+            audio_features (tuple(torch.Tensor)): local (B, S, D) and global (B, D) or None audio features
+            logit_scale (tuple(torch.Tensor)): local and global logit scales (1, )
+        '''
+        assert alpha == 0.0, f'alpha={alpha} not supported yet'
+        logit_scales = self.clamp_logit_scales()
+        vfeat, _, afeat, _ = self.encode_streams(vis, aud, for_loop, do_norm=True)
+
+        if world_size > 1 and self.gather_for_loss:  # gather all features
+            vfeat_all = torch.cat(torch.distributed.nn.all_gather(vfeat), dim=0)
+            afeat_all = torch.cat(torch.distributed.nn.all_gather(afeat), dim=0)
+        else:
+            vfeat_all = vfeat
+            afeat_all = afeat
+
+        loss_avc, _ = self.compute_loss(vfeat, afeat, vfeat_all.mT, afeat_all.mT, self.logit_scale, alpha=0)
+        out = {
+            'rgb_features': (vfeat, None), 'audio_features': (afeat, None),
+            'logit_scales': logit_scales,
+            'losses': {'segment_contrastive_loss': loss_avc},
+        }
+        return out
+
+    def compute_loss(self, vfeat, afeat, vfeat_all, afeat_all, scale, alpha=0.0, vfeat_m=None, afeat_m=None):
+        '''For Multi-level contrastive learning, the losses are made the same way for all levels'''
+        sim_v2a = vfeat @ afeat_all / scale
+        sim_a2v = afeat @ vfeat_all / scale
+        sim_v2a_t, sim_a2v_t = self._make_targets(sim_v2a, vfeat_all, afeat_all, scale, alpha, vfeat_m, afeat_m)
+        loss = self._loss(sim_v2a, sim_a2v, sim_v2a_t, sim_a2v_t)
+        return loss, (sim_v2a, sim_a2v)
+
+    @torch.no_grad()
+    def _make_targets(self, sim_v2a, vfeat_all, afeat_all, scale, alpha, vfeat_m, afeat_m):
+        # NOTE: for simplicity, we assume that sim_v2a.shape[0] == sim_a2v.shape[0]
+        # NOTE: sim_targets is not square (sim_v2a.shape is (bsize, bsize+Qsize) )
+        sim_targets = torch.eye(*sim_v2a.shape, device=sim_v2a.device, dtype=sim_v2a.dtype)
+        sim_v2a_targets = sim_targets
+        sim_a2v_targets = sim_targets
+        return sim_v2a_targets, sim_a2v_targets
+
+    def _loss(self, sim_v2a, sim_a2v, sim_v2a_targets, sim_a2v_targets):
+        loss_v2a = F.cross_entropy(sim_v2a, sim_v2a_targets)
+        loss_a2v = F.cross_entropy(sim_a2v, sim_a2v_targets)
+        return (loss_v2a + loss_a2v) / 2
+
+    def encode_streams(self, vis, aud, for_loop, do_norm=True):
+        # (B*S, D), (B, D) or None; because `flatten_to_2D = True`
+        flatten_to_2D = True
+        vfeat, _ = self.encode_stream(vis, self.v_encoder, self.vproj, do_norm, flatten_to_2D, for_loop)
+        afeat, _ = self.encode_stream(aud, self.a_encoder, self.aproj, do_norm, flatten_to_2D, for_loop)
+        return vfeat, None, afeat, None
+
+    def encode_stream(self, x, feat_extractor_fn, proj_fn, do_norm, flatten_to_2D, for_loop):
+        # x is (B, S, ...)
+        segment_x, _ = feat_extractor_fn(x, for_loop)  # segment_x: (B, S, D), global_x: (B, D)
+        if flatten_to_2D:
+            B, S, D = segment_x.shape
+            segment_x = segment_x.view(B*S, D)  # flatten batch and segment dims
+        segment_x = proj_fn(segment_x)
+        segment_x = F.normalize(segment_x, dim=-1) if do_norm else segment_x
+        # do_global is passed in to avoid computing global features when not needed (e.g. during eval)
+        return segment_x, None  # (B*S, D), (B, D) or None
+
+    def forward_for_logging(self, vis, aud, for_momentum=False, for_loop=False, do_norm=True):
+        '''
+        Runs the forward pass but keeps certain tensors in memory for logging purposes, ie code duplication.
+        NOTE: to be used outside of this module, most likely during logging
+
+        Args:
+            vis (torch.Tensor): RGB frames (B, S, C, Tv, H, W)
+            aud (torch.Tensor): audio spectrograms (B, S, Ta, F)
+        '''
+        flatten_to_2D = True
+        out = dict()
+
+        # factorizing self.encode_streams into encode_visual/encode_audio to avoid unnecessary computations
+        # (B*S, D), (B, D) or None;
+        # vfeat, _ = self.encode_stream(vis, self.v_encoder, self.vproj, do_norm, flatten_to_2D, for_loop)
+        # afeat, _ = self.encode_stream(aud, self.a_encoder, self.aproj, do_norm, flatten_to_2D, for_loop)
+        vfeat, _, afeat, _ = self.encode_streams(vis, aud, for_loop, do_norm)
+        # cache features (for 0-shot evaluation)
+        out['segment_vfeat'] = vfeat.clone()
+        out['segment_afeat'] = afeat.clone()
+        # and similiarity matrices (for visualization) (B*S, B*S)
+        out['segment_sim_v2a'] = out['segment_vfeat'] @ out['segment_afeat'].mT / self.logit_scale
+        out['segment_sim_a2v'] = out['segment_afeat'] @ out['segment_vfeat'].mT / self.logit_scale
+        # self
+        out['segment_sim_v2v'] = out['segment_vfeat'] @ out['segment_vfeat'].mT / self.logit_scale
+        out['segment_sim_a2a'] = out['segment_afeat'] @ out['segment_afeat'].mT / self.logit_scale
+
+        # compute losses
+        loss, _ = self.compute_loss(vfeat, afeat, vfeat.mT, afeat.mT, self.logit_scale)
+        out['segment_contrastive_loss'] = loss
+        return out
+
+    @torch.no_grad()
+    def clamp_logit_scales(self):
+        self.logit_scale.clamp_(self.clamp_scale_min, self.clamp_scale_max)
+        return (self.logit_scale, None)
+
+
+class MultilevelMoCoCLIP(nn.Module):
+
+    def __init__(self, n_embd: int, queue_size: int, momentum: float,
+                 afeat_extractor: OmegaConf, vfeat_extractor: OmegaConf, aproj: OmegaConf, vproj: OmegaConf,
+                 init_scale: float = 0.07, clamp_scale_min: float = 0.001, clamp_scale_max: float = 0.5):
+        super().__init__()
+        self.output_dict = True
+        self.n_embd = n_embd
+        self.momentum = momentum
+        self.to_add_global_repr = afeat_extractor.params.add_global_repr
+
+        # loading audio and rgb towers
+        self.v_encoder = instantiate_from_config(vfeat_extractor)
+        self.a_encoder = instantiate_from_config(afeat_extractor)
+        # loading audio and rgb towers and projection layers to account for different feature dimensions
+        self.segment_aproj = instantiate_from_config(aproj)
+        self.segment_vproj = instantiate_from_config(vproj)
+        self.global_aproj = instantiate_from_config(aproj) if self.to_add_global_repr else None
+        self.global_vproj = instantiate_from_config(vproj) if self.to_add_global_repr else None
+
+        self.clamp_scale_min, self.clamp_scale_max = clamp_scale_min, clamp_scale_max
+        self.init_scale = init_scale
+        self.segment_logit_scale = nn.Parameter(torch.ones([]) * self.init_scale)  # NOTE: exp(1/OpenCLIP)
+        self.global_logit_scale = nn.Parameter(torch.ones([]) * self.init_scale) if self.to_add_global_repr else None
+
+        # create momentum models
+        self.v_encoder_m = instantiate_from_config(vfeat_extractor)
+        self.a_encoder_m = instantiate_from_config(afeat_extractor)
+        self.segment_aproj_m = instantiate_from_config(aproj)
+        self.segment_vproj_m = instantiate_from_config(vproj)
+        self.global_aproj_m = instantiate_from_config(aproj) if self.to_add_global_repr else None
+        self.global_vproj_m = instantiate_from_config(vproj) if self.to_add_global_repr else None
+
+        self.model_pairs = [
+            [self.v_encoder, self.v_encoder_m], [self.segment_vproj, self.segment_vproj_m],
+            [self.a_encoder, self.a_encoder_m], [self.segment_aproj, self.segment_aproj_m],
+        ]
+        if self.to_add_global_repr:
+            self.model_pairs += [
+                [self.global_aproj, self.global_aproj_m], [self.global_vproj, self.global_vproj_m],
+            ]
+
+        self.copy_params()
+
+        self.segment_queue_size = queue_size * afeat_extractor.params.max_segments  # scaled by # of segments
+        self.global_queue_size = queue_size if self.to_add_global_repr else None
+        self.init_Qs(self.segment_queue_size, self.global_queue_size, self.n_embd)
+
+        # self.ln_final = text.ln_final  # perhaps only useful for transformer towers
+        # self.register_buffer('attn_mask', text.attn_mask, persistent=False)
+
+    def forward(self, vis: torch.Tensor, aud: torch.Tensor, alpha: float = 0.0, for_loop: bool = False,
+                world_size=None):
+        '''
+        Args:
+            vis (torch.Tensor): RGB frames (B, S, C, Tv, H, W)
+            aud (torch.Tensor): audio spectrograms (B, S, Ta, F)
+            alpha (float): linear interpolation coefficient for pseudo-targets (with 0.0 targets are 1-hot)
+            for_loop (bool): whether to process each segment in a for loop or all at once
+                             (speed-memory tradeoff)
+        Returns:
+            rgb_features (tuple(torch.Tensor)): local (B, S, D) and global (B, D) or None RGB features
+            audio_features (tuple(torch.Tensor)): local (B, S, D) and global (B, D) or None audio features
+            logit_scale (tuple(torch.Tensor)): local and global logit scales (1, )
+        '''
+        logit_scales = self.clamp_logit_scales()
+        to_add_global_repr = self.to_add_global_repr  # for readability only
+
+        feats = self.encode_streams(vis, aud, for_momentum=False, for_loop=for_loop, do_norm=True)
+        segment_vfeat, global_vfeat, segment_afeat, global_afeat = feats
+
+        # get momentum features
+        with torch.no_grad():
+            if self.training:
+                self._momentum_update()
+            feats_m = self.encode_streams(vis, aud, for_momentum=True, for_loop=for_loop, do_norm=True)
+            segment_vfeat_m, global_vfeat_m, segment_afeat_m, global_afeat_m = feats_m
+
+            # cat with queue to extend the list of negatives
+            segment_vfeat_all = torch.cat([segment_vfeat_m.t(), self.segment_v_queue.clone().detach()], dim=1)
+            segment_afeat_all = torch.cat([segment_afeat_m.t(), self.segment_a_queue.clone().detach()], dim=1)
+            if to_add_global_repr:
+                global_vfeat_all = torch.cat([global_vfeat_m.t(), self.global_v_queue.clone().detach()], dim=1)
+                global_afeat_all = torch.cat([global_afeat_m.t(), self.global_a_queue.clone().detach()], dim=1)
+
+        segment_loss_avc, _ = self.compute_loss(segment_vfeat, segment_afeat, segment_vfeat_all,
+                                                segment_afeat_all, self.segment_logit_scale,
+                                                alpha, segment_vfeat_m, segment_afeat_m)
+
+        global_loss_avc = None
+        if to_add_global_repr:
+            global_loss_avc, _ = self.compute_loss(global_vfeat, global_afeat, global_vfeat_all,
+                                                   global_afeat_all, self.global_logit_scale,
+                                                   alpha, global_vfeat_m, global_afeat_m)
+
+        if self.training:
+            self._multilevel_dequeue_and_enqueue(segment_vfeat_m, segment_afeat_m, global_vfeat_m, global_afeat_m)
+        else:
+            raise Exception('This module is used only during training. Use model.something instead.')
+
+        out = {
+            'rgb_features': (segment_vfeat, global_vfeat),
+            'audio_features': (segment_afeat, global_afeat), 'logit_scales': logit_scales,
+            'losses': {'segment_contrastive_loss': segment_loss_avc},
+        }
+        if global_loss_avc is not None:
+            out['losses']['global_contrastive_loss'] = global_loss_avc
+        return out
+
+    def compute_loss(self, vfeat, afeat, vfeat_all, afeat_all, scale, alpha=0.0, vfeat_m=None, afeat_m=None):
+        '''For Multi-level contrastive learning, the losses are made the same way for all levels'''
+        sim_v2a = vfeat @ afeat_all / scale
+        sim_a2v = afeat @ vfeat_all / scale
+        sim_v2a_t, sim_a2v_t = self._make_targets(sim_v2a, vfeat_all, afeat_all, scale, alpha, vfeat_m, afeat_m)
+        loss = self._loss(sim_v2a, sim_a2v, sim_v2a_t, sim_a2v_t)
+        return loss, (sim_v2a, sim_a2v)
+
+    @torch.no_grad()
+    def _make_targets(self, sim_v2a, vfeat_all, afeat_all, scale, alpha, vfeat_m, afeat_m):
+        # NOTE: for simplicity, we assume that sim_v2a.shape[0] == sim_a2v.shape[0]
+        # NOTE: sim_targets is not square (sim_v2a.shape is (bsize, bsize+Qsize) )
+        sim_targets = torch.eye(*sim_v2a.shape, device=sim_v2a.device, dtype=sim_v2a.dtype)
+        # the ALBEF alpha trick
+        if alpha > 0.0:
+            sim_v2a_m = vfeat_m @ afeat_all / scale
+            sim_a2v_m = afeat_m @ vfeat_all / scale
+            sim_v2a_targets = alpha * F.softmax(sim_v2a_m, dim=1) + (1 - alpha) * sim_targets
+            sim_a2v_targets = alpha * F.softmax(sim_a2v_m, dim=1) + (1 - alpha) * sim_targets
+        else:
+            sim_v2a_targets = sim_targets
+            sim_a2v_targets = sim_targets
+        return sim_v2a_targets, sim_a2v_targets
+
+    def _loss(self, sim_v2a, sim_a2v, sim_v2a_targets, sim_a2v_targets):
+        loss_v2a = F.cross_entropy(sim_v2a, sim_v2a_targets)
+        loss_a2v = F.cross_entropy(sim_a2v, sim_a2v_targets)
+        return (loss_v2a + loss_a2v) / 2
+
+    def encode_streams(self, vis, aud, for_momentum, for_loop, do_norm=True):
+        # (B*S, D), (B, D) or None; because `flatten_to_2D = True`
+        flatten_to_2D = True
+        segment_vfeat, global_vfeat = self.encode_visual(vis, for_momentum, self.to_add_global_repr, do_norm,
+                                                         for_loop=for_loop, flatten_to_2D=flatten_to_2D)
+        segment_afeat, global_afeat = self.encode_audio(aud, for_momentum, self.to_add_global_repr, do_norm,
+                                                        for_loop=for_loop, flatten_to_2D=flatten_to_2D)
+        return segment_vfeat, global_vfeat, segment_afeat, global_afeat
+
+    def encode_audio(self, x, for_momentum: bool = False, do_global: bool = True, do_norm: bool = True,
+                     flatten_to_2D=True, for_loop=False):
+        # define callables
+        encode_fn = self.a_encoder_m if for_momentum else self.a_encoder
+        segment_proj_fn = self.segment_aproj_m if for_momentum else self.segment_aproj
+        global_proj_fn = self.global_aproj_m if for_momentum else self.global_aproj
+        # do the encoding
+        return self.encode_stream(x, encode_fn, segment_proj_fn, global_proj_fn, do_global, do_norm,
+                                  flatten_to_2D, for_loop)
+
+    def encode_visual(self, x, for_momentum: bool = False, do_global: bool = True, do_norm: bool = True,
+                      flatten_to_2D=True, for_loop=False):
+        # define callables
+        encode_fn = self.v_encoder_m if for_momentum else self.v_encoder
+        segment_proj_fn = self.segment_vproj_m if for_momentum else self.segment_vproj
+        global_proj_fn = self.global_vproj_m if for_momentum else self.global_vproj
+        # do the encoding
+        return self.encode_stream(x, encode_fn, segment_proj_fn, global_proj_fn, do_global, do_norm,
+                                  flatten_to_2D, for_loop)
+
+    def encode_stream(self, x, feat_extractor_fn, segment_proj_fn, global_proj_fn, do_global, do_norm,
+                      flatten_to_2D, for_loop):
+        # x is (B, S, ...)
+        segment_x, global_x = feat_extractor_fn(x, for_loop)  # segment_x: (B, S, D), global_x: (B, D)
+        if flatten_to_2D:
+            B, S, D = segment_x.shape
+            segment_x = segment_x.view(B*S, D)  # flatten batch and segment dims
+        segment_x = segment_proj_fn(segment_x)
+        segment_x = F.normalize(segment_x, dim=-1) if do_norm else segment_x
+        # do_global is passed in to avoid computing global features when not needed (e.g. during eval)
+        if do_global and self.to_add_global_repr:
+            global_x = global_proj_fn(global_x)
+            global_x = F.normalize(global_x, dim=-1) if do_norm else global_x
+        return segment_x, global_x  # (B*S, D), (B, D) or None
+
+    def forward_for_logging(self, vis, aud, for_momentum=False, for_loop=False, do_norm=True):
+        '''
+        Runs the forward pass but keeps certain tensors in memory for logging purposes, ie code duplication.
+        NOTE: to be used outside of this module, most likely during logging
+
+        Args:
+            vis (torch.Tensor): RGB frames (B, S, C, Tv, H, W)
+            aud (torch.Tensor): audio spectrograms (B, S, Ta, F)
+        '''
+        flatten_to_2D = True
+
+        out = dict()
+
+        # factorizing self.encode_streams into encode_visual/encode_audio to avoid unnecessary computations
+        # (B*S, D), (B, D) or None;
+        segment_vfeat, global_vfeat = self.encode_visual(vis, for_momentum, self.to_add_global_repr, do_norm,
+                                                         flatten_to_2D, for_loop)
+        segment_afeat, global_afeat = self.encode_audio(aud, for_momentum, self.to_add_global_repr, do_norm,
+                                                        flatten_to_2D, for_loop)
+        # cache features (for 0-shot evaluation)
+        out['segment_vfeat'] = segment_vfeat.clone()
+        out['segment_afeat'] = segment_afeat.clone()
+        # and similiarity matrices (for visualization) (B*S, B*S)
+        out['segment_sim_v2a'] = out['segment_vfeat'] @ out['segment_afeat'].mT / self.segment_logit_scale
+        out['segment_sim_a2v'] = out['segment_afeat'] @ out['segment_vfeat'].mT / self.segment_logit_scale
+        # self
+        out['segment_sim_v2v'] = out['segment_vfeat'] @ out['segment_vfeat'].mT / self.segment_logit_scale
+        out['segment_sim_a2a'] = out['segment_afeat'] @ out['segment_afeat'].mT / self.segment_logit_scale
+
+        # compute losses
+        segment_loss, _ = self.compute_loss(segment_vfeat, segment_afeat, segment_vfeat.mT, segment_afeat.mT,
+                                            self.segment_logit_scale)
+        out['segment_contrastive_loss'] = segment_loss
+        if self.to_add_global_repr:
+            global_loss, _ = self.compute_loss(global_vfeat, global_afeat, global_vfeat.mT, global_afeat.mT,
+                                               self.global_logit_scale)
+            out['global_contrastive_loss'] = global_loss
+
+        return out
+
+
+    @torch.no_grad()
+    def clamp_logit_scales(self):
+        self.segment_logit_scale.clamp_(self.clamp_scale_min, self.clamp_scale_max)
+        if self.to_add_global_repr:
+            self.global_logit_scale.clamp_(self.clamp_scale_min, self.clamp_scale_max)
+        return (self.segment_logit_scale, self.global_logit_scale)
+
+    @torch.no_grad()
+    def copy_params(self):
+        for model_pair in self.model_pairs:
+            for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
+                param_m.data.copy_(param.data)  # initialize
+                param_m.requires_grad = False  # not update by gradient
+
+    @torch.no_grad()
+    def _momentum_update(self):
+        for model_pair in self.model_pairs:
+            for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
+                param_m.data = param_m.data * self.momentum + param.data * (1. - self.momentum)
+
+    @torch.no_grad()
+    def _multilevel_dequeue_and_enqueue(self, segment_vfeat_m, segment_afeat_m, global_vfeat_m, global_afeat_m):
+        if self.segment_queue_size > 0:
+            self._dequeue_and_enqueue(segment_vfeat_m, segment_afeat_m, 'segment_')
+        if self.to_add_global_repr and self.global_queue_size > 0:
+            self._dequeue_and_enqueue(global_vfeat_m, global_afeat_m, 'global_')
+
+    @torch.no_grad()
+    def _dequeue_and_enqueue(self, vfeat, afeat, level_prefix_: str):
+        # gather keys before updating queue
+        if torch.distributed.is_initialized():
+            vfeats = concat_all_gather(vfeat)
+            afeats = concat_all_gather(afeat)
+        else:
+            vfeats = vfeat
+            afeats = afeat
+
+        batch_size = vfeats.shape[0]
+        queue_size = getattr(self, level_prefix_ + 'queue_size')
+
+        # same as `ptr = int(self.segment_queue_ptr)` but allows accessing the attribute by string
+        ptr = int(getattr(self, level_prefix_ + 'queue_ptr'))
+        assert queue_size % batch_size == 0, f'For simplicity: {queue_size} % {batch_size} == 0'
+
+        # replace the keys at ptr (dequeue and enqueue)
+        getattr(self, level_prefix_ + 'v_queue')[:, ptr:ptr + batch_size] = vfeats.T
+        getattr(self, level_prefix_ + 'a_queue')[:, ptr:ptr + batch_size] = afeats.T
+        ptr = (ptr + batch_size) % queue_size  # move pointer
+
+        getattr(self, level_prefix_ + 'queue_ptr')[0] = ptr
+
+    def init_Qs(self, segment_queue_size: int, global_queue_size: int, n_embd: int):
+        # create the queues; TODO: flip the dimensions, yikes!
+        self.register_buffer('segment_v_queue', torch.randn(n_embd, segment_queue_size))
+        self.register_buffer('segment_a_queue', torch.randn(n_embd, segment_queue_size))
+        self.register_buffer('segment_queue_ptr', torch.zeros(1, dtype=torch.long))
+        self.segment_v_queue = nn.functional.normalize(self.segment_v_queue, dim=0)
+        self.segment_a_queue = nn.functional.normalize(self.segment_a_queue, dim=0)
+        if self.to_add_global_repr:
+            self.register_buffer('global_v_queue', torch.randn(n_embd, global_queue_size))
+            self.register_buffer('global_a_queue', torch.randn(n_embd, global_queue_size))
+            self.register_buffer('global_queue_ptr', torch.zeros(1, dtype=torch.long))
+            self.global_v_queue = nn.functional.normalize(self.global_v_queue, dim=0)
+            self.global_a_queue = nn.functional.normalize(self.global_a_queue, dim=0)
+
+@torch.no_grad()
+def concat_all_gather(tensor):
+    """
+    Performs all_gather operation on the provided tensors.
+    *** Warning ***: torch.distributed.all_gather has no gradient.
+    """
+    tensors_gather = [torch.ones_like(tensor) for _ in range(torch.distributed.get_world_size())]
+    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
+
+    output = torch.cat(tensors_gather, dim=0)
+    return output

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/RN101-quickgelu.json

@@ -0,0 +1,22 @@
+{
+    "embed_dim": 512,
+    "quick_gelu": true,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            23,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/RN101.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            23,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/RN50-quickgelu.json

@@ -0,0 +1,22 @@
+{
+    "embed_dim": 1024,
+    "quick_gelu": true,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            6,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/RN50.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            6,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/RN50x16.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 384,
+        "layers": [
+            6,
+            8,
+            18,
+            8
+        ],
+        "width": 96,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/RN50x4.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 288,
+        "layers": [
+            4,
+            6,
+            10,
+            6
+        ],
+        "width": 80,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/RN50x64.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 448,
+        "layers": [
+            3,
+            15,
+            36,
+            10
+        ],
+        "width": 128,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/ViT-B-16-plus-240.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 240,
+        "layers": 12,
+        "width": 896,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/ViT-B-16-plus.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 896,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/ViT-B-16.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/ViT-B-32-plus-256.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 256,
+        "layers": 12,
+        "width": 896,
+        "patch_size": 32
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/ViT-B-32-quickgelu.json

@@ -0,0 +1,17 @@
+{
+    "embed_dim": 512,
+    "quick_gelu": true,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "patch_size": 32
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/ViT-B-32.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "patch_size": 32
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/ViT-H-14.json

@@ -0,0 +1,17 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 32,
+        "width": 1280,
+        "head_width": 80,
+        "patch_size": 14
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 24
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/ViT-H-16.json

@@ -0,0 +1,17 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 32,
+        "width": 1280,
+        "head_width": 80,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 24
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/ViT-L-14-280.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 280,
+        "layers": 24,
+        "width": 1024,
+        "patch_size": 14
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/ViT-L-14-336.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 336,
+        "layers": 24,
+        "width": 1024,
+        "patch_size": 14
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/ViT-L-14.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 24,
+        "width": 1024,
+        "patch_size": 14
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/ViT-L-16-320.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 320,
+        "layers": 24,
+        "width": 1024,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    }
+}

modules/model/modules/feat_extractors/train_clip_src/open_clip/model_configs/ViT-L-16.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 24,
+        "width": 1024,
+        "patch_size": 16
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12
+    }
+}