Add the Github library to use custom modules

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+DenseAV/samples/puppies.mp4 filter=lfs diff=lfs merge=lfs -text

DenseAV/.gitignore

@@ -0,0 +1,5 @@
+# Created by .ignore support plugin (hsz.mobi)
+results/attention/*
+results/features/*
+
+.env

DenseAV/LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) Mark Hamilton. All rights reserved.
+
+Permission is hereby granted, free of charge, to any person obtaining a
+copy of this software and associated documentation files (the
+"Software"), to deal in the Software without restriction, including
+without limitation the rights to use, copy, modify, merge, publish,
+distribute, sublicense, and/or sell copies of the Software, and to
+permit persons to whom the Software is furnished to do so, subject to
+the following conditions:
+
+The above copyright notice and this permission notice shall be included
+in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
+LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

DenseAV/README.md

@@ -0,0 +1,172 @@
+# Separating the "Chirp" from the "Chat": Self-supervised Visual Grounding of Sound and Language
+###  CVPR 2024
+
+
+[![Website](https://img.shields.io/badge/DenseAV-%F0%9F%8C%90Website-purple?style=flat)](https://aka.ms/denseav) [![arXiv](https://img.shields.io/badge/arXiv-2406.05629-b31b1b.svg)](https://arxiv.org/abs/2406.05629) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mhamilton723/DenseAV/blob/main/demo.ipynb)
+
+[![Huggingface](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-DenseAV-orange)](https://huggingface.co/spaces/mhamilton723/DenseAV) 
+
+[//]: # ([![Huggingface](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Paper%20Page-orange)](https://huggingface.co/papers/2403.10516))
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/separating-the-chirp-from-the-chat-self/speech-prompted-semantic-segmentation-on)](https://paperswithcode.com/sota/speech-prompted-semantic-segmentation-on?p=separating-the-chirp-from-the-chat-self)
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/separating-the-chirp-from-the-chat-self/sound-prompted-semantic-segmentation-on)](https://paperswithcode.com/sota/sound-prompted-semantic-segmentation-on?p=separating-the-chirp-from-the-chat-self)
+
+
+[Mark Hamilton](https://mhamilton.net/),
+[Andrew Zisserman](https://www.robots.ox.ac.uk/~az/),
+[John R. Hershey](https://research.google/people/john-hershey/),
+[William T. Freeman](https://billf.mit.edu/about/bio)
+
+![DenseAV Overview Graphic](https://mhamilton.net/images/hero_fig_black.jpg)
+
+**TL;DR**:Our model, DenseAV, learns the meaning of words and the location of sounds (visual grounding) without supervision or text.
+
+https://github.com/mhamilton723/DenseAV/assets/6456637/ba908ab5-9618-42f9-8d7a-30ecb009091f
+
+
+## Contents
+<!--ts-->
+   * [Install](#install)
+   * [Model Zoo](#model-zoo)
+   * [Getting Datasets](#getting-atasets)
+   * [Evaluate Models](#evaluate-models)
+   * [Train a Model](#train-model)
+   * [Local Gradio Demo](#local-gradio-demo)
+   * [Coming Soon](coming-soon)
+   * [Citation](#citation)
+   * [Contact](#contact)
+<!--te-->
+
+## Install
+
+To use DenseAV locally clone the repository:
+
+```shell script
+git clone https://github.com/mhamilton723/DenseAV.git
+cd DenseAV
+pip install -e .
+```
+
+
+## Model Zoo
+
+To see examples of pretrained model usage please see our [Collab notebook](https://colab.research.google.com/github/mhamilton723/DenseAV/blob/main/demo.ipynb). We currently supply the following pretrained models:
+
+| Model Name                    | Checkpoint                                                                                                                       | Torch Hub Repository | Torch Hub Name     |
+|-------------------------------|----------------------------------------------------------------------------------------------------------------------------------|----------------------|--------------------|
+| Sound                         | [Download](https://marhamilresearch4.blob.core.windows.net/denseav-public/hub/denseav_sound.ckpt) | mhamilton723/DenseAV | sound              |
+| Language                      | [Download](https://marhamilresearch4.blob.core.windows.net/denseav-public/hub/denseav_language.ckpt) | mhamilton723/DenseAV | language           |
+| Sound + Language (Two Headed) | [Download](https://marhamilresearch4.blob.core.windows.net/denseav-public/hub/denseav_2head.ckpt)   | mhamilton723/DenseAV | sound_and_language |
+
+For example, to load the model trained on both sound and language:
+
+```python
+model = torch.hub.load("mhamilton723/DenseAV", 'sound_and_language')
+```
+
+### Load from HuggingFace
+
+```python
+from denseav.train import LitAVAligner
+
+model1 = LitAVAligner.from_pretrained("mhamilton723/DenseAV-sound")
+model2 = LitAVAligner.from_pretrained("mhamilton723/DenseAV-language")
+model3 = LitAVAligner.from_pretrained("mhamilton723/DenseAV-sound-language")
+```
+
+
+## Getting Datasets
+
+Our code assumes that all data lives in a common directory on your system, in these examples we use `/path/to/your/data`. Our code will often reference this directory as the `data_root`
+
+### Speech and Sound Prompted ADE20K
+
+To download our new Speech and Sound prompted ADE20K Dataset:
+
+```bash
+cd /path/to/your/data
+wget https://marhamilresearch4.blob.core.windows.net/denseav-public/datasets/ADE20KSoundPrompted.zip
+unzip ADE20KSoundPrompted.zip
+wget https://marhamilresearch4.blob.core.windows.net/denseav-public/datasets/ADE20KSpeechPrompted.zip
+unzip ADE20KSpeechPrompted.zip
+```
+
+### Places Audio
+
+First download the places audio dataset from its [original source](https://groups.csail.mit.edu/sls/downloads/placesaudio/downloads.cgi).
+
+To run the code the data will need to be processed to be of the form:
+
+```
+[Instructions coming soon]
+```
+
+### Audioset
+
+Because of copyright issues we cannot make [Audioset](https://research.google.com/audioset/dataset/index.html) easily availible to download.
+First download this dataset through appropriate means. [This other project](https://github.com/ktonal/audioset-downloader) appears to make this simple.
+
+To run the code the data will need to be processed to be of the form:
+
+```
+[Instructions coming soon]
+```
+
+
+## Evaluate Models
+
+To evaluate a trained model first clone the repository for
+[local development](#local-development). Then run
+
+```shell
+cd featup
+python evaluate.py
+```
+
+After evaluation, see the results in tensorboard's hparams tab. 
+
+```shell
+cd ../logs/evaluate
+tensorboard --logdir .
+```
+
+Then visit [https://localhost:6006](https://localhost:6006) and click on hparams to browse results. We report "advanced" speech metrics and "basic" sound metrics in our paper.
+
+
+## Train a Model
+
+```shell
+cd denseav
+python train.py
+```
+
+## Local Gradio Demo
+
+To run our [HuggingFace Spaces hosted DenseAV demo](https://huggingface.co/spaces/mhamilton723/FeatUp) locally first install DenseAV for local development. Then  run:
+
+```shell
+python gradio_app.py
+```
+
+Wait a few seconds for the demo to spin up, then navigate to [http://localhost:7860/](http://localhost:7860/) to view the demo.
+
+
+## Coming Soon:
+
+- Bigger models!
+
+## Citation
+
+```
+@misc{hamilton2024separating,
+      title={Separating the "Chirp" from the "Chat": Self-supervised Visual Grounding of Sound and Language}, 
+      author={Mark Hamilton and Andrew Zisserman and John R. Hershey and William T. Freeman},
+      year={2024},
+      eprint={2406.05629},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```
+
+## Contact
+
+For feedback, questions, or press inquiries please contact [Mark Hamilton](mailto:[email protected])

DenseAV/denseav/aggregators.py

@@ -0,0 +1,517 @@
+from abc import abstractmethod
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+
+from denseav.constants import *
+
+
+@torch.jit.script
+def masked_mean(x: torch.Tensor, mask: torch.Tensor, dim: int):
+    mask = mask.to(x)
+    return (x * mask).sum(dim, keepdim=True) / mask.sum(dim, keepdim=True).clamp_min(.001)
+
+
+@torch.jit.script
+def masked_max(x: torch.Tensor, mask: torch.Tensor, dim: int):
+    mask = mask.to(torch.bool)
+    eps = 1e7
+    # eps = torch.finfo(x.dtype).max
+    return (x - (~mask) * eps).max(dim, keepdim=True).values
+
+
+def masked_lse(x: torch.Tensor, mask: torch.Tensor, dim: int, temp):
+    x = x.to(torch.float32)
+    mask = mask.to(torch.float32)
+    x_masked = (x - (1 - mask) * torch.finfo(x.dtype).max)
+    return (torch.logsumexp(x_masked * temp, dim, keepdim=True) - torch.log(mask.sum(dim, keepdim=True))) / temp
+
+
+class BaseAggregator(torch.nn.Module):
+
+    def __init__(self, nonneg_sim, mask_silence, num_heads, head_agg, use_cls):
+        super().__init__()
+
+        self.nonneg_sim = nonneg_sim
+        self.mask_silence = mask_silence
+        self.num_heads = num_heads
+        self.head_agg = head_agg
+        self.use_cls = use_cls
+
+    @abstractmethod
+    def _agg_sim(self, sim, mask):
+        pass
+
+    def prepare_sims(self, sim, mask, agg_sim, agg_heads):
+        sim_size = sim.shape
+        assert len(mask.shape) == 2
+        assert len(sim_size) in {6, 7}, f"sim has wrong number of dimensions: {sim.shape}"
+        pairwise = len(sim_size) == 6
+
+        if self.mask_silence:
+            mask = mask
+        else:
+            mask = torch.ones_like(mask)
+
+        if self.nonneg_sim:
+            sim = sim.clamp_min(0)
+
+        if pairwise:
+            head_dim = 1
+        else:
+            head_dim = 2
+
+        if self.head_agg == "max_elementwise" and agg_heads:
+            sim = sim.max(head_dim, keepdim=True).values
+
+        if agg_sim:
+            sim = self._agg_sim(sim, mask)
+
+        if agg_heads:
+            if self.head_agg == "sum" or self.head_agg == "max_elementwise":
+                sim = sim.sum(head_dim)
+            elif self.head_agg == "max":
+                sim = sim.max(head_dim).values
+            else:
+                raise ValueError(f"Unknown head_agg: {self.head_agg}")
+
+        return sim
+
+    def _get_full_sims(self, preds, raw, agg_sim, agg_heads):
+        if agg_sim or agg_heads or raw:
+            assert (agg_sim or agg_heads) != raw, "Cannot have raw on at the same time as agg_sim or agg_heads"
+
+        audio_feats = preds[AUDIO_FEATS]
+        audio_mask = preds[AUDIO_MASK]
+        image_feats = preds[IMAGE_FEATS]
+
+        b1, c2, f, t1 = audio_feats.shape
+        b2, t2 = audio_mask.shape
+        d, c1, h, w = image_feats.shape
+        assert b1 == b2 and c1 == c2 and t1 == t2
+        assert c1 % self.num_heads == 0
+        new_c = c1 // self.num_heads
+        audio_feats = audio_feats.reshape(b1, self.num_heads, new_c, f, t1)
+        image_feats = image_feats.reshape(d, self.num_heads, new_c, h, w)
+        raw_sims = torch.einsum(
+            "akcft,vkchw->avkhwft",
+            audio_feats.to(torch.float32),
+            image_feats.to(torch.float32))
+
+        if self.use_cls:
+            audio_cls = preds[AUDIO_CLS].reshape(b1, self.num_heads, new_c)
+            image_cls = preds[IMAGE_CLS].reshape(d, self.num_heads, new_c)
+            cls_sims = torch.einsum(
+                "akc,vkc->avk",
+                audio_cls.to(torch.float32),
+                image_cls.to(torch.float32))
+            raw_sims += cls_sims.reshape(b1, d, self.num_heads, 1, 1, 1, 1)
+
+        if raw:
+            return raw_sims
+        else:
+            return self.prepare_sims(raw_sims, audio_mask, agg_sim, agg_heads)
+
+    def get_pairwise_sims(self, preds, raw, agg_sim, agg_heads):
+        if agg_sim or agg_heads or raw:
+            assert (agg_sim or agg_heads) != raw, "Cannot have raw on at the same time as agg_sim or agg_heads"
+
+        audio_feats = preds[AUDIO_FEATS]
+        audio_mask = preds[AUDIO_MASK]
+        image_feats = preds[IMAGE_FEATS]
+
+        a1, c1, f, t1 = audio_feats.shape
+        a2, t2 = audio_mask.shape
+
+        assert c1 % self.num_heads == 0
+        new_c = c1 // self.num_heads
+        audio_feats = audio_feats.reshape(a1, self.num_heads, new_c, f, t1)
+
+        if len(image_feats.shape) == 5:
+            print("Using similarity for video, should only be called during plotting")
+            v, vt, c2, h, w = image_feats.shape
+            image_feats = image_feats.reshape(v, vt, self.num_heads, new_c, h, w)
+            raw_sims = torch.einsum(
+                "bkcft,bskchw,bt->bskhwft",
+                audio_feats.to(torch.float32),
+                image_feats.to(torch.float32),
+                audio_mask.to(torch.float32))
+
+            if self.use_cls:
+                audio_cls = preds[AUDIO_CLS].reshape(v, self.num_heads, new_c)
+                image_cls = preds[IMAGE_CLS].reshape(v, vt, self.num_heads, new_c)
+                cls_sims = torch.einsum(
+                    "bkc,bskc->bsk",
+                    audio_cls.to(torch.float32),
+                    image_cls.to(torch.float32))
+                raw_sims += cls_sims.reshape(v, vt, self.num_heads, 1, 1, 1, 1)
+
+
+        elif len(image_feats.shape) == 4:
+            v, c2, h, w = image_feats.shape
+            image_feats = image_feats.reshape(v, self.num_heads, new_c, h, w)
+            raw_sims = torch.einsum(
+                "bkcft,bkchw,bt->bkhwft",
+                audio_feats.to(torch.float32),
+                image_feats.to(torch.float32),
+                audio_mask.to(torch.float32))
+
+            if self.use_cls:
+                audio_cls = preds[AUDIO_CLS].reshape(v, self.num_heads, new_c)
+                image_cls = preds[IMAGE_CLS].reshape(v, self.num_heads, new_c)
+                cls_sims = torch.einsum(
+                    "bkc,bkc->bk",
+                    audio_cls.to(torch.float32),
+                    image_cls.to(torch.float32))
+                raw_sims += cls_sims.reshape(v, self.num_heads, 1, 1, 1, 1)
+        else:
+            raise ValueError(f"Improper image shape: {image_feats.shape}")
+
+        assert a1 == a2 and c2 == c2 and t1 == t2
+
+        if raw:
+            return raw_sims
+        else:
+            return self.prepare_sims(raw_sims, audio_mask, agg_sim, agg_heads)
+
+    def forward(self, preds, agg_heads):
+        return self._get_full_sims(
+            preds, raw=False, agg_sim=True, agg_heads=agg_heads)
+
+    def forward_batched(self, preds, agg_heads, batch_size):
+        new_preds = {k: v for k, v in preds.items()}
+        big_image_feats = new_preds.pop(IMAGE_FEATS)
+        if self.use_cls:
+            big_image_cls = new_preds.pop(IMAGE_CLS)
+
+        n = big_image_feats.shape[0]
+        n_steps = math.ceil(n / batch_size)
+        outputs = []
+        for step in tqdm(range(n_steps), "Calculating Sim", leave=False):
+            new_preds[IMAGE_FEATS] = big_image_feats[step * batch_size:(step + 1) * batch_size].cuda()
+            if self.use_cls:
+                new_preds[IMAGE_CLS] = big_image_cls[step * batch_size:(step + 1) * batch_size].cuda()
+
+            sim = self.forward(new_preds, agg_heads=agg_heads)
+            outputs.append(sim.cpu())
+        return torch.cat(outputs, dim=1)
+
+
+class ImageThenAudioAggregator(BaseAggregator):
+
+    def __init__(self, image_agg_type, audio_agg_type, nonneg_sim, mask_silence, num_heads, head_agg, use_cls):
+        super().__init__(nonneg_sim, mask_silence, num_heads, head_agg, use_cls)
+        if image_agg_type == "max":
+            self.image_agg = lambda x, dim: x.max(dim=dim, keepdim=True).values
+        elif image_agg_type == "avg":
+            self.image_agg = lambda x, dim: x.mean(dim=dim, keepdim=True)
+        else:
+            raise ValueError(f"Unknown image_agg_type {image_agg_type}")
+
+        if audio_agg_type == "max":
+            self.time_agg = masked_max
+        elif audio_agg_type == "avg":
+            self.time_agg = masked_mean
+        else:
+            raise ValueError(f"Unknown audio_agg_type {audio_agg_type}")
+
+        self.freq_agg = lambda x, dim: x.mean(dim=dim, keepdim=True)
+
+    def _agg_sim(self, sim, mask):
+        sim_shape = sim.shape
+        new_mask_shape = [1] * len(sim_shape)
+        new_mask_shape[0] = sim_shape[0]
+        new_mask_shape[-1] = sim_shape[-1]
+        mask = mask.reshape(new_mask_shape)
+        sim = self.image_agg(sim, -3)
+        sim = self.image_agg(sim, -4)
+        sim = self.freq_agg(sim, -2)
+        sim = self.time_agg(sim, mask, -1)
+        return sim.squeeze(-1).squeeze(-1).squeeze(-1).squeeze(-1)
+
+
+class PairedAggregator(BaseAggregator):
+
+    def __init__(self, nonneg_sim, mask_silence, num_heads, head_agg, use_cls):
+        super().__init__(nonneg_sim, mask_silence, num_heads, head_agg, use_cls)
+        self.image_agg_max = lambda x, dim: x.max(dim=dim, keepdim=True).values
+        self.image_agg_mean = lambda x, dim: x.mean(dim=dim, keepdim=True)
+
+        self.time_agg_max = masked_max
+        self.time_agg_mean = masked_mean
+
+        self.freq_agg = lambda x, dim: x.mean(dim=dim, keepdim=True)
+
+    def _agg_sim(self, sim, mask):
+        sim_shape = sim.shape
+        new_mask_shape = [1] * len(sim_shape)
+        new_mask_shape[0] = sim_shape[0]
+        new_mask_shape[-1] = sim_shape[-1]
+        mask = mask.reshape(new_mask_shape)
+
+        sim_1 = self.image_agg_max(sim, -3)
+        sim_1 = self.image_agg_max(sim_1, -4)
+        sim_1 = self.freq_agg(sim_1, -2)
+        sim_1 = self.time_agg_mean(sim_1, mask, -1)
+
+        sim_2 = self.freq_agg(sim, -2)
+        sim_2 = self.time_agg_max(sim_2, mask, -1)
+        sim_2 = self.image_agg_mean(sim_2, -3)
+        sim_2 = self.image_agg_mean(sim_2, -4)
+
+        sim = 1 / 2 * (sim_1 + sim_2)
+
+        return sim.squeeze(-1).squeeze(-1).squeeze(-1).squeeze(-1)
+
+
+
+class CAVMAEAggregator(BaseAggregator):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(False, False, 1, "sum", False)
+
+    def _get_full_sims(self, preds, raw, agg_sim, agg_heads):
+        if agg_sim:
+            audio_feats = preds[AUDIO_FEATS]
+            image_feats = preds[IMAGE_FEATS]
+            pool_audio_feats = F.normalize(audio_feats.mean(dim=[-1, -2]), dim=1)
+            pool_image_feats = F.normalize(image_feats.mean(dim=[-1, -2]), dim=1)
+            sims = torch.einsum(
+                "bc,dc->bd",
+                pool_audio_feats.to(torch.float32),
+                pool_image_feats.to(torch.float32))
+            if agg_heads:
+                return sims
+            else:
+                return sims.unsqueeze(-1)
+
+        else:
+            return BaseAggregator._get_full_sims(self, preds, raw, agg_sim, agg_heads)
+
+    def get_pairwise_sims(self, preds, raw, agg_sim, agg_heads):
+        if agg_sim:
+            audio_feats = preds[AUDIO_FEATS]
+            image_feats = preds[IMAGE_FEATS]
+            pool_audio_feats = F.normalize(audio_feats.mean(dim=[-1, -2]), dim=1)
+            pool_image_feats = F.normalize(image_feats.mean(dim=[-1, -2]), dim=1)
+            sims = torch.einsum(
+                "bc,bc->b",
+                pool_audio_feats.to(torch.float32),
+                pool_image_feats.to(torch.float32))
+            if agg_heads:
+                return sims
+            else:
+                return sims.unsqueeze(-1)
+
+        else:
+            return BaseAggregator.get_pairwise_sims(self, preds, raw, agg_sim, agg_heads)
+
+
+class ImageBindAggregator(BaseAggregator):
+
+    def __init__(self, num_heads, *args, **kwargs):
+        super().__init__(False, False, num_heads, "sum", False)
+
+    def _get_full_sims(self, preds, raw, agg_sim, agg_heads):
+        if agg_sim:
+            sims = torch.einsum(
+                "bc,dc->bd",
+                preds[AUDIO_CLS].to(torch.float32),
+                preds[IMAGE_CLS].to(torch.float32))
+            if agg_heads:
+                return sims
+            else:
+                sims = sims.unsqueeze(-1)
+                return sims.repeat(*([1] * (sims.dim() - 1)), self.num_heads)
+
+
+        else:
+            return BaseAggregator._get_full_sims(self, preds, raw, agg_sim, agg_heads)
+
+    def get_pairwise_sims(self, preds, raw, agg_sim, agg_heads):
+        if agg_sim:
+            sims = torch.einsum(
+                "bc,dc->b",
+                preds[AUDIO_CLS].to(torch.float32),
+                preds[IMAGE_CLS].to(torch.float32))
+            if agg_heads:
+                return sims
+            else:
+                sims = sims.unsqueeze(-1)
+                return sims.repeat(*([1] * (sims.dim() - 1)), self.num_heads)
+
+        else:
+            return BaseAggregator.get_pairwise_sims(self, preds, raw, agg_sim, agg_heads)
+
+    def forward_batched(self, preds, agg_heads, batch_size):
+        return self.forward(preds, agg_heads)
+
+
+class SimPool(nn.Module):
+    def __init__(self, dim, num_heads=1, qkv_bias=False, qk_scale=None, gamma=None, use_beta=False):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.norm_patches = nn.LayerNorm(dim, eps=1e-6)
+
+        self.wq = nn.Linear(dim, dim, bias=qkv_bias)
+        self.wk = nn.Linear(dim, dim, bias=qkv_bias)
+
+        if gamma is not None:
+            self.gamma = torch.tensor([gamma])
+            if use_beta:
+                self.beta = nn.Parameter(torch.tensor([0.0]))
+        self.eps = torch.tensor([1e-6])
+
+        self.gamma = gamma
+        self.use_beta = use_beta
+
+    def prepare_input(self, x):
+        if len(x.shape) == 3:  # Transformer
+            # Input tensor dimensions:
+            # x: (B, N, d), where B is batch size, N are patch tokens, d is depth (channels)
+            B, N, d = x.shape
+            gap_cls = x.mean(-2)  # (B, N, d) -> (B, d)
+            gap_cls = gap_cls.unsqueeze(1)  # (B, d) -> (B, 1, d)
+            return gap_cls, x
+        if len(x.shape) == 4:  # CNN
+            # Input tensor dimensions:
+            # x: (B, d, H, W), where B is batch size, d is depth (channels), H is height, and W is width
+            B, d, H, W = x.shape
+            gap_cls = x.mean([-2, -1])  # (B, d, H, W) -> (B, d)
+            x = x.reshape(B, d, H * W).permute(0, 2, 1)  # (B, d, H, W) -> (B, d, H*W) -> (B, H*W, d)
+            gap_cls = gap_cls.unsqueeze(1)  # (B, d) -> (B, 1, d)
+            return gap_cls, x
+        else:
+            raise ValueError(f"Unsupported number of dimensions in input tensor: {len(x.shape)}")
+
+    def forward(self, x):
+        self.eps = self.eps.to(x.device)
+        # Prepare input tensor and perform GAP as initialization
+        gap_cls, x = self.prepare_input(x)
+
+        # Prepare queries (q), keys (k), and values (v)
+        q, k, v = gap_cls, self.norm_patches(x), self.norm_patches(x)
+
+        # Extract dimensions after normalization
+        Bq, Nq, dq = q.shape
+        Bk, Nk, dk = k.shape
+        Bv, Nv, dv = v.shape
+
+        # Check dimension consistency across batches and channels
+        assert Bq == Bk == Bv
+        assert dq == dk == dv
+
+        # Apply linear transformation for queries and keys then reshape
+        qq = self.wq(q).reshape(Bq, Nq, self.num_heads, dq // self.num_heads).permute(0, 2, 1,
+                                                                                      3)  # (Bq, Nq, dq) -> (B, num_heads, Nq, dq/num_heads)
+        kk = self.wk(k).reshape(Bk, Nk, self.num_heads, dk // self.num_heads).permute(0, 2, 1,
+                                                                                      3)  # (Bk, Nk, dk) -> (B, num_heads, Nk, dk/num_heads)
+
+        vv = v.reshape(Bv, Nv, self.num_heads, dv // self.num_heads).permute(0, 2, 1,
+                                                                             3)  # (Bv, Nv, dv) -> (B, num_heads, Nv, dv/num_heads)
+
+        # Compute attention scores
+        attn = (qq @ kk.transpose(-2, -1)) * self.scale
+        # Apply softmax for normalization
+        attn = attn.softmax(dim=-1)
+
+        # If gamma scaling is used
+        if self.gamma is not None:
+            # Apply gamma scaling on values and compute the weighted sum using attention scores
+            x = torch.pow(attn @ torch.pow((vv - vv.min() + self.eps), self.gamma),
+                          1 / self.gamma)  # (B, num_heads, Nv, dv/num_heads) -> (B, 1, 1, d)
+            # If use_beta, add a learnable translation
+            if self.use_beta:
+                x = x + self.beta
+        else:
+            # Compute the weighted sum using attention scores
+            x = (attn @ vv).transpose(1, 2).reshape(Bq, Nq, dq)
+
+        return x.squeeze()
+
+
+
+class SimPoolAggregator(BaseAggregator):
+
+    def __init__(self, num_heads, dim, *args, **kwargs):
+        super().__init__(False, False, num_heads, "sum", False)
+        self.pool = SimPool(dim, gamma=1.25)
+
+    def _get_full_sims(self, preds, raw, agg_sim, agg_heads):
+        if agg_sim:
+            device = self.pool.wq.weight.data.device
+            pooled_audio = self.pool(preds[AUDIO_FEATS].to(torch.float32).to(device))
+            pooled_image = self.pool(preds[IMAGE_FEATS].to(torch.float32).to(device))
+
+            sims = torch.einsum(
+                "bc,dc->bd",
+                pooled_audio,
+                pooled_image)
+            if agg_heads:
+                return sims
+            else:
+                sims = sims.unsqueeze(-1)
+                return sims.repeat(*([1] * (sims.dim() - 1)), self.num_heads)
+
+
+        else:
+            return BaseAggregator._get_full_sims(self, preds, raw, agg_sim, agg_heads)
+
+    def get_pairwise_sims(self, preds, raw, agg_sim, agg_heads):
+        if agg_sim:
+            device = self.pool.wq.weight.data.device
+            pooled_audio = self.pool(preds[AUDIO_FEATS].to(torch.float32).to(device))
+            pooled_image = self.pool(preds[IMAGE_FEATS].to(torch.float32).to(device))
+
+            sims = torch.einsum(
+                "bc,dc->b",
+                pooled_audio,
+                pooled_image)
+            if agg_heads:
+                return sims
+            else:
+                sims = sims.unsqueeze(-1)
+                return sims.repeat(*([1] * (sims.dim() - 1)), self.num_heads)
+
+        else:
+            return BaseAggregator.get_pairwise_sims(self, preds, raw, agg_sim, agg_heads)
+
+    def forward_batched(self, preds, agg_heads, batch_size):
+        return self.forward(preds, agg_heads)
+
+
+
+def get_aggregator(sim_agg_type, nonneg_sim, mask_silence, num_heads, head_agg, use_cls, dim):
+    shared_args = dict(
+        nonneg_sim=nonneg_sim,
+        mask_silence=mask_silence,
+        num_heads=num_heads,
+        head_agg=head_agg,
+        use_cls=use_cls,
+    )
+
+    if sim_agg_type == "paired":
+        agg1 = PairedAggregator(**shared_args)
+    elif sim_agg_type == "misa":
+        agg1 = ImageThenAudioAggregator("max", "avg", **shared_args)
+    elif sim_agg_type == "mima":
+        agg1 = ImageThenAudioAggregator("max", "max", **shared_args)
+    elif sim_agg_type == "sisa":
+        agg1 = ImageThenAudioAggregator("avg", "avg", **shared_args)
+    elif sim_agg_type == "cavmae":
+        agg1 = CAVMAEAggregator()
+    elif sim_agg_type == "imagebind":
+        agg1 = ImageBindAggregator(num_heads=shared_args["num_heads"])
+    elif sim_agg_type == "simpool":
+        agg1 = SimPoolAggregator(num_heads=shared_args["num_heads"], dim=dim)
+    else:
+        raise ValueError(f"Unknown loss_type {sim_agg_type}")
+
+    return agg1
+

DenseAV/denseav/aligners.py

@@ -0,0 +1,300 @@
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+from torch.nn import ModuleList
+
+from denseav.featurizers.DINO import Block
+
+
+class ChannelNorm(torch.nn.Module):
+
+    def __init__(self, dim, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.norm = torch.nn.LayerNorm(dim, eps=1e-4)
+
+    def forward_spatial(self, x):
+        return self.norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+    def forward(self, x, cls):
+        return self.forward_spatial(x), self.forward_cls(cls)
+
+    def forward_cls(self, cls):
+        if cls is not None:
+            return self.norm(cls)
+        else:
+            return None
+
+
+def id_conv(dim, strength=.9):
+    conv = torch.nn.Conv2d(dim, dim, 1, padding="same")
+    start_w = conv.weight.data
+    conv.weight.data = torch.nn.Parameter(
+        torch.eye(dim, device=start_w.device).unsqueeze(-1).unsqueeze(-1) * strength + start_w * (1 - strength))
+    conv.bias.data = torch.nn.Parameter(conv.bias.data * (1 - strength))
+    return conv
+
+
+class LinearAligner(torch.nn.Module):
+    def __init__(self, in_dim, out_dim, use_norm=True):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        if use_norm:
+            self.norm = ChannelNorm(in_dim)
+        else:
+            self.norm = Identity2()
+
+        if in_dim == out_dim:
+            self.layer = id_conv(in_dim, 0)
+        else:
+            self.layer = torch.nn.Conv2d(in_dim, out_dim, kernel_size=1, stride=1)
+
+        self.cls_layer = torch.nn.Linear(in_dim, out_dim)
+
+    def forward(self, spatial, cls):
+        norm_spatial, norm_cls = self.norm(spatial, cls)
+
+        if cls is not None:
+            aligned_cls = self.cls_layer(cls)
+        else:
+            aligned_cls = None
+
+        return self.layer(norm_spatial), aligned_cls
+
+class IdLinearAligner(torch.nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        assert self.out_dim == self.in_dim
+        self.layer = id_conv(in_dim, 1.0)
+    def forward(self, spatial, cls):
+        return self.layer(spatial), cls
+
+
+class FrequencyAvg(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, spatial, cls):
+        return spatial.mean(2, keepdim=True), cls
+
+
+class LearnedTimePool(torch.nn.Module):
+    def __init__(self, dim, width, maxpool):
+        super().__init__()
+        self.dim = dim
+        self.width = width
+        self.norm = ChannelNorm(dim)
+        if maxpool:
+            self.layer = torch.nn.Sequential(
+                torch.nn.Conv2d(dim, dim, kernel_size=width, stride=1, padding="same"),
+                torch.nn.MaxPool2d(kernel_size=(1, width), stride=(1, width))
+            )
+        else:
+            self.layer = torch.nn.Conv2d(dim, dim, kernel_size=(1, width), stride=(1, width))
+
+    def forward(self, spatial, cls):
+        norm_spatial, norm_cls = self.norm(spatial, cls)
+        return self.layer(norm_spatial), norm_cls
+
+
+class LearnedTimePool2(torch.nn.Module):
+    def __init__(self, in_dim, out_dim, width, maxpool, use_cls_layer):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.width = width
+
+        if maxpool:
+            self.layer = torch.nn.Sequential(
+                torch.nn.Conv2d(in_dim, out_dim, kernel_size=width, stride=1, padding="same"),
+                torch.nn.MaxPool2d(kernel_size=(1, width), stride=(1, width))
+            )
+        else:
+            self.layer = torch.nn.Conv2d(in_dim, out_dim, kernel_size=(1, width), stride=(1, width))
+
+        self.use_cls_layer = use_cls_layer
+        if use_cls_layer:
+            self.cls_layer = torch.nn.Linear(in_dim, out_dim)
+
+    def forward(self, spatial, cls):
+
+        if cls is not None:
+            if self.use_cls_layer:
+                aligned_cls = self.cls_layer(cls)
+            else:
+                aligned_cls = cls
+        else:
+            aligned_cls = None
+
+        return self.layer(spatial), aligned_cls
+
+
+class Sequential2(torch.nn.Module):
+
+    def __init__(self, *modules):
+        super().__init__()
+        self.mod_list = ModuleList(modules)
+
+    def forward(self, x, y):
+        results = (x, y)
+        for m in self.mod_list:
+            results = m(*results)
+        return results
+
+
+class ProgressiveGrowing(torch.nn.Module):
+
+    def __init__(self, stages, phase_lengths):
+        super().__init__()
+        self.stages = torch.nn.ModuleList(stages)
+        self.phase_lengths = torch.tensor(phase_lengths)
+        assert len(self.phase_lengths) + 1 == len(self.stages)
+        self.phase_boundaries = self.phase_lengths.cumsum(0)
+        self.register_buffer('phase', torch.tensor([1]))
+
+    def maybe_change_phase(self, global_step):
+        needed_phase = (global_step >= self.phase_boundaries).to(torch.int64).sum().item() + 1
+        if needed_phase != self.phase.item():
+            print(f"Changing aligner phase to {needed_phase}")
+            self.phase.copy_(torch.tensor([needed_phase]).to(self.phase.device))
+            return True
+        else:
+            return False
+
+    def parameters(self, recurse: bool = True):
+        phase = self.phase.item()
+        used_stages = self.stages[:phase]
+        print(f"Progressive Growing at stage {phase}")
+        all_params = []
+        for stage in used_stages:
+            all_params.extend(stage.parameters(recurse))
+        return iter(all_params)
+
+    def forward(self, spatial, cls):
+        pipeline = Sequential2(*self.stages[:self.phase.item()])
+        return pipeline(spatial, cls)
+
+
+class Identity2(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, y):
+        return x, y
+
+
+class SelfAttentionAligner(torch.nn.Module):
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+        self.num_heads = 6
+        if dim % self.num_heads != 0:
+            self.padding = self.num_heads - (dim % self.num_heads)
+        else:
+            self.padding = 0
+
+        self.block = Block(
+            dim + self.padding,
+            num_heads=self.num_heads,
+            mlp_ratio=4,
+            qkv_bias=True,
+            qk_scale=None,
+            drop=0.0,
+            attn_drop=0.0,
+            drop_path=0.0,
+            norm_layer=partial(torch.nn.LayerNorm, eps=1e-4))
+
+    def forward(self, spatial, cls):
+        padded_feats = F.pad(spatial, [0, 0, 0, 0, self.padding, 0])
+
+        B, C, H, W = padded_feats.shape
+        proj_feats = padded_feats.reshape(B, C, H * W).permute(0, 2, 1)
+
+        if cls is not None:
+            assert len(cls.shape) == 2
+            padded_cls = F.pad(cls, [self.padding, 0])
+            proj_feats = torch.cat([padded_cls.unsqueeze(1), proj_feats], dim=1)
+
+        aligned_feat, attn, qkv = self.block(proj_feats, return_qkv=True)
+
+        if cls is not None:
+            aligned_cls = aligned_feat[:, 0, :]
+            aligned_spatial = aligned_feat[:, 1:, :]
+        else:
+            aligned_cls = None
+            aligned_spatial = aligned_feat
+
+        aligned_spatial = aligned_spatial.reshape(B, H, W, self.dim + self.padding).permute(0, 3, 1, 2)
+
+        aligned_spatial = aligned_spatial[:, self.padding:, :, :]
+        if aligned_cls is not None:
+            aligned_cls = aligned_cls[:, self.padding:]
+
+        return aligned_spatial, aligned_cls
+
+
+def get_aligner(aligner_type, in_dim, out_dim, **kwargs):
+    if aligner_type is None:
+        return Identity2()
+
+    if "prog" in aligner_type:
+        phase_length = kwargs["phase_length"]
+
+    if aligner_type == "image_linear":
+        return LinearAligner(in_dim, out_dim)
+    elif aligner_type == "image_idlinear":
+        return IdLinearAligner(in_dim, out_dim)
+    elif aligner_type == "image_linear_no_norm":
+        return LinearAligner(in_dim, out_dim, use_norm=False)
+    elif aligner_type == "image_id":
+        return Identity2()
+    elif aligner_type == "image_norm":
+        return ChannelNorm(in_dim)
+    elif aligner_type == "audio_linear":
+        return Sequential2(
+            LinearAligner(in_dim, out_dim),
+            FrequencyAvg())
+    elif aligner_type == "audio_sa":
+        return Sequential2(
+            LinearAligner(in_dim, out_dim),
+            FrequencyAvg(),
+            SelfAttentionAligner(out_dim)
+        )
+    elif aligner_type == "audio_sa_sa":
+        return Sequential2(
+            FrequencyAvg(),
+            LinearAligner(in_dim, out_dim),
+            SelfAttentionAligner(out_dim),
+            SelfAttentionAligner(out_dim)
+        )
+    elif aligner_type == "audio_3_3_pool":
+        return Sequential2(
+            LinearAligner(in_dim, out_dim),
+            FrequencyAvg(),
+            LearnedTimePool(out_dim, 3, False),
+            LearnedTimePool(out_dim, 3, False),
+        )
+    elif aligner_type == "audio_sa_3_3_pool":
+        return Sequential2(
+            LinearAligner(in_dim, out_dim),
+            FrequencyAvg(),
+            LearnedTimePool(out_dim, 3, False),
+            LearnedTimePool(out_dim, 3, False),
+            SelfAttentionAligner(out_dim)
+        )
+    elif aligner_type == "audio_sa_3_3_pool_2":
+        return Sequential2(
+            FrequencyAvg(),
+            ChannelNorm(in_dim),
+            LearnedTimePool2(in_dim, out_dim, 3, False, True),
+            LearnedTimePool2(out_dim, out_dim, 3, False, False),
+            SelfAttentionAligner(out_dim)
+        )
+    else:
+        raise ValueError(f"Unknown aligner type {aligner_type}")

DenseAV/denseav/configs/av_align.yaml

@@ -0,0 +1,125 @@
+# Model args
+
+code_dim: 384
+image_model_type: "dino8"
+image_model_token_type: "token"
+image_aligner_type: "image_linear"
+image_pool_width: 2
+
+audio_model_type: "hubert"
+audio_aligner_type: "audio_sa_3_3_pool_2"
+audio_pool_width: 1
+
+learn_audio_cls: True
+
+#code_dim: 1024
+#image_model_type: "imagebind"
+#image_model_token_type: "token"
+#image_aligner_type: "image_linear"
+#image_pool_width: 1
+#
+#audio_model_type: "imagebind"
+#audio_aligner_type: "audio_sa"
+#audio_pool_width: 1
+#
+#learn_audio_cls: False
+
+audio_lora: False
+audio_lora_rank: 8
+image_lora: True
+image_lora_rank: 8
+
+
+spatial_dropout: 0.0
+channel_dropout: 0.0
+
+quad_mixup: 0.1
+bg_mixup: 0.0
+patch_mixup: 0.0
+mixup_weight: 0.1
+
+sim_agg_type: "misa"
+sim_agg_heads: 1
+sim_use_cls: False
+
+cal_init: 1.0
+cal_balance_weight: 0.1
+nonneg_sim: False
+nonneg_pressure: 0.01
+silence_l1: 0.01
+silence_l2: 0.0
+tv_weight: 0.01
+specialization_weight: 0.05
+head_agg: "max_elementwise"
+disentangle_weight: 0.0
+
+norm_vectors: False
+
+neg_audio: true
+neg_audio_weight: 0.01
+
+
+pretrain_steps: 3000
+pretrain_lr: .5e-4
+
+# Loss args
+lr: .5e-4
+lr_warmup: 1000
+
+#lr_warmup: 100
+
+lr_schedule: ~
+lr_cycle_length: 50000
+
+optimizer: "adam"
+gradient_clipping: 10.0
+adaptive_clipping: True
+gather_tensors: True
+loss_type: "nce"
+loss_leak: 0.0
+loss_margin: 0.0
+mask_silence: true
+extra_audio_masking: true
+max_steps: 1000001
+
+finetune_image_model: False
+finetune_audio_model: True
+
+# Checkpointing args
+load_strict: true
+starting_weights: ~
+auto_resume: false
+grouping_name: "foo"
+resume_prefix: "imagebind_exp2"
+
+# Data Args
+#dataset_name: "sample-audio"
+dataset_name: "places-audio"
+#dataset_name: "mixed"
+#dataset_name: "audio-set-full"
+use_extra_val_sets: true
+batch_size: 10
+load_size: 224
+image_aug: true
+audio_aug: false
+
+audio_level: false
+
+memory_buffer_size: 0
+
+val_check_interval: 10000 #0
+use_cached_embs: false
+num_workers: 12
+num_gpus: 4
+num_sanity_val_steps: 0 #-1
+seed: 0
+
+# Env args
+output_root: '../'
+pytorch_data_dir: '/pytorch-data/'
+submitting_to_aml: false
+
+hydra:
+  run:
+    dir: "."
+  output_subdir: ~

DenseAV/denseav/constants.py

@@ -0,0 +1,12 @@
+
+IMAGE_INPUT = "frames"
+IMAGE_FEATS = "image_feats"
+IMAGE_CLS = "image_cls"
+IMAGE_MASK = "image_masks"
+
+AUDIO_FEATS = "audio_feats"
+AUDIO_CLS = "audio_cls"
+AUDIO_MASK = "audio_mask"
+AUDIO_POS_MASK = "audio_pos_mask"
+
+DATA_SOURCE = "source"

DenseAV/denseav/data/AVDatasets.py

@@ -0,0 +1,1249 @@
+import glob
+import os
+from abc import ABC, abstractmethod
+from glob import glob
+from os.path import join
+from pathlib import Path
+from typing import List, Set
+
+import audioread
+import numpy as np
+import pandas as pd
+import pytorch_lightning as pl
+import torch
+import torch.nn.functional as F
+import torchaudio
+import torchvision.transforms as T
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader, default_collate, Subset, ConcatDataset
+from tqdm import tqdm
+
+from denseav.constants import AUDIO_MASK, AUDIO_POS_MASK, IMAGE_MASK, IMAGE_INPUT
+from denseav.data.make_tarballs import untar_all
+from denseav.shared import norm, prep_waveform
+
+
+def sample_choice(choices, probs):
+    # Check that probabilities sum to 1 and are non-negative
+    assert sum(probs) == 1, "Probabilities must sum to 1"
+    assert all(p >= 0 for p in probs), "Probabilities cannot be negative"
+
+    # Convert probs to a tensor
+    probs_tensor = torch.tensor(probs)
+
+    # Sample a choice according to the probabilities
+    index = torch.multinomial(probs_tensor, 1).item()
+
+    # Return the sampled choice
+    return choices[index]
+
+
+def grid_frames(frames):
+    top_row = torch.cat([frames[0], frames[1]], dim=2)
+    bottom_row = torch.cat([frames[2], frames[3]], dim=2)
+    return torch.cat([top_row, bottom_row], dim=3)
+
+
+def create_mixed_image(pos_frame, neg_frame, patch_size):
+    # Step 1: Check that patch_size evenly divides the image dimensions
+    b, c, h, w = pos_frame.shape
+    assert h % patch_size == 0 and w % patch_size == 0, "Patch size must evenly divide image dimensions"
+
+    # Step 2: Create a random binary mask with the same number of patches as the image
+    mask = torch.randint(0, 2, (b, 1, h // patch_size, w // patch_size))
+
+    # Step 3: Create a new image using patches from pos_frame and neg_frame according to the mask
+    # Upscale the mask to the size of the image
+    mask_upscaled = F.interpolate(mask.to(torch.float32), scale_factor=patch_size)
+
+    # Use the mask to create a mixed frame
+    mixed_frame = mask_upscaled * pos_frame + (1 - mask_upscaled) * neg_frame
+
+    return mixed_frame, mask_upscaled
+
+
+class AVDataset(ABC, Dataset):
+
+    @abstractmethod
+    def _dataset_folder(self) -> str:
+        pass
+
+    @abstractmethod
+    def _load_info(self, split) -> pd.DataFrame:
+        """
+        This function should return a dataframe with at least a column "id"
+        @return:
+        """
+        pass
+
+    @abstractmethod
+    def _missing_threshold(self) -> float:
+        pass
+
+    @abstractmethod
+    def default_target_length(self) -> int:
+        pass
+
+    def target_length(self):
+        if self.override_target_length is not None:
+            return self.override_target_length
+        else:
+            return self.default_target_length()
+
+    def _frame_root(self) -> str:
+        return join(self.root, "frames", self.split)
+
+    def _video_root(self) -> str:
+        return join(self.root, "videos", self.split)
+
+    def _audio_root(self) -> str:
+        return join(self.root, "audio", self.split)
+
+    def _semseg_root(self) -> str:
+        return join(self.root, "annotations", self.split)
+
+    def _embed_root(self) -> str:
+        return join(self.root, "embedding", self.audio_embed_model, self.split)
+
+    def _label_root(self) -> str:
+        return join(self.root, "pseudo-labels")
+
+    def _hn_root(self) -> str:
+        return join(self.root, "hard_negatives")
+
+    def _all_video_files(self) -> Set[str]:
+        return set(str(p) for p in Path(join(self._video_root())).rglob('*'))
+
+    def _all_frame_files(self) -> Set[str]:
+        return set(str(p) for p in Path(join(self._frame_root())).rglob('*'))
+
+    def _all_audio_files(self) -> Set[str]:
+        return set(str(p) for p in Path(join(self._audio_root())).rglob('*'))
+
+    def _all_embed_files(self) -> Set[str]:
+        return set(str(p) for p in Path(join(self._embed_root())).rglob('*'))
+
+    def _get_frame_files(self, row) -> List[str]:
+        return [self._frame_root() + "/" + row["id"] + f"_{i}.jpg" for i in range(self._expected_num_frames())]
+
+    def _get_semseg_file(self, row) -> str:
+        raise NotImplementedError("Class has not implemented _get_semseg_files")
+
+    def _get_audio_file(self, row) -> str:
+        return self._audio_root() + "/" + row["id"] + ".mp3"
+
+    def _get_video_file(self, row) -> str:
+        return self._video_root() + "/" + row["id"] + ".mp4"
+
+    def _get_embed_file(self, row) -> str:
+        return self._embed_root() + "/" + row["id"] + ".npz"
+
+    def _add_files_to_metadata(self, df) -> pd.DataFrame:
+        tqdm.pandas()
+
+        if self.use_audio_embed:
+            df["embed_file"] = df.progress_apply(self._get_embed_file, axis=1)
+
+        if self.use_audio or self.use_spec:
+            df["audio_file"] = df.progress_apply(self._get_audio_file, axis=1)
+
+        if self.use_frames:
+            df["frame_files"] = df.progress_apply(self._get_frame_files, axis=1)
+
+        if self.use_semseg:
+            df["semseg_file"] = df.progress_apply(self._get_semseg_file, axis=1)
+
+        df = self._filter_valid_metadata(df)
+
+        if self.use_hn:
+            loaded = np.load(join(self._hn_root(), "original", f"{self.split}_hard_negatives.npz"))
+            df["hn0"] = [t for t in torch.tensor(loaded["indices_0"])]
+            df["hn1"] = [t for t in torch.tensor(loaded["indices_1"])]
+
+        return df
+
+    def _split_name(self, split):
+        return split
+
+    def _filter_valid_metadata(self, df: pd.DataFrame) -> pd.DataFrame:
+
+        print("MY_DIR ", list(glob(join(self.root, "*"))))
+        if self.use_audio_embed:
+            missing_embed_files = set(df['embed_file']) - self.all_embed_files
+            valid_audio = ~df['embed_file'].isin(missing_embed_files)
+            print("ALL EMBED ", len(self.all_embed_files))
+        elif self.use_audio or self.use_spec:
+            missing_audio_files = set(df['audio_file']) - self.all_audio_files
+            valid_audio = ~df['audio_file'].isin(missing_audio_files)
+            print("ALL AUDIO ", len(self.all_audio_files))
+
+        if self.use_frames:
+            missing_frame_files = set(
+                item for sublist in df['frame_files'].tolist() for item in sublist) - self.all_frame_files
+            valid_frames = df['frame_files'].apply(lambda x: not any(file in missing_frame_files for file in x))
+            print("ALL FRAMES ", len(self.all_frame_files))
+            df["is_valid"] = valid_audio & valid_frames
+        else:
+            df["is_valid"] = valid_audio
+
+        percent_missing = (1 - (df["is_valid"].sum() / len(df)))
+
+        assert percent_missing <= self._missing_threshold(), \
+            f"Too many missing files: %{round(percent_missing * 100.0, 2)}"
+        assert len(df) > 0, "No files found"
+        return df[df["is_valid"]]
+
+    def __init__(
+            self,
+            root: str,
+            split: str = "train",
+            use_frames=False,
+            frame_transform=None,
+            use_audio=False,
+            use_spec=False,
+            use_audio_embed=False,
+            use_hn=False,
+            use_caption=False,
+            use_semseg=False,
+            neg_audio=False,
+            use_davenet_spec=False,
+            use_fnac_spec=False,
+            n_label_frames=196,
+            label_transform=None,
+            audio_embed_model="hubert",
+            n_frames=1,
+            audio_transform=None,
+            audio_aug=False,
+            spec_transform=None,
+            spec_mel_bins=128,
+            spec_mean=-6.6268077,
+            spec_std=5.358466,
+            sample_rate=16000,
+            override_target_length=None,
+            use_tags=False,
+            extra_audio_masking=False,
+            audio_level=False,
+            quad_mixup=0.0,
+            bg_mixup=0.0,
+            patch_mixup=0.0,
+            patch_size=8,
+    ):
+        super(AVDataset).__init__()
+        self.pytorch_data_dir = root
+        self.split = self._split_name(split)
+        self.root = join(root, self._dataset_folder())
+        self.use_frames = use_frames
+        self.frame_transform = frame_transform
+        self.use_audio = use_audio
+        self.use_spec = use_spec
+        self.use_audio_embed = use_audio_embed
+        self.use_davenet_spec = use_davenet_spec
+        self.use_fnac_spec = use_fnac_spec
+        self.use_hn = use_hn
+        self.use_caption = use_caption
+        self.label_transform = label_transform
+        self.audio_embed_model = audio_embed_model
+        self.audio_aug = audio_aug
+        self.n_frames = n_frames
+        self.audio_transform = audio_transform
+        self.spec_transform = spec_transform
+        self.spec_mel_bins = spec_mel_bins
+        self.spec_mean = spec_mean
+        self.spec_std = spec_std
+        self.use_semseg = use_semseg
+        self.override_target_length = override_target_length
+        self.use_tags = use_tags
+        self.extra_audio_masking = extra_audio_masking
+        self.neg_audio = neg_audio
+        self.audio_level = audio_level
+
+        self.quad_mixup = quad_mixup
+        self.bg_mixup = bg_mixup
+        self.patch_mixup = patch_mixup
+        self.patch_size = patch_size
+
+        self.sample_rate = sample_rate
+        self.n_label_frames = n_label_frames
+
+        if self.use_audio_embed:
+            self.all_embed_files = self._all_embed_files()
+
+        if self.use_audio or self.use_spec:
+            self.all_audio_files = self._all_audio_files()
+
+        if self.use_frames:
+            self.all_frame_files = self._all_frame_files()
+
+        self.metadata = self._add_files_to_metadata(self._load_info(self.split))
+
+        assert len(self.metadata) > 0
+
+    def __len__(self):
+        return len(self.metadata)
+
+    @abstractmethod
+    def _expected_num_frames(self) -> int:
+        pass
+
+    def get_audio_mask(self, real_length, padded_length, target_size):
+        if not isinstance(real_length, torch.Tensor):
+            real_length = torch.tensor(real_length)
+            padded_length = torch.tensor(padded_length)
+
+        n_frames = ((real_length / padded_length) * target_size).to(torch.int64)
+        oh = F.one_hot(n_frames, num_classes=target_size + 1)
+        if len(oh.shape) == 1:
+            oh = oh.unsqueeze(0)
+        return (1 - torch.cumsum(oh, dim=1))[:, :-1].to(torch.bool)
+
+    def _base_get_item(self, item):
+        id = self.metadata["id"].iloc[item]
+        data_dict = {"metadata": {"id": id, "index": item}}
+
+        if self.use_tags and "tags" in self.metadata:
+            tags = torch.tensor(self.metadata["tags"].iloc[item])
+            tag_oh = torch.zeros(self.num_tags, dtype=torch.float32)
+            tag_oh[tags] += 1
+            data_dict["tags"] = tag_oh
+
+        if self.use_audio or self.use_spec:
+            audio_file = self.metadata["audio_file"].iloc[item]
+            data_dict["metadata"]["audio_file"] = audio_file
+            loaded_waveform, obs_sr = torchaudio.load(audio_file)
+            loaded_waveform = loaded_waveform[0]
+
+            if self.neg_audio:
+                neg_audio_file = self.metadata["audio_file"].iloc[torch.randint(0, len(self), size=(1,)).item()]
+                data_dict["metadata"]["neg_audio_file"] = neg_audio_file
+                neg_waveform, neg_obs_sr = torchaudio.load(neg_audio_file)
+                neg_waveform = neg_waveform[0]
+            else:
+                neg_waveform, neg_obs_sr = None, None
+
+            (waveform,
+             spectrogram,
+             audio_length,
+             total_length,
+             original_length,
+             mask,
+             pos_mask) = prep_waveform(
+                loaded_waveform,
+                obs_sr,
+                self.target_length(),
+                self.spec_mel_bins,
+                self.spec_mean,
+                self.spec_std,
+                self.sample_rate,
+                self.use_spec,
+                False,
+                self.extra_audio_masking,
+                neg_waveform,
+                neg_obs_sr,
+                self.audio_level,
+                self.audio_aug
+            )
+
+            if self.spec_transform is not None and spectrogram is not None:
+                spectrogram = self.spec_transform(spectrogram)
+
+            if self.audio_transform is not None:
+                waveform = self.audio_transform(waveform)
+
+            data_dict["audio"] = waveform
+            data_dict[AUDIO_MASK] = mask
+            data_dict[AUDIO_POS_MASK] = pos_mask
+            data_dict["audio_length"] = audio_length
+            data_dict["original_length"] = original_length
+            data_dict["total_length"] = total_length
+            if spectrogram is not None:
+                data_dict["spec"] = spectrogram
+
+            if mask.mean() < .04:
+                return None
+
+        if self.use_davenet_spec:
+            from data.DavenetUtilities import davenet_load_audio
+            audio_file = self.metadata["audio_file"].iloc[item]
+            spec, n_frames = davenet_load_audio(audio_file)
+            data_dict["davenet_spec"] = spec
+
+        if self.use_fnac_spec:
+            from featurizers.FNACAVL import load_spectrogram as fnac_load_spectrogram
+            audio_file = self.metadata["audio_file"].iloc[item]
+            data_dict["fnac_spec"] = fnac_load_spectrogram(audio_file, 3)
+
+        if self.use_audio_embed:
+            loaded = np.load(self.metadata["embed_file"].iloc[item])
+            data_dict["audio_emb"] = loaded["feat"]
+            data_dict["audio_length"] = loaded["audio_length"]
+            data_dict["total_length"] = loaded["total_length"]
+            data_dict["original_length"] = loaded["original_length"]
+            data_dict[AUDIO_MASK] = self.get_audio_mask(
+                data_dict["audio_length"],
+                data_dict["total_length"],
+                data_dict["audio_emb"].shape[-1]) \
+                .squeeze().to(torch.float32)
+            data_dict[AUDIO_POS_MASK] = data_dict[AUDIO_MASK].to(torch.float32)
+
+        if self.use_frames:
+
+            def get_frames(item):
+                file_group = self.metadata["frame_files"].iloc[item]
+                if self.n_frames is not None:
+                    selected_frames = torch.randperm(len(file_group))[:self.n_frames]
+                    file_group = [file_group[i] for i in selected_frames]
+                data_dict["metadata"]["frame_files"] = file_group
+                images = [Image.open(file).convert("RGB") for file in file_group]
+
+                if self.frame_transform is not None:
+                    images = torch.cat([self.frame_transform(img).unsqueeze(0) for img in images], dim=0)
+
+                return images, file_group
+
+            no_mixup = 1.0 - (self.bg_mixup + self.quad_mixup + self.patch_mixup)
+
+            mixup_type = sample_choice(
+                ["quad", "bg", "patch", None],
+                [self.quad_mixup, self.bg_mixup, self.patch_mixup, no_mixup]
+            )
+
+            if mixup_type == "quad":
+                indices = [item] + torch.randint(0, len(self), size=(3,)).numpy().tolist()
+                frames_and_files = [get_frames(i) for i in indices]
+                file_group = frames_and_files[0][1]
+                perm = torch.randperm(4)
+                all_frames = [F.interpolate(frames_and_files[i][0], scale_factor=0.5, mode="bilinear") for i in
+                              perm]
+                b, c, h, w = all_frames[0].shape
+                indices = [indices[p] for p in perm]
+                masks = [(torch.ones(b, 1, h, w) if index == item else torch.zeros(b, 1, h, w)) for index in
+                         indices]
+
+                data_dict[IMAGE_INPUT] = grid_frames(all_frames)
+                data_dict[IMAGE_MASK] = grid_frames(masks)
+            elif mixup_type == "bg":
+                neg_item = torch.randint(0, len(self), size=(1,)).item()
+                neg_frame, _ = get_frames(neg_item)
+                pos_frame, file_group = get_frames(item)
+
+                b, c, h, w = neg_frame.shape
+                neg_mask = torch.zeros(b, 1, h, w)
+                pos_mask = torch.ones(b, 1, h, w)
+
+                if torch.rand(1).item() > 0.5:
+                    bg_frame = neg_frame
+                    bg_mask = neg_mask
+                    fg_frame = F.interpolate(pos_frame, scale_factor=0.5, mode="bilinear")
+                    fg_mask = F.interpolate(pos_mask, scale_factor=0.5, mode="bilinear")
+                else:
+                    bg_frame = pos_frame
+                    bg_mask = pos_mask
+                    fg_frame = F.interpolate(neg_frame, scale_factor=0.5, mode="bilinear")
+                    fg_mask = F.interpolate(neg_mask, scale_factor=0.5, mode="bilinear")
+
+                start_h = torch.randint(0, h // 2, size=(1,))
+                start_w = torch.randint(0, w // 2, size=(1,))
+                bg_frame[:, :, start_h:start_h + fg_frame.shape[2], start_w:start_w + fg_frame.shape[3]] = fg_frame
+                bg_mask[:, :, start_h:start_h + fg_frame.shape[2], start_w:start_w + fg_frame.shape[3]] = fg_mask
+
+                data_dict["frames"] = bg_frame
+                data_dict["image_masks"] = bg_mask
+
+            elif mixup_type == "patch":
+                neg_item = torch.randint(0, len(self), size=(1,)).item()
+                neg_frame, _ = get_frames(neg_item)
+                pos_frame, file_group = get_frames(item)
+                frames, masks = create_mixed_image(pos_frame, neg_frame, self.patch_size)
+                data_dict["frames"] = frames
+                data_dict["image_masks"] = masks
+
+            elif mixup_type is None:
+                frames, file_group = get_frames(item)
+
+                data_dict["frames"] = frames
+                b, c, h, w = frames.shape
+                data_dict["image_masks"] = torch.ones(b, 1, h, w)
+            else:
+                raise ValueError(f"Unknown mixup type {mixup_type}")
+
+            if "original_length" in data_dict:
+                if self._expected_num_frames() == 1:
+                    frame_nums = torch.tensor([0])
+                else:
+                    frame_nums = torch.tensor([
+                        int(f.split("/")[-1].split("_")[-1].split(".")[0]) for f in file_group])
+
+                data_dict["frame_nums"] = frame_nums
+                frame_fracs = ((frame_nums + .5) / (self._expected_num_frames()))
+                frame_position = (frame_fracs * data_dict["original_length"]) / data_dict["total_length"]
+                data_dict["frame_position"] = frame_position
+
+        if self.use_caption:
+            if "word" in self.metadata:
+                words = self.metadata["word"].iloc[item]
+                start = self.metadata["start"].iloc[item]
+                end = self.metadata["end"].iloc[item]
+                if isinstance(words, float):
+                    words = [""]
+                    start = [0.0]
+                    end = [-1.0]
+
+                data_dict["caption"] = {
+                    "words": words,
+                    "start": start,
+                    "end": end,
+                }
+            if "text" in self.metadata:
+                data_dict["text"] = self.metadata["text"].iloc[item]
+
+        if self.use_semseg:
+            semseg_path = join(self._semseg_root(), self.metadata["semseg_file"].iloc[item])
+            semseg = Image.open(semseg_path)
+            if self.label_transform is not None:
+                semseg = np.array(self.label_transform(semseg))
+            data_dict["semseg"] = semseg
+            data_dict["metadata"]["semseg_file"] = semseg_path
+
+            # if hasattr(self, "num_classes"):
+            #     data_dict["num_pixels_per_class"] = F.one_hot(
+            #         torch.tensor(semseg).to(torch.int64), self.num_classes() + 1).sum(dim=[0, 1])
+
+        return data_dict
+
+    def __getitem__(self, item):
+        try:
+            data_dict = self._base_get_item(item)
+            if self.use_hn:
+                indices = torch.cat([self.metadata["hn0"].iloc[item], self.metadata["hn1"].iloc[item]], dim=0)
+                neg_index = indices[torch.randint(0, indices.shape[0], (1,))]
+                negative_dict = self._base_get_item(neg_index)
+                data_dict["negatives"] = negative_dict
+            return data_dict
+        except (audioread.exceptions.NoBackendError, EOFError) as e:
+            # raise e
+            bad_path = self.metadata["audio_file"].iloc[item]
+            print(e)
+            print(f"Removing bad audio file {bad_path}")
+            # os.remove(bad_path)
+            return None
+        except ValueError as e:
+            # raise e
+            bad_path = self.metadata["audio_file"].iloc[item]
+            if "Input signal length=0" in str(e):
+                print(e)
+                print(f"Removing bad file {bad_path} due to input signal length=0")
+            #     os.remove(bad_path)
+            return None
+        except OSError as e:
+            # raise e
+            bad_paths = self.metadata["frame_files"].iloc[item]
+            for bad_path in bad_paths:
+                print(e)
+                print(f"Removing bad frame file {bad_path}")
+            return None
+        except RuntimeError as e:
+            # raise e
+            bad_path = self.metadata["audio_file"].iloc[item]
+            print(e)
+            print(f"Removing bad audio file {bad_path}")
+            # os.remove(bad_path)
+            return None
+
+
+class PlacesAudio(AVDataset):
+
+    def _load_info(self, split) -> pd.DataFrame:
+        df = pd.read_json(join(os.path.dirname(self._audio_root()), "metadata", f"{split}.json"))
+        df["id"] = df["data"].apply(lambda d: d["wav"][5:-4])
+
+        if self.use_caption:
+            if split == "train":
+                word_df = pd.read_json(
+                    join(os.path.dirname(self._audio_root()), "metadata", f"word-alignment-{split}.json")
+                )
+            else:
+                word_df = pd.read_csv(
+                    join(os.path.dirname(self._audio_root()), "metadata", f"word-alignment-{split}.csv")) \
+                    .groupby("id").aggregate(lambda g: list(g)).reset_index().drop("Unnamed: 0", axis=1)
+            df = pd.merge(df, word_df, on="id", how="outer")
+        return df
+
+    def _missing_threshold(self) -> float:
+        # return 0.0
+        return 0.97  # TODO fix
+
+    def _expected_num_frames(self):
+        return 1
+
+    def default_target_length(self) -> int:
+        return 20
+
+    def _frame_root(self) -> str:
+        return join(os.path.dirname(self.root), "places_subset")
+
+    def _audio_root(self) -> str:
+        return join(self.root, "wavs")
+
+    def _embed_root(self) -> str:
+        return join(self.root, "embedding", self.audio_embed_model)
+
+    def _dataset_folder(self) -> str:
+        return "PlacesAudio_400k_distro"
+
+    def _get_audio_file(self, row) -> str:
+        return join(self._audio_root(), row["id"] + ".wav")
+
+    def _get_frame_files(self, row) -> List[str]:
+        return [join(self._frame_root(), row["data"]["image"])]
+
+    def _get_embed_file(self, row) -> str:
+        return join(self._embed_root(), row["id"] + ".npz")
+
+
+class AudioSet(AVDataset):
+    def _expected_num_frames(self):
+        return 10
+
+    def default_target_length(self) -> int:
+        return 20
+
+    def _dataset_folder(self) -> str:
+        return "audioset-raw"
+
+    def _missing_threshold(self) -> float:
+        if self.split == "val" or self.split == "test":
+            return 0.02
+        else:
+            return 0.17
+
+    def train_seg_file(self):
+        return "unbalanced_train_segments.csv"
+
+    def _load_info(self, split) -> pd.DataFrame:
+        if split == "train":
+            df = pd.read_csv(join(self.root, "metadata", self.train_seg_file()))
+        elif split == "val" or split == "test":
+            df = pd.read_csv(join(self.root, "metadata", "eval_segments_subset.csv"))
+        else:
+            raise ValueError(f"Unknown split {split}")
+
+        labels = pd.read_csv(join(self.root, "metadata", "class_labels_indices.csv"))
+        mid_to_index = dict(zip(labels["mid"], labels["index"]))
+        df["tags"] = df["positive_labels"].apply(lambda l: [mid_to_index[e] for e in l.strip('"').split(",")])
+
+        self.num_tags = max(*[i for k, i in mid_to_index.items()]) + 1
+        df["id"] = df.apply(lambda r: f"{r.YTID}_{r.start_seconds}_{r.end_seconds}", axis=1)
+        return df
+
+    def _frame_root(self) -> str:
+        return join(self.root, "frames")
+
+    def _audio_root(self) -> str:
+        return join(self.root, "audio")
+
+    def _all_frame_files(self) -> Set[str]:
+        frame_files = set()
+
+        for entry in os.scandir(self._frame_root()):
+            if entry.is_file():
+                frame_files.add(entry.path)
+            elif entry.is_dir():
+                for subentry in os.scandir(entry.path):
+                    if subentry.is_file():
+                        frame_files.add(subentry.path)
+
+        return frame_files
+
+    def _all_audio_files(self) -> Set[str]:
+        return set(entry.path for entry in os.scandir(self._audio_root()) if entry.is_file())
+
+    def _all_embed_files(self) -> Set[str]:
+        return set(entry.path for entry in os.scandir(self._embed_root()) if entry.is_file())
+
+    def _embed_root(self) -> str:
+        return join(self.root, "embedding", self.audio_embed_model)
+
+    def prefix(self):
+        return ""
+
+    def _get_audio_file(self, row) -> str:
+        return f"{self.root}/audio/{self.prefix()}{row.id}.mp3"
+
+    def _get_frame_files(self, row) -> List[str]:
+        return [f"{self.root}/frames/frame_{fn}/{self.prefix()}{row.id}.jpg" for fn in range(10)]
+
+    def _get_embed_file(self, row) -> str:
+        return f"{self.root}/embedding/{self.audio_embed_model}/{self.prefix()}{row.id}.npz"
+
+
+class AudioSetEval(AudioSet):
+
+    def _dataset_folder(self) -> str:
+        return "audioset-eval"
+
+    def _get_frame_files(self, row) -> List[str]:
+        base_path = f"{self.root}/frames/{self.prefix()}{row.id}_"
+        return [base_path + f"{fn}.jpg" for fn in range(10)]
+
+    def prefix(self):
+        return ""
+
+
+class ADE20K(AVDataset):
+
+    def _split_name(self, split):
+        if split == "val":
+            return "validation"
+        elif split == "train":
+            return "training"
+        else:
+            raise ValueError(f"Unknown split name {split}")
+
+    def _load_info(self, split) -> pd.DataFrame:
+        df = pd.read_json(join(self.root, "metadata_with_caption_dedup.json"))
+        df["id"] = df["image"]
+        df = df[df["image"].apply(lambda f: f.split("/")[0] == split)]
+
+        if self.use_caption:
+            df["word"] = df["caption"].apply(lambda c: c["words"])
+            df["start"] = df["caption"].apply(lambda c: c["start"])
+            df["end"] = df["caption"].apply(lambda c: c["end"])
+            df["text"] = df["word"].apply(lambda l: " ".join(l))
+        return df
+
+    def _missing_threshold(self) -> float:
+        return 0.03
+
+    def _expected_num_frames(self):
+        return 1
+
+    def default_target_length(self) -> int:
+        return 20
+
+    def _dataset_folder(self) -> str:
+        return "ADE20K"
+
+    def _frame_root(self) -> str:
+        return join(self.root, "frames")
+
+    def _audio_root(self) -> str:
+        return join(self.root, "audio")
+
+    def _semseg_root(self) -> str:
+        return join(self.root, "annotations")
+
+    def _embed_root(self) -> str:
+        return join(self.root, "embedding", self.audio_embed_model)
+
+    def _get_audio_file(self, row) -> str:
+        return join(self._audio_root(), row["audio"])
+
+    def _get_frame_files(self, row) -> List[str]:
+        return [join(self._frame_root(), row["image"])]
+
+    def _get_semseg_file(self, row) -> str:
+        return join(self._semseg_root(), row["seg"])
+
+    def _get_embed_file(self, row) -> str:
+        return join(self._embed_root(), row["image"].replace(".jpg", ".npz"))
+
+    def num_classes(self):
+        return 3662
+
+
+class ADE20KPromptedBase(AVDataset):
+
+    def _expected_num_frames(self):
+        return 1
+
+    def default_target_length(self) -> int:
+        return 20
+
+    def _frame_root(self) -> str:
+        return join(self.root, "frames")
+
+    def _audio_root(self) -> str:
+        return join(self.root, "audio")
+
+    def _semseg_root(self) -> str:
+        return join(self.root, "annotations")
+
+    def _embed_root(self) -> str:
+        return join(self.root, "embedding", self.audio_embed_model)
+
+    def _get_frame_files(self, row) -> List[str]:
+        return [join(self._frame_root(), row["image_location"])]
+
+    def _get_semseg_file(self, row) -> str:
+        return join(self._semseg_root(), row["image_location"].replace(".jpg", "_seg.png"))
+
+    def _get_embed_file(self, row) -> str:
+        return join(self._embed_root(), row["image_location"].replace(".jpg", ".npz"))
+
+    def num_classes(self):
+        return 3662
+
+    def _missing_threshold(self) -> float:
+        return 0.0
+
+
+class ADE20KSpeechPrompted(ADE20KPromptedBase):
+
+    def _get_audio_file(self, row) -> str:
+        return join(self._audio_root(), row["speech_prompt_file"].split("/")[-1])
+
+    def _dataset_folder(self) -> str:
+        return "ADE20KSpeechPrompted"
+
+    def _audio_root(self) -> str:
+        # return join(self.root, "audio-noise-10") # TODO Remove
+        return join(self.root, "audio")  # TODO Remove
+
+    def _load_info(self, split) -> pd.DataFrame:
+        df = pd.read_csv(join(self.root, "prompted_segmentation.csv"))
+        df = df[df["speech_prompt_file"].apply(lambda s: isinstance(s, str))]
+        df = df[df["ade_class_id"].apply(lambda id: id != 0)]
+        df["id"] = df["image_location"]
+        return df
+
+
+class ADE20KSoundPrompted(ADE20KPromptedBase):
+
+    def _get_audio_file(self, row) -> str:
+        return join(self._audio_root(), row["vggsound_file"].split("/")[-1])
+
+    def _dataset_folder(self) -> str:
+        return "ADE20KSoundPrompted"
+
+    def _load_info(self, split) -> pd.DataFrame:
+        df = pd.read_csv(join(self.root, "prompted_segmentation.csv"))
+        df = df[df["vggsound_file"].apply(lambda s: isinstance(s, str))]
+        df = df[df["ade_class_id"].apply(lambda id: id != 0)]
+        df["id"] = df["image_location"]
+        return df
+
+
+class PlacesAndAudioSet(Dataset):
+
+    def __init__(self, **kwargs):
+        self.ds1 = PlacesAudio(**kwargs, n_frames=1)
+        self.ds2 = AudioSet(**kwargs, n_frames=1)
+
+    def __len__(self):
+        return len(self.ds1)
+
+    def __getitem__(self, item):
+        if torch.rand(1).item() > .5:
+            d = self.ds2[torch.randint(0, len(self.ds2) - 1, size=(1,)).item()]
+            if d is not None:
+                d["source"] = 1
+        else:
+            d = self.ds1[item]
+            if d is not None:
+                d["source"] = 0
+        return d
+
+
+class AVDataModule(pl.LightningDataModule):
+    def __init__(self,
+                 dataset_name,
+                 load_size,
+                 image_aug,
+                 audio_aug,
+                 extra_audio_masking,
+                 audio_model_type,
+                 pytorch_data_dir,
+                 use_cached_embs,
+                 batch_size,
+                 num_workers,
+                 audio_level,
+                 neg_audio,
+                 data_for_plotting,
+                 use_original_val_set,
+                 use_extra_val_sets,
+                 quad_mixup,
+                 bg_mixup,
+                 patch_mixup,
+                 patch_size,
+                 **kwargs):
+
+        super().__init__()
+        self.dataset_name = dataset_name
+        self.load_size = load_size
+        self.image_aug = image_aug
+        self.audio_aug = audio_aug
+        self.extra_audio_masking = extra_audio_masking
+        self.audio_model_type = audio_model_type
+        self.pytorch_data_dir = pytorch_data_dir
+        self.use_cached_embs = use_cached_embs
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.data_for_plotting = data_for_plotting
+        self.audio_level = audio_level
+        self.neg_audio = neg_audio
+
+        self.quad_mixup = quad_mixup
+        self.bg_mixup = bg_mixup
+        self.patch_mixup = patch_mixup
+        self.patch_size = patch_size
+
+        self.loader_args = dict(
+            num_workers=self.num_workers,
+            batch_size=self.batch_size,
+        )
+        self.save_hyperparameters()
+        self.extra_args = kwargs
+
+        self.use_original_val_set = use_original_val_set
+        self.use_extra_val_sets = use_extra_val_sets
+
+    def maybe_unpack(self, remove_source):
+        targets = [
+            (
+                join(self.pytorch_data_dir, "audioset-subset", "frame_archives"),
+                join(self.pytorch_data_dir, "audioset-subset", "frames"),
+                1
+            ),
+            (
+                join(self.pytorch_data_dir, "audioset-raw", "frame_archives"),
+                join(self.pytorch_data_dir, "audioset-raw", "frames"),
+                4
+            ),
+            (
+                join(self.pytorch_data_dir, "audioset-raw", "audio_archives"),
+                join(self.pytorch_data_dir, "audioset-raw", "audio"),
+                1
+            ),
+
+        ]
+
+        for (archive_dir, target_dir, n_parts) in targets:
+            if not os.path.exists(target_dir) and os.path.exists(archive_dir):
+                print(f"Could not find {target_dir}, attempting to unpack archives")
+                if os.path.exists(archive_dir):
+                    untar_all(archive_dir, target_dir, remove_source)
+                else:
+                    raise RuntimeError(f"Could not find archive folder: {archive_dir}")
+
+    def get_dataset_by_name(self, name, stage, data_for_plotting, n_frames=None):
+
+        if name == "vggss":
+            resize_op = T.Resize((self.load_size, self.load_size), Image.BILINEAR)
+        else:
+            resize_op = T.Resize(self.load_size, Image.BILINEAR)
+
+        img_transform = T.Compose([
+            resize_op,
+            T.CenterCrop(self.load_size),
+            T.ToTensor(),
+            norm])
+
+        if self.image_aug:
+            train_img_transform = T.Compose([
+                T.RandomResizedCrop(self.load_size),
+                T.RandomHorizontalFlip(),
+                T.ColorJitter(.2, .2, .2, .2),
+                T.RandomGrayscale(),
+                T.ToTensor(),
+                norm])
+            val_img_transform = img_transform
+        else:
+            train_img_transform = img_transform
+            val_img_transform = img_transform
+
+        if self.audio_aug:
+            train_audio_aug = True
+            val_audio_aug = False
+        else:
+            train_audio_aug = False
+            val_audio_aug = False
+
+        if self.audio_model_type == "hubert":
+            from featurizers.Hubert import HubertAudioTransform
+            audio_transform = HubertAudioTransform()
+        else:
+            audio_transform = None
+
+        if self.audio_model_type == "passt":
+            sample_rate = 32000
+        else:
+            sample_rate = 16000
+
+        if not self.use_cached_embs:
+            if self.audio_model_type == "hubert":
+                self.extra_args["use_audio"] = True
+            elif self.audio_model_type in {"audiomae", "audiomae-finetuned", "cavmae", "cavmae-mixed", "imagebind"}:
+                self.extra_args["use_spec"] = True
+            elif self.audio_model_type == "davenet":
+                self.extra_args["use_audio"] = True
+                self.extra_args["use_davenet_spec"] = True
+            elif self.audio_model_type == "fnac":
+                self.extra_args["use_audio"] = True
+                self.extra_args["use_fnac_spec"] = True
+            else:
+                raise ValueError(f"Unknown audio model type {self.audio_model_type}")
+
+            if self.audio_model_type == "cavmae" or self.audio_model_type == "cavmae-mixed":
+                self.extra_args["spec_mean"] = -5.081
+                self.extra_args["spec_std"] = 4.4849
+            elif self.audio_model_type == "imagebind":
+                self.extra_args["spec_mean"] = -4.268
+                self.extra_args["spec_std"] = 9.138
+
+        # if self.audio_model_type in {"audiomae", "audiomae-finetune", "cavmae"} \
+        #         and "override_target_length" not in self.extra_args:
+        if "override_target_length" not in self.extra_args:
+            self.extra_args["override_target_length"] = 10
+
+        data_args = dict(
+            root=self.pytorch_data_dir,
+            use_frames=True,
+            audio_transform=audio_transform,
+            sample_rate=sample_rate,
+            audio_level=self.audio_level,
+            **self.extra_args
+        )
+
+        if n_frames is not None:
+            data_args["n_frames"] = n_frames
+
+        train_args = dict(
+            frame_transform=train_img_transform,
+            extra_audio_masking=self.extra_audio_masking,
+            neg_audio=self.neg_audio,
+            quad_mixup=self.quad_mixup,
+            bg_mixup=self.bg_mixup,
+            patch_mixup=self.patch_mixup,
+            patch_size=self.patch_size,
+            audio_aug=train_audio_aug
+        )
+        val_args = dict(
+            frame_transform=val_img_transform,
+            audio_aug=val_audio_aug
+        )
+
+        if data_for_plotting:
+            val_args["use_audio"] = True
+            val_args["use_spec"] = True
+
+        if "ade" in name:
+            label_transform = T.Compose([
+                T.Resize(self.load_size, Image.NEAREST),
+                T.CenterCrop(self.load_size),
+                prep_ade_label
+            ])
+        else:
+            label_transform = T.Compose([
+                T.Resize(self.load_size, Image.NEAREST),
+                T.CenterCrop(self.load_size)
+            ])
+
+        val_args["use_audio"] = True
+        val_args["label_transform"] = label_transform
+
+        if name == "places-audio":
+            dataset_constructor = PlacesAudio
+        elif name == "mixed-full":
+            dataset_constructor = PlacesAndAudioSet
+        elif name == "audio-set-full":
+            dataset_constructor = AudioSet
+        elif name == "audio-set-eval":
+            dataset_constructor = AudioSetEval
+        elif name == "ade":
+            val_args["use_semseg"] = True
+            dataset_constructor = ADE20K
+        elif name == "ade-speech-prompted":
+            val_args["use_semseg"] = True
+            dataset_constructor = ADE20KSpeechPrompted
+        elif name == "ade-sound-prompted":
+            val_args["use_semseg"] = True
+            dataset_constructor = ADE20KSoundPrompted
+        else:
+            raise ValueError(f"Unknown dataset name {name}")
+
+        data_args["use_audio_embed"] = self.use_cached_embs
+        data_args["audio_embed_model"] = self.audio_model_type
+
+        if stage == "full":
+            val_dataset = dataset_constructor(split="val", **{**data_args, **val_args})
+            train_dataset = dataset_constructor(split="train", **{**data_args, **val_args})
+            return ConcatDataset([train_dataset, val_dataset])
+        elif stage == "fit":
+            return dataset_constructor(split="train", **{**data_args, **train_args})
+        elif stage == "validate":
+            return dataset_constructor(split="val", **{**data_args, **val_args})
+        else:
+            raise ValueError(f"Unknown stage: {stage}")
+
+    def _maybe_subset(self, dataset, length):
+        if len(dataset) > length and self.dataset_name not in {"ade-sound-prompted", "ade-speech-prompted", "vggss"}:
+            print("Using a subset of validation data")
+            return Subset(dataset, generate_subset(len(dataset), length))
+        else:
+            print("Not using val subset")
+            return dataset
+
+    def _make_val_datasets(self):
+        val_sets = []
+        if self.use_original_val_set:
+            val_sets.append(self._maybe_subset(self.get_dataset_by_name(
+                self.dataset_name, "validate", self.data_for_plotting), 1000))
+
+        if self.use_extra_val_sets:
+            val_sets.append(self._maybe_subset(self.get_dataset_by_name(
+                "places-audio", "validate", self.data_for_plotting), 1000))
+            val_sets.append(self._maybe_subset(self.get_dataset_by_name(
+                "audio-set-eval", "validate", False, n_frames=1), 1000))
+            val_sets.append(self.get_dataset_by_name(
+                "ade-speech-prompted", "validate", True))
+            val_sets.append(self.get_dataset_by_name(
+                "ade-sound-prompted", "validate", self.data_for_plotting))
+
+        return val_sets
+
+    def setup(self, stage: str):
+        if stage == "full":
+            self.full_dataset = self.get_dataset_by_name(self.dataset_name, stage, self.data_for_plotting)
+        elif stage == "fit":
+            self.train_dataset = self.get_dataset_by_name(self.dataset_name, stage, self.data_for_plotting)
+            self.val_datasets = self._make_val_datasets()
+        elif stage == "validate":
+            self.val_datasets = self._make_val_datasets()
+        else:
+            raise ValueError(f"Unknown stage: {stage}")
+
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset, shuffle=True, **self.loader_args, collate_fn=custom_coallate)
+
+    def subsampled_train_dataloader(self, k=5000):
+        if len(self.train_dataset) > k:
+            ds = Subset(self.train_dataset, generate_subset(len(self.train_dataset), k))
+        else:
+            ds = self.train_dataset
+
+        return DataLoader(ds, shuffle=True, **self.loader_args, collate_fn=custom_coallate)
+
+    def val_dataloader(self):
+        return [
+            DataLoader(dataset, shuffle=False, **self.loader_args, collate_fn=custom_coallate)
+            for dataset in self.val_datasets
+        ]
+
+    def full_dataloader(self):
+        return DataLoader(self.full_dataset, shuffle=False, **self.loader_args, collate_fn=custom_coallate)
+
+
+def generate_subset(n, batch, seed=0):
+    np.random.seed(seed)
+    return np.random.permutation(n)[:batch]
+
+
+def prep_ade_label(img):
+    seg = np.array(img)
+    class_labels = (seg[:, :, 0] / 10).astype(np.int32) * 256 + (seg[:, :, 1].astype(np.int32))
+    return class_labels
+
+
+def maybe_replace(e, not_none):
+    if e is not None:
+        return e
+    else:
+        print("Warning found a None in the dataset indicitive of a loading failure, replacing it with another item")
+        return not_none[0]
+
+
+empty_caption = {
+    "words": [],
+    "start": [],
+    "end": [],
+}
+
+
+def custom_coallate(l):
+    if l is None:
+        return l
+
+    not_none = [e for e in l if e is not None]
+    assert len(not_none) > 0
+
+    l = [maybe_replace(e, not_none) for e in l]
+
+    to_merge = {}
+
+    def pop_or_default(dict, k, default):
+        if k in dict:
+            return dict.pop(k)
+        else:
+            print(f"WARNING: Could not find {k}, using {default}")
+            return default
+
+    if "caption" in l[0]:
+        to_merge["caption"] = [pop_or_default(l[i], "caption", empty_caption) for i in range(len(l))]
+
+    if "text" in l[0]:
+        to_merge["text"] = [pop_or_default(l[i], "text", "") for i in range(len(l))]
+
+    result = default_collate(l)
+
+    return {**result, **to_merge}
+
+
+if __name__ == "__main__":
+
+    from featurizers.Hubert import HubertAudioTransform
+
+    pytorch_data_dir = "/pytorch-data"
+    dataset_constructor = PlacesAudio
+    split = "val"
+
+    img_transform = T.Compose([
+        T.Resize(224, Image.BILINEAR),
+        T.CenterCrop(224),
+        T.ToTensor(),
+        norm])
+
+    video_transform = T.Compose([
+        T.Resize(224, Image.BILINEAR),
+        T.CenterCrop(224),
+        norm])
+
+    label_transform = T.Compose([
+        T.Resize(224, Image.NEAREST),
+        T.CenterCrop(224)
+    ])
+
+    audio_transform = HubertAudioTransform()
+
+    data_args = dict(
+        root=pytorch_data_dir,
+        frame_transform=img_transform,
+        use_frames=True,
+        use_spec=True,
+        use_audio=True,
+        use_caption=False,
+        use_semseg=False,
+        label_transform=label_transform,
+        audio_transform=audio_transform,
+        use_audio_embed=False,
+        audio_embed_model="audiomae",
+        extra_audio_masking=False,
+        neg_audio=False,
+        override_target_length=10,
+        audio_level=False,
+        quad_mixup=.3,
+        patch_mixup=.3,
+        bg_mixup=.3,
+    )
+
+
+    def return_datasets(dataset_constructor, split):
+        dataset = dataset_constructor(split=split, **data_args)
+        return dataset
+
+
+    train_ds = return_datasets(dataset_constructor, split)
+
+    print(len(train_ds))
+    train_loader = DataLoader(train_ds, batch_size=1, shuffle=False, num_workers=36, collate_fn=custom_coallate)
+    for batch in tqdm(train_loader):
+        pass

DenseAV/denseav/data/make_tarballs.py

@@ -0,0 +1,108 @@
+import glob
+import os
+import tarfile
+from glob import glob
+from io import BytesIO
+from os.path import join
+
+from torch.utils.data import Dataset, DataLoader
+from tqdm import tqdm
+from pathlib import Path
+
+from denseav.shared import batch
+
+import tempfile
+import shutil
+
+
+class Tarballer(Dataset):
+
+    def __init__(self, source, target, n):
+        source_path = Path(source)
+        self.frames = [f.relative_to(source_path) for f in source_path.rglob('*') if f.is_file()]
+        assert (len(self.frames) > 0)
+        self.source = source
+        self.target_dir = target
+        self.batched = list(batch(self.frames, n))
+        os.makedirs(self.target_dir, exist_ok=True)
+
+    def __len__(self):
+        return len(self.batched)
+
+    def __getitem__(self, item):
+        with tarfile.open(join(self.target_dir, f"{item}.tar"), "w") as tar:
+            for relpath in self.batched[item]:
+                abs_path = os.path.join(self.source, str(relpath))  # Convert to string here
+                with open(abs_path, "rb") as file:
+                    file_content = file.read()
+                info = tarfile.TarInfo(name=str(relpath))  # Convert to string here
+                info.size = len(file_content)
+                tar.addfile(info, fileobj=BytesIO(file_content))
+
+        return 0
+
+
+class UnTarballer:
+
+    def __init__(self, archive_dir, target_dir, remove_source=False):
+        self.tarballs = sorted(glob(join(archive_dir, "*.tar")))
+        self.target_dir = target_dir
+        self.remove_source = remove_source  # New flag to determine if source tarball should be removed
+        os.makedirs(self.target_dir, exist_ok=True)
+
+    def __len__(self):
+        return len(self.tarballs)
+
+    def __getitem__(self, item):
+        with tarfile.open(self.tarballs[item], "r") as tar:
+            # Create a unique temporary directory inside the target directory
+            with tempfile.TemporaryDirectory(dir=self.target_dir) as tmpdirname:
+                tar.extractall(tmpdirname)  # Extract to the temporary directory
+
+                # Move contents from temporary directory to final target directory
+                for src_dir, dirs, files in os.walk(tmpdirname):
+                    dst_dir = src_dir.replace(tmpdirname, self.target_dir, 1)
+                    os.makedirs(dst_dir, exist_ok=True)
+                    for file_ in files:
+                        src_file = os.path.join(src_dir, file_)
+                        dst_file = os.path.join(dst_dir, file_)
+                        shutil.move(src_file, dst_file)
+
+        # Remove the source tarball if the flag is set to True
+        if self.remove_source:
+            os.remove(self.tarballs[item])
+
+        return 0
+
+def untar_all(archive_dir, target_dir, remove_source):
+    loader = DataLoader(UnTarballer(archive_dir, target_dir, remove_source), num_workers=24)
+    for _ in tqdm(loader):
+        pass
+
+
+if __name__ == "__main__":
+    # loader = DataLoader(Tarballer(
+    #     join("/pytorch-data", "audioset-raw", "audio"),
+    #     join("/pytorch-data", "audioset-raw", "audio_archives")
+    # ), num_workers=24)
+
+    # loader = DataLoader(Tarballer(
+    #     join("/pytorch-data", "audioset-raw", "frames"),
+    #     join("/pytorch-data", "audioset-raw", "frame_archives"),
+    #     5000
+    # ), num_workers=24)
+
+    # loader = DataLoader(Tarballer(
+    #     join("/pytorch-data", "ADE20KLabels"),
+    #     join("/pytorch-data", "ADE20KLabelsAr"),
+    #     100
+    # ), num_workers=24)
+    #
+    # for _ in tqdm(loader):
+    #     pass
+    #
+    # #
+    #
+    untar_all(
+        join("/pytorch-data", "audioset-raw", "frame_archives"),
+        join("/pytorch-data", "audioset-raw", "frames_4"))

DenseAV/denseav/eval_utils.py

@@ -0,0 +1,135 @@
+import json
+from collections import defaultdict
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torchmetrics.functional.classification import binary_average_precision
+from tqdm import tqdm
+
+from constants import *
+from denseav.shared import unnorm, remove_axes
+
+
+def prep_heatmap(sims, masks, h, w):
+    masks = masks.to(torch.float32)
+    hm = torch.einsum("bhwt,bt->bhw", sims, masks) / masks.sum(-1).reshape(-1, 1, 1)
+    hm -= hm.min()
+    hm /= hm.max()
+    return F.interpolate(hm.unsqueeze(1), (h, w), mode="bilinear").squeeze(1)
+
+
+def iou(prediction, target):
+    prediction = prediction > 0.0
+    target = target > 0.5
+    intersection = torch.logical_and(prediction, target).sum().float()
+    union = torch.logical_or(prediction, target).sum().float()
+    if union == 0:
+        return 1.0
+    return (intersection / union).item()  # Convert to Python scalar
+
+
+def multi_iou(prediction, target, k=20):
+    prediction = torch.tensor(prediction)
+    target = torch.tensor(target)
+    target = target > 0.5
+
+    thresholds = torch.linspace(prediction.min(), prediction.max(), k)
+    hard_pred = prediction.unsqueeze(0) > thresholds.reshape(k, 1, 1, 1, 1)
+    target = torch.broadcast_to(target.unsqueeze(0), hard_pred.shape)
+
+    # Calculate IoU for each threshold
+    intersection = torch.logical_and(hard_pred, target).sum(dim=(1, 2, 3, 4)).float()
+    union = torch.logical_or(hard_pred, target).sum(dim=(1, 2, 3, 4)).float()
+    union = torch.where(union == 0, torch.tensor(1.0), union)  # Avoid division by zero
+    iou_scores = intersection / union
+
+    # Find the best IoU and corresponding threshold
+    best_iou, best_idx = torch.max(iou_scores, dim=0)
+    # best_threshold = thresholds[best_idx]
+    # print(best_threshold)
+    return best_iou  # , best_threshold.item()
+
+
+def get_paired_heatmaps(
+        model,
+        results,
+        class_ids,
+        timing,
+        class_names=None):
+    sims = model.sim_agg.get_pairwise_sims(
+        results,
+        raw=False,
+        agg_sim=False,
+        agg_heads=True
+    ).squeeze(1).mean(-2)
+
+    prompt_classes = torch.tensor(list(class_ids))
+    gt = results["semseg"] == prompt_classes.reshape(-1, 1, 1)
+    basic_masks = results[AUDIO_MASK]  # BxT
+    _, fullh, fullw = gt.shape
+    basic_heatmaps = prep_heatmap(sims, basic_masks, fullh, fullw)
+
+    if timing is not None:
+        prompt_timing = np.array(list(timing))
+        raw_timing = torch.tensor([json.loads(t) for t in prompt_timing])
+        timing = torch.clone(raw_timing)
+        timing[:, 0] -= .2
+        timing[:, 1] += .2
+        total_length = (results['total_length'] / 16000)[0]
+        fracs = timing / total_length
+        bounds = basic_masks.shape[1] * fracs
+        bounds[:, 0] = bounds[:, 0].floor()
+        bounds[:, 1] = bounds[:, 1].ceil()
+        bounds = bounds.to(torch.int64)
+        advanced_masks = (F.one_hot(bounds, basic_masks.shape[1]).cumsum(-1).sum(-2) == 1).to(basic_masks)
+        advanced_heatmaps = prep_heatmap(sims, advanced_masks, fullh, fullw)
+
+    metrics = defaultdict(list)
+    unique_classes = torch.unique(prompt_classes)
+
+    should_plot = class_names is not None
+
+    if should_plot:
+        prompt_names = np.array(list(class_names))
+
+    for prompt_class in tqdm(unique_classes):
+        subset = torch.where(prompt_classes == prompt_class)[0]
+        gt_subset = gt[subset]
+        basic_subset = basic_heatmaps[subset]
+        metrics["basic_ap"].append(binary_average_precision(basic_subset.flatten(), gt_subset.flatten()))
+        metrics["basic_iou"].append(multi_iou(basic_subset.flatten(), gt_subset.flatten()))
+
+        if timing is not None:
+            advanced_subset = advanced_heatmaps[subset]
+            metrics["advanced_ap"].append(binary_average_precision(advanced_subset.flatten(), gt_subset.flatten()))
+            metrics["advanced_iou"].append(multi_iou(advanced_subset.flatten(), gt_subset.flatten()))
+
+        if should_plot:
+            prompt_class_subset = prompt_classes[subset]
+            name_subset = prompt_names[subset]
+            print(prompt_class, name_subset, prompt_class_subset)
+            n_imgs = min(len(subset), 5)
+            if n_imgs > 1:
+                fig, axes = plt.subplots(n_imgs, 5, figsize=(4 * 5, n_imgs * 3))
+                frame_subset = unnorm(results[IMAGE_INPUT][subset].squeeze(1)).permute(0, 2, 3, 1)
+                semseg_subset = results["semseg"][subset]
+                for img_num in range(n_imgs):
+                    axes[img_num, 0].imshow(frame_subset[img_num])
+                    axes[img_num, 1].imshow(basic_subset[img_num])
+                    axes[img_num, 2].imshow(advanced_subset[img_num])
+                    axes[img_num, 3].imshow(gt_subset[img_num])
+                    axes[img_num, 4].imshow(semseg_subset[img_num], cmap="tab20", interpolation='none')
+
+                axes[0, 0].set_title("Image")
+                class_name = name_subset[0].split(",")[0]
+                axes[0, 1].set_title(f"{class_name} Basic Heatmap")
+                axes[0, 2].set_title(f"{class_name} Advanced Heatmap")
+                axes[0, 3].set_title("True Mask")
+                axes[0, 4].set_title("True Seg")
+                remove_axes(axes)
+                plt.tight_layout()
+                plt.show()
+
+    return metrics, unique_classes

DenseAV/denseav/evaluate.py

@@ -0,0 +1,87 @@
+from os.path import join
+import hydra
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+from pytorch_lightning import seed_everything
+from pytorch_lightning.loggers import TensorBoardLogger
+from denseav.data.AVDatasets import AVDataModule
+from denseav.shared import load_trained_model
+
+
+@hydra.main(config_path="configs", config_name="av_align.yaml")
+def my_app(cfg: DictConfig) -> None:
+    from saved_models import saved_model_dict
+
+    seed_everything(0)
+    print(OmegaConf.to_yaml(cfg))
+
+    models_to_eval = [
+        "denseav_language",
+        "denseav_sound",
+    ]
+
+    checkpoint_dir = "../checkpoints"
+    saved_models = saved_model_dict(checkpoint_dir)
+    for model_name in models_to_eval:
+        model_info = saved_models[model_name]
+        extra_data_args = model_info["data_args"] if "data_args" in model_info else {}
+        model_info["extra_args"]["output_root"] = "../"
+        model_info["extra_args"]["neg_audio"] = False
+        model_info["extra_args"]["image_mixup"] = 0.0
+
+        model = load_trained_model(join(checkpoint_dir, model_info["chkpt_name"]), model_info["extra_args"])
+        model.set_full_train(True)
+
+        if model.image_model_type == "dinov2":
+            load_size = cfg.load_size * 2
+        else:
+            load_size = cfg.load_size
+
+        if model.image_model_type == "davenet":
+            batch_size = cfg.batch_size // 2
+        elif model.image_model_type == "imagebind":
+            batch_size = cfg.batch_size
+        else:
+            batch_size = cfg.batch_size
+
+        print(load_size)
+
+        data_args = dict(
+            dataset_name=cfg.dataset_name,
+            load_size=load_size,
+            image_aug=cfg.image_aug,
+            audio_aug=cfg.audio_aug,
+            audio_model_type=model.audio_model_type,
+            pytorch_data_dir=cfg.pytorch_data_dir,
+            use_cached_embs=model.use_cached_embs,
+            batch_size=batch_size,
+            num_workers=cfg.num_workers,
+            extra_audio_masking=False,
+            use_original_val_set=False,
+            use_extra_val_sets=True,
+            use_caption=True,
+            data_for_plotting=False,
+            n_frames=None,
+            audio_level=False,
+            neg_audio=False,
+            quad_mixup=0.0,
+            bg_mixup=0.0,
+            patch_mixup=0.0,
+            patch_size=8,
+        )
+        data_args = {**data_args, **extra_data_args}
+
+        datamodule = AVDataModule(**data_args)
+        log_dir = join(cfg.output_root, "logs", "evaluate", model_name)
+        print(log_dir)
+        tb_logger = TensorBoardLogger(log_dir, default_hp_metric=False)
+        trainer = Trainer(
+            accelerator='gpu',
+            strategy="ddp",
+            devices=cfg.num_gpus,
+            logger=tb_logger)
+        trainer.validate(model, datamodule)
+
+
+if __name__ == "__main__":
+    my_app()

DenseAV/denseav/featurizers/AudioMAE.py

@@ -0,0 +1,570 @@
+import math
+import os
+import warnings
+from functools import partial
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from timm.models.layers import to_2tuple
+from torch.utils.data import Dataset
+from torchaudio.functional import resample
+import pickle
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    r"""Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.patch_hw = (img_size[1] // patch_size[1], img_size[0] // patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        # assert H == self.img_size[0] and W == self.img_size[1], \
+        #    f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class HybridEmbed(nn.Module):
+    """ CNN Feature Map Embedding
+    Extract feature map from CNN, flatten, project to embedding dim.
+    """
+
+    def __init__(self, backbone, img_size=224, feature_size=None, in_chans=3, embed_dim=768):
+        super().__init__()
+        assert isinstance(backbone, nn.Module)
+        img_size = to_2tuple(img_size)
+        self.img_size = img_size
+        self.backbone = backbone
+        if feature_size is None:
+            with torch.no_grad():
+                # FIXME this is hacky, but most reliable way of determining the exact dim of the output feature
+                # map for all networks, the feature metadata has reliable channel and stride info, but using
+                # stride to calc feature dim requires info about padding of each stage that isn't captured.
+                training = backbone.training
+                if training:
+                    backbone.eval()
+                o = self.backbone(torch.zeros(1, in_chans, img_size[0], img_size[1]))[-1]
+                feature_size = o.shape[-2:]
+                feature_dim = o.shape[1]
+                backbone.train(training)
+        else:
+            feature_size = to_2tuple(feature_size)
+            feature_dim = self.backbone.feature_info.channels()[-1]
+        self.num_patches = feature_size[0] * feature_size[1]
+        self.proj = nn.Linear(feature_dim, embed_dim)
+
+    def forward(self, x):
+        x = self.backbone(x)[-1]
+        x = x.flatten(2).transpose(1, 2)
+        x = self.proj(x)
+        return x
+
+
+class TimmVisionTransformer(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., hybrid_backbone=None, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+
+        if hybrid_backbone is not None:
+            self.patch_embed = HybridEmbed(
+                hybrid_backbone, img_size=img_size, in_chans=in_chans, embed_dim=embed_dim)
+        else:
+            self.patch_embed = PatchEmbed(
+                img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+
+        # NOTE as per official impl, we could have a pre-logits representation dense layer + tanh here
+        # self.repr = nn.Linear(embed_dim, representation_size)
+        # self.repr_act = nn.Tanh()
+
+        # Classifier head
+        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def get_classifier(self):
+        return self.head
+
+    def reset_classifier(self, num_classes, global_pool=''):
+        self.num_classes = num_classes
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    def forward_features(self, x):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+
+        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x = self.norm(x)
+        return x[:, 0]
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head(x)
+        return x
+
+
+class VisionTransformer(TimmVisionTransformer):
+    """ Vision Transformer with support for global average pooling
+    """
+
+    def __init__(self, **kwargs):
+        super(VisionTransformer, self).__init__(**kwargs)
+        norm_layer = kwargs['norm_layer']
+        embed_dim = kwargs['embed_dim']
+        self.fc_norm = norm_layer(embed_dim)
+        del self.norm  # remove the original norm
+
+    def interpolate_pos_encoding(self, x, embed):
+        new_patches = x.shape[1]
+        old_patches = embed.shape[1]
+
+        w = 8
+        h = int(new_patches / w)
+        if new_patches == old_patches:
+            return embed
+
+        dim = x.shape[-1]
+        pos_embed = nn.functional.interpolate(
+            embed.reshape(1, 64, 8, dim).permute(0, 3, 1, 2),
+            size=(h, w),
+            mode='bicubic',
+        )
+        pos_embed = pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return pos_embed
+
+    def forward(self, x):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+
+        x = x + self.interpolate_pos_encoding(x, self.pos_embed[:, 1:, :])
+
+        cls_token = self.cls_token + self.pos_embed[:, :1, :]
+        cls_tokens = cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = self.pos_drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        # x = x[:, 1:, :].mean(dim=1)  # global pool without cls token
+        # outcome = self.fc_norm(x)
+
+        return x[:, 1:, :].reshape(B, -1, 8, 768).permute(0, 3, 2, 1), x[:, 0]
+
+
+class NewPatchEmbed(nn.Module):
+    """ Flexible Image to Patch Embedding
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, stride=10):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        stride = to_2tuple(stride)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride)  # with overlapped patches
+        _, _, h, w = self.get_output_shape(img_size)  # n, emb_dim, h, w
+        self.patch_hw = (h, w)
+        self.num_patches = h * w
+
+    def get_output_shape(self, img_size):
+        # todo: don't be lazy..
+        return self.proj(torch.randn(1, 1, img_size[0], img_size[1])).shape
+
+    def forward(self, x):
+        x = self.proj(x)
+        x = x.flatten(2).transpose(1, 2)
+        return x
+
+
+def pca(image_feats_list, dim=3, fit_pca=None):
+    from sklearn.decomposition import PCA
+
+    device = image_feats_list[0].device
+
+    def flatten(tensor, target_size=None):
+        if target_size is not None and fit_pca is None:
+            F.interpolate(tensor, (target_size, target_size), mode="bilinear")
+        B, C, H, W = tensor.shape
+        return feats.permute(1, 0, 2, 3).reshape(C, B * H * W).permute(1, 0).detach().cpu()
+
+    if len(image_feats_list) > 1 and fit_pca is None:
+        target_size = image_feats_list[0].shape[2]
+    else:
+        target_size = None
+
+    flattened_feats = []
+    for feats in image_feats_list:
+        flattened_feats.append(flatten(feats, target_size))
+    x = torch.cat(flattened_feats, dim=0)
+
+    if fit_pca is None:
+        fit_pca = PCA(n_components=dim, svd_solver="arpack").fit(np.nan_to_num(x.detach().numpy()))
+
+    reduced_feats = []
+    for feats in image_feats_list:
+        x_red = torch.from_numpy(fit_pca.transform(flatten(feats)))
+        x_red -= x_red.min(dim=0, keepdim=True).values
+        x_red /= x_red.max(dim=0, keepdim=True).values
+        B, C, H, W = feats.shape
+        reduced_feats.append(x_red.reshape(B, H, W, dim).permute(0, 3, 1, 2).to(device))
+
+    return reduced_feats, fit_pca
+
+
+class AudiosetDataset(Dataset):
+    def __init__(self, audio_conf):
+        self.audio_conf = audio_conf
+        self.melbins = self.audio_conf.get('num_mel_bins')
+        self.dataset = self.audio_conf.get('dataset')
+        self.norm_mean = self.audio_conf.get('mean')
+        self.norm_std = self.audio_conf.get('std')
+
+        print('Dataset: {}, mean {:.3f} and std {:.3f}'.format(self.dataset, self.norm_mean, self.norm_std))
+        print(f'size of dataset {self.__len__()}')
+
+    def _wav2fbank(self, filename):
+        sample_rate = 16000
+        target_length = 10
+        waveform, obs_sr = torchaudio.load(filename)
+        waveform = waveform[0]
+        if obs_sr != sample_rate:
+            waveform = resample(waveform, obs_sr, sample_rate)
+
+        original_length = waveform.shape[0]
+        padding = target_length * sample_rate - original_length
+
+        if padding > 0:
+            m = torch.nn.ZeroPad2d((0, padding))
+            waveform = m(waveform)
+        else:
+            waveform = waveform[:target_length * sample_rate]
+
+
+        waveform = waveform - waveform.mean()
+
+        # 498 128, 998, 128
+        fbank = torchaudio.compliance.kaldi.fbank(
+            waveform.unsqueeze(0),
+            htk_compat=True,
+            sample_frequency=sample_rate,
+            use_energy=False,
+            window_type='hanning',
+            num_mel_bins=128,
+            dither=0.0,
+            frame_shift=10)
+
+        normed_fbank = (fbank - self.norm_mean) / (self.norm_std * 2)
+
+        return normed_fbank
+
+    def __getitem__(self, index):
+        datum = {"wav": "../../samples/example.wav"}
+        fbank = self._wav2fbank(datum['wav'])
+        fbank = fbank.transpose(0, 1).unsqueeze(0)  # 1, 128, 1024 (...,freq,time)
+        fbank = torch.transpose(fbank.squeeze(), 0, 1)  # time, freq
+        # the output fbank shape is [time_frame_num, frequency_bins], e.g., [1024, 128]
+        return fbank.unsqueeze(0)
+
+    def __len__(self):
+        return 1
+
+
+class AudioMAE(nn.Module):
+
+    def __init__(self, output_path, finetuned):
+        super().__init__()
+        # build model
+        model = VisionTransformer(
+            patch_size=16,
+            embed_dim=768,
+            depth=12,
+            num_heads=12,
+            mlp_ratio=4,
+            qkv_bias=True,
+            norm_layer=partial(nn.LayerNorm, eps=1e-6),
+            num_classes=527,
+            drop_path_rate=0.1)
+
+        img_size = (1024, 128)  # 1024, 128
+        emb_dim = 768
+        model.patch_embed = NewPatchEmbed(
+            img_size=img_size, patch_size=(16, 16), in_chans=1, embed_dim=emb_dim, stride=16)
+        num_patches = model.patch_embed.num_patches
+        model.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, emb_dim), requires_grad=False)
+
+        if finetuned:
+            fn = "audiomae_finetuned.pth"
+        else:
+            fn = "audiomae.pth"
+
+        checkpoint = torch.load(os.path.join(output_path, 'models', fn), map_location='cpu')
+
+        checkpoint_model = checkpoint['model']
+        state_dict = model.state_dict()
+        for k in ['head.weight', 'head.bias']:
+            if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape:
+                print(f"Removing key {k} from pretrained checkpoint")
+                del checkpoint_model[k]
+        msg = model.load_state_dict(checkpoint_model, strict=False)
+        print(msg)
+
+        model = model.eval()
+        self.model = model
+        self.config = dict(output_path=output_path, finetuned=finetuned)
+
+    def forward(self, audio, include_cls):
+        patch_tokens, cls_token = self.model(audio)
+
+        if include_cls:
+            return patch_tokens, cls_token
+        else:
+            return patch_tokens
+
+
+if __name__ == '__main__':
+    import os
+
+    device = torch.device("cuda:2")
+
+    torch.manual_seed(0)
+    np.random.seed(0)
+
+    model = AudioMAE("../../", True).to(device)
+
+    audio_conf_val = {
+        'num_mel_bins': 128,
+        'target_length': 1024,
+        'dataset': "audioset",
+        'mode': 'val',
+        'mean': -4.2677393,
+        'std': 4.5689974,
+    }
+
+    dataset = AudiosetDataset(audio_conf=audio_conf_val)
+
+    batch = dataset[0].unsqueeze(0).to(device)
+
+    embeddings = model(batch, include_cls=False)
+
+    import matplotlib.pyplot as plt
+
+    with torch.no_grad():
+        [pca_feats], _ = pca([embeddings])
+        plt.imshow(pca_feats.cpu().squeeze(0).permute(1, 2, 0))
+        plt.show()
+        print("here")
+
+    print("here")

DenseAV/denseav/featurizers/CAVMAE.py

@@ -0,0 +1,1082 @@
+import random
+
+import numpy as np
+import timm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+import torchvision.transforms as T
+from PIL import Image
+from timm.models.layers import to_2tuple, DropPath
+from timm.models.vision_transformer import Mlp, PatchEmbed, Block
+import os
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_h_size, grid_w_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_h_size, dtype=float)
+    grid_w = np.arange(grid_w_size, dtype=float)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_w_size, grid_h_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=float)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000 ** omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+# --------------------------------------------------------
+# Interpolate position embeddings for high-resolution
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+def interpolate_pos_embed(model, checkpoint_model):
+    if 'pos_embed' in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model['pos_embed']
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model['pos_embed'] = new_pos_embed
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.norm1_a = norm_layer(dim)
+        self.norm1_v = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        self.norm2_a = norm_layer(dim)
+        self.norm2_v = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x, modality=None):
+        if modality == None:
+            x = x + self.drop_path(self.attn(self.norm1(x)))
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        elif modality == 'a':
+            x = x + self.drop_path(self.attn(self.norm1_a(x)))
+            x = x + self.drop_path(self.mlp(self.norm2_a(x)))
+        elif modality == 'v':
+            x = x + self.drop_path(self.attn(self.norm1_v(x)))
+            x = x + self.drop_path(self.mlp(self.norm2_v(x)))
+        return x
+
+
+# our main proposed model, for pretraining only, for finetuning, use CAVMAEFT class
+class CAVMAE(nn.Module):
+    """ CAV-MAE Model
+    """
+
+    def __init__(self, img_size=224, audio_length=1024, patch_size=16, in_chans=3,
+                 embed_dim=768, modality_specific_depth=11, num_heads=12,
+                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+                 mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=False, tr_pos=False):
+        super().__init__()
+        print('A CAV-MAE Model')
+        print('Use norm_pix_loss: ', norm_pix_loss)
+        print('Learnable Positional Embedding: ', tr_pos)
+
+        # the encoder part
+        # overide the timm package
+        timm.models.vision_transformer.PatchEmbed = PatchEmbed
+        timm.models.vision_transformer.Block = Block
+
+        self.patch_embed_a = PatchEmbed(img_size, patch_size, 1, embed_dim)
+        self.patch_embed_v = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
+
+        self.patch_embed_a.num_patches = int(audio_length * 128 / 256)
+        print('Number of Audio Patches: {:d}, Visual Patches: {:d}'.format(self.patch_embed_a.num_patches,
+                                                                           self.patch_embed_v.num_patches))
+
+        self.modality_a = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.modality_v = nn.Parameter(torch.zeros(1, 1, embed_dim))
+
+        self.pos_embed_a = nn.Parameter(torch.zeros(1, self.patch_embed_a.num_patches, embed_dim),
+                                        requires_grad=tr_pos)  # fixed sin-cos embedding
+        self.pos_embed_v = nn.Parameter(torch.zeros(1, self.patch_embed_v.num_patches, embed_dim),
+                                        requires_grad=tr_pos)  # fixed sin-cos embedding
+
+        # audio-branch
+        self.blocks_a = nn.ModuleList(
+            [Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer) for i in
+             range(modality_specific_depth)])
+        # visual-branch
+        self.blocks_v = nn.ModuleList(
+            [Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer) for i in
+             range(modality_specific_depth)])
+        # unified branch
+        self.blocks_u = nn.ModuleList(
+            [Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer) for i in
+             range(12 - modality_specific_depth)])
+
+        # independent normalization layer for audio, visual, and audio-visual
+        self.norm_a, self.norm_v, self.norm = norm_layer(embed_dim), norm_layer(embed_dim), norm_layer(embed_dim)
+
+        # the decoder part
+        # Project to lower dimension for the decoder
+        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
+
+        # token used for masking
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
+
+        self.decoder_modality_a = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
+        self.decoder_modality_v = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
+
+        self.decoder_pos_embed_a = nn.Parameter(torch.zeros(1, self.patch_embed_a.num_patches, decoder_embed_dim),
+                                                requires_grad=tr_pos)  # fixed sin-cos embedding
+        self.decoder_pos_embed_v = nn.Parameter(torch.zeros(1, self.patch_embed_v.num_patches, decoder_embed_dim),
+                                                requires_grad=tr_pos)  # fixed sin-cos embedding
+
+        self.decoder_blocks = nn.ModuleList(
+            [Block(decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer)
+             for i in range(decoder_depth)])
+
+        self.decoder_norm = norm_layer(decoder_embed_dim)
+
+        # project channel is different for two modality, use two projection head
+        self.decoder_pred_a = nn.Linear(decoder_embed_dim, patch_size ** 2 * 1, bias=True)  # decoder to patch
+        self.decoder_pred_v = nn.Linear(decoder_embed_dim, patch_size ** 2 * in_chans, bias=True)  # decoder to patch
+
+        self.norm_pix_loss = norm_pix_loss
+
+        self.initialize_weights()
+
+        print('Audio Positional Embedding Shape:', self.pos_embed_a.shape)
+        print('Visual Positional Embedding Shape:', self.pos_embed_v.shape)
+
+    def initialize_weights(self):
+        # initialize (and freeze) pos_embed by sin-cos embedding, opt the cls token, add by myself
+        pos_embed_a = get_2d_sincos_pos_embed(self.pos_embed_a.shape[-1], 8, int(self.patch_embed_a.num_patches / 8),
+                                              cls_token=False)
+        self.pos_embed_a.data.copy_(torch.from_numpy(pos_embed_a).float().unsqueeze(0))
+
+        pos_embed_v = get_2d_sincos_pos_embed(self.pos_embed_v.shape[-1], int(self.patch_embed_v.num_patches ** .5),
+                                              int(self.patch_embed_v.num_patches ** .5), cls_token=False)
+        self.pos_embed_v.data.copy_(torch.from_numpy(pos_embed_v).float().unsqueeze(0))
+
+        decoder_pos_embed_a = get_2d_sincos_pos_embed(self.decoder_pos_embed_a.shape[-1], 8,
+                                                      int(self.patch_embed_a.num_patches / 8), cls_token=False)
+        self.decoder_pos_embed_a.data.copy_(torch.from_numpy(decoder_pos_embed_a).float().unsqueeze(0))
+
+        decoder_pos_embed_v = get_2d_sincos_pos_embed(self.decoder_pos_embed_v.shape[-1],
+                                                      int(self.patch_embed_v.num_patches ** .5),
+                                                      int(self.patch_embed_v.num_patches ** .5), cls_token=False)
+        self.decoder_pos_embed_v.data.copy_(torch.from_numpy(decoder_pos_embed_v).float().unsqueeze(0))
+
+        # initialize patch_embed like nn.Linear (instead of nn.Conv2d)
+        w = self.patch_embed_a.proj.weight.data
+        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        w = self.patch_embed_v.proj.weight.data
+        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+
+        # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
+        torch.nn.init.normal_(self.modality_a, std=.02)
+        torch.nn.init.normal_(self.modality_v, std=.02)
+        torch.nn.init.normal_(self.decoder_modality_a, std=.02)
+        torch.nn.init.normal_(self.decoder_modality_v, std=.02)
+        torch.nn.init.normal_(self.mask_token, std=.02)
+
+        # initialize nn.Linear and nn.LayerNorm
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            # we use xavier_uniform following official JAX ViT:
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def patchify(self, imgs, c, h, w, p=16):
+        """
+        imgs: (N, 3, H, W)
+        x: (N, L, patch_size**2 *3)
+        """
+        x = imgs.reshape(shape=(imgs.shape[0], c, h, p, w, p))
+        x = torch.einsum('nchpwq->nhwpqc', x)
+        x = x.reshape(shape=(imgs.shape[0], h * w, p ** 2 * c))
+        return x
+
+    def unpatchify(self, x, c, h, w, p=16):
+        """
+        x: (N, L, patch_size**2 *3)
+        imgs: (N, 3, H, W)
+        """
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p))
+        return imgs
+
+    def random_masking_unstructured(self, x, mask_ratio):
+        """
+        Perform per-sample random masking by per-sample shuffling.
+        Per-sample shuffling is done by argsort random noise.
+        x: [N, L, D], sequence
+        """
+        N, L, D = x.shape  # batch, length, dim
+        len_keep = int(L * (1 - mask_ratio))
+
+        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]
+
+        # sort noise for each sample
+        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        # keep the first subset
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+        # generate the binary mask: 0 is keep, 1 is remove
+        mask = torch.ones([N, L], device=x.device)
+        mask[:, :len_keep] = 0
+        # unshuffle to get the binary mask
+        mask = torch.gather(mask, dim=1, index=ids_restore)
+
+        return x_masked, mask, ids_restore
+
+    def random_masking_structured(self, x, mask_ratio, t=64, f=8, mode='time'):
+        """
+        Perform per-sample random masking by per-sample shuffling.
+        Per-sample shuffling is done by argsort random noise.
+        x: [N, L, D], sequence
+        """
+        N, L, D = x.shape  # batch, length, dim
+        len_keep = int(L * (1 - mask_ratio))
+
+        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]
+        assert L == f * t
+        noise = noise.reshape(N, f, t)  # the audio patch is in shape [f,t], not [t,f]
+        if mode == 'time':
+            for i in range(N):
+                mask_t_list = random.sample(range(t), int(t * mask_ratio))
+                for k in mask_t_list:
+                    noise[i, :, k] = 1.1  # large value will be removed
+        elif mode == 'freq':
+            for i in range(N):
+                mask_f_list = random.sample(range(f), int(f * mask_ratio))
+                for k in mask_f_list:
+                    noise[i, k, :] = 1.1  # large value will be removed
+        elif mode == 'tf':
+            for i in range(N):
+                mask_t_list = random.sample(range(t), int(t * mask_ratio * 0.7))
+                for k in mask_t_list:
+                    noise[i, :, k] = 1.1  # large value will be removed
+            for i in range(N):
+                mask_f_list = random.sample(range(f), int(f * mask_ratio * 0.7))
+                for k in mask_f_list:
+                    noise[i, k, :] = 1.1  # large value will be removed
+        noise = noise.reshape(N, L)
+
+        # sort noise for each sample, only need to manuplate these two ids_shuffle, ids_restore
+        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        # keep the first subset
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+        # generate the binary mask: 0 is keep, 1 is remove
+        mask = torch.ones([N, L], device=x.device)
+        mask[:, :len_keep] = 0
+        # unshuffle to get the binary mask
+        mask = torch.gather(mask, dim=1, index=ids_restore)
+
+        return x_masked, mask, ids_restore
+
+    def forward_encoder(self, a, v, mask_ratio_a, mask_ratio_v, mask_mode='unstructured'):
+        # embed patches
+        a = a.unsqueeze(1)
+        a = a.transpose(2, 3)
+        a = self.patch_embed_a(a)
+        a = a + self.pos_embed_a
+        a = a + self.modality_a
+
+        v = self.patch_embed_v(v)
+        v = v + self.pos_embed_v
+        v = v + self.modality_v
+
+        # by default, we always use unstructured masking
+        if mask_mode == 'unstructured':
+            a, mask_a, ids_restore_a = self.random_masking_unstructured(a, mask_ratio_a)
+        # in ablation study, we tried time/freq/tf masking. mode in ['freq', 'time', 'tf']
+        else:
+            a, mask_a, ids_restore_a = self.random_masking_structured(a, mask_ratio_a, t=64, f=8, mode=mask_mode)
+
+        # visual branch always use unstructured masking
+        v, mask_v, ids_restore_v = self.random_masking_unstructured(v, mask_ratio_v)
+
+        # audio and visual stream, independent blocks
+        for blk in self.blocks_a:
+            a = blk(a)
+
+        for blk in self.blocks_v:
+            v = blk(v)
+
+        x = torch.cat((a, v), dim=1)
+
+        # unified stream, shared blocks_u, but independent normalization layers
+        for blk in self.blocks_u:
+            x = blk(x)
+        x = self.norm(x)
+
+        for blk in self.blocks_u:
+            ca = blk(a, 'a')
+        ca = self.norm_a(ca)
+
+        for blk in self.blocks_u:
+            cv = blk(v, 'v')
+        cv = self.norm_v(cv)
+
+        return x, mask_a, ids_restore_a, mask_v, ids_restore_v, ca, cv
+
+    def forward_decoder(self, x, mask_a, ids_restore_a, mask_v, ids_restore_v):
+
+        x = self.decoder_embed(x)
+
+        # append mask tokens to sequence
+        # mask_tokens_a in shape [B, #a_mask_token, mask_token_dim], get the number of masked samples from mask_a[0], which is the first example of the batch, all samples should have same number of masked tokens
+        mask_tokens_a = self.mask_token.repeat(x.shape[0], int(mask_a[0].sum()), 1)
+        a_ = torch.cat([x[:, :self.patch_embed_a.num_patches - int(mask_a[0].sum()), :], mask_tokens_a],
+                       dim=1)  # no cls token
+        a_ = torch.gather(a_, dim=1, index=ids_restore_a.unsqueeze(-1).repeat(1, 1, x.shape[2]))  # unshuffle
+
+        # similar for the visual modality
+        mask_tokens_v = self.mask_token.repeat(x.shape[0], int(mask_v[0].sum()), 1)
+        v_ = torch.cat([x[:, self.patch_embed_a.num_patches - int(mask_a[0].sum()):, :], mask_tokens_v],
+                       dim=1)  # no cls token
+        v_ = torch.gather(v_, dim=1, index=ids_restore_v.unsqueeze(-1).repeat(1, 1, x.shape[2]))  # unshuffle
+
+        # concatenate audio and visual tokens
+        x = torch.cat([a_, v_], dim=1)
+
+        decoder_pos_embed = torch.cat([self.decoder_pos_embed_a, self.decoder_pos_embed_v], dim=1)
+        x = x + decoder_pos_embed
+
+        # add modality indication tokens
+        x[:, 0:self.patch_embed_a.num_patches, :] = x[:, 0:self.patch_embed_a.num_patches, :] + self.decoder_modality_a
+        x[:, self.patch_embed_a.num_patches:, :] = x[:, self.patch_embed_a.num_patches:, :] + self.decoder_modality_v
+
+        # apply Transformer blocks
+        for blk in self.decoder_blocks:
+            x = blk(x)
+        x = self.decoder_norm(x)
+
+        # predictor projection
+        x_a = self.decoder_pred_a(x[:, :self.patch_embed_a.num_patches, :])
+        x_v = self.decoder_pred_v(x[:, self.patch_embed_a.num_patches:, :])
+
+        # return audio and video tokens
+        return x_a, x_v
+
+    def forward_contrastive(self, audio_rep, video_rep, bidirect_contrast=False):
+        # calculate nce loss for mean-visual representation and mean-audio representation
+
+        audio_rep = torch.nn.functional.normalize(audio_rep, dim=-1)
+        video_rep = torch.nn.functional.normalize(video_rep, dim=-1)
+
+        total = torch.mm(audio_rep, torch.transpose(video_rep, 0, 1)) / 0.05
+
+        # by default we use single directional
+        if bidirect_contrast == False:
+            nce = -torch.mean(torch.diag(torch.nn.functional.log_softmax(total, dim=0)))
+            c_acc = torch.sum(torch.eq(torch.argmax(torch.nn.functional.softmax(total, dim=0), dim=0),
+                                       torch.arange(0, total.shape[0], device=audio_rep.device))) / total.shape[0]
+            return nce, c_acc
+        else:
+            nce_1 = -torch.mean(torch.diag(torch.nn.functional.log_softmax(total, dim=0)))
+            nce_2 = -torch.mean(torch.diag(torch.nn.functional.log_softmax(total.t(), dim=0)))
+            c_acc_1 = torch.sum(torch.eq(torch.argmax(torch.nn.functional.softmax(total, dim=0), dim=0),
+                                         torch.arange(0, total.shape[0], device=audio_rep.device))) / total.shape[0]
+            c_acc_2 = torch.sum(torch.eq(torch.argmax(torch.nn.functional.softmax(total.t(), dim=0), dim=0),
+                                         torch.arange(0, total.shape[0], device=audio_rep.device))) / total.shape[0]
+            nce = (nce_1 + nce_2) / 2
+            c_acc = (c_acc_1 + c_acc_2) / 2
+            return nce, c_acc
+
+    def forward_mae_loss(self, input, pred, mask, modality):
+        if modality == 'a':
+            # for audio, need to adjust the shape
+            input = input.unsqueeze(1)
+            input = input.transpose(2, 3)
+            target = self.patchify(input, 1, int(input.shape[2] / self.patch_embed_a.patch_size[0]),
+                                   int(input.shape[3] / self.patch_embed_a.patch_size[1]), 16)
+        elif modality == 'v':
+            target = self.patchify(input, 3, int(input.shape[2] / self.patch_embed_v.patch_size[0]),
+                                   int(input.shape[3] / self.patch_embed_v.patch_size[1]), 16)
+
+        # patch-wise normalization might minorly improve the classification performance, but will make the model lose inpainting function
+        if self.norm_pix_loss:
+            mean = target.mean(dim=-1, keepdim=True)
+            var = target.var(dim=-1, keepdim=True)
+            target = (target - mean) / (var + 1.e-6) ** .5
+
+        loss = (pred - target) ** 2
+        loss = loss.mean(dim=-1)  # [N, L], mean loss per patch
+
+        loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches
+        return loss
+
+    def forward(self, audio, imgs, mask_ratio_a=0.75, mask_ratio_v=0.75, mae_loss_weight=1., contrast_loss_weight=0.01,
+                mask_mode='unstructured'):
+        # latent is used for reconstruction (mae), latent_c_{a,v} are used for contrastive learning
+        latent, mask_a, ids_restore_a, mask_v, ids_restore_v, latent_c_a, latent_c_v = self.forward_encoder(audio, imgs,
+                                                                                                            mask_ratio_a,
+                                                                                                            mask_ratio_v,
+                                                                                                            mask_mode=mask_mode)
+        # if mae loss is used
+        if mae_loss_weight != 0:
+            pred_a, pred_v = self.forward_decoder(latent, mask_a, ids_restore_a, mask_v, ids_restore_v)
+            loss_mae_a = self.forward_mae_loss(audio, pred_a, mask_a, 'a')
+            loss_mae_v = self.forward_mae_loss(imgs, pred_v, mask_v, 'v')
+            loss_mae = mae_loss_weight * (loss_mae_a + loss_mae_v)
+        else:
+            loss_mae_a, loss_mae_v, loss_mae = torch.tensor(0.0, device=audio.device), torch.tensor(0.0,
+                                                                                                    device=audio.device), torch.tensor(
+                0.0, device=audio.device)
+
+        # if contrastive loss is used
+        if contrast_loss_weight != 0:
+            # note this is single directional
+            loss_c, c_acc = self.forward_contrastive(latent_c_a.mean(dim=1), latent_c_v.mean(dim=1))
+            loss_c = contrast_loss_weight * loss_c
+        else:
+            loss_c, c_acc = torch.tensor(0.0, device=audio.device), torch.tensor(0.0, device=audio.device)
+
+        loss = loss_mae + loss_c
+
+        return loss, loss_mae, loss_mae_a, loss_mae_v, loss_c, mask_a, mask_v, c_acc
+
+    # used only for inpainting, ignore if inpainting is not of interest
+    def forward_inpaint(self, audio, imgs, mask_ratio_a=0.75, mask_ratio_v=0.75, mask_mode='unstructured'):
+        latent, mask_a, ids_restore_a, mask_v, ids_restore_v, latent_c_a, latent_c_v = self.forward_encoder(audio, imgs,
+                                                                                                            mask_ratio_a,
+                                                                                                            mask_ratio_v,
+                                                                                                            mask_mode=mask_mode)
+        pred_a, pred_v = self.forward_decoder(latent, mask_a, ids_restore_a, mask_v, ids_restore_v)  # [N, L, p*p*3]
+        loss_pixel_a = self.forward_mae_loss(audio, pred_a, mask_a, 'a')
+        loss_pixel_v = self.forward_mae_loss(imgs, pred_v, mask_v, 'v')
+        return pred_a, pred_v, mask_a, mask_v, loss_pixel_a, loss_pixel_v
+
+    # used for retrieval, ignore if retrieval is not of interest
+    def forward_feat(self, a, v):
+        # embed patches
+        a = a.unsqueeze(1)
+        a = a.transpose(2, 3)
+        a = self.patch_embed_a(a)
+        a = a + self.pos_embed_a
+        a = a + self.modality_a
+
+        v = self.patch_embed_v(v)
+        v = v + self.pos_embed_v
+        v = v + self.modality_v
+
+        # the modality-specific stream
+        for blk in self.blocks_a:
+            a = blk(a)
+
+        for blk in self.blocks_v:
+            v = blk(v)
+
+        # use modality specific normalization,
+        for blk in self.blocks_u:
+            a = blk(a, 'a')
+        a = self.norm_a(a)
+
+        for blk in self.blocks_u:
+            v = blk(v, 'v')
+        v = self.norm_v(v)
+        return a, v
+
+    def forward_audio(self, a):
+        # embed patches
+        a = a.unsqueeze(1)
+        a = a.transpose(2, 3)
+        a = self.patch_embed_a(a)
+        a = a + self.pos_embed_a
+        a = a + self.modality_a
+
+        # the modality-specific stream
+        for blk in self.blocks_a:
+            a = blk(a)
+
+        # use modality specific normalization,
+        for blk in self.blocks_u:
+            a = blk(a, 'a')
+        a = self.norm_a(a)
+
+        return a.reshape(a.shape[0], 128 // 16, 1024 // 16, 768).permute(0, 3, 1, 2)
+
+    def forward_video(self, v):
+        v = self.patch_embed_v(v)
+        v = v + self.pos_embed_v
+        v = v + self.modality_v
+
+        for blk in self.blocks_v:
+            v = blk(v)
+
+        for blk in self.blocks_u:
+            v = blk(v, 'v')
+        v = self.norm_v(v)
+        return v.reshape(v.shape[0], 224 // 16, 224 // 16, 768).permute(0, 3, 1, 2)
+
+
+# the finetuned CAV-MAE model
+class CAVMAEFT(nn.Module):
+    def __init__(self, label_dim, img_size=224, audio_length=1024, patch_size=16, in_chans=3,
+                 embed_dim=768, modality_specific_depth=11, num_heads=12, mlp_ratio=4., norm_layer=nn.LayerNorm,
+                 norm_pix_loss=False, tr_pos=True):
+        super().__init__()
+        timm.models.vision_transformer.Block = Block
+        print('Use norm_pix_loss: ', norm_pix_loss)
+
+        timm.models.vision_transformer.PatchEmbed = PatchEmbed
+        timm.models.vision_transformer.Block = Block
+
+        self.patch_embed_a = PatchEmbed(img_size, patch_size, 1, embed_dim)
+        self.patch_embed_v = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
+
+        self.patch_embed_a.num_patches = int(audio_length * 128 / 256)
+        print('Number of Audio Patches: {:d}, Visual Patches: {:d}'.format(self.patch_embed_a.num_patches,
+                                                                           self.patch_embed_v.num_patches))
+
+        self.modality_a = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.modality_v = nn.Parameter(torch.zeros(1, 1, embed_dim))
+
+        self.pos_embed_a = nn.Parameter(torch.zeros(1, self.patch_embed_a.num_patches, embed_dim),
+                                        requires_grad=tr_pos)  # fixed sin-cos embedding
+        self.pos_embed_v = nn.Parameter(torch.zeros(1, self.patch_embed_v.num_patches, embed_dim),
+                                        requires_grad=tr_pos)  # fixed sin-cos embedding
+
+        self.blocks_a = nn.ModuleList(
+            [Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer) for i in
+             range(modality_specific_depth)])
+        self.blocks_v = nn.ModuleList(
+            [Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer) for i in
+             range(modality_specific_depth)])
+        self.blocks_u = nn.ModuleList(
+            [Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer) for i in
+             range(12 - modality_specific_depth)])
+
+        self.norm_a = norm_layer(embed_dim)
+        self.norm_v = norm_layer(embed_dim)
+        self.norm = norm_layer(embed_dim)
+
+        self.mlp_head = nn.Sequential(nn.LayerNorm(embed_dim), nn.Linear(embed_dim, label_dim))
+
+        self.initialize_weights()
+
+        print('Audio Positional Embedding Shape:', self.pos_embed_a.shape)
+        print('Visual Positional Embedding Shape:', self.pos_embed_v.shape)
+
+    def get_patch_num(self, input_shape, stride):
+        test_input = torch.zeros(1, 1, input_shape[0], input_shape[1])
+        test_proj = torch.nn.Conv2d(1, 4, kernel_size=(16, 16), stride=(stride, stride))
+        test_output = test_proj(test_input)
+        print(test_output.shape)
+        return test_output.shape[2], test_output[3], test_output[2] * test_output[2]
+
+    def initialize_weights(self):
+        pos_embed_a = get_2d_sincos_pos_embed(self.pos_embed_a.shape[-1], 8, int(self.patch_embed_a.num_patches / 8),
+                                              cls_token=False)
+        self.pos_embed_a.data.copy_(torch.from_numpy(pos_embed_a).float().unsqueeze(0))
+
+        pos_embed_v = get_2d_sincos_pos_embed(self.pos_embed_v.shape[-1], int(self.patch_embed_v.num_patches ** .5),
+                                              int(self.patch_embed_v.num_patches ** .5), cls_token=False)
+        self.pos_embed_v.data.copy_(torch.from_numpy(pos_embed_v).float().unsqueeze(0))
+
+        w = self.patch_embed_a.proj.weight.data
+        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        w = self.patch_embed_v.proj.weight.data
+        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+
+        torch.nn.init.normal_(self.modality_a, std=.02)
+        torch.nn.init.normal_(self.modality_v, std=.02)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            # we use xavier_uniform following official JAX ViT:
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, a, v, mode):
+        # multi-modal fine-tuning, our default method for fine-tuning
+        if mode == 'multimodal':
+            a = a.unsqueeze(1)
+            a = a.transpose(2, 3)
+            a = self.patch_embed_a(a)
+            a = a + self.pos_embed_a
+            a = a + self.modality_a
+
+            v = self.patch_embed_v(v)
+            v = v + self.pos_embed_v
+            v = v + self.modality_v
+
+            for blk in self.blocks_a:
+                a = blk(a)
+
+            for blk in self.blocks_v:
+                v = blk(v)
+
+            x = torch.cat((a, v), dim=1)
+
+            for blk in self.blocks_u:
+                x = blk(x)
+            x = self.norm(x)
+
+            x = x.mean(dim=1)
+            x = self.mlp_head(x)
+            return x
+
+        # finetune with only audio (and inference with only audio when the model is finetuned with only audio)
+        elif mode == 'audioonly':
+            a = a.unsqueeze(1)
+            a = a.transpose(2, 3)
+            a = self.patch_embed_a(a)
+            a = a + self.pos_embed_a
+            a = a + self.modality_a
+
+            for blk in self.blocks_a:
+                a = blk(a)
+
+            # note here uses the 'a' normalization, it is used in both training and inference, so it is fine
+            for blk in self.blocks_u:
+                a = blk(a, 'a')
+            a = self.norm_a(a)
+            x = a.mean(dim=1)
+            x = self.mlp_head(x)
+            return x
+
+        # finetune with only image (and inference with only audio when the model is finetuned with only image)
+        elif mode == 'videoonly':
+            v = self.patch_embed_v(v)
+            v = v + self.pos_embed_v
+            v = v + self.modality_v
+
+            for blk in self.blocks_v:
+                v = blk(v)
+
+            # note here uses the 'v' normalization, it is used in both training and inference, so it is fine
+            for blk in self.blocks_u:
+                v = blk(v, 'v')
+            v = self.norm_v(v)
+            x = v.mean(dim=1)
+            x = self.mlp_head(x)
+            return x
+
+        # used in case that the model is finetuned with both modality, but in inference only audio is given
+        elif mode == 'missingaudioonly':
+            a = a.unsqueeze(1)
+            a = a.transpose(2, 3)
+            a = self.patch_embed_a(a)
+            a = a + self.pos_embed_a
+            a = a + self.modality_a
+
+            for blk in self.blocks_a:
+                a = blk(a)
+
+            # two forward passes to the block_u, one with modality-specific normalization, another with unified normalization
+            u = a
+            for blk in self.blocks_u:
+                u = blk(u)  # note here use unified normalization
+            u = self.norm(u)
+            u = u.mean(dim=1)
+
+            for blk in self.blocks_u:
+                a = blk(a, 'a')  # note here use modality-specific normalization
+            a = self.norm_a(a)
+            a = a.mean(dim=1)
+
+            # average the output of the two forward passes
+            x = (u + a) / 2
+            x = self.mlp_head(x)
+            return x
+
+        # used in case that the model is fine-tuned with both modality, but in inference only image is given
+        elif mode == 'missingvideoonly':
+            v = self.patch_embed_v(v)
+            v = v + self.pos_embed_v
+            v = v + self.modality_v
+
+            for blk in self.blocks_v:
+                v = blk(v)
+
+            # two forward passes to the block_u, one with modality-specific normalization, another with unified normalization
+            u = v
+            for blk in self.blocks_u:
+                u = blk(u)  # note here use unified normalization
+            u = self.norm(u)
+            u = u.mean(dim=1)
+
+            for blk in self.blocks_u:
+                v = blk(v, 'v')  # note here use modality-specific normalization
+            v = self.norm_v(v)
+            v = v.mean(dim=1)
+
+            # average the output of the two forward passes
+            x = (u + v) / 2
+            x = self.mlp_head(x)
+            return x
+
+    # for retrieval
+    def forward_feat(self, a, v, mode='av'):
+        # return both audio and visual
+        if mode == 'av':
+            a = a.unsqueeze(1)
+            a = a.transpose(2, 3)
+            a = self.patch_embed_a(a)
+            a = a + self.pos_embed_a
+            a = a + self.modality_a
+
+            v = self.patch_embed_v(v)
+            v = v + self.pos_embed_v
+            v = v + self.modality_v
+
+            for blk in self.blocks_a:
+                a = blk(a)
+
+            for blk in self.blocks_v:
+                v = blk(v)
+
+            for blk in self.blocks_u:
+                a = blk(a, 'a')
+            a = self.norm_a(a)
+
+            for blk in self.blocks_u:
+                v = blk(v, 'v')
+
+            v = self.norm_v(v)
+            return a, v
+
+        # return only audio
+        if mode == 'a':
+            a = a.unsqueeze(1)
+            a = a.transpose(2, 3)
+            a = self.patch_embed_a(a)
+            a = a + self.pos_embed_a
+            a = a + self.modality_a
+
+            for blk in self.blocks_a:
+                a = blk(a)
+
+            for blk in self.blocks_u:
+                a = blk(a, 'a')
+
+            a = self.norm_a(a)
+            return a
+
+
+def _wav2fbank(filename):
+    waveform, sr = torchaudio.load(filename)
+    waveform = torchaudio.functional.resample(
+        waveform, orig_freq=sr, new_freq=16000
+    )
+
+    waveform = waveform - waveform.mean()
+    waveform
+    print(sr)
+
+    fbank = torchaudio.compliance.kaldi.fbank(
+        waveform,
+        htk_compat=True,
+        sample_frequency=sr,
+        use_energy=False,
+        window_type='hanning',
+        num_mel_bins=128,
+        dither=0.0,
+        frame_shift=10)
+
+    target_length = 1024
+    n_frames = fbank.shape[0]
+
+    p = target_length - n_frames
+
+    # cut and pad
+    if p > 0:
+        m = torch.nn.ZeroPad2d((0, 0, 0, p))
+        fbank = m(fbank)
+    elif p < 0:
+        fbank = fbank[0:target_length, :]
+
+    return fbank
+
+
+def pca(image_feats_list, dim=3, fit_pca=None):
+    from sklearn.decomposition import PCA
+
+    device = image_feats_list[0].device
+
+    def flatten(tensor, target_size=None):
+        if target_size is not None and fit_pca is None:
+            F.interpolate(tensor, (target_size, target_size), mode="bilinear")
+        B, C, H, W = tensor.shape
+        return feats.permute(1, 0, 2, 3).reshape(C, B * H * W).permute(1, 0).detach().cpu()
+
+    if len(image_feats_list) > 1 and fit_pca is None:
+        target_size = image_feats_list[0].shape[2]
+    else:
+        target_size = None
+
+    flattened_feats = []
+    for feats in image_feats_list:
+        flattened_feats.append(flatten(feats, target_size))
+    x = torch.cat(flattened_feats, dim=0)
+
+    if fit_pca is None:
+        fit_pca = PCA(n_components=dim).fit(x)
+
+    reduced_feats = []
+    for feats in image_feats_list:
+        x_red = torch.from_numpy(fit_pca.transform(flatten(feats)))
+        x_red -= x_red.min(dim=0, keepdim=True).values
+        x_red /= x_red.max(dim=0, keepdim=True).values
+        B, C, H, W = feats.shape
+        reduced_feats.append(x_red.reshape(B, H, W, dim).permute(0, 3, 1, 2).to(device))
+
+    return reduced_feats, fit_pca
+
+
+class CAVMAEAudioFeaturizer(nn.Module):
+
+    def __init__(self, output_path, model_name="base", model=None):
+        super().__init__()
+        if model is not None:
+            self.model = model
+        else:
+            if model_name == "base":
+                model_path = os.path.join(output_path, 'models/audio_model.21.pth')
+            else:
+                raise ValueError(f"Unknown model type {model_name}")
+
+            audio_model = CAVMAE(
+                audio_length=1024,
+                modality_specific_depth=11,
+                norm_pix_loss=True,
+                tr_pos=False)
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            mdl_weight = torch.load(model_path, map_location=device)
+            audio_model = torch.nn.DataParallel(audio_model)
+            audio_model.load_state_dict(mdl_weight, strict=True)
+            self.model = audio_model.module.cuda()
+
+    def forward(self, audio, include_cls):
+        cls_token = None
+        patch_tokens = self.model.forward_audio(audio.squeeze(1))
+
+        if include_cls:
+            return patch_tokens, cls_token
+        else:
+            return patch_tokens
+
+
+class CAVMAEImageFeaturizer(nn.Module):
+
+    def __init__(self, output_path, model=None, model_name="base"):
+        super().__init__()
+        if model is not None:
+            self.model: CAVMAE = model
+        else:
+            if model_name == "base":
+                model_path = os.path.join(output_path, 'models/audio_model.21.pth')
+            else:
+                raise ValueError(f"Unknown model type {model_name}")
+
+            audio_model = CAVMAE(
+                audio_length=1024,
+                modality_specific_depth=11,
+                norm_pix_loss=True,
+                tr_pos=False)
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            mdl_weight = torch.load(model_path, map_location=device)
+            audio_model = torch.nn.DataParallel(audio_model)
+            audio_model.load_state_dict(mdl_weight, strict=True)
+            self.model: CAVMAE = audio_model.module.cuda()
+
+    def forward(self, image, include_cls):
+        cls_token = None
+        patch_tokens = self.model.forward_video(image)
+
+        if include_cls:
+            return patch_tokens, cls_token
+        else:
+            return patch_tokens
+
+
+if __name__ == "__main__":
+    model_path = os.path.join("../../", 'models/audio_model.21.pth')
+    audio_model = CAVMAE(
+        audio_length=1024,
+        modality_specific_depth=11,
+        norm_pix_loss=True,
+        tr_pos=False)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    mdl_weight = torch.load(model_path, map_location=device)
+    audio_model = torch.nn.DataParallel(audio_model)
+    audio_model.load_state_dict(mdl_weight, strict=True)
+    model: CAVMAE = audio_model.module.cuda()
+
+    image_paths = ["../../samples/dog_image.jpg", "../../samples/car_image.jpg", "../../samples/bird_image.jpg"]
+    audio_paths = ["../../samples/dog_audio.wav", "../../samples/car_audio.wav", "../../samples/bird_audio.wav"]
+
+    images = []
+    audios = []
+
+    for image_path in image_paths:
+        image = Image.open(image_path).convert("RGB")
+        preprocess = T.Compose([
+            T.Resize(224, interpolation=Image.BICUBIC),
+            T.CenterCrop(224),
+            T.ToTensor(),
+            T.Normalize(
+                mean=[0.4850, 0.4560, 0.4060],
+                std=[0.2290, 0.2240, 0.2250]
+            )])
+        images.append(preprocess(image).unsqueeze(0).cuda())
+
+    for audio_path in audio_paths:
+        a = _wav2fbank(audio_path).cuda().unsqueeze(0)
+        a = (a + 5.081) / (4.4849)
+        audios.append(a)
+
+    audio_feats, image_feats = model.forward_feat(
+        torch.cat(audios, dim=0), torch.cat(images, dim=0))
+
+    audio_feats = F.normalize(audio_feats.mean(1), dim=1)
+    image_feats = F.normalize(image_feats.mean(1), dim=1)
+
+    sims = torch.einsum("bc,dc->bd", image_feats, audio_feats)
+    print(sims)
+
+    print("here")
+
+    # a_feat = F.normalize(a_feat, dim=1)
+    # v_feat = F.normalize(v_feat, dim=1)
+
+    # [red_v_feat, red_a_feat], fit_pca = pca([v_feat, a_feat])
+    #
+    # [red_v_feat], fit_pca = pca([v_feat])
+    # [red_a_feat], fit_pca = pca([a_feat])
+    #
+    # import matplotlib.pyplot as plt
+    #
+    # fig, ax = plt.subplots(1, 2, figsize=(2 * 5, 5))
+    # ax[0].imshow(red_v_feat[0].permute(1, 2, 0).cpu())
+    # ax[1].imshow(red_a_feat[0].permute(1, 2, 0).cpu())
+    # plt.tight_layout()
+    # plt.show()
+    # print("here")

DenseAV/denseav/featurizers/CLIP.py

@@ -0,0 +1,50 @@
+import clip
+import torch
+from torch import nn
+
+
+class CLIPFeaturizer(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.model, self.preprocess = clip.load("ViT-B/16", device="cpu")
+        self.model.eval().cuda()
+        self.config = {}
+
+    def get_cls_token(self, img):
+        return self.model.encode_image(img).to(torch.float32)
+
+    def forward(self, img, include_cls):
+        features = self.model.get_visual_features(img, include_cls)
+        new_features = []
+        for i in range(2):
+            t = features[i]
+            if isinstance(t, torch.Tensor):
+                new_features.append(t.to(torch.float32))
+            else:
+                new_features.append(t)
+
+        return new_features
+
+
+if __name__ == "__main__":
+    import torchvision.transforms as T
+    from PIL import Image
+    from shared import norm, crop_to_divisor
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    image = Image.open("../samples/lex1.jpg")
+    load_size = 224  # * 3
+    transform = T.Compose([
+        T.Resize(load_size, Image.BILINEAR),
+        # T.CenterCrop(load_size),
+        T.ToTensor(),
+        lambda x: crop_to_divisor(x, 16),
+        norm])
+
+    model = CLIPFeaturizer().cuda()
+
+    results = model(transform(image).cuda().unsqueeze(0))
+
+    print(clip.available_models())

DenseAV/denseav/featurizers/DAVENet.py

@@ -0,0 +1,162 @@
+# Author: David Harwath
+import torch
+import torch.nn as nn
+import torch.nn.functional
+import torch.nn.functional
+import torch.nn.functional as F
+import torch.utils.model_zoo as model_zoo
+import torchvision.models as imagemodels
+
+
+class Davenet(nn.Module):
+    def __init__(self, embedding_dim=1024):
+        super(Davenet, self).__init__()
+        self.embedding_dim = embedding_dim
+        self.batchnorm1 = nn.BatchNorm2d(1)
+        self.conv1 = nn.Conv2d(1, 128, kernel_size=(40, 1), stride=(1, 1), padding=(0, 0))
+        self.conv2 = nn.Conv2d(128, 256, kernel_size=(1, 11), stride=(1, 1), padding=(0, 5))
+        self.conv3 = nn.Conv2d(256, 512, kernel_size=(1, 17), stride=(1, 1), padding=(0, 8))
+        self.conv4 = nn.Conv2d(512, 512, kernel_size=(1, 17), stride=(1, 1), padding=(0, 8))
+        self.conv5 = nn.Conv2d(512, embedding_dim, kernel_size=(1, 17), stride=(1, 1), padding=(0, 8))
+        self.pool = nn.MaxPool2d(kernel_size=(1, 3), stride=(1, 2), padding=(0, 1))
+
+    def forward(self, x):
+        if x.dim() == 3:
+            x = x.unsqueeze(1)
+        x = self.batchnorm1(x)
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = self.pool(x)
+        x = F.relu(self.conv3(x))
+        x = self.pool(x)
+        x = F.relu(self.conv4(x))
+        x = self.pool(x)
+        x = F.relu(self.conv5(x))
+        x = self.pool(x)
+        x = x.squeeze(2)
+        return x
+
+
+class Resnet18(imagemodels.ResNet):
+    def __init__(self, embedding_dim=1024, pretrained=False):
+        super(Resnet18, self).__init__(imagemodels.resnet.BasicBlock, [2, 2, 2, 2])
+        if pretrained:
+            self.load_state_dict(model_zoo.load_url(imagemodels.resnet.model_urls['resnet18']))
+        self.avgpool = None
+        self.fc = None
+        self.embedder = nn.Conv2d(512, embedding_dim, kernel_size=1, stride=1, padding=0)
+        self.embedding_dim = embedding_dim
+        self.pretrained = pretrained
+
+    def forward(self, x):
+        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)
+        x = self.embedder(x)
+        return x
+
+
+class Resnet34(imagemodels.ResNet):
+    def __init__(self, embedding_dim=1024, pretrained=False):
+        super(Resnet34, self).__init__(imagemodels.resnet.BasicBlock, [3, 4, 6, 3])
+        if pretrained:
+            self.load_state_dict(model_zoo.load_url(imagemodels.resnet.model_urls['resnet34']))
+        self.avgpool = None
+        self.fc = None
+        self.embedder = nn.Conv2d(512, embedding_dim, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        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)
+        x = self.embedder(x)
+        return x
+
+
+class Resnet50(imagemodels.ResNet):
+    def __init__(self, embedding_dim=1024, pretrained=False):
+        super(Resnet50, self).__init__(imagemodels.resnet.Bottleneck, [3, 4, 6, 3])
+        if pretrained:
+            self.load_state_dict(model_zoo.load_url(imagemodels.resnet.model_urls['resnet50']))
+        self.avgpool = None
+        self.fc = None
+        self.embedder = nn.Conv2d(2048, embedding_dim, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        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)
+        x = self.embedder(x)
+        return x
+
+
+class VGG16(nn.Module):
+    def __init__(self, embedding_dim=1024, pretrained=False):
+        super(VGG16, self).__init__()
+        seed_model = imagemodels.__dict__['vgg16'](pretrained=pretrained).features
+        seed_model = nn.Sequential(*list(seed_model.children())[:-1])  # remove final maxpool
+        last_layer_index = len(list(seed_model.children()))
+        seed_model.add_module(str(last_layer_index),
+                              nn.Conv2d(512, embedding_dim, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)))
+        self.image_model = seed_model
+
+    def forward(self, x):
+        x = self.image_model(x)
+        return x
+
+
+def prep(dict):
+    return {k.replace("module.", ""): v for k, v in dict.items()}
+
+
+class DavenetAudioFeaturizer(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.audio_model = Davenet()
+        self.audio_model.load_state_dict(prep(torch.load("../models/davenet_pt_audio.pth")))
+
+    def forward(self, audio, include_cls):
+        patch_tokens = self.audio_model(audio).unsqueeze(2)
+
+        if include_cls:
+            return patch_tokens, None
+        else:
+            return patch_tokens
+
+    def get_last_params(self):
+        return []
+
+
+class DavenetImageFeaturizer(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.image_model = VGG16()
+        self.image_model.load_state_dict(prep(torch.load("../models/davenet_pt_image.pth")))
+
+    def forward(self, image, include_cls):
+        patch_tokens = self.image_model(image)
+
+        if include_cls:
+            return patch_tokens, None
+        else:
+            return patch_tokens
+
+    def get_last_params(self):
+        return []

DenseAV/denseav/featurizers/DINO.py

@@ -0,0 +1,451 @@
+import math
+import warnings
+from functools import partial
+
+import timm
+import torch
+import torch.nn as nn
+
+eps = 1e-4
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, return_qkv=False):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x, attn, qkv
+
+
+class Block(nn.Module):
+    def __init__(self, dim,
+                 num_heads,
+                 mlp_ratio=4.,
+                 qkv_bias=False,
+                 qk_scale=None,
+                 drop=0.,
+                 attn_drop=0.,
+                 drop_path=0.,
+                 act_layer=nn.GELU,
+                 norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x, return_attention=False, return_qkv=False):
+        y, attn, qkv = self.attn(self.norm1(x))
+        if return_attention:
+            return attn
+        x = x + self.drop_path(y)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        if return_qkv:
+            return x, attn, qkv
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        num_patches = (img_size // patch_size) * (img_size // patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer """
+
+    def __init__(self,
+                 img_size=[224],
+                 patch_size=16,
+                 in_chans=3,
+                 num_classes=0,
+                 embed_dim=768,
+                 depth=12,
+                 num_heads=12,
+                 mlp_ratio=4.,
+                 qkv_bias=False,
+                 qk_scale=None,
+                 drop_rate=0.,
+                 attn_drop_rate=0.,
+                 drop_path_rate=0.,
+                 norm_layer=nn.LayerNorm,
+                 **kwargs):
+        super().__init__()
+
+        self.num_features = self.embed_dim = embed_dim
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size[0], patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+
+        # Classifier head
+        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        class_pos_embed = self.pos_embed[:, 0]
+        patch_pos_embed = self.pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_embed.patch_size
+        h0 = h // self.patch_embed.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            mode='bicubic',
+        )
+        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).reshape(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def prepare_tokens(self, x):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)  # patch linear embedding
+
+        # add the [CLS] token to the embed patch tokens
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # add positional encoding to each token
+        x = x + self.interpolate_pos_encoding(x, w, h)
+
+        return self.pos_drop(x)
+
+    def forward(self, x):
+        x = self.prepare_tokens(x)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        return x[:, 0]
+
+    def forward_feats(self, x):
+        x = self.prepare_tokens(x)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        return x
+
+    def get_intermediate_feat(self, x, n=1, norm=True):
+        x = self.prepare_tokens(x)
+        # we return the output tokens from the `n` last blocks
+        feat = []
+        attns = []
+        qkvs = []
+        for i, blk in enumerate(self.blocks):
+            x, attn, qkv = blk(x, return_qkv=True)
+            if len(self.blocks) - i <= n:
+                if norm:
+                    feat.append(self.norm(x))
+                else:
+                    feat.append(x)
+                qkvs.append(qkv)
+                attns.append(attn)
+        return feat, attns, qkvs
+
+    def get_last_selfattention(self, x):
+        x = self.prepare_tokens(x)
+        for i, blk in enumerate(self.blocks):
+            if i < len(self.blocks) - 1:
+                x = blk(x)
+            else:
+                # return attention of the last block
+                return blk(x, return_attention=True)
+
+    def get_intermediate_layers(self, x, n=1):
+        x = self.prepare_tokens(x)
+        # we return the output tokens from the `n` last blocks
+        output = []
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if len(self.blocks) - i <= n:
+                output.append(self.norm(x))
+        return output
+
+
+def vit_tiny(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=eps), **kwargs)
+    return model
+
+
+def vit_small(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=eps), **kwargs)
+    return model
+
+
+def vit_base(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=eps), **kwargs)
+    return model
+
+
+class DINOHead(nn.Module):
+    def __init__(self, in_dim, out_dim, use_bn=False, norm_last_layer=True, nlayers=3, hidden_dim=2048,
+                 bottleneck_dim=256):
+        super().__init__()
+        nlayers = max(nlayers, 1)
+        if nlayers == 1:
+            self.mlp = nn.Linear(in_dim, bottleneck_dim)
+        else:
+            layers = [nn.Linear(in_dim, hidden_dim)]
+            if use_bn:
+                layers.append(nn.BatchNorm1d(hidden_dim))
+            layers.append(nn.GELU())
+            for _ in range(nlayers - 2):
+                layers.append(nn.Linear(hidden_dim, hidden_dim))
+                if use_bn:
+                    layers.append(nn.BatchNorm1d(hidden_dim))
+                layers.append(nn.GELU())
+            layers.append(nn.Linear(hidden_dim, bottleneck_dim))
+            self.mlp = nn.Sequential(*layers)
+        self.apply(self._init_weights)
+        self.last_layer = nn.utils.weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))
+        self.last_layer.weight_g.data.fill_(1)
+        if norm_last_layer:
+            self.last_layer.weight_g.requires_grad = False
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x = self.mlp(x)
+        x = nn.functional.normalize(x, dim=-1, p=2)
+        x = self.last_layer(x)
+        return x
+
+
+
+class DINOFeaturizer(nn.Module):
+
+    def __init__(self, arch, patch_size, feat_type):
+        super().__init__()
+        self.arch = arch
+        self.patch_size = patch_size
+        self.feat_type = feat_type
+
+        self.config = {
+            "arch": arch,
+            "patch_size": patch_size,
+            "feat_type": feat_type
+        }
+
+        self.model = vit_small(
+            patch_size=patch_size,
+            num_classes=0)
+
+        if "3d-dino" in arch:
+            state_dict = torch.load("../models/3d-dino-co3d.pth")["teacher"]
+            state_dict = {k.replace("module.", "").replace("backbone.", ""): v for k, v in state_dict.items()}
+            state_dict = {k: v for k, v in state_dict.items() if "head." not in k}
+        elif "iarpa-dino" in arch:
+            state_dict = torch.load("../models/dino_iarpa.pth")["teacher"]
+            state_dict = {k.replace("module.", "").replace("backbone.", ""): v for k, v in state_dict.items()}
+            state_dict = {k: v for k, v in state_dict.items() if "head." not in k}
+        elif "chk-dino" in arch:
+            state_dict = torch.load("../models/dino_deitsmall16_pretrain_full_checkpoint.pth")["teacher"]
+            state_dict = {k.replace("module.", "").replace("backbone.", ""): v for k, v in state_dict.items()}
+            state_dict = {k: v for k, v in state_dict.items() if "head." not in k}
+        elif "ft_dino" in arch:
+            arch = "_".join(arch.split("_")[:-1])
+            state_dict = torch.load("../models/{}.pth".format(arch))["teacher"]
+            state_dict = {k.replace("module.", "").replace("backbone.", ""): v for k, v in state_dict.items()}
+            state_dict = {k: v for k, v in state_dict.items() if "head." not in k}
+        elif "dino" in arch:
+            state_dict = torch.hub.load('facebookresearch/dino:main', self.arch).state_dict()
+        else:  # model from timm -- load weights from timm to dino model (enables working on arbitrary size images).
+            temp_model = timm.create_model(self.arch, pretrained=True)
+            state_dict = temp_model.state_dict()
+            del state_dict['head.weight']
+            del state_dict['head.bias']
+
+        self.model.load_state_dict(state_dict, strict=True)
+
+        if arch == "vit_small":
+            self.n_feats = 384
+        else:
+            self.n_feats = 768
+
+    def get_cls_token(self, img):
+        return self.model.forward(img)
+
+    def forward(self, img, n=1, include_cls=False):
+        assert (img.shape[2] % self.patch_size == 0)
+        assert (img.shape[3] % self.patch_size == 0)
+
+        feat, attn, qkv = self.model.get_intermediate_feat(img, n=n)
+        feat, attn, qkv = feat[0], attn[0], qkv[0]
+
+        feat_h = img.shape[2] // self.patch_size
+        feat_w = img.shape[3] // self.patch_size
+
+        if self.feat_type == "token":
+            image_feat = feat[:, 1:, :].reshape(feat.shape[0], feat_h, feat_w, -1).permute(0, 3, 1, 2)
+            cls_feat = feat[:, 0, :]
+        elif self.feat_type == "key":
+            x = qkv[1, :, :, 1:, :]  # remove cls token
+            desc = x.permute(0, 2, 3, 1).flatten(start_dim=-2, end_dim=-1)
+            image_feat = desc.reshape(desc.shape[0], feat_h, feat_w, desc.shape[2]) \
+                .permute(0, 3, 1, 2)
+            cls_feat = None
+        else:
+            raise ValueError("Unknown feat type:{}".format(self.feat_type))
+
+        if include_cls:
+            return image_feat, cls_feat
+
+        return image_feat

DenseAV/denseav/featurizers/DINOv2.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+
+
+class DINOv2Featurizer(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14').cuda()
+        # self.model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14_reg')
+        self.model.eval()
+        self.config = {}
+
+    def get_cls_token(self, img):
+        pass
+
+    def forward(self, img, include_cls):
+        feature_dict = self.model.forward_features(img)
+        _, _, h, w = img.shape
+        new_h, new_w = h // 14, w // 14
+        b, _, c = feature_dict["x_norm_patchtokens"].shape
+        spatial_tokens = feature_dict["x_norm_patchtokens"].permute(0, 2, 1).reshape(b, c, new_h, new_w)
+
+        if include_cls:
+            return spatial_tokens, feature_dict["x_norm_clstoken"]
+        else:
+            return spatial_tokens
+
+
+if __name__ == "__main__":
+    import torchvision.transforms as T
+    from PIL import Image
+    from shared import norm, crop_to_divisor
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    image = Image.open("../../samples/dog_man_1_crop.jpg")
+    load_size = 224  # * 3
+    transform = T.Compose([
+        T.Resize(load_size, Image.BILINEAR),
+        T.CenterCrop(load_size),
+        T.ToTensor(),
+        norm])
+
+    model = DINOv2Featurizer().cuda()
+
+    results = model(transform(image).cuda().unsqueeze(0), include_cls=False)
+
+    print(results.shape)

DenseAV/denseav/featurizers/Hubert.py

@@ -0,0 +1,70 @@
+import torch
+import torch.nn as nn
+from transformers import Wav2Vec2Processor, HubertModel, HubertConfig
+from transformers.pytorch_utils import Conv1D
+
+class HubertAudioTransform():
+
+    def __init__(self):
+        self.processor = Wav2Vec2Processor.from_pretrained("facebook/hubert-large-ls960-ft")
+
+    def __call__(self, audio):
+        return self.processor(audio, return_tensors="pt", sampling_rate=16000).input_values.squeeze(0)
+
+
+def copy_conv(l):
+    new_l = Conv1D()
+
+
+class Hubert(nn.Module):
+    def __init__(self):
+        super().__init__()
+        model1 = HubertModel.from_pretrained("facebook/hubert-large-ls960-ft")
+        config = model1.config
+        del model1
+        config.layer_norm_eps = 1e-4
+        self.model = HubertModel.from_pretrained("facebook/hubert-large-ls960-ft", config=config)
+        self.config = dict()
+
+
+    def forward(self, audio, include_cls):
+        outputs = self.model(audio)
+        # outputs = deepspeed.checkpointing.checkpoint(self.model, audio)
+
+        patch_tokens = outputs.last_hidden_state.permute(0, 2, 1).unsqueeze(2)
+
+        # return patch_tokens
+        if include_cls:
+            return patch_tokens, None
+        else:
+            return patch_tokens
+
+    def get_last_params(self):
+        return self.model.encoder.layers[-1].parameters()
+
+
+if __name__ == "__main__":
+    import librosa
+    from shared import pca, remove_axes
+    import matplotlib.pyplot as plt
+    from pytorch_lightning import seed_everything
+
+    audio, _ = librosa.load("../../samples/example.wav", sr=16000)
+    audio = torch.from_numpy(audio).unsqueeze(0).to("cuda")
+
+    model = Hubert().to("cuda")
+    embeddings = model.forward(audio, include_cls=False)
+
+    print(embeddings.shape)
+    seed_everything(0)
+
+    with torch.no_grad():
+        [pca_feats], _ = pca([embeddings])
+        pca_feats = torch.broadcast_to(
+            pca_feats, (pca_feats.shape[0], pca_feats.shape[1], 25, pca_feats.shape[3]))
+        fig, axes = plt.subplots(2, 1, figsize=(10, 7))
+        axes[1].imshow(pca_feats.cpu().squeeze(0).permute(1, 2, 0))
+        remove_axes(axes)
+        plt.tight_layout()
+        plt.show()
+        print("here")

DenseAV/denseav/featurizers/ImageBind.py

@@ -0,0 +1,2033 @@
+import gzip
+import html
+import io
+import logging
+import math
+import os
+from functools import lru_cache
+from functools import partial
+from types import SimpleNamespace
+from typing import Callable, List
+from typing import Optional
+
+import einops
+import ftfy
+import numpy as np
+import regex as re
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+import torchaudio
+import torchvision.transforms as T
+from PIL import Image
+from timm.models.layers import DropPath, trunc_normal_
+from torchvision import transforms
+import matplotlib.pyplot as plt
+from iopath.common.file_io import g_pathmgr
+
+
+class Attention(nn.Module):
+    def __init__(
+            self,
+            dim,
+            num_heads=8,
+            qkv_bias=False,
+            qk_scale=None,
+            attn_drop=0.0,
+            proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version,
+        # can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )  # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Mlp(nn.Module):
+    def __init__(
+            self,
+            in_features,
+            hidden_features=None,
+            out_features=None,
+            act_layer=nn.GELU,
+            drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class MultiheadAttention(nn.MultiheadAttention):
+    def forward(self, x: torch.Tensor, attn_mask: torch.Tensor):
+        return super().forward(x, x, x, need_weights=False, attn_mask=attn_mask)[0]
+
+
+class ViTAttention(Attention):
+    def forward(self, x: torch.Tensor, attn_mask: torch.Tensor):
+        assert attn_mask is None
+        return super().forward(x)
+
+
+class BlockWithMasking(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            attn_target: Callable,
+            mlp_ratio: int = 4,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = nn.LayerNorm,
+            ffn_dropout_rate: float = 0.0,
+            drop_path: float = 0.0,
+            layer_scale_type: str = None,
+            layer_scale_init_value: float = 1e-4,
+    ):
+        super().__init__()
+
+        assert not isinstance(
+            attn_target, nn.Module
+        ), "attn_target should be a Callable. Otherwise attn_target is shared across blocks!"
+        self.attn = attn_target()
+        if drop_path > 0.0:
+            self.drop_path = DropPath(drop_path)
+        else:
+            self.drop_path = nn.Identity()
+        self.norm_1 = norm_layer(dim)
+        mlp_hidden_dim = int(mlp_ratio * dim)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=ffn_dropout_rate,
+        )
+        self.norm_2 = norm_layer(dim)
+        self.layer_scale_type = layer_scale_type
+        if self.layer_scale_type is not None:
+            assert self.layer_scale_type in [
+                "per_channel",
+                "scalar",
+            ], f"Found Layer scale type {self.layer_scale_type}"
+            if self.layer_scale_type == "per_channel":
+                # one gamma value per channel
+                gamma_shape = [1, 1, dim]
+            elif self.layer_scale_type == "scalar":
+                # single gamma value for all channels
+                gamma_shape = [1, 1, 1]
+            # two gammas: for each part of the fwd in the encoder
+            self.layer_scale_gamma1 = nn.Parameter(
+                torch.ones(size=gamma_shape) * layer_scale_init_value,
+                requires_grad=True,
+            )
+            self.layer_scale_gamma2 = nn.Parameter(
+                torch.ones(size=gamma_shape) * layer_scale_init_value,
+                requires_grad=True,
+            )
+
+    def forward(self, x: torch.Tensor, attn_mask: torch.Tensor):
+        if self.layer_scale_type is None:
+            x = x + self.drop_path(self.attn(self.norm_1(x), attn_mask))
+            x = x + self.drop_path(self.mlp(self.norm_2(x)))
+        else:
+            x = (
+                    x
+                    + self.drop_path(self.attn(self.norm_1(x), attn_mask))
+                    * self.layer_scale_gamma1
+            )
+            x = x + self.drop_path(self.mlp(self.norm_2(x))) * self.layer_scale_gamma2
+        return x
+
+
+_LAYER_NORM = partial(nn.LayerNorm, eps=1e-6)
+
+
+class SimpleTransformer(nn.Module):
+    def __init__(
+            self,
+            attn_target: Callable,
+            embed_dim: int,
+            num_blocks: int,
+            block: Callable = BlockWithMasking,
+            pre_transformer_layer: Callable = None,
+            post_transformer_layer: Callable = None,
+            drop_path_rate: float = 0.0,
+            drop_path_type: str = "progressive",
+            norm_layer: Callable = _LAYER_NORM,
+            mlp_ratio: int = 4,
+            ffn_dropout_rate: float = 0.0,
+            layer_scale_type: str = None,  # from cait; possible values are None, "per_channel", "scalar"
+            layer_scale_init_value: float = 1e-4,  # from cait; float
+            weight_init_style: str = "jax",  # possible values jax or pytorch
+    ):
+        """
+        Simple Transformer with the following features
+        1. Supports masked attention
+        2. Supports DropPath
+        3. Supports LayerScale
+        4. Supports Dropout in Attention and FFN
+        5. Makes few assumptions about the input except that it is a Tensor
+        """
+        super().__init__()
+        self.pre_transformer_layer = pre_transformer_layer
+        if drop_path_type == "progressive":
+            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, num_blocks)]
+        elif drop_path_type == "uniform":
+            dpr = [drop_path_rate for i in range(num_blocks)]
+        else:
+            raise ValueError(f"Unknown drop_path_type: {drop_path_type}")
+
+        self.blocks = nn.Sequential(
+            *[
+                block(
+                    dim=embed_dim,
+                    attn_target=attn_target,
+                    mlp_ratio=mlp_ratio,
+                    ffn_dropout_rate=ffn_dropout_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                    layer_scale_type=layer_scale_type,
+                    layer_scale_init_value=layer_scale_init_value,
+                )
+                for i in range(num_blocks)
+            ]
+        )
+        self.post_transformer_layer = post_transformer_layer
+        self.weight_init_style = weight_init_style
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            if self.weight_init_style == "jax":
+                # Based on MAE and official Jax ViT implementation
+                torch.nn.init.xavier_uniform_(m.weight)
+            elif self.weight_init_style == "pytorch":
+                # PyTorch ViT uses trunc_normal_
+                trunc_normal_(m.weight, std=0.02)
+
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (nn.LayerNorm)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(
+            self,
+            tokens: torch.Tensor,
+            attn_mask: torch.Tensor = None,
+            use_checkpoint: bool = False,
+            checkpoint_every_n: int = 1,
+            checkpoint_blk_ids: List[int] = None,
+    ):
+        """
+        Inputs
+        - tokens: data of shape N x L x D (or L x N x D depending on the attention implementation)
+        - attn: mask of shape L x L
+
+        Output
+        - x: data of shape N x L x D (or L x N x D depending on the attention implementation)
+        """
+        if self.pre_transformer_layer:
+            tokens = self.pre_transformer_layer(tokens)
+        if use_checkpoint and checkpoint_blk_ids is None:
+            checkpoint_blk_ids = [
+                blk_id
+                for blk_id in range(len(self.blocks))
+                if blk_id % checkpoint_every_n == 0
+            ]
+        if checkpoint_blk_ids:
+            checkpoint_blk_ids = set(checkpoint_blk_ids)
+        for blk_id, blk in enumerate(self.blocks):
+            if use_checkpoint and blk_id in checkpoint_blk_ids:
+                tokens = checkpoint.checkpoint(
+                    blk, tokens, attn_mask, use_reentrant=False
+                )
+            else:
+                tokens = blk(tokens, attn_mask=attn_mask)
+        if self.post_transformer_layer:
+            tokens = self.post_transformer_layer(tokens)
+        return tokens
+
+
+def get_sinusoid_encoding_table(n_position, d_hid):
+    """Sinusoid position encoding table"""
+
+    # TODO: make it with torch instead of numpy
+    def get_position_angle_vec(position):
+        return [
+            position / np.power(10000, 2 * (hid_j // 2) / d_hid)
+            for hid_j in range(d_hid)
+        ]
+
+    sinusoid_table = np.array(
+        [get_position_angle_vec(pos_i) for pos_i in range(n_position)]
+    )
+    sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
+    sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1
+
+    return torch.FloatTensor(sinusoid_table).unsqueeze(0)
+
+
+def interpolate_pos_encoding_2d(target_spatial_size, pos_embed):
+    N = pos_embed.shape[1]
+    if N == target_spatial_size:
+        return pos_embed
+    dim = pos_embed.shape[-1]
+    # nn.functional.interpolate doesn't work with bfloat16 so we cast to float32
+    pos_embed, updated = cast_if_src_dtype(pos_embed, torch.bfloat16, torch.float32)
+    pos_embed = nn.functional.interpolate(
+        pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(
+            0, 3, 1, 2
+        ),
+        scale_factor=math.sqrt(target_spatial_size / N),
+        mode="bicubic",
+    )
+    if updated:
+        pos_embed, _ = cast_if_src_dtype(pos_embed, torch.float32, torch.bfloat16)
+    pos_embed = pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+    return pos_embed
+
+
+def interpolate_pos_encoding(
+        npatch_per_img,
+        pos_embed,
+        patches_layout,
+        input_shape=None,
+        first_patch_idx=1,
+):
+    assert first_patch_idx == 0 or first_patch_idx == 1, "there is 1 CLS token or none"
+    N = pos_embed.shape[1] - first_patch_idx  # since it's 1 if cls_token exists
+    if npatch_per_img == N:
+        return pos_embed
+
+    # assert (
+    #         patches_layout[-1] == patches_layout[-2]
+    # ), "Interpolation of pos embed not supported for non-square layouts"
+
+    class_emb = pos_embed[:, :first_patch_idx]
+    pos_embed = pos_embed[:, first_patch_idx:]
+
+    if input_shape is None or patches_layout[0] == 1:
+        # simple 2D pos embedding, no temporal component
+        pos_embed = interpolate_pos_encoding_2d(npatch_per_img, pos_embed)
+    elif patches_layout[0] > 1:
+        # pos embed has a temporal component
+        assert len(input_shape) == 4, "temporal interpolation not supported"
+        # we only support 2D interpolation in this case
+        num_frames = patches_layout[0]
+        num_spatial_tokens = patches_layout[1] * patches_layout[2]
+        pos_embed = pos_embed.view(1, num_frames, num_spatial_tokens, -1)
+        # interpolate embedding for zeroth frame
+        pos_embed = interpolate_pos_encoding_2d(
+            npatch_per_img, pos_embed[0, 0, ...].unsqueeze(0)
+        )
+    else:
+        raise ValueError("This type of interpolation isn't implemented")
+
+    return torch.cat((class_emb, pos_embed), dim=1)
+
+
+def _get_pos_embedding(
+        npatch_per_img,
+        pos_embed,
+        patches_layout,
+        input_shape,
+        first_patch_idx=1,
+):
+    pos_embed = interpolate_pos_encoding(
+        npatch_per_img,
+        pos_embed,
+        patches_layout,
+        input_shape=input_shape,
+        first_patch_idx=first_patch_idx,
+    )
+    return pos_embed
+
+
+class VerboseNNModule(nn.Module):
+    """
+    Wrapper around nn.Module that prints registered buffers and parameter names.
+    """
+
+    @staticmethod
+    def get_readable_tensor_repr(name: str, tensor: torch.Tensor) -> str:
+        st = (
+                "("
+                + name
+                + "): "
+                + "tensor("
+                + str(tuple(tensor[1].shape))
+                + ", requires_grad="
+                + str(tensor[1].requires_grad)
+                + ")\n"
+        )
+        return st
+
+    def extra_repr(self) -> str:
+        named_modules = set()
+        for p in self.named_modules():
+            named_modules.update([p[0]])
+        named_modules = list(named_modules)
+
+        string_repr = ""
+        for p in self.named_parameters():
+            name = p[0].split(".")[0]
+            if name not in named_modules:
+                string_repr += self.get_readable_tensor_repr(name, p)
+
+        for p in self.named_buffers():
+            name = p[0].split(".")[0]
+            string_repr += self.get_readable_tensor_repr(name, p)
+
+        return string_repr
+
+
+class PatchEmbedGeneric(nn.Module):
+    """
+    PatchEmbed from Hydra
+    """
+
+    def __init__(self, proj_stem, norm_layer: Optional[nn.Module] = None):
+        super().__init__()
+
+        if len(proj_stem) > 1:
+            self.proj = nn.Sequential(*proj_stem)
+        else:
+            # Special case to be able to load pre-trained models that were
+            # trained with a standard stem
+            self.proj = proj_stem[0]
+        self.norm_layer = norm_layer
+
+    def get_patch_layout(self, img_size):
+        with torch.no_grad():
+            dummy_img = torch.zeros(
+                [
+                    1,
+                ]
+                + img_size
+            )
+            dummy_out = self.proj(dummy_img)
+        embed_dim = dummy_out.shape[1]
+        patches_layout = tuple(dummy_out.shape[2:])
+        num_patches = np.prod(patches_layout)
+        return patches_layout, num_patches, embed_dim
+
+    def forward(self, x):
+        x = self.proj(x)
+        # B C (T) H W -> B (T)HW C
+        x = x.flatten(2).transpose(1, 2)
+        if self.norm_layer is not None:
+            x = self.norm_layer(x)
+        return x
+
+
+class SpatioTemporalPosEmbeddingHelper(VerboseNNModule):
+    def __init__(
+            self,
+            patches_layout: List,
+            num_patches: int,
+            num_cls_tokens: int,
+            embed_dim: int,
+            learnable: bool,
+    ) -> None:
+        super().__init__()
+        self.num_cls_tokens = num_cls_tokens
+        self.patches_layout = patches_layout
+        self.num_patches = num_patches
+        self.num_tokens = num_cls_tokens + num_patches
+        self.learnable = learnable
+        if self.learnable:
+            self.pos_embed = nn.Parameter(torch.zeros(1, self.num_tokens, embed_dim))
+            trunc_normal_(self.pos_embed, std=0.02)
+        else:
+            self.register_buffer(
+                "pos_embed", get_sinusoid_encoding_table(self.num_tokens, embed_dim)
+            )
+
+    def get_pos_embedding(self, vision_input, all_vision_tokens):
+        input_shape = vision_input.shape
+        pos_embed = _get_pos_embedding(
+            all_vision_tokens.size(1) - self.num_cls_tokens,
+            pos_embed=self.pos_embed,
+            patches_layout=self.patches_layout,
+            input_shape=input_shape,
+            first_patch_idx=self.num_cls_tokens,
+        )
+        return pos_embed
+
+
+class RGBDTPreprocessor(VerboseNNModule):
+    def __init__(
+            self,
+            rgbt_stem: PatchEmbedGeneric,
+            depth_stem: PatchEmbedGeneric,
+            img_size: List = (3, 224, 224),
+            num_cls_tokens: int = 1,
+            pos_embed_fn: Callable = None,
+            use_type_embed: bool = False,
+            init_param_style: str = "openclip",
+    ) -> None:
+        super().__init__()
+        stem = rgbt_stem if rgbt_stem is not None else depth_stem
+        (
+            self.patches_layout,
+            self.num_patches,
+            self.embed_dim,
+        ) = stem.get_patch_layout(img_size)
+        self.rgbt_stem = rgbt_stem
+        self.depth_stem = depth_stem
+        self.use_pos_embed = pos_embed_fn is not None
+        self.use_type_embed = use_type_embed
+        self.num_cls_tokens = num_cls_tokens
+
+        if self.use_pos_embed:
+            self.pos_embedding_helper = pos_embed_fn(
+                patches_layout=self.patches_layout,
+                num_cls_tokens=num_cls_tokens,
+                num_patches=self.num_patches,
+                embed_dim=self.embed_dim,
+            )
+        if self.num_cls_tokens > 0:
+            self.cls_token = nn.Parameter(
+                torch.zeros(1, self.num_cls_tokens, self.embed_dim)
+            )
+        if self.use_type_embed:
+            self.type_embed = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+
+        self.init_parameters(init_param_style)
+
+    @torch.no_grad()
+    def init_parameters(self, init_param_style):
+        if init_param_style == "openclip":
+            # OpenCLIP style initialization
+            scale = self.embed_dim ** -0.5
+            if self.use_pos_embed:
+                nn.init.normal_(self.pos_embedding_helper.pos_embed)
+                self.pos_embedding_helper.pos_embed *= scale
+
+            if self.num_cls_tokens > 0:
+                nn.init.normal_(self.cls_token)
+                self.cls_token *= scale
+        elif init_param_style == "vit":
+            self.cls_token.data.fill_(0)
+        else:
+            raise ValueError(f"Unknown init {init_param_style}")
+
+        if self.use_type_embed:
+            nn.init.normal_(self.type_embed)
+
+    def get_pos_emb_2(self, input, stem):
+        patches = stem.proj(input)
+        target_size = patches.shape[-2:]
+        original_size = list(self.pos_embedding_helper.patches_layout)[-2:]
+
+        orig_ce = self.pos_embedding_helper.pos_embed[:, 0, :]
+        orig_pe = ((self.pos_embedding_helper.pos_embed[:, 1:, :]
+                    .reshape(1, *original_size, self.embed_dim))
+                   .permute(0, 3, 1, 2))
+
+        new_pe = F.interpolate(orig_pe, size=target_size, mode="bicubic")
+
+        new_full_pe = torch.cat([orig_ce.unsqueeze(1), new_pe.permute(0, 2, 3, 1).reshape(1, -1, self.embed_dim)],
+                                dim=1)
+
+        return new_full_pe
+
+    def tokenize_input_and_cls_pos(self, input, stem, mask):
+        # tokens is of shape B x L x D
+        tokens = stem(input)
+        assert tokens.ndim == 3
+        assert tokens.shape[2] == self.embed_dim
+        B = tokens.shape[0]
+        if self.num_cls_tokens > 0:
+            class_tokens = self.cls_token.expand(
+                B, -1, -1
+            )  # stole class_tokens impl from Phil Wang, thanks
+            tokens = torch.cat((class_tokens, tokens), dim=1)
+        if self.use_pos_embed:
+            pos_embed = self.pos_embedding_helper.get_pos_embedding(input, tokens)
+            # pos_embed = self.get_pos_emb_2(input, stem)
+            tokens = tokens + pos_embed
+        if self.use_type_embed:
+            tokens = tokens + self.type_embed.expand(B, -1, -1)
+        return tokens
+
+    def forward(self, vision=None, depth=None, patch_mask=None):
+        if patch_mask is not None:
+            raise NotImplementedError()
+
+        if vision is not None:
+            vision_tokens = self.tokenize_input_and_cls_pos(
+                vision, self.rgbt_stem, patch_mask
+            )
+
+        if depth is not None:
+            depth_tokens = self.tokenize_input_and_cls_pos(
+                depth, self.depth_stem, patch_mask
+            )
+
+        # aggregate tokens
+        if vision is not None and depth is not None:
+            final_tokens = vision_tokens + depth_tokens
+        else:
+            final_tokens = vision_tokens if vision is not None else depth_tokens
+        return_dict = {
+            "trunk": {
+                "tokens": final_tokens,
+            },
+            "head": {},
+        }
+        return return_dict
+
+
+class AudioPreprocessor(RGBDTPreprocessor):
+    def __init__(self, audio_stem: PatchEmbedGeneric, **kwargs) -> None:
+        super().__init__(rgbt_stem=audio_stem, depth_stem=None, **kwargs)
+
+    def forward(self, audio=None):
+        return super().forward(vision=audio)
+
+
+class ThermalPreprocessor(RGBDTPreprocessor):
+    def __init__(self, thermal_stem: PatchEmbedGeneric, **kwargs) -> None:
+        super().__init__(rgbt_stem=thermal_stem, depth_stem=None, **kwargs)
+
+    def forward(self, thermal=None):
+        return super().forward(vision=thermal)
+
+
+def build_causal_attention_mask(context_length):
+    # lazily create causal attention mask, with full attention between the vision tokens
+    # pytorch uses additive attention mask; fill with -inf
+    mask = torch.empty(context_length, context_length, requires_grad=False)
+    mask.fill_(float("-inf"))
+    mask.triu_(1)  # zero out the lower diagonal
+    return mask
+
+
+class TextPreprocessor(VerboseNNModule):
+    def __init__(
+            self,
+            vocab_size: int,
+            context_length: int,
+            embed_dim: int,
+            causal_masking: bool,
+            supply_seq_len_to_head: bool = True,
+            num_cls_tokens: int = 0,
+            init_param_style: str = "openclip",
+    ) -> None:
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.context_length = context_length
+        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
+        self.pos_embed = nn.Parameter(
+            torch.empty(1, self.context_length + num_cls_tokens, embed_dim)
+        )
+        self.causal_masking = causal_masking
+        if self.causal_masking:
+            mask = build_causal_attention_mask(self.context_length)
+            # register the mask as a buffer, so it can be moved to the right device
+            self.register_buffer("mask", mask)
+
+        self.supply_seq_len_to_head = supply_seq_len_to_head
+        self.num_cls_tokens = num_cls_tokens
+        self.embed_dim = embed_dim
+        if num_cls_tokens > 0:
+            assert self.causal_masking is False, "Masking + CLS token isn't implemented"
+            self.cls_token = nn.Parameter(
+                torch.zeros(1, self.num_cls_tokens, embed_dim)
+            )
+
+        self.init_parameters(init_param_style)
+
+    @torch.no_grad()
+    def init_parameters(self, init_param_style="openclip"):
+        # OpenCLIP style initialization
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.pos_embed, std=0.01)
+
+        if init_param_style == "openclip":
+            # OpenCLIP style initialization
+            scale = self.embed_dim ** -0.5
+            if self.num_cls_tokens > 0:
+                nn.init.normal_(self.cls_token)
+                self.cls_token *= scale
+        elif init_param_style == "vit":
+            self.cls_token.data.fill_(0)
+        else:
+            raise ValueError(f"Unknown init {init_param_style}")
+
+    def forward(self, text):
+        # text tokens are of shape B x L x D
+        text_tokens = self.token_embedding(text)
+        # concat CLS tokens if any
+        if self.num_cls_tokens > 0:
+            B = text_tokens.shape[0]
+            class_tokens = self.cls_token.expand(
+                B, -1, -1
+            )  # stole class_tokens impl from Phil Wang, thanks
+            text_tokens = torch.cat((class_tokens, text_tokens), dim=1)
+        text_tokens = text_tokens + self.pos_embed
+        return_dict = {
+            "trunk": {
+                "tokens": text_tokens,
+            },
+            "head": {},
+        }
+        # Compute sequence length after adding CLS tokens
+        if self.supply_seq_len_to_head:
+            text_lengths = text.argmax(dim=-1)
+            return_dict["head"] = {
+                "seq_len": text_lengths,
+            }
+        if self.causal_masking:
+            return_dict["trunk"].update({"attn_mask": self.mask})
+        return return_dict
+
+
+class Im2Video(nn.Module):
+    """Convert an image into a trivial video."""
+
+    def __init__(self, time_dim=2):
+        super().__init__()
+        self.time_dim = time_dim
+
+    def forward(self, x):
+        if x.ndim == 4:
+            # B, C, H, W -> B, C, T, H, W
+            return x.unsqueeze(self.time_dim)
+        elif x.ndim == 5:
+            return x
+        else:
+            raise ValueError(f"Dimension incorrect {x.shape}")
+
+
+class PadIm2Video(Im2Video):
+    def __init__(self, ntimes, pad_type, time_dim=2):
+        super().__init__(time_dim=time_dim)
+        assert ntimes > 0
+        assert pad_type in ["zero", "repeat"]
+        self.ntimes = ntimes
+        self.pad_type = pad_type
+
+    def forward(self, x):
+        x = super().forward(x)
+        if x.shape[self.time_dim] == 1:
+            if self.pad_type == "repeat":
+                new_shape = [1] * len(x.shape)
+                new_shape[self.time_dim] = self.ntimes
+                x = x.repeat(new_shape)
+            elif self.pad_type == "zero":
+                padarg = [0, 0] * len(x.shape)
+                padarg[2 * self.time_dim + 1] = self.ntimes - x.shape[self.time_dim]
+                x = nn.functional.pad(x, padarg)
+        return x
+
+
+# Modified from github.com/openai/CLIP
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a signficant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = (
+            list(range(ord("!"), ord("~") + 1))
+            + list(range(ord("¡"), ord("¬") + 1))
+            + list(range(ord("®"), ord("ÿ") + 1))
+    )
+    cs = bs[:]
+    n = 0
+    for b in range(2 ** 8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2 ** 8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r"\s+", " ", text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer(object):
+    def __init__(self, bpe_path: str, context_length=77):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+
+        with g_pathmgr.open(bpe_path, "rb") as fh:
+            bpe_bytes = io.BytesIO(fh.read())
+            merges = gzip.open(bpe_bytes).read().decode("utf-8").split("\n")
+        merges = merges[1: 49152 - 256 - 2 + 1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v + "</w>" for v in vocab]
+        for merge in merges:
+            vocab.append("".join(merge))
+        vocab.extend(["<|startoftext|>", "<|endoftext|>"])
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {
+            "<|startoftext|>": "<|startoftext|>",
+            "<|endoftext|>": "<|endoftext|>",
+        }
+        self.pat = re.compile(
+            r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
+            re.IGNORECASE,
+        )
+        self.context_length = context_length
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + (token[-1] + "</w>",)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token + "</w>"
+
+        while True:
+            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
+                    new_word.append(first + second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = " ".join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = "".join(self.byte_encoder[b] for b in token.encode("utf-8"))
+            bpe_tokens.extend(
+                self.encoder[bpe_token] for bpe_token in self.bpe(token).split(" ")
+            )
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = "".join([self.decoder[token] for token in tokens])
+        text = (
+            bytearray([self.byte_decoder[c] for c in text])
+            .decode("utf-8", errors="replace")
+            .replace("</w>", " ")
+        )
+        return text
+
+    def __call__(self, texts, context_length=None):
+        if not context_length:
+            context_length = self.context_length
+
+        if isinstance(texts, str):
+            texts = [texts]
+
+        sot_token = self.encoder["<|startoftext|>"]
+        eot_token = self.encoder["<|endoftext|>"]
+        all_tokens = [[sot_token] + self.encode(text) + [eot_token] for text in texts]
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+        for i, tokens in enumerate(all_tokens):
+            tokens = tokens[:context_length]
+            result[i, : len(tokens)] = torch.tensor(tokens)
+
+        if len(result) == 1:
+            return result[0]
+        return result
+
+
+class Normalize(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        return torch.nn.functional.normalize(x, dim=self.dim, p=2)
+
+
+class LearnableLogitScaling(nn.Module):
+    def __init__(
+            self,
+            logit_scale_init: float = 1 / 0.07,
+            learnable: bool = True,
+            max_logit_scale: float = 100,
+    ) -> None:
+        super().__init__()
+        self.max_logit_scale = max_logit_scale
+        self.logit_scale_init = logit_scale_init
+        self.learnable = learnable
+        log_logit_scale = torch.ones([]) * np.log(self.logit_scale_init)
+        if learnable:
+            self.log_logit_scale = nn.Parameter(log_logit_scale)
+        else:
+            self.register_buffer("log_logit_scale", log_logit_scale)
+
+    def forward(self, x):
+        return torch.clip(self.log_logit_scale.exp(), max=self.max_logit_scale) * x
+
+    def extra_repr(self):
+        st = f"logit_scale_init={self.logit_scale_init},learnable={self.learnable}, max_logit_scale={self.max_logit_scale}"
+        return st
+
+
+class EinOpsRearrange(nn.Module):
+    def __init__(self, rearrange_expr: str, **kwargs) -> None:
+        super().__init__()
+        self.rearrange_expr = rearrange_expr
+        self.kwargs = kwargs
+
+    def forward(self, x):
+        assert isinstance(x, torch.Tensor)
+        return einops.rearrange(x, self.rearrange_expr, **self.kwargs)
+
+
+class IMUPreprocessor(VerboseNNModule):
+    def __init__(
+            self,
+            kernel_size: int,
+            imu_stem: PatchEmbedGeneric,
+            embed_dim: int,
+            img_size: List = (6, 2000),
+            num_cls_tokens: int = 1,
+            pos_embed_fn: Callable = None,
+            init_param_style: str = "openclip",
+    ) -> None:
+        super().__init__()
+        stem = imu_stem
+        self.imu_stem = imu_stem
+        self.embed_dim = embed_dim
+        self.use_pos_embed = pos_embed_fn is not None
+        self.num_cls_tokens = num_cls_tokens
+        self.kernel_size = kernel_size
+        self.pos_embed = nn.Parameter(
+            torch.empty(1, (img_size[1] // kernel_size) + num_cls_tokens, embed_dim)
+        )
+
+        if self.num_cls_tokens > 0:
+            self.cls_token = nn.Parameter(
+                torch.zeros(1, self.num_cls_tokens, self.embed_dim)
+            )
+
+        self.init_parameters(init_param_style)
+
+    @torch.no_grad()
+    def init_parameters(self, init_param_style):
+        nn.init.normal_(self.pos_embed, std=0.01)
+
+        if init_param_style == "openclip":
+            # OpenCLIP style initialization
+            scale = self.embed_dim ** -0.5
+
+            if self.num_cls_tokens > 0:
+                nn.init.normal_(self.cls_token)
+                self.cls_token *= scale
+        elif init_param_style == "vit":
+            self.cls_token.data.fill_(0)
+        else:
+            raise ValueError(f"Unknown init {init_param_style}")
+
+    def tokenize_input_and_cls_pos(self, input, stem):
+        # tokens is of shape B x L x D
+        tokens = stem.norm_layer(stem.proj(input))
+        assert tokens.ndim == 3
+        assert tokens.shape[2] == self.embed_dim
+        B = tokens.shape[0]
+        if self.num_cls_tokens > 0:
+            class_tokens = self.cls_token.expand(
+                B, -1, -1
+            )  # stole class_tokens impl from Phil Wang, thanks
+            tokens = torch.cat((class_tokens, tokens), dim=1)
+        if self.use_pos_embed:
+            tokens = tokens + self.pos_embed
+        return tokens
+
+    def forward(self, imu):
+        # Patchify
+        imu = imu.unfold(
+            -1,
+            self.kernel_size,
+            self.kernel_size,
+        ).permute(0, 2, 1, 3)
+        imu = imu.reshape(imu.size(0), imu.size(1), -1)
+
+        imu_tokens = self.tokenize_input_and_cls_pos(
+            imu,
+            self.imu_stem,
+        )
+
+        return_dict = {
+            "trunk": {
+                "tokens": imu_tokens,
+            },
+            "head": {},
+        }
+        return return_dict
+
+
+def cast_if_src_dtype(
+        tensor: torch.Tensor, src_dtype: torch.dtype, tgt_dtype: torch.dtype
+):
+    updated = False
+    if tensor.dtype == src_dtype:
+        tensor = tensor.to(dtype=tgt_dtype)
+        updated = True
+    return tensor, updated
+
+
+class QuickGELU(nn.Module):
+    # From https://github.com/openai/CLIP/blob/d50d76daa670286dd6cacf3bcd80b5e4823fc8e1/clip/model.py#L166
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class SelectElement(nn.Module):
+    def __init__(self, index) -> None:
+        super().__init__()
+        self.index = index
+
+    def forward(self, x):
+        assert x.ndim >= 3
+        return x[:, self.index, ...]
+
+
+class ReshapeSpatial(nn.Module):
+    def __init__(self, shape) -> None:
+        super().__init__()
+        self.h, self.w = shape
+
+    def forward(self, x):
+        assert x.ndim >= 3
+        return x[:, 1:, ...].reshape(x.shape[0], self.h, self.w, -1), x[:, 0, :]
+
+
+class ReshapeAudio(nn.Module):
+    def __init__(self, shape) -> None:
+        super().__init__()
+        self.h, self.w = shape
+
+    def forward(self, x):
+        assert x.ndim == 3
+        return x[:, 1:, :].reshape(-1, 5, self.h, self.w, x.shape[-1]), x[:, 0, :]
+
+
+class ApplyTwice(nn.Module):
+    def __init__(self, module) -> None:
+        super().__init__()
+        self.module = module
+
+    def forward(self, pair):
+        return self.module(pair[0]), self.module(pair[1])
+
+
+class SelectEOSAndProject(nn.Module):
+    """
+    Text Pooling used in OpenCLIP
+    """
+
+    def __init__(self, proj: nn.Module) -> None:
+        super().__init__()
+        self.proj = proj
+
+    def forward(self, x, seq_len):
+        assert x.ndim == 3
+        # x is of shape B x L x D
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), seq_len]
+        x = self.proj(x)
+        return x
+
+
+ModalityType = SimpleNamespace(
+    VISION="vision",
+    TEXT="text",
+    AUDIO="audio",
+    THERMAL="thermal",
+    DEPTH="depth",
+    IMU="imu",
+)
+
+
+class ImageBindModel(nn.Module):
+    def __init__(
+            self,
+            video_frames=2,
+            kernel_size=(2, 14, 14),
+            audio_kernel_size=16,
+            audio_stride=10,
+            out_embed_dim=768,
+            vision_embed_dim=1024,
+            vision_num_blocks=24,
+            vision_num_heads=16,
+            audio_embed_dim=768,
+            audio_num_blocks=12,
+            audio_num_heads=12,
+            audio_num_mel_bins=128,
+            audio_target_len=204,
+            audio_drop_path=0.1,
+            text_embed_dim=768,
+            text_num_blocks=12,
+            text_num_heads=12,
+            depth_embed_dim=384,
+            depth_kernel_size=16,
+            depth_num_blocks=12,
+            depth_num_heads=8,
+            depth_drop_path=0.0,
+            thermal_embed_dim=768,
+            thermal_kernel_size=16,
+            thermal_num_blocks=12,
+            thermal_num_heads=12,
+            thermal_drop_path=0.0,
+            imu_embed_dim=512,
+            imu_kernel_size=8,
+            imu_num_blocks=6,
+            imu_num_heads=8,
+            imu_drop_path=0.7,
+    ):
+        super().__init__()
+
+        self.modality_preprocessors = self._create_modality_preprocessors(
+            video_frames,
+            vision_embed_dim,
+            kernel_size,
+            text_embed_dim,
+            audio_embed_dim,
+            audio_kernel_size,
+            audio_stride,
+            audio_num_mel_bins,
+            audio_target_len,
+            depth_embed_dim,
+            depth_kernel_size,
+            thermal_embed_dim,
+            thermal_kernel_size,
+            imu_embed_dim,
+        )
+
+        self.modality_trunks = self._create_modality_trunks(
+            vision_embed_dim,
+            vision_num_blocks,
+            vision_num_heads,
+            text_embed_dim,
+            text_num_blocks,
+            text_num_heads,
+            audio_embed_dim,
+            audio_num_blocks,
+            audio_num_heads,
+            audio_drop_path,
+            depth_embed_dim,
+            depth_num_blocks,
+            depth_num_heads,
+            depth_drop_path,
+            thermal_embed_dim,
+            thermal_num_blocks,
+            thermal_num_heads,
+            thermal_drop_path,
+            imu_embed_dim,
+            imu_num_blocks,
+            imu_num_heads,
+            imu_drop_path,
+        )
+
+        self.modality_heads = self._create_modality_heads(
+            out_embed_dim,
+            vision_embed_dim,
+            text_embed_dim,
+            audio_embed_dim,
+            depth_embed_dim,
+            thermal_embed_dim,
+            imu_embed_dim,
+        )
+
+        self.modality_postprocessors = self._create_modality_postprocessors(
+            out_embed_dim
+        )
+
+    def _create_modality_preprocessors(
+            self,
+            video_frames=2,
+            vision_embed_dim=1024,
+            kernel_size=(2, 14, 14),
+            text_embed_dim=768,
+            audio_embed_dim=768,
+            audio_kernel_size=16,
+            audio_stride=10,
+            audio_num_mel_bins=128,
+            audio_target_len=204,
+            depth_embed_dim=768,
+            depth_kernel_size=16,
+            thermal_embed_dim=768,
+            thermal_kernel_size=16,
+            imu_embed_dim=512,
+    ):
+        rgbt_stem = PatchEmbedGeneric(
+            proj_stem=[
+                PadIm2Video(pad_type="repeat", ntimes=2),
+                nn.Conv3d(
+                    in_channels=3,
+                    kernel_size=kernel_size,
+                    out_channels=vision_embed_dim,
+                    stride=kernel_size,
+                    bias=False,
+                ),
+            ]
+        )
+        rgbt_preprocessor = RGBDTPreprocessor(
+            img_size=[3, video_frames, 224, 224],
+            num_cls_tokens=1,
+            pos_embed_fn=partial(SpatioTemporalPosEmbeddingHelper, learnable=True),
+            rgbt_stem=rgbt_stem,
+            depth_stem=None,
+        )
+
+        text_preprocessor = TextPreprocessor(
+            context_length=77,
+            vocab_size=49408,
+            embed_dim=text_embed_dim,
+            causal_masking=True,
+        )
+
+        audio_stem = PatchEmbedGeneric(
+            proj_stem=[
+                nn.Conv2d(
+                    in_channels=1,
+                    kernel_size=audio_kernel_size,
+                    stride=audio_stride,
+                    out_channels=audio_embed_dim,
+                    bias=False,
+                ),
+            ],
+            norm_layer=nn.LayerNorm(normalized_shape=audio_embed_dim),
+        )
+        audio_preprocessor = AudioPreprocessor(
+            img_size=[1, audio_num_mel_bins, audio_target_len],
+            num_cls_tokens=1,
+            pos_embed_fn=partial(SpatioTemporalPosEmbeddingHelper, learnable=True),
+            audio_stem=audio_stem,
+        )
+
+        # depth_stem = PatchEmbedGeneric(
+        #     [
+        #         nn.Conv2d(
+        #             kernel_size=depth_kernel_size,
+        #             in_channels=1,
+        #             out_channels=depth_embed_dim,
+        #             stride=depth_kernel_size,
+        #             bias=False,
+        #         ),
+        #     ],
+        #     norm_layer=nn.LayerNorm(normalized_shape=depth_embed_dim),
+        # )
+        #
+        # depth_preprocessor = RGBDTPreprocessor(
+        #     img_size=[1, 224, 224],
+        #     num_cls_tokens=1,
+        #     pos_embed_fn=partial(SpatioTemporalPosEmbeddingHelper, learnable=True),
+        #     rgbt_stem=None,
+        #     depth_stem=depth_stem,
+        # )
+        #
+        # thermal_stem = PatchEmbedGeneric(
+        #     [
+        #         nn.Conv2d(
+        #             kernel_size=thermal_kernel_size,
+        #             in_channels=1,
+        #             out_channels=thermal_embed_dim,
+        #             stride=thermal_kernel_size,
+        #             bias=False,
+        #         ),
+        #     ],
+        #     norm_layer=nn.LayerNorm(normalized_shape=thermal_embed_dim),
+        # )
+        # thermal_preprocessor = ThermalPreprocessor(
+        #     img_size=[1, 224, 224],
+        #     num_cls_tokens=1,
+        #     pos_embed_fn=partial(SpatioTemporalPosEmbeddingHelper, learnable=True),
+        #     thermal_stem=thermal_stem,
+        # )
+        #
+        # imu_stem = PatchEmbedGeneric(
+        #     [
+        #         nn.Linear(
+        #             in_features=48,
+        #             out_features=imu_embed_dim,
+        #             bias=False,
+        #         ),
+        #     ],
+        #     norm_layer=nn.LayerNorm(normalized_shape=imu_embed_dim),
+        # )
+        #
+        # imu_preprocessor = IMUPreprocessor(
+        #     img_size=[6, 2000],
+        #     num_cls_tokens=1,
+        #     kernel_size=8,
+        #     embed_dim=imu_embed_dim,
+        #     pos_embed_fn=partial(SpatioTemporalPosEmbeddingHelper, learnable=True),
+        #     imu_stem=imu_stem,
+        # )
+
+        modality_preprocessors = {
+            ModalityType.VISION: rgbt_preprocessor,
+            ModalityType.TEXT: text_preprocessor,
+            ModalityType.AUDIO: audio_preprocessor,
+            # ModalityType.DEPTH: depth_preprocessor,
+            # ModalityType.THERMAL: thermal_preprocessor,
+            # ModalityType.IMU: imu_preprocessor,
+        }
+
+        return nn.ModuleDict(modality_preprocessors)
+
+    def _create_modality_trunks(
+            self,
+            vision_embed_dim=1024,
+            vision_num_blocks=24,
+            vision_num_heads=16,
+            text_embed_dim=768,
+            text_num_blocks=12,
+            text_num_heads=12,
+            audio_embed_dim=768,
+            audio_num_blocks=12,
+            audio_num_heads=12,
+            audio_drop_path=0.0,
+            depth_embed_dim=768,
+            depth_num_blocks=12,
+            depth_num_heads=12,
+            depth_drop_path=0.0,
+            thermal_embed_dim=768,
+            thermal_num_blocks=12,
+            thermal_num_heads=12,
+            thermal_drop_path=0.0,
+            imu_embed_dim=512,
+            imu_num_blocks=6,
+            imu_num_heads=8,
+            imu_drop_path=0.7,
+    ):
+        def instantiate_trunk(
+                embed_dim, num_blocks, num_heads, pre_transformer_ln, add_bias_kv, drop_path
+        ):
+            return SimpleTransformer(
+                embed_dim=embed_dim,
+                num_blocks=num_blocks,
+                ffn_dropout_rate=0.0,
+                drop_path_rate=drop_path,
+                attn_target=partial(
+                    MultiheadAttention,
+                    embed_dim=embed_dim,
+                    num_heads=num_heads,
+                    bias=True,
+                    add_bias_kv=add_bias_kv,
+                ),
+                pre_transformer_layer=nn.Sequential(
+                    nn.LayerNorm(embed_dim, eps=1e-6)
+                    if pre_transformer_ln
+                    else nn.Identity(),
+                    EinOpsRearrange("b l d -> l b d"),
+                ),
+                post_transformer_layer=EinOpsRearrange("l b d -> b l d"),
+            )
+
+        modality_trunks = {}
+        modality_trunks[ModalityType.VISION] = instantiate_trunk(
+            vision_embed_dim,
+            vision_num_blocks,
+            vision_num_heads,
+            pre_transformer_ln=True,
+            add_bias_kv=False,
+            drop_path=0.0,
+        )
+        modality_trunks[ModalityType.TEXT] = instantiate_trunk(
+            text_embed_dim,
+            text_num_blocks,
+            text_num_heads,
+            pre_transformer_ln=False,
+            add_bias_kv=False,
+            drop_path=0.0,
+        )
+        modality_trunks[ModalityType.AUDIO] = instantiate_trunk(
+            audio_embed_dim,
+            audio_num_blocks,
+            audio_num_heads,
+            pre_transformer_ln=False,
+            add_bias_kv=True,
+            drop_path=audio_drop_path,
+        )
+        # modality_trunks[ModalityType.DEPTH] = instantiate_trunk(
+        #     depth_embed_dim,
+        #     depth_num_blocks,
+        #     depth_num_heads,
+        #     pre_transformer_ln=False,
+        #     add_bias_kv=True,
+        #     drop_path=depth_drop_path,
+        # )
+        # modality_trunks[ModalityType.THERMAL] = instantiate_trunk(
+        #     thermal_embed_dim,
+        #     thermal_num_blocks,
+        #     thermal_num_heads,
+        #     pre_transformer_ln=False,
+        #     add_bias_kv=True,
+        #     drop_path=thermal_drop_path,
+        # )
+        # modality_trunks[ModalityType.IMU] = instantiate_trunk(
+        #     imu_embed_dim,
+        #     imu_num_blocks,
+        #     imu_num_heads,
+        #     pre_transformer_ln=False,
+        #     add_bias_kv=True,
+        #     drop_path=imu_drop_path,
+        # )
+
+        return nn.ModuleDict(modality_trunks)
+
+    def _create_modality_heads(
+            self,
+            out_embed_dim,
+            vision_embed_dim,
+            text_embed_dim,
+            audio_embed_dim,
+            depth_embed_dim,
+            thermal_embed_dim,
+            imu_embed_dim,
+    ):
+        modality_heads = {}
+
+        modality_heads[ModalityType.VISION] = nn.Sequential(
+            nn.LayerNorm(normalized_shape=vision_embed_dim, eps=1e-6),
+            SelectElement(index=0),
+            nn.Linear(vision_embed_dim, out_embed_dim, bias=False),
+        )
+
+        modality_heads[ModalityType.TEXT] = SelectEOSAndProject(
+            proj=nn.Sequential(
+                nn.LayerNorm(normalized_shape=text_embed_dim, eps=1e-6),
+                nn.Linear(text_embed_dim, out_embed_dim, bias=False),
+            )
+        )
+
+        modality_heads[ModalityType.AUDIO] = nn.Sequential(
+            nn.LayerNorm(normalized_shape=audio_embed_dim, eps=1e-6),
+            SelectElement(index=0),
+            nn.Linear(audio_embed_dim, out_embed_dim, bias=False),
+        )
+
+        # modality_heads[ModalityType.DEPTH] = nn.Sequential(
+        #     nn.LayerNorm(normalized_shape=depth_embed_dim, eps=1e-6),
+        #     SelectElement(index=0),
+        #     nn.Linear(depth_embed_dim, out_embed_dim, bias=False),
+        # )
+        #
+        # modality_heads[ModalityType.THERMAL] = nn.Sequential(
+        #     nn.LayerNorm(normalized_shape=thermal_embed_dim, eps=1e-6),
+        #     SelectElement(index=0),
+        #     nn.Linear(thermal_embed_dim, out_embed_dim, bias=False),
+        # )
+        #
+        # modality_heads[ModalityType.IMU] = nn.Sequential(
+        #     nn.LayerNorm(normalized_shape=imu_embed_dim, eps=1e-6),
+        #     SelectElement(index=0),
+        #     nn.Dropout(p=0.5),
+        #     nn.Linear(imu_embed_dim, out_embed_dim, bias=False),
+        # )
+
+        return nn.ModuleDict(modality_heads)
+
+    def _create_modality_postprocessors(self, out_embed_dim):
+        modality_postprocessors = {}
+
+        modality_postprocessors[ModalityType.VISION] = Normalize(dim=-1)
+        modality_postprocessors[ModalityType.TEXT] = nn.Sequential(
+            Normalize(dim=-1), LearnableLogitScaling(learnable=True)
+        )
+        modality_postprocessors[ModalityType.AUDIO] = nn.Sequential(
+            Normalize(dim=-1),
+            LearnableLogitScaling(logit_scale_init=20.0, learnable=False),
+        )
+        # modality_postprocessors[ModalityType.DEPTH] = nn.Sequential(
+        #     Normalize(dim=-1),
+        #     LearnableLogitScaling(logit_scale_init=5.0, learnable=False),
+        # )
+        # modality_postprocessors[ModalityType.THERMAL] = nn.Sequential(
+        #     Normalize(dim=-1),
+        #     LearnableLogitScaling(logit_scale_init=10.0, learnable=False),
+        # )
+        # modality_postprocessors[ModalityType.IMU] = nn.Sequential(
+        #     Normalize(dim=-1),
+        #     LearnableLogitScaling(logit_scale_init=5.0, learnable=False),
+        # )
+
+        return nn.ModuleDict(modality_postprocessors)
+
+    def forward(self, inputs):
+        outputs = {}
+        for modality_key, modality_value in inputs.items():
+            reduce_list = (
+                    modality_value.ndim >= 5
+            )  # Audio and Video inputs consist of multiple clips
+            if reduce_list:
+                B, S = modality_value.shape[:2]
+                modality_value = modality_value.reshape(
+                    B * S, *modality_value.shape[2:]
+                )
+
+            if modality_value is not None:
+                modality_value = self.modality_preprocessors[modality_key](
+                    **{modality_key: modality_value}
+                )
+                trunk_inputs = modality_value["trunk"]
+                head_inputs = modality_value["head"]
+                modality_value = self.modality_trunks[modality_key](**trunk_inputs)
+                modality_value = self.modality_heads[modality_key](
+                    modality_value, **head_inputs
+                )
+                modality_value = self.modality_postprocessors[modality_key](
+                    modality_value
+                )
+
+                if reduce_list:
+                    modality_value = modality_value.reshape(B, S, -1)
+                    modality_value = modality_value.mean(dim=1)
+
+                outputs[modality_key] = modality_value
+
+        return outputs
+
+    def reconfigure_head(self, k, v):
+        if k == ModalityType.AUDIO:
+            return torch.nn.Sequential(v[0], v[2])
+        elif k == ModalityType.VISION:
+            return torch.nn.Sequential(v[0], v[2])
+        else:
+            return v
+
+    def forward_features(self, inputs):
+        outputs = {}
+
+        reconfigured_heads = {k: self.reconfigure_head(k, v) for k, v in self.modality_heads.items()}
+
+        for modality_key, modality_value in inputs.items():
+            reduce_list = (
+                    modality_value.ndim >= 5
+            )  # Audio and Video inputs consist of multiple clips
+            if reduce_list:
+                B, S = modality_value.shape[:2]
+                modality_value = modality_value.reshape(
+                    B * S, *modality_value.shape[2:]
+                )
+
+            if modality_value is not None:
+                modality_value = self.modality_preprocessors[modality_key](
+                    **{modality_key: modality_value}
+                )
+                trunk_inputs = modality_value["trunk"]
+                head_inputs = modality_value["head"]
+                modality_value = self.modality_trunks[modality_key](**trunk_inputs)
+                modality_value = reconfigured_heads[modality_key](
+                    modality_value, **head_inputs
+                )
+                modality_value = self.modality_postprocessors[modality_key](
+                    modality_value
+                )
+                if modality_key == ModalityType.AUDIO:
+                    modality_value = ReshapeAudio((12, 19))(modality_value)
+                elif modality_key == ModalityType.VISION:
+                    modality_value = ReshapeSpatial((16, 16))(modality_value)
+
+                outputs[modality_key] = modality_value
+
+        return outputs
+
+
+def imagebind_huge(output_path, pretrained=False):
+    model = ImageBindModel(
+        vision_embed_dim=1280,
+        vision_num_blocks=32,
+        vision_num_heads=16,
+        text_embed_dim=1024,
+        text_num_blocks=24,
+        text_num_heads=16,
+        out_embed_dim=1024,
+        audio_drop_path=0.1,
+        imu_drop_path=0.7,
+    )
+
+    if pretrained:
+        path = os.path.join(output_path, 'models/imagebind_huge.pth')
+
+        if not os.path.exists(path):
+            print(f"Downloading imagebind weights to {path} ...")
+            os.makedirs(os.path.dirname(path), exist_ok=True)
+            torch.hub.download_url_to_file(
+                "https://dl.fbaipublicfiles.com/imagebind/imagebind_huge.pth",
+                path,
+                progress=True,
+            )
+
+        model.load_state_dict(torch.load(path), strict=False)
+
+    return model
+
+
+DEFAULT_AUDIO_FRAME_SHIFT_MS = 10  # in milliseconds
+
+
+def waveform2melspec(waveform, sample_rate, num_mel_bins, target_length):
+    # Based on https://github.com/YuanGongND/ast/blob/d7d8b4b8e06cdaeb6c843cdb38794c1c7692234c/src/dataloader.py#L102
+    waveform -= waveform.mean()
+    fbank = torchaudio.compliance.kaldi.fbank(
+        waveform,
+        htk_compat=True,
+        sample_frequency=sample_rate,
+        use_energy=False,
+        window_type="hanning",
+        num_mel_bins=num_mel_bins,
+        dither=0.0,
+        frame_length=25,
+        frame_shift=DEFAULT_AUDIO_FRAME_SHIFT_MS,
+    )
+    # Convert to [mel_bins, num_frames] shape
+    fbank = fbank.transpose(0, 1)
+    # Pad to target_length
+    n_frames = fbank.size(1)
+    p = target_length - n_frames
+    # if p is too large (say >20%), flash a warning
+    if abs(p) / n_frames > 0.2:
+        logging.warning(
+            "Large gap between audio n_frames(%d) and "
+            "target_length (%d). Is the audio_target_length "
+            "setting correct?",
+            n_frames,
+            target_length,
+        )
+    # cut and pad
+    if p > 0:
+        fbank = torch.nn.functional.pad(fbank, (0, p), mode="constant", value=0)
+    elif p < 0:
+        fbank = fbank[:, 0:target_length]
+    # Convert to [1, mel_bins, num_frames] shape, essentially like a 1
+    # channel image
+    fbank = fbank.unsqueeze(0)
+    return fbank
+
+
+def get_clip_timepoints(clip_sampler, duration):
+    # Read out all clips in this video
+    all_clips_timepoints = []
+    is_last_clip = False
+    end = 0.0
+    while not is_last_clip:
+        start, end, _, _, is_last_clip = clip_sampler(end, duration, annotation=None)
+        all_clips_timepoints.append((start, end))
+    return all_clips_timepoints
+
+
+def load_and_transform_vision_data(image_paths, device):
+    if image_paths is None:
+        return None
+
+    image_ouputs = []
+    for image_path in image_paths:
+        data_transform = transforms.Compose(
+            [
+                transforms.Resize(
+                    224, interpolation=transforms.InterpolationMode.BICUBIC
+                ),
+                transforms.CenterCrop(224),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=(0.48145466, 0.4578275, 0.40821073),
+                    std=(0.26862954, 0.26130258, 0.27577711),
+                ),
+            ]
+        )
+        with open(image_path, "rb") as fopen:
+            image = Image.open(fopen).convert("RGB")
+
+        image = data_transform(image).to(device)
+        image_ouputs.append(image)
+    return torch.stack(image_ouputs, dim=0)
+
+
+def load_and_transform_audio_data(
+        audio_paths,
+        device,
+        num_mel_bins=128,
+        target_length=204,
+        sample_rate=16000,
+        clip_duration=2,
+        clips_per_video=3,
+        mean=-4.268,
+        std=9.138,
+):
+    from pytorchvideo.data.clip_sampling import ConstantClipsPerVideoSampler
+
+    if audio_paths is None:
+        return None
+
+    audio_outputs = []
+    clip_sampler = ConstantClipsPerVideoSampler(
+        clip_duration=clip_duration, clips_per_video=clips_per_video
+    )
+
+    for audio_path in audio_paths:
+        waveform, sr = torchaudio.load(audio_path)
+        if sample_rate != sr:
+            waveform = torchaudio.functional.resample(
+                waveform, orig_freq=sr, new_freq=sample_rate
+            )
+        all_clips_timepoints = get_clip_timepoints(
+            clip_sampler, waveform.size(1) / sample_rate
+        )
+        all_clips = []
+        for clip_timepoints in all_clips_timepoints:
+            waveform_clip = waveform[
+                            :,
+                            int(clip_timepoints[0] * sample_rate): int(
+                                clip_timepoints[1] * sample_rate
+                            ),
+                            ]
+            waveform_melspec = waveform2melspec(
+                waveform_clip, sample_rate, num_mel_bins, target_length
+            )
+            all_clips.append(waveform_melspec)
+
+        normalize = transforms.Normalize(mean=mean, std=std)
+        all_clips = [normalize(ac).to(device) for ac in all_clips]
+
+        all_clips = torch.stack(all_clips, dim=0)
+        audio_outputs.append(all_clips)
+
+    return torch.stack(audio_outputs, dim=0)
+
+
+class UnNormalize(object):
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, image):
+        image2 = torch.clone(image)
+        for t, m, s in zip(image2, self.mean, self.std):
+            t.mul_(s).add_(m)
+        return image2
+
+
+norm = T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+unnorm = UnNormalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+
+
+class TorchPCA(object):
+
+    def __init__(self, n_components):
+        self.n_components = n_components
+
+    def fit(self, X):
+        self.mean_ = X.mean(dim=0)
+        unbiased = X - self.mean_.unsqueeze(0)
+        U, S, V = torch.pca_lowrank(unbiased, q=self.n_components, center=False, niter=4)
+        self.components_ = V.T
+        self.singular_values_ = S
+        return self
+
+    def transform(self, X):
+        t0 = X - self.mean_.unsqueeze(0)
+        projected = t0 @ self.components_.T
+        return projected
+
+
+def pca(image_feats_list, dim=3, fit_pca=None):
+    # from sklearn.decomposition import PCA
+
+    device = image_feats_list[0].device
+
+    def flatten(tensor, target_size=None):
+        if target_size is not None and fit_pca is None:
+            F.interpolate(tensor, (target_size, target_size), mode="bilinear")
+        B, C, H, W = tensor.shape
+        return feats.permute(1, 0, 2, 3).reshape(C, B * H * W).permute(1, 0).detach().cpu()
+
+    if len(image_feats_list) > 1 and fit_pca is None:
+        target_size = image_feats_list[0].shape[2]
+    else:
+        target_size = None
+
+    flattened_feats = []
+    for feats in image_feats_list:
+        flattened_feats.append(flatten(feats, target_size))
+    x = torch.cat(flattened_feats, dim=0)
+
+    if fit_pca is None:
+        # fit_pca = PCA(n_components=dim, svd_solver='arpack').fit(np.nan_to_num(x.detach().numpy()))
+        fit_pca = TorchPCA(n_components=dim).fit(x)
+
+    reduced_feats = []
+    for feats in image_feats_list:
+        # x_red = torch.from_numpy(fit_pca.transform(flatten(feats)))
+        x_red = fit_pca.transform(flatten(feats))
+        x_red -= x_red.min(dim=0, keepdim=True).values
+        x_red /= x_red.max(dim=0, keepdim=True).values
+        B, C, H, W = feats.shape
+        reduced_feats.append(x_red.reshape(B, H, W, dim).permute(0, 3, 1, 2).to(device))
+
+    return reduced_feats, fit_pca
+
+
+def my_load_audio(audio_file):
+    loaded_waveform, obs_sr = torchaudio.load(audio_file)
+    loaded_waveform = loaded_waveform[0]
+
+    neg_waveform, neg_obs_sr = None, None
+    from data.AVDatasets import prep_waveform
+
+    (waveform,
+     spectrogram,
+     audio_length,
+     total_length,
+     original_length,
+     mask,
+     pos_mask) = prep_waveform(
+        loaded_waveform,
+        obs_sr,
+        10,
+        128,
+        -4.268,
+        9.138,
+        16000,
+        True,
+        False,
+        False,
+        neg_waveform,
+        neg_obs_sr,
+        False,
+    )
+
+    patch_size = 204
+    n_tiles = spectrogram.shape[1] // patch_size
+    assert n_tiles == 5
+
+    patches = []
+    for i in range(n_tiles):
+        patches.append(spectrogram[:, i * patch_size:(i + 1) * patch_size, :])
+
+    patches = torch.cat(patches, dim=0).permute(0, 2, 1).unsqueeze(1)
+    return patches
+
+
+class ImageBindImageFeaturizer(nn.Module):
+
+    def __init__(self, output_path, model=None):
+        super().__init__()
+        if model is not None:
+            self.model = model
+        else:
+            self.model = imagebind_huge(output_path, pretrained=True).cuda()
+
+    def forward(self, image, include_cls):
+        inputs = {
+            ModalityType.VISION: image,
+        }
+
+        patch_tokens, cls_tokens = self.model.forward_features(inputs)[ModalityType.VISION]
+        patch_tokens = patch_tokens.permute(0, 3, 1, 2)
+
+        if include_cls:
+            return patch_tokens, cls_tokens
+        else:
+            return patch_tokens
+
+
+class ImageBindAudioFeaturizer(nn.Module):
+
+    def __init__(self, output_path, model=None):
+        super().__init__()
+        if model is not None:
+            self.model = model
+        else:
+            self.model = imagebind_huge(output_path, pretrained=True).cuda()
+
+    def forward(self, spec, include_cls):
+
+        patch_size = 204
+        n_tiles = spec.shape[2] // patch_size
+        assert n_tiles == 5
+
+        patches = []
+        for i in range(n_tiles):
+            patches.append(spec[:, :, i * patch_size:(i + 1) * patch_size, :])
+
+        patches = torch.cat(patches, dim=1).permute(0, 1, 3, 2).unsqueeze(2)
+
+        inputs = {
+            ModalityType.AUDIO: patches,
+        }
+
+        patch_tokens, cls_token = self.model.forward_features(inputs)[ModalityType.AUDIO]
+
+        patch_tokens = patch_tokens.permute(0, 4, 2, 1, 3)
+        b, c, h, p, w = patch_tokens.shape
+        patch_tokens = patch_tokens.reshape(b, c, h, w * p)
+
+        cls_token = cls_token.reshape(b, p, -1).mean(1)
+
+        if include_cls:
+            return patch_tokens, cls_token
+        else:
+            return patch_tokens
+
+
+if __name__ == "__main__":
+    image_paths = ["../../samples/dog_image.jpg", "../../samples/car_image.jpg", "../../samples/bird_image.jpg"]
+    audio_paths = ["../../samples/dog_audio.wav", "../../samples/car_audio.wav", "../../samples/bird_audio.wav"]
+
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+    # Instantiate model
+    model = imagebind_huge("../../", pretrained=True)
+    model.eval()
+    model.to(device)
+
+    audio_inputs = torch.cat([my_load_audio(af).unsqueeze(0) for af in audio_paths], dim=0).cuda()
+    # Load data
+    inputs = {
+        ModalityType.VISION: load_and_transform_vision_data(image_paths, device),
+        # ModalityType.AUDIO: load_and_transform_audio_data(audio_paths, device, clip_duration=2, clips_per_video=5),
+        ModalityType.AUDIO: audio_inputs,
+
+    }
+
+    with torch.no_grad():
+        embeddings = model.forward_features(inputs)
+        cls_tokens = model.forward(inputs)
+
+    audio_cls_token = embeddings["audio"][1].reshape(3, 5, -1).mean(1)
+
+    sims1 = torch.einsum(
+        "bc,dc->bd",
+        embeddings["vision"][1],
+        audio_cls_token)
+
+    print(torch.softmax(sims1, dim=1).cpu().numpy())
+    #
+    # sims2 = torch.einsum(
+    #     "bc,dc->bd",
+    #     embeddings["vision"].mean(1).mean(1),
+    #     embeddings["audio"].mean(1).mean(1).mean(1)
+    # )
+    #
+    # print(torch.softmax(sims2, dim=1).cpu().numpy())
+    #
+    #
+    # img_num = 0
+    # img_feats = F.normalize(embeddings["vision"].permute(0, 3, 1, 2), dim=1)
+    # [red_img_feats], fit_pca = pca([img_feats])
+    #
+    # fig, axes = plt.subplots(2, 2, figsize=(4 * 2, 4 * 2))
+    # axes[0][0].imshow(unnorm(inputs["vision"][0].unsqueeze(0))[0].permute(1, 2, 0).detach().cpu())
+    # axes[0][1].imshow(unnorm(inputs["vision"][1].unsqueeze(0))[0].permute(1, 2, 0).detach().cpu())
+    # axes[1][0].imshow(red_img_feats[0].permute(1, 2, 0).detach().cpu())
+    # axes[1][1].imshow(red_img_feats[1].permute(1, 2, 0).detach().cpu())
+    # plt.tight_layout()
+    # plt.show()
+    #
+    audio_embs = F.normalize(embeddings["audio"][0], dim=-1)
+    b, n, h, w, c = audio_embs.shape
+
+    audio_embs = audio_embs.permute(0, 4, 2, 1, 3).reshape(b, c, h, w * n)
+
+    b, n, c, h, w = inputs["audio"].shape
+    audio_inputs = inputs["audio"].permute(0, 2, 3, 1, 4).reshape(b, c, h, w * n)
+
+    print("here")
+
+    for img_num in range(3):
+        [red_audio], fit_pca = pca([audio_embs[img_num].unsqueeze(0)])
+        fig, axes = plt.subplots(2, 1, figsize=(10 * 1, 4 * 2))
+        axes[0].imshow(audio_inputs[img_num, 0].detach().cpu())
+        axes[1].imshow(red_audio[0].permute(1, 2, 0).detach().cpu())
+        plt.tight_layout()
+        plt.show()
+
+    print("here")

DenseAV/denseav/plotting.py

@@ -0,0 +1,244 @@
+import os
+from collections import defaultdict
+
+import matplotlib.colors as mcolors
+import matplotlib.pyplot as plt
+import numpy as np
+import scipy.io.wavfile as wavfile
+import torch
+import torch.nn.functional as F
+import torchvision
+from moviepy.editor import VideoFileClip, AudioFileClip
+from base64 import b64encode
+from denseav.shared import pca
+
+
+def write_video_with_audio(video_frames, audio_array, video_fps, audio_fps, output_path):
+    """
+    Writes video frames and audio to a specified path.
+
+    Parameters:
+    - video_frames: torch.Tensor of shape (num_frames, height, width, channels)
+    - audio_array: torch.Tensor of shape (num_samples, num_channels)
+    - video_fps: int, frames per second of the video
+    - audio_fps: int, sample rate of the audio
+    - output_path: str, path to save the final video with audio
+    """
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    temp_video_path = output_path.replace('.mp4', '_temp.mp4')
+    temp_audio_path = output_path.replace('.mp4', '_temp_audio.wav')
+    video_options = {
+        'crf': '23',
+        'preset': 'slow',
+        'bit_rate': '1000k'}
+
+    if audio_array is not None:
+        torchvision.io.write_video(
+            filename=temp_video_path,
+            video_array=video_frames,
+            fps=video_fps,
+            options=video_options
+        )
+
+        wavfile.write(temp_audio_path, audio_fps, audio_array.cpu().to(torch.float64).permute(1, 0).numpy())
+        video_clip = VideoFileClip(temp_video_path)
+        audio_clip = AudioFileClip(temp_audio_path)
+        final_clip = video_clip.set_audio(audio_clip)
+        final_clip.write_videofile(output_path, codec='libx264', verbose=False)
+        os.remove(temp_video_path)
+        os.remove(temp_audio_path)
+    else:
+        torchvision.io.write_video(
+            filename=output_path,
+            video_array=video_frames,
+            fps=video_fps,
+            options=video_options
+        )
+
+
+def alpha_blend_layers(layers):
+    blended_image = layers[0]
+    for layer in layers[1:]:
+        rgb1, alpha1 = blended_image[:, :3, :, :], blended_image[:, 3:4, :, :]
+        rgb2, alpha2 = layer[:, :3, :, :], layer[:, 3:4, :, :]
+        alpha_out = alpha2 + alpha1 * (1 - alpha2)
+        rgb_out = (rgb2 * alpha2 + rgb1 * alpha1 * (1 - alpha2)) / alpha_out.clamp(min=1e-7)
+        blended_image = torch.cat([rgb_out, alpha_out], dim=1)
+    return (blended_image[:, :3] * 255).clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1)
+
+
+def _prep_sims_for_plotting(sim_by_head, frames):
+    with torch.no_grad():
+        results = defaultdict(list)
+        n_frames, _, vh, vw = frames.shape
+
+        sims = sim_by_head.max(dim=1).values
+
+        n_audio_feats = sims.shape[-1]
+        for frame_num in range(n_frames):
+            selected_audio_feat = int((frame_num / n_frames) * n_audio_feats)
+
+            selected_sim = F.interpolate(
+                sims[frame_num, :, :, selected_audio_feat].unsqueeze(0).unsqueeze(0),
+                size=(vh, vw),
+                mode="bicubic")
+
+            results["sims_all"].append(selected_sim)
+
+            for head in range(sim_by_head.shape[1]):
+                selected_sim = F.interpolate(
+                    sim_by_head[frame_num, head, :, :, selected_audio_feat].unsqueeze(0).unsqueeze(0),
+                    size=(vh, vw),
+                    mode="bicubic")
+                results[f"sims_{head + 1}"].append(selected_sim)
+
+        results = {k: torch.cat(v, dim=0) for k, v in results.items()}
+        return results
+
+
+def get_plasma_with_alpha():
+    plasma = plt.cm.plasma(np.linspace(0, 1, 256))
+    alphas = np.linspace(0, 1, 256)
+    plasma_with_alpha = np.zeros((256, 4))
+    plasma_with_alpha[:, 0:3] = plasma[:, 0:3]
+    plasma_with_alpha[:, 3] = alphas
+    return mcolors.ListedColormap(plasma_with_alpha)
+
+
+def get_inferno_with_alpha_2(alpha=0.5, k=30):
+    k_fraction = k / 100.0
+    custom_cmap = np.zeros((256, 4))
+    threshold_index = int(k_fraction * 256)
+    custom_cmap[:threshold_index, :3] = 0  # RGB values for black
+    custom_cmap[:threshold_index, 3] = alpha  # Alpha value
+    remaining_inferno = plt.cm.inferno(np.linspace(0, 1, 256 - threshold_index))
+    custom_cmap[threshold_index:, :3] = remaining_inferno[:, :3]
+    custom_cmap[threshold_index:, 3] = alpha  # Alpha value
+    return mcolors.ListedColormap(custom_cmap)
+
+
+def get_inferno_with_alpha():
+    plasma = plt.cm.inferno(np.linspace(0, 1, 256))
+    alphas = np.linspace(0, 1, 256)
+    plasma_with_alpha = np.zeros((256, 4))
+    plasma_with_alpha[:, 0:3] = plasma[:, 0:3]
+    plasma_with_alpha[:, 3] = alphas
+    return mcolors.ListedColormap(plasma_with_alpha)
+
+
+red_cmap = mcolors.LinearSegmentedColormap('RedMap', segmentdata={
+    'red': [(0.0, 1.0, 1.0), (1.0, 1.0, 1.0)],
+    'green': [(0.0, 0.0, 0.0), (1.0, 0.0, 0.0)],
+    'blue': [(0.0, 0.0, 0.0), (1.0, 0.0, 0.0)],
+    'alpha': [(0.0, 0.0, 0.0), (1.0, 1.0, 1.0)]
+})
+
+blue_cmap = mcolors.LinearSegmentedColormap('BlueMap', segmentdata={
+    'red': [(0.0, 0.0, 0.0), (1.0, 0.0, 0.0)],
+    'green': [(0.0, 0.0, 0.0), (1.0, 0.0, 0.0)],
+    'blue': [(0.0, 1.0, 1.0), (1.0, 1.0, 1.0)],
+    'alpha': [(0.0, 0.0, 0.0), (1.0, 1.0, 1.0)]
+})
+
+
+def plot_attention_video(sims_by_head, frames, audio, video_fps, audio_fps, output_filename):
+    prepped_sims = _prep_sims_for_plotting(sims_by_head, frames)
+    n_frames, _, vh, vw = frames.shape
+    sims_all = prepped_sims["sims_all"].clamp_min(0)
+    sims_all -= sims_all.min()
+    sims_all = sims_all / sims_all.max()
+    cmap = get_inferno_with_alpha()
+    layer1 = torch.cat([frames, torch.ones(n_frames, 1, vh, vw)], axis=1)
+    layer2 = torch.tensor(cmap(sims_all.squeeze().detach().cpu())).permute(0, 3, 1, 2)
+    write_video_with_audio(
+        alpha_blend_layers([layer1, layer2]),
+        audio,
+        video_fps,
+        audio_fps,
+        output_filename)
+
+
+def plot_2head_attention_video(sims_by_head, frames, audio, video_fps, audio_fps, output_filename):
+    prepped_sims = _prep_sims_for_plotting(sims_by_head, frames)
+    sims_1 = prepped_sims["sims_1"]
+    sims_2 = prepped_sims["sims_2"]
+
+    n_frames, _, vh, vw = frames.shape
+
+    mask = sims_1 > sims_2
+    sims_1 *= mask
+    sims_2 *= (~mask)
+
+    sims_1 = sims_1.clamp_min(0)
+    sims_1 -= sims_1.min()
+    sims_1 = sims_1 / sims_1.max()
+
+    sims_2 = sims_2.clamp_min(0)
+    sims_2 -= sims_2.min()
+    sims_2 = sims_2 / sims_2.max()
+
+    layer1 = torch.cat([frames, torch.ones(n_frames, 1, vh, vw)], axis=1)
+    layer2_head1 = torch.tensor(red_cmap(sims_1.squeeze().detach().cpu())).permute(0, 3, 1, 2)
+    layer2_head2 = torch.tensor(blue_cmap(sims_2.squeeze().detach().cpu())).permute(0, 3, 1, 2)
+
+    write_video_with_audio(
+        alpha_blend_layers([layer1, layer2_head1, layer2_head2]),
+        audio,
+        video_fps,
+        audio_fps,
+        output_filename)
+
+
+def plot_feature_video(image_feats,
+                       audio_feats,
+                       frames,
+                       audio,
+                       video_fps,
+                       audio_fps,
+                       video_filename,
+                       audio_filename):
+    with torch.no_grad():
+        image_feats_ = image_feats.cpu()
+        audio_feats_ = audio_feats.cpu()
+        [red_img_feats, red_audio_feats], _ = pca([
+            image_feats_,
+            audio_feats_,  # .tile(image_feats_.shape[0], 1, 1, 1)
+        ])
+        _, _, vh, vw = frames.shape
+        red_img_feats = F.interpolate(red_img_feats, size=(vh, vw), mode="bicubic")
+        red_audio_feats = red_audio_feats[0].unsqueeze(0)
+        red_audio_feats = F.interpolate(red_audio_feats, size=(50, red_img_feats.shape[0]), mode="bicubic")
+
+    write_video_with_audio(
+        (red_img_feats.permute(0, 2, 3, 1) * 255).clamp(0, 255).to(torch.uint8),
+        audio,
+        video_fps,
+        audio_fps,
+        video_filename)
+
+    red_audio_feats_expanded = red_audio_feats.tile(red_img_feats.shape[0], 1, 1, 1)
+    red_audio_feats_expanded = F.interpolate(red_audio_feats_expanded, scale_factor=6, mode="bicubic")
+    for i in range(red_img_feats.shape[0]):
+        center_index = i * 6
+        min_index = max(center_index - 2, 0)
+        max_index = min(center_index + 2, red_audio_feats_expanded.shape[-1])
+        red_audio_feats_expanded[i, :, :, min_index:max_index] = 1
+
+    write_video_with_audio(
+        (red_audio_feats_expanded.permute(0, 2, 3, 1) * 255).clamp(0, 255).to(torch.uint8),
+        audio,
+        video_fps,
+        audio_fps,
+        audio_filename)
+
+
+def display_video_in_notebook(path):
+    from IPython.display import HTML, display
+    mp4 = open(path, 'rb').read()
+    data_url = "data:video/mp4;base64," + b64encode(mp4).decode()
+    display(HTML("""
+  <video width=400 controls>
+        <source src="%s" type="video/mp4">
+  </video>
+  """ % data_url))

DenseAV/denseav/saved_models.py

@@ -0,0 +1,262 @@
+import os
+import re
+from os.path import join
+
+import torch
+
+
+
+def get_latest(name, checkpoint_dir, extra_args=None):
+    if extra_args is None:
+        extra_args = dict()
+    files = os.listdir(join(checkpoint_dir, name))
+    steps = torch.tensor([int(f.split("step=")[-1].split(".")[0]) for f in files])
+    selected = files[steps.argmax()]
+    return dict(
+        chkpt_name=os.path.join(name, selected),
+        extra_args=extra_args)
+
+
+DS_PARAM_REGEX = r'_forward_module\.(.+)'
+
+
+def convert_deepspeed_checkpoint(deepspeed_ckpt_path: str, pl_ckpt_path: str = None):
+    '''
+    Creates a PyTorch Lightning checkpoint from the DeepSpeed checkpoint directory, while patching
+    in parameters which are improperly loaded by the DeepSpeed conversion utility.
+    deepspeed_ckpt_path: Path to the DeepSpeed checkpoint folder.
+    pl_ckpt_path: Path to the reconstructed PyTorch Lightning checkpoint. If not specified, will be
+        placed in the same directory as the DeepSpeed checkpoint directory with the same name but
+        a .pt extension.
+    Returns: path to the converted checkpoint.
+    '''
+    from pytorch_lightning.utilities.deepspeed import convert_zero_checkpoint_to_fp32_state_dict
+
+
+    if not (deepspeed_ckpt_path.endswith('.ckpt') and os.path.isdir(deepspeed_ckpt_path)):
+        raise ValueError(
+            'args.ckpt_dir should point to the checkpoint directory'
+            ' output by DeepSpeed (e.g. "last.ckpt" or "epoch=4-step=39150.ckpt").'
+        )
+
+    # Convert state dict to PyTorch format
+    if not pl_ckpt_path:
+        pl_ckpt_path = f'{deepspeed_ckpt_path[:-4]}pt'  # .ckpt --> .pt
+
+    if not os.path.exists(pl_ckpt_path):
+        convert_zero_checkpoint_to_fp32_state_dict(deepspeed_ckpt_path, pl_ckpt_path)
+
+    # Patch in missing parameters that failed to be converted by DeepSpeed utility
+    pl_ckpt = _merge_deepspeed_weights(deepspeed_ckpt_path, pl_ckpt_path)
+    torch.save(pl_ckpt, pl_ckpt_path)
+
+    return pl_ckpt_path
+
+
+def get_optim_files(checkpoint_dir):
+    files = sorted([f for f in os.listdir(checkpoint_dir) if "optim" in f])
+    return [join(checkpoint_dir, f) for f in files]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    f = [f for f in os.listdir(checkpoint_dir) if "model_states" in f][0]
+    return join(checkpoint_dir, f)
+
+
+def _merge_deepspeed_weights(deepspeed_ckpt_path: str, fp32_ckpt_path: str):
+    '''
+    Merges tensors with keys in the DeepSpeed checkpoint but not in the fp32_checkpoint
+    into the fp32 state dict.
+    deepspeed_ckpt_path: Path to the DeepSpeed checkpoint folder.
+    fp32_ckpt_path: Path to the reconstructed
+    '''
+    from pytorch_lightning.utilities.deepspeed import ds_checkpoint_dir
+
+
+    # This first part is based on pytorch_lightning.utilities.deepspeed.convert_zero_checkpoint_to_fp32_state_dict
+    checkpoint_dir = ds_checkpoint_dir(deepspeed_ckpt_path)
+    optim_files = get_optim_files(checkpoint_dir)
+    optim_state = torch.load(optim_files[0], map_location='cpu')
+    zero_stage = optim_state["optimizer_state_dict"]["zero_stage"]
+    deepspeed_model_file = get_model_state_file(checkpoint_dir, zero_stage)
+
+    # Start adding all parameters from DeepSpeed ckpt to generated PyTorch Lightning ckpt
+    ds_ckpt = torch.load(deepspeed_model_file, map_location='cpu')
+    ds_sd = ds_ckpt['module']
+
+    fp32_ckpt = torch.load(fp32_ckpt_path, map_location='cpu')
+    fp32_sd = fp32_ckpt['state_dict']
+
+    for k, v in ds_sd.items():
+        try:
+            match = re.match(DS_PARAM_REGEX, k)
+            param_name = match.group(1)
+        except:
+            print(f'Failed to extract parameter from DeepSpeed key {k}')
+            continue
+
+        v = v.to(torch.float32)
+        if param_name not in fp32_sd:
+            print(f'Adding parameter {param_name} from DeepSpeed state_dict to fp32_sd')
+            fp32_sd[param_name] = v
+        else:
+            assert torch.allclose(v, fp32_sd[param_name].to(torch.float32), atol=1e-2)
+
+    return fp32_ckpt
+
+
+def get_version_and_step(f, i):
+    step = f.split("step=")[-1].split(".")[0]
+    if "-v" in step:
+        [step, version] = step.split("-v")
+    else:
+        step, version = step, 0
+
+    return int(version), int(step), i
+
+
+def get_latest_ds(name, extra_args=None):
+    if extra_args is None:
+        extra_args = dict()
+    files = os.listdir(f"../checkpoints/{name}")
+    latest = sorted([get_version_and_step(f, i) for i, f in enumerate(files)], reverse=True)[0]
+    selected = files[latest[-1]]
+    # print(f"Selecting file: {selected}")
+    ds_chkpt = join(name, selected)
+    reg_chkpt = join(name + "_fp32", selected)
+    reg_chkpt_path = join("../checkpoints", reg_chkpt)
+    if not os.path.exists(reg_chkpt_path):
+        os.makedirs(os.path.dirname(reg_chkpt_path), exist_ok=True)
+        print(f"Checkpoint {reg_chkpt} does not exist, converting from deepspeed")
+        convert_deepspeed_checkpoint(join("../checkpoints", ds_chkpt), reg_chkpt_path)
+    return dict(
+        chkpt_name=reg_chkpt,
+        extra_args=extra_args)
+
+
+def get_all_models_in_dir(name, checkpoint_dir, extra_args=None):
+    ret = {}
+    for model_dir in os.listdir(join(checkpoint_dir, name)):
+        full_name = f"{name}/{model_dir}/train"
+        # print(f'"{full_name}",')
+        ret[full_name] = get_latest(full_name, checkpoint_dir, extra_args)
+    return ret
+
+
+def saved_model_dict(checkpoint_dir):
+    model_info = {
+
+        **get_all_models_in_dir(
+            "9-5-23-mixed",
+            checkpoint_dir,
+            extra_args=dict(
+                mixup_weight=0.0,
+                sim_use_cls=False,
+                audio_pool_width=1,
+                memory_buffer_size=0,
+                loss_leak=0.0)
+        ),
+
+        **get_all_models_in_dir(
+            "1-23-24-rebuttal-heads",
+            checkpoint_dir,
+            extra_args=dict(
+                loss_leak=0.0)
+        ),
+
+        **get_all_models_in_dir(
+            "11-8-23",
+            checkpoint_dir,
+            extra_args=dict(loss_leak=0.0)),
+
+        **get_all_models_in_dir(
+            "10-30-23-3",
+            checkpoint_dir,
+            extra_args=dict(loss_leak=0.0)),
+
+        "davenet": dict(
+            chkpt_name=None,
+            extra_args=dict(
+                audio_blur=1,
+                image_model_type="davenet",
+                image_aligner_type=None,
+                audio_model_type="davenet",
+                audio_aligner_type=None,
+                audio_input="davenet_spec",
+                use_cached_embs=False,
+                dropout=False,
+                sim_agg_heads=1,
+                nonneg_sim=False,
+                audio_lora=False,
+                image_lora=False,
+                norm_vectors=False,
+            ),
+            data_args=dict(
+                use_cached_embs=False,
+                use_davenet_spec=True,
+                override_target_length=20,
+                audio_model_type="davenet",
+            ),
+        ),
+
+        "cavmae": dict(
+            chkpt_name=None,
+            extra_args=dict(
+                audio_blur=1,
+                image_model_type="cavmae",
+                image_aligner_type=None,
+                audio_model_type="cavmae",
+                audio_aligner_type=None,
+                audio_input="spec",
+                use_cached_embs=False,
+                sim_agg_heads=1,
+                dropout=False,
+                nonneg_sim=False,
+                audio_lora=False,
+                image_lora=False,
+                norm_vectors=False,
+                learn_audio_cls=False,
+                sim_agg_type="cavmae",
+            ),
+            data_args=dict(
+                use_cached_embs=False,
+                use_davenet_spec=True,
+                audio_model_type="cavmae",
+                override_target_length=10,
+            ),
+        ),
+
+        "imagebind": dict(
+            chkpt_name=None,
+            extra_args=dict(
+                audio_blur=1,
+                image_model_type="imagebind",
+                image_aligner_type=None,
+                audio_model_type="imagebind",
+                audio_aligner_type=None,
+                audio_input="spec",
+                use_cached_embs=False,
+                sim_agg_heads=1,
+                dropout=False,
+                nonneg_sim=False,
+                audio_lora=False,
+                image_lora=False,
+                norm_vectors=False,
+                learn_audio_cls=False,
+                sim_agg_type="imagebind",
+            ),
+            data_args=dict(
+                use_cached_embs=False,
+                use_davenet_spec=True,
+                audio_model_type="imagebind",
+                override_target_length=10,
+            ),
+        ),
+
+    }
+
+    model_info["denseav_language"] = model_info["10-30-23-3/places_base/train"]
+    model_info["denseav_sound"] = model_info["11-8-23/hubert_1h_asf_cls_full_image_train_small_lr/train"]
+    model_info["denseav_2head"] = model_info["1-23-24-rebuttal-heads/mixed-2h/train"]
+
+    return model_info

DenseAV/denseav/shared.py

@@ -0,0 +1,739 @@
+import random
+from collections import defaultdict, deque
+from typing import Any
+
+import math
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+import torchaudio
+import torchvision.transforms as T
+from PIL import Image
+from torch.utils.data import Dataset
+from torchaudio.functional import resample
+
+
+class UnNormalize(object):
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, image):
+        image2 = torch.clone(image)
+        for t, m, s in zip(image2, self.mean, self.std):
+            t.mul_(s).add_(m)
+        return image2
+
+
+class SliceDataset(Dataset):
+
+    def __init__(self, ds, start, end):
+        self.ds = ds
+        self.start = start
+        self.end = end
+
+    def __len__(self):
+        return self.end - self.start
+
+    def __getitem__(self, item):
+        return self.ds[item + self.start]
+
+
+class SubsetDataset(Dataset):
+
+    def __init__(self, ds, subset):
+        self.ds = ds
+        self.subset = subset
+
+    def __len__(self):
+        return len(self.subset)
+
+    def __getitem__(self, item):
+        return self.ds[self.subset[item]]
+
+
+norm = T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+unnorm = UnNormalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+
+
+def crop_to_divisor(x, patch_size):
+    if len(x.shape) == 3:
+        C, H, W = x.shape
+        return x[:, :(patch_size * (H // patch_size)), :(patch_size * (W // patch_size))]
+    elif len(x.shape) == 4:
+        B, C, H, W = x.shape
+        return x[:, :, :(patch_size * (H // patch_size)), :(patch_size * (W // patch_size))]
+    else:
+        raise ValueError("x should have 3 or 4 dimensions")
+
+
+def _remove_axes(ax):
+    ax.xaxis.set_major_formatter(plt.NullFormatter())
+    ax.yaxis.set_major_formatter(plt.NullFormatter())
+    ax.set_xticks([])
+    ax.set_yticks([])
+
+
+def remove_axes(axes):
+    if len(axes.shape) == 2:
+        for ax1 in axes:
+            for ax in ax1:
+                _remove_axes(ax)
+    else:
+        for ax in axes:
+            _remove_axes(ax)
+
+
+def get_image_featurizer(name, token_type="key", **kwargs):
+    name = name.lower()
+
+    if name == "vit":
+        from denseav.featurizers.DINO import DINOFeaturizer
+        patch_size = 16
+        model = DINOFeaturizer("vit_small_patch16_224", patch_size, token_type)
+        dim = 384
+    elif name == "dino16":
+        from denseav.featurizers.DINO import DINOFeaturizer
+        patch_size = 16
+        model = DINOFeaturizer("dino_vits16", patch_size, token_type)
+        dim = 384
+    elif name == "dino8":
+        from denseav.featurizers.DINO import DINOFeaturizer
+        patch_size = 8
+        model = DINOFeaturizer("dino_vits8", patch_size, token_type)
+        dim = 384
+    elif name == "clip":
+        from denseav.featurizers.CLIP import CLIPFeaturizer
+        patch_size = 16
+        model = CLIPFeaturizer()
+        dim = 512
+    elif name == "cavmae":
+        from denseav.featurizers.CAVMAE import CAVMAEImageFeaturizer
+        model = CAVMAEImageFeaturizer(kwargs["output_root"], model=kwargs.get("model"))
+        dim = 768
+        patch_size = 16
+    elif name == "fnac":
+        from denseav.featurizers.FNACAVL import FNACImageFeaturizer
+        model = FNACImageFeaturizer(kwargs["output_root"], model=kwargs.get("model"))
+        dim = 512
+        patch_size = 16
+    elif name == "imagebind":
+        from denseav.featurizers.ImageBind import ImageBindImageFeaturizer
+        model = ImageBindImageFeaturizer(kwargs["output_root"], model=kwargs.get("model"))
+        dim = 1024
+        patch_size = 16
+    elif name == "resnet50":
+        from torchvision import models
+        model = models.resnet50(pretrained=True)
+        model = torch.nn.Sequential(*list(model.children())[:-2])
+        patch_size = 1
+        dim = 2048
+    elif name == "davenet":
+        from fdenseav.eaturizers.DAVENet import DavenetImageFeaturizer
+        model = DavenetImageFeaturizer()
+        patch_size = 1
+        dim = 1024
+    elif name == "dinov2":
+        from denseav.featurizers.DINOv2 import DINOv2Featurizer
+        model = DINOv2Featurizer()
+        patch_size = 14
+        dim = 768
+    else:
+        raise ValueError("unknown model: {}".format(name))
+    return model, patch_size, dim
+
+
+def get_audio_featurizer(name, **kwargs):
+    if name == "davenet":
+        from denseav.featurizers.DAVENet import DavenetAudioFeaturizer
+        model = DavenetAudioFeaturizer()
+        dim = 1024
+    elif name == "dino8":
+        model, _, dim = get_image_featurizer("dino8")
+    elif name == "hubert":
+        from denseav.featurizers.Hubert import Hubert
+        model = Hubert()
+        dim = 1024
+    elif name == "cavmae":
+        from denseav.featurizers.CAVMAE import CAVMAEAudioFeaturizer
+        model = CAVMAEAudioFeaturizer(kwargs["output_root"], model=kwargs.get("model"))
+        dim = 768
+    elif name == "imagebind":
+        from denseav.featurizers.ImageBind import ImageBindAudioFeaturizer
+        model = ImageBindAudioFeaturizer(kwargs["output_root"], model=kwargs.get("model"))
+        dim = 1024
+    elif name == "audiomae":
+        from denseav.featurizers.AudioMAE import AudioMAE
+        model = AudioMAE(kwargs["output_root"], False)
+        dim = 768
+    elif name == "audiomae-finetuned":
+        from denseav.featurizers.AudioMAE import AudioMAE
+        model = AudioMAE(kwargs["output_root"], True)
+        dim = 768
+    else:
+        raise ValueError("Unknown audio model type")
+
+    return model, dim
+
+
+def load_img(image_path, transform):
+    return transform(Image.open(image_path)).unsqueeze(0)
+
+
+def pytorch_to_pil(tensor):
+    return Image.fromarray((unnorm(tensor).permute(0, 2, 3, 1).cpu() * 255)
+                           .clamp(0, 255).to(torch.uint8).detach().numpy()[0])
+
+
+def _get_random_window(waveform, mask, min_size, max_size):
+    effective_size = mask.sum().to(torch.int64)
+    if effective_size <= min_size:
+        return waveform, mask
+    else:
+        window_size = min(torch.randint(low=min_size, high=min(effective_size, max_size), size=()), waveform.shape[0])
+        if window_size == waveform.shape[0]:
+            window_start = 0
+        else:
+            window_start = torch.randint(low=0, high=effective_size - window_size, size=())
+
+        new_waveform = torch.zeros_like(waveform)
+        new_mask = torch.zeros_like(mask)
+        new_waveform[window_start:window_start + window_size] = waveform[window_start:window_start + window_size]
+        new_mask[window_start:window_start + window_size] = mask[window_start:window_start + window_size]
+        return new_waveform, new_mask
+
+
+def _splice_clips(clip1, clip2, loc, easing_size):
+    assert loc >= 0 and loc < len(clip1), "Invalid location"
+    assert easing_size > 0 and easing_size <= len(clip2), "Invalid easing size"
+
+    try:
+        assert loc + clip2.shape[0] < clip1.shape[0]
+    except Exception as e:
+        print(loc, clip2.shape[0], clip1.shape[0])
+        raise e
+
+    # Split clip1 into three parts: before splice, easing region, after splice
+    before_splice = clip1[:loc]
+    after_splice = clip1[loc + clip2.shape[0]:]
+
+    # Compute the fading weights for the easing region
+    # fade_in_weights = torch.cos(torch.linspace(1, 0, easing_size, device=clip1.device))
+    fade_in_weights = 0.5 * (1 + torch.cos(math.pi * torch.linspace(0, 1, easing_size)))
+    fade_out_weights = 1 - fade_in_weights
+
+    clip1_ease = torch.cat([
+        fade_in_weights,
+        torch.zeros(clip2.shape[0] - easing_size * 2),
+        fade_out_weights,
+    ])
+
+    mask = torch.cat([torch.ones(loc), clip1_ease, torch.ones(clip1.shape[0] - (loc + clip2.shape[0]))])
+
+    # Apply fading weights to clip1 and clip2 within the easing region
+    splice = clip1_ease * clip1[loc:loc + clip2.shape[0]] + (1 - clip1_ease) * clip2
+
+    # Concatenate all parts back together
+    spliced_clip = torch.cat((before_splice, splice, after_splice))
+
+    return spliced_clip, mask
+
+
+def _generate_random_subset(waveform, low, high):
+    length = len(waveform)
+
+    # If waveform is smaller than low or has zero length, return unmodified
+    if length < low or length == 0:
+        return waveform
+
+    # Generate random start index within valid range
+    start = random.randint(0, length - low)
+
+    # Generate random subset size within valid range
+    subset_size = random.randint(low, min(high, length - start))
+
+    # Extract the random subset from the waveform
+    subset = waveform[start: start + subset_size]
+
+    return subset
+
+
+def level_audio(waveform):
+    waveform -= waveform.mean()
+    waveform /= waveform.abs.max().valus.clamp_min(.0001)
+    return waveform
+
+
+def prep_waveform(waveform,
+                  obs_sr,
+                  target_length,
+                  spec_mel_bins,
+                  spec_mean,
+                  spec_std,
+                  sample_rate,
+                  return_spec,
+                  random_clip,
+                  extra_audio_masking,
+                  neg_waveform,
+                  neg_obs_sr,
+                  audio_level,
+                  audio_aug,
+                  ):
+    if obs_sr != sample_rate:
+        waveform = resample(waveform, obs_sr, sample_rate)
+        if audio_level:
+            waveform = level_audio(waveform)
+
+    if neg_obs_sr is not None and neg_obs_sr != sample_rate:
+        neg_waveform = resample(neg_waveform, neg_obs_sr, sample_rate)
+        if audio_level:
+            neg_waveform = level_audio(neg_waveform)
+
+    if neg_obs_sr is not None:  # and random.random() > .5:
+        neg_waveform_clip = _generate_random_subset(neg_waveform, sample_rate, sample_rate * 4)
+        if waveform.shape[0] - neg_waveform_clip.shape[0] > 0:
+            start = random.randint(0, waveform.shape[0] - neg_waveform_clip.shape[0] - 1)
+            easing = max(int(neg_waveform_clip.shape[0] * 1 / 4), sample_rate // 2)
+            easing = min(int(neg_waveform_clip.shape[0] * 1 / 2), easing)
+            waveform, pos_mask = _splice_clips(waveform, neg_waveform_clip, start, easing_size=easing)
+        else:
+            waveform, pos_mask = waveform, torch.ones_like(waveform)
+    else:
+        waveform, pos_mask = waveform, torch.ones_like(waveform)
+
+    mask = torch.ones_like(waveform)
+    original_length = waveform.shape[0]
+
+    if target_length == 10:
+        target_samples = 164200  # Result is 1024 after spec
+    else:
+        target_samples = int(target_length * sample_rate)
+
+    padding = target_samples - original_length
+
+    if padding > 0:
+        p = torch.nn.ZeroPad2d((0, padding))
+        waveform = p(waveform)
+        mask = p(mask)
+        pos_mask = p(pos_mask)
+    else:
+        if random_clip:
+            start = torch.randint(0, waveform.shape[0] - target_samples, size=())
+        else:
+            start = 0
+        end = start + target_samples
+        waveform = waveform[start:end]
+        mask = mask[start:end]
+        pos_mask = pos_mask[start:end]
+
+    audio_length = min(original_length, target_samples)
+    total_length = target_samples
+
+    if extra_audio_masking:
+        min_size = sample_rate // 2
+        max_size = total_length
+        if original_length > min_size and random.random() > .5:
+            waveform, mask = _get_random_window(waveform, mask, min_size, max_size)
+
+    if audio_aug:
+        import torchaudio_augmentations as AA
+        from torchvision.transforms import RandomApply, Compose
+
+        transform = Compose([
+            RandomApply([AA.PolarityInversion()], p=0.5),
+            RandomApply([AA.Noise(min_snr=0.001, max_snr=0.005)], p=0.2),
+            RandomApply([AA.Gain()], p=0.2),
+            RandomApply([AA.HighLowPass(sample_rate=sample_rate)], p=0.2),
+            RandomApply([AA.PitchShift(n_samples=waveform.shape[-1], sample_rate=sample_rate)], p=0.2),
+            RandomApply([AA.Reverb(sample_rate=sample_rate)], p=0.2)
+        ])
+        waveform = transform(waveform.unsqueeze(0)).squeeze(0)
+
+    if return_spec:
+        spectrogram = torchaudio.compliance.kaldi.fbank(
+            waveform.unsqueeze(0) - waveform.mean(),
+            htk_compat=True,
+            sample_frequency=sample_rate,
+            use_energy=False,
+            window_type='hanning',
+            num_mel_bins=spec_mel_bins,
+            dither=0.0,
+            frame_shift=10)
+
+        spectrogram = ((spectrogram - spec_mean) / spec_std).unsqueeze(0)
+    else:
+        spectrogram = None
+
+    if mask.mean() < .04:
+        print(f"Bad entry: {mask.mean()}")
+
+    return waveform, spectrogram, audio_length, total_length, original_length, mask, pos_mask
+
+
+class ToTargetTensor(object):
+    def __call__(self, target):
+        return torch.as_tensor(np.array(target), dtype=torch.int64).unsqueeze(0)
+
+
+def show_heatmap(ax,
+                 image,
+                 heatmap,
+                 cmap="bwr",
+                 color=False,
+                 center=False,
+                 show_negative=False,
+                 cax=None,
+                 vmax=None,
+                 vmin=None):
+    frame = []
+
+    if color:
+        frame.append(ax.imshow(image))
+    else:
+        bw = np.dot(np.array(image)[..., :3] / 255, [0.2989, 0.5870, 0.1140])
+        bw = np.ones_like(image) * np.expand_dims(bw, -1)
+        frame.append(ax.imshow(bw))
+
+    if center:
+        heatmap -= heatmap.mean()
+
+    if not show_negative:
+        heatmap = heatmap.clamp_min(0)
+
+    heatmap = F.interpolate(heatmap.unsqueeze(0).unsqueeze(0), (image.shape[0], image.shape[1])) \
+        .squeeze(0).squeeze(0)
+
+    if vmax is None:
+        vmax = np.abs(heatmap).max()
+    if vmin is None:
+        vmin = -vmax
+
+    hm = ax.imshow(heatmap, alpha=.5, cmap=cmap, vmax=vmax, vmin=vmin)
+    if cax is not None:
+        plt.colorbar(hm, cax=cax, orientation='vertical')
+
+    frame.extend([hm])
+    return frame
+
+
+class TorchPCA(object):
+
+    def __init__(self, n_components):
+        self.n_components = n_components
+
+    def fit(self, X):
+        self.mean_ = X.mean(dim=0)
+        unbiased = X - self.mean_.unsqueeze(0)
+        U, S, V = torch.pca_lowrank(unbiased, q=self.n_components, center=False, niter=4)
+        self.components_ = V.T
+        self.singular_values_ = S
+        return self
+
+    def transform(self, X):
+        t0 = X - self.mean_.unsqueeze(0)
+        projected = t0 @ self.components_.T
+        return projected
+
+
+def pca(image_feats_list, dim=3, fit_pca=None):
+    device = image_feats_list[0].device
+
+    def flatten(tensor, target_size=None):
+        if target_size is not None and fit_pca is None:
+            F.interpolate(tensor, (target_size, target_size), mode="bilinear")
+        B, C, H, W = tensor.shape
+        return feats.permute(1, 0, 2, 3).reshape(C, B * H * W).permute(1, 0).detach().cpu()
+
+    if len(image_feats_list) > 1 and fit_pca is None:
+        target_size = image_feats_list[0].shape[2]
+    else:
+        target_size = None
+
+    flattened_feats = []
+    for feats in image_feats_list:
+        flattened_feats.append(flatten(feats, target_size))
+    x = torch.cat(flattened_feats, dim=0)
+
+    if fit_pca is None:
+        # fit_pca = PCA(n_components=dim, svd_solver='arpack').fit(np.nan_to_num(x.detach().numpy()))
+        fit_pca = TorchPCA(n_components=dim).fit(x)
+
+    reduced_feats = []
+    for feats in image_feats_list:
+        # x_red = torch.from_numpy(fit_pca.transform(flatten(feats)))
+        x_red = fit_pca.transform(flatten(feats))
+        x_red -= x_red.min(dim=0, keepdim=True).values
+        x_red /= x_red.max(dim=0, keepdim=True).values
+        B, C, H, W = feats.shape
+        reduced_feats.append(x_red.reshape(B, H, W, dim).permute(0, 3, 1, 2).to(device))
+
+    return reduced_feats, fit_pca
+
+
+def merge_col(fig, axes, col):
+    gs = axes[0, col].get_gridspec()
+    for ax in axes[:, col]:
+        ax.remove()
+    return fig.add_subplot(gs[:, col])
+
+
+def visualize_av_features(
+        audio,
+        video,
+        feat_a,
+        feat_v,
+        att_a,
+        n_frames,
+        norm_before_pca=True,
+        axes=None,
+        fig=None,
+        modify_fig=True,
+        video_time=0,
+        fit_pca=None
+):
+    assert (len(audio.shape) == 3)  # C, F, T
+    assert (len(video.shape) == 4)  # T, C, H, W
+    assert (len(feat_a.shape) == 2)  # C, T
+    assert (len(feat_v.shape) == 4)  # T, C, H, W
+    assert (len(att_a.shape) == 2)  # F, T
+
+    ac, af, at = audio.shape
+    fac, fat = feat_a.shape
+
+    if modify_fig:
+        if axes is None:
+            fig, axes = plt.subplots(3, 3, figsize=(5 * 3, 5))
+            fig.tight_layout()
+
+        bigax1 = merge_col(fig, axes, 0)
+        bigax2 = merge_col(fig, axes, 1)
+        _remove_axes(bigax1)
+        _remove_axes(bigax2)
+        remove_axes(axes[:, 2])
+    else:
+        bigax1 = fig.axes[-2]
+        bigax2 = fig.axes[-1]
+
+    frame_v = unnorm(video).permute(0, 2, 3, 1).detach().cpu()
+    frame_v -= frame_v.min()
+    frame_v /= frame_v.max()
+
+    frame_a = audio.detach().cpu()
+    frame_a -= frame_a.min()
+    frame_a /= frame_a.max()
+
+    if norm_before_pca:
+        [red_feat_v], fit_pca = pca([F.normalize(feat_v, dim=1)], fit_pca=fit_pca)
+        [red_feat_a], _ = pca([F.normalize(feat_a.unsqueeze(0).unsqueeze(-1), dim=1)], fit_pca=fit_pca)
+    else:
+        [red_feat_v], fit_pca = pca([feat_v], fit_pca=fit_pca)
+        [red_feat_a], _ = pca([feat_a.unsqueeze(0).unsqueeze(-1)], fit_pca=fit_pca)
+
+    red_feat_v = red_feat_v.permute(0, 2, 3, 1).detach().cpu()
+    red_feat_a = red_feat_a.permute(0, 2, 3, 1)[0].detach().cpu()
+
+    if red_feat_a.shape[0] == 1:
+        new_height = int((frame_a.shape[0] / frame_a.shape[1]) * red_feat_a.shape[1])
+        red_feat_a = torch.broadcast_to(
+            red_feat_a, (new_height, red_feat_a.shape[1], red_feat_a.shape[2]))
+        plt_att_a = torch.broadcast_to(att_a, (new_height, att_a.shape[1]))
+    else:
+        plt_att_a = att_a
+
+    frac_signal = n_frames / fat
+    n_at = int(at * frac_signal)
+
+    return [bigax1.imshow(frame_v[video_time]),
+            bigax2.imshow(red_feat_v[video_time]),
+            axes[0, 2].imshow(frame_a[:, :n_at]),
+            axes[0, 2].set_title("Spectrogram"),
+            axes[1, 2].imshow(red_feat_a[:, :n_frames]),
+            axes[1, 2].set_title("Audio Features"),
+            axes[2, 2].imshow(plt_att_a[:, :n_frames], vmin=0),
+            axes[2, 2].set_title("Audio Attention")], fig, fit_pca
+
+
+def create_label_tensor(labels, starts, ends, max_time, n_steps):
+    assert isinstance(starts, torch.Tensor)
+    assert isinstance(ends, torch.Tensor)
+
+    ends[ends < 0] = max_time
+    fps = n_steps / max_time
+    times = (torch.arange(0, n_steps, device=labels.device, dtype=torch.float32) + .5) / fps
+    after_start = starts.unsqueeze(1) <= times.unsqueeze(0)
+    before_end = ends.unsqueeze(1) >= times.unsqueeze(0)
+    # Find when you are inside of a word
+    in_word = (after_start * before_end)
+    # Find which word you are inside of
+    word_to_use = in_word.to(torch.float32).argmax(0)
+    # Get the label for that word, or mask out the label if in no word
+    final_labels = labels[word_to_use] * in_word.any(0).reshape(-1, 1, 1)
+    return final_labels
+
+
+def generate_subset(n, batch, seed=0):
+    np.random.seed(seed)
+    return np.random.permutation(n)[:batch]
+
+
+def channel_blur(t, window=5, std_dev=1):
+    tb, tc, th, tw = t.shape
+    x = torch.linspace(-2, 2, window, device=t.device, dtype=torch.float32)
+    k = torch.exp((-x ** 2 / (2 * std_dev ** 2)))
+    k = k / k.sum()
+    pad = window // 2
+    t_pad = F.pad(t, [0, 0, 0, 0, pad, pad], mode="replicate")
+    tpb, tpc, tph, tpw = t_pad.shape
+    flattened_t = t_pad.permute(0, 2, 3, 1).reshape(tpb * tph * tpw, 1, -1)
+    return F.conv1d(flattened_t, k.reshape(1, 1, window)).reshape(tpb, tph, tpw, tc).permute(0, 3, 1, 2)
+
+
+def time_blur(t, window=5, std_dev=1):
+    tb, tc, tt = t.shape
+    with torch.no_grad():
+        x = torch.linspace(-2, 2, window, device=t.device, dtype=torch.float32)
+        k = torch.exp((-x ** 2 / (2 * std_dev ** 2)))
+        k = k / k.sum()
+        k = k.reshape(1, 1, window).detach()
+    pad = window // 2
+    t_pad = F.pad(t, [pad, pad], mode="replicate")
+    return F.conv1d(t_pad.reshape(tb * tc, 1, -1), k).reshape(tb, tc, tt)
+
+
+def create_model_from_cfg(clazz, cfg, extra_args):
+    import inspect
+    expected_args = inspect.getfullargspec(clazz.__init__).args[1:]
+    new_args = {k: v for k, v in {**cfg, **extra_args}.items() if k in expected_args}
+    return clazz(**new_args)
+
+
+def load_trained_model(chkpt_dir, extra_args, strict=True):
+    from train_av_alignment import LitAVAligner
+    model = LitAVAligner.load_from_checkpoint(chkpt_dir, **extra_args, strict=strict).cuda()
+    return model
+
+
+def flatten(l):
+    return [item for sublist in l for item in sublist]
+
+
+def flatten_preds(preds):
+    results = {}
+    for k in preds[0].keys():
+        if k == "caption_labels":
+            continue
+        if isinstance(preds[0][k], torch.Tensor):
+            results[k] = torch.cat([p[k] for p in preds], dim=0)
+    if "caption" in preds[0]:
+        results["caption"] = flatten([p["caption"] for p in preds])
+
+    if "metadata" in preds[0]:
+        results["frame_files"] = flatten([list(p["metadata"]["frame_files"][0]) for p in preds])
+        results["audio_file"] = flatten([list(p["metadata"]["audio_file"]) for p in preds])
+        results["id"] = flatten([list(p["metadata"]["id"]) for p in preds])
+        results["index"] = torch.tensor(flatten([list(p["metadata"]["index"]) for p in preds]))
+
+    return results
+
+
+def batch(iterable, n=1):
+    l = len(iterable)
+    for ndx in range(0, l, n):
+        yield iterable[ndx:min(ndx + n, l)]
+
+
+class GatherLayer(torch.autograd.Function):
+    """Gather tensors from all process, supporting backward propagation."""
+
+    @staticmethod
+    def jvp(ctx: Any, *grad_inputs: Any) -> Any:
+        pass
+
+    @staticmethod
+    def forward(ctx, inputs):
+        ctx.save_for_backward(inputs)
+        output = [torch.zeros_like(inputs) for _ in range(dist.get_world_size())]
+        dist.all_gather(output, inputs)
+        return tuple(output)
+
+    @staticmethod
+    def backward(ctx, *grads):
+        (inputs,) = ctx.saved_tensors
+        grad_out = torch.zeros_like(inputs)
+        grad_out[:] = grads[dist.get_rank()]
+        return grad_out
+
+
+class RollingAvg:
+
+    def __init__(self, length, nonzero=False):
+        self.length = length
+        self.nonzero = nonzero
+        self.metrics = defaultdict(lambda: deque(maxlen=self.length))
+
+    def add(self, name, metric):
+        if self.nonzero and metric == 0:
+            return
+        if isinstance(metric, torch.Tensor):
+            metric = metric.detach()
+
+        self.metrics[name].append(metric)
+
+    def get(self, name):
+        with torch.no_grad():
+            return torch.tensor(list(self.metrics[name])).mean()
+
+    def get_all(self):
+        return {k: self.get(k) for k in self.metrics.keys()}
+
+    def add_all(self, values):
+        for k, v in values.items():
+            self.add(k, v)
+
+    def logall(self, log_func):
+        for k in self.metrics.keys():
+            log_func(k, self.get(k))
+
+
+def gaussian_kernel(k, sigma):
+    kernel = torch.tensor([math.exp(-0.5 * (x - (k // 2)) ** 2 / sigma ** 2) for x in range(k)], dtype=torch.float32)
+    kernel /= kernel.sum()  # Normalize the kernel
+    return kernel
+
+
+def blur_dim(t, window=5, std_dev=1, dim=-1):
+    shape = t.shape
+    n_dims = len(shape)
+
+    # Create the Gaussian kernel
+    with torch.no_grad():
+        x = torch.linspace(-2, 2, window, device=t.device, dtype=torch.float32)
+        k = torch.exp(-x ** 2 / (2 * std_dev ** 2))
+        k = k / k.sum()
+        k = k.view(1, 1, window).detach()
+
+    # Calculate padding
+    pad = window // 2
+
+    # Move the target dimension to the end
+    permute_order = list(range(n_dims))
+    permute_order.append(permute_order.pop(dim))
+    t_permuted = t.permute(permute_order)
+
+    # Flatten all dimensions except the last one
+    new_shape = (-1, t_permuted.size(-1))
+    t_flattened = t_permuted.reshape(new_shape)
+
+    # Pad the tensor
+    t_padded = F.pad(t_flattened.unsqueeze(1), (pad, pad), mode="replicate")
+
+    # Apply convolution
+    blurred = F.conv1d(t_padded, k)
+
+    # Reshape back to original
+    blurred = blurred.squeeze(1).reshape(*t_permuted.shape)
+    blurred = blurred.permute([permute_order.index(i) for i in range(n_dims)])
+
+    return blurred

DenseAV/denseav/train.py

@@ -0,0 +1,1213 @@
+import os
+from collections import deque
+from itertools import combinations
+from os.path import join
+
+import hydra
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from omegaconf import DictConfig, OmegaConf
+from peft import get_peft_model, LoraConfig
+from pytorch_lightning import Trainer
+from pytorch_lightning import seed_everything
+from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
+from pytorch_lightning.loggers import TensorBoardLogger
+from pytorch_lightning.utilities import grad_norm
+from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, SequentialLR, LambdaLR
+from torchmetrics.functional.classification import binary_average_precision
+
+from huggingface_hub import PyTorchModelHubMixin
+
+from denseav.aggregators import get_aggregator
+from denseav.aligners import get_aligner, ProgressiveGrowing
+from denseav.constants import *
+from denseav.data.AVDatasets import AVDataModule
+from denseav.shared import flatten_preds, GatherLayer, \
+    get_image_featurizer, get_audio_featurizer, RollingAvg, create_model_from_cfg
+
+torch.multiprocessing.set_sharing_strategy('file_system')
+
+
+def _imposter_indices_helper(true_indices: torch.Tensor, samples: torch.Tensor):
+    mask = (true_indices == samples).to(torch.int64)
+    n = mask.shape[0]
+
+    if not mask.any():
+        return samples
+    else:
+        new_samples = torch.randint(0, n, size=(n,), device=true_indices.device)
+        comb_samples = mask * new_samples + (1 - mask) * samples
+        return _imposter_indices_helper(true_indices, comb_samples)
+
+
+def imposter_indices(n, device):
+    return _imposter_indices_helper(
+        torch.arange(0, n, device=device),
+        torch.randint(0, n, size=(n,), device=device))
+
+
+def get_sim_per_row(image_outputs, audio_outputs, n_frames, sim_type):
+    max_t = audio_outputs.shape[-1]
+    oh = F.one_hot(n_frames - 1, num_classes=max_t)
+    audio_mask = 1 - torch.cumsum(oh, dim=1)
+    audio_mask = F.pad(audio_mask, [1, 0], value=1)[:, :max_t].to(audio_outputs.dtype)
+
+    full_sim = torch.einsum("bct,bchw->bthw", audio_outputs, image_outputs)
+    expanded_am = audio_mask.unsqueeze(-1).unsqueeze(-1)
+
+    if sim_type.endswith("mi"):
+        offset = 10 * (full_sim.max() - full_sim.min())
+        full_sim = (full_sim - ((1 - expanded_am) * offset)).max(1, keepdim=True).values
+
+    if sim_type.startswith("mi"):
+        full_sim = full_sim.max(-1, keepdim=True).values.max(-2, keepdim=True).values
+
+    if sim_type.endswith("sa"):
+        full_sim = (full_sim * (expanded_am / expanded_am.sum(1, keepdim=True).clamp_min(1))).sum(1, keepdim=True)
+
+    return full_sim.mean(dim=[1, 2, 3])
+
+
+def sampled_margin_rank_loss(image_outputs, audio_outputs, n_frames, sim_type, margin=1.):
+    """
+    Computes the triplet margin ranking loss for each anchor image/caption pair
+    The impostor image/caption is randomly sampled from the minibatch
+    """
+    assert (image_outputs.dim() == 4)
+    assert (audio_outputs.dim() == 3)
+    n = image_outputs.size(0)
+    imp_ind_i = imposter_indices(n, image_outputs.device)
+    imp_ind_a = imposter_indices(n, image_outputs.device)
+    true_sim = get_sim_per_row(image_outputs, audio_outputs, n_frames, sim_type)
+    imp_sim_i = get_sim_per_row(image_outputs[imp_ind_i], audio_outputs, n_frames, sim_type)
+    imp_sim_a = get_sim_per_row(image_outputs, audio_outputs[imp_ind_a], n_frames[imp_ind_a], sim_type)
+    a2i_loss = (margin + imp_sim_i - true_sim).clamp_min(0)
+    i2a_loss = (margin + imp_sim_a - true_sim).clamp_min(0)
+    return (a2i_loss + i2a_loss).mean() / 2
+
+
+class SimilarityCalibrator(torch.nn.Module):
+
+    def __init__(self, cal_init, max_w=100, min_w=.01, subtract_mean=True, use_bias=False):
+        super().__init__()
+        self.max_w = max_w
+        self.min_w = min_w
+        self.w = torch.nn.Parameter(torch.tensor([cal_init]).log())
+
+        self.use_bias = use_bias
+        if self.use_bias:
+            self.b = torch.nn.Parameter(torch.tensor([0.0]))
+
+        self.subtract_mean = subtract_mean
+
+    def get_w(self):
+        return torch.exp(self.w).clamp_max(self.max_w).clamp_min(self.min_w)
+
+    def forward(self, x):
+        sims = self.get_w() * x
+
+        if self.use_bias:
+            sims = sims + self.b
+
+        if self.subtract_mean:
+            return sims - sims.mean()
+        else:
+            return sims
+
+
+class SpatialDropout(torch.nn.Module):
+
+    def __init__(self, p, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.p = p
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        dropout = torch.rand((b, 1, h, w), dtype=x.dtype, device=x.device) > self.p
+
+        if self.training:
+            return x * dropout
+        else:
+            return x
+
+
+class LitAVAligner(pl.LightningModule, PyTorchModelHubMixin, repo_url="https://github.com/mhamilton723/DenseAV", license="mit", tags=["denseav"]):
+    def __init__(self,
+                 code_dim,
+                 image_model_type,
+                 image_model_token_type,
+                 image_aligner_type,
+                 image_pool_width,
+                 audio_model_type,
+                 audio_aligner_type,
+                 audio_pool_width,
+                 audio_lora,
+                 audio_lora_rank,
+                 image_lora,
+                 image_lora_rank,
+                 gradient_clipping,
+                 learn_audio_cls,
+                 silence_l1,
+                 silence_l2,
+                 tv_weight,
+                 nonneg_sim,
+                 nonneg_pressure,
+                 pretrain_lr,
+                 lr,
+                 lr_warmup,
+                 lr_schedule,
+                 lr_cycle_length,
+                 optimizer,
+                 gather_tensors,
+                 sim_agg_type,
+                 sim_agg_heads,
+                 sim_use_cls,
+                 disentangle_weight,
+                 norm_vectors,
+                 cal_init,
+                 cal_balance_weight,
+                 loss_type,
+                 loss_margin,
+                 mask_silence,
+                 finetune_image_model,
+                 finetune_audio_model,
+                 use_cached_embs,
+                 output_root,
+                 neg_audio,
+                 neg_audio_weight,
+                 head_agg,
+                 adaptive_clipping,
+                 specialization_weight,
+                 spatial_dropout,
+                 channel_dropout,
+                 mixup_weight,
+                 memory_buffer_size,
+                 loss_leak,
+                 ):
+        super().__init__()
+
+        self.code_dim = code_dim
+        self.image_model_type = image_model_type
+        self.image_model_token_type = image_model_token_type
+        self.image_aligner_type = image_aligner_type
+        self.image_pool_width = image_pool_width
+        self.audio_model_type = audio_model_type
+        self.audio_aligner_type = audio_aligner_type
+        self.audio_pool_width = audio_pool_width
+
+        self.gradient_clipping = gradient_clipping
+        self.learn_audio_cls = learn_audio_cls
+        self.silence_l1 = silence_l1
+        self.silence_l2 = silence_l2
+
+        self.tv_weight = tv_weight
+        self.nonneg_sim = nonneg_sim
+        self.nonneg_pressure = nonneg_pressure
+        self.pretrain_lr = pretrain_lr
+        self.lr = lr
+        self.lr_warmup = lr_warmup
+        self.lr_schedule = lr_schedule
+        self.lr_cycle_length = lr_cycle_length
+        self.optimizer = optimizer
+        self.gather_tensors = gather_tensors
+        self.sim_agg_type = sim_agg_type
+        self.sim_agg_heads = sim_agg_heads
+        self.sim_use_cls = sim_use_cls
+        self.disentangle_weight = disentangle_weight
+
+        self.norm_vectors = norm_vectors
+        self.cal_init = cal_init
+        self.cal_balance_weight = cal_balance_weight
+        self.loss_type = loss_type
+        self.loss_margin = loss_margin
+        self.mask_silence = mask_silence
+        self.finetune_image_model = finetune_image_model
+        self.finetune_audio_model = finetune_audio_model
+        self.use_cached_embs = use_cached_embs
+        self.output_root = output_root
+        self.audio_lora = audio_lora
+        self.audio_lora_rank = audio_lora_rank
+        self.image_lora = image_lora
+        self.image_lora_rank = image_lora_rank
+        self.neg_audio = neg_audio
+        self.neg_audio_weight = neg_audio_weight
+        self.head_agg = head_agg
+
+        self.adaptive_clipping = adaptive_clipping
+        self.specialization_weight = specialization_weight
+        self.spatial_dropout = spatial_dropout
+        self.channel_dropout = channel_dropout
+        self.mixup_weight = mixup_weight
+
+        self.memory_buffer_size = memory_buffer_size
+        self.memory_buffer = deque(maxlen=self.memory_buffer_size)
+        self.loss_leak = loss_leak
+
+        if self.audio_model_type in {"audiomae", "audiomae-finetuned", "cavmae", "cavmae-mixed", "imagebind"}:
+            self.audio_input = "spec"
+        elif self.audio_model_type == "davenet":
+            self.audio_input = "davenet_spec"
+        elif self.audio_model_type == "fnac":
+            self.audio_input = "fnac_spec"
+        else:
+            self.audio_input = "audio"
+
+        extra_model_args = dict(output_root=output_root)
+
+        self.image_model, _, self.image_feat_dim = get_image_featurizer(
+            image_model_type, token_type=self.image_model_token_type, **extra_model_args)
+
+        self.image_model.eval()
+        if not self.finetune_image_model:
+            for param in self.image_model.parameters():
+                param.requires_grad = False
+
+        if image_model_type in {"cavmae", "cavmae-mixed", "imagebind", "fnac"}:
+            extra_model_args["model"] = self.image_model.model
+
+        if use_cached_embs:
+            _, self.audio_feat_dim = get_audio_featurizer(audio_model_type, **extra_model_args)
+        else:
+            self.audio_model, self.audio_feat_dim = get_audio_featurizer(audio_model_type, **extra_model_args)
+
+            self.audio_model.eval()
+            if not self.finetune_audio_model:
+                for param in self.audio_model.parameters():
+                    param.requires_grad = False
+
+        if self.image_lora:
+            if self.image_model_type in {"sam", "dino8", "dinov2", "cavmae", "cavmae-mixed"}:
+                target_modules = ["qkv"]
+            elif self.image_model_type == "clip":
+                target_modules = ["out_proj"]
+            elif self.image_model_type == "imagebind":
+                target_modules = ["out_proj", "fc1", "fc2"]
+            else:
+                target_modules = ["q", "k", "v"]
+
+            peft_config = LoraConfig(
+                target_modules=target_modules,
+                inference_mode=False,
+                r=image_lora_rank,
+                lora_alpha=32,
+                lora_dropout=0.1
+            )
+            self.image_model = get_peft_model(self.image_model, peft_config)
+            self.image_model.print_trainable_parameters()
+
+        if self.audio_lora:
+            if self.audio_model_type == "hubert":
+                target_modules = ["q_proj", "k_proj", "v_proj"]
+            else:
+                target_modules = ["q", "k", "v"]
+
+            peft_config = LoraConfig(
+                inference_mode=False,
+                target_modules=target_modules,
+                r=audio_lora_rank,
+                lora_alpha=32,
+                lora_dropout=0.1
+            )
+            self.audio_model = get_peft_model(self.audio_model, peft_config)
+            self.audio_model.print_trainable_parameters()
+
+        shared_aligner_args = dict(out_dim=self.code_dim)
+
+        self.audio_aligner = get_aligner(
+            self.audio_aligner_type, self.audio_feat_dim, **shared_aligner_args)
+        self.image_aligner = get_aligner(
+            self.image_aligner_type, self.image_feat_dim, **shared_aligner_args)
+
+        if self.loss_type == "nce":
+            self.sim_cal = SimilarityCalibrator(self.cal_init, subtract_mean=True, use_bias=False)
+        else:
+            self.sim_cal = SimilarityCalibrator(self.cal_init, subtract_mean=False, use_bias=True)
+
+        if self.learn_audio_cls:
+            self.audio_cls = torch.nn.Parameter(torch.randn(self.audio_feat_dim))
+
+        if self.spatial_dropout > 0.0:
+            self.spatial_dropout_layer = SpatialDropout(self.spatial_dropout)
+
+        if self.channel_dropout > 0.0:
+            self.channel_dropout_layer = torch.nn.Dropout2d(self.channel_dropout)
+
+        self.sim_agg = get_aggregator(
+            self.sim_agg_type,
+            self.nonneg_sim,
+            self.mask_silence,
+            self.sim_agg_heads,
+            self.head_agg,
+            self.sim_use_cls,
+            dim=self.image_feat_dim
+        )
+
+        self.hparams_logged = False
+        self.rolling_avg = RollingAvg(50)
+        self.grad_avg = RollingAvg(50, nonzero=True)
+
+        self.save_hyperparameters()
+
+    def set_full_train(self, full_train):
+        self.full_train = full_train
+
+    def prep_feats(self, feats, is_audio):
+
+        if not is_audio and self.training and self.image_pool_width > 1:
+            feats = torch.nn.AvgPool2d(self.image_pool_width)(feats)
+
+        if is_audio and self.training and self.audio_pool_width > 1:
+            feats = torch.nn.AvgPool2d((1, self.audio_pool_width))(feats)
+
+        if self.norm_vectors:
+            feats = F.normalize(feats, dim=1)
+
+        return feats
+
+    def on_before_optimizer_step(self, optimizer, optimizer_idx):
+        norms = grad_norm(self, norm_type=2)
+        avg_grads = self.grad_avg.get_all()
+        params = {
+            f"grad_2.0_norm/{name}": p
+            for name, p in self.named_parameters()
+            if p.grad is not None
+        }
+
+        if self.adaptive_clipping:
+            for k in norms.keys():
+                if k in params:
+                    avg_grad = max(avg_grads.get(k, norms[k]), 1e-5)
+                    if self.global_step > 10 and norms[k] > avg_grad * 5:
+                        print(f"Bad grad for {k}: {norms[k]} scaling to {avg_grad * 5}")
+                        torch.nn.utils.clip_grad_norm_(params[k], avg_grad * 5)
+                        norms[k] = avg_grad * 5
+
+                    if norms[k] > self.gradient_clipping:
+                        # print(f"Bad grad for {k}: {norms[k]} scaling to {self.gradient_clipping}")
+                        torch.nn.utils.clip_grad_norm_(params[k], self.gradient_clipping)
+
+        # self.grad_avg.add_all(norms)
+        # self.log_dict(norms)
+
+    def interpolate_mask(self, mask, target_length, discrete):
+        b, t = mask.shape
+
+        mask = F.interpolate(mask.reshape(b, 1, 1, t), (1, target_length), mode="bilinear") \
+            .reshape(b, target_length)
+
+        if discrete:
+            mask = mask > 0.01
+            sums = mask.sum(1)
+            all_zeros = torch.where(sums == 0)[0]
+            if len(all_zeros) > 0:
+                print("Fixing a bad mask")
+                for entry in all_zeros:
+                    mask[entry, torch.randint(0, target_length - 1, size=())] = True
+        else:
+            return mask
+        return mask
+
+    def forward_audio(self, batch):
+        if self.use_cached_embs:
+            audio_feats = batch["audio_emb"]
+            if "audio_cls" in batch:
+                audio_cls = batch["audio_cls"]
+            else:
+                audio_cls = None
+        else:
+            audio = batch[self.audio_input]
+
+            if self.full_train:
+                audio_feats, audio_cls = self.audio_model(audio, include_cls=True)
+            else:
+                with torch.no_grad():
+                    audio_feats, audio_cls = self.audio_model(audio, include_cls=True)
+
+        mask = batch[AUDIO_MASK] if AUDIO_MASK in batch else torch.ones_like(audio)
+        pos_mask = batch[AUDIO_POS_MASK] if AUDIO_POS_MASK in batch else torch.ones_like(audio)
+
+        if self.learn_audio_cls:
+            assert audio_cls is None
+            audio_cls = torch.broadcast_to(self.audio_cls.unsqueeze(0), (audio_feats.shape[0], audio_feats.shape[1]))
+
+        aligned_audio_feats, aligned_audio_cls = self.audio_aligner(audio_feats, audio_cls)
+
+        if self.channel_dropout > 0.0:
+            aligned_audio_feats = self.channel_dropout_layer(aligned_audio_feats)
+
+        aligned_audio_feats = self.prep_feats(aligned_audio_feats, is_audio=True)
+        audio_mask = self.interpolate_mask(mask, aligned_audio_feats.shape[-1], True)
+        audio_pos_mask = self.interpolate_mask(pos_mask, aligned_audio_feats.shape[-1], False)
+
+        ret = {
+            AUDIO_MASK: audio_mask,
+            AUDIO_POS_MASK: audio_pos_mask,
+            AUDIO_FEATS: aligned_audio_feats,
+        }
+
+        if aligned_audio_cls is not None:
+            ret[AUDIO_CLS] = aligned_audio_cls
+
+        return ret
+
+    # @autocast(device_type="cuda", enabled=False)
+    def forward_image(self, batch, max_batch_size=None):
+
+        with torch.no_grad():
+            image = batch[IMAGE_INPUT]
+            b, nf, c, h, w = image.shape
+            image = image.reshape(b * nf, c, h, w)
+
+            if max_batch_size is None:
+                max_batch_size = image.shape[0]
+
+            chunks = [image[i:i + max_batch_size] for i in range(0, image.shape[0], max_batch_size)]
+
+            all_image_feats = []
+            all_image_cls = []
+
+            for chunk in chunks:
+                if self.full_train:
+                    image_feats, image_cls = self.image_model(chunk, include_cls=True)
+                else:
+                    with torch.no_grad():
+                        image_feats, image_cls = self.image_model(chunk, include_cls=True)
+
+                aligned_image_feats, aligned_image_cls = self.image_aligner(image_feats, image_cls)
+
+                all_image_feats.append(aligned_image_feats)
+                all_image_cls.append(aligned_image_cls)
+
+            # Stitch the chunks back together
+            aligned_image_feats = torch.cat(all_image_feats, dim=0)
+            aligned_image_cls = torch.cat(all_image_cls, dim=0)
+
+        if self.channel_dropout > 0.0:
+            aligned_image_feats = self.channel_dropout_layer(aligned_image_feats)
+
+        if self.spatial_dropout > 0.0:
+            aligned_image_feats = self.spatial_dropout_layer(aligned_image_feats)
+
+        aligned_image_feats = self.prep_feats(aligned_image_feats, is_audio=False)
+        ret = {IMAGE_FEATS: aligned_image_feats}
+
+        if IMAGE_MASK in batch:
+            with torch.no_grad():
+                mask = batch[IMAGE_MASK]
+                mask = mask.reshape(b * nf, 1, h, w)
+                b, c, h, w = aligned_image_feats.shape
+                mask = F.adaptive_avg_pool2d(mask.to(aligned_image_feats), output_size=(h, w))
+                ret[IMAGE_MASK] = mask
+
+        if aligned_image_cls is not None:
+            ret[IMAGE_CLS] = aligned_image_cls
+
+        return ret
+
+    def forward(self, batch):
+        audio_feat_dict = self.forward_audio(batch)
+        image_feat_dict = self.forward_image(batch)
+        return {**image_feat_dict, **audio_feat_dict}
+
+    def contrast_loss(self, sims):
+        b = sims.shape[0]
+        sims = sims - torch.eye(b, b, device=sims.device) * self.loss_margin
+        sims_1 = sims
+        sims_2 = sims.permute(1, 0)
+
+        if self.loss_leak > 0.0:
+            id = torch.eye(sims_1.shape[0], sims_1.shape[1], device=sims.device, dtype=sims.dtype)
+            label_mask = id * (1 - self.loss_leak)
+            label_mask += (1 - id) * self.loss_leak / (sims_1.shape[0] - 1)
+            label_mask /= label_mask.sum(dim=1, keepdim=True)
+        else:
+            label_mask = torch.eye(sims_1.shape[0], sims_1.shape[1], device=sims.device, dtype=sims.dtype)
+
+        labels = torch.arange(0, sims.shape[0], device=sims.device)
+        self.rolling_avg.add(f"acc/1", (sims.argmax(dim=1) == labels).to(sims).mean())
+        self.rolling_avg.add(f"acc/2", (sims.argmax(dim=0) == labels).to(sims).mean())
+
+        if self.loss_type == "margin":
+            margin_loss_tensor = (sims - torch.diag(sims)).clamp_min(0)
+            margin_loss = margin_loss_tensor.mean()
+            self.rolling_avg.add(f"loss/frac_nonzero", (margin_loss_tensor > 0).to(sims).mean())
+            self.rolling_avg.add(f"loss/margin", margin_loss)
+            return margin_loss
+        elif self.loss_type == "ce":
+            ce_loss = 1 / 2 * F.cross_entropy(sims_1, labels) + \
+                      1 / 2 * F.cross_entropy(sims_2, labels)
+            self.rolling_avg.add(f"loss/ce", ce_loss)
+            return ce_loss
+        elif self.loss_type == "bce":
+            bce_loss = F.binary_cross_entropy_with_logits(sims_1.flatten(), label_mask.flatten())
+            self.rolling_avg.add(f"loss/bce", bce_loss)
+            return bce_loss
+        elif self.loss_type == "nce":
+            nce_loss = 1 / 2 * (-F.log_softmax(sims_1, dim=-1) * label_mask).sum(1).mean() + \
+                       1 / 2 * (-F.log_softmax(sims_2, dim=-1) * label_mask).sum(1).mean()
+            self.rolling_avg.add(f"loss/nce", nce_loss)
+            return nce_loss
+        else:
+            raise ValueError(f"Unknown loss type {self.loss_type}")
+
+    def loss(self, preds):
+        image_feats = preds[IMAGE_FEATS]
+        audio_feats = preds[AUDIO_FEATS]
+        audio_mask = preds[AUDIO_MASK]
+        image_mask = preds[IMAGE_MASK]
+        audio_pos_mask = preds[AUDIO_POS_MASK]
+        if DATA_SOURCE in preds:
+            source = preds[DATA_SOURCE].to(torch.int64)
+        else:
+            source = None
+
+        uncal_sims = self.sim_agg(preds, agg_heads=True)
+        sims = self.sim_cal(uncal_sims)
+
+        _mask = 1 - torch.eye(sims.shape[0], device=sims.device)
+        self.log(f"sim/pos", torch.diag(sims).mean())
+        self.log(f"sim/neg", (sims * _mask).sum() / (_mask.sum()))
+        self.log(f"sim/uncal_pos", torch.diag(uncal_sims).mean())
+        self.log(f"sim/uncal_neg", (uncal_sims * _mask).sum() / (_mask.sum()))
+
+        b, c, h, w = image_feats.shape
+        b, c, f, t = audio_feats.shape
+        n_samples = 250
+
+        nh = self.sim_agg_heads
+        image_feats_by_head = image_feats.reshape(b, self.sim_agg_heads, c // nh, h, w)
+        audio_feats_by_head = audio_feats.reshape(b, self.sim_agg_heads, c // nh, f, t)
+
+        def maybe_clamp(t):
+            return t.clamp_min(0) if self.nonneg_sim else t
+
+        paired_sim_raw = self.sim_agg.get_pairwise_sims(preds, raw=True, agg_sim=False, agg_heads=False)
+        paired_sim = maybe_clamp(paired_sim_raw)
+
+        loss = 0.0
+
+        if self.nonneg_pressure:
+            afb, afk, afc, aff, aft = audio_feats_by_head.shape
+            ifb, ifk, ifc, ifh, ifw = image_feats_by_head.shape
+            assert (afb == ifb)
+
+            device = audio_feats_by_head.device
+            random_b = torch.randint(0, afb, size=(n_samples,), device=device)
+            random_t = torch.randint(0, aft, size=(n_samples,), device=device)
+            random_f = torch.randint(0, aff, size=(n_samples,), device=device)
+            random_h = torch.randint(0, ifh, size=(n_samples,), device=device)
+            random_w = torch.randint(0, ifw, size=(n_samples,), device=device)
+
+            random_audio_feats = audio_feats_by_head[random_b, :, :, random_f, random_t]
+            random_image_feats = image_feats_by_head[random_b, :, :, random_h, random_w]
+            random_sim_raw = torch.einsum("bkc,dkc->bdk", random_audio_feats, random_image_feats)
+
+            nonneg_loss = random_sim_raw.clamp_max(0).square().mean()
+            self.rolling_avg.add(f"loss/nonneg", nonneg_loss)
+            loss += nonneg_loss * self.nonneg_pressure
+
+        if self.silence_l1 > 0 or self.silence_l2 > 0:
+            masked_b, masked_t = torch.where(~audio_mask)
+            if len(masked_b) > n_samples:
+                subset = torch.randperm(len(masked_b))[:n_samples]
+                masked_b = masked_b[subset]
+                masked_t = masked_t[subset]
+
+            if len(masked_b) == n_samples:
+                silent_audio_feats = audio_feats_by_head[masked_b, :, :, :, masked_t].mean(-1)  # d k c
+                silence_tensor = maybe_clamp(
+                    torch.einsum("bkchw,dkc->bkdhw", image_feats_by_head, silent_audio_feats))
+
+                silence_l1_loss = silence_tensor.abs().mean()
+                self.rolling_avg.add(f"loss/silence_l1", silence_l1_loss)
+                loss += silence_l1_loss * self.silence_l1
+
+                silence_l2_loss = silence_tensor.square().mean()
+                self.rolling_avg.add(f"loss/silence_l2", silence_l2_loss)
+                loss += silence_l2_loss * self.silence_l2
+            else:
+                pass
+
+        if self.neg_audio_weight > 0 and self.neg_audio:
+            b, t = audio_pos_mask.shape
+            negative_weight = ((1 - audio_pos_mask) * audio_mask.to(sims)).reshape(b, 1, 1, 1, 1, t)
+            negative_weight = torch.broadcast_to(negative_weight, paired_sim.shape)
+            if negative_weight.sum() > 0:
+                neg_audio_loss = (paired_sim.square() * negative_weight).sum() \
+                                 / negative_weight.sum().clamp_min(0.1)
+                self.rolling_avg.add(f"loss/neg_audio", neg_audio_loss)
+                self.rolling_avg.add(f"loss/neg_weight_avg", negative_weight.mean())
+                loss += neg_audio_loss * self.neg_audio_weight
+            else:
+                print("WARNING: No negative samples found in batch")
+
+        if self.tv_weight > 0:
+            tv_loss = (paired_sim[:, :, :, :, :, 1:] - paired_sim[:, :, :, :, :, :-1]).square().mean()
+            self.rolling_avg.add(f"loss/tv", tv_loss)
+            loss += tv_loss * self.tv_weight
+
+        self.log(f"cal/w", self.sim_cal.get_w())
+        if self.cal_balance_weight > 0.0:
+            cal_balance = (np.log(self.cal_init) - torch.log(self.sim_cal.get_w().clamp_min(.00000001))) \
+                .clamp_min(0).square().mean()
+            self.rolling_avg.add(f"loss/cal_balance", cal_balance)
+            loss += cal_balance * self.cal_balance_weight
+
+        if self.disentangle_weight > 0.0:
+            assert source is not None
+            assert self.sim_agg_heads % 2 == 0
+
+            dilation = self.sim_agg_heads // 2
+            sources_oh = F.one_hot(source, num_classes=2)
+            b, h = sources_oh.shape
+            sources_mask = 1 - torch.broadcast_to(sources_oh.unsqueeze(-1), (b, h, dilation)) \
+                .reshape(b, h * dilation).to(paired_sim)
+            disentangle_loss = torch.einsum("bkhwft,bk->bhwft", paired_sim, sources_mask).square().mean()
+            self.rolling_avg.add(f"loss/disentangle", disentangle_loss)
+            loss += disentangle_loss * self.disentangle_weight
+
+        if self.specialization_weight > 0.0 and self.sim_agg_heads > 1:
+            total_specialization_loss = 0.0
+            combos = list(combinations(range(self.sim_agg_heads), 2))
+            for i, j in combos:
+                specialization_loss_pair = (paired_sim[:, i].abs() * paired_sim[:, j].abs()).mean()
+                total_specialization_loss += specialization_loss_pair
+            avg_specialization_loss = total_specialization_loss / len(combos)
+            self.rolling_avg.add(f"loss/specialize", avg_specialization_loss)
+            loss += avg_specialization_loss * self.specialization_weight
+
+        if self.mixup_weight > 0.0:
+            b, _, h, w = image_mask.shape
+            neg_img_mask = torch.broadcast_to(
+                1 - image_mask.to(paired_sim).reshape(b, 1, h, w, 1, 1),
+                paired_sim.shape)
+            image_mixup_loss = (paired_sim * neg_img_mask).square().sum() / neg_img_mask.sum().clamp_min(0.1)
+            self.rolling_avg.add(f"loss/image_mixup", image_mixup_loss)
+            loss += image_mixup_loss * self.mixup_weight
+
+        sims = sims
+        loss += self.contrast_loss(sims)
+        self.rolling_avg.add(f"loss/total", loss)
+
+        return loss
+
+    def setup_hparams(self):
+        recalls = ['A_r1', 'A_r5', 'A_r10', 'I_r1', 'I_r5', 'I_r10']
+
+        if self.trainer.datamodule.use_extra_val_sets:
+            datasets = ["Places", "AudioSet"]
+        else:
+            datasets = ["Val"]
+
+        heads = ["total"]
+
+        metric_names = [
+            "hp/speech_basic_ap", "hp/speech_advanced_ap", "hp/sound_basic_ap",
+            "hp/speech_basic_iou", "hp/speech_advanced_iou", "hp/sound_basic_iou",
+        ]
+        for dataset in datasets:
+            for head in heads:
+                for recall in recalls:
+                    metric_names.append(f"hp/{dataset}/{head}/{recall}")
+
+        if self.sim_agg_heads == 2:
+            metric_names.extend(["hp/ap_dis", "hp/act_dis"])
+
+        if hasattr(self.trainer, "datamodule"):
+            all_hparams = {**self.hparams, **self.trainer.datamodule.hparams}
+        else:
+            all_hparams = self.hparams
+
+        starting_values = {n: torch.nan for n in metric_names}
+        self.logger.log_hyperparams(all_hparams, starting_values)
+
+    def on_train_start(self):
+        self.setup_hparams()
+        self.hparams_logged = True
+
+    def on_train_batch_start(self, batch, batch_idx):
+        remake_optimizers = False
+
+        if isinstance(self.image_aligner, ProgressiveGrowing):
+            should_remake = self.image_aligner.maybe_change_phase(self.global_step)
+            remake_optimizers = remake_optimizers or should_remake
+        if isinstance(self.audio_aligner, ProgressiveGrowing):
+            should_remake = self.audio_aligner.maybe_change_phase(self.global_step)
+            remake_optimizers = remake_optimizers or should_remake
+
+        if remake_optimizers:
+            raise NotImplementedError()
+
+    def _combine_preds(self, all_preds):
+        temp = {}
+        new_preds = {}
+
+        # Collect tensors for each key into lists
+        for d in all_preds:
+            for key, value in d.items():
+                if isinstance(value, torch.Tensor):
+                    if key not in temp:
+                        temp[key] = []
+                    temp[key].append(value)
+
+        # Concatenate all tensors for each key using a single call to torch.cat
+        for key, tensor_list in temp.items():
+            new_preds[key] = torch.cat(tensor_list)
+        return new_preds
+
+    def training_step(self, batch, batch_idx):
+        assert batch[IMAGE_INPUT].shape[1] == 1
+
+        preds = self.forward(batch)
+        if DATA_SOURCE in batch:
+            preds[DATA_SOURCE] = batch[DATA_SOURCE]
+
+        if self.trainer.world_size > 1 and self.gather_tensors:
+            for k, v in preds.items():
+                new_v = v.contiguous()
+                preds[k] = torch.cat(GatherLayer.apply(new_v), dim=0)
+
+        if self.memory_buffer_size > 0:
+            new_preds = self._combine_preds(list(self.memory_buffer) + [preds])
+        else:
+            new_preds = preds
+
+        loss = self.loss(new_preds)
+
+        if self.memory_buffer_size > 0:
+            self.memory_buffer.append(self._recursive_detach(preds, gather=False))
+
+        if self.trainer.is_global_zero and self.global_step % 50 == 1:
+            writer = self.logger.experiment
+            self.rolling_avg.logall(lambda k, v: writer.add_scalar(k, v, global_step=self.global_step))
+
+        if self.trainer.scaler is not None:
+            self.log("loss_scale", self.trainer.scaler.get_scale())
+
+        if self.global_step % 10000 == 0 and self.global_step > 0:
+            print("RESETTING TFEVENT FILE")
+            self.logger.experiment.close()
+            self.logger.experiment._get_file_writer()
+
+        return loss
+
+    def on_validation_start(self) -> None:
+        if not self.hparams_logged:
+            self.setup_hparams()
+            self.hparams_logged = True
+
+    def _auto_gather(self, t):
+        if t.dtype == torch.bool:
+            t = t.to(torch.float)
+
+        if self.trainer.num_devices == 1:
+            return t.cpu()
+
+        t = torch.clone(t).contiguous()
+        if self.trainer.is_global_zero:
+            gather_list = [torch.zeros_like(t) for _ in range(dist.get_world_size())]
+            dist.gather(t, gather_list)
+            return torch.cat(gather_list, dim=0).cpu()
+        else:
+            dist.gather(t)
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+
+        with torch.no_grad():
+            preds = self.forward(batch)
+
+            ret = {}
+            for k in preds.keys():
+                if k in preds:
+                    ret[k] = self._auto_gather(preds[k])
+
+            batch_keys = [IMAGE_INPUT, "spec", "semseg", "num_pixels_per_class", 'total_length']
+            for k in batch_keys:
+                if k in batch:
+                    ret[k] = self._auto_gather(batch[k])
+
+            if "metadata" in batch:
+                if isinstance(batch["metadata"]["id"], torch.Tensor):
+                    ret["id"] = self._auto_gather(batch["metadata"]["id"])
+                ret["index"] = self._auto_gather(batch["metadata"]["index"])
+
+            return ret
+
+    def _calc_recalls(self, sim):
+        top_10_a = sim.topk(10, 0).indices == torch.arange(sim.shape[0]).unsqueeze(0)
+        top_10_i = (sim.topk(10, 1).indices == torch.arange(sim.shape[0]).unsqueeze(1)).permute(1, 0)
+        a_recall = lambda p: top_10_a[0:p].any(0).to(sim).mean()
+        i_recall = lambda p: top_10_i[0:p].any(0).to(sim).mean()
+        return {'A_r1': a_recall(1),
+                'A_r5': a_recall(5),
+                'A_r10': a_recall(10),
+                'I_r1': i_recall(1),
+                'I_r5': i_recall(5),
+                'I_r10': i_recall(10)}
+
+    def calc_recalls(self, preds, dataset):
+        sim = self.sim_agg.forward_batched(
+            preds=preds,
+            agg_heads=False,
+            batch_size=4,
+        ).cpu()
+
+        all_metrics = dict()
+        for k, v in self._calc_recalls(sim.sum(-1)).items():
+            all_metrics[f"hp/{dataset}/total/" + k] = v
+
+        return all_metrics
+
+    def retrieval_validation(self, outputs, dataset_name):
+        if len(outputs) == 0:
+            return
+
+        if self.trainer.is_global_zero:
+            results = flatten_preds(outputs)
+            if not self.trainer.sanity_checking:
+                print(results[IMAGE_FEATS].shape[0])
+                # assert (results[IMAGE_FEATS].shape[0] == 1000)
+            results[IMAGE_FEATS] = results[IMAGE_FEATS].cpu()
+            results[AUDIO_FEATS] = results[AUDIO_FEATS].cuda()
+            if self.sim_use_cls:
+                results[AUDIO_CLS] = results[AUDIO_CLS].cuda()
+                results[AUDIO_CLS] = results[AUDIO_CLS].cuda()
+
+            results[AUDIO_MASK] = results[AUDIO_MASK].cuda()
+
+            recalls = self.calc_recalls(results, dataset_name)
+
+            results[IMAGE_FEATS] = results[IMAGE_FEATS].cuda()
+
+            writer = self.logger.experiment
+            print("here")
+            for name, v in recalls.items():
+                writer.add_scalar(f"{name}", v, self.global_step + 1)
+
+    def semseg_validation(self, speech_preds, sound_preds):
+
+        if self.trainer.is_global_zero:
+            from eval_utils import get_paired_heatmaps
+            def prep_preds(preds, loader):
+                results = flatten_preds(preds)
+                metadata = loader.dataset.metadata
+                ordered_metadata = metadata.iloc[results["index"].numpy(), :].copy()
+                ordered_metadata["order"] = range(len(ordered_metadata))
+                return results, ordered_metadata
+
+            [_, _, speech_loader, sound_loader] = self.trainer.val_dataloaders
+            speech_results, speech_metadata = prep_preds(speech_preds, speech_loader)
+            sound_results, sound_metadata = prep_preds(sound_preds, sound_loader)
+
+            self.sound_metrics, unique_sound_indices = get_paired_heatmaps(
+                self, sound_results, sound_metadata["ade_class_id"], None)
+
+            self.speech_metrics, unique_word_indices = get_paired_heatmaps(
+                self, speech_results, speech_metadata["ade_class_id"], speech_metadata["timing"])
+
+            writer = self.logger.experiment
+
+            all_metrics = {
+                **{"sound_" + k: v for k, v in self.sound_metrics.items()},
+                **{"speech_" + k: v for k, v in self.speech_metrics.items()},
+            }
+
+            for k, v in all_metrics.items():
+                writer.add_scalar(f"hp/{k}", torch.tensor(v).mean(), self.global_step + 1)
+
+    def disentangle_validation(self, word_preds, sound_preds):
+
+        if len(word_preds) == 0 or len(sound_preds) == 0:
+            return
+
+        if self.trainer.is_global_zero:
+            word_preds = flatten_preds(word_preds)
+            sound_preds = flatten_preds(sound_preds)
+
+            word_scores = self.sim_agg.get_pairwise_sims(
+                word_preds,
+                raw=False,
+                agg_sim=True,
+                agg_heads=False,
+            )
+
+            sound_scores = self.sim_agg.get_pairwise_sims(
+                sound_preds,
+                raw=False,
+                agg_sim=True,
+                agg_heads=False,
+            )
+
+            all_scores = torch.cat([word_scores, sound_scores], dim=0)
+            all_scores -= all_scores.min(dim=0, keepdim=True).values
+            all_scores /= all_scores.max(dim=0, keepdim=True).values.clamp_min(.0001)
+
+            is_words = torch.cat([
+                torch.ones(word_scores.shape[0]),
+                torch.zeros(sound_scores.shape[0])], dim=0).to(torch.bool)
+
+            assert all_scores.shape[1] == 2
+            ap_matrix = torch.zeros(2, 2)
+            act_matrix = torch.zeros(2, 2)
+
+            for head in range(2):
+                # writer.add_histogram(f"h{head}_all_scores", all_scores[:, head])
+                for dataset_num in range(2):
+                    if dataset_num == 0:
+                        labels = is_words
+                    else:
+                        labels = ~is_words
+
+                    ap_matrix[head, dataset_num] = binary_average_precision(
+                        all_scores[:, head].cpu(), labels.to(torch.int64).cpu())
+
+                    act_matrix[head, dataset_num] = 1 - (all_scores[:, head][labels]).mean()
+
+            ap_dis = max(.5 * (ap_matrix[0, 0] + ap_matrix[1, 1]),
+                         .5 * (ap_matrix[0, 1] + ap_matrix[1, 0]))
+
+            act_dis = max(.5 * (act_matrix[0, 0] + act_matrix[1, 1]),
+                          .5 * (act_matrix[0, 1] + act_matrix[1, 0]))
+
+            print("AP", ap_matrix)
+            print("AP dis", ap_dis)
+            print("Act", act_matrix)
+            print("Act dis", act_dis)
+
+            writer = self.logger.experiment
+            writer.add_scalar("hp/ap_dis", ap_dis, self.global_step + 1)
+            writer.add_scalar("hp/act_dis", act_dis, self.global_step + 1)
+
+    def validation_epoch_end(self, outputs) -> None:
+        print("Val end")
+        with torch.no_grad():
+            if self.trainer.datamodule.use_extra_val_sets:
+                if self.sim_agg_heads == 2:
+                    self.disentangle_validation(outputs[0], outputs[1])
+                self.retrieval_validation(outputs[0], "Places")
+                self.retrieval_validation(outputs[1], "AudioSet")
+                self.semseg_validation(outputs[2], outputs[3])
+
+            else:
+                print("HERE!")
+                self.retrieval_validation(outputs, "Val")
+
+        writer = self.logger.experiment
+        writer.flush()
+
+    def _recursive_detach(self, obj, gather=True):
+        if isinstance(obj, torch.Tensor):
+            if gather:
+                return self._auto_gather(obj)
+            else:
+                obj.detach()
+        elif isinstance(obj, dict):
+            return {k: self._recursive_detach(v, gather) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [self._recursive_detach(v, gather) for v in obj]
+        else:
+            return obj
+
+    def predict_step(self, batch, batch_idx: int, dataloader_idx: int = 0):
+        with torch.no_grad():
+            predictions = {}
+            for k, v in batch.items():
+                predictions[k] = self._recursive_detach(v)
+            for k, v in self.forward(batch).items():
+                predictions[k] = self._auto_gather(v)
+
+            return predictions
+
+    def _configure_optimizers(self, full_train, lr):
+        params = [
+            *self.audio_aligner.parameters(),
+            *self.image_aligner.parameters(),
+            *self.sim_cal.parameters(),
+            *self.sim_agg.parameters()
+        ]
+
+        if (self.finetune_image_model or self.image_lora) and full_train:
+            params.extend(self.image_model.parameters())
+
+        if (self.finetune_audio_model or self.audio_lora) and full_train:
+            params.extend(self.audio_model.parameters())
+
+        if self.learn_audio_cls:
+            params.append(self.audio_cls)
+
+        last_epoch = self.global_step - 1
+        if self.optimizer == "adam":
+            opt = torch.optim.Adam(params, lr=lr, eps=1e-7)
+        elif self.optimizer == "nadam":
+            opt = torch.optim.NAdam(params, lr=lr, eps=1e-7)
+        else:
+            raise ValueError(f"Unknown optimizer {self.optimizer}")
+
+        if self.lr_schedule == "sgdr":
+            scheduler = CosineAnnealingWarmRestarts(
+                opt, self.lr_cycle_length, 2, eta_min=lr * 2e-2, last_epoch=last_epoch)
+        else:
+            scheduler = LambdaLR(opt, lr_lambda=lambda step: 1.0, last_epoch=last_epoch)
+
+        if self.lr_warmup > 0:
+            warmup = LambdaLR(
+                opt,
+                lr_lambda=lambda step: min(max(float(step), 0.0) / self.lr_warmup, 1.0),
+                last_epoch=last_epoch,
+            )
+            scheduler = SequentialLR(
+                opt,
+                schedulers=[warmup, scheduler],
+                milestones=[self.lr_warmup],
+                last_epoch=last_epoch)
+
+        scheduler = {"scheduler": scheduler, "interval": "step"}
+
+        return [opt], [scheduler]
+
+    def configure_optimizers(self):
+        if self.full_train:
+            return self._configure_optimizers(self.full_train, self.lr)
+        else:
+            return self._configure_optimizers(self.full_train, self.pretrain_lr)
+
+
+@hydra.main(config_path="configs", config_name="av_align.yaml", version_base=None)
+def my_app(cfg: DictConfig) -> None:
+    print(OmegaConf.to_yaml(cfg))
+    seed_everything(cfg.seed, workers=True)
+
+    exp_name = f"{cfg.resume_prefix}"
+
+    if cfg.image_model_type == "dino8":
+        patch_size = 8 * cfg.image_pool_width
+    elif cfg.image_model_type == "cavmae":
+        patch_size = 16 * cfg.image_pool_width
+    elif cfg.image_model_type == "imagebind":
+        patch_size = 16 * cfg.image_pool_width
+    elif cfg.image_model_type == "clip":
+        patch_size = 16 * cfg.image_pool_width
+    elif cfg.image_model_type == "cavmae-mixed":
+        patch_size = 16 * cfg.image_pool_width
+    elif cfg.image_model_type == "dinov2":
+        patch_size = 14 * cfg.image_pool_width
+    else:
+        raise ValueError(f"Unknown patch size for model {cfg.image_model_type}")
+
+    datamodule = AVDataModule(
+        dataset_name=cfg.dataset_name,
+        load_size=cfg.load_size,
+        image_aug=cfg.image_aug,
+        audio_aug=cfg.audio_aug,
+        extra_audio_masking=cfg.extra_audio_masking,
+        audio_model_type=cfg.audio_model_type,
+        pytorch_data_dir=cfg.pytorch_data_dir,
+        use_cached_embs=cfg.use_cached_embs,
+        batch_size=cfg.batch_size,
+        num_workers=cfg.num_workers,
+        audio_level=cfg.audio_level,
+        neg_audio=cfg.neg_audio,
+        use_original_val_set=not cfg.use_extra_val_sets,
+        use_extra_val_sets=cfg.use_extra_val_sets,
+        data_for_plotting=False,
+        quad_mixup=cfg.quad_mixup,
+        bg_mixup=cfg.bg_mixup,
+        patch_mixup=cfg.patch_mixup,
+        patch_size=patch_size
+    )
+    datamodule.maybe_unpack(remove_source=cfg.submitting_to_aml)
+
+    aligner = create_model_from_cfg(LitAVAligner, cfg, {})
+
+    if cfg.starting_weights is not None:
+        loaded = torch.load(join(cfg.output_root, cfg.starting_weights), map_location='cpu')
+        state = loaded["state_dict"]
+        aligner.load_state_dict(state, strict=cfg.load_strict)
+        del state
+        del loaded
+
+    if cfg.num_gpus > 1:
+        # strategy = "ddp_sharded"  # _find_unused_parameters_true"
+        strategy = "ddp"  # _find_unused_parameters_true"
+    else:
+        strategy = "auto"
+
+    if cfg.dataset_name in {"places-audio", "mixed", "audio-set", "mixed-full"}:
+        val_args = dict(check_val_every_n_epoch=2)
+    elif cfg.dataset_name in {"dolphin"}:
+        val_args = dict(check_val_every_n_epoch=5)
+    else:
+        val_args = dict(val_check_interval=10000)
+
+    # val_args = dict(val_check_interval=1000)
+
+    def maybe_get_ckpt(ckpt_dir):
+        if cfg.auto_resume and os.path.exists(ckpt_dir):
+            print(f"Attempting to resume from {ckpt_dir}")
+            candidates = os.listdir(ckpt_dir)
+            assert (len(candidates) == 1)
+            return join(ckpt_dir, candidates[0])
+        elif cfg.auto_resume:
+            print(f"Could not find checkpoint at {ckpt_dir}")
+            return None
+        else:
+            return None
+
+    log_dir = join(cfg.output_root, "logs", cfg.grouping_name, exp_name)
+    ckpt_dir = join(cfg.output_root, "checkpoints", cfg.grouping_name, exp_name)
+
+    import gc
+    torch.cuda.empty_cache()
+    gc.collect()
+
+    def run_exp(aligner, full_train):
+        trainer_args = dict(
+            accelerator='gpu',
+            strategy=strategy,
+            devices=cfg.num_gpus,
+            num_sanity_val_steps=cfg.num_sanity_val_steps,
+            log_every_n_steps=50,
+            reload_dataloaders_every_n_epochs=10,
+            precision="16",
+            # profiler="simple",
+            # precision="bf16",
+            max_steps=cfg.max_steps,
+            **val_args)
+
+        aligner.set_full_train(full_train)
+        if full_train:
+            suffix = "train"
+        else:
+            suffix = "pretrain"
+            trainer_args["max_steps"] = cfg.pretrain_steps
+
+        print(f"Starting {suffix} phase")
+
+        logger = TensorBoardLogger(join(log_dir, suffix), default_hp_metric=False)
+        callbacks = [
+            ModelCheckpoint(join(ckpt_dir, suffix), every_n_epochs=1),
+            LearningRateMonitor(logging_interval='step'),
+        ]
+        Trainer(logger=logger,
+                callbacks=callbacks,
+                **trainer_args).fit(
+            aligner,
+            datamodule=datamodule,
+            ckpt_path=maybe_get_ckpt(join(ckpt_dir, suffix)))
+
+    train_chkpt = maybe_get_ckpt(join(ckpt_dir, "train"))
+
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    if cfg.pretrain_steps > 0 and train_chkpt is None:
+        run_exp(aligner, full_train=False)
+    run_exp(aligner, full_train=True)
+
+
+if __name__ == "__main__":
+    my_app()

DenseAV/gradio_app.py

@@ -0,0 +1,196 @@
+import csv
+import os
+import tempfile
+
+import gradio as gr
+import requests
+import torch
+import torchvision
+import torchvision.transforms as T
+from PIL import Image
+from featup.util import norm
+from torchaudio.functional import resample
+
+from denseav.train import LitAVAligner
+from denseav.plotting import plot_attention_video, plot_2head_attention_video, plot_feature_video
+from denseav.shared import norm, crop_to_divisor, blur_dim
+from os.path import join
+
+if __name__ == "__main__":
+
+    mode = "local"
+
+    if mode == "local":
+        sample_videos_dir = "samples"
+    else:
+        os.environ['TORCH_HOME'] = '/tmp/.cache'
+        os.environ['HF_HOME'] = '/tmp/.cache'
+        os.environ['HF_DATASETS_CACHE'] = '/tmp/.cache'
+        os.environ['TRANSFORMERS_CACHE'] = '/tmp/.cache'
+        os.environ['GRADIO_EXAMPLES_CACHE'] = '/tmp/gradio_cache'
+        sample_videos_dir = "/tmp/samples"
+
+
+    def download_video(url, save_path):
+        response = requests.get(url)
+        with open(save_path, 'wb') as file:
+            file.write(response.content)
+
+
+    base_url = "https://marhamilresearch4.blob.core.windows.net/denseav-public/samples/"
+    sample_videos_urls = {
+        "puppies.mp4": base_url + "puppies.mp4",
+        "peppers.mp4": base_url + "peppers.mp4",
+        "boat.mp4": base_url + "boat.mp4",
+        "elephant2.mp4": base_url + "elephant2.mp4",
+
+    }
+
+    # Ensure the directory for sample videos exists
+    os.makedirs(sample_videos_dir, exist_ok=True)
+
+    # Download each sample video
+    for filename, url in sample_videos_urls.items():
+        save_path = os.path.join(sample_videos_dir, filename)
+        # Download the video if it doesn't already exist
+        if not os.path.exists(save_path):
+            print(f"Downloading {filename}...")
+            download_video(url, save_path)
+        else:
+            print(f"{filename} already exists. Skipping download.")
+
+    csv.field_size_limit(100000000)
+    options = ['language', "sound-language", "sound"]
+    load_size = 224
+    plot_size = 224
+
+    video_input = gr.Video(label="Choose a video to featurize", height=480)
+    model_option = gr.Radio(options, value="language", label='Choose a model')
+
+    video_output1 = gr.Video(label="Audio Video Attention", height=480)
+    video_output2 = gr.Video(label="Multi-Head Audio Video Attention (Only Availible for sound_and_language)",
+                             height=480)
+    video_output3 = gr.Video(label="Visual Features", height=480)
+
+    models = {o: LitAVAligner.from_pretrained(f"mhamilton723/DenseAV-{o}") for o in options}
+
+
+    def process_video(video, model_option):
+        model = models[model_option].cuda()
+
+        original_frames, audio, info = torchvision.io.read_video(video, end_pts=10, pts_unit='sec')
+        sample_rate = 16000
+
+        if info["audio_fps"] != sample_rate:
+            audio = resample(audio, info["audio_fps"], sample_rate)
+        audio = audio[0].unsqueeze(0)
+
+        img_transform = T.Compose([
+            T.Resize(load_size, Image.BILINEAR),
+            lambda x: crop_to_divisor(x, 8),
+            lambda x: x.to(torch.float32) / 255,
+            norm])
+
+        frames = torch.cat([img_transform(f.permute(2, 0, 1)).unsqueeze(0) for f in original_frames], axis=0)
+
+        plotting_img_transform = T.Compose([
+            T.Resize(plot_size, Image.BILINEAR),
+            lambda x: crop_to_divisor(x, 8),
+            lambda x: x.to(torch.float32) / 255])
+
+        frames_to_plot = plotting_img_transform(original_frames.permute(0, 3, 1, 2))
+
+        with torch.no_grad():
+            audio_feats = model.forward_audio({"audio": audio.cuda()})
+            audio_feats = {k: v.cpu() for k, v in audio_feats.items()}
+            image_feats = model.forward_image({"frames": frames.unsqueeze(0).cuda()}, max_batch_size=2)
+            image_feats = {k: v.cpu() for k, v in image_feats.items()}
+
+            sim_by_head = model.sim_agg.get_pairwise_sims(
+                {**image_feats, **audio_feats},
+                raw=False,
+                agg_sim=False,
+                agg_heads=False
+            ).mean(dim=-2).cpu()
+
+            sim_by_head = blur_dim(sim_by_head, window=3, dim=-1)
+            print(sim_by_head.shape)
+
+        temp_video_path_1 = tempfile.mktemp(suffix='.mp4')
+
+        plot_attention_video(
+            sim_by_head,
+            frames_to_plot,
+            audio,
+            info["video_fps"],
+            sample_rate,
+            temp_video_path_1)
+
+        if model_option == "sound_and_language":
+            temp_video_path_2 = tempfile.mktemp(suffix='.mp4')
+
+            plot_2head_attention_video(
+                sim_by_head,
+                frames_to_plot,
+                audio,
+                info["video_fps"],
+                sample_rate,
+                temp_video_path_2)
+
+        else:
+            temp_video_path_2 = None
+
+        temp_video_path_3 = tempfile.mktemp(suffix='.mp4')
+        temp_video_path_4 = tempfile.mktemp(suffix='.mp4')
+
+        plot_feature_video(
+            image_feats["image_feats"].cpu(),
+            audio_feats['audio_feats'].cpu(),
+            frames_to_plot,
+            audio,
+            info["video_fps"],
+            sample_rate,
+            temp_video_path_3,
+            temp_video_path_4,
+        )
+        # return temp_video_path_1, temp_video_path_2, temp_video_path_3, temp_video_path_4
+
+        return temp_video_path_1, temp_video_path_2, temp_video_path_3
+
+
+    with gr.Blocks() as demo:
+        with gr.Column():
+            gr.Markdown("## Visualizing Sound and Language with DenseAV")
+            gr.Markdown(
+                "This demo allows you to explore the inner attention maps of DenseAV's dense multi-head contrastive operator.")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    model_option.render()
+                with gr.Column(scale=3):
+                    video_input.render()
+            with gr.Row():
+                submit_button = gr.Button("Submit")
+            with gr.Row():
+                gr.Examples(
+                    examples=[
+                        [join(sample_videos_dir, "puppies.mp4"), "sound_and_language"],
+                        [join(sample_videos_dir, "peppers.mp4"), "language"],
+                        [join(sample_videos_dir, "elephant2.mp4"), "language"],
+                        [join(sample_videos_dir, "boat.mp4"), "language"]
+
+                    ],
+                    inputs=[video_input, model_option]
+                )
+            with gr.Row():
+                video_output1.render()
+                video_output2.render()
+                video_output3.render()
+
+        submit_button.click(fn=process_video, inputs=[video_input, model_option],
+                            outputs=[video_output1, video_output2, video_output3])
+
+
+    if mode == "local":
+        demo.launch(server_name="0.0.0.0", server_port=6006, debug=True)
+    else:
+        demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)

DenseAV/hubconf.py

@@ -0,0 +1,25 @@
+# hubconf.py
+from denseav.train import LitAVAligner
+
+dependencies = ['torch', 'torchvision', 'PIL', 'denseav']  # List any dependencies here
+
+
+def _load_base(model_name):
+    model = LitAVAligner.load_from_checkpoint(
+        f"https://marhamilresearch4.blob.core.windows.net/denseav-public/hub/{model_name}.ckpt",
+        **{'loss_leak': 0.0, 'use_cached_embs': False},
+        strict=True)
+    model.set_full_train(True)
+    return model
+
+
+def sound_and_language():
+    return _load_base("denseav_2head")
+
+
+def language():
+    return _load_base("denseav_language")
+
+
+def sound():
+    return _load_base("denseav_sound")

DenseAV/samples/puppies.mp4

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

DenseAV/setup.py

@@ -0,0 +1,37 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='denseav',
+    version='0.1.0',
+    packages=find_packages(),
+    install_requires=[
+        'torch',
+        'kornia',
+        'omegaconf',
+        'pytorch-lightning',
+        'torchvision',
+        'tqdm',
+        'torchmetrics',
+        'scikit-learn',
+        'numpy',
+        'matplotlib',
+        'timm==0.4.12',
+        'moviepy',
+        'hydra-core',
+        'peft==0.5.0',
+        'av',
+        'audioread'
+    ],
+    author='Mark Hamilton',
+    author_email='[email protected]',
+    description='Offical code for the CVPR 2024 Paper: Separating the "Chirp" from the "Chat": Self-supervised Visual Grounding of Sound and Language',
+    long_description=open('README.md').read(),
+    long_description_content_type='text/markdown',
+    url='https://github.com/mhamilton723/DenseAV',
+    classifiers=[
+        'Programming Language :: Python :: 3',
+        'License :: OSI Approved :: MIT License',
+        'Operating System :: OS Independent',
+    ],
+    python_requires='>=3.6'
+)