Upload 21 files

README.md

@@ -1,14 +1,132 @@
----
-title: Taro
-emoji: 👁
-colorFrom: green
-colorTo: green
-sdk: gradio
-sdk_version: 6.9.0
-python_version: '3.12'
-app_file: app.py
-pinned: false
-short_description: 'TARO: Video-to-Audio Synthesis (ICCV 2025)'
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# [ICCV'25] TARO: Timestep-Adaptive Representation Alignment with Onset-Aware Conditioning for Synchronized Video-to-Audio Synthesis
+<br>
+
+**[Tri Ton](https://triton99.github.io/)<sup>1</sup>, [Ji Woo Hong](https://jiwoohong93.github.io/)<sup>1</sup>, [Chang D. Yoo](https://sanctusfactory.com/family.php)<sup>1†</sup>** 
+<br>
+<sup>1</sup>KAIST, South Korea
+<br>
+†Corresponding authors
+
+<p align="center">
+        <a href="https://triton99.github.io/taro-site/" target='_blank'>
+        <img src="https://img.shields.io/badge/🐳-Project%20Page-blue">
+        </a>
+        <a href="https://arxiv.org/abs/2504.05684" target='_blank'>
+        <img src="https://img.shields.io/badge/arXiv-2312.13528-b31b1b.svg">
+        </a>
+        <img alt="GitHub Repo stars" src="https://img.shields.io/github/stars/triton99/TARO">
+</p>
+
+## 📣 News
+- **[09/2025]**: Training & Inference code released.
+- **[06/2025]**: TARO accepted to ICCV 2025 🎉.
+- **[04/2024]**: Paper uploaded to arXiv. Check out the manuscript [here](https://arxiv.org/abs/2504.05684).(https://arxiv.org/abs/2504.05684).
+
+## To-Dos
+- [x] Release model weights on Google Drive.
+- [x] Release inference code
+- [x] Release training code & dataset preparation
+
+## ⚙️ Environmental Setups
+1. Clone TARO.
+```bash
+git clone https://github.com/triton99/TARO
+cd TARO
+```
+
+2. Create the environment.
+```bash
+conda create -n taro python==3.10
+conda activate taro
+pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121
+
+# Training
+pip install --force pip==24.0
+git clone https://github.com/pytorch/fairseq
+cd fairseq
+pip install --editable ./ --no-build-isolation
+cd ..
+
+git clone https://github.com/cwx-worst-one/EAT.git
+
+# Inference
+pip3 install -r requirements.txt
+```
+
+## 📁 Data Preparations
+Please download the [VGGSound dataset](https://www.robots.ox.ac.uk/~vgg/data/vggsound/), extract the videos, and organize them into two folders: one with .mp4 files and one with corresponding .wav files (matching base filenames). 
+
+Update the path variables at the top of the preprocessing scripts to point to your folders, then run:
+```bash
+./preprocess_video.sh
+
+./preprocess_audio.sh
+```
+
+After processing, the data will have the following structure:
+```bash
+VGGSound/train
+    ├── videos
+    │   ├── abc.mp4
+    │   └── ...
+    ├── audios
+    │   ├── abc.wav
+    │   └── ...
+    ├── cavp_feats
+    │   ├── abc.npz
+    │   └── ...
+    ├── onset_feats
+    │   ├── abc.npz
+    │   └── ...
+    ├── melspec
+    │   ├── abc.npy
+    │   └── ...
+    └── fbank
+    │   ├── abc.npy
+    │   └── ...
+```
+
+
+## 🚀 Getting Started
+
+### Download Checkpoints
+
+The pretrained TARO checkpoint can be downloaded on [Google Drive](https://drive.google.com/drive/folders/1YqLsEtVYeSchhAh-wKS-BWuB6MK6_mJB?usp=sharing).
+
+The CAVP checkpoint can be downloaded from [Diff-Foley](https://github.com/luosiallen/Diff-Foley).
+
+The onset checkpoint can be downloaded from [SyncFusion](https://github.com/mcomunita/syncfusion).
+
+### Training
+```bash
+./train.sh
+```
+
+### Inference
+To run the inference code, you can use the following command:
+```bash
+python infer.py \
+    --video_path ./test.mp4 \
+    --save_folder_path ./output \
+    --cavp_config_path ./cavp/model/cavp.yaml \
+    --cavp_ckpt_path ./cavp_epoch66.ckpt \
+    --onset_ckpt_path ./onset_model.ckpt \
+    --model_ckpt_path ./taro_ckpt.pt
+```
+
+## 📖 Citing TARO
+
+If you find our repository useful, please consider giving it a star ⭐ and citing our paper in your work:
+
+```bibtex
+@inproceedings{ton2025taro,
+  title     = {TARO: Timestep-Adaptive Representation Alignment with Onset-Aware Conditioning for Synchronized Video-to-Audio Synthesis},
+  author    = {Ton, Tri and Hong, Ji Woo and Yoo, Chang D},
+  year      = {2025},
+  booktitle = {International Conference on Computer Vision (ICCV)},
+}
+```
+
+## 🤗 Acknowledgements
+
+Our code is based on [REPA](https://github.com/sihyun-yu/REPA), [Diff-Foley](https://github.com/luosiallen/Diff-Foley), and [SyncFusion](https://github.com/mcomunita/syncfusion). We thank the authors for their excellent work!

cavp/cavp.yaml

@@ -0,0 +1,9 @@
+
+model:
+  target: cavp.model.cavp_model.CAVP_Inference
+  params:
+    video_encode: Slowonly_pool
+    spec_encode: cnn14_pool
+    embed_dim: 512
+    video_pretrained: True
+    audio_pretrained: True

cavp/model/cavp_model.py

@@ -0,0 +1,96 @@
+
+
+from .cavp_modules import ResNet3dSlowOnly, Cnn14
+import torch.nn as nn 
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+class CAVP_Inference(nn.Module):
+
+    def __init__(
+            self,
+            video_encode,
+            spec_encode,
+            embed_dim: int,
+            video_pretrained: bool = False,
+            audio_pretrained: bool = False,
+    ):
+        super().__init__()
+
+        self.video_encode = video_encode
+        self.spec_encode = spec_encode
+
+
+        # 1). Video Encoder:
+        assert self.video_encode == "Slowonly_pool"
+        self.video_encoder = ResNet3dSlowOnly(depth=50, pretrained=None)    # Doesn't matter to set pretrained=None, since we will load CAVP weight outside.
+        
+        # Video Project & Pooling Head:
+        self.video_project_head = nn.Linear(2048, embed_dim)
+        self.video_pool = nn.MaxPool1d(kernel_size=16)
+
+
+        # 2). Spec Encoder:
+        assert self.spec_encode == "cnn14_pool"     # Pretrained
+        self.spec_encoder = Cnn14(embed_dim=512)
+        
+        # Spec Project & Pooling Head:
+        self.spec_project_head = nn.Identity()
+        self.spec_pool = nn.MaxPool1d(kernel_size=16)
+
+        # 3). Logit Scale:
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        
+        
+
+    def encode_video(self, video, normalize: bool = False, train=False, pool=True):
+
+        # Video: B x T x 3 x H x W
+        assert self.video_encode == "Slowonly_pool"
+        video = video.permute(0, 2, 1, 3, 4)
+        video_feat = self.video_encoder(video)
+        bs, c, t, _, _ = video_feat.shape
+        video_feat = video_feat.reshape(bs, c, t).permute(0, 2, 1)
+        video_feat = self.video_project_head(video_feat)
+        
+        # Pooling:
+        if pool:
+            video_feat = self.video_pool(video_feat.permute(0,2,1)).squeeze(2)
+        
+        # Normalize:
+        if normalize:
+            video_feat = F.normalize(video_feat, dim=-1)
+
+        return video_feat
+
+
+    def encode_spec(self, spec, normalize: bool = False, pool=True):
+        # spec: B x Mel_num x T
+        assert self.spec_encode == "cnn14_pool"
+        spec = spec.unsqueeze(1)                        # B x 1 x Mel x T
+        spec = spec.permute(0, 1, 3, 2)                 # B x 1 x T x Mel
+        spec_feat = self.spec_encoder(spec)             # B x T x C
+        spec_feat = self.spec_project_head(spec_feat)
+        
+        # Pooling:
+        if pool:
+            spec_feat = self.spec_pool(spec_feat.permute(0, 2, 1)).squeeze(2)
+
+        # Normalize:
+        if normalize:
+            spec_feat = F.normalize(spec_feat, dim=-1)
+        
+        return spec_feat
+
+
+    def forward(self, video, spec, output_dict=True):
+        video_features = self.encode_video(video, normalize=True)
+        spec_features = self.encode_spec(spec, normalize=True)
+        if output_dict:
+            return {
+                "video_features": video_features,
+                "spec_features": spec_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return video_features, spec_features, self.logit_scale.exp()

cavp/model/cavp_modules.py

@@ -0,0 +1,1545 @@
+import warnings
+
+import torch
+import torch.nn as nn
+from mmcv.cnn import ConvModule, kaiming_init
+from mmcv.runner import _load_checkpoint, load_checkpoint
+from mmcv.utils import print_log
+
+import warnings
+import torch.nn as nn
+import torch.utils.checkpoint as cp
+from mmcv.cnn import (ConvModule, NonLocal3d, build_activation_layer,
+                      constant_init, kaiming_init)
+from mmcv.runner import _load_checkpoint, load_checkpoint
+from mmcv.utils import _BatchNorm
+from torch.nn.modules.utils import _ntuple, _triple
+
+from itertools import repeat
+import collections.abc
+from typing import Callable, Optional, Sequence, Tuple
+from torch.utils.checkpoint import checkpoint
+
+from torch.nn import functional as F
+from collections import OrderedDict
+
+class BasicBlock3d(nn.Module):
+    """BasicBlock 3d block for ResNet3D.
+    Args:
+        inplanes (int): Number of channels for the input in first conv3d layer.
+        planes (int): Number of channels produced by some norm/conv3d layers.
+        spatial_stride (int): Spatial stride in the conv3d layer. Default: 1.
+        temporal_stride (int): Temporal stride in the conv3d layer. Default: 1.
+        dilation (int): Spacing between kernel elements. Default: 1.
+        downsample (nn.Module | None): Downsample layer. Default: None.
+        style (str): ``pytorch`` or ``caffe``. If set to "pytorch", the
+            stride-two layer is the 3x3 conv layer, otherwise the stride-two
+            layer is the first 1x1 conv layer. Default: 'pytorch'.
+        inflate (bool): Whether to inflate kernel. Default: True.
+        non_local (bool): Determine whether to apply non-local module in this
+            block. Default: False.
+        non_local_cfg (dict): Config for non-local module. Default: ``dict()``.
+        conv_cfg (dict): Config dict for convolution layer.
+            Default: ``dict(type='Conv3d')``.
+        norm_cfg (dict): Config for norm layers. required keys are ``type``,
+            Default: ``dict(type='BN3d')``.
+        act_cfg (dict): Config dict for activation layer.
+            Default: ``dict(type='ReLU')``.
+        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
+            memory while slowing down the training speed. Default: False.
+    """
+    expansion = 1
+
+    def __init__(self,
+                 inplanes,
+                 planes,
+                 spatial_stride=1,
+                 temporal_stride=1,
+                 dilation=1,
+                 downsample=None,
+                 style='pytorch',
+                 inflate=True,
+                 non_local=False,
+                 non_local_cfg=dict(),
+                 conv_cfg=dict(type='Conv3d'),
+                 norm_cfg=dict(type='BN3d'),
+                 act_cfg=dict(type='ReLU'),
+                 with_cp=False,
+                 **kwargs):
+        super().__init__()
+        assert style in ['pytorch', 'caffe']
+        # make sure that only ``inflate_style`` is passed into kwargs
+        assert set(kwargs).issubset(['inflate_style'])
+
+        self.inplanes = inplanes
+        self.planes = planes
+        self.spatial_stride = spatial_stride
+        self.temporal_stride = temporal_stride
+        self.dilation = dilation
+        self.style = style
+        self.inflate = inflate
+        self.conv_cfg = conv_cfg
+        self.norm_cfg = norm_cfg
+        self.act_cfg = act_cfg
+        self.with_cp = with_cp
+        self.non_local = non_local
+        self.non_local_cfg = non_local_cfg
+
+        self.conv1_stride_s = spatial_stride
+        self.conv2_stride_s = 1
+        self.conv1_stride_t = temporal_stride
+        self.conv2_stride_t = 1
+
+        if self.inflate:
+            conv1_kernel_size = (3, 3, 3)
+            conv1_padding = (1, dilation, dilation)
+            conv2_kernel_size = (3, 3, 3)
+            conv2_padding = (1, 1, 1)
+        else:
+            conv1_kernel_size = (1, 3, 3)
+            conv1_padding = (0, dilation, dilation)
+            conv2_kernel_size = (1, 3, 3)
+            conv2_padding = (0, 1, 1)
+
+        self.conv1 = ConvModule(
+            inplanes,
+            planes,
+            conv1_kernel_size,
+            stride=(self.conv1_stride_t, self.conv1_stride_s,
+                    self.conv1_stride_s),
+            padding=conv1_padding,
+            dilation=(1, dilation, dilation),
+            bias=False,
+            conv_cfg=self.conv_cfg,
+            norm_cfg=self.norm_cfg,
+            act_cfg=self.act_cfg)
+
+        self.conv2 = ConvModule(
+            planes,
+            planes * self.expansion,
+            conv2_kernel_size,
+            stride=(self.conv2_stride_t, self.conv2_stride_s,
+                    self.conv2_stride_s),
+            padding=conv2_padding,
+            bias=False,
+            conv_cfg=self.conv_cfg,
+            norm_cfg=self.norm_cfg,
+            act_cfg=None)
+
+        self.downsample = downsample
+        self.relu = build_activation_layer(self.act_cfg)
+
+        if self.non_local:
+            self.non_local_block = NonLocal3d(self.conv2.norm.num_features,
+                                              **self.non_local_cfg)
+
+    def forward(self, x):
+        """Defines the computation performed at every call."""
+
+        def _inner_forward(x):
+            """Forward wrapper for utilizing checkpoint."""
+            identity = x
+
+            out = self.conv1(x)
+            out = self.conv2(out)
+
+            if self.downsample is not None:
+                identity = self.downsample(x)
+
+            out = out + identity
+            return out
+
+        if self.with_cp and x.requires_grad:
+            out = cp.checkpoint(_inner_forward, x)
+        else:
+            out = _inner_forward(x)
+        out = self.relu(out)
+
+        if self.non_local:
+            out = self.non_local_block(out)
+
+        return out
+
+
+class Bottleneck3d(nn.Module):
+    """Bottleneck 3d block for ResNet3D.
+    Args:
+        inplanes (int): Number of channels for the input in first conv3d layer.
+        planes (int): Number of channels produced by some norm/conv3d layers.
+        spatial_stride (int): Spatial stride in the conv3d layer. Default: 1.
+        temporal_stride (int): Temporal stride in the conv3d layer. Default: 1.
+        dilation (int): Spacing between kernel elements. Default: 1.
+        downsample (nn.Module | None): Downsample layer. Default: None.
+        style (str): ``pytorch`` or ``caffe``. If set to "pytorch", the
+            stride-two layer is the 3x3 conv layer, otherwise the stride-two
+            layer is the first 1x1 conv layer. Default: 'pytorch'.
+        inflate (bool): Whether to inflate kernel. Default: True.
+        inflate_style (str): ``3x1x1`` or ``3x3x3``. which determines the
+            kernel sizes and padding strides for conv1 and conv2 in each block.
+            Default: '3x1x1'.
+        non_local (bool): Determine whether to apply non-local module in this
+            block. Default: False.
+        non_local_cfg (dict): Config for non-local module. Default: ``dict()``.
+        conv_cfg (dict): Config dict for convolution layer.
+            Default: ``dict(type='Conv3d')``.
+        norm_cfg (dict): Config for norm layers. required keys are ``type``,
+            Default: ``dict(type='BN3d')``.
+        act_cfg (dict): Config dict for activation layer.
+            Default: ``dict(type='ReLU')``.
+        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
+            memory while slowing down the training speed. Default: False.
+    """
+    expansion = 4
+
+    def __init__(self,
+                 inplanes,
+                 planes,
+                 spatial_stride=1,
+                 temporal_stride=1,
+                 dilation=1,
+                 downsample=None,
+                 style='pytorch',
+                 inflate=True,
+                 inflate_style='3x1x1',
+                 non_local=False,
+                 non_local_cfg=dict(),
+                 conv_cfg=dict(type='Conv3d'),
+                 norm_cfg=dict(type='BN3d'),
+                 act_cfg=dict(type='ReLU'),
+                 with_cp=False):
+        super().__init__()
+        assert style in ['pytorch', 'caffe']
+        assert inflate_style in ['3x1x1', '3x3x3']
+
+        self.inplanes = inplanes
+        self.planes = planes
+        self.spatial_stride = spatial_stride
+        self.temporal_stride = temporal_stride
+        self.dilation = dilation
+        self.style = style
+        self.inflate = inflate
+        self.inflate_style = inflate_style
+        self.norm_cfg = norm_cfg
+        self.conv_cfg = conv_cfg
+        self.act_cfg = act_cfg
+        self.with_cp = with_cp
+        self.non_local = non_local
+        self.non_local_cfg = non_local_cfg
+
+        if self.style == 'pytorch':
+            self.conv1_stride_s = 1
+            self.conv2_stride_s = spatial_stride
+            self.conv1_stride_t = 1
+            self.conv2_stride_t = temporal_stride
+        else:
+            self.conv1_stride_s = spatial_stride
+            self.conv2_stride_s = 1
+            self.conv1_stride_t = temporal_stride
+            self.conv2_stride_t = 1
+
+        if self.inflate:
+            if inflate_style == '3x1x1':
+                conv1_kernel_size = (3, 1, 1)
+                conv1_padding = (1, 0, 0)
+                conv2_kernel_size = (1, 3, 3)
+                conv2_padding = (0, dilation, dilation)
+            else:
+                conv1_kernel_size = (1, 1, 1)
+                conv1_padding = (0, 0, 0)
+                conv2_kernel_size = (3, 3, 3)
+                conv2_padding = (1, dilation, dilation)
+        else:
+            conv1_kernel_size = (1, 1, 1)
+            conv1_padding = (0, 0, 0)
+            conv2_kernel_size = (1, 3, 3)
+            conv2_padding = (0, dilation, dilation)
+
+        self.conv1 = ConvModule(
+            inplanes,
+            planes,
+            conv1_kernel_size,
+            stride=(self.conv1_stride_t, self.conv1_stride_s,
+                    self.conv1_stride_s),
+            padding=conv1_padding,
+            bias=False,
+            conv_cfg=self.conv_cfg,
+            norm_cfg=self.norm_cfg,
+            act_cfg=self.act_cfg)
+
+        self.conv2 = ConvModule(
+            planes,
+            planes,
+            conv2_kernel_size,
+            stride=(self.conv2_stride_t, self.conv2_stride_s,
+                    self.conv2_stride_s),
+            padding=conv2_padding,
+            dilation=(1, dilation, dilation),
+            bias=False,
+            conv_cfg=self.conv_cfg,
+            norm_cfg=self.norm_cfg,
+            act_cfg=self.act_cfg)
+
+        self.conv3 = ConvModule(
+            planes,
+            planes * self.expansion,
+            1,
+            bias=False,
+            conv_cfg=self.conv_cfg,
+            norm_cfg=self.norm_cfg,
+            # No activation in the third ConvModule for bottleneck
+            act_cfg=None)
+
+        self.downsample = downsample
+        self.relu = build_activation_layer(self.act_cfg)
+
+        if self.non_local:
+            self.non_local_block = NonLocal3d(self.conv3.norm.num_features,
+                                              **self.non_local_cfg)
+
+    def forward(self, x):
+        """Defines the computation performed at every call."""
+
+        def _inner_forward(x):
+            """Forward wrapper for utilizing checkpoint."""
+            identity = x
+
+            out = self.conv1(x)
+            out = self.conv2(out)
+            out = self.conv3(out)
+
+            if self.downsample is not None:
+                identity = self.downsample(x)
+
+            out = out + identity
+            return out
+
+        if self.with_cp and x.requires_grad:
+            out = cp.checkpoint(_inner_forward, x)
+        else:
+            out = _inner_forward(x)
+        out = self.relu(out)
+
+        if self.non_local:
+            out = self.non_local_block(out)
+
+        return out
+
+
+class ResNet3d(nn.Module):
+    """ResNet 3d backbone.
+    Args:
+        depth (int): Depth of resnet, from {18, 34, 50, 101, 152}.
+        pretrained (str | None): Name of pretrained model.
+        stage_blocks (tuple | None): Set number of stages for each res layer.
+            Default: None.
+        pretrained2d (bool): Whether to load pretrained 2D model.
+            Default: True.
+        in_channels (int): Channel num of input features. Default: 3.
+        base_channels (int): Channel num of stem output features. Default: 64.
+        out_indices (Sequence[int]): Indices of output feature. Default: (3, ).
+        num_stages (int): Resnet stages. Default: 4.
+        spatial_strides (Sequence[int]):
+            Spatial strides of residual blocks of each stage.
+            Default: ``(1, 2, 2, 2)``.
+        temporal_strides (Sequence[int]):
+            Temporal strides of residual blocks of each stage.
+            Default: ``(1, 1, 1, 1)``.
+        dilations (Sequence[int]): Dilation of each stage.
+            Default: ``(1, 1, 1, 1)``.
+        conv1_kernel (Sequence[int]): Kernel size of the first conv layer.
+            Default: ``(3, 7, 7)``.
+        conv1_stride_s (int): Spatial stride of the first conv layer.
+            Default: 2.
+        conv1_stride_t (int): Temporal stride of the first conv layer.
+            Default: 1.
+        pool1_stride_s (int): Spatial stride of the first pooling layer.
+            Default: 2.
+        pool1_stride_t (int): Temporal stride of the first pooling layer.
+            Default: 1.
+        with_pool2 (bool): Whether to use pool2. Default: True.
+        style (str): `pytorch` or `caffe`. If set to "pytorch", the stride-two
+            layer is the 3x3 conv layer, otherwise the stride-two layer is
+            the first 1x1 conv layer. Default: 'pytorch'.
+        frozen_stages (int): Stages to be frozen (all param fixed). -1 means
+            not freezing any parameters. Default: -1.
+        inflate (Sequence[int]): Inflate Dims of each block.
+            Default: (1, 1, 1, 1).
+        inflate_style (str): ``3x1x1`` or ``3x3x3``. which determines the
+            kernel sizes and padding strides for conv1 and conv2 in each block.
+            Default: '3x1x1'.
+        conv_cfg (dict): Config for conv layers. required keys are ``type``
+            Default: ``dict(type='Conv3d')``.
+        norm_cfg (dict): Config for norm layers. required keys are ``type`` and
+            ``requires_grad``.
+            Default: ``dict(type='BN3d', requires_grad=True)``.
+        act_cfg (dict): Config dict for activation layer.
+            Default: ``dict(type='ReLU', inplace=True)``.
+        norm_eval (bool): Whether to set BN layers to eval mode, namely, freeze
+            running stats (mean and var). Default: False.
+        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
+            memory while slowing down the training speed. Default: False.
+        non_local (Sequence[int]): Determine whether to apply non-local module
+            in the corresponding block of each stages. Default: (0, 0, 0, 0).
+        non_local_cfg (dict): Config for non-local module. Default: ``dict()``.
+        zero_init_residual (bool):
+            Whether to use zero initialization for residual block,
+            Default: True.
+        kwargs (dict, optional): Key arguments for "make_res_layer".
+    """
+
+    arch_settings = {
+        18: (BasicBlock3d, (2, 2, 2, 2)),
+        34: (BasicBlock3d, (3, 4, 6, 3)),
+        50: (Bottleneck3d, (3, 4, 6, 3)),
+        101: (Bottleneck3d, (3, 4, 23, 3)),
+        152: (Bottleneck3d, (3, 8, 36, 3))
+    }
+
+    def __init__(self,
+                 depth,
+                 pretrained,
+                 stage_blocks=None,
+                 pretrained2d=True,
+                 in_channels=3,
+                 num_stages=4,
+                 base_channels=64,
+                 out_indices=(3, ),
+                 spatial_strides=(1, 2, 2, 2),
+                 temporal_strides=(1, 1, 1, 1),
+                 dilations=(1, 1, 1, 1),
+                 conv1_kernel=(3, 7, 7),
+                 conv1_stride_s=2,
+                 conv1_stride_t=1,
+                 pool1_stride_s=2,
+                 pool1_stride_t=1,
+                 with_pool1=True,
+                 with_pool2=True,
+                 style='pytorch',
+                 frozen_stages=-1,
+                 inflate=(1, 1, 1, 1),
+                 inflate_style='3x1x1',
+                 conv_cfg=dict(type='Conv3d'),
+                 norm_cfg=dict(type='BN3d', requires_grad=True),
+                 act_cfg=dict(type='ReLU', inplace=True),
+                 norm_eval=False,
+                 with_cp=False,
+                 non_local=(0, 0, 0, 0),
+                 non_local_cfg=dict(),
+                 zero_init_residual=True,
+                 **kwargs):
+        super().__init__()
+        if depth not in self.arch_settings:
+            raise KeyError(f'invalid depth {depth} for resnet')
+        self.depth = depth
+        self.pretrained = pretrained
+        self.pretrained2d = pretrained2d
+        self.in_channels = in_channels
+        self.base_channels = base_channels
+        self.num_stages = num_stages
+        assert 1 <= num_stages <= 4
+        self.stage_blocks = stage_blocks
+        self.out_indices = out_indices
+        assert max(out_indices) < num_stages
+        self.spatial_strides = spatial_strides
+        self.temporal_strides = temporal_strides
+        self.dilations = dilations
+        assert len(spatial_strides) == len(temporal_strides) == len(
+            dilations) == num_stages
+        if self.stage_blocks is not None:
+            assert len(self.stage_blocks) == num_stages
+
+        self.conv1_kernel = conv1_kernel
+        self.conv1_stride_s = conv1_stride_s
+        self.conv1_stride_t = conv1_stride_t
+        self.pool1_stride_s = pool1_stride_s
+        self.pool1_stride_t = pool1_stride_t
+        self.with_pool1 = with_pool1
+        self.with_pool2 = with_pool2
+        self.style = style
+        self.frozen_stages = frozen_stages
+        self.stage_inflations = _ntuple(num_stages)(inflate)
+        self.non_local_stages = _ntuple(num_stages)(non_local)
+        self.inflate_style = inflate_style
+        self.conv_cfg = conv_cfg
+        self.norm_cfg = norm_cfg
+        self.act_cfg = act_cfg
+        self.norm_eval = norm_eval
+        self.with_cp = with_cp
+        self.zero_init_residual = zero_init_residual
+
+        self.block, stage_blocks = self.arch_settings[depth]
+
+        if self.stage_blocks is None:
+            self.stage_blocks = stage_blocks[:num_stages]
+
+        self.inplanes = self.base_channels
+
+        self.non_local_cfg = non_local_cfg
+
+        self._make_stem_layer()
+
+        self.res_layers = []
+        for i, num_blocks in enumerate(self.stage_blocks):
+            spatial_stride = spatial_strides[i]
+            temporal_stride = temporal_strides[i]
+            dilation = dilations[i]
+            planes = self.base_channels * 2**i
+            res_layer = self.make_res_layer(
+                self.block,
+                self.inplanes,
+                planes,
+                num_blocks,
+                spatial_stride=spatial_stride,
+                temporal_stride=temporal_stride,
+                dilation=dilation,
+                style=self.style,
+                norm_cfg=self.norm_cfg,
+                conv_cfg=self.conv_cfg,
+                act_cfg=self.act_cfg,
+                non_local=self.non_local_stages[i],
+                non_local_cfg=self.non_local_cfg,
+                inflate=self.stage_inflations[i],
+                inflate_style=self.inflate_style,
+                with_cp=with_cp,
+                **kwargs)
+            self.inplanes = planes * self.block.expansion
+            layer_name = f'layer{i + 1}'
+            self.add_module(layer_name, res_layer)
+            self.res_layers.append(layer_name)
+
+        self.feat_dim = self.block.expansion * self.base_channels * 2**(
+            len(self.stage_blocks) - 1)
+        
+
+        # Adaptive Pool:
+        self.adaptive_pool = nn.AdaptiveAvgPool2d((1,1))
+
+    @staticmethod
+    def make_res_layer(block,
+                       inplanes,
+                       planes,
+                       blocks,
+                       spatial_stride=1,
+                       temporal_stride=1,
+                       dilation=1,
+                       style='pytorch',
+                       inflate=1,
+                       inflate_style='3x1x1',
+                       non_local=0,
+                       non_local_cfg=dict(),
+                       norm_cfg=None,
+                       act_cfg=None,
+                       conv_cfg=None,
+                       with_cp=False,
+                       **kwargs):
+        """Build residual layer for ResNet3D.
+        Args:
+            block (nn.Module): Residual module to be built.
+            inplanes (int): Number of channels for the input feature
+                in each block.
+            planes (int): Number of channels for the output feature
+                in each block.
+            blocks (int): Number of residual blocks.
+            spatial_stride (int | Sequence[int]): Spatial strides in
+                residual and conv layers. Default: 1.
+            temporal_stride (int | Sequence[int]): Temporal strides in
+                residual and conv layers. Default: 1.
+            dilation (int): Spacing between kernel elements. Default: 1.
+            style (str): ``pytorch`` or ``caffe``. If set to ``pytorch``,
+                the stride-two layer is the 3x3 conv layer, otherwise
+                the stride-two layer is the first 1x1 conv layer.
+                Default: ``pytorch``.
+            inflate (int | Sequence[int]): Determine whether to inflate
+                for each block. Default: 1.
+            inflate_style (str): ``3x1x1`` or ``3x3x3``. which determines
+                the kernel sizes and padding strides for conv1 and conv2
+                in each block. Default: '3x1x1'.
+            non_local (int | Sequence[int]): Determine whether to apply
+                non-local module in the corresponding block of each stages.
+                Default: 0.
+            non_local_cfg (dict): Config for non-local module.
+                Default: ``dict()``.
+            conv_cfg (dict | None): Config for norm layers. Default: None.
+            norm_cfg (dict | None): Config for norm layers. Default: None.
+            act_cfg (dict | None): Config for activate layers. Default: None.
+            with_cp (bool | None): Use checkpoint or not. Using checkpoint
+                will save some memory while slowing down the training speed.
+                Default: False.
+        Returns:
+            nn.Module: A residual layer for the given config.
+        """
+        inflate = inflate if not isinstance(inflate,
+                                            int) else (inflate, ) * blocks
+        non_local = non_local if not isinstance(
+            non_local, int) else (non_local, ) * blocks
+        assert len(inflate) == blocks and len(non_local) == blocks
+        downsample = None
+        if spatial_stride != 1 or inplanes != planes * block.expansion:
+            downsample = ConvModule(
+                inplanes,
+                planes * block.expansion,
+                kernel_size=1,
+                stride=(temporal_stride, spatial_stride, spatial_stride),
+                bias=False,
+                conv_cfg=conv_cfg,
+                norm_cfg=norm_cfg,
+                act_cfg=None)
+
+        layers = []
+        layers.append(
+            block(
+                inplanes,
+                planes,
+                spatial_stride=spatial_stride,
+                temporal_stride=temporal_stride,
+                dilation=dilation,
+                downsample=downsample,
+                style=style,
+                inflate=(inflate[0] == 1),
+                inflate_style=inflate_style,
+                non_local=(non_local[0] == 1),
+                non_local_cfg=non_local_cfg,
+                norm_cfg=norm_cfg,
+                conv_cfg=conv_cfg,
+                act_cfg=act_cfg,
+                with_cp=with_cp,
+                **kwargs))
+        inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(
+                block(
+                    inplanes,
+                    planes,
+                    spatial_stride=1,
+                    temporal_stride=1,
+                    dilation=dilation,
+                    style=style,
+                    inflate=(inflate[i] == 1),
+                    inflate_style=inflate_style,
+                    non_local=(non_local[i] == 1),
+                    non_local_cfg=non_local_cfg,
+                    norm_cfg=norm_cfg,
+                    conv_cfg=conv_cfg,
+                    act_cfg=act_cfg,
+                    with_cp=with_cp,
+                    **kwargs))
+
+        return nn.Sequential(*layers)
+
+    @staticmethod
+    def _inflate_conv_params(conv3d, state_dict_2d, module_name_2d,
+                             inflated_param_names):
+        """Inflate a conv module from 2d to 3d.
+        Args:
+            conv3d (nn.Module): The destination conv3d module.
+            state_dict_2d (OrderedDict): The state dict of pretrained 2d model.
+            module_name_2d (str): The name of corresponding conv module in the
+                2d model.
+            inflated_param_names (list[str]): List of parameters that have been
+                inflated.
+        """
+        weight_2d_name = module_name_2d + '.weight'
+
+        conv2d_weight = state_dict_2d[weight_2d_name]
+        kernel_t = conv3d.weight.data.shape[2]
+
+        new_weight = conv2d_weight.data.unsqueeze(2).expand_as(
+            conv3d.weight) / kernel_t
+        conv3d.weight.data.copy_(new_weight)
+        inflated_param_names.append(weight_2d_name)
+
+        if getattr(conv3d, 'bias') is not None:
+            bias_2d_name = module_name_2d + '.bias'
+            conv3d.bias.data.copy_(state_dict_2d[bias_2d_name])
+            inflated_param_names.append(bias_2d_name)
+
+    @staticmethod
+    def _inflate_bn_params(bn3d, state_dict_2d, module_name_2d,
+                           inflated_param_names):
+        """Inflate a norm module from 2d to 3d.
+        Args:
+            bn3d (nn.Module): The destination bn3d module.
+            state_dict_2d (OrderedDict): The state dict of pretrained 2d model.
+            module_name_2d (str): The name of corresponding bn module in the
+                2d model.
+            inflated_param_names (list[str]): List of parameters that have been
+                inflated.
+        """
+        for param_name, param in bn3d.named_parameters():
+            param_2d_name = f'{module_name_2d}.{param_name}'
+            param_2d = state_dict_2d[param_2d_name]
+            if param.data.shape != param_2d.shape:
+                warnings.warn(f'The parameter of {module_name_2d} is not'
+                              'loaded due to incompatible shapes. ')
+                return
+
+            param.data.copy_(param_2d)
+            inflated_param_names.append(param_2d_name)
+
+        for param_name, param in bn3d.named_buffers():
+            param_2d_name = f'{module_name_2d}.{param_name}'
+            # some buffers like num_batches_tracked may not exist in old
+            # checkpoints
+            if param_2d_name in state_dict_2d:
+                param_2d = state_dict_2d[param_2d_name]
+                param.data.copy_(param_2d)
+                inflated_param_names.append(param_2d_name)
+
+    @staticmethod
+    def _inflate_weights(self, logger):
+        """Inflate the resnet2d parameters to resnet3d.
+        The differences between resnet3d and resnet2d mainly lie in an extra
+        axis of conv kernel. To utilize the pretrained parameters in 2d model,
+        the weight of conv2d models should be inflated to fit in the shapes of
+        the 3d counterpart.
+        Args:
+            logger (logging.Logger): The logger used to print
+                debugging information.
+        """
+
+        state_dict_r2d = _load_checkpoint(self.pretrained)
+        if 'state_dict' in state_dict_r2d:
+            state_dict_r2d = state_dict_r2d['state_dict']
+
+        inflated_param_names = []
+        for name, module in self.named_modules():
+            if isinstance(module, ConvModule):
+                # we use a ConvModule to wrap conv+bn+relu layers, thus the
+                # name mapping is needed
+                if 'downsample' in name:
+                    # layer{X}.{Y}.downsample.conv->layer{X}.{Y}.downsample.0
+                    original_conv_name = name + '.0'
+                    # layer{X}.{Y}.downsample.bn->layer{X}.{Y}.downsample.1
+                    original_bn_name = name + '.1'
+                else:
+                    # layer{X}.{Y}.conv{n}.conv->layer{X}.{Y}.conv{n}
+                    original_conv_name = name
+                    # layer{X}.{Y}.conv{n}.bn->layer{X}.{Y}.bn{n}
+                    original_bn_name = name.replace('conv', 'bn')
+                if original_conv_name + '.weight' not in state_dict_r2d:
+                    logger.warning(f'Module not exist in the state_dict_r2d'
+                                   f': {original_conv_name}')
+                else:
+                    shape_2d = state_dict_r2d[original_conv_name +
+                                              '.weight'].shape
+                    shape_3d = module.conv.weight.data.shape
+                    if shape_2d != shape_3d[:2] + shape_3d[3:]:
+                        logger.warning(f'Weight shape mismatch for '
+                                       f': {original_conv_name} : '
+                                       f'3d weight shape: {shape_3d}; '
+                                       f'2d weight shape: {shape_2d}. ')
+                    else:
+                        self._inflate_conv_params(module.conv, state_dict_r2d,
+                                                  original_conv_name,
+                                                  inflated_param_names)
+
+                if original_bn_name + '.weight' not in state_dict_r2d:
+                    logger.warning(f'Module not exist in the state_dict_r2d'
+                                   f': {original_bn_name}')
+                else:
+                    self._inflate_bn_params(module.bn, state_dict_r2d,
+                                            original_bn_name,
+                                            inflated_param_names)
+
+        # check if any parameters in the 2d checkpoint are not loaded
+        remaining_names = set(
+            state_dict_r2d.keys()) - set(inflated_param_names)
+        if remaining_names:
+            logger.info(f'These parameters in the 2d checkpoint are not loaded'
+                        f': {remaining_names}')
+
+    def inflate_weights(self, logger):
+        self._inflate_weights(self, logger)
+
+    def _make_stem_layer(self):
+        """Construct the stem layers consists of a conv+norm+act module and a
+        pooling layer."""
+        self.conv1 = ConvModule(
+            self.in_channels,
+            self.base_channels,
+            kernel_size=self.conv1_kernel,
+            stride=(self.conv1_stride_t, self.conv1_stride_s,
+                    self.conv1_stride_s),
+            padding=tuple([(k - 1) // 2 for k in _triple(self.conv1_kernel)]),
+            bias=False,
+            conv_cfg=self.conv_cfg,
+            norm_cfg=self.norm_cfg,
+            act_cfg=self.act_cfg)
+
+        self.maxpool = nn.MaxPool3d(
+            kernel_size=(1, 3, 3),
+            stride=(self.pool1_stride_t, self.pool1_stride_s,
+                    self.pool1_stride_s),
+            padding=(0, 1, 1))
+
+        self.pool2 = nn.MaxPool3d(kernel_size=(2, 1, 1), stride=(2, 1, 1))
+
+    def _freeze_stages(self):
+        """Prevent all the parameters from being optimized before
+        ``self.frozen_stages``."""
+        if self.frozen_stages >= 0:
+            self.conv1.eval()
+            for param in self.conv1.parameters():
+                param.requires_grad = False
+
+        for i in range(1, self.frozen_stages + 1):
+            m = getattr(self, f'layer{i}')
+            m.eval()
+            for param in m.parameters():
+                param.requires_grad = False
+
+    @staticmethod
+    def _init_weights(self, pretrained=None):
+        """Initiate the parameters either from existing checkpoint or from
+        scratch.
+        Args:
+            pretrained (str | None): The path of the pretrained weight. Will
+                override the original `pretrained` if set. The arg is added to
+                be compatible with mmdet. Default: None.
+        """
+        if pretrained:
+            self.pretrained = pretrained
+        if isinstance(self.pretrained, str):
+            logger = get_root_logger()
+            logger.info(f'load model from: {self.pretrained}')
+
+            if self.pretrained2d:
+                # Inflate 2D model into 3D model.
+                self.inflate_weights(logger)
+
+            else:
+                # Directly load 3D model.
+                load_checkpoint(
+                    self, self.pretrained, strict=False, logger=logger)
+
+        elif self.pretrained is None:
+            for m in self.modules():
+                if isinstance(m, nn.Conv3d):
+                    kaiming_init(m)
+                elif isinstance(m, _BatchNorm):
+                    constant_init(m, 1)
+
+            if self.zero_init_residual:
+                for m in self.modules():
+                    if isinstance(m, Bottleneck3d):
+                        constant_init(m.conv3.bn, 0)
+                    elif isinstance(m, BasicBlock3d):
+                        constant_init(m.conv2.bn, 0)
+        else:
+            raise TypeError('pretrained must be a str or None')
+
+    def init_weights(self, pretrained=None):
+        self._init_weights(self, pretrained)
+
+    def forward(self, x):
+        """Defines the computation performed at every call.
+        Args:
+            x (torch.Tensor): The input data.
+        Returns:
+            torch.Tensor: The feature of the input
+            samples extracted by the backbone.
+        """
+        x = self.conv1(x)
+        if self.with_pool1:
+            x = self.maxpool(x)
+        outs = []
+        for i, layer_name in enumerate(self.res_layers):
+            res_layer = getattr(self, layer_name)
+            x = res_layer(x)
+            if i == 0 and self.with_pool2:
+                x = self.pool2(x)
+            if i in self.out_indices:
+                outs.append(x)
+        if len(outs) == 1:
+            out = outs[0]
+            out = self.adaptive_pool(out)
+            return out
+
+        return tuple(outs)
+
+    def train(self, mode=True):
+        """Set the optimization status when training."""
+        super().train(mode)
+        self._freeze_stages()
+        if mode and self.norm_eval:
+            for m in self.modules():
+                if isinstance(m, _BatchNorm):
+                    m.eval()
+
+
+
+
+class ResNet3dPathway(ResNet3d):
+    """A pathway of Slowfast based on ResNet3d.
+    Args:
+        *args (arguments): Arguments same as :class:``ResNet3d``.
+        lateral (bool): Determines whether to enable the lateral connection
+            from another pathway. Default: False.
+        speed_ratio (int): Speed ratio indicating the ratio between time
+            dimension of the fast and slow pathway, corresponding to the
+            ``alpha`` in the paper. Default: 8.
+        channel_ratio (int): Reduce the channel number of fast pathway
+            by ``channel_ratio``, corresponding to ``beta`` in the paper.
+            Default: 8.
+        fusion_kernel (int): The kernel size of lateral fusion.
+            Default: 5.
+        **kwargs (keyword arguments): Keywords arguments for ResNet3d.
+    """
+
+    def __init__(self,
+                 *args,
+                 lateral=False,
+                 lateral_norm=False,
+                 speed_ratio=8,
+                 channel_ratio=8,
+                 fusion_kernel=5,
+                 **kwargs):
+        self.lateral = lateral
+        self.lateral_norm = lateral_norm
+        self.speed_ratio = speed_ratio
+        self.channel_ratio = channel_ratio
+        self.fusion_kernel = fusion_kernel
+        super().__init__(*args, **kwargs)
+        self.inplanes = self.base_channels
+        if self.lateral:
+            self.conv1_lateral = ConvModule(
+                self.inplanes // self.channel_ratio,
+                # https://arxiv.org/abs/1812.03982, the
+                # third type of lateral connection has out_channel:
+                # 2 * \beta * C
+                self.inplanes * 2 // self.channel_ratio,
+                kernel_size=(fusion_kernel, 1, 1),
+                stride=(self.speed_ratio, 1, 1),
+                padding=((fusion_kernel - 1) // 2, 0, 0),
+                bias=False,
+                conv_cfg=self.conv_cfg,
+                norm_cfg=self.norm_cfg if self.lateral_norm else None,
+                act_cfg=self.act_cfg if self.lateral_norm else None)
+
+        self.lateral_connections = []
+        for i in range(len(self.stage_blocks)):
+            planes = self.base_channels * 2**i
+            self.inplanes = planes * self.block.expansion
+
+            if lateral and i != self.num_stages - 1:
+                # no lateral connection needed in final stage
+                lateral_name = f'layer{(i + 1)}_lateral'
+                setattr(
+                    self, lateral_name,
+                    ConvModule(
+                        self.inplanes // self.channel_ratio,
+                        self.inplanes * 2 // self.channel_ratio,
+                        kernel_size=(fusion_kernel, 1, 1),
+                        stride=(self.speed_ratio, 1, 1),
+                        padding=((fusion_kernel - 1) // 2, 0, 0),
+                        bias=False,
+                        conv_cfg=self.conv_cfg,
+                        norm_cfg=self.norm_cfg if self.lateral_norm else None,
+                        act_cfg=self.act_cfg if self.lateral_norm else None))
+                self.lateral_connections.append(lateral_name)
+
+    def make_res_layer(self,
+                       block,
+                       inplanes,
+                       planes,
+                       blocks,
+                       spatial_stride=1,
+                       temporal_stride=1,
+                       dilation=1,
+                       style='pytorch',
+                       inflate=1,
+                       inflate_style='3x1x1',
+                       non_local=0,
+                       non_local_cfg=dict(),
+                       conv_cfg=None,
+                       norm_cfg=None,
+                       act_cfg=None,
+                       with_cp=False):
+        """Build residual layer for Slowfast.
+        Args:
+            block (nn.Module): Residual module to be built.
+            inplanes (int): Number of channels for the input
+                feature in each block.
+            planes (int): Number of channels for the output
+                feature in each block.
+            blocks (int): Number of residual blocks.
+            spatial_stride (int | Sequence[int]): Spatial strides
+                in residual and conv layers. Default: 1.
+            temporal_stride (int | Sequence[int]): Temporal strides in
+                residual and conv layers. Default: 1.
+            dilation (int): Spacing between kernel elements. Default: 1.
+            style (str): ``pytorch`` or ``caffe``. If set to ``pytorch``,
+                the stride-two layer is the 3x3 conv layer,
+                otherwise the stride-two layer is the first 1x1 conv layer.
+                Default: ``pytorch``.
+            inflate (int | Sequence[int]): Determine whether to inflate
+                for each block. Default: 1.
+            inflate_style (str): ``3x1x1`` or ``3x3x3``. which determines
+                the kernel sizes and padding strides for conv1 and
+                conv2 in each block. Default: ``3x1x1``.
+            non_local (int | Sequence[int]): Determine whether to apply
+                non-local module in the corresponding block of each stages.
+                Default: 0.
+            non_local_cfg (dict): Config for non-local module.
+                Default: ``dict()``.
+            conv_cfg (dict | None): Config for conv layers. Default: None.
+            norm_cfg (dict | None): Config for norm layers. Default: None.
+            act_cfg (dict | None): Config for activate layers. Default: None.
+            with_cp (bool): Use checkpoint or not. Using checkpoint will save
+                some memory while slowing down the training speed.
+                Default: False.
+        Returns:
+            nn.Module: A residual layer for the given config.
+        """
+        inflate = inflate if not isinstance(inflate,
+                                            int) else (inflate, ) * blocks
+        non_local = non_local if not isinstance(
+            non_local, int) else (non_local, ) * blocks
+        assert len(inflate) == blocks and len(non_local) == blocks
+        if self.lateral:
+            lateral_inplanes = inplanes * 2 // self.channel_ratio
+        else:
+            lateral_inplanes = 0
+        if (spatial_stride != 1
+                or (inplanes + lateral_inplanes) != planes * block.expansion):
+            downsample = ConvModule(
+                inplanes + lateral_inplanes,
+                planes * block.expansion,
+                kernel_size=1,
+                stride=(temporal_stride, spatial_stride, spatial_stride),
+                bias=False,
+                conv_cfg=conv_cfg,
+                norm_cfg=norm_cfg,
+                act_cfg=None)
+        else:
+            downsample = None
+
+        layers = []
+        layers.append(
+            block(
+                inplanes + lateral_inplanes,
+                planes,
+                spatial_stride,
+                temporal_stride,
+                dilation,
+                downsample,
+                style=style,
+                inflate=(inflate[0] == 1),
+                inflate_style=inflate_style,
+                non_local=(non_local[0] == 1),
+                non_local_cfg=non_local_cfg,
+                conv_cfg=conv_cfg,
+                norm_cfg=norm_cfg,
+                act_cfg=act_cfg,
+                with_cp=with_cp))
+        inplanes = planes * block.expansion
+
+        for i in range(1, blocks):
+            layers.append(
+                block(
+                    inplanes,
+                    planes,
+                    1,
+                    1,
+                    dilation,
+                    style=style,
+                    inflate=(inflate[i] == 1),
+                    inflate_style=inflate_style,
+                    non_local=(non_local[i] == 1),
+                    non_local_cfg=non_local_cfg,
+                    conv_cfg=conv_cfg,
+                    norm_cfg=norm_cfg,
+                    act_cfg=act_cfg,
+                    with_cp=with_cp))
+
+        return nn.Sequential(*layers)
+
+    def inflate_weights(self, logger):
+        """Inflate the resnet2d parameters to resnet3d pathway.
+        The differences between resnet3d and resnet2d mainly lie in an extra
+        axis of conv kernel. To utilize the pretrained parameters in 2d model,
+        the weight of conv2d models should be inflated to fit in the shapes of
+        the 3d counterpart. For pathway the ``lateral_connection`` part should
+        not be inflated from 2d weights.
+        Args:
+            logger (logging.Logger): The logger used to print
+                debugging information.
+        """
+
+        state_dict_r2d = _load_checkpoint(self.pretrained)
+        if 'state_dict' in state_dict_r2d:
+            state_dict_r2d = state_dict_r2d['state_dict']
+
+        inflated_param_names = []
+        for name, module in self.named_modules():
+            if 'lateral' in name:
+                continue
+            if isinstance(module, ConvModule):
+                # we use a ConvModule to wrap conv+bn+relu layers, thus the
+                # name mapping is needed
+                if 'downsample' in name:
+                    # layer{X}.{Y}.downsample.conv->layer{X}.{Y}.downsample.0
+                    original_conv_name = name + '.0'
+                    # layer{X}.{Y}.downsample.bn->layer{X}.{Y}.downsample.1
+                    original_bn_name = name + '.1'
+                else:
+                    # layer{X}.{Y}.conv{n}.conv->layer{X}.{Y}.conv{n}
+                    original_conv_name = name
+                    # layer{X}.{Y}.conv{n}.bn->layer{X}.{Y}.bn{n}
+                    original_bn_name = name.replace('conv', 'bn')
+                if original_conv_name + '.weight' not in state_dict_r2d:
+                    logger.warning(f'Module not exist in the state_dict_r2d'
+                                   f': {original_conv_name}')
+                else:
+                    self._inflate_conv_params(module.conv, state_dict_r2d,
+                                              original_conv_name,
+                                              inflated_param_names)
+                if original_bn_name + '.weight' not in state_dict_r2d:
+                    logger.warning(f'Module not exist in the state_dict_r2d'
+                                   f': {original_bn_name}')
+                else:
+                    self._inflate_bn_params(module.bn, state_dict_r2d,
+                                            original_bn_name,
+                                            inflated_param_names)
+
+        # check if any parameters in the 2d checkpoint are not loaded
+        remaining_names = set(
+            state_dict_r2d.keys()) - set(inflated_param_names)
+        if remaining_names:
+            logger.info(f'These parameters in the 2d checkpoint are not loaded'
+                        f': {remaining_names}')
+
+    def _inflate_conv_params(self, conv3d, state_dict_2d, module_name_2d,
+                             inflated_param_names):
+        """Inflate a conv module from 2d to 3d.
+        The differences of conv modules betweene 2d and 3d in Pathway
+        mainly lie in the inplanes due to lateral connections. To fit the
+        shapes of the lateral connection counterpart, it will expand
+        parameters by concatting conv2d parameters and extra zero paddings.
+        Args:
+            conv3d (nn.Module): The destination conv3d module.
+            state_dict_2d (OrderedDict): The state dict of pretrained 2d model.
+            module_name_2d (str): The name of corresponding conv module in the
+                2d model.
+            inflated_param_names (list[str]): List of parameters that have been
+                inflated.
+        """
+        weight_2d_name = module_name_2d + '.weight'
+        conv2d_weight = state_dict_2d[weight_2d_name]
+        old_shape = conv2d_weight.shape
+        new_shape = conv3d.weight.data.shape
+        kernel_t = new_shape[2]
+
+        if new_shape[1] != old_shape[1]:
+            if new_shape[1] < old_shape[1]:
+                warnings.warn(f'The parameter of {module_name_2d} is not'
+                              'loaded due to incompatible shapes. ')
+                return
+            # Inplanes may be different due to lateral connections
+            new_channels = new_shape[1] - old_shape[1]
+            pad_shape = old_shape
+            pad_shape = pad_shape[:1] + (new_channels, ) + pad_shape[2:]
+            # Expand parameters by concat extra channels
+            conv2d_weight = torch.cat(
+                (conv2d_weight,
+                 torch.zeros(pad_shape).type_as(conv2d_weight).to(
+                     conv2d_weight.device)),
+                dim=1)
+
+        new_weight = conv2d_weight.data.unsqueeze(2).expand_as(
+            conv3d.weight) / kernel_t
+        conv3d.weight.data.copy_(new_weight)
+        inflated_param_names.append(weight_2d_name)
+
+        if getattr(conv3d, 'bias') is not None:
+            bias_2d_name = module_name_2d + '.bias'
+            conv3d.bias.data.copy_(state_dict_2d[bias_2d_name])
+            inflated_param_names.append(bias_2d_name)
+
+    def _freeze_stages(self):
+        """Prevent all the parameters from being optimized before
+        `self.frozen_stages`."""
+        if self.frozen_stages >= 0:
+            self.conv1.eval()
+            for param in self.conv1.parameters():
+                param.requires_grad = False
+
+        for i in range(1, self.frozen_stages + 1):
+            m = getattr(self, f'layer{i}')
+            m.eval()
+            for param in m.parameters():
+                param.requires_grad = False
+
+            if i != len(self.res_layers) and self.lateral:
+                # No fusion needed in the final stage
+                lateral_name = self.lateral_connections[i - 1]
+                conv_lateral = getattr(self, lateral_name)
+                conv_lateral.eval()
+                for param in conv_lateral.parameters():
+                    param.requires_grad = False
+
+    def init_weights(self, pretrained=None):
+        """Initiate the parameters either from existing checkpoint or from
+        scratch."""
+        if pretrained:
+            self.pretrained = pretrained
+
+        # Override the init_weights of i3d
+        super().init_weights()
+        for module_name in self.lateral_connections:
+            layer = getattr(self, module_name)
+            for m in layer.modules():
+                if isinstance(m, (nn.Conv3d, nn.Conv2d)):
+                    kaiming_init(m)
+
+
+pathway_cfg = {
+    'resnet3d': ResNet3dPathway,
+    # TODO: BNInceptionPathway
+}
+
+
+def build_pathway(cfg, *args, **kwargs):
+    """Build pathway.
+    Args:
+        cfg (None or dict): cfg should contain:
+            - type (str): identify conv layer type.
+    Returns:
+        nn.Module: Created pathway.
+    """
+    if not (isinstance(cfg, dict) and 'type' in cfg):
+        raise TypeError('cfg must be a dict containing the key "type"')
+    cfg_ = cfg.copy()
+
+    pathway_type = cfg_.pop('type')
+    if pathway_type not in pathway_cfg:
+        raise KeyError(f'Unrecognized pathway type {pathway_type}')
+
+    pathway_cls = pathway_cfg[pathway_type]
+    pathway = pathway_cls(*args, **kwargs, **cfg_)
+
+    return pathway
+
+
+
+"""
+CAVP: Video Encoder
+"""
+
+
+class ResNet3dSlowOnly(ResNet3dPathway):
+    """SlowOnly backbone based on ResNet3dPathway.
+    Args:
+        *args (arguments): Arguments same as :class:`ResNet3dPathway`.
+        conv1_kernel (Sequence[int]): Kernel size of the first conv layer.
+            Default: (1, 7, 7).
+        conv1_stride_t (int): Temporal stride of the first conv layer.
+            Default: 1.
+        pool1_stride_t (int): Temporal stride of the first pooling layer.
+            Default: 1.
+        inflate (Sequence[int]): Inflate Dims of each block.
+            Default: (0, 0, 1, 1).
+        **kwargs (keyword arguments): Keywords arguments for
+            :class:`ResNet3dPathway`.
+    """
+
+    def __init__(self,
+                 *args,
+                 lateral=False,
+                 conv1_kernel=(1, 7, 7),
+                 conv1_stride_t=1,
+                 pool1_stride_t=1,
+                 inflate=(0, 0, 1, 1),
+                 with_pool2=False,
+                 **kwargs):
+        super().__init__(
+            *args,
+            lateral=lateral,
+            conv1_kernel=conv1_kernel,
+            conv1_stride_t=conv1_stride_t,
+            pool1_stride_t=pool1_stride_t,
+            inflate=inflate,
+            with_pool2=with_pool2,
+            **kwargs)
+
+        assert not self.lateral
+
+
+class BasicBlock(nn.Module):
+    """Basic Block for resnet 18 and resnet 34
+    """
+
+    #BasicBlock and BottleNeck block
+    #have different output size
+    #we use class attribute expansion
+    #to distinct
+    expansion = 1
+
+    def __init__(self, in_channels, out_channels, stride=1):
+        super().__init__()
+
+        #residual function
+        self.residual_function = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels * BasicBlock.expansion, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels * BasicBlock.expansion)
+        )
+
+        #shortcut
+        self.shortcut = nn.Sequential()
+
+        #the shortcut output dimension is not the same with residual function
+        #use 1*1 convolution to match the dimension
+        if stride != 1 or in_channels != BasicBlock.expansion * out_channels:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels * BasicBlock.expansion, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels * BasicBlock.expansion)
+            )
+
+    def forward(self, x):
+        return nn.ReLU(inplace=True)(self.residual_function(x) + self.shortcut(x))
+
+class BottleNeck(nn.Module):
+    """Residual block for resnet over 50 layers
+    """
+    expansion = 4
+    def __init__(self, in_channels, out_channels, stride=1):
+        super().__init__()
+        self.residual_function = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, stride=stride, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels * BottleNeck.expansion, kernel_size=1, bias=False),
+            nn.BatchNorm2d(out_channels * BottleNeck.expansion),
+        )
+
+        self.shortcut = nn.Sequential()
+
+        if stride != 1 or in_channels != out_channels * BottleNeck.expansion:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels * BottleNeck.expansion, stride=stride, kernel_size=1, bias=False),
+                nn.BatchNorm2d(out_channels * BottleNeck.expansion)
+            )
+
+    def forward(self, x):
+        return nn.ReLU(inplace=True)(self.residual_function(x) + self.shortcut(x))
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, num_block, num_classes=100, truncate_sec=4):
+        super().__init__()
+
+        self.in_channels = 64
+
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(1, 64, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True))
+        #we use a different inputsize than the original paper
+        #so conv2_x's stride is 1
+        self.conv2_x = self._make_layer(block, 64, num_block[0], 2)
+        self.conv3_x = self._make_layer(block, 128, num_block[1], 2)
+        self.conv4_x = self._make_layer(block, 256, num_block[2], 2)
+        self.conv5_x = self._make_layer(block, 512, num_block[3], 2)
+
+        assert truncate_sec == 4 or truncate_sec == 8 or truncate_sec == 10
+        if truncate_sec == 4:
+            self.avg_pool = nn.AdaptiveAvgPool2d((1, 16))
+        elif truncate_sec == 8:
+            self.avg_pool = nn.AdaptiveAvgPool2d((1, 32))
+        elif truncate_sec == 10:
+            self.avg_pool = nn.AdaptiveAvgPool2d((1, 40))
+
+
+    def _make_layer(self, block, out_channels, num_blocks, stride):
+        """make resnet layers(by layer i didnt mean this 'layer' was the
+        same as a neuron netowork layer, ex. conv layer), one layer may
+        contain more than one residual block
+        Args:
+            block: block type, basic block or bottle neck block
+            out_channels: output depth channel number of this layer
+            num_blocks: how many blocks per layer
+            stride: the stride of the first block of this layer
+        Return:
+            return a resnet layer
+        """
+
+        # we have num_block blocks per layer, the first block
+        # could be 1 or 2, other blocks would always be 1
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_channels, out_channels, stride))
+            self.in_channels = out_channels * block.expansion
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        output = self.conv1(x)
+        output = self.conv2_x(output)
+        output = self.conv3_x(output)
+        output = self.conv4_x(output)
+        output = self.conv5_x(output)
+        output = self.avg_pool(output)
+        bs, c, _, t = output.shape
+        output = output.view(bs, c, t)
+        return output
+
+
+
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable):
+            return x
+        return tuple(repeat(x, n))
+    return parse
+
+
+
+"""
+Cnn14: Spec Encoder
+"""
+
+def interpolate(x, ratio):
+    """Interpolate data in time domain. This is used to compensate the
+    resolution reduction in downsampling of a CNN.
+    Args:
+      x: (batch_size, time_steps, classes_num)
+      ratio: int, ratio to interpolate
+    Returns:
+      upsampled: (batch_size, time_steps * ratio, classes_num)
+    """
+    (batch_size, time_steps, classes_num) = x.shape
+    upsampled = x[:, :, None, :].repeat(1, 1, ratio, 1)
+    upsampled = upsampled.reshape(batch_size, time_steps * ratio, classes_num)
+    return upsampled
+
+def init_bn(bn):
+    """Initialize a Batchnorm layer. """
+    bn.bias.data.fill_(0.)
+    bn.weight.data.fill_(1.)
+
+
+def init_layer(layer):
+    """Initialize a Linear or Convolutional layer. """
+    nn.init.xavier_uniform_(layer.weight)
+
+    if hasattr(layer, 'bias'):
+        if layer.bias is not None:
+            layer.bias.data.fill_(0.)
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+
+        super(ConvBlock, self).__init__()
+
+        self.conv1 = nn.Conv2d(in_channels=in_channels,
+                              out_channels=out_channels,
+                              kernel_size=(3, 3), stride=(1, 1),
+                              padding=(1, 1), bias=False)
+
+        self.conv2 = nn.Conv2d(in_channels=out_channels,
+                              out_channels=out_channels,
+                              kernel_size=(3, 3), stride=(1, 1),
+                              padding=(1, 1), bias=False)
+
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+
+        self.init_weight()
+
+    def init_weight(self):
+        init_layer(self.conv1)
+        init_layer(self.conv2)
+        init_bn(self.bn1)
+        init_bn(self.bn2)
+
+
+    def forward(self, input, pool_size=(2, 2), pool_type='avg'):
+
+        x = input
+        x = F.relu_(self.bn1(self.conv1(x)))
+        x = F.relu_(self.bn2(self.conv2(x)))
+        if pool_type == 'max':
+            x = F.max_pool2d(x, kernel_size=pool_size)
+        elif pool_type == 'avg':
+            x = F.avg_pool2d(x, kernel_size=pool_size)
+        elif pool_type == 'avg+max':
+            x1 = F.avg_pool2d(x, kernel_size=pool_size)
+            x2 = F.max_pool2d(x, kernel_size=pool_size)
+            x = x1 + x2
+        else:
+            raise Exception('Incorrect argument!')
+
+        return x
+
+
+
+class Cnn14(nn.Module):
+    def __init__(self, embed_dim, enable_fusion=False, fusion_type='None'):
+        super(Cnn14, self).__init__()
+
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
+        
+        self.bn = nn.BatchNorm2d(128)
+
+        if (self.enable_fusion) and (self.fusion_type == 'channel_map'):
+            self.conv_block1 = ConvBlock(in_channels=4, out_channels=64)
+        else:
+            self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
+        self.conv_block2 = ConvBlock(in_channels=64, out_channels=128)
+        self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)
+        self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)
+        self.conv_block5 = ConvBlock(in_channels=512, out_channels=1024)
+        self.conv_block6 = ConvBlock(in_channels=1024, out_channels=2048)
+
+        self.fc1 = nn.Linear(2048, 2048, bias=True)
+        self.final_project = nn.Linear(2048, embed_dim, bias=True)
+
+        self.init_weight()
+
+    def init_weight(self):
+        init_bn(self.bn)
+        init_layer(self.fc1)
+        init_layer(self.final_project)
+
+    def forward(self, input, mixup_lambda=None, device=None):
+        """
+        Input: (batch_size, data_length)"""
+
+        x = input
+        x = x.transpose(1, 3)
+        x = self.bn(x)
+        x = x.transpose(1, 3)
+
+        x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block2(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block3(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block4(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block5(x, pool_size=(1, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block6(x, pool_size=(1, 1), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = torch.mean(x, dim=3)
+
+        latent_x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1)
+        latent_x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1)
+        latent_x = latent_x1 + latent_x2
+        latent_x = latent_x.transpose(1, 2)
+        latent_x = F.relu_(self.fc1(latent_x))
+        x = F.relu_(self.fc1(latent_x))
+        output = self.final_project(x)
+        return output
+
+

cavp_util.py

@@ -0,0 +1,150 @@
+import torch
+import subprocess
+from pathlib import Path
+import os
+import cv2
+import numpy as np
+import torchvision.transforms as transforms
+from PIL import Image
+from tqdm import tqdm
+from omegaconf import OmegaConf
+import importlib
+
+
+def which_ffmpeg() -> str:
+    '''Determines the path to ffmpeg library
+
+    Returns:
+        str -- path to the library
+    '''
+    result = subprocess.run(['which', 'ffmpeg'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    ffmpeg_path = result.stdout.decode('utf-8').replace('\n', '')
+    return ffmpeg_path
+
+def reencode_video_with_diff_fps(video_path: str, tmp_path: str, extraction_fps: int, start_second, truncate_second) -> str:
+    '''Reencodes the video given the path and saves it to the tmp_path folder.
+
+    Args:
+        video_path (str): original video
+        tmp_path (str): the folder where tmp files are stored (will be appended with a proper filename).
+        extraction_fps (int): target fps value
+
+    Returns:
+        str: The path where the tmp file is stored. To be used to load the video from
+    '''
+    assert which_ffmpeg() != '', 'Is ffmpeg installed? Check if the conda environment is activated.'
+    os.makedirs(tmp_path, exist_ok=True)
+
+    # form the path to tmp directory
+    new_path = os.path.join(tmp_path, f'{Path(video_path).stem}_new_fps_{str(extraction_fps)}_truncate_{start_second}_{truncate_second}.mp4')
+    cmd = f'{which_ffmpeg()} -hide_banner -loglevel panic '
+    cmd += f'-y -ss {start_second} -t {truncate_second} -i {video_path} -an -filter:v fps=fps={extraction_fps} {new_path}'
+    subprocess.call(cmd.split())
+    return new_path
+
+def instantiate_from_config(config, reload=False):
+    if not "target" in config:
+        if config == '__is_first_stage__':
+            return None
+        elif config == "__is_unconditional__":
+            return None
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"], reload=reload)(**config.get("params", dict()))
+
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+class Extract_CAVP_Features(torch.nn.Module):
+
+    def __init__(self, device=None, tmp_path="./", video_shape=(224,224), config_path=None, ckpt_path=None):
+        super(Extract_CAVP_Features, self).__init__()
+        self.fps = 4
+        self.batch_size = 40
+        self.device = device
+        self.tmp_path = tmp_path
+
+        # Initalize CAVP model:
+        config = OmegaConf.load(config_path)
+        self.stage1_model = instantiate_from_config(config.model).to(device)
+
+        # Loading Model from:
+        assert ckpt_path is not None
+        self.init_first_from_ckpt(ckpt_path)
+        self.stage1_model.eval()
+        
+        # Transform:
+        self.img_transform = transforms.Compose([
+            transforms.Resize(video_shape),
+            transforms.ToTensor(),
+        ])
+
+
+    def init_first_from_ckpt(self, path):
+        model = torch.load(path, map_location="cpu")
+        if "state_dict" in list(model.keys()):
+            model = model["state_dict"]
+        # Remove: module prefix
+        new_model = {}
+        for key in model.keys():
+            new_key = key.replace("module.","")
+            new_model[new_key] = model[key]
+        self.stage1_model.load_state_dict(new_model, strict=False)
+
+
+    @torch.no_grad()
+    def forward(self, video_path, tmp_path="./tmp_folder"):
+        start_second = 0
+        truncate_second = 10
+        self.tmp_path = tmp_path
+
+        # Load the video, change fps:
+        video_path_low_fps = reencode_video_with_diff_fps(video_path, self.tmp_path, self.fps, start_second, truncate_second)
+
+        # read the video:
+        cap = cv2.VideoCapture(video_path_low_fps)
+
+        feat_batch_list = []
+        video_feats = []
+        first_frame = True
+        # pbar = tqdm(cap.get(7))
+        i = 0
+        while cap.isOpened():
+            i += 1
+            # pbar.set_description("Processing Frames: {} Total: {}".format(i, cap.get(7)))
+            frames_exists, rgb = cap.read()
+            
+            if first_frame:
+                if not frames_exists:
+                    continue
+            first_frame = False
+
+            if frames_exists:
+                rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
+                rgb_tensor = self.img_transform(Image.fromarray(rgb)).unsqueeze(0).to(self.device)
+                feat_batch_list.append(rgb_tensor)      # 32 x 3 x 224 x 224
+                
+                # Forward:
+                if len(feat_batch_list) == self.batch_size:
+                    # Stage1 Model:
+                    input_feats = torch.cat(feat_batch_list,0).unsqueeze(0).to(self.device)
+                    contrastive_video_feats = self.stage1_model.encode_video(input_feats, normalize=True, pool=False)
+                    video_feats.extend(contrastive_video_feats.detach().cpu().numpy())
+                    feat_batch_list = []
+            else:
+                if len(feat_batch_list) != 0:
+                    input_feats = torch.cat(feat_batch_list,0).unsqueeze(0).to(self.device)
+                    contrastive_video_feats = self.stage1_model.encode_video(input_feats, normalize=True, pool=False)
+                    video_feats.extend(contrastive_video_feats.detach().cpu().numpy())
+                cap.release()
+                break
+
+        # Remove the file
+        os.remove(video_path_low_fps)
+        video_contrastive_feats = np.concatenate(video_feats)
+        return video_contrastive_feats
+

dataset.py

@@ -0,0 +1,273 @@
+import os
+
+import numpy as np
+import torch
+import random
+import math
+
+from torch.utils.data import Dataset
+
+class audio_video_spec_fullset_Dataset(Dataset):
+    # Only Load audio dataset: for training Stage1: Audio Npy Dataset
+    def __init__(self, split, data_dir):
+        super().__init__()
+        debug_num=False
+
+        if split == "train":
+            self.split = "train"
+        elif split == "valid" or split == 'test':
+            self.split = "test"
+
+        # Default params:
+        self.min_duration = 2
+        self.sr = 16000
+        self.duration = 10
+        self.truncate = 130560
+        self.fps = 4
+        self.fix_frames = False
+        self.hop_len = 160
+        self.onset_truncate = 120
+
+
+        # spec_dir: spectrogram path
+        # feat_dir: CAVP feature path
+        # fbank_dir: fbank feature path
+        # onset_dir: onset feature path
+        dataset_spec_dir = os.path.join(data_dir, "melspec", self.split)
+        dataset_feat_dir = os.path.join(data_dir, "cavp_feats", self.split)
+        dataset_fbank_dir = os.path.join(data_dir, "fbank", self.split)
+        dataset_onset_dir = os.path.join(data_dir, "onset_feats", "train")
+        list_onset = os.listdir(dataset_onset_dir)
+        list_onset = list(map(lambda x: x.split('.')[0], list_onset))
+
+
+        with open(os.path.join(data_dir, '{}_list.txt'.format(self.split)), "r") as f:
+            data_list = f.readlines()
+            data_list = list(map(lambda x: x.strip(), data_list))
+            data_list = list(set(data_list) & set(list_onset))
+
+            spec_list = list(map(lambda x: os.path.join(dataset_spec_dir, x) + ".npy", data_list))      # spec
+            feat_list = list(map(lambda x: os.path.join(dataset_feat_dir, x) + ".npz",     data_list))      # feat
+            fbank_list = list(map(lambda x: os.path.join(dataset_fbank_dir, x) + ".npy",     data_list))      # fbank
+            onset_list = list(map(lambda x: os.path.join(dataset_onset_dir, x) + ".npy",     data_list))      # onset
+
+
+        # Merge Data:
+        self.data_list = data_list
+        self.spec_list = spec_list 
+        self.feat_list = feat_list
+        self.fbank_list = fbank_list
+        self.onset_list = onset_list
+
+
+        assert len(self.data_list) == len(self.spec_list) == len(self.feat_list)
+
+
+        shuffle_idx = np.random.permutation(np.arange(len(self.data_list)))
+        self.data_list = [self.data_list[i] for i in shuffle_idx]
+        self.spec_list = [self.spec_list[i] for i in shuffle_idx]
+        self.feat_list = [self.feat_list[i] for i in shuffle_idx]
+        self.fbank_list = [self.fbank_list[i] for i in shuffle_idx]
+        self.onset_list = [self.onset_list[i] for i in shuffle_idx]
+
+
+        if debug_num:
+            self.data_list = self.data_list[:debug_num]
+            self.spec_list = self.spec_list[:debug_num]
+            self.feat_list = self.feat_list[:debug_num]
+            self.fbank_list = self.fbank_list[:debug_num]
+            self.onset_list = self.onset_list[:debug_num]
+        print('Split: {}  Sample Num: {}'.format(split, len(self.data_list)))
+
+
+    def __len__(self):
+        return len(self.data_list)
+    
+
+    def load_spec_and_feat(self, spec_path, video_feat_path, fbank_path, onset_path):
+        """Load audio spec and video feat"""
+        spec_raw = np.load(spec_path).astype(np.float32).T                    # channel: 1
+        video_feat = np.load(video_feat_path)['arr_0'].astype(np.float32)
+        fbank = np.load(fbank_path).astype(np.float32)
+        onset = np.load(onset_path).astype(np.float32).reshape(-1)
+
+
+        # Padding the samples:
+        spec_len = self.sr * self.duration / self.hop_len
+        fbank_len = int(spec_len / spec_raw.shape[1] * len(fbank))
+        if spec_raw.shape[1] < spec_len:
+            fbank = np.tile(fbank, (math.ceil(spec_len / spec_raw.shape[1]), 1))
+            spec_raw = np.tile(spec_raw, math.ceil(spec_len / spec_raw.shape[1]))
+        spec_raw = spec_raw[:, :int(spec_len)]
+        fbank = fbank[:fbank_len]
+        
+        feat_len = self.fps * self.duration
+        if video_feat.shape[0] < feat_len:
+            video_feat = np.tile(video_feat, (math.ceil(feat_len / video_feat.shape[0]), 1))
+        video_feat = video_feat[:int(feat_len)]
+
+        onset_len = 15 * self.duration
+        if onset.shape[0] < onset_len:
+            onset = np.tile(onset, (math.ceil(onset_len / onset.shape[0])))
+        onset = onset[:int(onset_len)]
+
+        return spec_raw, video_feat, fbank, onset
+
+
+    def mix_audio_and_feat(self, spec1=None, spec2=None, video_feat1=None, video_feat2=None, fbank1=None, fbank2=None, onset1=None, onset2=None, video_info_dict={}, mode='single'):
+        """ Return Mix Spec and Mix video feat"""
+        if mode == "single":
+            # spec1:
+            if not self.fix_frames:
+                start_idx = random.randint(0, self.sr * self.duration - self.truncate - 1)  # audio start
+            else:
+                start_idx = 0
+
+            start_frame = int(self.fps * start_idx / self.sr)
+            truncate_frame = int(self.fps * self.truncate / self.sr)
+
+            start_onset = int(15 * start_idx / self.sr)
+            truncate_onset = self.onset_truncate
+
+            # Spec Start & Truncate:
+            spec_start = int(start_idx / self.hop_len)
+            spec_truncate = int(self.truncate / self.hop_len)
+
+            # Fbank Start & Truncate:
+            fbank_start = int((spec_start / spec1.shape[1]) * len(fbank1))
+            fbank_truncate = int((spec_truncate / spec1.shape[1]) * len(fbank1))
+
+            spec1 = spec1[:, spec_start : spec_start + spec_truncate]
+            video_feat1 = video_feat1[start_frame: start_frame + truncate_frame]
+            fbank1 = fbank1[fbank_start: fbank_start + fbank_truncate]
+            onset1 = onset1[start_onset: start_onset + truncate_onset]
+
+            # info_dict:
+            video_info_dict['video_time1'] = str(start_frame) + '_' + str(start_frame+truncate_frame)   # Start frame, end frame
+            video_info_dict['video_time2'] = ""
+            return spec1, video_feat1, fbank1, onset1, video_info_dict
+        
+        elif mode == "concat":
+            total_spec_len = int(self.truncate / self.hop_len)
+            # Random Trucate len:
+            spec1_truncate_len = random.randint(self.min_duration * self.sr // self.hop_len, total_spec_len - self.min_duration * self.sr // self.hop_len - 1)
+            spec2_truncate_len = total_spec_len - spec1_truncate_len
+
+            # Sample spec clip:
+            spec_start1 = random.randint(0, total_spec_len - spec1_truncate_len - 1)
+            spec_start2 = random.randint(0, total_spec_len - spec2_truncate_len - 1)
+            spec_end1, spec_end2 = spec_start1 + spec1_truncate_len, spec_start2 + spec2_truncate_len
+
+            start1_fbank, truncate1_fbank = int((spec_start1 / spec1.shape[1]) * len(fbank1)), int((spec1_truncate_len / spec1.shape[1]) * len(fbank1))
+            start2_fbank, truncate2_fbank = int((spec_start2 / spec2.shape[1]) * len(fbank2)), int((spec2_truncate_len / spec2.shape[1]) * len(fbank2))
+
+            # concat spec:
+            spec1, spec2 = spec1[:, spec_start1 : spec_end1], spec2[:, spec_start2 : spec_end2]
+            concat_audio_spec = np.concatenate([spec1, spec2], axis=1)  
+
+            # Concat Video Feat:
+            start1_frame, truncate1_frame = int(self.fps * spec_start1 * self.hop_len / self.sr), int(self.fps * spec1_truncate_len * self.hop_len / self.sr)
+            start2_frame, truncate2_frame = int(self.fps * spec_start2 * self.hop_len / self.sr), int(self.fps * self.truncate / self.sr) - truncate1_frame
+            video_feat1, video_feat2 = video_feat1[start1_frame : start1_frame + truncate1_frame], video_feat2[start2_frame : start2_frame + truncate2_frame]
+            concat_video_feat = np.concatenate([video_feat1, video_feat2])
+
+            # Concat Fbank:
+            fbank1, fbank2 = fbank1[start1_fbank : start1_fbank + truncate1_fbank], fbank2[start2_fbank : start2_fbank + truncate2_fbank]
+            concat_fbank = np.concatenate([fbank1, fbank2])
+
+            # Concat Onset:
+            start1_onset, truncate1_onset = int(15 * spec_start1 * self.hop_len / self.sr), int(15 * spec1_truncate_len * self.hop_len / self.sr)
+            start2_onset, truncate2_onset = int(15 * spec_start2 * self.hop_len / self.sr), self.onset_truncate - truncate1_onset
+            onset_feat1, onset_feat2 = onset1[start1_onset : start1_onset + truncate1_onset], onset2[start2_onset : start2_onset + truncate2_onset]
+            concat_onset = np.concatenate([onset_feat1, onset_feat2])
+
+            video_info_dict['video_time1'] = str(start1_frame) + '_' + str(start1_frame+truncate1_frame)   # Start frame, end frame
+            video_info_dict['video_time2'] = str(start2_frame) + '_' + str(start2_frame+truncate2_frame)
+            return concat_audio_spec, concat_video_feat, concat_fbank, concat_onset, video_info_dict
+
+
+
+    def __getitem__(self, idx):
+        audio_name1 = self.data_list[idx]
+        spec_npy_path1 = self.spec_list[idx]
+        video_feat_path1 = self.feat_list[idx]
+        fbank_path1 = self.fbank_list[idx]
+        onset_path1 = self.onset_list[idx]
+
+
+        # select other video:
+        flag = False
+        if random.uniform(0, 1) < 0.5:
+            flag = True
+            random_idx = idx
+            while random_idx == idx:
+                random_idx = random.randint(0, len(self.data_list)-1)
+            audio_name2 = self.data_list[random_idx]
+            spec_npy_path2 = self.spec_list[random_idx]
+            video_feat_path2 = self.feat_list[random_idx]
+            fbank_path2 = self.fbank_list[random_idx]
+            onset_path2 = self.onset_list[random_idx]
+
+
+        # Load the Spec and Feat:
+        spec1, video_feat1, fbank1, onset1 = self.load_spec_and_feat(spec_npy_path1, video_feat_path1, fbank_path1, onset_path1)
+
+        if flag:
+            spec2, video_feat2, fbank2, onset2 = self.load_spec_and_feat(spec_npy_path2, video_feat_path2, fbank_path2, onset_path2)
+            video_info_dict = {'audio_name1':audio_name1, 'audio_name2': audio_name2}
+            mix_spec, mix_video_feat, mix_fbank, mix_onset, mix_info = self.mix_audio_and_feat(spec1, spec2, video_feat1, video_feat2, fbank1, fbank2, onset1, onset2, video_info_dict, mode='concat')
+        else:
+            video_info_dict = {'audio_name1':audio_name1, 'audio_name2': ""}
+            mix_spec, mix_video_feat, mix_fbank, mix_onset, mix_info = self.mix_audio_and_feat(spec1=spec1, video_feat1=video_feat1, fbank1=fbank1, onset1=onset1, video_info_dict=video_info_dict, mode='single')
+
+
+        norm_mean = -4.268
+        norm_std = 4.569
+        target_length = 1024
+        n_frames = mix_fbank.shape[0]
+        mix_fbank = torch.from_numpy(mix_fbank).contiguous()
+        diff = target_length - n_frames
+        if diff > 0:
+            m = torch.nn.ZeroPad2d((0, 0, 0, diff))
+            mix_fbank = m(mix_fbank)
+            mix_fbank[n_frames:] = (mix_fbank[n_frames:] - norm_mean) / (norm_std * 2)
+        elif diff < 0:
+            mix_fbank = mix_fbank[0:target_length, :]
+
+        mix_spec = mix_spec[None]
+        mix_spec = torch.from_numpy(mix_spec).contiguous()
+        mix_video_feat = torch.from_numpy(mix_video_feat).contiguous()
+        mix_onset = torch.from_numpy(mix_onset).contiguous()
+
+        data_dict = {}
+        data_dict['mix_spec'] = mix_spec
+        data_dict['mix_video_feat'] = mix_video_feat
+        data_dict['mix_fbank'] = mix_fbank
+        data_dict['mix_onset'] = mix_onset
+        data_dict['mix_info_dict'] = mix_info     
+        return data_dict
+
+
+
+class audio_video_spec_fullset_Dataset_Train(audio_video_spec_fullset_Dataset):
+    def __init__(self, data_dir):
+        super().__init__(split='train', data_dir=data_dir)
+
+
+
+def collate_fn_taro(data):
+    mix_spec = torch.stack([example["mix_spec"] for example in data])
+    mix_video_feat = torch.stack([example["mix_video_feat"] for example in data])
+    mix_fbank = torch.stack([example["mix_fbank"] for example in data])
+    mix_onset = torch.stack([example["mix_onset"] for example in data])
+    mix_info_dict = [example["mix_info_dict"] for example in data]   
+
+    return {
+        "mix_spec": mix_spec,
+        "mix_video_feat": mix_video_feat,
+        "mix_fbank": mix_fbank,
+        "mix_onset": mix_onset,
+        "mix_info_dict": mix_info_dict,
+    }
+
+

infer.py

@@ -0,0 +1,151 @@
+import torch
+import os
+import numpy as np
+import random
+import soundfile as sf
+import ffmpeg
+
+from argparse import ArgumentParser
+from diffusers import AudioLDM2Pipeline
+from models import MMDiT
+from samplers import euler_sampler, euler_maruyama_sampler
+from cavp_util import Extract_CAVP_Features
+from onset_util import VideoOnsetNet, extract_onset
+
+def set_global_seed(seed):
+    np.random.seed(seed % (2**32))
+    random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+def main():
+    parser = ArgumentParser(description="Inference script parameters")
+    parser.add_argument("--video_path", type=str, default="./test.mp4", required=True, help="Path to the input video file")
+    parser.add_argument("--save_folder_path", type=str, default="./output", help="Folder to save output files")
+    parser.add_argument("--cavp_config_path", type=str, default="./cavp.yaml", help="Path to CAVP config file")
+    parser.add_argument("--cavp_ckpt_path", type=str, default="./cavp_epoch66.ckpt", help="Path to CAVP checkpoint file")
+    parser.add_argument("--onset_ckpt_path", type=str, default="./onset_model.ckpt", help="Path to onset model checkpoint file")
+    parser.add_argument("--model_ckpt_path", type=str, default="./taro_ckpt.pt", help="Path to MMDiT model checkpoint file")
+
+    args = parser.parse_args()
+    os.makedirs(args.save_folder_path, exist_ok=True)
+
+    seed = 0
+    set_global_seed(seed)
+    torch.set_grad_enabled(False)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    weight_dtype = torch.bfloat16
+
+    # Load models
+    extract_cavp = Extract_CAVP_Features(device=device, config_path=args.cavp_config_path, ckpt_path=args.cavp_ckpt_path)
+
+    # Load the pre-trained onset detection model
+    state_dict = torch.load(args.onset_ckpt_path)["state_dict"]
+    new_state_dict = {}
+    for key, value in state_dict.items():
+        if "model.net.model" in key:
+            new_key = key.replace("model.net.model", "net.model")  # Adjust the key as needed
+        elif "model.fc." in key:
+            new_key = key.replace("model.fc", "fc")  # Adjust the key as needed
+        new_state_dict[new_key] = value
+    onset_model = VideoOnsetNet(False).to(device)
+    onset_model.load_state_dict(new_state_dict)
+    onset_model.eval()
+
+    model = MMDiT(
+        adm_in_channels=120,
+        z_dims = [768],
+        encoder_depth=4,
+    ).to(device)
+
+    state_dict = torch.load(args.model_ckpt_path, map_location=device)['ema']
+    model.load_state_dict(state_dict)
+    model.eval()
+    model.to(weight_dtype)
+    model_audioldm = AudioLDM2Pipeline.from_pretrained("cvssp/audioldm2")
+    vae = model_audioldm.vae.to(device)
+    vae.eval()
+
+    vocoder = model_audioldm.vocoder.to(device)
+    
+    # Extract Features
+    video_name = os.path.basename(args.video_path).split(".")[0]
+
+    cavp_feats = extract_cavp(args.video_path, tmp_path=args.save_folder_path)
+    onset_feats = extract_onset(args.video_path, onset_model, tmp_path=args.save_folder_path, device=device)
+
+    # Parameters for inference
+    sr = 16000
+    truncate = 131072
+    fps = 4
+
+    truncate_frame = int(fps * truncate / sr)
+    truncate_onset = 120
+
+    cfg_scale = 8
+    mode = "sde"
+    num_steps = 25
+    heun = False
+    guidance_low = 0.0
+    guidance_high = 0.7
+    path_type = "linear"
+
+    latent_size = (204, 16)
+    latents_scale = torch.tensor(
+        [0.18215, 0.18215, 0.18215, 0.18215, 0.18215, 0.18215, 0.18215, 0.18215]
+        ).view(1, 8, 1, 1).to(device)
+
+    # Start inference
+    video_feats = torch.from_numpy(cavp_feats[:truncate_frame]).unsqueeze(0).to(device).to(weight_dtype)
+    onset_feats = torch.from_numpy(onset_feats[:truncate_onset]).unsqueeze(0).to(device).to(weight_dtype)
+
+    z = torch.randn(len(video_feats), model.in_channels, latent_size[0], latent_size[1], device=device).to(weight_dtype)
+
+    # Sample audios
+    sampling_kwargs = dict(
+        model=model, 
+        latents=z,
+        y=onset_feats,
+        context=video_feats,
+        num_steps=num_steps, 
+        heun=heun,
+        cfg_scale=cfg_scale,
+        guidance_low=guidance_low,
+        guidance_high=guidance_high,
+        path_type=path_type,
+    )
+
+    with torch.no_grad():
+        if mode == "sde":
+            samples = euler_maruyama_sampler(**sampling_kwargs)
+        elif mode == "ode":
+            samples = euler_sampler(**sampling_kwargs)
+        else:
+            raise NotImplementedError()
+
+        samples = vae.decode(samples / latents_scale).sample
+        wav_samples = vocoder(samples.squeeze()).detach().cpu().numpy()
+
+        # Save the audio
+        sf.write(os.path.join(args.save_folder_path, video_name + ".wav"), wav_samples, sr)
+
+        # Save the video with the generated audio
+        trimmed_video_file_path = os.path.join(args.save_folder_path, video_name + "_trimmed.mp4")
+        trimmed_audio_file_path = os.path.join(args.save_folder_path, video_name + ".wav")
+        output_path = os.path.join(args.save_folder_path, video_name + "_wa.mp4")
+
+        # Trim the video to match the audio duration
+        ffmpeg.input(args.video_path, ss=0, t=truncate / sr).output(trimmed_video_file_path, vcodec='libx264', an=None).run(overwrite_output=True)
+
+        # Combine trimmed video and generated audio
+        input_video = ffmpeg.input(trimmed_video_file_path)
+        input_audio = ffmpeg.input(trimmed_audio_file_path)
+        ffmpeg.output(input_video, input_audio, output_path, vcodec='libx264', acodec='aac', strict='experimental').run(overwrite_output=True)
+        os.remove(trimmed_video_file_path)
+
+    print("========================================FINISH INFERENCE===========================================")
+
+if __name__ == "__main__":
+    main()
+

loss.py

@@ -0,0 +1,94 @@
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+def mean_flat(x):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return torch.mean(x, dim=list(range(1, len(x.size()))))
+
+def sum_flat(x):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return torch.sum(x, dim=list(range(1, len(x.size()))))
+
+class SILoss:
+    def __init__(
+            self,
+            prediction='v',
+            path_type="linear",
+            weighting="uniform",
+            encoders=[], 
+            accelerator=None, 
+            latents_scale=None, 
+            latents_bias=None,
+            ):
+        self.prediction = prediction
+        self.weighting = weighting
+        self.path_type = path_type
+        self.encoders = encoders
+        self.accelerator = accelerator
+        self.latents_scale = latents_scale
+        self.latents_bias = latents_bias
+
+    def interpolant(self, t):
+        if self.path_type == "linear":
+            alpha_t = 1 - t
+            sigma_t = t
+            d_alpha_t = -1
+            d_sigma_t =  1
+        elif self.path_type == "cosine":
+            alpha_t = torch.cos(t * np.pi / 2)
+            sigma_t = torch.sin(t * np.pi / 2)
+            d_alpha_t = -np.pi / 2 * torch.sin(t * np.pi / 2)
+            d_sigma_t =  np.pi / 2 * torch.cos(t * np.pi / 2)
+        else:
+            raise NotImplementedError()
+
+        return alpha_t, sigma_t, d_alpha_t, d_sigma_t
+
+    def __call__(self, model, images, model_kwargs=None, zs=None):
+        if model_kwargs == None:
+            model_kwargs = {}
+        # sample timesteps
+        if self.weighting == "uniform":
+            time_input = torch.rand((images.shape[0], 1, 1, 1))
+        elif self.weighting == "lognormal":
+            # sample timestep according to log-normal distribution of sigmas following EDM
+            rnd_normal = torch.randn((images.shape[0], 1 ,1, 1))
+            sigma = rnd_normal.exp()
+            if self.path_type == "linear":
+                time_input = sigma / (1 + sigma)
+            elif self.path_type == "cosine":
+                time_input = 2 / np.pi * torch.atan(sigma)
+                
+        time_input = time_input.to(device=images.device, dtype=images.dtype)
+        
+        noises = torch.randn_like(images)
+        alpha_t, sigma_t, d_alpha_t, d_sigma_t = self.interpolant(time_input)
+            
+        model_input = alpha_t * images + sigma_t * noises
+        if self.prediction == 'v':
+            model_target = d_alpha_t * images + d_sigma_t * noises
+        else:
+            raise NotImplementedError() # TODO: add x or eps prediction
+        model_output, zs_tilde  = model(model_input, time_input.flatten(), **model_kwargs)
+        denoising_loss = mean_flat((model_output - model_target) ** 2)
+
+        epsilon = 1e-8
+        t_weight = torch.sigmoid(torch.log(alpha_t / (sigma_t + epsilon))).squeeze()
+
+        # projection loss
+        proj_loss = 0.
+        if len(zs) > 0:
+            bsz = zs[0].shape[0]
+            for i, (z, z_tilde) in enumerate(zip(zs, zs_tilde)):
+                for j, (z_j, z_tilde_j) in enumerate(zip(z, z_tilde)):
+                    z_tilde_j = torch.nn.functional.normalize(z_tilde_j, dim=-1) 
+                    z_j = torch.nn.functional.normalize(z_j, dim=-1)
+                    proj_loss += mean_flat(-(z_j * z_tilde_j).sum(dim=-1)) * t_weight[j]
+            proj_loss /= (len(zs) * bsz)
+
+        return denoising_loss, proj_loss

models.py

@@ -0,0 +1,747 @@
+import math
+from typing import Dict, Optional
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+
+import torch, math
+from torch import nn
+
+def prob_mask_like(shape, prob, device):
+    if prob == 1:
+        return torch.ones(shape, device = device, dtype = torch.bool)
+    elif prob == 0:
+        return torch.zeros(shape, device = device, dtype = torch.bool)
+    else:
+        return torch.zeros(shape, device = device).float().uniform_(0, 1) < prob
+
+
+def attention(q, k, v, heads, mask=None):
+    """Convenience wrapper around a basic attention operation"""
+    b, _, dim_head = q.shape
+    dim_head //= heads
+    q, k, v = map(lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2), (q, k, v))
+    out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
+    return out.transpose(1, 2).reshape(b, -1, heads * dim_head)
+
+
+class Mlp(nn.Module):
+    """ MLP as used in Vision Transformer, MLP-Mixer and related networks"""
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, bias=True, dtype=None, device=None):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias, dtype=dtype, device=device)
+        self.act = act_layer
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias, dtype=dtype, device=device)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        return x
+
+def build_mlp(hidden_size, projector_dim, z_dim):
+    return nn.Sequential(
+                nn.Conv1d(in_channels=816, out_channels=416, kernel_size=1), 
+                nn.SiLU(),
+                nn.Linear(hidden_size, projector_dim),
+                nn.SiLU(),
+                nn.Linear(projector_dim, projector_dim),
+                nn.SiLU(),
+                nn.Linear(projector_dim, z_dim),
+            )
+
+class PatchEmbed(nn.Module):
+    """ 2D Image to Patch Embedding"""
+    def __init__(
+            self,
+            img_size: Optional[int] = 224,
+            patch_size: int = 16,
+            in_chans: int = 3,
+            embed_dim: int = 768,
+            flatten: bool = True,
+            bias: bool = True,
+            strict_img_size: bool = True,
+            dynamic_img_pad: bool = False,
+            dtype=None,
+            device=None,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        if img_size is not None:
+            self.img_size = img_size
+            self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+            self.num_patches = self.grid_size[0] * self.grid_size[1]
+        else:
+            self.img_size = None
+            self.grid_size = None
+            self.num_patches = None
+
+        # flatten spatial dim and transpose to channels last, kept for bwd compat
+        self.flatten = flatten
+        self.strict_img_size = strict_img_size
+        self.dynamic_img_pad = dynamic_img_pad
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias, dtype=dtype, device=device)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # NCHW -> NLC
+        return x
+
+
+def modulate(x, shift, scale):
+    if shift is None:
+        shift = torch.zeros_like(scale)
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+#################################################################################
+#                   Sine/Cosine Positional Embedding Functions                  #
+#################################################################################
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size_1, grid_size_2, cls_token=False, extra_tokens=0, scaling_factor=None, offset=None):
+    """
+    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_1, dtype=np.float32)
+    grid_w = np.arange(grid_size_2, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+    if scaling_factor is not None:
+        grid = grid / scaling_factor
+    if offset is not None:
+        grid = grid - offset
+    grid = grid.reshape([2, 1, grid_size_1, grid_size_2])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, 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)
+    return np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+
+
+#################################################################################
+#               Embedding Layers for Timesteps and Class Labels                 #
+#################################################################################
+
+
+class TimestepEmbedder(nn.Module):
+    """Embeds scalar timesteps into vector representations."""
+
+    def __init__(self, hidden_size, frequency_embedding_size=256, dtype=None, device=None):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True, dtype=dtype, device=device),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True, dtype=dtype, device=device),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period)
+            * torch.arange(start=0, end=half, dtype=torch.float32)
+            / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        if torch.is_floating_point(t):
+            embedding = embedding.to(dtype=t.dtype)
+        return embedding
+
+    def forward(self, t, dtype, **kwargs):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class VectorEmbedder(nn.Module):
+    """Embeds a flat vector of dimension input_dim"""
+
+    def __init__(self, input_dim: int, hidden_size: int, dtype=None, device=None):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(input_dim, hidden_size, bias=True, dtype=dtype, device=device),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True, dtype=dtype, device=device),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.mlp(x)
+
+
+#################################################################################
+#                                 Core DiT Model                                #
+#################################################################################
+
+
+def split_qkv(qkv, head_dim):
+    qkv = qkv.reshape(qkv.shape[0], qkv.shape[1], 3, -1, head_dim).movedim(2, 0)
+    return qkv[0], qkv[1], qkv[2]
+
+def optimized_attention(qkv, num_heads):
+    return attention(qkv[0], qkv[1], qkv[2], num_heads)
+
+class SelfAttention(nn.Module):
+    ATTENTION_MODES = ("xformers", "torch", "torch-hb", "math", "debug")
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_scale: Optional[float] = None,
+        attn_mode: str = "xformers",
+        pre_only: bool = False,
+        qk_norm: Optional[str] = None,
+        rmsnorm: bool = False,
+        dtype=None,
+        device=None,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias, dtype=dtype, device=device)
+        if not pre_only:
+            self.proj = nn.Linear(dim, dim, dtype=dtype, device=device)
+        assert attn_mode in self.ATTENTION_MODES
+        self.attn_mode = attn_mode
+        self.pre_only = pre_only
+
+        if qk_norm == "rms":
+            self.ln_q = RMSNorm(self.head_dim, elementwise_affine=True, eps=1.0e-6, dtype=dtype, device=device)
+            self.ln_k = RMSNorm(self.head_dim, elementwise_affine=True, eps=1.0e-6, dtype=dtype, device=device)
+        elif qk_norm == "ln":
+            self.ln_q = nn.LayerNorm(self.head_dim, elementwise_affine=True, eps=1.0e-6, dtype=dtype, device=device)
+            self.ln_k = nn.LayerNorm(self.head_dim, elementwise_affine=True, eps=1.0e-6, dtype=dtype, device=device)
+        elif qk_norm is None:
+            self.ln_q = nn.Identity()
+            self.ln_k = nn.Identity()
+        else:
+            raise ValueError(qk_norm)
+
+    def pre_attention(self, x: torch.Tensor):
+        B, L, C = x.shape
+        qkv = self.qkv(x)
+        q, k, v = split_qkv(qkv, self.head_dim)
+        q = self.ln_q(q).reshape(q.shape[0], q.shape[1], -1)
+        k = self.ln_k(k).reshape(q.shape[0], q.shape[1], -1)
+        return (q, k, v)
+
+    def post_attention(self, x: torch.Tensor) -> torch.Tensor:
+        assert not self.pre_only
+        x = self.proj(x)
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        (q, k, v) = self.pre_attention(x)
+        x = attention(q, k, v, self.num_heads)
+        x = self.post_attention(x)
+        return x
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(
+        self, dim: int, elementwise_affine: bool = False, eps: float = 1e-6, device=None, dtype=None
+    ):
+        """
+        Initialize the RMSNorm normalization layer.
+        Args:
+            dim (int): The dimension of the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+        Attributes:
+            eps (float): A small value added to the denominator for numerical stability.
+            weight (nn.Parameter): Learnable scaling parameter.
+        """
+        super().__init__()
+        self.eps = eps
+        self.learnable_scale = elementwise_affine
+        if self.learnable_scale:
+            self.weight = nn.Parameter(torch.empty(dim, device=device, dtype=dtype))
+        else:
+            self.register_parameter("weight", None)
+
+    def _norm(self, x):
+        """
+        Apply the RMSNorm normalization to the input tensor.
+        Args:
+            x (torch.Tensor): The input tensor.
+        Returns:
+            torch.Tensor: The normalized tensor.
+        """
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        """
+        Forward pass through the RMSNorm layer.
+        Args:
+            x (torch.Tensor): The input tensor.
+        Returns:
+            torch.Tensor: The output tensor after applying RMSNorm.
+        """
+        x = self._norm(x)
+        if self.learnable_scale:
+            return x * self.weight.to(device=x.device, dtype=x.dtype)
+        else:
+            return x
+
+
+class SwiGLUFeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+        multiple_of: int,
+        ffn_dim_multiplier: Optional[float] = None,
+    ):
+        """
+        Initialize the FeedForward module.
+
+        Args:
+            dim (int): Input dimension.
+            hidden_dim (int): Hidden dimension of the feedforward layer.
+            multiple_of (int): Value to ensure hidden dimension is a multiple of this value.
+            ffn_dim_multiplier (float, optional): Custom multiplier for hidden dimension. Defaults to None.
+
+        Attributes:
+            w1 (ColumnParallelLinear): Linear transformation for the first layer.
+            w2 (RowParallelLinear): Linear transformation for the second layer.
+            w3 (ColumnParallelLinear): Linear transformation for the third layer.
+
+        """
+        super().__init__()
+        hidden_dim = int(2 * hidden_dim / 3)
+        # custom dim factor multiplier
+        if ffn_dim_multiplier is not None:
+            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x):
+        return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))
+
+
+class DismantledBlock(nn.Module):
+    """A DiT block with gated adaptive layer norm (adaLN) conditioning."""
+
+    ATTENTION_MODES = ("xformers", "torch", "torch-hb", "math", "debug")
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_mode: str = "xformers",
+        qkv_bias: bool = False,
+        pre_only: bool = False,
+        rmsnorm: bool = False,
+        scale_mod_only: bool = False,
+        swiglu: bool = False,
+        qk_norm: Optional[str] = None,
+        dtype=None,
+        device=None,
+        **block_kwargs,
+    ):
+        super().__init__()
+        assert attn_mode in self.ATTENTION_MODES
+        if not rmsnorm:
+            self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+        else:
+            self.norm1 = RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, attn_mode=attn_mode, pre_only=pre_only, qk_norm=qk_norm, rmsnorm=rmsnorm, dtype=dtype, device=device)
+        if not pre_only:
+            if not rmsnorm:
+                self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+            else:
+                self.norm2 = RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        if not pre_only:
+            if not swiglu:
+                self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=nn.GELU(approximate="tanh"), dtype=dtype, device=device)
+            else:
+                self.mlp = SwiGLUFeedForward(dim=hidden_size, hidden_dim=mlp_hidden_dim, multiple_of=256)
+        self.scale_mod_only = scale_mod_only
+        if not scale_mod_only:
+            n_mods = 6 if not pre_only else 2
+        else:
+            n_mods = 4 if not pre_only else 1
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, n_mods * hidden_size, bias=True, dtype=dtype, device=device))
+        self.pre_only = pre_only
+
+    def pre_attention(self, x: torch.Tensor, c: torch.Tensor):
+        assert x is not None, "pre_attention called with None input"
+        if not self.pre_only:
+            if not self.scale_mod_only:
+                shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
+            else:
+                shift_msa = None
+                shift_mlp = None
+                scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(4, dim=1)
+            qkv = self.attn.pre_attention(modulate(self.norm1(x), shift_msa, scale_msa))
+            return qkv, (x, gate_msa, shift_mlp, scale_mlp, gate_mlp)
+        else:
+            if not self.scale_mod_only:
+                shift_msa, scale_msa = self.adaLN_modulation(c).chunk(2, dim=1)
+            else:
+                shift_msa = None
+                scale_msa = self.adaLN_modulation(c)
+            qkv = self.attn.pre_attention(modulate(self.norm1(x), shift_msa, scale_msa))
+            return qkv, None
+
+    def post_attention(self, attn, x, gate_msa, shift_mlp, scale_mlp, gate_mlp):
+        assert not self.pre_only
+        x = x + gate_msa.unsqueeze(1) * self.attn.post_attention(attn)
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+    def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        assert not self.pre_only
+        (q, k, v), intermediates = self.pre_attention(x, c)
+        attn = attention(q, k, v, self.attn.num_heads)
+        return self.post_attention(attn, *intermediates)
+
+
+def block_mixing(context, x, context_block, x_block, c):
+    assert context is not None, "block_mixing called with None context"
+    context_qkv, context_intermediates = context_block.pre_attention(context, c)
+
+    x_qkv, x_intermediates = x_block.pre_attention(x, c)
+
+    o = []
+    for t in range(3):
+        o.append(torch.cat((context_qkv[t], x_qkv[t]), dim=1))
+    q, k, v = tuple(o)
+
+    attn = attention(q, k, v, x_block.attn.num_heads)
+    context_attn, x_attn = (attn[:, : context_qkv[0].shape[1]], attn[:, context_qkv[0].shape[1] :])
+
+    if not context_block.pre_only:
+        context = context_block.post_attention(context_attn, *context_intermediates)
+    else:
+        context = None
+    x = x_block.post_attention(x_attn, *x_intermediates)
+    return context, x
+
+
+class JointBlock(nn.Module):
+    """just a small wrapper to serve as a fsdp unit"""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        pre_only = kwargs.pop("pre_only")
+        qk_norm = kwargs.pop("qk_norm", None)
+        self.context_block = DismantledBlock(*args, pre_only=pre_only, qk_norm=qk_norm, **kwargs)
+        self.x_block = DismantledBlock(*args, pre_only=False, qk_norm=qk_norm, **kwargs)
+
+    def forward(self, *args, **kwargs):
+        return block_mixing(*args, context_block=self.context_block, x_block=self.x_block, **kwargs)
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT.
+    """
+
+    def __init__(self, hidden_size: int, patch_size, out_channels: int, total_out_channels: Optional[int] = None, dtype=None, device=None):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+        self.linear = (
+            nn.Linear(hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True, dtype=dtype, device=device)
+            if (total_out_channels is None)
+            else nn.Linear(hidden_size, total_out_channels, bias=True, dtype=dtype, device=device)
+        )
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True, dtype=dtype, device=device))
+
+    def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class MMDiT(nn.Module):
+    """Diffusion model with a Transformer backbone."""
+
+    def __init__(
+        self,
+        input_size=(204, 16),
+        patch_size=(2, 2),
+        in_channels: int = 8,
+        depth: int = 12,
+        mlp_ratio: float = 4.0,
+        learn_sigma: bool = False,
+        adm_in_channels: Optional[int] = None,
+        context_embedder_config: Optional[Dict] = None,
+        register_length: int = 0,
+        attn_mode: str = "torch",
+        rmsnorm: bool = False,
+        scale_mod_only: bool = False,
+        swiglu: bool = False,
+        out_channels: Optional[int] = None,
+        pos_embed_scaling_factor: Optional[float] = None,
+        pos_embed_offset: Optional[float] = None,
+        pos_embed_max_size: Optional[int] = None,
+        num_patches = None,
+        qk_norm: Optional[str] = None,
+        qkv_bias: bool = True,
+        dtype = None,
+        device = None,
+        encoder_depth = 4,
+        z_dims=[768],
+        projector_dim=2048,
+    ):
+        super().__init__()
+        print(f"mmdit initializing with: {input_size=}, {patch_size=}, {in_channels=}, {depth=}, {mlp_ratio=}, {learn_sigma=}, {adm_in_channels=}, {context_embedder_config=}, {register_length=}, {attn_mode=}, {rmsnorm=}, {scale_mod_only=}, {swiglu=}, {out_channels=}, {pos_embed_scaling_factor=}, {pos_embed_offset=}, {pos_embed_max_size=}, {num_patches=}, {qk_norm=}, {qkv_bias=}, {dtype=}, {device=}")
+        self.dtype = dtype
+        self.learn_sigma = learn_sigma
+        self.in_channels = in_channels
+        default_out_channels = in_channels * 2 if learn_sigma else in_channels
+        self.out_channels = out_channels if out_channels is not None else default_out_channels
+        self.patch_size = patch_size
+        self.pos_embed_scaling_factor = pos_embed_scaling_factor
+        self.pos_embed_offset = pos_embed_offset
+        self.pos_embed_max_size = pos_embed_max_size = (102, 8)
+
+
+        # apply magic --> this defines a head_size of 64
+        hidden_size = 64 * depth
+        # hidden_size = 32 * depth
+        num_heads = depth
+
+        self.num_heads = num_heads
+
+        self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True, strict_img_size=self.pos_embed_max_size is None, dtype=dtype, device=device)
+        self.t_embedder = TimestepEmbedder(hidden_size, dtype=dtype, device=device)
+
+        if adm_in_channels is not None:
+            assert isinstance(adm_in_channels, int)
+            self.y_embedder = VectorEmbedder(adm_in_channels, hidden_size, dtype=dtype, device=device)
+        else:
+            self.y_embedder = None
+
+        self.context_embedder = nn.Identity()
+        # TODO: hand coded
+        context_embedder_config = {"params": {"in_features": 512, "out_features": hidden_size}, "target": "torch.nn.Linear"}
+        if context_embedder_config is not None:
+            if context_embedder_config["target"] == "torch.nn.Linear":
+                self.context_embedder = nn.Linear(**context_embedder_config["params"], dtype=dtype, device=device)
+
+        self.register_length = register_length
+        if self.register_length > 0:
+            self.register = nn.Parameter(torch.randn(1, register_length, hidden_size, dtype=dtype, device=device))
+
+        num_patches = self.x_embedder.num_patches
+        # Will use fixed sin-cos embedding:
+        # just use a buffer already
+        if num_patches is not None:
+            self.register_buffer(
+                "pos_embed",
+                torch.zeros(1, num_patches, hidden_size, dtype=dtype, device=device),
+            )
+        else:
+            self.pos_embed = None
+
+        self.joint_blocks = nn.ModuleList(
+            [
+                JointBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, attn_mode=attn_mode, pre_only=i == depth - 1, rmsnorm=rmsnorm, scale_mod_only=scale_mod_only, swiglu=swiglu, qk_norm=qk_norm, dtype=dtype, device=device)
+                for i in range(depth)
+            ]
+        )
+
+        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels, dtype=dtype, device=device)
+
+        # REPA
+        self.encoder_depth = encoder_depth
+        self.projectors = nn.ModuleList([
+            build_mlp(hidden_size, projector_dim, z_dim) for z_dim in z_dims
+            ])
+
+        # Initialize (and freeze) pos_embed by sin-cos embedding:
+        grid_size_1 = 102
+        grid_size_2 = 8
+        pos_embed = get_2d_sincos_pos_embed(
+            self.pos_embed.shape[-1], grid_size_1, grid_size_2
+            )
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        nn.init.normal_(self.context_embedder.weight, std=0.02)
+
+        if self.y_embedder is not None:
+            nn.init.normal_(self.y_embedder.mlp[0].weight, std=0.02)
+            nn.init.normal_(self.y_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers in DiT blocks:
+        for block in self.joint_blocks:
+            nn.init.constant_(block.context_block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.context_block.adaLN_modulation[-1].bias, 0)
+            nn.init.constant_(block.x_block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.x_block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def cropped_pos_embed(self, hw):
+        assert self.pos_embed_max_size is not None
+        p1, p2 = self.x_embedder.patch_size
+        h, w = hw
+        # patched size
+        h = h // p1
+        w = w // p2
+
+        assert h <= self.pos_embed_max_size[0], (h, self.pos_embed_max_size[0])
+        assert w <= self.pos_embed_max_size[1], (w, self.pos_embed_max_size[1])
+        top = (self.pos_embed_max_size[0] - h) // 2
+        left = (self.pos_embed_max_size[1] - w) // 2
+        spatial_pos_embed = rearrange(
+            self.pos_embed,
+            "1 (h w) c -> 1 h w c",
+            h=self.pos_embed_max_size[0],
+            w=self.pos_embed_max_size[1],
+        )
+        spatial_pos_embed = spatial_pos_embed[:, top : top + h, left : left + w, :]
+        spatial_pos_embed = rearrange(spatial_pos_embed, "1 h w c -> 1 (h w) c")
+        return spatial_pos_embed
+
+    def unpatchify(self, x, hw=None):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        c = self.out_channels
+        p1, p2 = self.x_embedder.patch_size
+        h, w = hw
+        # patched size
+        h = h // p1
+        w = w // p2
+        assert h * w == x.shape[1]
+
+        h_1, w_1 = self.x_embedder.img_size
+
+        x = x.reshape(shape=(x.shape[0], h, w, p1, p2, c))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], c, h_1, w_1))
+
+        return imgs
+
+    def forward_core_with_concat(
+            self, x: torch.Tensor, c_mod: torch.Tensor, context: Optional[torch.Tensor] = None,
+            detach: Optional[bool] = False) -> torch.Tensor:
+        if self.register_length > 0:
+            context = torch.cat((repeat(self.register, "1 ... -> b ...", b=x.shape[0]), context if context is not None else torch.Tensor([]).type_as(x)), 1)
+
+        # context is B, L', D
+        # x is B, L, D
+        B, L, D = x.shape
+        for i, block in enumerate(self.joint_blocks):
+            context, x = block(context, x, c=c_mod)
+
+            if (i + 1) == self.encoder_depth:
+                zs = [projector(x) for projector in self.projectors]
+
+        x = self.final_layer(x, c_mod)  # (N, T, patch_size ** 2 * out_channels)
+        return x, zs
+
+    def forward(
+            self, x: torch.Tensor, t: torch.Tensor, y: Optional[torch.Tensor] = None, context: Optional[torch.Tensor] = None, do_guidance=False,
+            detach: Optional[bool] = False) -> torch.Tensor:
+        """
+        Forward pass of DiT.
+        x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
+        t: (N,) tensor of diffusion timesteps
+        y: (N,) tensor of class labels
+        """
+        hw = x.shape[-2:]
+        x = self.x_embedder(x) + self.cropped_pos_embed(hw)
+        c = self.t_embedder(t, dtype=x.dtype)  # (N, D)
+
+        context = self.context_embedder(context)
+
+        if self.training and not do_guidance:
+            cond_mask = prob_mask_like((context.shape[0],), prob = 1 - 0.1, device = context.device) # classifier free guidance
+            cond_mask = cond_mask.to(context.dtype)
+            context = cond_mask.view(-1, 1, 1) * context
+        elif do_guidance:
+            N = x.shape[0]
+            half_bs = N // 2
+            cond_mask = torch.cat((torch.ones(half_bs), torch.zeros(N - half_bs))).to(context.device)
+            cond_mask = cond_mask.to(context.dtype)
+            context = cond_mask.view(-1, 1, 1) * context
+        else:
+            cond_mask = torch.ones(context.shape[0], device = context.device, dtype = torch.bool)
+
+        if y is not None:
+            y = self.y_embedder(y)
+            y = cond_mask.view(-1, 1) * y
+            c = c + y  # (N, D)
+
+        x, zs = self.forward_core_with_concat(x, c, context, detach)
+
+        x = self.unpatchify(x, hw=hw)  # (N, out_channels, H, W)
+        return x, zs

onset_util.py

@@ -0,0 +1,446 @@
+import torch
+import torch.nn as nn
+from torchvision.io import read_video
+import os
+from einops import rearrange
+import torchvision.transforms as transforms
+from cavp_util import reencode_video_with_diff_fps
+
+
+def extract_onset(video_path, onset_model, tmp_path, device="cuda"):
+    """Extract onset features from video using a pre-trained onset detection model."""
+    # Preprocess the video frames
+    transform = transforms.Compose([
+        transforms.Resize((128, 128), antialias=True),
+        transforms.CenterCrop((112, 112)),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+    start_second = 0
+    truncate_second = 10
+    # Load the video, change fps:
+    video_path_low_fps = reencode_video_with_diff_fps(video_path, tmp_path, 15, start_second, truncate_second)
+    frames, _, _ = read_video(video_path_low_fps, pts_unit="sec", output_format="TCHW")
+    if frames.shape[0] >= 150:
+        frames = frames[:150]
+    elif frames.shape[0] >= 120:
+        frames = frames[:120]
+
+    # Transform frames
+    frames = frames / 255.0
+    frames = transform(frames)
+
+    frames = rearrange(frames, '(b t) c h w -> b c t h w', t=30).to(device)
+
+    # Forward pass through the model to get onset features
+    with torch.no_grad():
+        onset_features = onset_model(frames).reshape(-1)
+
+    # Remove the file
+    os.remove(video_path_low_fps)
+    return onset_features.detach().cpu().numpy()
+
+#################################################################################
+#                                     ResNet                                    #
+#################################################################################
+
+__all__ = ['r3d_18', 'mc3_18', 'r2plus1d_18']
+
+model_urls = {
+    'r3d_18': 'https://download.pytorch.org/models/r3d_18-b3b3357e.pth',
+    'mc3_18': 'https://download.pytorch.org/models/mc3_18-a90a0ba3.pth',
+    'r2plus1d_18': 'https://download.pytorch.org/models/r2plus1d_18-91a641e6.pth',
+}
+
+
+class Conv3DSimple(nn.Conv3d):
+    def __init__(self,
+                 in_planes,
+                 out_planes,
+                 midplanes=None,
+                 stride=1,
+                 padding=1):
+
+        super(Conv3DSimple, self).__init__(
+            in_channels=in_planes,
+            out_channels=out_planes,
+            kernel_size=(3, 3, 3),
+            stride=stride,
+            padding=padding,
+            bias=False)
+
+    @staticmethod
+    def get_downsample_stride(stride):
+        return stride, stride, stride
+
+
+class Conv2Plus1D(nn.Sequential):
+
+    def __init__(self,
+                 in_planes,
+                 out_planes,
+                 midplanes,
+                 stride=1,
+                 padding=1):
+        super(Conv2Plus1D, self).__init__(
+            nn.Conv3d(in_planes, midplanes, kernel_size=(1, 3, 3),
+                      stride=(1, stride, stride), padding=(0, padding, padding),
+                      bias=False),
+            nn.BatchNorm3d(midplanes),
+            nn.ReLU(inplace=True),
+            nn.Conv3d(midplanes, out_planes, kernel_size=(3, 1, 1),
+                      stride=(stride, 1, 1), padding=(padding, 0, 0),
+                      bias=False))
+
+    @staticmethod
+    def get_downsample_stride(stride):
+        return stride, stride, stride
+
+
+class Conv3DNoTemporal(nn.Conv3d):
+
+    def __init__(self,
+                 in_planes,
+                 out_planes,
+                 midplanes=None,
+                 stride=1,
+                 padding=1):
+
+        super(Conv3DNoTemporal, self).__init__(
+            in_channels=in_planes,
+            out_channels=out_planes,
+            kernel_size=(1, 3, 3),
+            stride=(1, stride, stride),
+            padding=(0, padding, padding),
+            bias=False)
+
+    @staticmethod
+    def get_downsample_stride(stride):
+        return 1, stride, stride
+
+
+class BasicBlock(nn.Module):
+
+    expansion = 1
+
+    def __init__(self, inplanes, planes, conv_builder, stride=1, downsample=None):
+        midplanes = (inplanes * planes * 3 * 3 *
+                     3) // (inplanes * 3 * 3 + 3 * planes)
+
+        super(BasicBlock, self).__init__()
+        self.conv1 = nn.Sequential(
+            conv_builder(inplanes, planes, midplanes, stride),
+            nn.BatchNorm3d(planes),
+            nn.ReLU(inplace=True)
+        )
+        self.conv2 = nn.Sequential(
+            conv_builder(planes, planes, midplanes),
+            nn.BatchNorm3d(planes)
+        )
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.conv2(out)
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, conv_builder, stride=1, downsample=None):
+
+        super(Bottleneck, self).__init__()
+        midplanes = (inplanes * planes * 3 * 3 *
+                     3) // (inplanes * 3 * 3 + 3 * planes)
+
+        # 1x1x1
+        self.conv1 = nn.Sequential(
+            nn.Conv3d(inplanes, planes, kernel_size=1, bias=False),
+            nn.BatchNorm3d(planes),
+            nn.ReLU(inplace=True)
+        )
+        # Second kernel
+        self.conv2 = nn.Sequential(
+            conv_builder(planes, planes, midplanes, stride),
+            nn.BatchNorm3d(planes),
+            nn.ReLU(inplace=True)
+        )
+
+        # 1x1x1
+        self.conv3 = nn.Sequential(
+            nn.Conv3d(planes, planes * self.expansion,
+                      kernel_size=1, bias=False),
+            nn.BatchNorm3d(planes * self.expansion)
+        )
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.conv2(out)
+        out = self.conv3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class BasicStem(nn.Sequential):
+    """The default conv-batchnorm-relu stem
+    """
+
+    def __init__(self):
+        super(BasicStem, self).__init__(
+            nn.Conv3d(3, 64, kernel_size=(3, 7, 7), stride=(1, 2, 2),
+                      padding=(1, 3, 3), bias=False),
+            nn.BatchNorm3d(64),
+            nn.ReLU(inplace=True))
+
+
+class R2Plus1dStem(nn.Sequential):
+    """R(2+1)D stem is different than the default one as it uses separated 3D convolution
+    """
+
+    def __init__(self):
+        super(R2Plus1dStem, self).__init__(
+            nn.Conv3d(3, 45, kernel_size=(1, 7, 7),
+                      stride=(1, 2, 2), padding=(0, 3, 3),
+                      bias=False),
+            nn.BatchNorm3d(45),
+            nn.ReLU(inplace=True),
+            nn.Conv3d(45, 64, kernel_size=(3, 1, 1),
+                      stride=(1, 1, 1), padding=(1, 0, 0),
+                      bias=False),
+            nn.BatchNorm3d(64),
+            nn.ReLU(inplace=True))
+
+
+class VideoResNet(nn.Module):
+
+    def __init__(self, block, conv_makers, layers,
+                 stem, num_classes=400,
+                 zero_init_residual=False):
+        """Generic resnet video generator.
+        Args:
+            block (nn.Module): resnet building block
+            conv_makers (list(functions)): generator function for each layer
+            layers (List[int]): number of blocks per layer
+            stem (nn.Module, optional): Resnet stem, if None, defaults to conv-bn-relu. Defaults to None.
+            num_classes (int, optional): Dimension of the final FC layer. Defaults to 400.
+            zero_init_residual (bool, optional): Zero init bottleneck residual BN. Defaults to False.
+        """
+        super(VideoResNet, self).__init__()
+        self.inplanes = 64
+
+        self.stem = stem()
+
+        self.layer1 = self._make_layer(
+            block, conv_makers[0], 64, layers[0], stride=1)
+        self.layer2 = self._make_layer(
+            block, conv_makers[1], 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(
+            block, conv_makers[2], 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(
+            block, conv_makers[3], 512, layers[3], stride=2)
+
+        self.avgpool = nn.AdaptiveAvgPool3d((1, 1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        # init weights
+        self._initialize_weights()
+
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+
+    def forward(self, x):
+        x = self.stem(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        # Flatten the layer to fc
+        # x = x.flatten(1)
+        # x = self.fc(x)
+        N = x.shape[0]
+        x = x.squeeze()
+        if N == 1:
+            x = x[None]
+
+        return x
+
+    def _make_layer(self, block, conv_builder, planes, blocks, stride=1):
+        downsample = None
+
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            ds_stride = conv_builder.get_downsample_stride(stride)
+            downsample = nn.Sequential(
+                nn.Conv3d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=ds_stride, bias=False),
+                nn.BatchNorm3d(planes * block.expansion)
+            )
+        layers = []
+        layers.append(block(self.inplanes, planes,
+                      conv_builder, stride, downsample))
+
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, conv_builder))
+
+        return nn.Sequential(*layers)
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv3d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out',
+                                        nonlinearity='relu')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm3d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.constant_(m.bias, 0)
+
+
+def _video_resnet(arch, pretrained=False, progress=True, **kwargs):
+    model = VideoResNet(**kwargs)
+
+    if pretrained:
+        state_dict = load_state_dict_from_url(model_urls[arch],
+                                              progress=progress)
+        model.load_state_dict(state_dict)
+    return model
+
+
+def r3d_18(pretrained=False, progress=True, **kwargs):
+    """Construct 18 layer Resnet3D model as in
+    https://arxiv.org/abs/1711.11248
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on Kinetics-400
+        progress (bool): If True, displays a progress bar of the download to stderr
+    Returns:
+        nn.Module: R3D-18 network
+    """
+
+    return _video_resnet('r3d_18',
+                         pretrained, progress,
+                         block=BasicBlock,
+                         conv_makers=[Conv3DSimple] * 4,
+                         layers=[2, 2, 2, 2],
+                         stem=BasicStem, **kwargs)
+
+
+def mc3_18(pretrained=False, progress=True, **kwargs):
+    """Constructor for 18 layer Mixed Convolution network as in
+    https://arxiv.org/abs/1711.11248
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on Kinetics-400
+        progress (bool): If True, displays a progress bar of the download to stderr
+    Returns:
+        nn.Module: MC3 Network definition
+    """
+    return _video_resnet('mc3_18',
+                         pretrained, progress,
+                         block=BasicBlock,
+                         conv_makers=[Conv3DSimple] + [Conv3DNoTemporal] * 3,
+                         layers=[2, 2, 2, 2],
+                         stem=BasicStem, **kwargs)
+
+
+def r2plus1d_18(pretrained=False, progress=True, **kwargs):
+    """Constructor for the 18 layer deep R(2+1)D network as in
+    https://arxiv.org/abs/1711.11248
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on Kinetics-400
+        progress (bool): If True, displays a progress bar of the download to stderr
+    Returns:
+        nn.Module: R(2+1)D-18 network
+    """
+    return _video_resnet('r2plus1d_18',
+                         pretrained, progress,
+                         block=BasicBlock,
+                         conv_makers=[Conv2Plus1D] * 4,
+                         layers=[2, 2, 2, 2],
+                         stem=R2Plus1dStem, **kwargs)
+
+
+#################################################################################
+#                                   Onset Net                                   #
+#################################################################################
+
+class R2plus1d18KeepTemp(nn.Module):
+
+    def __init__(self, pretrained=True):
+        super().__init__()
+
+        self.model = r2plus1d_18(pretrained=pretrained)
+
+        self.model.layer2[0].conv1[0][3] = nn.Conv3d(230, 128, kernel_size=(3, 1, 1),
+                                                     stride=(1, 1, 1), padding=(1, 0, 0), bias=False)
+        self.model.layer2[0].downsample = nn.Sequential(
+            nn.Conv3d(64, 128, kernel_size=(1, 1, 1), stride=(1, 2, 2), bias=False),
+            nn.BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+        )
+        self.model.layer3[0].conv1[0][3] = nn.Conv3d(460, 256, kernel_size=(3, 1, 1),
+                                                     stride=(1, 1, 1), padding=(1, 0, 0), bias=False)
+        self.model.layer3[0].downsample = nn.Sequential(
+            nn.Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(1, 2, 2), bias=False),
+            nn.BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+        )
+        self.model.layer4[0].conv1[0][3] = nn.Conv3d(921, 512, kernel_size=(3, 1, 1),
+                                                     stride=(1, 1, 1), padding=(1, 0, 0), bias=False)
+        self.model.layer4[0].downsample = nn.Sequential(
+            nn.Conv3d(256, 512, kernel_size=(1, 1, 1), stride=(1, 2, 2), bias=False),
+            nn.BatchNorm3d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+        )
+        self.model.avgpool = nn.AdaptiveAvgPool3d((None, 1, 1))
+        self.model.fc = nn.Identity()
+
+    def forward(self, x):
+        x = self.model(x)
+        return x
+
+
+class VideoOnsetNet(nn.Module):
+    # Video Onset detection network
+    def __init__(self, pretrained=False):
+        super(VideoOnsetNet, self).__init__()
+        self.net = R2plus1d18KeepTemp(pretrained=pretrained)
+        self.fc = nn.Sequential(
+            nn.Linear(512, 128),
+            nn.ReLU(True),
+            nn.Linear(128, 1)
+        )
+
+    def forward(self, x):
+        x = self.net(x)
+        x = x.transpose(-1, -2)
+        x = self.fc(x)
+        x = x.squeeze(-1)
+
+        return x
+

preprocess/extract_cavp.py

@@ -0,0 +1,42 @@
+import os
+import torch
+import numpy as np
+from tqdm import tqdm
+from argparse import ArgumentParser
+
+import sys
+sys.path.append(os.getcwd())
+from cavp_util import Extract_CAVP_Features
+
+def main():
+    parser = ArgumentParser(description="Inference script parameters")
+    parser.add_argument("--video_folder_path", type=str, default="./input_videos", required=True, help="Path to the input video folder")
+    parser.add_argument("--save_folder_path", type=str, default="./output", help="Folder to save output files")
+    parser.add_argument("--cavp_config_path", type=str, default="./cavp.yaml", help="Path to CAVP config file")
+    parser.add_argument("--cavp_ckpt_path", type=str, default="./cavp_epoch66.ckpt", help="Path to CAVP checkpoint file")
+
+    args = parser.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    extract_cavp = Extract_CAVP_Features(device=device, config_path=args.cavp_config_path, ckpt_path=args.cavp_ckpt_path)
+
+    os.makedirs(os.path.join(args.save_folder_path, "cavp_feats"), exist_ok=True)
+
+    data_list = [file for file in os.listdir(args.video_folder_path) if file.endswith(".mp4")]
+    data_list = sorted(data_list)
+
+    for _, video_file in enumerate(tqdm(data_list, desc="Extracting CAVP features", total=len(data_list))):
+        video_path = os.path.join(args.video_folder_path, video_file)
+        try:
+            cavp_feats = extract_cavp(video_path, tmp_path=args.save_folder_path)
+            # Save cavp_feats as npz file
+            base_name = os.path.splitext(os.path.basename(video_file))[0]
+            np.savez(os.path.join(args.save_folder_path, "cavp_feats", f"{base_name}.npz"), cavp_feats)
+        except Exception as e:
+            print(f"Error processing {video_file}: {e}")
+
+    print("========================================FINISH CAVP EXTRACTION===========================================")
+
+
+if __name__ == "__main__":
+    main()

preprocess/extract_fbank.py

@@ -0,0 +1,58 @@
+import os
+import torch
+import torchaudio
+import numpy as np
+from tqdm import tqdm
+import soundfile as sf
+from argparse import ArgumentParser
+
+def main():
+    parser = ArgumentParser(description="Inference script parameters")
+    parser.add_argument("--wav_folder_path", type=str, default="./input_wavs", required=True, help="Path to the input video folder")
+    parser.add_argument("--save_folder_path", type=str, default="./output", help="Folder to save output files")
+
+    args = parser.parse_args()
+
+    os.makedirs(os.path.join(args.save_folder_path, "fbank"), exist_ok=True)
+
+    target_length = 1024
+    norm_mean = -4.268
+    norm_std = 4.569
+
+    # Loop over all .wav files in the audio folder
+    for filename in tqdm(os.listdir(args.wav_folder_path)):
+        if filename.endswith('.wav'):
+            # Load the audio file
+            source_file = os.path.join(args.wav_folder_path, filename)
+            wav, sr = sf.read(source_file)
+            if len(wav.shape) > 1:
+                wav = wav[:, 0]
+
+            source = torch.from_numpy(wav).float()
+            if not sr == 16e3:
+                source = torchaudio.functional.resample(source, orig_freq=sr, new_freq=16000).float()
+
+            source = source - source.mean()
+            source = source.unsqueeze(dim=0)
+            source = torchaudio.compliance.kaldi.fbank(source, htk_compat=True, sample_frequency=16000, use_energy=False,
+                window_type='hanning', num_mel_bins=128, dither=0.0, frame_shift=10).unsqueeze(dim=0)
+
+            n_frames = source.shape[1]
+            diff = target_length - n_frames
+            if diff > 0:
+                m = torch.nn.ZeroPad2d((0, 0, 0, diff)) 
+                source = m(source)
+            elif diff < 0:
+                source = source[:,0:target_length, :]
+            source = (source - norm_mean) / (norm_std * 2)
+
+            # Save the spectrogram as .npy file
+            output_filename = os.path.splitext(filename)[0] + '.npy'
+            output_path = os.path.join(args.save_folder_path, "fbank", output_filename)
+
+            np.save(output_path, source.squeeze(0).numpy())
+
+    print("========================================FINISH FBANK EXTRACTION===========================================")
+
+if __name__ == "__main__":
+    main()

preprocess/extract_mel.py

@@ -0,0 +1,333 @@
+import os
+import torchaudio
+import numpy as np
+from tqdm import tqdm
+import torch
+import torch.nn.functional as F
+from scipy.signal import get_window
+from librosa.util import pad_center, tiny, normalize
+from librosa.filters import mel as librosa_mel_fn
+from argparse import ArgumentParser
+
+def window_sumsquare(
+    window,
+    n_frames,
+    hop_length,
+    win_length,
+    n_fft,
+    dtype=np.float32,
+    norm=None,
+):
+    if win_length is None:
+        win_length = n_fft
+
+    n = n_fft + hop_length * (n_frames - 1)
+    x = np.zeros(n, dtype=dtype)
+
+    # Compute the squared window at the desired length
+    win_sq = get_window(window, win_length, fftbins=True)
+    win_sq = normalize(win_sq, norm=norm) ** 2
+    win_sq = pad_center(win_sq, size=n_fft)
+
+    # Fill the envelope
+    for i in range(n_frames):
+        sample = i * hop_length
+        x[sample : min(n, sample + n_fft)] += win_sq[: max(0, min(n_fft, n - sample))]
+    return x
+
+def dynamic_range_compression(x, normalize_fun=torch.log, C=1, clip_val=1e-5):
+    """
+    PARAMS
+    ------
+    C: compression factor
+    """
+    return normalize_fun(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression(x, C=1):
+    """
+    PARAMS
+    ------
+    C: compression factor used to compress
+    """
+    return torch.exp(x) / C
+
+
+class STFT(torch.nn.Module):
+    """adapted from Prem Seetharaman's https://github.com/pseeth/pytorch-stft"""
+
+    def __init__(self, filter_length, hop_length, win_length, window="hann"):
+        super(STFT, self).__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.window = window
+        self.forward_transform = None
+        scale = self.filter_length / self.hop_length
+        fourier_basis = np.fft.fft(np.eye(self.filter_length))
+
+        cutoff = int((self.filter_length / 2 + 1))
+        fourier_basis = np.vstack(
+            [np.real(fourier_basis[:cutoff, :]), np.imag(fourier_basis[:cutoff, :])]
+        )
+
+        forward_basis = torch.FloatTensor(fourier_basis[:, None, :])
+        inverse_basis = torch.FloatTensor(
+            np.linalg.pinv(scale * fourier_basis).T[:, None, :]
+        )
+
+        if window is not None:
+            assert filter_length >= win_length
+            # get window and zero center pad it to filter_length
+            fft_window = get_window(window, win_length, fftbins=True)
+            fft_window = pad_center(fft_window, size=filter_length)
+            fft_window = torch.from_numpy(fft_window).float()
+
+            # window the bases
+            forward_basis *= fft_window
+            inverse_basis *= fft_window
+
+        self.register_buffer("forward_basis", forward_basis.float())
+        self.register_buffer("inverse_basis", inverse_basis.float())
+
+    def transform(self, input_data):
+        num_batches = input_data.size(0)
+        num_samples = input_data.size(1)
+
+        self.num_samples = num_samples
+
+        # similar to librosa, reflect-pad the input
+        input_data = input_data.view(num_batches, 1, num_samples)
+        input_data = F.pad(
+            input_data.unsqueeze(1),
+            (int(self.filter_length / 2), int(self.filter_length / 2), 0, 0),
+            mode="reflect",
+        )
+        input_data = input_data.squeeze(1)
+
+        forward_transform = F.conv1d(
+            input_data,
+            torch.autograd.Variable(self.forward_basis, requires_grad=False),
+            stride=self.hop_length,
+            padding=0,
+        ).cpu()
+
+        cutoff = int((self.filter_length / 2) + 1)
+        real_part = forward_transform[:, :cutoff, :]
+        imag_part = forward_transform[:, cutoff:, :]
+
+        magnitude = torch.sqrt(real_part**2 + imag_part**2)
+        phase = torch.autograd.Variable(torch.atan2(imag_part.data, real_part.data))
+
+        return magnitude, phase
+
+    def inverse(self, magnitude, phase):
+        recombine_magnitude_phase = torch.cat(
+            [magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1
+        )
+
+        inverse_transform = F.conv_transpose1d(
+            recombine_magnitude_phase,
+            torch.autograd.Variable(self.inverse_basis, requires_grad=False),
+            stride=self.hop_length,
+            padding=0,
+        )
+
+        if self.window is not None:
+            window_sum = window_sumsquare(
+                self.window,
+                magnitude.size(-1),
+                hop_length=self.hop_length,
+                win_length=self.win_length,
+                n_fft=self.filter_length,
+                dtype=np.float32,
+            )
+            # remove modulation effects
+            approx_nonzero_indices = torch.from_numpy(
+                np.where(window_sum > tiny(window_sum))[0]
+            )
+            window_sum = torch.autograd.Variable(
+                torch.from_numpy(window_sum), requires_grad=False
+            )
+            window_sum = window_sum
+            inverse_transform[:, :, approx_nonzero_indices] /= window_sum[
+                approx_nonzero_indices
+            ]
+
+            # scale by hop ratio
+            inverse_transform *= float(self.filter_length) / self.hop_length
+
+        inverse_transform = inverse_transform[:, :, int(self.filter_length / 2) :]
+        inverse_transform = inverse_transform[:, :, : -int(self.filter_length / 2) :]
+
+        return inverse_transform
+
+    def forward(self, input_data):
+        self.magnitude, self.phase = self.transform(input_data)
+        reconstruction = self.inverse(self.magnitude, self.phase)
+        return reconstruction
+
+
+class TacotronSTFT(torch.nn.Module):
+    def __init__(
+        self,
+        filter_length,
+        hop_length,
+        win_length,
+        n_mel_channels,
+        sampling_rate,
+        mel_fmin,
+        mel_fmax,
+    ):
+        super(TacotronSTFT, self).__init__()
+        self.n_mel_channels = n_mel_channels
+        self.sampling_rate = sampling_rate
+        self.stft_fn = STFT(filter_length, hop_length, win_length)
+        mel_basis = librosa_mel_fn(
+            sr=sampling_rate, n_fft=filter_length, n_mels=n_mel_channels, fmin=mel_fmin, fmax=mel_fmax
+        )
+        mel_basis = torch.from_numpy(mel_basis).float()
+        self.register_buffer("mel_basis", mel_basis)
+
+    def spectral_normalize(self, magnitudes, normalize_fun):
+        output = dynamic_range_compression(magnitudes, normalize_fun)
+        return output
+
+    def spectral_de_normalize(self, magnitudes):
+        output = dynamic_range_decompression(magnitudes)
+        return output
+
+    def mel_spectrogram(self, y, normalize_fun=torch.log):
+        assert torch.min(y.data) >= -1, torch.min(y.data)
+        assert torch.max(y.data) <= 1, torch.max(y.data)
+
+        magnitudes, phases = self.stft_fn.transform(y)
+        magnitudes = magnitudes.data
+        mel_output = torch.matmul(self.mel_basis, magnitudes)
+        mel_output = self.spectral_normalize(mel_output, normalize_fun)
+        energy = torch.norm(magnitudes, dim=1)
+
+        log_magnitudes = self.spectral_normalize(magnitudes, normalize_fun)
+
+        return mel_output, log_magnitudes, energy
+
+def get_mel_from_wav(audio, _stft):
+    audio = torch.clip(torch.FloatTensor(audio).unsqueeze(0), -1, 1)
+    audio = torch.autograd.Variable(audio, requires_grad=False)
+    melspec, log_magnitudes_stft, energy = _stft.mel_spectrogram(audio)
+    melspec = torch.squeeze(melspec, 0).numpy().astype(np.float32)
+    log_magnitudes_stft = (
+        torch.squeeze(log_magnitudes_stft, 0).numpy().astype(np.float32)
+    )
+    energy = torch.squeeze(energy, 0).numpy().astype(np.float32)
+    return melspec, log_magnitudes_stft, energy
+
+def _pad_spec(fbank, target_length=1024):
+    n_frames = fbank.shape[0]
+    p = target_length - n_frames
+    # cut and pad
+    if p > 0:
+        m = torch.nn.ZeroPad2d((0, 0, 0, p))
+        fbank = m(fbank)
+    elif p < 0:
+        fbank = fbank[0:target_length, :]
+
+    if fbank.size(-1) % 2 != 0:
+        fbank = fbank[..., :-1]
+
+    return fbank
+
+def pad_wav(waveform, segment_length):
+    waveform_length = waveform.shape[-1]
+    assert waveform_length > 100, "Waveform is too short, %s" % waveform_length
+    if segment_length is None or waveform_length == segment_length:
+        return waveform
+    elif waveform_length > segment_length:
+        return waveform[:segment_length]
+    elif waveform_length < segment_length:
+        temp_wav = np.zeros((1, segment_length))
+        temp_wav[:, :waveform_length] = waveform
+    return temp_wav
+
+def normalize_wav(waveform):
+    waveform = waveform - np.mean(waveform)
+    waveform = waveform / (np.max(np.abs(waveform)) + 1e-8)
+    return waveform * 0.5
+
+def read_wav_file(filename, segment_length):
+    waveform, sr = torchaudio.load(filename)
+    waveform = torchaudio.functional.resample(waveform, orig_freq=sr, new_freq=16000)
+    waveform = waveform.numpy()[0, ...]
+    waveform = normalize_wav(waveform)
+    waveform = waveform[None, ...]
+    waveform = pad_wav(waveform, segment_length)
+    
+    waveform = waveform / np.max(np.abs(waveform))
+    waveform = 0.5 * waveform
+    
+    return waveform
+
+def wav_to_fbank(filename, target_length=1024, fn_STFT=None):
+    assert fn_STFT is not None
+
+    # mixup
+    waveform = read_wav_file(filename, target_length * 160)  # hop size is 160
+
+    waveform = waveform[0, ...]
+    waveform = torch.FloatTensor(waveform)
+
+    fbank, log_magnitudes_stft, energy = get_mel_from_wav(waveform, fn_STFT)
+
+    fbank = torch.FloatTensor(fbank.T)
+    log_magnitudes_stft = torch.FloatTensor(log_magnitudes_stft.T)
+
+    fbank, log_magnitudes_stft = _pad_spec(fbank, target_length), _pad_spec(
+        log_magnitudes_stft, target_length
+    )
+
+    return fbank, log_magnitudes_stft, waveform
+
+def main():
+    parser = ArgumentParser(description="Inference script parameters")
+    parser.add_argument("--wav_folder_path", type=str, default="./input_wavs", required=True, help="Path to the input video folder")
+    parser.add_argument("--save_folder_path", type=str, default="./output", help="Folder to save output files")
+
+    args = parser.parse_args()
+
+    os.makedirs(os.path.join(args.save_folder_path, "melspec"), exist_ok=True)
+
+    # Parameters
+    filter_length = 1024
+    hop_length = 160
+    win_length = 1024
+    n_mel_channels = 64
+    sampling_rate = 16000
+    mel_fmin = 0
+    mel_fmax = 8000
+    duration = 10
+
+    fn_STFT = TacotronSTFT(
+        filter_length,
+        hop_length,
+        win_length,
+        n_mel_channels,
+        sampling_rate,
+        mel_fmin,
+        mel_fmax,
+    )
+
+    for filename in tqdm(os.listdir(args.wav_folder_path)):
+        if filename.endswith('.wav'):
+            original_audio_file_path = os.path.join(args.wav_folder_path, filename)
+            mel, _, _ = wav_to_fbank(
+                original_audio_file_path, target_length=int(duration * 102.4), fn_STFT=fn_STFT
+            )
+            output_filename = os.path.splitext(filename)[0] + '.npy'
+            output_path = os.path.join(args.save_folder_path, "melspec", output_filename)
+            np.save(output_path, mel.numpy())
+
+    print("========================================FINISH MELSPEC EXTRACTION===========================================")
+
+
+if __name__ == "__main__":
+    main()

preprocess/extract_onset.py

@@ -0,0 +1,52 @@
+import os
+import torch
+import numpy as np
+from tqdm import tqdm
+from argparse import ArgumentParser
+
+import sys
+sys.path.append(os.getcwd())
+from onset_util import VideoOnsetNet, extract_onset
+
+def main():
+    parser = ArgumentParser(description="Inference script parameters")
+    parser.add_argument("--video_folder_path", type=str, default="./input_videos", required=True, help="Path to the input video folder")
+    parser.add_argument("--save_folder_path", type=str, default="./output", help="Folder to save output files")
+    parser.add_argument("--onset_ckpt_path", type=str, default="./onset_ckpt.ckpt", help="Path to onset checkpoint")
+
+    args = parser.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # Load the pre-trained onset detection model
+    state_dict = torch.load(args.onset_ckpt_path)["state_dict"]
+    new_state_dict = {}
+    for key, value in state_dict.items():
+        if "model.net.model" in key:
+            new_key = key.replace("model.net.model", "net.model")  # Adjust the key as needed
+        elif "model.fc." in key:
+            new_key = key.replace("model.fc", "fc")  # Adjust the key as needed
+        new_state_dict[new_key] = value
+    onset_model = VideoOnsetNet(False).to(device)
+    onset_model.load_state_dict(new_state_dict)
+    onset_model.eval()
+
+    os.makedirs(os.path.join(args.save_folder_path, "onset_feats"), exist_ok=True)
+
+    data_list = [file for file in os.listdir(args.video_folder_path) if file.endswith(".mp4")]
+    data_list = sorted(data_list)
+
+    for _, video_file in enumerate(tqdm(data_list, desc="Extracting Onset features", total=len(data_list))):
+        video_path = os.path.join(args.video_folder_path, video_file)
+        try:
+            onset_feats = extract_onset(video_path, onset_model, tmp_path=args.save_folder_path, device=device)
+            # Save cavp_feats as npz file
+            base_name = os.path.splitext(os.path.basename(video_file))[0]
+            np.savez(os.path.join(args.save_folder_path, "onset_feats", f"{base_name}.npz"), onset_feats)
+        except Exception as e:
+            print(f"Error processing {video_file}: {e}")
+
+    print("========================================FINISH CAVP EXTRACTION===========================================")
+
+
+if __name__ == "__main__":
+    main()

preprocess_audio.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+
+# Set common paths
+WAV_FOLDER_PATH="./VGGSound/audios"
+SAVE_FOLDER_PATH="./output"
+
+echo "Starting audio preprocessing..."
+echo "WAV folder: $WAV_FOLDER_PATH"
+echo "Save folder: $SAVE_FOLDER_PATH"
+
+# Extract mel spectrograms
+echo "Extracting mel spectrograms..."
+CUDA_VISIBLE_DEVICES=7 python preprocess/extract_mel.py \
+    --wav_folder_path $WAV_FOLDER_PATH \
+    --save_folder_path $SAVE_FOLDER_PATH
+
+# Extract fbank features
+echo "Extracting fbank features..."
+CUDA_VISIBLE_DEVICES=7 python preprocess/extract_fbank.py \
+    --wav_folder_path $WAV_FOLDER_PATH \
+    --save_folder_path $SAVE_FOLDER_PATH
+
+echo "Audio preprocessing completed!"

preprocess_video.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+
+# Set common paths
+VIDEO_FOLDER_PATH="./VGGSound/videos"
+SAVE_FOLDER_PATH="./output"
+
+echo "Starting video preprocessing..."
+echo "Video folder: $VIDEO_FOLDER_PATH"
+echo "Save folder: $SAVE_FOLDER_PATH"
+
+# Extract CAVP features
+echo "Extracting CAVP features..."
+CUDA_VISIBLE_DEVICES=0 python preprocess/extract_cavp.py \
+    --video_folder_path $VIDEO_FOLDER_PATH \
+    --save_folder_path $SAVE_FOLDER_PATH \
+    --cavp_config_path ./cavp/cavp.yaml \
+    --cavp_ckpt_path ./ckpts/cavp_epoch66.ckpt
+
+# Extract onset features
+echo "Extracting onset features..."
+CUDA_VISIBLE_DEVICES=0 python preprocess/extract_onset.py \
+    --video_folder_path $VIDEO_FOLDER_PATH \
+    --save_folder_path $SAVE_FOLDER_PATH \
+    --onset_ckpt_path ./ckpts/onset_model.ckpt
+
+echo "Video preprocessing completed!"

requirements.txt

@@ -0,0 +1,14 @@
+accelerate==1.2.0
+av==10.0.0
+diffusers==0.31.0
+einops==0.8.0
+ffmpeg-python==0.2.0
+h5py==3.10.0
+librosa==0.10.2.post1
+mmcv==1.7.0
+numpy==1.23.5
+opencv-python==4.5.5.64
+soundfile==0.12.1
+timm==1.0.12
+transformers==4.47.0
+wandb==0.19.0

samplers.py

@@ -0,0 +1,198 @@
+import torch
+import numpy as np
+
+
+def expand_t_like_x(t, x_cur):
+    """Function to reshape time t to broadcastable dimension of x
+    Args:
+      t: [batch_dim,], time vector
+      x: [batch_dim,...], data point
+    """
+    dims = [1] * (len(x_cur.size()) - 1)
+    t = t.view(t.size(0), *dims)
+    return t
+
+def get_score_from_velocity(vt, xt, t, path_type="linear"):
+    """Wrapper function: transfrom velocity prediction model to score
+    Args:
+        velocity: [batch_dim, ...] shaped tensor; velocity model output
+        x: [batch_dim, ...] shaped tensor; x_t data point
+        t: [batch_dim,] time tensor
+    """
+    t = expand_t_like_x(t, xt)
+    if path_type == "linear":
+        alpha_t, d_alpha_t = 1 - t, torch.ones_like(xt, device=xt.device) * -1
+        sigma_t, d_sigma_t = t, torch.ones_like(xt, device=xt.device)
+    elif path_type == "cosine":
+        alpha_t = torch.cos(t * np.pi / 2)
+        sigma_t = torch.sin(t * np.pi / 2)
+        d_alpha_t = -np.pi / 2 * torch.sin(t * np.pi / 2)
+        d_sigma_t =  np.pi / 2 * torch.cos(t * np.pi / 2)
+    else:
+        raise NotImplementedError
+
+    mean = xt
+    reverse_alpha_ratio = alpha_t / d_alpha_t
+    var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t
+    score = (reverse_alpha_ratio * vt - mean) / var
+
+    return score
+
+
+def compute_diffusion(t_cur):
+    return 2 * t_cur
+
+
+def euler_sampler(
+        model,
+        latents,
+        y,
+        context,
+        num_steps=20,
+        heun=False,
+        cfg_scale=1.0,
+        guidance_low=0.0,
+        guidance_high=1.0,
+        path_type="linear", # not used, just for compatability
+        ):
+    # setup conditioning
+    if cfg_scale > 1.0:
+        y_null = torch.zeros_like(y).to(y.device)
+        context_null = torch.zeros_like(context).to(context.device)
+    _dtype = latents.dtype    
+    t_steps = torch.linspace(1, 0, num_steps+1, dtype=torch.bfloat16)
+    x_next = latents.to(torch.bfloat16)
+    device = x_next.device
+
+    with torch.no_grad():
+        for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])):
+            x_cur = x_next
+            if cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low:
+                model_input = torch.cat([x_cur] * 2, dim=0)
+                y_cur = torch.cat([y, y_null], dim=0)
+                context_cur = torch.cat([context, context_null], dim=0)
+            else:
+                model_input = x_cur
+                y_cur = y            
+                context_cur = context
+            do_guidance = (cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low)
+            kwargs = dict(y=y_cur, context=context_cur, do_guidance=do_guidance)
+            time_input = torch.ones(model_input.size(0)).to(device=device, dtype=torch.bfloat16) * t_cur
+            d_cur = model(
+                model_input.to(dtype=_dtype), time_input.to(dtype=_dtype), **kwargs
+                )[0]
+            if cfg_scale > 1. and t_cur <= guidance_high and t_cur >= guidance_low:
+                d_cur_cond, d_cur_uncond = d_cur.chunk(2)
+                d_cur = d_cur_uncond + cfg_scale * (d_cur_cond - d_cur_uncond)                
+            x_next = x_cur + (t_next - t_cur) * d_cur
+            if heun and (i < num_steps - 1):
+                if cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low:
+                    model_input = torch.cat([x_next] * 2)
+                    y_cur = torch.cat([y, y_null], dim=0)
+                    context_cur = torch.cat([context, context_null], dim=0)
+                else:
+                    model_input = x_next
+                    y_cur = y
+                    context_cur = context
+                do_guidance = (cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low)
+                kwargs = dict(y=y_cur, context=context_cur, do_guidance=do_guidance)
+                time_input = torch.ones(model_input.size(0)).to(
+                    device=model_input.device, dtype=torch.bfloat16
+                    ) * t_next
+                d_prime = model(
+                    model_input.to(dtype=_dtype), time_input.to(dtype=_dtype), **kwargs
+                    )[0]
+                if cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low:
+                    d_prime_cond, d_prime_uncond = d_prime.chunk(2)
+                    d_prime = d_prime_uncond + cfg_scale * (d_prime_cond - d_prime_uncond)
+                x_next = x_cur + (t_next - t_cur) * (0.5 * d_cur + 0.5 * d_prime)
+                
+    return x_next
+
+
+def euler_maruyama_sampler(
+        model,
+        latents,
+        y,
+        context,
+        num_steps=20,
+        heun=False,  # not used, just for compatability
+        cfg_scale=1.0,
+        guidance_low=0.0,
+        guidance_high=1.0,
+        path_type="linear",
+        ):
+    # setup conditioning
+    if cfg_scale > 1.0:
+        y_null = torch.zeros_like(y).to(y.device)
+        context_null = torch.zeros_like(context).to(context.device)
+
+    _dtype = latents.dtype
+    t_steps = torch.linspace(1., 0.04, num_steps, dtype=torch.bfloat16)
+    t_steps = torch.cat([t_steps, torch.tensor([0.], dtype=torch.bfloat16)])
+    x_next = latents.to(torch.bfloat16)
+    device = x_next.device
+
+    with torch.no_grad():
+        for i, (t_cur, t_next) in enumerate(zip(t_steps[:-2], t_steps[1:-1])):
+            dt = t_next - t_cur
+            x_cur = x_next
+            if cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low:
+                model_input = torch.cat([x_cur] * 2, dim=0)
+                y_cur = torch.cat([y, y_null], dim=0)
+                context_cur = torch.cat([context, context_null], dim=0)
+            else:
+                model_input = x_cur
+                y_cur = y            
+                context_cur = context
+            do_guidance = (cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low)
+            kwargs = dict(y=y_cur, context=context_cur, do_guidance=do_guidance)
+            time_input = torch.ones(model_input.size(0)).to(device=device, dtype=torch.bfloat16) * t_cur
+            diffusion = compute_diffusion(t_cur)            
+            eps_i = torch.randn_like(x_cur).to(device)
+            deps = eps_i * torch.sqrt(torch.abs(dt))
+
+            # compute drift
+            v_cur = model(
+                model_input.to(dtype=_dtype), time_input.to(dtype=_dtype), **kwargs
+                )[0]
+            s_cur = get_score_from_velocity(v_cur, model_input, time_input, path_type=path_type)
+            d_cur = v_cur - 0.5 * diffusion * s_cur
+            if cfg_scale > 1. and t_cur <= guidance_high and t_cur >= guidance_low:
+                d_cur_cond, d_cur_uncond = d_cur.chunk(2)
+                d_cur = d_cur_uncond + cfg_scale * (d_cur_cond - d_cur_uncond)
+
+            x_next =  x_cur + d_cur * dt + torch.sqrt(diffusion) * deps
+    
+    # last step
+    t_cur, t_next = t_steps[-2], t_steps[-1]
+    dt = t_next - t_cur
+    x_cur = x_next
+    if cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low:
+        model_input = torch.cat([x_cur] * 2, dim=0)
+        y_cur = torch.cat([y, y_null], dim=0)
+        context_cur = torch.cat([context, context_null], dim=0)
+    else:
+        model_input = x_cur
+        y_cur = y
+        context_cur = context
+    do_guidance = (cfg_scale > 1.0 and t_cur <= guidance_high and t_cur >= guidance_low)
+    kwargs = dict(y=y_cur, context=context_cur, do_guidance=do_guidance)
+    time_input = torch.ones(model_input.size(0)).to(
+        device=device, dtype=torch.bfloat16
+        ) * t_cur
+    
+    # compute drift
+    v_cur = model(
+        model_input.to(dtype=_dtype), time_input.to(dtype=_dtype), **kwargs
+        )[0]
+    s_cur = get_score_from_velocity(v_cur, model_input, time_input, path_type=path_type)
+    diffusion = compute_diffusion(t_cur)
+    d_cur = v_cur - 0.5 * diffusion * s_cur
+    if cfg_scale > 1. and t_cur <= guidance_high and t_cur >= guidance_low:
+        d_cur_cond, d_cur_uncond = d_cur.chunk(2)
+        d_cur = d_cur_uncond + cfg_scale * (d_cur_cond - d_cur_uncond)
+
+    mean_x = x_cur + dt * d_cur
+
+    return mean_x

train.py

@@ -0,0 +1,403 @@
+import argparse
+import copy
+from copy import deepcopy
+import logging
+import os
+from pathlib import Path
+from collections import OrderedDict
+import json
+import fairseq
+import fairseq.utils
+from dataclasses import dataclass
+
+import torch
+from tqdm.auto import tqdm
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import ProjectConfiguration, set_seed
+
+from models import MMDiT
+from loss import SILoss
+
+from dataset import audio_video_spec_fullset_Dataset_Train, collate_fn_taro
+from diffusers import AudioLDM2Pipeline 
+import wandb
+
+logger = get_logger(__name__)
+
+def prob_mask_like(shape, prob, device):
+    if prob == 1:
+        return torch.ones(shape, device = device, dtype = torch.bool)
+    elif prob == 0:
+        return torch.zeros(shape, device = device, dtype = torch.bool)
+    else:
+        return torch.zeros(shape, device = device).float().uniform_(0, 1) < prob
+
+
+@torch.no_grad()
+def sample_posterior(moments, latents_scale=1., latents_bias=0.):
+    mean, std = torch.chunk(moments, 2, dim=1)
+    z = mean + std * torch.randn_like(mean)
+    z = (z * latents_scale + latents_bias) 
+    return z 
+
+
+@torch.no_grad()
+def update_ema(ema_model, model, decay=0.9999):
+    """
+    Step the EMA model towards the current model.
+    """
+    ema_params = OrderedDict(ema_model.named_parameters())
+    model_params = OrderedDict(model.named_parameters())
+
+    for name, param in model_params.items():
+        name = name.replace("module.", "")
+        # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
+        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
+
+
+def create_logger(logging_dir):
+    """
+    Create a logger that writes to a log file and stdout.
+    """
+    logging.basicConfig(
+        level=logging.INFO,
+        format='[\033[34m%(asctime)s\033[0m] %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S',
+        handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
+    )
+    logger = logging.getLogger(__name__)
+    return logger
+
+@dataclass
+class UserDirModule:
+    user_dir: str
+
+def requires_grad(model, flag=True):
+    """
+    Set requires_grad flag for all parameters in a model.
+    """
+    for p in model.parameters():
+        p.requires_grad = flag
+
+
+#################################################################################
+#                                  Training Loop                                #
+#################################################################################
+
+def main(args):    
+    # set accelerator
+    logging_dir = Path(args.output_dir, args.logging_dir)
+    accelerator_project_config = ProjectConfiguration(
+        project_dir=args.output_dir, logging_dir=logging_dir
+        )
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with=args.report_to,
+        project_config=accelerator_project_config,
+    )
+
+    if accelerator.is_main_process:
+        os.makedirs(args.output_dir, exist_ok=True)  # Make results folder (holds all experiment subfolders)
+        save_dir = os.path.join(args.output_dir, args.exp_name)
+        os.makedirs(save_dir, exist_ok=True)
+        args_dict = vars(args)
+        # Save to a JSON file
+        json_dir = os.path.join(save_dir, "args.json")
+        with open(json_dir, 'w') as f:
+            json.dump(args_dict, f, indent=4)
+        checkpoint_dir = f"{save_dir}/checkpoints"  # Stores saved model checkpoints
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        logger = create_logger(save_dir)
+        logger.info(f"Experiment directory created at {save_dir}")
+    device = accelerator.device
+    if torch.backends.mps.is_available():
+        accelerator.native_amp = False    
+    if args.seed is not None:
+        set_seed(args.seed + accelerator.process_index)
+
+    # Create model:
+    assert args.resolution % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder)."
+
+    if args.enc_type == "eat-base":
+        model_dir = 'EAT'
+        model_path = UserDirModule(model_dir)
+        fairseq.utils.import_user_module(model_path)
+        checkpoint_dir_eat = "/home/tton/workspace/SiT_Foley/audio_encoder/EAT-base_epoch30_pt.pt"
+        model_eat, cfg_eat, task_eat = fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_dir_eat])
+        model_eat = model_eat[0].to(device)
+        encoders, encoder_types, architectures = [model_eat], ['eat-base'], ['vit']
+        z_dims = [768]
+    elif args.enc_type == "None":
+        encoders, encoder_types, architectures = [None], [None], [None]
+        z_dims = [0]
+
+    model = MMDiT(
+        adm_in_channels=120,
+        z_dims = z_dims,
+        encoder_depth=args.encoder_depth,
+    )
+
+    model = model.to(device)
+    ema = deepcopy(model).to(device)  # Create an EMA of the model for use after training
+    requires_grad(ema, False)
+    model_audioldm = AudioLDM2Pipeline.from_pretrained("cvssp/audioldm2")
+    vae = model_audioldm.vae.to(device)
+    vae.eval()
+    for param in vae.parameters():
+        param.requires_grad = False
+
+    scale_factor = 0.18215
+    latents_scale = scale_factor
+    latents_bias = 0.
+
+    # create loss function
+    loss_fn = SILoss(
+        prediction=args.prediction,
+        path_type=args.path_type, 
+        encoders=encoders,
+        accelerator=accelerator,
+        latents_scale=latents_scale,
+        latents_bias=latents_bias,
+        weighting=args.weighting
+    )
+    if accelerator.is_main_process:
+        logger.info(f"SiT Parameters: {sum(p.numel() for p in model.parameters()):,}")
+    
+    # Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )    
+    
+    # Setup data:
+    train_dataset = audio_video_spec_fullset_Dataset_Train(args.data_dir)
+    local_batch_size = int(args.batch_size // accelerator.num_processes)
+    train_dataloader = DataLoader(
+        train_dataset,
+        batch_size=local_batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=True,
+        collate_fn=collate_fn_taro,
+    )
+    if accelerator.is_main_process:
+        logger.info(f"Dataset contains {len(train_dataset):,} images ({args.data_dir})")
+    
+    # Prepare models for training:
+    update_ema(ema, model, decay=0)  # Ensure EMA is initialized with synced weights
+    model.train()  # important! This enables embedding dropout for classifier-free guidance
+    ema.eval()  # EMA model should always be in eval mode
+    
+    # resume:
+    global_step = 0
+    if args.resume_step > 0:
+        ckpt_name = str(args.resume_step).zfill(7) +'.pt'
+        ckpt = torch.load(
+            f'{os.path.join(args.output_dir, args.exp_name)}/checkpoints/{ckpt_name}',
+            map_location='cpu',
+            )
+        model.load_state_dict(ckpt['model'])
+        ema.load_state_dict(ckpt['ema'])
+        optimizer.load_state_dict(ckpt['opt'])
+        global_step = ckpt['steps']
+
+    model, optimizer, train_dataloader = accelerator.prepare(
+        model, optimizer, train_dataloader
+    )
+
+    if accelerator.is_main_process:
+        tracker_config = vars(copy.deepcopy(args))
+        accelerator.init_trackers(
+            project_name="REPA", 
+            config=tracker_config,
+            init_kwargs={
+                "wandb": {"name": f"{args.exp_name}"}
+            },
+        )
+        
+    progress_bar = tqdm(
+        range(0, args.max_train_steps),
+        initial=global_step,
+        desc="Steps",
+        # Only show the progress bar once on each machine.
+        disable=not accelerator.is_local_main_process,
+    )
+
+    # Labels to condition the model with (feel free to change):
+    sample_batch_size = 64 // accelerator.num_processes
+    batch_gt_xs = next(iter(train_dataloader))
+    gt_xs = batch_gt_xs["mix_spec"]
+
+    with torch.no_grad():
+        gt_xs = gt_xs[:sample_batch_size]
+        encoder_posterior_gt_xs = vae.encode(gt_xs.to(device))[0]
+        gt_xs = encoder_posterior_gt_xs.sample() * scale_factor
+    ys = batch_gt_xs["mix_onset"][:sample_batch_size]
+    ys = ys.to(device)
+    contexts = batch_gt_xs["mix_video_feat"][:sample_batch_size]
+    contexts = contexts.to(device)
+
+    for epoch in range(args.epochs):
+        model.train()
+        for batch_idx in train_dataloader:
+            raw_spec = batch_idx["mix_spec"]
+            context = batch_idx["mix_video_feat"]
+            y = batch_idx["mix_onset"]
+            with torch.no_grad():
+                raw_spec = raw_spec.permute(0, 1, 3, 2)
+                encoder_posterior = vae.encode(raw_spec.to(device), return_dict=True)[0]
+                x = encoder_posterior.sample() * scale_factor
+            raw_image = batch_idx["mix_fbank"]
+
+            x = x.squeeze(dim=1).to(device)
+            context = context.to(device)
+            y = y.to(device)
+            z = None
+            labels = context
+            with torch.no_grad():
+                zs = []
+                with accelerator.autocast():
+                    for encoder, encoder_type, arch in zip(encoders, encoder_types, architectures):
+                        if encoder is not None:
+                            raw_image = raw_image[:, :834].unsqueeze(1)
+                            z = encoder.extract_features(raw_image, padding_mask=None,mask=False, remove_extra_tokens=True)['x'] # B x 416 x 768
+                            zs.append(z)
+
+            with accelerator.accumulate(model):
+                model_kwargs = dict(context=labels, y=y)
+                loss, proj_loss = loss_fn(model, x, model_kwargs, zs=zs)
+                loss_mean = loss.mean()
+                if len(zs) > 0:
+                    proj_loss_mean = proj_loss.mean()
+                    loss = loss_mean + proj_loss_mean * args.proj_coeff
+                else:
+                    proj_loss_mean = torch.tensor(0., device=device)
+                    loss = loss_mean
+                    
+                ## optimization
+                accelerator.backward(loss)
+                if accelerator.sync_gradients:
+                    params_to_clip = model.parameters()
+                    grad_norm = accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+                optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+
+                if accelerator.sync_gradients:
+                    update_ema(ema, model) # change ema function
+            
+            ### enter
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1                
+            if global_step % args.checkpointing_steps == 0 and global_step > 0:
+                if accelerator.is_main_process:
+                    checkpoint = {
+                        "model": model.state_dict(),
+                        "ema": ema.state_dict(),
+                        "opt": optimizer.state_dict(),
+                        "args": args,
+                        "steps": global_step,
+                    }
+                    checkpoint_path = f"{checkpoint_dir}/{global_step:07d}.pt"
+                    torch.save(checkpoint, checkpoint_path)
+                    logger.info(f"Saved checkpoint to {checkpoint_path}")
+
+            logs = {
+                "loss": accelerator.gather(loss_mean).mean().detach().item(), 
+                "proj_loss": accelerator.gather(proj_loss_mean).mean().detach().item(),
+                "grad_norm": accelerator.gather(grad_norm).mean().detach().item()
+            }
+            progress_bar.set_postfix(**logs)
+            accelerator.log(logs, step=global_step)
+
+            if global_step >= args.max_train_steps:
+                break
+        if global_step >= args.max_train_steps:
+            break
+
+    model.eval()  # important! This disables randomized embedding dropout
+    # do any sampling/FID calculation/etc. with ema (or model) in eval mode ...
+    
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+        logger.info("Done!")
+    accelerator.end_training()
+
+def parse_args(input_args=None):
+    parser = argparse.ArgumentParser(description="Training")
+
+    # logging:
+    parser.add_argument("--output-dir", type=str, default="exps")
+    parser.add_argument("--exp-name", type=str, required=True)
+    parser.add_argument("--logging-dir", type=str, default="logs")
+    parser.add_argument("--report-to", type=str, default="wandb")
+    parser.add_argument("--sampling-steps", type=int, default=10000)
+    parser.add_argument("--resume-step", type=int, default=0)
+
+    # model
+    parser.add_argument("--model", type=str)
+    parser.add_argument("--num-classes", type=int, default=1000)
+    parser.add_argument("--encoder-depth", type=int, default=8)
+    parser.add_argument("--fused-attn", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--qk-norm",  action=argparse.BooleanOptionalAction, default=False)
+
+    # dataset
+    parser.add_argument("--data-dir", type=str, default="/home/tton/tton_data/data/VGGSound")
+    parser.add_argument("--resolution", type=int, choices=[256], default=256)
+    parser.add_argument("--batch-size", type=int, default=64)
+
+    # precision
+    parser.add_argument("--allow-tf32", action="store_true")
+    parser.add_argument("--mixed-precision", type=str, default="fp16", choices=["no", "fp16", "bf16"])
+
+    # optimization
+    parser.add_argument("--epochs", type=int, default=1400)
+    parser.add_argument("--max-train-steps", type=int, default=1000000)
+    parser.add_argument("--checkpointing-steps", type=int, default=50000)
+    parser.add_argument("--gradient-accumulation-steps", type=int, default=1)
+    parser.add_argument("--learning-rate", type=float, default=1e-4)
+    parser.add_argument("--adam-beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam-beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
+    parser.add_argument("--adam-weight-decay", type=float, default=0., help="Weight decay to use.")
+    parser.add_argument("--adam-epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
+    parser.add_argument("--max-grad-norm", default=1.0, type=float, help="Max gradient norm.")
+
+    # seed
+    parser.add_argument("--seed", type=int, default=0)
+
+    # cpu
+    parser.add_argument("--num-workers", type=int, default=4)
+
+    # loss
+    parser.add_argument("--path-type", type=str, default="linear", choices=["linear", "cosine"])
+    parser.add_argument("--prediction", type=str, default="v", choices=["v"]) # currently we only support v-prediction
+    parser.add_argument("--cfg-prob", type=float, default=0.1)
+    parser.add_argument("--enc-type", type=str, default='dinov2-vit-b')
+    parser.add_argument("--proj-coeff", type=float, default=0.5)
+    parser.add_argument("--weighting", default="uniform", type=str, help="Max gradient norm.")
+    parser.add_argument("--legacy", action=argparse.BooleanOptionalAction, default=False)
+
+    if input_args is not None:
+        args = parser.parse_args(input_args)
+    else:
+        args = parser.parse_args()
+        
+    return args
+
+if __name__ == "__main__":
+    args = parse_args()
+    
+    main(args)

train.sh

@@ -0,0 +1,15 @@
+CUDA_VISIBLE_DEVICES=0 accelerate launch train.py \
+  --report-to="wandb" \
+  --allow-tf32 \
+  --mixed-precision="fp16" \
+  --seed=0 \
+  --path-type="linear" \
+  --prediction="v" \
+  --weighting="uniform" \
+  --model="SiT-B/2" \
+  --enc-type="eat-base" \
+  --proj-coeff=0.5 \
+  --encoder-depth=4 \
+  --output-dir="exps" \
+  --exp-name="taro-output" \
+  --data-dir="./VGGSound" \