Upload 32 files

DenseAV/denseav/aggregators.py

@@ -6,7 +6,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from tqdm import tqdm
 
-from constants import *
+from DenseAV.denseav.constants import *
 
 
 @torch.jit.script

DenseAV/denseav/aligners.py

@@ -4,7 +4,7 @@ import torch
 import torch.nn.functional as F
 from torch.nn import ModuleList
 
-from featurizers.DINO import Block
+from DenseAV.denseav.featurizers.DINO import Block
 
 
 class ChannelNorm(torch.nn.Module):

DenseAV/denseav/eval_utils.py

@@ -9,7 +9,7 @@ from torchmetrics.functional.classification import binary_average_precision
 from tqdm import tqdm
 
 from constants import *
-from shared import unnorm, remove_axes
+from DenseAV.denseav.shared import unnorm, remove_axes
 
 
 def prep_heatmap(sims, masks, h, w):

DenseAV/denseav/evaluate.py

@@ -4,8 +4,8 @@ from omegaconf import DictConfig, OmegaConf
 from pytorch_lightning import Trainer
 from pytorch_lightning import seed_everything
 from pytorch_lightning.loggers import TensorBoardLogger
-from data.AVDatasets import AVDataModule
-from shared import load_trained_model
+from DenseAV.denseav.data.AVDatasets import AVDataModule
+from DenseAV.denseav.shared import load_trained_model
 
 
 @hydra.main(config_path="configs", config_name="av_align.yaml")

DenseAV/denseav/plotting.py

@@ -10,7 +10,7 @@ import torch.nn.functional as F
 import torchvision
 from moviepy.editor import VideoFileClip, AudioFileClip
 from base64 import b64encode
-from shared import pca
+from DenseAV.denseav.shared import pca
 
 
 def write_video_with_audio(video_frames, audio_array, video_fps, audio_fps, output_path):

DenseAV/denseav/shared.py

@@ -90,37 +90,37 @@ def get_image_featurizer(name, token_type="key", **kwargs):
     name = name.lower()
 
     if name == "vit":
-        from featurizers.DINO import DINOFeaturizer
+        from DenseAV.denseav.featurizers.DINO import DINOFeaturizer
         patch_size = 16
         model = DINOFeaturizer("vit_small_patch16_224", patch_size, token_type)
         dim = 384
     elif name == "dino16":
-        from featurizers.DINO import DINOFeaturizer
+        from DenseAV.denseav.featurizers.DINO import DINOFeaturizer
         patch_size = 16
         model = DINOFeaturizer("dino_vits16", patch_size, token_type)
         dim = 384
     elif name == "dino8":
-        from featurizers.DINO import DINOFeaturizer
+        from DenseAV.denseav.featurizers.DINO import DINOFeaturizer
         patch_size = 8
         model = DINOFeaturizer("dino_vits8", patch_size, token_type)
         dim = 384
     elif name == "clip":
-        from featurizers.CLIP import CLIPFeaturizer
+        from DenseAV.denseav.featurizers.CLIP import CLIPFeaturizer
         patch_size = 16
         model = CLIPFeaturizer()
         dim = 512
     elif name == "cavmae":
-        from featurizers.CAVMAE import CAVMAEImageFeaturizer
+        from DenseAV.denseav.featurizers.CAVMAE import CAVMAEImageFeaturizer
         model = CAVMAEImageFeaturizer(kwargs["output_root"], model=kwargs.get("model"))
         dim = 768
         patch_size = 16
     elif name == "fnac":
-        from featurizers.FNACAVL import FNACImageFeaturizer
+        from DenseAV.denseav.featurizers.FNACAVL import FNACImageFeaturizer
         model = FNACImageFeaturizer(kwargs["output_root"], model=kwargs.get("model"))
         dim = 512
         patch_size = 16
     elif name == "imagebind":
-        from featurizers.ImageBind import ImageBindImageFeaturizer
+        from DenseAV.denseav.featurizers.ImageBind import ImageBindImageFeaturizer
         model = ImageBindImageFeaturizer(kwargs["output_root"], model=kwargs.get("model"))
         dim = 1024
         patch_size = 16
@@ -131,12 +131,12 @@ def get_image_featurizer(name, token_type="key", **kwargs):
         patch_size = 1
         dim = 2048
     elif name == "davenet":
-        from featurizers.DAVENet import DavenetImageFeaturizer
+        from fDenseAV.denseav.eaturizers.DAVENet import DavenetImageFeaturizer
         model = DavenetImageFeaturizer()
         patch_size = 1
         dim = 1024
     elif name == "dinov2":
-        from featurizers.DINOv2 import DINOv2Featurizer
+        from DenseAV.denseav.featurizers.DINOv2 import DINOv2Featurizer
         model = DINOv2Featurizer()
         patch_size = 14
         dim = 768
@@ -147,29 +147,29 @@ def get_image_featurizer(name, token_type="key", **kwargs):
 
 def get_audio_featurizer(name, **kwargs):
     if name == "davenet":
-        from featurizers.DAVENet import DavenetAudioFeaturizer
+        from DenseAV.denseav.featurizers.DAVENet import DavenetAudioFeaturizer
         model = DavenetAudioFeaturizer()
         dim = 1024
     elif name == "dino8":
         model, _, dim = get_image_featurizer("dino8")
     elif name == "hubert":
-        from featurizers.Hubert import Hubert
+        from DenseAV.denseav.featurizers.Hubert import Hubert
         model = Hubert()
         dim = 1024
     elif name == "cavmae":
-        from featurizers.CAVMAE import CAVMAEAudioFeaturizer
+        from DenseAV.denseav.featurizers.CAVMAE import CAVMAEAudioFeaturizer
         model = CAVMAEAudioFeaturizer(kwargs["output_root"], model=kwargs.get("model"))
         dim = 768
     elif name == "imagebind":
-        from featurizers.ImageBind import ImageBindAudioFeaturizer
+        from DenseAV.denseav.featurizers.ImageBind import ImageBindAudioFeaturizer
         model = ImageBindAudioFeaturizer(kwargs["output_root"], model=kwargs.get("model"))
         dim = 1024
     elif name == "audiomae":
-        from featurizers.AudioMAE import AudioMAE
+        from DenseAV.denseav.featurizers.AudioMAE import AudioMAE
         model = AudioMAE(kwargs["output_root"], False)
         dim = 768
     elif name == "audiomae-finetuned":
-        from featurizers.AudioMAE import AudioMAE
+        from DenseAV.denseav.featurizers.AudioMAE import AudioMAE
         model = AudioMAE(kwargs["output_root"], True)
         dim = 768
     else:

DenseAV/denseav/train.py

@@ -1,1222 +1,1222 @@
-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.denseav.aggregators import get_aggregator
-from aligners import get_aligner, ProgressiveGrowing
-from constants import *
-from data.AVDatasets import AVDataModule
-from 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
-
-        self.full_train = False # Added by me
-
-        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:
-        print("---"*10)
-        print("Setup with full_train = False")
-        run_exp(aligner, full_train=False)
-        print("---"*10)
-    else:
-        print("---"*10)
-        print("Setup with full_train = False")
-        run_exp(aligner, full_train=True)
-        print("---"*10)
-
-
-if __name__ == "__main__":
-    my_app()
+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.denseav.aggregators import get_aggregator
+from DenseAV.denseav.aligners import get_aligner, ProgressiveGrowing
+from DenseAV.denseav.constants import *
+from DenseAV.denseav.data.AVDatasets import AVDataModule
+from DenseAV.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
+
+        self.full_train = False # Added by me
+
+        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:
+        print("---"*10)
+        print("Setup with full_train = False")
+        run_exp(aligner, full_train=False)
+        print("---"*10)
+    else:
+        print("---"*10)
+        print("Setup with full_train = False")
+        run_exp(aligner, full_train=True)
+        print("---"*10)
+
+
+if __name__ == "__main__":
+    my_app()