Upload model

audio_extractor.py

@@ -0,0 +1,235 @@
+from math import prod
+
+import torch
+from torch import nn
+
+from einops.layers.torch import Rearrange
+from einops import rearrange
+
+from typing import List, Optional
+
+from abc import ABC, abstractmethod
+
+class Extractor(ABC):
+    """Abstract base class for encoders."""
+    
+    # Just declare that implementers should have this attribute
+    embedding_dim: int
+    
+    @abstractmethod
+    def forward(self, x : torch.Tensor) -> torch.Tensor:
+        """Forward pass through the encoder."""
+        pass
+    
+    @abstractmethod
+    def total_patches(self, time: int) -> int:
+        """Returns the total patches given the time dimension of the input."""
+        pass
+
+
+
+class ConvChannelFeatureExtractor(Extractor, nn.Module):
+    """
+    Convolutional feature encoder for the audio data.
+
+    Computes successive 1D convolutions (with activations) over the time
+    dimension of the audio signal. This encoder also uses different kernels for each time signal. 
+    Therefore, in_channels argument is necessary!
+
+    Inspiration from https://github.com/facebookresearch/fairseq/blob/main/fairseq/models/wav2vec/wav2vec2.py
+    and https://github.com/SPOClab-ca/BENDR/blob/main/dn3_ext.py
+
+    Args:
+        conv_layers_spec: list of tuples (dim, k, stride) where:
+            * dim: number of output channels of the layer (unrelated to EEG channels);
+            * k: temporal length of the layer's kernel;
+            * stride: temporal stride of the layer's kernel.
+
+        in_channels: int
+            Number of audio channels.
+        dropout: float
+        mode: str
+            Normalisation mode. Either``default`` or ``layer_norm``.
+        conv_bias: bool
+        depthwise: bool 
+            Perform depthwise convolutions rather than the full convolution.
+    """
+
+    def __init__(
+        self,
+        *args,
+        conv_layers_spec: list[tuple[int, int, int]],
+        in_channels : int = 2,
+        dropout: float = 0.0,
+        mode: str = "default",
+        conv_bias: bool = False,
+        depthwise : bool = False,
+        share_weights_over_channels : bool = False,
+        **kwargs,
+    ):
+        assert mode in {"default", "layer_norm"}
+        super().__init__() # type: ignore
+
+        def block(
+            n_in : int,
+            n_out : int,
+            k : int,
+            stride : int,
+            is_layer_norm : bool =False,
+            is_group_norm : bool =False,
+            conv_bias : bool =False,
+            depthwise : bool = True,
+        ):
+
+            def make_conv():
+                if depthwise:
+                    assert n_out % n_in == 0, f"For depthwise signals we can not have non-multipler of {n_out} and {n_in}"
+                    conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias, groups = n_in)
+                else:
+                    conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias)
+
+                nn.init.kaiming_normal_(conv.weight)
+                return conv
+
+            assert not (
+                is_layer_norm and is_group_norm
+            ), "layer norm and group norm are exclusive"
+
+            if is_layer_norm:
+                return nn.Sequential(
+                    make_conv(),
+                    nn.Dropout(p=dropout),
+                    nn.Sequential(
+                        Rearrange("... channels time -> ... time channels"),
+                        nn.LayerNorm(n_out, elementwise_affine=True),  # Fixed: use n_out instead of dim
+                        Rearrange("... time channels -> ... channels time"),
+                    ),
+                    nn.GELU(),
+                )
+            elif is_group_norm:
+                return nn.Sequential(
+                    make_conv(),
+                    nn.Dropout(p=dropout),
+                    nn.GroupNorm(n_out, n_out, affine=True),  # Fixed: use n_out instead of dim
+                    nn.GELU(),
+                )
+            else:
+                return nn.Sequential(make_conv(), nn.Dropout(p=dropout), nn.GELU())
+
+        self.in_channels = in_channels
+        self.depthwise = depthwise
+        self.conv_layers_spec = conv_layers_spec
+        self.cnns = nn.ModuleList()
+
+        if share_weights_over_channels:
+            in_d = 1
+            conv_layers = []
+            for i, cl in enumerate(conv_layers_spec):
+                assert len(cl) == 3, "invalid conv definition: " + str(cl)
+                (dim, k, stride) = cl
+                conv_layers.append( # type: ignore
+                    block(
+                        in_d,
+                        dim,
+                        k,
+                        stride,
+                        is_layer_norm=mode == "layer_norm",
+                        is_group_norm=mode == "default" and i == 0,
+                        conv_bias=conv_bias,
+                        depthwise=self.depthwise
+                    )
+                )
+                in_d = dim
+            cnn : nn.Module = nn.Sequential(*conv_layers) # type: ignore
+            self.embedding_dim = conv_layers_spec[-1][0]
+            self.cnns.append(cnn)
+        else:
+            for channels in range(self.in_channels):
+                in_d = 1
+                conv_layers = []
+                for i, cl in enumerate(conv_layers_spec):
+                    assert len(cl) == 3, "invalid conv definition: " + str(cl)
+                    (dim, k, stride) = cl
+                    conv_layers.append( # type: ignore
+                        block(
+                            in_d,
+                            dim,
+                            k,
+                            stride,
+                            is_layer_norm=mode == "layer_norm",
+                            is_group_norm=mode == "default" and i == 0,
+                            conv_bias=conv_bias,
+                            depthwise=self.depthwise
+                        )
+                    )
+                    in_d = dim
+                cnn : nn.Module = nn.Sequential(*conv_layers) # type: ignore
+                self.cnns.append(cnn)
+        
+        self.embedding_dim = self.conv_layers_spec[-1][0]
+        self.weight_sharing = share_weights_over_channels
+
+    def forward(self, x : torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: (batch_size, n_chans, n_times)
+                    Batched EEG signal.
+
+        Returns:
+            local_features: (batch_size, emb_dim, n_times_out)
+                Local features extracted from the audio signal.
+                ``emb_dim`` corresponds to the ``dim`` of the last element of
+                ``conv_layers_spec``.
+        """
+
+        out = []
+        for channel_index in range(self.in_channels):
+            # If we are sharing weights over the channels, use one CNN for all the channel dimensions
+            if self.weight_sharing:
+                module = self.cnns[0]
+            else:
+                module = self.cnns[channel_index]
+            processed = module(x[:, [channel_index], ...])
+            processed = rearrange(processed, "batch_size n_channels n_time -> batch_size n_time n_channels")
+            out.append(processed)
+        processed = torch.stack(out, dim = 1)
+        processed = torch.flatten(processed, start_dim = 1, end_dim = 2)
+        return processed
+
+    def total_patches(self, time: int) -> int:
+        """Calculate the number of output time steps for a given input length."""
+        x = torch.zeros((1, self.in_channels, time))
+        processed = self.forward(x)
+        return processed.shape[1]  # Return time dimension size
+
+    @property
+    def receptive_fields(self) -> List[int]:
+        rf = 1
+        receptive_fields = [rf]
+        for _, width, stride in reversed(self.conv_layers_spec):
+            rf = (rf - 1) * stride + width  # assumes no padding and no dilation
+            receptive_fields.append(rf)
+        return list(reversed(receptive_fields))
+
+    def description(self, sfreq : Optional[int] = None, dummy_time : Optional[int] = None) -> str:
+        dims, _, strides = zip(*self.conv_layers_spec)
+        receptive_fields = self.receptive_fields
+        rf = receptive_fields[0]
+        desc = f"Receptive field: {rf} samples"
+        if sfreq is not None:
+            desc += f", {rf / sfreq:.2f} seconds"
+
+        ds_factor = prod(strides)
+        desc += f" | Downsampled by {ds_factor}"
+        if sfreq is not None:
+            desc += f", new sfreq: {sfreq / ds_factor:.2f} Hz"
+        desc += f" | Overlap of {rf - ds_factor} samples"
+        if dummy_time is not None:
+            n_times_out = self.total_patches(dummy_time)
+            desc += f" | {n_times_out} encoded samples/trial"
+
+        n_features = [
+            f"{dim}*{rf}" for dim, rf in zip([self.in_channels] + list(dims), receptive_fields)
+        ]
+        desc += f" | #features/sample at each layer (n_channels*n_times): [{', '.join(n_features)}] = {[eval(x) for x in n_features]}"
+        return desc

config.json

@@ -0,0 +1,53 @@
+{
+  "architectures": [
+    "WavJEPANatModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_wavjepa_nat.WavJEPANatConfig",
+    "AutoModel": "modeling_wavjepa_nat.WavJEPANatModel"
+  },
+  "decoder_cfg": {
+    "enable_nested_tensor": false,
+    "mask_check": true,
+    "num_layers": 12
+  },
+  "decoder_layers_cfg": {
+    "activation": "gelu",
+    "batch_first": true,
+    "bias": true,
+    "d_model": 384,
+    "dim_feedforward": 1536,
+    "dropout": 0.0,
+    "layer_norm_eps": 1e-06,
+    "nhead": 12,
+    "norm_first": true
+  },
+  "dtype": "float32",
+  "encoder_cfg": {
+    "enable_nested_tensor": false,
+    "mask_check": true,
+    "num_layers": 12
+  },
+  "encoder_layers_cfg": {
+    "activation": "gelu",
+    "batch_first": true,
+    "bias": true,
+    "d_model": 768,
+    "dim_feedforward": 3072,
+    "dropout": 0.0,
+    "layer_norm_eps": 1e-06,
+    "nhead": 12,
+    "norm_first": true
+  },
+  "extractor_config": {
+    "conv_bias": false,
+    "conv_layers_spec": "[(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512,2,2)]",
+    "depthwise": false,
+    "dropout": 0.0,
+    "in_channels": 2,
+    "mode": "default",
+    "share_weights_over_channels": false
+  },
+  "model_type": "wavjepa-nat-base",
+  "transformers_version": "4.57.1"
+}

configuration_wavjepa_nat.py

@@ -0,0 +1,83 @@
+from transformers import PretrainedConfig
+from .types import TransformerLayerCFG, TransformerEncoderCFG
+
+
+
+class WavJEPANatConfig(PretrainedConfig):
+    model_type = "wavjepa-nat-base"
+    model_size = "base"
+    in_channels: int = 2
+
+    def __init__(
+        self,
+        extractor_layers_spec: str = "[(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512,2,2)]",
+        extractor_dropout : float = 0.0,
+        extractor_mode : str = "default", 
+        extractor_conv_bias : bool = False, 
+        extractor_depthwise: bool = False,
+        encoder_d_model: int = 768,
+        encoder_nhead : int = 12,
+        encoder_batch_first = True,
+        encoder_norm_first = True, 
+        encoder_bias = True,
+        encoder_mlp_ratio = 4.0,
+        encoder_dropout = 0.0,
+        encoder_num_layers: int = 12,
+        encoder_enable_nested_tensor = False, 
+        encoder_mask_check = True,
+        encoder_activation = 'gelu',
+        decoder_d_model: int = 384,
+        decoder_nhead : int = 12,
+        decoder_batch_first = True,
+        decoder_norm_first = True, 
+        decoder_bias = True,
+        decoder_mlp_ratio = 4.0,
+        decoder_dropout = 0.0,
+        decoder_num_layers: int = 12,
+        decoder_enable_nested_tensor = False, 
+        decoder_mask_check = True,
+        decoder_activation = 'gelu',
+        **kwargs
+    ):
+        self.encoder_cfg = TransformerEncoderCFG.create(
+                        num_layers = encoder_num_layers,
+                        enable_nested_tensor = encoder_enable_nested_tensor,
+                        mask_check = encoder_mask_check,
+        )
+        self.decoder_cfg = TransformerEncoderCFG.create(
+                        num_layers = decoder_num_layers,
+                        enable_nested_tensor = decoder_enable_nested_tensor,
+                        mask_check = decoder_mask_check,
+        )
+        self.encoder_layers_cfg = TransformerLayerCFG.create(
+                        d_model = encoder_d_model,
+                        nhead = encoder_nhead,
+                        batch_first = encoder_batch_first,
+                        norm_first = encoder_norm_first,
+                        bias = encoder_bias,
+                        mlp_ratio = encoder_mlp_ratio,
+                        dropout = encoder_dropout,
+                        activation = encoder_activation,
+                        layer_norm_eps = 1e-6
+        )
+        self.decoder_layers_cfg = TransformerLayerCFG.create(
+                        d_model = decoder_d_model,
+                        nhead = decoder_nhead,
+                        batch_first = decoder_batch_first,
+                        norm_first = decoder_norm_first,
+                        bias = decoder_bias,
+                        mlp_ratio = decoder_mlp_ratio,
+                        dropout = decoder_dropout,
+                        activation = decoder_activation,
+                        layer_norm_eps = 1e-6
+        )
+        self.extractor_config = dict(
+                conv_layers_spec = extractor_layers_spec,
+                in_channels = self.in_channels,
+                dropout = extractor_dropout,
+                mode = extractor_mode,
+                conv_bias = extractor_conv_bias,
+                depthwise = extractor_depthwise,
+                share_weights_over_channels=False)
+
+        super().__init__(**kwargs)

model.py

@@ -0,0 +1,181 @@
+import copy
+import numpy as np 
+
+from typing import Any, Optional
+
+import torch
+from torch import nn
+
+from einops import repeat, rearrange
+
+from .pos_embed import get_1d_sincos_pos_embed_from_grid, get_2d_sincos_pos_embed, get_binaural_pos_embed
+from .audio_extractor import Extractor
+from .types import TransformerLayerCFG, TransformerEncoderCFG
+from .utils import normalize, calculate_padding_mask, get_timestamps
+
+class WavJEPANat(nn.Module):
+    """
+    Joint-Embedding Predictive Architecture (JEPA).
+
+    This implementation is inspired by:
+        * I-JEPA http://arxiv.org/abs/2301.08243
+        * Data2vec 2.0 http://arxiv.org/abs/2212.07525
+    """
+
+    teacher_encoder: nn.Module
+    sample_rate : int = 16000
+    process_audio_seconds : float = 2.01
+    in_channels : int = 2
+
+    
+    def __init__(
+        self,
+        feature_extractor: Extractor,
+        transformer_encoder_layers_cfg : TransformerLayerCFG,
+        transformer_encoder_cfg : TransformerEncoderCFG,
+        transformer_decoder_layers_cfg : TransformerLayerCFG,
+        transformer_decoder_cfg : TransformerEncoderCFG,
+        size : str = "base",
+        **kwargs : dict[str, Any],
+    ):
+        super().__init__(**kwargs)
+    
+        self.is_spectrogram = False
+        self.target_length = int(self.sample_rate * self.process_audio_seconds)
+        self.extract_audio = feature_extractor
+        self.total_patches = 400
+        self.output_steps = self.total_patches // self.in_channels
+        self.feature_norms : nn.Module = nn.LayerNorm(self.extract_audio.embedding_dim)
+
+        self.n_encoder_heads = transformer_encoder_layers_cfg["nhead"]
+        self.encoder_embedding_dim = transformer_encoder_layers_cfg["d_model"]
+        self.n_decoder_heads = transformer_decoder_layers_cfg["nhead"]
+        self.decoder_embedding_dim = transformer_decoder_layers_cfg["d_model"]
+
+        encoder_layer = nn.TransformerEncoderLayer(**transformer_encoder_layers_cfg)
+        self.encoder = nn.TransformerEncoder(encoder_layer, norm = nn.LayerNorm(self.encoder_embedding_dim), **transformer_encoder_cfg)
+        self.post_extraction_mapper : Optional[nn.Module] = nn.Linear(feature_extractor.embedding_dim, self.encoder_embedding_dim) if feature_extractor.embedding_dim != self.encoder_embedding_dim else None
+        decoder_layer = nn.TransformerEncoderLayer(**transformer_decoder_layers_cfg)
+        self.decoder = nn.TransformerEncoder(decoder_layer, norm = nn.LayerNorm(self.decoder_embedding_dim), **transformer_decoder_cfg)
+        self.decoder_to_encoder_mapper = nn.Linear(self.decoder_embedding_dim, self.encoder_embedding_dim, bias=True)
+        self.encoder_to_decoder_mapper = nn.Linear(self.encoder_embedding_dim, self.decoder_embedding_dim)
+
+        # For the autocast add batch dimensions.
+        self.mask_token = nn.Parameter(
+            torch.zeros(1, 1, self.decoder_embedding_dim, requires_grad=True)
+        )
+        torch.nn.init.normal_(self.mask_token, std=0.02)
+        self.pos_encoding_encoder = self._get_pos_embed_params(self.encoder_embedding_dim)
+        self.pos_encoding_decoder = self._get_pos_embed_params(self.decoder_embedding_dim)
+        self._init_teacher()
+
+
+    def _get_pos_embed_params(self, embedding_dim):
+        """Calculates the pos embedding embedding parameters and returns them."""
+        # Update positional embedding
+        pos_embed = nn.Parameter(
+            torch.zeros(
+                1,
+                self.total_patches,
+                embedding_dim,
+            ),
+            requires_grad=False,
+        )
+        positions = np.arange(self.total_patches, dtype=np.float64)
+        if self.is_spectrogram:
+            # If it is a spectrogram, we use 2d sincos embeddings.
+            pos_embed_data = get_2d_sincos_pos_embed(
+                embedding_dim, self.extract_audio.grid_size, cls_token_num=0
+            )
+        #TODO! Remove this total patches later.
+        elif not self.is_spectrogram and self.in_channels == 2 and (self.total_patches == 400):
+            # We use 1D sincos embeddings with channel number indicated on the last 384 dimensions.
+            pos_embed_data = get_binaural_pos_embed(embedding_dim, time_steps=self.total_patches // self.in_channels)
+        elif not self.is_spectrogram and self.in_channels == 2 and (self.total_patches == 200):
+            #Use 1D pos_embeddings if channel-mixing feature extractor
+            pos_embed_data = get_1d_sincos_pos_embed_from_grid(
+                embedding_dim,
+                positions,
+            )     
+        elif not self.is_spectrogram and self.in_channels == 1 and (self.total_patches == 200):
+            # IF it is plain audio, we used 1d sincos embeddings
+            pos_embed_data = get_1d_sincos_pos_embed_from_grid(
+                embedding_dim,
+                positions,
+            )
+        else:
+            raise Exception(f"Not implemented for more in_channels, {self.in_channels}, {self.total_patches}")
+        pos_embed.data.copy_(torch.from_numpy(pos_embed_data).float().unsqueeze(0))
+        return pos_embed
+
+    def _init_teacher(self):
+        self.teacher_encoder = copy.deepcopy(self.encoder)
+        self.teacher_encoder.requires_grad_(False)
+
+
+
+    @torch.inference_mode()
+    def _get_segment_representation(self, audio : torch.Tensor, padding_mask : torch.tensor):
+        # Get the audio representatin of waveform x.
+        self.eval()
+        local_features = self.extract_audio(audio)
+        local_features = self.feature_norms(local_features)
+        if self.post_extraction_mapper:
+            local_features = self.post_extraction_mapper(local_features)
+        local_features = local_features + self.pos_encoding_encoder
+        # Encoder and decoder forward
+        contextual_features = self.encoder(local_features, src_key_padding_mask = padding_mask)
+        return contextual_features
+
+    @torch.inference_mode()
+    def get_audio_representation(self, audio : torch.Tensor):
+        B = audio.shape[0]
+        input_audio_len = audio.shape[-1]
+        # Assert audio is of correct shape
+        if audio.ndim != 3:
+            raise ValueError(
+                "audio input tensor must be 2D with shape (n_sounds, n_channels, num_samples)"
+            )
+        cur_frames = audio.shape[-1]
+        pad_frames = self.target_length - (cur_frames % self.target_length)
+        
+        if pad_frames > 0:
+            # Padding with constant 0s
+            pad_arg = (
+                0,
+                pad_frames,
+            )  # (channel, channel, height, height, width, width)
+            audio = torch.nn.functional.pad(audio, pad_arg, mode="constant")
+        embeddings = []
+        padding_mask, cut_off = calculate_padding_mask(pad_frames = pad_frames, 
+                                        total_frames = audio.shape[-1], 
+                                        sr = self.sample_rate,
+                                        output_steps = self.output_steps,
+                                        process_seconds = self.target_length // self.sample_rate, 
+                                        device = audio.device, 
+                                        B = B)
+        mask_idx = 0
+        masked_mean = torch.zeros(audio.shape, dtype = torch.bool)
+        masked_mean[..., cur_frames:] = True
+        mt = torch.masked.masked_tensor(audio, masked_mean)
+        # Now get the embeddings o the model.
+        for i in range(audio.shape[-1] // self.target_length):
+            mt = audio[..., i * self.target_length : (i + 1) * self.target_length]
+            mask = padding_mask[..., mask_idx : mask_idx + self.output_steps]
+            with torch.no_grad():
+                # We do not include padding tokens in the mean and std calculation.
+                mask = repeat(mask, "B E -> B (C E)", C = self.in_channels)
+                embedding = self._get_segment_representation(
+                    normalize(mt),
+                    mask
+                )
+                embedding = rearrange(embedding, "B (C S) E -> B C S E", C = self.in_channels)
+            mask_idx = mask_idx + self.output_steps
+            embeddings.append(embedding)
+        
+
+        x = torch.concatenate(embeddings, axis = 2)
+        x = x[:, :, :cut_off, :]
+        ts = get_timestamps(self.sample_rate, B, input_audio_len, x)
+        assert ts.shape[-1] == x.shape[2]
+        return x, ts 

model.safetensors

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

modeling_wavjepa_nat.py

@@ -0,0 +1,34 @@
+from transformers import PreTrainedModel
+
+from .model import WavJEPANat
+from .configuration_wavjepa_nat import WavJEPANatConfig
+from .audio_extractor import ConvChannelFeatureExtractor
+import torch 
+from typing import Union 
+
+class WavJEPANatModel(PreTrainedModel):
+    config_class = WavJEPANatConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.model = WavJEPANat(
+                feature_extractor = ConvChannelFeatureExtractor(
+                    conv_layers_spec = eval(config.extractor_config['conv_layers_spec']),
+                    in_channels = config.extractor_config['in_channels'],
+                    dropout = config.extractor_config['dropout'],
+                    mode = config.extractor_config['mode'],
+                    conv_bias = config.extractor_config['conv_bias'],
+                    depthwise = config.extractor_config['depthwise'],
+                    share_weights_over_channels = config.extractor_config['share_weights_over_channels']
+                    ),
+                transformer_encoder_layers_cfg = config.encoder_layers_cfg,
+                transformer_encoder_cfg = config.encoder_cfg,
+                transformer_decoder_layers_cfg = config.decoder_layers_cfg,
+                transformer_decoder_cfg = config.decoder_cfg,
+                size = config.model_size,
+        )
+
+    def forward(self, tensor) -> Union[torch.Tensor, torch.Tensor]:
+        return self.model.get_audio_representation(tensor)
+

pos_embed.py

@@ -0,0 +1,267 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# Position embedding utils
+# --------------------------------------------------------
+
+
+# https://github.com/facebookresearch/AudioMAE/blob/main/util/pos_embed.py
+import numpy as np
+import torch
+
+
+# --------------------------------------------------------
+# 2D sine-cosine position embedding
+# References:
+# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
+# MoCo v3: https://github.com/facebookresearch/moco-v3
+# --------------------------------------------------------
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token_num):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    if grid_size is int:
+        gH = grid_size
+        gW = grid_size
+    else:
+        gH = grid_size[0]
+        gW = grid_size[1]
+    grid_h = np.arange(gH, dtype=np.float64)
+    grid_w = np.arange(gW, dtype=np.float64)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, gH, gW])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    for _ in range(cls_token_num):
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_flexible(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size[0], dtype=np.float64)
+    grid_w = np.arange(grid_size[1], dtype=np.float64)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size[0], grid_size[1]])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed(embed_dim, length):
+    """
+    Create 1D sinusoidal positional embeddings.
+    
+    Args:
+        embed_dim: embedding dimension
+        length: sequence length
+    
+    Returns:
+        pos_embed: [length, embed_dim]
+    """
+    assert embed_dim % 2 == 0
+    
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+    
+    pos = np.arange(length, dtype=np.float64)  # (length,)
+    out = np.einsum("m,d->md", pos, omega)  # (length, D/2)
+    
+    emb_sin = np.sin(out)  # (length, D/2)
+    emb_cos = np.cos(out)  # (length, D/2)
+    
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (length, D)
+    return emb
+
+def get_binaural_pos_embed(embed_dim, time_steps=100):
+    """
+    Create positional embeddings for binaural audio.
+    Same time encoding, different channel encoding.
+    
+    Args:
+        embed_dim: embedding dimension
+        time_steps: number of time steps per channel
+    
+    Returns:
+        pos_embed: [2*time_steps, embed_dim] - for concatenated L+R channels
+    """
+    assert embed_dim % 2 == 0
+    
+    # Time dimension encoding (same for both channels)
+    time_embed = get_1d_sincos_pos_embed(embed_dim // 2, time_steps)
+    
+    # Channel dimension encoding (different for L and R)
+    channel_embed_left = np.zeros((time_steps, embed_dim // 2))  # Left channel = 0
+    channel_embed_right = get_1d_sincos_pos_embed(embed_dim // 2, 1)  # Right channel = different
+    channel_embed_right = np.tile(channel_embed_right, (time_steps, 1))
+    
+    # Combine time and channel embeddings
+    left_pos_embed = np.concatenate([time_embed, channel_embed_left], axis=1)
+    right_pos_embed = np.concatenate([time_embed, channel_embed_right], axis=1)
+    
+    # Concatenate left and right channel embeddings
+    binaural_pos_embed = np.concatenate([left_pos_embed, right_pos_embed], axis=0)
+    
+    return binaural_pos_embed
+    
+# --------------------------------------------------------
+# Interpolate position embeddings for high-resolution
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+def interpolate_pos_embed(model, checkpoint_model):
+    if "pos_embed" in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model["pos_embed"]
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches**0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print(
+                "Position interpolate from %dx%d to %dx%d"
+                % (orig_size, orig_size, new_size, new_size)
+            )
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(
+                -1, orig_size, orig_size, embedding_size
+            ).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens,
+                size=(new_size, new_size),
+                mode="bicubic",
+                align_corners=False,
+            )
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model["pos_embed"] = new_pos_embed
+
+
+def interpolate_pos_embed_img2audio(model, checkpoint_model, orig_size, new_size):
+    if "pos_embed" in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model["pos_embed"]
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        # orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        # new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print(
+                "Position interpolate from %dx%d to %dx%d"
+                % (orig_size[0], orig_size[1], new_size[0], new_size[1])
+            )
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(
+                -1, orig_size[0], orig_size[1], embedding_size
+            ).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens,
+                size=(new_size[0], new_size[1]),
+                mode="bicubic",
+                align_corners=False,
+            )
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model["pos_embed"] = new_pos_embed
+
+
+def interpolate_pos_embed_audio(model, checkpoint_model, orig_size, new_size):
+    if "pos_embed" in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model["pos_embed"]
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        if orig_size != new_size:
+            print(
+                "Position interpolate from %dx%d to %dx%d"
+                % (orig_size[0], orig_size[1], new_size[0], new_size[1])
+            )
+            # extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            cls_token = pos_embed_checkpoint[:, 0, :].unsqueeze(1)
+            pos_tokens = pos_embed_checkpoint[:, 1:, :]  # remove
+            pos_tokens = pos_tokens.reshape(
+                -1, orig_size[0], orig_size[1], embedding_size
+            )  # .permute(0, 3, 1, 2)
+            # pos_tokens = torch.nn.functional.interpolate(
+            #    pos_tokens, size=(new_size[0], new_size[1]), mode='bicubic', align_corners=False)
+
+            # pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            pos_tokens = pos_tokens[:, :, : new_size[1], :]  # assume only time diff
+            pos_tokens = pos_tokens.flatten(1, 2)
+            new_pos_embed = torch.cat((cls_token, pos_tokens), dim=1)
+            checkpoint_model["pos_embed"] = new_pos_embed
+
+
+def interpolate_patch_embed_audio(
+    model,
+    checkpoint_model,
+    orig_channel,
+    new_channel=1,
+    kernel_size=(16, 16),
+    stride=(16, 16),
+    padding=(0, 0),
+):
+    if orig_channel != new_channel:
+        if "patch_embed.proj.weight" in checkpoint_model:
+            # aggregate 3 channels in rgb ckpt to 1 channel for audio
+            new_proj_weight = torch.nn.Parameter(
+                torch.sum(checkpoint_model["patch_embed.proj.weight"], dim=1).unsqueeze(
+                    1
+                )
+            )
+            checkpoint_model["patch_embed.proj.weight"] = new_proj_weight

types.py

@@ -0,0 +1,51 @@
+from typing import TypedDict
+from torch import nn
+
+
+class TransformerLayerCFG(TypedDict):
+    d_model : int
+    nhead : int
+    batch_first : bool
+    norm_first : bool
+    bias : bool
+    dim_feedforward : int
+    dropout : float
+    activation : nn.Module
+    layer_norm_eps : float
+
+    @classmethod
+    def create(cls, 
+    d_model : int = 768,
+    nhead : int = 12,
+    batch_first : bool = True,
+    norm_first : bool = False,
+    bias : bool = True,
+    mlp_ratio : float = 4.0,
+    dropout : float = 0.0,
+    activation : nn.Module = nn.GELU(),
+    layer_norm_eps : float = 1e-6) -> 'TransformerLayerCFG':
+        return TransformerLayerCFG(d_model = d_model, 
+                                   nhead = nhead, 
+                                   batch_first = batch_first,
+                                   norm_first = norm_first, 
+                                   bias = bias, 
+                                   dim_feedforward = int(d_model * mlp_ratio),
+                                   dropout = dropout, 
+                                   activation = activation, 
+                                   layer_norm_eps = layer_norm_eps)
+
+
+# Norm needs to be defined by the user!
+class TransformerEncoderCFG(TypedDict):
+    num_layers : int
+    enable_nested_tensor: bool
+    mask_check: bool
+
+    @classmethod
+    def create(cls, 
+    num_layers : int = 12, 
+    enable_nested_tensor: bool = False,
+    mask_check: bool = True) -> 'TransformerEncoderCFG':
+        return TransformerEncoderCFG(num_layers=num_layers, 
+                enable_nested_tensor = enable_nested_tensor,
+                mask_check = mask_check)

utils.py

@@ -0,0 +1,31 @@
+import torch 
+
+def get_timestamps(sample_rate, B, input_audio_len, x):
+    audio_len = input_audio_len
+    sec = audio_len / sample_rate
+    x_len = x.shape[2]
+    step = sec / x_len * 1000  # sec -> ms
+    ts = torch.tensor([step * i for i in range(x_len)]).unsqueeze(0)
+    ts = ts.repeat(B, 1)
+    return ts
+    
+def normalize(audio):
+    mean = audio.mean(dim=(-2, -1), keepdim=True)
+    std = audio.std(dim=(-2, -1), keepdim=True)
+    audio = (audio - mean) / (std + 1e-5) # Add epsilon for stability
+    return audio
+
+def calculate_padding_mask(pad_frames, total_frames, sr, output_steps, process_seconds, device, B):
+    # How many 2 seconds chunks does this audio have?
+    # Find it and then multiply by the output_steps.
+    total_frames = int((total_frames / sr) / process_seconds)
+    total_output_steps = output_steps * total_frames
+    mask = torch.zeros((B, total_output_steps), dtype = torch.bool, device = device)
+
+    # Check the number of padding tokens that we have in the audio.
+    output_sr = int(output_steps / process_seconds)
+    pad_seconds = pad_frames / sr
+    pad_steps = int(pad_seconds * output_sr)
+    # Create the mask
+    mask[..., total_output_steps - pad_steps:] = True 
+    return mask, total_output_steps - pad_steps