initial commit

.gitignore

@@ -0,0 +1,146 @@
+run_*.sh
+log/
+saves
+saves/
+weights/
+weights
+output/
+output
+pretrained/
+workspace
+workspace/
+ext_weights/
+ext_weights
+.checkpoints/
+.vscode/
+training/example_output/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

Dockerfile

@@ -0,0 +1,39 @@
+FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
+
+WORKDIR /code
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    python3.9 \
+    python3-pip \
+    git \
+    ffmpeg \
+    libsm6 \
+    libxext6 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Python dependencies
+COPY requirements.txt .
+RUN pip3 install --no-cache-dir -r requirements.txt
+
+# Clone and install MMAudio
+RUN git clone https://github.com/hkchengrex/MMAudio.git && \
+    cd MMAudio && \
+    pip3 install -e .
+
+# Copy the application files
+COPY app.py .
+
+# Create output directory
+RUN mkdir -p output/gradio && chmod 777 output/gradio
+
+# Set environment variables for Hugging Face Spaces
+ENV PYTHONUNBUFFERED=1
+ENV GRADIO_SERVER_NAME=0.0.0.0
+ENV GRADIO_SERVER_PORT=7860
+
+# Expose the port
+EXPOSE 7860
+
+# Run the Gradio app
+CMD ["python3", "app.py"]

LICENSE

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

README.md

@@ -0,0 +1,56 @@
+---
+title: Sonisphere
+emoji: 🐢
+colorFrom: green
+colorTo: gray
+sdk: gradio
+sdk_version: 5.20.0
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# Sonisphere Demo
+
+This is a Hugging Face Spaces demo for [MMAudio](https://hkchengrex.com/MMAudio/), a powerful model for generating realistic audio for videos.
+
+## 🎥 Features
+
+- Upload any video and generate matching audio
+- Control the generation with text prompts
+- Adjust generation parameters like steps and guidance strength
+- Process videos up to 30 seconds in length
+
+## 🚀 Usage
+
+1. Upload a video or use one of the example videos
+2. Enter a text prompt describing the desired audio
+3. (Optional) Add a negative prompt to specify what you don't want
+4. Adjust the generation parameters if needed
+5. Click "Submit" and wait for the generation to complete
+
+## ⚙️ Parameters
+
+- **Prompt**: Describe the audio you want to generate
+- **Negative prompt**: Specify what you don't want in the audio (default: "music")
+- **Seed**: Control randomness (-1 for random seed)
+- **Number of steps**: More steps = better quality but slower (default: 25)
+- **Guidance Strength**: Controls how closely the generation follows the prompt (default: 4.5)
+- **Duration**: Length of the generated audio in seconds (default: 8)
+
+## 📝 Notes
+
+- Processing high-resolution videos (>384px on shorter side) takes longer and doesn't improve results
+- The model works best with videos between 5-30 seconds
+- Generation time depends on video length and number of steps
+
+## 🔗 Links
+
+- [Project Page](https://hkchengrex.com/MMAudio/)
+- [GitHub Repository](https://github.com/hkchengrex/MMAudio)
+- [Paper](https://arxiv.org/abs/2401.09774)
+
+## 📜 License
+
+This demo uses the MMAudio model which is released under the [MIT license](https://github.com/hkchengrex/MMAudio/blob/main/LICENSE).

app.py

@@ -0,0 +1,128 @@
+import gc
+import logging
+from datetime import datetime
+from fractions import Fraction
+from pathlib import Path
+
+import gradio as gr
+import torch
+import torchaudio
+
+from mmaudio.eval_utils import (ModelConfig, VideoInfo, all_model_cfg, generate, load_image,
+                               load_video, make_video, setup_eval_logging)
+from mmaudio.model.flow_matching import FlowMatching
+from mmaudio.model.networks import MMAudio, get_my_mmaudio
+from mmaudio.model.sequence_config import SequenceConfig
+from mmaudio.model.utils.features_utils import FeaturesUtils
+
+# Setup logging
+setup_eval_logging()
+log = logging.getLogger()
+
+# Configure device and dtype
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+if device == 'cpu':
+    log.warning('CUDA is not available, running on CPU')
+dtype = torch.bfloat16
+
+# Configure model and paths
+model: ModelConfig = all_model_cfg['large_44k_v2']
+model.download_if_needed()
+output_dir = Path('./output/gradio')
+output_dir.mkdir(exist_ok=True, parents=True)
+
+def get_model() -> tuple[MMAudio, FeaturesUtils, SequenceConfig]:
+    seq_cfg = model.seq_cfg
+    
+    net: MMAudio = get_my_mmaudio(model.model_name).to(device, dtype).eval()
+    net.load_weights(torch.load(model.model_path, map_location=device, weights_only=True))
+    log.info(f'Loaded weights from {model.model_path}')
+    
+    feature_utils = FeaturesUtils(tod_vae_ckpt=model.vae_path,
+                                synchformer_ckpt=model.synchformer_ckpt,
+                                enable_conditions=True,
+                                mode=model.mode,
+                                bigvgan_vocoder_ckpt=model.bigvgan_16k_path,
+                                need_vae_encoder=False)
+    feature_utils = feature_utils.to(device, dtype).eval()
+    
+    return net, feature_utils, seq_cfg
+
+# Load model once at startup
+net, feature_utils, seq_cfg = get_model()
+
+@torch.inference_mode()
+def video_to_audio(video: gr.Video, prompt: str, negative_prompt: str, seed: int, num_steps: int,
+                   cfg_strength: float, duration: float):
+    try:
+        rng = torch.Generator(device=device)
+        if seed >= 0:
+            rng.manual_seed(seed)
+        else:
+            rng.seed()
+        fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=num_steps)
+
+        video_info = load_video(video, duration)
+        clip_frames = video_info.clip_frames.unsqueeze(0)
+        sync_frames = video_info.sync_frames.unsqueeze(0)
+        duration = video_info.duration_sec
+        
+        seq_cfg.duration = duration
+        net.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
+
+        audios = generate(clip_frames, sync_frames, [prompt],
+                        negative_text=[negative_prompt],
+                        feature_utils=feature_utils,
+                        net=net,
+                        fm=fm,
+                        rng=rng,
+                        cfg_strength=cfg_strength)
+        audio = audios.float().cpu()[0]
+
+        current_time_string = datetime.now().strftime('%Y%m%d_%H%M%S')
+        video_save_path = output_dir / f'{current_time_string}.mp4'
+        make_video(video_info, video_save_path, audio, sampling_rate=seq_cfg.sampling_rate)
+        
+        gc.collect()
+        torch.cuda.empty_cache()
+        
+        return video_save_path
+    except Exception as e:
+        log.error(f"Error in video_to_audio: {str(e)}")
+        raise gr.Error(f"An error occurred: {str(e)}")
+
+# Create the Gradio interface
+demo = gr.Interface(
+    fn=video_to_audio,
+    title="MMAudio — Video-to-Audio Synthesis",
+    description="""
+    Generate realistic audio for your videos using MMAudio!
+    
+    Project page: [MMAudio](https://hkchengrex.com/MMAudio/)
+    Code: [GitHub](https://github.com/hkchengrex/MMAudio)
+    
+    Note: Processing high-resolution videos (>384px on shorter side) takes longer and doesn't improve results.
+    """,
+    inputs=[
+        gr.Video(label="Upload Video"),
+        gr.Text(label="Prompt", placeholder="Describe the audio you want to generate..."),
+        gr.Text(label="Negative prompt", value="music", placeholder="What you don't want in the audio..."),
+        gr.Number(label="Seed (-1: random)", value=-1, precision=0, minimum=-1),
+        gr.Number(label="Number of steps", value=25, precision=0, minimum=1),
+        gr.Slider(label="Guidance Strength", value=4.5, minimum=1, maximum=10, step=0.5),
+        gr.Slider(label="Duration (seconds)", value=8, minimum=1, maximum=30, step=1),
+    ],
+    outputs=gr.Video(label="Generated Result"),
+    examples=[
+        ["https://huggingface.co/hkchengrex/MMAudio/resolve/main/examples/sora_beach.mp4",
+         "waves, seagulls", "", 0, 25, 4.5, 10],
+        ["https://huggingface.co/hkchengrex/MMAudio/resolve/main/examples/sora_serpent.mp4",
+         "", "music", 0, 25, 4.5, 10],
+    ],
+    cache_examples=True,
+)
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+

batch_eval.py

@@ -0,0 +1,110 @@
+import logging
+import os
+from pathlib import Path
+
+import hydra
+import torch
+import torch.distributed as distributed
+import torchaudio
+from hydra.core.hydra_config import HydraConfig
+from omegaconf import DictConfig
+from tqdm import tqdm
+
+from mmaudio.data.data_setup import setup_eval_dataset
+from mmaudio.eval_utils import ModelConfig, all_model_cfg, generate
+from mmaudio.model.flow_matching import FlowMatching
+from mmaudio.model.networks import MMAudio, get_my_mmaudio
+from mmaudio.model.utils.features_utils import FeaturesUtils
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+local_rank = int(os.environ['LOCAL_RANK'])
+world_size = int(os.environ['WORLD_SIZE'])
+log = logging.getLogger()
+
+
+@torch.inference_mode()
+@hydra.main(version_base='1.3.2', config_path='config', config_name='eval_config.yaml')
+def main(cfg: DictConfig):
+    device = 'cuda'
+    torch.cuda.set_device(local_rank)
+
+    if cfg.model not in all_model_cfg:
+        raise ValueError(f'Unknown model variant: {cfg.model}')
+    model: ModelConfig = all_model_cfg[cfg.model]
+    model.download_if_needed()
+    seq_cfg = model.seq_cfg
+
+    run_dir = Path(HydraConfig.get().run.dir)
+    if cfg.output_name is None:
+        output_dir = run_dir / cfg.dataset
+    else:
+        output_dir = run_dir / f'{cfg.dataset}-{cfg.output_name}'
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # load a pretrained model
+    seq_cfg.duration = cfg.duration_s
+    net: MMAudio = get_my_mmaudio(cfg.model).to(device).eval()
+    net.load_weights(torch.load(model.model_path, map_location=device, weights_only=True))
+    log.info(f'Loaded weights from {model.model_path}')
+    net.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
+    log.info(f'Latent seq len: {seq_cfg.latent_seq_len}')
+    log.info(f'Clip seq len: {seq_cfg.clip_seq_len}')
+    log.info(f'Sync seq len: {seq_cfg.sync_seq_len}')
+
+    # misc setup
+    rng = torch.Generator(device=device)
+    rng.manual_seed(cfg.seed)
+    fm = FlowMatching(cfg.sampling.min_sigma,
+                      inference_mode=cfg.sampling.method,
+                      num_steps=cfg.sampling.num_steps)
+
+    feature_utils = FeaturesUtils(tod_vae_ckpt=model.vae_path,
+                                  synchformer_ckpt=model.synchformer_ckpt,
+                                  enable_conditions=True,
+                                  mode=model.mode,
+                                  bigvgan_vocoder_ckpt=model.bigvgan_16k_path,
+                                  need_vae_encoder=False)
+    feature_utils = feature_utils.to(device).eval()
+
+    if cfg.compile:
+        net.preprocess_conditions = torch.compile(net.preprocess_conditions)
+        net.predict_flow = torch.compile(net.predict_flow)
+        feature_utils.compile()
+
+    dataset, loader = setup_eval_dataset(cfg.dataset, cfg)
+
+    with torch.amp.autocast(enabled=cfg.amp, dtype=torch.bfloat16, device_type=device):
+        for batch in tqdm(loader):
+            audios = generate(batch.get('clip_video', None),
+                              batch.get('sync_video', None),
+                              batch.get('caption', None),
+                              feature_utils=feature_utils,
+                              net=net,
+                              fm=fm,
+                              rng=rng,
+                              cfg_strength=cfg.cfg_strength,
+                              clip_batch_size_multiplier=64,
+                              sync_batch_size_multiplier=64)
+            audios = audios.float().cpu()
+            names = batch['name']
+            for audio, name in zip(audios, names):
+                torchaudio.save(output_dir / f'{name}.flac', audio, seq_cfg.sampling_rate)
+
+
+def distributed_setup():
+    distributed.init_process_group(backend="nccl")
+    local_rank = distributed.get_rank()
+    world_size = distributed.get_world_size()
+    log.info(f'Initialized: local_rank={local_rank}, world_size={world_size}')
+    return local_rank, world_size
+
+
+if __name__ == '__main__':
+    distributed_setup()
+
+    main()
+
+    # clean-up
+    distributed.destroy_process_group()

config/base_config.yaml

@@ -0,0 +1,62 @@
+defaults:
+  - data: base
+  - eval_data: base
+  - override hydra/job_logging: custom-simplest
+  - _self_
+
+hydra:
+  run:
+    dir: ./output/${exp_id}
+  output_subdir: ${now:%Y-%m-%d_%H-%M-%S}-hydra
+
+enable_email: False
+
+model: small_16k
+
+exp_id: default
+debug: False
+cudnn_benchmark: True
+compile: True
+amp: True
+weights: null
+checkpoint: null
+seed: 14159265
+num_workers: 10 # per-GPU
+pin_memory: False # set to True if your system can handle it, i.e., have enough memory
+
+# NOTE: This DOSE NOT affect the model during inference in any way
+# they are just for the dataloader to fill in the missing data in multi-modal loading
+# to change the sequence length for the model, see networks.py
+data_dim:
+  text_seq_len: 77
+  clip_dim: 1024
+  sync_dim: 768
+  text_dim: 1024
+
+# ema configuration
+ema:
+  enable: True 
+  sigma_rels: [0.05, 0.1]
+  update_every: 1
+  checkpoint_every: 5_000
+  checkpoint_folder: ${hydra:run.dir}/ema_ckpts
+  default_output_sigma: 0.05
+
+
+# sampling
+sampling:
+  mean: 0.0
+  scale: 1.0
+  min_sigma: 0.0
+  method: euler
+  num_steps: 25
+
+# classifier-free guidance
+null_condition_probability: 0.1
+cfg_strength: 4.5
+
+# checkpoint paths to external modules
+vae_16k_ckpt: ./ext_weights/v1-16.pth
+vae_44k_ckpt: ./ext_weights/v1-44.pth
+bigvgan_vocoder_ckpt: ./ext_weights/best_netG.pt
+synchformer_ckpt: ./ext_weights/synchformer_state_dict.pth

config/data/base.yaml

@@ -0,0 +1,70 @@
+VGGSound:
+  root: ../data/video
+  subset_name: sets/vgg3-train.tsv
+  fps: 8
+  height: 384
+  width: 384
+  sample_duration_sec: 8.0
+
+VGGSound_test:
+  root: ../data/video
+  subset_name: sets/vgg3-test.tsv
+  fps: 8
+  height: 384
+  width: 384
+  sample_duration_sec: 8.0
+
+VGGSound_val:
+  root: ../data/video
+  subset_name: sets/vgg3-val.tsv
+  fps: 8
+  height: 384
+  width: 384
+  sample_duration_sec: 8.0
+
+ExtractedVGG:
+  tsv: ../data/v1-16-memmap/vgg-train.tsv
+  memmap_dir: ../data/v1-16-memmap/vgg-train
+
+ExtractedVGG_test:
+  tag: test
+  gt_cache: ../data/eval-cache/vggsound-test
+  output_subdir: null
+  tsv: ../data/v1-16-memmap/vgg-test.tsv
+  memmap_dir: ../data/v1-16-memmap/vgg-test
+
+ExtractedVGG_val:
+  tag: val
+  gt_cache: ../data/eval-cache/vggsound-val
+  output_subdir: val
+  tsv: ../data/v1-16-memmap/vgg-val.tsv
+  memmap_dir: ../data/v1-16-memmap/vgg-val
+
+AudioCaps:
+  tsv: ../data/v1-16-memmap/audiocaps.tsv
+  memmap_dir: ../data/v1-16-memmap/audiocaps
+
+AudioSetSL:
+  tsv: ../data/v1-16-memmap/audioset_sl.tsv
+  memmap_dir: ../data/v1-16-memmap/audioset_sl
+
+BBCSound:
+  tsv: ../data/v1-16-memmap/bbcsound.tsv
+  memmap_dir: ../data/v1-16-memmap/bbcsound
+
+FreeSound:
+  tsv: ../data/v1-16-memmap/freesound.tsv
+  memmap_dir: ../data/v1-16-memmap/freesound
+
+Clotho:
+  tsv: ../data/v1-16-memmap/clotho.tsv
+  memmap_dir: ../data/v1-16-memmap/clotho
+
+Example_video:
+  tsv: ./training/example_output/memmap/vgg-example.tsv
+  memmap_dir: ./training/example_output/memmap/vgg-example
+
+Example_audio:
+  tsv: ./training/example_output/memmap/audio-example.tsv
+  memmap_dir: ./training/example_output/memmap/audio-example
+  

config/eval_config.yaml

@@ -0,0 +1,17 @@
+defaults:
+  - base_config
+  - override hydra/job_logging: custom-simplest
+  - _self_
+
+hydra:
+  run:
+    dir: ./output/${exp_id}
+  output_subdir: eval-${now:%Y-%m-%d_%H-%M-%S}-hydra
+
+exp_id: ${model}
+dataset: audiocaps
+duration_s: 8.0
+
+# for inference, this is the per-GPU batch size
+batch_size: 16
+output_name: null

config/eval_data/base.yaml

@@ -0,0 +1,22 @@
+AudioCaps:
+  audio_path: ../data/AudioCaps-test-audioldm-ver
+  # a csv file, with a header row of 'name' and 'caption'
+  # name should match the audio file name without extension
+  # Can be downloaded here: https://github.com/hkchengrex/MMAudio/releases/download/v0.1/AudioCaps_audioldm_data.csv
+  csv_path: ../data/AudioCaps-test-audioldm-ver/data.csv
+
+AudioCaps_full:
+  audio_path: ../data/AudioCaps-test-full-ver
+  # a csv file, with a header row of 'name' and 'caption'
+  # name should match the audio file name without extension
+  # Can be downloaded here: https://github.com/hkchengrex/MMAudio/releases/download/v0.1/AudioCaps_full_data.csv
+  csv_path: ../data/AudioCaps-test-full-ver/data.csv
+
+MovieGen:
+  video_path: ../data/MovieGen/MovieGenAudioBenchSfx/video_with_audio
+  jsonl_path: ../data/MovieGen/MovieGenAudioBenchSfx/metadata
+  
+VGGSound:
+  video_path: ../data/test-videos
+  # from the officially released csv file
+  csv_path: ../data/vggsound.csv

config/hydra/job_logging/custom-eval.yaml

@@ -0,0 +1,32 @@
+# python logging configuration for tasks
+version: 1
+formatters:
+  simple:
+    format: '[%(asctime)s][%(levelname)s][r${oc.env:LOCAL_RANK}] - %(message)s'
+    datefmt: '%Y-%m-%d %H:%M:%S'
+  colorlog:
+    '()': 'colorlog.ColoredFormatter'
+    format: '[%(cyan)s%(asctime)s%(reset)s][%(log_color)s%(levelname)s%(reset)s] - %(message)s'
+    datefmt: '%Y-%m-%d %H:%M:%S'
+    log_colors:
+      DEBUG: purple
+      INFO: green
+      WARNING: yellow
+      ERROR: red
+      CRITICAL: red
+handlers:
+  console:
+    class: logging.StreamHandler
+    formatter: colorlog
+    stream: ext://sys.stdout
+  file:
+    class: logging.FileHandler
+    formatter: simple
+    # absolute file path
+    filename: ${hydra.runtime.output_dir}/eval-${now:%Y-%m-%d_%H-%M-%S}-rank${oc.env:LOCAL_RANK}.log
+    mode: w
+root:
+  level: INFO
+  handlers: [console, file]
+
+disable_existing_loggers: false

config/hydra/job_logging/custom-no-rank.yaml

@@ -0,0 +1,32 @@
+# python logging configuration for tasks
+version: 1
+formatters:
+  simple:
+    format: '[%(asctime)s][%(levelname)s] - %(message)s'
+    datefmt: '%Y-%m-%d %H:%M:%S'
+  colorlog:
+    '()': 'colorlog.ColoredFormatter'
+    format: '[%(cyan)s%(asctime)s%(reset)s][%(log_color)s%(levelname)s%(reset)s] - %(message)s'
+    datefmt: '%Y-%m-%d %H:%M:%S'
+    log_colors:
+      DEBUG: purple
+      INFO: green
+      WARNING: yellow
+      ERROR: red
+      CRITICAL: red
+handlers:
+  console:
+    class: logging.StreamHandler
+    formatter: colorlog
+    stream: ext://sys.stdout
+  file:
+    class: logging.FileHandler
+    formatter: simple
+    # absolute file path
+    filename: ${hydra.runtime.output_dir}/${now:%Y-%m-%d_%H-%M-%S}-eval.log
+    mode: w
+root:
+  level: INFO
+  handlers: [console, file]
+
+disable_existing_loggers: false

config/hydra/job_logging/custom-simplest.yaml

@@ -0,0 +1,26 @@
+# python logging configuration for tasks
+version: 1
+formatters:
+  simple:
+    format: '[%(asctime)s][%(levelname)s] - %(message)s'
+    datefmt: '%Y-%m-%d %H:%M:%S'
+  colorlog:
+    '()': 'colorlog.ColoredFormatter'
+    format: '[%(cyan)s%(asctime)s%(reset)s][%(log_color)s%(levelname)s%(reset)s] - %(message)s'
+    datefmt: '%Y-%m-%d %H:%M:%S'
+    log_colors:
+      DEBUG: purple
+      INFO: green
+      WARNING: yellow
+      ERROR: red
+      CRITICAL: red
+handlers:
+  console:
+    class: logging.StreamHandler
+    formatter: colorlog
+    stream: ext://sys.stdout
+root:
+  level: INFO
+  handlers: [console]
+
+disable_existing_loggers: false

config/hydra/job_logging/custom.yaml

@@ -0,0 +1,33 @@
+# @package hydra.job_logging
+# python logging configuration for tasks
+version: 1
+formatters:
+  simple:
+    format: '[%(asctime)s][%(levelname)s][r${oc.env:LOCAL_RANK}] - %(message)s'
+    datefmt: '%Y-%m-%d %H:%M:%S'
+  colorlog:
+    '()': 'colorlog.ColoredFormatter'
+    format: '[%(cyan)s%(asctime)s%(reset)s][%(blue)sr${oc.env:LOCAL_RANK}%(reset)s][%(log_color)s%(levelname)s%(reset)s] - %(message)s'
+    datefmt: '%Y-%m-%d %H:%M:%S'
+    log_colors:
+      DEBUG: purple
+      INFO: green
+      WARNING: yellow
+      ERROR: red
+      CRITICAL: red
+handlers:
+  console:
+    class: logging.StreamHandler
+    formatter: colorlog
+    stream: ext://sys.stdout
+  file:
+    class: logging.FileHandler
+    formatter: simple
+    # absolute file path
+    filename: ${hydra.runtime.output_dir}/train-${now:%Y-%m-%d_%H-%M-%S}-rank${oc.env:LOCAL_RANK}.log
+    mode: w
+root:
+  level: INFO
+  handlers: [console, file]
+
+disable_existing_loggers: false

config/train_config.yaml

@@ -0,0 +1,41 @@
+defaults:
+  - base_config
+  - override data: base
+  - override hydra/job_logging: custom
+  - _self_
+
+hydra:
+  run:
+    dir: ./output/${exp_id}
+  output_subdir: train-${now:%Y-%m-%d_%H-%M-%S}-hydra
+
+ema: 
+  start: 0
+
+mini_train: False
+example_train: False
+enable_grad_scaler: False
+vgg_oversample_rate: 5
+
+log_text_interval: 200
+log_extra_interval: 20_000
+val_interval: 5_000
+eval_interval: 20_000
+save_eval_interval: 40_000
+save_weights_interval: 10_000
+save_checkpoint_interval: 10_000
+save_copy_iterations: []
+
+batch_size: 512
+eval_batch_size: 256 # per-GPU
+
+num_iterations: 300_000
+learning_rate: 1.0e-4
+linear_warmup_steps: 1_000
+
+lr_schedule: step
+lr_schedule_steps: [240_000, 270_000]
+lr_schedule_gamma: 0.1
+
+clip_grad_norm: 1.0
+weight_decay: 1.0e-6

demo.py

@@ -0,0 +1,141 @@
+import logging
+from argparse import ArgumentParser
+from pathlib import Path
+
+import torch
+import torchaudio
+
+from mmaudio.eval_utils import (ModelConfig, all_model_cfg, generate, load_video, make_video,
+                                setup_eval_logging)
+from mmaudio.model.flow_matching import FlowMatching
+from mmaudio.model.networks import MMAudio, get_my_mmaudio
+from mmaudio.model.utils.features_utils import FeaturesUtils
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+log = logging.getLogger()
+
+
+@torch.inference_mode()
+def main():
+    setup_eval_logging()
+
+    parser = ArgumentParser()
+    parser.add_argument('--variant',
+                        type=str,
+                        default='large_44k_v2',
+                        help='small_16k, small_44k, medium_44k, large_44k, large_44k_v2')
+    parser.add_argument('--video', type=Path, help='Path to the video file')
+    parser.add_argument('--prompt', type=str, help='Input prompt', default='')
+    parser.add_argument('--negative_prompt', type=str, help='Negative prompt', default='')
+    parser.add_argument('--duration', type=float, default=8.0)
+    parser.add_argument('--cfg_strength', type=float, default=4.5)
+    parser.add_argument('--num_steps', type=int, default=25)
+
+    parser.add_argument('--mask_away_clip', action='store_true')
+
+    parser.add_argument('--output', type=Path, help='Output directory', default='./output')
+    parser.add_argument('--seed', type=int, help='Random seed', default=42)
+    parser.add_argument('--skip_video_composite', action='store_true')
+    parser.add_argument('--full_precision', action='store_true')
+
+    args = parser.parse_args()
+
+    if args.variant not in all_model_cfg:
+        raise ValueError(f'Unknown model variant: {args.variant}')
+    model: ModelConfig = all_model_cfg[args.variant]
+    model.download_if_needed()
+    seq_cfg = model.seq_cfg
+
+    if args.video:
+        video_path: Path = Path(args.video).expanduser()
+    else:
+        video_path = None
+    prompt: str = args.prompt
+    negative_prompt: str = args.negative_prompt
+    output_dir: str = args.output.expanduser()
+    seed: int = args.seed
+    num_steps: int = args.num_steps
+    duration: float = args.duration
+    cfg_strength: float = args.cfg_strength
+    skip_video_composite: bool = args.skip_video_composite
+    mask_away_clip: bool = args.mask_away_clip
+
+    device = 'cpu'
+    if torch.cuda.is_available():
+        device = 'cuda'
+    elif torch.backends.mps.is_available():
+        device = 'mps'
+    else:
+        log.warning('CUDA/MPS are not available, running on CPU')
+    dtype = torch.float32 if args.full_precision else torch.bfloat16
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # load a pretrained model
+    net: MMAudio = get_my_mmaudio(model.model_name).to(device, dtype).eval()
+    net.load_weights(torch.load(model.model_path, map_location=device, weights_only=True))
+    log.info(f'Loaded weights from {model.model_path}')
+
+    # misc setup
+    rng = torch.Generator(device=device)
+    rng.manual_seed(seed)
+    fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=num_steps)
+
+    feature_utils = FeaturesUtils(tod_vae_ckpt=model.vae_path,
+                                  synchformer_ckpt=model.synchformer_ckpt,
+                                  enable_conditions=True,
+                                  mode=model.mode,
+                                  bigvgan_vocoder_ckpt=model.bigvgan_16k_path,
+                                  need_vae_encoder=False)
+    feature_utils = feature_utils.to(device, dtype).eval()
+
+    if video_path is not None:
+        log.info(f'Using video {video_path}')
+        video_info = load_video(video_path, duration)
+        clip_frames = video_info.clip_frames
+        sync_frames = video_info.sync_frames
+        duration = video_info.duration_sec
+        if mask_away_clip:
+            clip_frames = None
+        else:
+            clip_frames = clip_frames.unsqueeze(0)
+        sync_frames = sync_frames.unsqueeze(0)
+    else:
+        log.info('No video provided -- text-to-audio mode')
+        clip_frames = sync_frames = None
+
+    seq_cfg.duration = duration
+    net.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
+
+    log.info(f'Prompt: {prompt}')
+    log.info(f'Negative prompt: {negative_prompt}')
+
+    audios = generate(clip_frames,
+                      sync_frames, [prompt],
+                      negative_text=[negative_prompt],
+                      feature_utils=feature_utils,
+                      net=net,
+                      fm=fm,
+                      rng=rng,
+                      cfg_strength=cfg_strength)
+    audio = audios.float().cpu()[0]
+    if video_path is not None:
+        save_path = output_dir / f'{video_path.stem}.flac'
+    else:
+        safe_filename = prompt.replace(' ', '_').replace('/', '_').replace('.', '')
+        save_path = output_dir / f'{safe_filename}.flac'
+    torchaudio.save(save_path, audio, seq_cfg.sampling_rate)
+
+    log.info(f'Audio saved to {save_path}')
+    if video_path is not None and not skip_video_composite:
+        video_save_path = output_dir / f'{video_path.stem}.mp4'
+        make_video(video_info, video_save_path, audio, sampling_rate=seq_cfg.sampling_rate)
+        log.info(f'Video saved to {output_dir / video_save_path}')
+
+    log.info('Memory usage: %.2f GB', torch.cuda.max_memory_allocated() / (2**30))
+
+
+if __name__ == '__main__':
+    main()

docs/EVAL.md

@@ -0,0 +1,22 @@
+# Evaluation
+
+## Batch Evaluation
+
+To evaluate the model on a dataset, use the `batch_eval.py` script. It is significantly more efficient in large-scale evaluation compared to `demo.py`, supporting batched inference, multi-GPU inference, torch compilation, and skipping video compositions.
+
+An example of running this script with four GPUs is as follows:
+
+```bash
+OMP_NUM_THREADS=4 torchrun --standalone --nproc_per_node=4  batch_eval.py duration_s=8 dataset=vggsound model=small_16k num_workers=8
+```
+
+You may need to update the data paths in `config/eval_data/base.yaml`. 
+More configuration options can be found in `config/base_config.yaml` and `config/eval_config.yaml`.
+
+## Precomputed Results
+
+Precomputed results for VGGSound, AudioCaps, and MovieGen are available here: https://huggingface.co/datasets/hkchengrex/MMAudio-precomputed-results
+
+## Obtaining Quantitative Metrics
+
+Our evaluation code is available here: https://github.com/hkchengrex/av-benchmark

docs/MODELS.md

@@ -0,0 +1,50 @@
+# Pretrained models
+
+The models will be downloaded automatically when you run the demo script. MD5 checksums are provided in `mmaudio/utils/download_utils.py`.
+The models are also available at https://huggingface.co/hkchengrex/MMAudio/tree/main
+
+| Model    | Download link | File size |
+| -------- | ------- | ------- |
+| Flow prediction network, small 16kHz | <a href="https://huggingface.co/hkchengrex/MMAudio/resolve/main/weights/mmaudio_small_16k.pth" download="mmaudio_small_16k.pth">mmaudio_small_16k.pth</a> | 601M |
+| Flow prediction network, small 44.1kHz | <a href="https://huggingface.co/hkchengrex/MMAudio/resolve/main/weights/mmaudio_small_44k.pth" download="mmaudio_small_44k.pth">mmaudio_small_44k.pth</a> | 601M |
+| Flow prediction network, medium 44.1kHz | <a href="https://huggingface.co/hkchengrex/MMAudio/resolve/main/weights/mmaudio_medium_44k.pth" download="mmaudio_medium_44k.pth">mmaudio_medium_44k.pth</a> | 2.4G |
+| Flow prediction network, large 44.1kHz | <a href="https://huggingface.co/hkchengrex/MMAudio/resolve/main/weights/mmaudio_large_44k.pth" download="mmaudio_large_44k.pth">mmaudio_large_44k.pth</a> | 3.9G |
+| Flow prediction network, large 44.1kHz, v2 **(recommended)** | <a href="https://huggingface.co/hkchengrex/MMAudio/resolve/main/weights/mmaudio_large_44k_v2.pth" download="mmaudio_large_44k_v2.pth">mmaudio_large_44k_v2.pth</a> | 3.9G |
+| 16kHz VAE | <a href="https://github.com/hkchengrex/MMAudio/releases/download/v0.1/v1-16.pth">v1-16.pth</a> | 655M |
+| 16kHz BigVGAN vocoder (from Make-An-Audio 2) |<a href="https://github.com/hkchengrex/MMAudio/releases/download/v0.1/best_netG.pt">best_netG.pt</a> | 429M |
+| 44.1kHz VAE |<a href="https://github.com/hkchengrex/MMAudio/releases/download/v0.1/v1-44.pth">v1-44.pth</a> | 1.2G | 
+| Synchformer visual encoder |<a href="https://github.com/hkchengrex/MMAudio/releases/download/v0.1/synchformer_state_dict.pth">synchformer_state_dict.pth</a> | 907M |
+
+To run the model, you need four components: a flow prediction network, visual feature extractors (Synchformer and CLIP, CLIP will be downloaded automatically), a VAE, and a vocoder. VAEs and vocoders are specific to the sampling rate (16kHz or 44.1kHz) and not model sizes.
+The 44.1kHz vocoder will be downloaded automatically.
+The `_v2` model performs worse in benchmarking (e.g., in  Fréchet distance), but, in my experience, generalizes better to new data.
+
+The expected directory structure (full):
+
+```bash
+MMAudio
+├── ext_weights
+│   ├── best_netG.pt
+│   ├── synchformer_state_dict.pth
+│   ├── v1-16.pth
+│   └── v1-44.pth
+├── weights
+│   ├── mmaudio_small_16k.pth
+│   ├── mmaudio_small_44k.pth
+│   ├── mmaudio_medium_44k.pth
+│   ├── mmaudio_large_44k.pth
+│   └── mmaudio_large_44k_v2.pth
+└── ...
+```
+
+The expected directory structure (minimal, for the recommended model only):
+
+```bash
+MMAudio
+├── ext_weights
+│   ├── synchformer_state_dict.pth
+│   └── v1-44.pth
+├── weights
+│   └── mmaudio_large_44k_v2.pth
+└── ...
+```

docs/TRAINING.md

@@ -0,0 +1,160 @@
+# Training
+
+## Overview
+
+We have put a large emphasis on making training as fast as possible.
+Consequently, some pre-processing steps are required.
+
+Namely, before starting any training, we
+
+1. Obtain training data as videos, audios, and captions.
+2. Encode training audios into spectrograms and then with VAE into mean/std
+3. Extract CLIP and synchronization features from videos
+4. Extract CLIP features from text (captions)
+5. Encode all extracted features into [MemoryMappedTensors](https://pytorch.org/tensordict/main/reference/generated/tensordict.MemoryMappedTensor.html) with [TensorDict](https://pytorch.org/tensordict/main/reference/tensordict.html)
+
+**NOTE:** for maximum training speed (e.g., when training the base model with 2*H100s), you would need around 3~5 GB/s of random read speed. Spinning disks would not be able to catch up and most consumer-grade SSDs would struggle. In my experience, the best bet is to have a large enough system memory such that the OS can cache the data. This way, the data is read from RAM instead of disk.
+
+The current training script does not support `_v2` training.
+
+## Recommended Hardware Configuration
+
+These are what I recommend for a smooth and efficient training experience. These are not minimum requirements.
+
+- Single-node machine. We did not implement multi-node training
+- GPUs: for the small model, two 80G-H100s or above; for the large model, eight 80G-H100s or above
+- System memory: for 16kHz training, 600GB+; for 44kHz training, 700GB+
+- Storage: >2TB of fast NVMe storage. If you have enough system memory, OS caching will help and the storage does not need to be as fast.
+
+## Prerequisites
+
+1. Install [av-benchmark](https://github.com/hkchengrex/av-benchmark). We use this library to automatically evaluate on the validation set during training, and on the test set after training.
+2. Extract features for evaluation using [av-benchmark](https://github.com/hkchengrex/av-benchmark) for the validation and test set as a [validation cache](https://github.com/hkchengrex/MMAudio/blob/34bf089fdd2e457cd5ef33be96c0e1c8a0412476/config/data/base.yaml#L38) and a [test cache](https://github.com/hkchengrex/MMAudio/blob/34bf089fdd2e457cd5ef33be96c0e1c8a0412476/config/data/base.yaml#L31). You can also download the precomputed evaluation cache [here](https://huggingface.co/datasets/hkchengrex/MMAudio-precomputed-results/tree/main).
+
+3. You will need ffmpeg to extract frames from videos. Note that `torchaudio` imposes a maximum version limit (`ffmpeg<7`). You can install it as follows:
+
+```bash
+conda install -c conda-forge 'ffmpeg<7'
+```
+
+4. Download the training datasets. We used [VGGSound](https://arxiv.org/abs/2004.14368), [AudioCaps](https://audiocaps.github.io/), and [WavCaps](https://arxiv.org/abs/2303.17395). Note that the audio files in the huggingface release of WavCaps have been downsampled to 32kHz. To the best of our ability, we located the original (high-sampling rate) audio files and used them instead to prevent artifacts during 44.1kHz training. We did not use the "SoundBible" portion of WavCaps, since it is a small set with many short audio unsuitable for our training.
+
+5. Download the corresponding VAE (`v1-16.pth` for 16kHz training, and `v1-44.pth` for 44.1kHz training), vocoder models (`best_netG.pt` for 16kHz training; the vocoder for 44.1kHz training will be downloaded automatically), the [empty string encoding](https://github.com/hkchengrex/MMAudio/releases/download/v0.1/empty_string.pth), and Synchformer weights from [MODELS.md](https://github.com/hkchengrex/MMAudio/blob/main/docs/MODELS.md) place them in `ext_weights/`.
+
+## Preparing Audio-Video-Text Features
+
+We have prepared some example data in `training/example_videos`.
+`training/extract_video_training_latents.py` extracts audio, video, and text features and save them as a `TensorDict` with a `.tsv` file containing metadata to `output_dir`.
+
+To run this script, use the `torchrun` utility:
+
+```bash
+torchrun --standalone training/extract_video_training_latents.py
+```
+
+You can run this script with multiple GPUs (with `--nproc_per_node=<n>` after `--standalone` and before the script name) to speed up extraction.
+Modify the definitions near the top of the script to switch between 16kHz/44.1kHz extraction.
+Change the data path definitions in `data_cfg` if necessary.
+
+Arguments:
+
+- `latent_dir` -- where intermediate latent outputs are saved. It is safe to delete this directory afterwards.
+- `output_dir` -- where TensorDict and the metadata file are saved.
+
+Outputs produced in `output_dir`:
+
+1. A directory named `vgg-{split}` (i.e., in the TensorDict format), containing
+    a. `mean.memmap` mean values predicted by the VAE encoder (number of videos X sequence length X channel size)
+    b. `std.memmap` standard deviation values predicted by the VAE encoder (number of videos X sequence length X channel size)
+    c. `text_features.memmap` text features extracted from CLIP (number of videos X 77 (sequence length) X 1024)
+    d. `clip_features.memmap` clip features extracted from CLIP (number of videos X 64 (8 fps) X 1024)
+    e. `sync_features.memmap` synchronization features extracted from Synchformer (number of videos X 192 (24 fps) X 768)
+    f. `meta.json` that contains the metadata for the above memory mappings
+2. A tab-separated values file named `vgg-{split}.tsv` that contains two columns: `id` containing video file names without extension, and `label` containing corresponding text labels (i.e., captions)
+
+## Preparing Audio-Text Features
+
+We have prepared some example data in `training/example_audios`.
+
+1. Run `training/partition_clips` to partition each audio file into clips (by finding start and end points; we do not save the partitioned audio onto the disk to save disk space)
+2. Run `training/extract_audio_training_latents.py` to extract each clip's audio and text features and save them as a `TensorDict` with a `.tsv` file containing metadata to `output_dir`.
+
+### Partitioning the audio files
+
+Run
+
+```bash
+python training/partition_clips.py
+```
+
+Arguments:
+
+- `data_dir` -- path to a directory containing the audio files (`.flac` or `.wav`)
+- `output_dir` -- path to the output `.csv` file
+- `start` -- optional; useful when you need to run multiple processes to speed up processing -- this defines the beginning of the chunk to be processed
+- `end` -- optional; useful when you need to run multiple processes to speed up processing -- this defines the end of the chunk to be processed
+
+### Extracting audio and text features
+
+Run
+
+```bash
+torchrun --standalone training/extract_audio_training_latents.py
+```
+
+You can run this with multiple GPUs (with `--nproc_per_node=<n>`) to speed up extraction.
+Modify the definitions near the top of the script to switch between 16kHz/44.1kHz extraction.
+
+Arguments:
+
+- `data_dir` -- path to a directory containing the audio files (`.flac` or `.wav`), same as the previous step
+- `captions_tsv` -- path to the captions file, a tab-separated values (tsv) file at least with columns `id` and `caption`
+- `clips_tsv` -- path to the clips file, generated in the last step
+- `latent_dir` -- where intermediate latent outputs are saved. It is safe to delete this directory afterwards.
+- `output_dir` -- where TensorDict and the metadata file are saved.
+
+**Reference tsv files (with overlaps removed as mentioned in the paper) can be found [here](https://github.com/hkchengrex/MMAudio/releases/tag/v0.1).**
+Note that these reference tsv files are the **outputs** of `extract_audio_training_latents.py`, which means the `id` column might contain duplicate entries (one per clip). You can still use it as the `captions_tsv` input though -- the script will handle duplicates gracefully.
+Among these reference tsv files, `audioset_sl.tsv`, `bbcsound.tsv`, and `freesound.tsv` are subsets that are parts of WavCaps. These subsets might be smaller than the original datasets.
+The Clotho data contains both the development set and the validation set.
+
+Outputs produced in `output_dir`:
+
+1. A directory named `{basename(output_dir)}` (i.e., in the TensorDict format), containing
+    a. `mean.memmap` mean values predicted by the VAE encoder (number of audios X sequence length X channel size)
+    b. `std.memmap` standard deviation values predicted by the VAE encoder (number of audios X sequence length X channel size)
+    c. `text_features.memmap` text features extracted from CLIP (number of audios X 77 (sequence length) X 1024)
+    f. `meta.json` that contains the metadata for the above memory mappings
+2. A tab-separated values file named `{basename(output_dir)}.tsv` that contains two columns: `id` containing audio file names without extension, and `label` containing corresponding text labels (i.e., captions)
+
+## Training on Extracted Features
+
+We use Distributed Data Parallel (DDP) for training.
+First, specify the data path in `config/data/base.yaml`. If you used the default parameters in the scripts above to extract features for the example data, the `Example_video` and `Example_audio` items should already be correct.
+
+To run training on the example data, use the following command:
+
+```bash
+OMP_NUM_THREADS=4 torchrun --standalone --nproc_per_node=1 train.py exp_id=debug compile=False  debug=True example_train=True  batch_size=1
+```
+
+This will not train a useful model, but it will check if everything is set up correctly.
+
+For full training on the base model with two GPUs, use the following command:
+
+```bash
+OMP_NUM_THREADS=4 torchrun --standalone --nproc_per_node=2 train.py exp_id=exp_1 model=small_16k
+```
+
+Any outputs from training will be stored in `output/<exp_id>`.
+
+More configuration options can be found in `config/base_config.yaml` and `config/train_config.yaml`.
+For the medium and large models, specify `vgg_oversample_rate` to be `3` to reduce overfitting.
+
+## Checkpoints
+
+Model checkpoints, including optimizer states and the latest EMA weights, are available here: https://huggingface.co/hkchengrex/MMAudio
+
+---
+
+Godspeed!

docs/index.html

@@ -0,0 +1,149 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <!-- Google tag (gtag.js) -->
+    <script async src="https://www.googletagmanager.com/gtag/js?id=G-0JKBJ3WRJZ"></script>
+    <script>
+    window.dataLayer = window.dataLayer || [];
+    function gtag(){dataLayer.push(arguments);}
+    gtag('js', new Date());
+    gtag('config', 'G-0JKBJ3WRJZ');
+    </script>
+
+    <link rel="preconnect" href="https://fonts.googleapis.com">
+    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+    <link href="https://fonts.googleapis.com/css2?family=Source+Sans+3&display=swap" rel="stylesheet">
+    <meta charset="UTF-8">
+    <title>MMAudio</title>
+
+    <link rel="icon" type="image/png" href="images/icon.png">
+
+    <meta name="viewport" content="width=device-width, initial-scale=1">
+    <!-- CSS only -->
+    <link href="https://cdn.jsdelivr.net/npm/[email protected]/dist/css/bootstrap.min.css" rel="stylesheet"
+        integrity="sha384-+0n0xVW2eSR5OomGNYDnhzAbDsOXxcvSN1TPprVMTNDbiYZCxYbOOl7+AMvyTG2x" crossorigin="anonymous">
+    <script src="https://ajax.googleapis.com/ajax/libs/jquery/3.5.1/jquery.min.js"></script>
+
+    <link rel="stylesheet" href="style.css">
+</head>
+<body>
+
+    <body>
+        <br><br><br><br>
+        <div class="container">
+            <div class="row text-center" style="font-size:38px">
+                <div class="col strong">
+                    Taming Multimodal Joint Training for High-Quality <br>Video-to-Audio Synthesis
+                </div>
+            </div>
+    
+            <br>
+            <div class="row text-center" style="font-size:28px">
+                <div class="col">
+                    CVPR 2025
+                </div>
+            </div>
+            <br>
+    
+            <div class="h-100 row text-center heavy justify-content-md-center" style="font-size:22px;">
+                <div class="col-sm-auto px-lg-2">
+                    <a href="https://hkchengrex.github.io/">Ho Kei Cheng<sup>1</sup></a>
+                </div>
+                <div class="col-sm-auto px-lg-2">
+                    <nobr><a href="https://scholar.google.co.jp/citations?user=RRIO1CcAAAAJ">Masato Ishii<sup>2</sup></a></nobr>
+                </div>
+                <div class="col-sm-auto px-lg-2">
+                    <nobr><a href="https://scholar.google.com/citations?user=sXAjHFIAAAAJ">Akio Hayakawa<sup>2</sup></a></nobr>
+                </div>
+                <div class="col-sm-auto px-lg-2">
+                    <nobr><a href="https://scholar.google.com/citations?user=XCRO260AAAAJ">Takashi Shibuya<sup>2</sup></a></nobr>
+                </div>
+                <div class="col-sm-auto px-lg-2">
+                    <nobr><a href="https://www.alexander-schwing.de/">Alexander Schwing<sup>1</sup></a></nobr>
+                </div>
+                <div class="col-sm-auto px-lg-2" >
+                    <nobr><a href="https://www.yukimitsufuji.com/">Yuki Mitsufuji<sup>2,3</sup></a></nobr>
+                </div>
+            </div>
+
+            <div class="h-100 row text-center heavy justify-content-md-center" style="font-size:22px;">
+                <div class="col-sm-auto px-lg-2">
+                    <sup>1</sup>University of Illinois Urbana-Champaign
+                </div>
+                <div class="col-sm-auto px-lg-2">
+                    <sup>2</sup>Sony AI
+                </div>
+                <div class="col-sm-auto px-lg-2">
+                    <sup>3</sup>Sony Group Corporation
+                </div>
+            </div>
+    
+            <br>
+    
+            <br>
+    
+            <div class="h-100 row text-center justify-content-md-center" style="font-size:20px;">
+                <div class="col-sm-2">
+                    <a href="https://arxiv.org/abs/2412.15322">[Paper]</a>
+                </div>
+                <div class="col-sm-2">
+                    <a href="https://github.com/hkchengrex/MMAudio">[Code]</a>
+                </div>
+                <div class="col-sm-3">
+                    <a href="https://huggingface.co/spaces/hkchengrex/MMAudio">[Huggingface Demo]</a>
+                </div>
+                <div class="col-sm-2">
+                    <a href="https://colab.research.google.com/drive/1TAaXCY2-kPk4xE4PwKB3EqFbSnkUuzZ8?usp=sharing">[Colab Demo]</a>
+                </div>
+                <div class="col-sm-3">
+                    <a href="https://replicate.com/zsxkib/mmaudio">[Replicate Demo]</a>
+                </div>
+            </div>
+    
+            <br>
+    
+            <hr>
+    
+            <div class="row" style="font-size:32px">
+                <div class="col strong">
+                    TL;DR
+                </div>
+            </div>
+            <br>
+            <div class="row">
+                <div class="col">
+                    <p class="light" style="text-align: left;">
+                        MMAudio generates synchronized audio given video and/or text inputs.
+                    </p>
+                </div>
+            </div>
+    
+            <br>
+            <hr>
+            <br>
+    
+            <div class="row" style="font-size:32px">
+                <div class="col strong">
+                    Demo
+                </div>
+            </div>
+            <br>
+            <div class="row" style="font-size:48px">
+                <div class="col strong text-center">
+                    <a href="video_main.html" style="text-decoration: underline;">&lt;More results&gt;</a>
+                </div>
+            </div>
+            <br>
+            <div class="video-container" style="text-align: center;">
+                <iframe src="https://youtube.com/embed/YElewUT2M4M"></iframe>
+                </div>
+
+            <br>
+    
+            <br><br>
+            <br><br>
+    
+        </div>
+
+</body>
+</html>

docs/style.css

@@ -0,0 +1,78 @@
+body {
+    font-family: 'Source Sans 3', sans-serif;
+    font-size: 18px;
+    margin-left: auto;
+    margin-right: auto;
+    font-weight: 400;
+    height: 100%;
+    max-width: 1000px;
+}
+
+table {
+    width: 100%;
+    border-collapse: collapse;
+}
+th, td {
+    border: 1px solid #ddd;
+    padding: 8px;
+    text-align: center;
+}
+th {
+    background-color: #f2f2f2;
+}
+video {
+    width: 100%;
+    height: auto;
+}
+p {
+    font-size: 28px;
+}
+h2 {
+    font-size: 36px;
+}
+
+.strong {
+    font-weight: 700;
+}
+
+.light {
+    font-weight: 100;
+}
+
+.heavy {
+    font-weight: 900;
+}
+
+.column {
+    float: left;
+}
+
+a:link,
+a:visited {
+    color: #05538f;
+    text-decoration: none;
+}
+
+a:hover {
+    color: #63cbdd;
+}
+
+hr {
+    border: 0;
+    height: 1px;
+    background-image: linear-gradient(to right, rgba(0, 0, 0, 0), rgba(0, 0, 0, 0.75), rgba(0, 0, 0, 0));
+}
+
+.video-container {
+    position: relative;
+    padding-bottom: 56.25%; /* 16:9 */
+    height: 0;
+  }
+  
+.video-container iframe {
+    position: absolute;
+    top: 0;
+    left: 0;
+    width: 100%;
+    height: 100%;
+}

docs/style_videos.css

@@ -0,0 +1,52 @@
+body {
+    font-family: 'Source Sans 3', sans-serif;
+    font-size: 1.5vh;
+    font-weight: 400;
+}
+
+table {
+    width: 100%;
+    border-collapse: collapse;
+}
+th, td {
+    border: 1px solid #ddd;
+    padding: 8px;
+    text-align: center;
+}
+th {
+    background-color: #f2f2f2;
+}
+video {
+    width: 100%;
+    height: auto;
+}
+p {
+    font-size: 1.5vh;
+    font-weight: bold;
+}
+h2 {
+    font-size: 2vh;
+    font-weight: bold;
+}
+
+.video-container {
+    position: relative;
+    padding-bottom: 56.25%; /* 16:9 */
+    height: 0;
+  }
+  
+.video-container iframe {
+    position: absolute;
+    top: 0;
+    left: 0;
+    width: 100%;
+    height: 100%;
+}
+
+.video-header {
+    background-color: #f2f2f2;
+    text-align: center;
+    font-size: 1.5vh;
+    font-weight: bold;
+    padding: 8px;
+}

docs/video_gen.html

@@ -0,0 +1,254 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <!-- Google tag (gtag.js) -->
+    <script async src="https://www.googletagmanager.com/gtag/js?id=G-0JKBJ3WRJZ"></script>
+    <script>
+    window.dataLayer = window.dataLayer || [];
+    function gtag(){dataLayer.push(arguments);}
+    gtag('js', new Date());
+    gtag('config', 'G-0JKBJ3WRJZ');
+    </script>
+
+    <link href='https://fonts.googleapis.com/css?family=Source+Sans+Pro' rel='stylesheet' type='text/css'>
+    <meta charset="UTF-8">
+    <title>MMAudio</title>
+
+    <link rel="icon" type="image/png" href="images/icon.png">
+
+    <meta name="viewport" content="width=device-width, initial-scale=1">
+    <!-- CSS only -->
+    <link href="https://cdn.jsdelivr.net/npm/[email protected]/dist/css/bootstrap.min.css" rel="stylesheet"
+        integrity="sha384-+0n0xVW2eSR5OomGNYDnhzAbDsOXxcvSN1TPprVMTNDbiYZCxYbOOl7+AMvyTG2x" crossorigin="anonymous">
+    <script src="https://ajax.googleapis.com/ajax/libs/jquery/3.7.1/jquery.min.js"></script>
+
+    <link rel="stylesheet" href="style_videos.css">
+</head>
+<body>
+
+    <div id="moviegen_all">
+    <h2 id="moviegen" style="text-align: center;">Comparisons with Movie Gen Audio on Videos Generated by MovieGen</h2>
+    <p id="moviegen1" style="overflow: hidden;">
+        Example 1: Ice cracking with sharp snapping sound, and metal tool scraping against the ice surface. 
+        <span style="float: right;"><a href="#index">Back to index</a></span>
+    </p> 
+
+    <div class="row g-1">
+        <div class="col-sm-6">
+            <div class="video-header">Movie Gen Audio</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/d7Lb0ihtGcE"></iframe>
+            </div> 
+        </div>
+        <div class="col-sm-6">
+            <div class="video-header">Ours</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/F4JoJ2r2m8U"></iframe>
+                </div> 
+        </div>
+    </div>
+    <br>
+
+    <!-- <p id="moviegen2">Example 2: Rhythmic splashing and lapping of water. <span style="float:right;"><a href="#index">Back to index</a></span> </p> 
+
+    <table>
+        <thead>
+            <tr>
+                <th>Movie Gen Audio</th>
+                <th>Ours</th>
+            </tr>
+        </thead>
+        <tbody>
+            <tr>
+                <td width="50%">
+                    <div class="video-container">
+                    <iframe src="https://youtube.com/embed/5gQNPK99CIk"></iframe>
+                    </div>
+                </td>
+                <td width="50%">
+                    <div class="video-container">
+                    <iframe src="https://youtube.com/embed/AbwnTzG-BpA"></iframe>
+                    </div>
+                </td>
+            </tr>
+        </tbody>
+    </table> -->
+
+    <p id="moviegen2" style="overflow: hidden;">
+        Example 2: Rhythmic splashing and lapping of water. 
+        <span style="float:right;"><a href="#index">Back to index</a></span> 
+    </p> 
+    <div class="row g-1">
+        <div class="col-sm-6">
+            <div class="video-header">Movie Gen Audio</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/5gQNPK99CIk"></iframe>
+            </div> 
+        </div>
+        <div class="col-sm-6">
+            <div class="video-header">Ours</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/AbwnTzG-BpA"></iframe>
+                </div> 
+        </div>
+    </div>
+    <br>
+
+    <p id="moviegen3" style="overflow: hidden;">
+        Example 3: Shovel scrapes against dry earth. 
+        <span style="float:right;"><a href="#index">Back to index</a></span> 
+    </p> 
+    <div class="row g-1">
+        <div class="col-sm-6">
+            <div class="video-header">Movie Gen Audio</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/PUKGyEve7XQ"></iframe>
+            </div> 
+        </div>
+        <div class="col-sm-6">
+            <div class="video-header">Ours</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/CNn7i8VNkdc"></iframe>
+            </div> 
+        </div>
+    </div>
+    <br>
+    
+
+    <p id="moviegen4" style="overflow: hidden;">
+        (Failure case) Example 4: Creamy sound of mashed potatoes being scooped. 
+        <span style="float:right;"><a href="#index">Back to index</a></span> 
+    </p> 
+    <div class="row g-1">
+        <div class="col-sm-6">
+            <div class="video-header">Movie Gen Audio</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/PJv1zxR9JjQ"></iframe>
+            </div> 
+        </div>
+        <div class="col-sm-6">
+            <div class="video-header">Ours</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/c3-LJ1lNsPQ"></iframe>
+            </div> 
+        </div>
+    </div>
+    <br>
+
+    </div>
+
+    <div id="hunyuan_sora_all">
+
+    <h2 id="hunyuan" style="text-align: center;">Results on Videos Generated by Hunyuan</h2>
+    <p style="overflow: hidden;">
+        <span style="float:right;"><a href="#index">Back to index</a></span> 
+    </p> 
+    <div class="row g-1">
+        <div class="col-sm-6">
+            <div class="video-header">Typing</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/8ln_9hhH_nk"></iframe>
+            </div> 
+        </div>
+        <div class="col-sm-6">
+            <div class="video-header">Water is rushing down a stream and pouring</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/5df1FZFQj30"></iframe>
+            </div> 
+        </div>
+    </div>
+    <div class="row g-1">
+        <div class="col-sm-6">
+            <div class="video-header">Waves on beach</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/7wQ9D5WgpFc"></iframe>
+            </div> 
+        </div>
+        <div class="col-sm-6">
+            <div class="video-header">Water droplet</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/q7M2nsalGjM"></iframe>
+            </div> 
+        </div>
+    </div>
+    <br>
+
+    <h2 id="sora" style="text-align: center;">Results on Videos Generated by Sora</h2>
+    <p style="overflow: hidden;">
+        <span style="float:right;"><a href="#index">Back to index</a></span> 
+    </p> 
+    <div class="row g-1">
+        <div class="col-sm-6">
+            <div class="video-header">Ships riding waves</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/JbgQzHHytk8"></iframe>
+            </div> 
+        </div>
+        <div class="col-sm-6">
+            <div class="video-header">Train (no text prompt given)</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/xOW7zrjpWC8"></iframe>
+            </div> 
+        </div>
+    </div>
+    <div class="row g-1">
+        <div class="col-sm-6">
+            <div class="video-header">Seashore (no text prompt given)</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/fIuw5Y8ZZ9E"></iframe>
+            </div> 
+        </div>
+        <div class="col-sm-6">
+            <div class="video-header">Surfing (failure: unprompted music)</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/UcSTk-v0M_s"></iframe>
+            </div> 
+        </div>
+    </div>
+    <br>
+
+    <div id="mochi_ltx_all">
+    <h2 id="mochi" style="text-align: center;">Results on Videos Generated by Mochi 1</h2>
+    <p style="overflow: hidden;">
+        <span style="float:right;"><a href="#index">Back to index</a></span> 
+    </p> 
+    <div class="row g-1">
+        <div class="col-sm-6">
+            <div class="video-header">Magical fire and lightning (no text prompt given)</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/tTlRZaSMNwY"></iframe>
+            </div> 
+        </div>
+        <div class="col-sm-6">
+            <div class="video-header">Storm (no text prompt given)</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/4hrZTMJUy3w"></iframe>
+            </div> 
+        </div>
+    </div>
+    <br>
+
+    <h2 id="ltx" style="text-align: center;">Results on Videos Generated by LTX-Video</h2>
+    <p style="overflow: hidden;">
+        <span style="float:right;"><a href="#index">Back to index</a></span> 
+    </p> 
+    <div class="row g-1">
+        <div class="col-sm-6">
+            <div class="video-header">Firewood burning and cracking</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/P7_DDpgev0g"></iframe>
+            </div> 
+        </div>
+        <div class="col-sm-6">
+            <div class="video-header">Waterfall, water splashing</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/4MvjceYnIO0"></iframe>
+            </div> 
+        </div>
+    </div>
+    <br>
+
+    </div>
+
+</body>
+</html>

docs/video_main.html

@@ -0,0 +1,98 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <!-- Google tag (gtag.js) -->
+    <script async src="https://www.googletagmanager.com/gtag/js?id=G-0JKBJ3WRJZ"></script>
+    <script>
+    window.dataLayer = window.dataLayer || [];
+    function gtag(){dataLayer.push(arguments);}
+    gtag('js', new Date());
+    gtag('config', 'G-0JKBJ3WRJZ');
+    </script>
+
+    <link href='https://fonts.googleapis.com/css?family=Source+Sans+Pro' rel='stylesheet' type='text/css'>
+    <meta charset="UTF-8">
+    <title>MMAudio</title>
+
+    <link rel="icon" type="image/png" href="images/icon.png">
+
+    <meta name="viewport" content="width=device-width, initial-scale=1, maximum-scale=1, user-scalable=no">
+    <!-- CSS only -->
+    <link href="https://cdn.jsdelivr.net/npm/[email protected]/dist/css/bootstrap.min.css" rel="stylesheet"
+        integrity="sha384-+0n0xVW2eSR5OomGNYDnhzAbDsOXxcvSN1TPprVMTNDbiYZCxYbOOl7+AMvyTG2x" crossorigin="anonymous">
+    <script src="https://ajax.googleapis.com/ajax/libs/jquery/3.7.1/jquery.min.js"></script>
+
+    <link rel="stylesheet" href="style_videos.css">
+
+    <script type="text/javascript">
+        $(document).ready(function(){
+            $("#content").load("video_gen.html #moviegen_all");
+            $("#load_moveigen").click(function(){
+                $("#content").load("video_gen.html #moviegen_all");
+            });
+            $("#load_hunyuan_sora").click(function(){
+                $("#content").load("video_gen.html #hunyuan_sora_all");
+            });
+            $("#load_mochi_ltx").click(function(){
+                $("#content").load("video_gen.html #mochi_ltx_all");
+            });
+            $("#load_vgg1").click(function(){
+                $("#content").load("video_vgg.html #vgg1");
+            });
+            $("#load_vgg2").click(function(){
+                $("#content").load("video_vgg.html #vgg2");
+            });
+            $("#load_vgg3").click(function(){
+                $("#content").load("video_vgg.html #vgg3");
+            });
+            $("#load_vgg4").click(function(){
+                $("#content").load("video_vgg.html #vgg4");
+            });
+            $("#load_vgg5").click(function(){
+                $("#content").load("video_vgg.html #vgg5");
+            });
+            $("#load_vgg6").click(function(){
+                $("#content").load("video_vgg.html #vgg6");
+            });
+            $("#load_vgg_extra").click(function(){
+                $("#content").load("video_vgg.html #vgg_extra");
+            });
+        });
+    </script>
+</head>
+<body>
+    <h1 id="index" style="text-align: center;">Index</h1>
+    <p><b>(Click on the links to load the corresponding videos)</b> <span style="float:right;"><a href="index.html">Back to project page</a></span></p>
+
+    <ol>
+        <li>
+            <a href="#" id="load_moveigen">Comparisons with Movie Gen Audio on Videos Generated by MovieGen</a>
+        </li>
+        <li>
+            <a href="#" id="load_hunyuan_sora">Results on Videos Generated by Hunyuan and Sora</a>
+        </li>
+        <li>
+            <a href="#" id="load_mochi_ltx">Results on Videos Generated by Mochi 1 and LTX-Video</a>
+        </li>
+        <li>
+            On VGGSound
+            <ol>
+                <li><a id='load_vgg1' href="#">Example 1: Wolf howling</a></li>
+                <li><a id='load_vgg2' href="#">Example 2: Striking a golf ball</a></li>
+                <li><a id='load_vgg3' href="#">Example 3: Hitting a drum</a></li>
+                <li><a id='load_vgg4' href="#">Example 4: Dog barking</a></li>
+                <li><a id='load_vgg5' href="#">Example 5: Playing a string instrument</a></li>
+                <li><a id='load_vgg6' href="#">Example 6: A group of people playing tambourines</a></li>
+                <li><a id='load_vgg_extra' href="#">Extra results & failure cases</a></li>
+            </ol>
+        </li>
+    </ol>
+
+    <div id="content" class="container-fluid">
+
+    </div>
+    <br>
+    <br>
+
+</body>
+</html>

docs/video_vgg.html

@@ -0,0 +1,452 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <!-- Google tag (gtag.js) -->
+    <script async src="https://www.googletagmanager.com/gtag/js?id=G-0JKBJ3WRJZ"></script>
+    <script>
+    window.dataLayer = window.dataLayer || [];
+    function gtag(){dataLayer.push(arguments);}
+    gtag('js', new Date());
+    gtag('config', 'G-0JKBJ3WRJZ');
+    </script>
+
+    <link href='https://fonts.googleapis.com/css?family=Source+Sans+Pro' rel='stylesheet' type='text/css'>
+    <meta charset="UTF-8">
+    <title>MMAudio</title>
+
+    <meta name="viewport" content="width=device-width, initial-scale=1">
+    <!-- CSS only -->
+    <link href="https://cdn.jsdelivr.net/npm/[email protected]/dist/css/bootstrap.min.css" rel="stylesheet"
+        integrity="sha384-+0n0xVW2eSR5OomGNYDnhzAbDsOXxcvSN1TPprVMTNDbiYZCxYbOOl7+AMvyTG2x" crossorigin="anonymous">
+    <script src="https://ajax.googleapis.com/ajax/libs/jquery/3.5.1/jquery.min.js"></script>
+
+    <link rel="stylesheet" href="style_videos.css">
+</head>
+<body>
+
+    <div id="vgg1">
+    <h2 style="text-align: center;">Comparisons with state-of-the-art methods in VGGSound</h2>
+    <p style="overflow: hidden;">
+        Example 1: Wolf howling. 
+        <span style="float:right;"><a href="#index">Back to index</a></span> 
+    </p> 
+        <div class="row g-1">
+            <div class="col-sm-3">
+                <div class="video-header">Ground-truth</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/9J_V74gqMUA"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">Ours</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/P6O8IpjErPc"></iframe>
+                    </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">V2A-Mapper</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/w-5eyqepvTk"></iframe>
+                    </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">FoleyCrafter</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/VOLfoZlRkzo"></iframe>
+                    </div> 
+            </div>
+        </div>
+        <div class="row g-1">
+            <div class="col-sm-3">
+                <div class="video-header">Frieren</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/49owKyA5Pa8"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">VATT</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/QVtrFgbeGDM"></iframe>
+                    </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">V-AURA</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/8r0uEfSNjvI"></iframe>
+                    </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">Seeing and Hearing</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/bn-sLg2qulk"></iframe>
+                    </div> 
+            </div>
+        </div>
+    </div>
+
+    <div id="vgg2">
+        <h2 style="text-align: center;">Comparisons with state-of-the-art methods in VGGSound</h2>
+        <p style="overflow: hidden;">
+            Example 2: Striking a golf ball.
+            <span style="float:right;"><a href="#index">Back to index</a></span>
+        </p>
+
+        <div class="row g-1">
+            <div class="col-sm-3">
+                <div class="video-header">Ground-truth</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/1hwSu42kkho"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">Ours</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/kZibDoDCNxI"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">V2A-Mapper</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/jgKfLBLhh7Y"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">FoleyCrafter</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/Lfsx8mOPcJo"></iframe>
+                </div>
+            </div>
+        </div>
+        <div class="row g-1">
+            <div class="col-sm-3">
+                <div class="video-header">Frieren</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/tz-LpbB0MBc"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">VATT</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/RTDUHMi08n4"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">V-AURA</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/N-3TDOsPnZQ"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">Seeing and Hearing</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/QnsHnLn4gB0"></iframe>
+                </div>
+            </div>
+        </div>
+    </div>
+
+    <div id="vgg3">
+        <h2 style="text-align: center;">Comparisons with state-of-the-art methods in VGGSound</h2>
+        <p style="overflow: hidden;">
+            Example 3: Hitting a drum. 
+            <span style="float:right;"><a href="#index">Back to index</a></span> 
+        </p> 
+
+        <div class="row g-1">
+            <div class="col-sm-3">
+                <div class="video-header">Ground-truth</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/0oeIwq77w0Q"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">Ours</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/-UtPV9ohuIM"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">V2A-Mapper</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/9yivkgN-zwc"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">FoleyCrafter</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/kkCsXPOlBvY"></iframe>
+                </div> 
+            </div>
+        </div>
+        <div class="row g-1">
+            <div class="col-sm-3">
+                <div class="video-header">Frieren</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/MbNKsVsuvig"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">VATT</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/2yYviBjrpBw"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">V-AURA</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/9yivkgN-zwc"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">Seeing and Hearing</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/6dnyQt4Fuhs"></iframe>
+                </div> 
+            </div>
+        </div>
+    </div>
+    </div>
+
+    <div id="vgg4">
+        <h2 style="text-align: center;">Comparisons with state-of-the-art methods in VGGSound</h2>
+        <p style="overflow: hidden;">
+            Example 4: Dog barking. 
+            <span style="float:right;"><a href="#index">Back to index</a></span> 
+        </p> 
+
+        <div class="row g-1">
+            <div class="col-sm-3">
+                <div class="video-header">Ground-truth</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/ckaqvTyMYAw"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">Ours</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/_aRndFZzZ-I"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">V2A-Mapper</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/mNCISP3LBl0"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">FoleyCrafter</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/phZBQ3L7foE"></iframe>
+                </div> 
+            </div>
+        </div>
+        <div class="row g-1">
+            <div class="col-sm-3">
+                <div class="video-header">Frieren</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/Sb5Mg1-ORao"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">VATT</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/eHmAGOmtDDg"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">V-AURA</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/NEGa3krBrm0"></iframe>
+                </div> 
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">Seeing and Hearing</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/aO0EAXlwE7A"></iframe>
+                </div> 
+            </div>
+        </div>
+    </div>
+    
+    <div id="vgg5">
+        <h2 style="text-align: center;">Comparisons with state-of-the-art methods in VGGSound</h2>
+        <p style="overflow: hidden;">
+            Example 5: Playing a string instrument.
+            <span style="float:right;"><a href="#index">Back to index</a></span>
+        </p>
+
+        <div class="row g-1">
+            <div class="col-sm-3">
+                <div class="video-header">Ground-truth</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/KP1QhWauIOc"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">Ours</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/ovaJhWSquYE"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">V2A-Mapper</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/N723FS9lcy8"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">FoleyCrafter</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/t0N4ZAAXo58"></iframe>
+                </div>
+            </div>
+        </div>
+        <div class="row g-1">
+            <div class="col-sm-3">
+                <div class="video-header">Frieren</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/8YSRs03QNNA"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">VATT</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/vOpMz55J1kY"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">V-AURA</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/9JHC75vr9h0"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">Seeing and Hearing</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/9w0JckNzXmY"></iframe>
+                </div>
+            </div>
+        </div>
+    </div>
+    
+    <div id="vgg6">
+        <h2 style="text-align: center;">Comparisons with state-of-the-art methods in VGGSound</h2>
+        <p style="overflow: hidden;">
+            Example 6: A group of people playing tambourines.
+            <span style="float:right;"><a href="#index">Back to index</a></span>
+        </p>
+
+        <div class="row g-1">
+            <div class="col-sm-3">
+                <div class="video-header">Ground-truth</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/mx6JLxzUkRc"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">Ours</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/oLirHhP9Su8"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">V2A-Mapper</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/HkLkHMqptv0"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">FoleyCrafter</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/rpHiiODjmNU"></iframe>
+                </div>
+            </div>
+        </div>
+        <div class="row g-1">
+            <div class="col-sm-3">
+                <div class="video-header">Frieren</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/1mVD3fJ0LpM"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">VATT</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/yjVFnJiEJlw"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">V-AURA</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/neVeMSWtRkU"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-3">
+                <div class="video-header">Seeing and Hearing</div>
+                <div class="video-container">
+                    <iframe src="https://youtube.com/embed/EUE7YwyVWz8"></iframe>
+                </div>
+            </div>
+        </div>
+    </div>
+    
+    <div id="vgg_extra">
+        <h2 style="text-align: center;">Comparisons with state-of-the-art methods in VGGSound</h2>
+        <p style="overflow: hidden;">
+            <span style="float:right;"><a href="#index">Back to index</a></span>
+        </p>
+
+        <div class="row g-1">
+            <div class="col-sm-3">
+            <div class="video-header">Moving train</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/Ta6H45rBzJc"></iframe>
+            </div>
+            </div>
+            <div class="col-sm-3">
+            <div class="video-header">Water splashing</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/hl6AtgHXpb4"></iframe>
+            </div>
+            </div>
+            <div class="col-sm-3">
+            <div class="video-header">Skateboarding</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/n4sCNi_9buI"></iframe>
+            </div>
+            </div>
+            <div class="col-sm-3">
+            <div class="video-header">Synchronized clapping</div>
+            <div class="video-container">
+                <iframe src="https://youtube.com/embed/oxexfpLn7FE"></iframe>
+            </div>
+            </div>
+        </div>
+
+        <br><br>
+    
+        <div id="extra-failure">
+            <h2 style="text-align: center;">Failure cases</h2>
+            <p style="overflow: hidden;">
+            <span style="float:right;"><a href="#index">Back to index</a></span>
+            </p>
+
+            <div class="row g-1">
+            <div class="col-sm-6">
+                <div class="video-header">Human speech</div>
+                <div class="video-container">
+                <iframe src="https://youtube.com/embed/nx0CyrDu70Y"></iframe>
+                </div>
+            </div>
+            <div class="col-sm-6">
+                <div class="video-header">Unfamiliar vision input</div>
+                <div class="video-container">
+                <iframe src="https://youtube.com/embed/hfnAqmK3X7w"></iframe>
+                </div>
+            </div>
+            </div>
+        </div>
+        </div>
+
+</body>
+</html>

gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

gradio_demo.py

@@ -0,0 +1,343 @@
+import gc
+import logging
+from argparse import ArgumentParser
+from datetime import datetime
+from fractions import Fraction
+from pathlib import Path
+
+import gradio as gr
+import torch
+import torchaudio
+
+from mmaudio.eval_utils import (ModelConfig, VideoInfo, all_model_cfg, generate, load_image,
+                                load_video, make_video, setup_eval_logging)
+from mmaudio.model.flow_matching import FlowMatching
+from mmaudio.model.networks import MMAudio, get_my_mmaudio
+from mmaudio.model.sequence_config import SequenceConfig
+from mmaudio.model.utils.features_utils import FeaturesUtils
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+log = logging.getLogger()
+
+device = 'cpu'
+if torch.cuda.is_available():
+    device = 'cuda'
+elif torch.backends.mps.is_available():
+    device = 'mps'
+else:
+    log.warning('CUDA/MPS are not available, running on CPU')
+dtype = torch.bfloat16
+
+model: ModelConfig = all_model_cfg['large_44k_v2']
+model.download_if_needed()
+output_dir = Path('./output/gradio')
+
+setup_eval_logging()
+
+
+def get_model() -> tuple[MMAudio, FeaturesUtils, SequenceConfig]:
+    seq_cfg = model.seq_cfg
+
+    net: MMAudio = get_my_mmaudio(model.model_name).to(device, dtype).eval()
+    net.load_weights(torch.load(model.model_path, map_location=device, weights_only=True))
+    log.info(f'Loaded weights from {model.model_path}')
+
+    feature_utils = FeaturesUtils(tod_vae_ckpt=model.vae_path,
+                                  synchformer_ckpt=model.synchformer_ckpt,
+                                  enable_conditions=True,
+                                  mode=model.mode,
+                                  bigvgan_vocoder_ckpt=model.bigvgan_16k_path,
+                                  need_vae_encoder=False)
+    feature_utils = feature_utils.to(device, dtype).eval()
+
+    return net, feature_utils, seq_cfg
+
+
+net, feature_utils, seq_cfg = get_model()
+
+
+@torch.inference_mode()
+def video_to_audio(video: gr.Video, prompt: str, negative_prompt: str, seed: int, num_steps: int,
+                   cfg_strength: float, duration: float):
+
+    rng = torch.Generator(device=device)
+    if seed >= 0:
+        rng.manual_seed(seed)
+    else:
+        rng.seed()
+    fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=num_steps)
+
+    video_info = load_video(video, duration)
+    clip_frames = video_info.clip_frames
+    sync_frames = video_info.sync_frames
+    duration = video_info.duration_sec
+    clip_frames = clip_frames.unsqueeze(0)
+    sync_frames = sync_frames.unsqueeze(0)
+    seq_cfg.duration = duration
+    net.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
+
+    audios = generate(clip_frames,
+                      sync_frames, [prompt],
+                      negative_text=[negative_prompt],
+                      feature_utils=feature_utils,
+                      net=net,
+                      fm=fm,
+                      rng=rng,
+                      cfg_strength=cfg_strength)
+    audio = audios.float().cpu()[0]
+
+    current_time_string = datetime.now().strftime('%Y%m%d_%H%M%S')
+    output_dir.mkdir(exist_ok=True, parents=True)
+    video_save_path = output_dir / f'{current_time_string}.mp4'
+    make_video(video_info, video_save_path, audio, sampling_rate=seq_cfg.sampling_rate)
+    gc.collect()
+    return video_save_path
+
+
+@torch.inference_mode()
+def image_to_audio(image: gr.Image, prompt: str, negative_prompt: str, seed: int, num_steps: int,
+                   cfg_strength: float, duration: float):
+
+    rng = torch.Generator(device=device)
+    if seed >= 0:
+        rng.manual_seed(seed)
+    else:
+        rng.seed()
+    fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=num_steps)
+
+    image_info = load_image(image)
+    clip_frames = image_info.clip_frames
+    sync_frames = image_info.sync_frames
+    clip_frames = clip_frames.unsqueeze(0)
+    sync_frames = sync_frames.unsqueeze(0)
+    seq_cfg.duration = duration
+    net.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
+
+    audios = generate(clip_frames,
+                      sync_frames, [prompt],
+                      negative_text=[negative_prompt],
+                      feature_utils=feature_utils,
+                      net=net,
+                      fm=fm,
+                      rng=rng,
+                      cfg_strength=cfg_strength,
+                      image_input=True)
+    audio = audios.float().cpu()[0]
+
+    current_time_string = datetime.now().strftime('%Y%m%d_%H%M%S')
+    output_dir.mkdir(exist_ok=True, parents=True)
+    video_save_path = output_dir / f'{current_time_string}.mp4'
+    video_info = VideoInfo.from_image_info(image_info, duration, fps=Fraction(1))
+    make_video(video_info, video_save_path, audio, sampling_rate=seq_cfg.sampling_rate)
+    gc.collect()
+    return video_save_path
+
+
+@torch.inference_mode()
+def text_to_audio(prompt: str, negative_prompt: str, seed: int, num_steps: int, cfg_strength: float,
+                  duration: float):
+
+    rng = torch.Generator(device=device)
+    if seed >= 0:
+        rng.manual_seed(seed)
+    else:
+        rng.seed()
+    fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=num_steps)
+
+    clip_frames = sync_frames = None
+    seq_cfg.duration = duration
+    net.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
+
+    audios = generate(clip_frames,
+                      sync_frames, [prompt],
+                      negative_text=[negative_prompt],
+                      feature_utils=feature_utils,
+                      net=net,
+                      fm=fm,
+                      rng=rng,
+                      cfg_strength=cfg_strength)
+    audio = audios.float().cpu()[0]
+
+    current_time_string = datetime.now().strftime('%Y%m%d_%H%M%S')
+    output_dir.mkdir(exist_ok=True, parents=True)
+    audio_save_path = output_dir / f'{current_time_string}.flac'
+    torchaudio.save(audio_save_path, audio, seq_cfg.sampling_rate)
+    gc.collect()
+    return audio_save_path
+
+
+video_to_audio_tab = gr.Interface(
+    fn=video_to_audio,
+    description="""
+    Project page: <a href="https://hkchengrex.com/MMAudio/">https://hkchengrex.com/MMAudio/</a><br>
+    Code: <a href="https://github.com/hkchengrex/MMAudio">https://github.com/hkchengrex/MMAudio</a><br>
+
+    NOTE: It takes longer to process high-resolution videos (>384 px on the shorter side). 
+    Doing so does not improve results.
+    """,
+    inputs=[
+        gr.Video(),
+        gr.Text(label='Prompt'),
+        gr.Text(label='Negative prompt', value='music'),
+        gr.Number(label='Seed (-1: random)', value=-1, precision=0, minimum=-1),
+        gr.Number(label='Num steps', value=25, precision=0, minimum=1),
+        gr.Number(label='Guidance Strength', value=4.5, minimum=1),
+        gr.Number(label='Duration (sec)', value=8, minimum=1),
+    ],
+    outputs='playable_video',
+    cache_examples=False,
+    title='MMAudio — Video-to-Audio Synthesis',
+    examples=[
+        [
+            'https://huggingface.co/hkchengrex/MMAudio/resolve/main/examples/sora_beach.mp4',
+            'waves, seagulls',
+            '',
+            0,
+            25,
+            4.5,
+            10,
+        ],
+        [
+            'https://huggingface.co/hkchengrex/MMAudio/resolve/main/examples/sora_serpent.mp4',
+            '',
+            'music',
+            0,
+            25,
+            4.5,
+            10,
+        ],
+        [
+            'https://huggingface.co/hkchengrex/MMAudio/resolve/main/examples/sora_seahorse.mp4',
+            'bubbles',
+            '',
+            0,
+            25,
+            4.5,
+            10,
+        ],
+        [
+            'https://huggingface.co/hkchengrex/MMAudio/resolve/main/examples/sora_india.mp4',
+            'Indian holy music',
+            '',
+            0,
+            25,
+            4.5,
+            10,
+        ],
+        [
+            'https://huggingface.co/hkchengrex/MMAudio/resolve/main/examples/sora_galloping.mp4',
+            'galloping',
+            '',
+            0,
+            25,
+            4.5,
+            10,
+        ],
+        [
+            'https://huggingface.co/hkchengrex/MMAudio/resolve/main/examples/sora_kraken.mp4',
+            'waves, storm',
+            '',
+            0,
+            25,
+            4.5,
+            10,
+        ],
+        [
+            'https://huggingface.co/hkchengrex/MMAudio/resolve/main/examples/mochi_storm.mp4',
+            'storm',
+            '',
+            0,
+            25,
+            4.5,
+            10,
+        ],
+        [
+            'https://huggingface.co/hkchengrex/MMAudio/resolve/main/examples/hunyuan_spring.mp4',
+            '',
+            '',
+            0,
+            25,
+            4.5,
+            10,
+        ],
+        [
+            'https://huggingface.co/hkchengrex/MMAudio/resolve/main/examples/hunyuan_typing.mp4',
+            'typing',
+            '',
+            0,
+            25,
+            4.5,
+            10,
+        ],
+        [
+            'https://huggingface.co/hkchengrex/MMAudio/resolve/main/examples/hunyuan_wake_up.mp4',
+            '',
+            '',
+            0,
+            25,
+            4.5,
+            10,
+        ],
+        [
+            'https://huggingface.co/hkchengrex/MMAudio/resolve/main/examples/sora_nyc.mp4',
+            '',
+            '',
+            0,
+            25,
+            4.5,
+            10,
+        ],
+    ])
+
+text_to_audio_tab = gr.Interface(
+    fn=text_to_audio,
+    description="""
+    Project page: <a href="https://hkchengrex.com/MMAudio/">https://hkchengrex.com/MMAudio/</a><br>
+    Code: <a href="https://github.com/hkchengrex/MMAudio">https://github.com/hkchengrex/MMAudio</a><br>
+    """,
+    inputs=[
+        gr.Text(label='Prompt'),
+        gr.Text(label='Negative prompt'),
+        gr.Number(label='Seed (-1: random)', value=-1, precision=0, minimum=-1),
+        gr.Number(label='Num steps', value=25, precision=0, minimum=1),
+        gr.Number(label='Guidance Strength', value=4.5, minimum=1),
+        gr.Number(label='Duration (sec)', value=8, minimum=1),
+    ],
+    outputs='audio',
+    cache_examples=False,
+    title='MMAudio — Text-to-Audio Synthesis',
+)
+
+image_to_audio_tab = gr.Interface(
+    fn=image_to_audio,
+    description="""
+    Project page: <a href="https://hkchengrex.com/MMAudio/">https://hkchengrex.com/MMAudio/</a><br>
+    Code: <a href="https://github.com/hkchengrex/MMAudio">https://github.com/hkchengrex/MMAudio</a><br>
+
+    NOTE: It takes longer to process high-resolution images (>384 px on the shorter side). 
+    Doing so does not improve results.
+    """,
+    inputs=[
+        gr.Image(type='filepath'),
+        gr.Text(label='Prompt'),
+        gr.Text(label='Negative prompt'),
+        gr.Number(label='Seed (-1: random)', value=-1, precision=0, minimum=-1),
+        gr.Number(label='Num steps', value=25, precision=0, minimum=1),
+        gr.Number(label='Guidance Strength', value=4.5, minimum=1),
+        gr.Number(label='Duration (sec)', value=8, minimum=1),
+    ],
+    outputs='playable_video',
+    cache_examples=False,
+    title='MMAudio — Image-to-Audio Synthesis (experimental)',
+)
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument('--port', type=int, default=7860)
+    args = parser.parse_args()
+
+    gr.TabbedInterface([video_to_audio_tab, text_to_audio_tab, image_to_audio_tab],
+                       ['Video-to-Audio', 'Text-to-Audio', 'Image-to-Audio (experimental)']).launch(
+                           server_port=args.port, allowed_paths=[output_dir])

mmaudio/data/av_utils.py

@@ -0,0 +1,162 @@
+from dataclasses import dataclass
+from fractions import Fraction
+from pathlib import Path
+from typing import Optional
+
+import av
+import numpy as np
+import torch
+from av import AudioFrame
+
+
+@dataclass
+class VideoInfo:
+    duration_sec: float
+    fps: Fraction
+    clip_frames: torch.Tensor
+    sync_frames: torch.Tensor
+    all_frames: Optional[list[np.ndarray]]
+
+    @property
+    def height(self):
+        return self.all_frames[0].shape[0]
+
+    @property
+    def width(self):
+        return self.all_frames[0].shape[1]
+
+    @classmethod
+    def from_image_info(cls, image_info: 'ImageInfo', duration_sec: float,
+                        fps: Fraction) -> 'VideoInfo':
+        num_frames = int(duration_sec * fps)
+        all_frames = [image_info.original_frame] * num_frames
+        return cls(duration_sec=duration_sec,
+                   fps=fps,
+                   clip_frames=image_info.clip_frames,
+                   sync_frames=image_info.sync_frames,
+                   all_frames=all_frames)
+
+
+@dataclass
+class ImageInfo:
+    clip_frames: torch.Tensor
+    sync_frames: torch.Tensor
+    original_frame: Optional[np.ndarray]
+
+    @property
+    def height(self):
+        return self.original_frame.shape[0]
+
+    @property
+    def width(self):
+        return self.original_frame.shape[1]
+
+
+def read_frames(video_path: Path, list_of_fps: list[float], start_sec: float, end_sec: float,
+                need_all_frames: bool) -> tuple[list[np.ndarray], list[np.ndarray], Fraction]:
+    output_frames = [[] for _ in list_of_fps]
+    next_frame_time_for_each_fps = [0.0 for _ in list_of_fps]
+    time_delta_for_each_fps = [1 / fps for fps in list_of_fps]
+    all_frames = []
+
+    # container = av.open(video_path)
+    with av.open(video_path) as container:
+        stream = container.streams.video[0]
+        fps = stream.guessed_rate
+        stream.thread_type = 'AUTO'
+        for packet in container.demux(stream):
+            for frame in packet.decode():
+                frame_time = frame.time
+                if frame_time < start_sec:
+                    continue
+                if frame_time > end_sec:
+                    break
+
+                frame_np = None
+                if need_all_frames:
+                    frame_np = frame.to_ndarray(format='rgb24')
+                    all_frames.append(frame_np)
+
+                for i, _ in enumerate(list_of_fps):
+                    this_time = frame_time
+                    while this_time >= next_frame_time_for_each_fps[i]:
+                        if frame_np is None:
+                            frame_np = frame.to_ndarray(format='rgb24')
+
+                        output_frames[i].append(frame_np)
+                        next_frame_time_for_each_fps[i] += time_delta_for_each_fps[i]
+
+    output_frames = [np.stack(frames) for frames in output_frames]
+    return output_frames, all_frames, fps
+
+
+def reencode_with_audio(video_info: VideoInfo, output_path: Path, audio: torch.Tensor,
+                        sampling_rate: int):
+    container = av.open(output_path, 'w')
+    output_video_stream = container.add_stream('h264', video_info.fps)
+    output_video_stream.codec_context.bit_rate = 10 * 1e6  # 10 Mbps
+    output_video_stream.width = video_info.width
+    output_video_stream.height = video_info.height
+    output_video_stream.pix_fmt = 'yuv420p'
+
+    output_audio_stream = container.add_stream('aac', sampling_rate)
+
+    # encode video
+    for image in video_info.all_frames:
+        image = av.VideoFrame.from_ndarray(image)
+        packet = output_video_stream.encode(image)
+        container.mux(packet)
+
+    for packet in output_video_stream.encode():
+        container.mux(packet)
+
+    # convert float tensor audio to numpy array
+    audio_np = audio.numpy().astype(np.float32)
+    audio_frame = AudioFrame.from_ndarray(audio_np, format='flt', layout='mono')
+    audio_frame.sample_rate = sampling_rate
+
+    for packet in output_audio_stream.encode(audio_frame):
+        container.mux(packet)
+
+    for packet in output_audio_stream.encode():
+        container.mux(packet)
+
+    container.close()
+
+
+def remux_with_audio(video_path: Path, audio: torch.Tensor, output_path: Path, sampling_rate: int):
+    """
+    NOTE: I don't think we can get the exact video duration right without re-encoding
+    so we are not using this but keeping it here for reference
+    """
+    video = av.open(video_path)
+    output = av.open(output_path, 'w')
+    input_video_stream = video.streams.video[0]
+    output_video_stream = output.add_stream(template=input_video_stream)
+    output_audio_stream = output.add_stream('aac', sampling_rate)
+
+    duration_sec = audio.shape[-1] / sampling_rate
+
+    for packet in video.demux(input_video_stream):
+        # We need to skip the "flushing" packets that `demux` generates.
+        if packet.dts is None:
+            continue
+        # We need to assign the packet to the new stream.
+        packet.stream = output_video_stream
+        output.mux(packet)
+
+    # convert float tensor audio to numpy array
+    audio_np = audio.numpy().astype(np.float32)
+    audio_frame = av.AudioFrame.from_ndarray(audio_np, format='flt', layout='mono')
+    audio_frame.sample_rate = sampling_rate
+
+    for packet in output_audio_stream.encode(audio_frame):
+        output.mux(packet)
+
+    for packet in output_audio_stream.encode():
+        output.mux(packet)
+
+    video.close()
+    output.close()
+
+    output.close()

mmaudio/data/data_setup.py

@@ -0,0 +1,174 @@
+import logging
+import random
+
+import numpy as np
+import torch
+from omegaconf import DictConfig
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.dataloader import default_collate
+from torch.utils.data.distributed import DistributedSampler
+
+from mmaudio.data.eval.audiocaps import AudioCapsData
+from mmaudio.data.eval.video_dataset import MovieGen, VGGSound
+from mmaudio.data.extracted_audio import ExtractedAudio
+from mmaudio.data.extracted_vgg import ExtractedVGG
+from mmaudio.data.mm_dataset import MultiModalDataset
+from mmaudio.utils.dist_utils import local_rank
+
+log = logging.getLogger()
+
+
+# Re-seed randomness every time we start a worker
+def worker_init_fn(worker_id: int):
+    worker_seed = torch.initial_seed() % (2**31) + worker_id + local_rank * 1000
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
+    log.debug(f'Worker {worker_id} re-seeded with seed {worker_seed} in rank {local_rank}')
+
+
+def load_vgg_data(cfg: DictConfig, data_cfg: DictConfig) -> Dataset:
+    dataset = ExtractedVGG(tsv_path=data_cfg.tsv,
+                           data_dim=cfg.data_dim,
+                           premade_mmap_dir=data_cfg.memmap_dir)
+
+    return dataset
+
+
+def load_audio_data(cfg: DictConfig, data_cfg: DictConfig) -> Dataset:
+    dataset = ExtractedAudio(tsv_path=data_cfg.tsv,
+                             data_dim=cfg.data_dim,
+                             premade_mmap_dir=data_cfg.memmap_dir)
+
+    return dataset
+
+
+def setup_training_datasets(cfg: DictConfig) -> tuple[Dataset, DistributedSampler, DataLoader]:
+    if cfg.mini_train:
+        vgg = load_vgg_data(cfg, cfg.data.ExtractedVGG_val)
+        audiocaps = load_audio_data(cfg, cfg.data.AudioCaps)
+        dataset = MultiModalDataset([vgg], [audiocaps])
+    if cfg.example_train:
+        video = load_vgg_data(cfg, cfg.data.Example_video)
+        audio = load_audio_data(cfg, cfg.data.Example_audio)
+        dataset = MultiModalDataset([video], [audio])
+    else:
+        # load the largest one first
+        freesound = load_audio_data(cfg, cfg.data.FreeSound)
+        vgg = load_vgg_data(cfg, cfg.data.ExtractedVGG)
+        audiocaps = load_audio_data(cfg, cfg.data.AudioCaps)
+        audioset_sl = load_audio_data(cfg, cfg.data.AudioSetSL)
+        bbcsound = load_audio_data(cfg, cfg.data.BBCSound)
+        clotho = load_audio_data(cfg, cfg.data.Clotho)
+        dataset = MultiModalDataset([vgg] * cfg.vgg_oversample_rate,
+                                    [audiocaps, audioset_sl, bbcsound, freesound, clotho])
+
+    batch_size = cfg.batch_size
+    num_workers = cfg.num_workers
+    pin_memory = cfg.pin_memory
+    sampler, loader = construct_loader(dataset,
+                                       batch_size,
+                                       num_workers,
+                                       shuffle=True,
+                                       drop_last=True,
+                                       pin_memory=pin_memory)
+
+    return dataset, sampler, loader
+
+
+def setup_test_datasets(cfg):
+    dataset = load_vgg_data(cfg, cfg.data.ExtractedVGG_test)
+
+    batch_size = cfg.batch_size
+    num_workers = cfg.num_workers
+    pin_memory = cfg.pin_memory
+    sampler, loader = construct_loader(dataset,
+                                       batch_size,
+                                       num_workers,
+                                       shuffle=False,
+                                       drop_last=False,
+                                       pin_memory=pin_memory)
+
+    return dataset, sampler, loader
+
+
+def setup_val_datasets(cfg: DictConfig) -> tuple[Dataset, DataLoader, DataLoader]:
+    if cfg.example_train:
+        dataset = load_vgg_data(cfg, cfg.data.Example_video)
+    else:
+        dataset = load_vgg_data(cfg, cfg.data.ExtractedVGG_val)
+
+    val_batch_size = cfg.batch_size
+    val_eval_batch_size = cfg.eval_batch_size
+    num_workers = cfg.num_workers
+    pin_memory = cfg.pin_memory
+    _, val_loader = construct_loader(dataset,
+                                     val_batch_size,
+                                     num_workers,
+                                     shuffle=False,
+                                     drop_last=False,
+                                     pin_memory=pin_memory)
+    _, eval_loader = construct_loader(dataset,
+                                      val_eval_batch_size,
+                                      num_workers,
+                                      shuffle=False,
+                                      drop_last=False,
+                                      pin_memory=pin_memory)
+
+    return dataset, val_loader, eval_loader
+
+
+def setup_eval_dataset(dataset_name: str, cfg: DictConfig) -> tuple[Dataset, DataLoader]:
+    if dataset_name.startswith('audiocaps_full'):
+        dataset = AudioCapsData(cfg.eval_data.AudioCaps_full.audio_path,
+                                cfg.eval_data.AudioCaps_full.csv_path)
+    elif dataset_name.startswith('audiocaps'):
+        dataset = AudioCapsData(cfg.eval_data.AudioCaps.audio_path,
+                                cfg.eval_data.AudioCaps.csv_path)
+    elif dataset_name.startswith('moviegen'):
+        dataset = MovieGen(cfg.eval_data.MovieGen.video_path,
+                           cfg.eval_data.MovieGen.jsonl_path,
+                           duration_sec=cfg.duration_s)
+    elif dataset_name.startswith('vggsound'):
+        dataset = VGGSound(cfg.eval_data.VGGSound.video_path,
+                           cfg.eval_data.VGGSound.csv_path,
+                           duration_sec=cfg.duration_s)
+    else:
+        raise ValueError(f'Invalid dataset name: {dataset_name}')
+
+    batch_size = cfg.batch_size
+    num_workers = cfg.num_workers
+    pin_memory = cfg.pin_memory
+    _, loader = construct_loader(dataset,
+                                 batch_size,
+                                 num_workers,
+                                 shuffle=False,
+                                 drop_last=False,
+                                 pin_memory=pin_memory,
+                                 error_avoidance=True)
+    return dataset, loader
+
+
+def error_avoidance_collate(batch):
+    batch = list(filter(lambda x: x is not None, batch))
+    return default_collate(batch)
+
+
+def construct_loader(dataset: Dataset,
+                     batch_size: int,
+                     num_workers: int,
+                     *,
+                     shuffle: bool = True,
+                     drop_last: bool = True,
+                     pin_memory: bool = False,
+                     error_avoidance: bool = False) -> tuple[DistributedSampler, DataLoader]:
+    train_sampler = DistributedSampler(dataset, rank=local_rank, shuffle=shuffle)
+    train_loader = DataLoader(dataset,
+                              batch_size,
+                              sampler=train_sampler,
+                              num_workers=num_workers,
+                              worker_init_fn=worker_init_fn,
+                              drop_last=drop_last,
+                              persistent_workers=num_workers > 0,
+                              pin_memory=pin_memory,
+                              collate_fn=error_avoidance_collate if error_avoidance else None)
+    return train_sampler, train_loader

mmaudio/data/eval/audiocaps.py

@@ -0,0 +1,39 @@
+import logging
+import os
+from collections import defaultdict
+from pathlib import Path
+from typing import Union
+
+import pandas as pd
+import torch
+from torch.utils.data.dataset import Dataset
+
+log = logging.getLogger()
+
+
+class AudioCapsData(Dataset):
+
+    def __init__(self, audio_path: Union[str, Path], csv_path: Union[str, Path]):
+        df = pd.read_csv(csv_path).to_dict(orient='records')
+
+        audio_files = sorted(os.listdir(audio_path))
+        audio_files = set(
+            [Path(f).stem for f in audio_files if f.endswith('.wav') or f.endswith('.flac')])
+
+        self.data = []
+        for row in df:
+            self.data.append({
+                'name': row['name'],
+                'caption': row['caption'],
+            })
+
+        self.audio_path = Path(audio_path)
+        self.csv_path = Path(csv_path)
+
+        log.info(f'Found {len(self.data)} matching audio files in {self.audio_path}')
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        return self.data[idx]
+
+    def __len__(self):
+        return len(self.data)

mmaudio/data/eval/moviegen.py

@@ -0,0 +1,131 @@
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Union
+
+import torch
+from torch.utils.data.dataset import Dataset
+from torchvision.transforms import v2
+from torio.io import StreamingMediaDecoder
+
+from mmaudio.utils.dist_utils import local_rank
+
+log = logging.getLogger()
+
+_CLIP_SIZE = 384
+_CLIP_FPS = 8.0
+
+_SYNC_SIZE = 224
+_SYNC_FPS = 25.0
+
+
+class MovieGenData(Dataset):
+
+    def __init__(
+        self,
+        video_root: Union[str, Path],
+        sync_root: Union[str, Path],
+        jsonl_root: Union[str, Path],
+        *,
+        duration_sec: float = 10.0,
+        read_clip: bool = True,
+    ):
+        self.video_root = Path(video_root)
+        self.sync_root = Path(sync_root)
+        self.jsonl_root = Path(jsonl_root)
+        self.read_clip = read_clip
+
+        videos = sorted(os.listdir(self.video_root))
+        videos = [v[:-4] for v in videos]  # remove extensions
+        self.captions = {}
+
+        for v in videos:
+            with open(self.jsonl_root / (v + '.jsonl')) as f:
+                data = json.load(f)
+                self.captions[v] = data['audio_prompt']
+
+        if local_rank == 0:
+            log.info(f'{len(videos)} videos found in {video_root}')
+
+        self.duration_sec = duration_sec
+
+        self.clip_expected_length = int(_CLIP_FPS * self.duration_sec)
+        self.sync_expected_length = int(_SYNC_FPS * self.duration_sec)
+
+        self.clip_augment = v2.Compose([
+            v2.Resize((_CLIP_SIZE, _CLIP_SIZE), interpolation=v2.InterpolationMode.BICUBIC),
+            v2.ToImage(),
+            v2.ToDtype(torch.float32, scale=True),
+        ])
+
+        self.sync_augment = v2.Compose([
+            v2.Resize((_SYNC_SIZE, _SYNC_SIZE), interpolation=v2.InterpolationMode.BICUBIC),
+            v2.CenterCrop(_SYNC_SIZE),
+            v2.ToImage(),
+            v2.ToDtype(torch.float32, scale=True),
+            v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ])
+
+        self.videos = videos
+
+    def sample(self, idx: int) -> dict[str, torch.Tensor]:
+        video_id = self.videos[idx]
+        caption = self.captions[video_id]
+
+        reader = StreamingMediaDecoder(self.video_root / (video_id + '.mp4'))
+        reader.add_basic_video_stream(
+            frames_per_chunk=int(_CLIP_FPS * self.duration_sec),
+            frame_rate=_CLIP_FPS,
+            format='rgb24',
+        )
+        reader.add_basic_video_stream(
+            frames_per_chunk=int(_SYNC_FPS * self.duration_sec),
+            frame_rate=_SYNC_FPS,
+            format='rgb24',
+        )
+
+        reader.fill_buffer()
+        data_chunk = reader.pop_chunks()
+
+        clip_chunk = data_chunk[0]
+        sync_chunk = data_chunk[1]
+        if clip_chunk is None:
+            raise RuntimeError(f'CLIP video returned None {video_id}')
+        if clip_chunk.shape[0] < self.clip_expected_length:
+            raise RuntimeError(f'CLIP video too short {video_id}')
+
+        if sync_chunk is None:
+            raise RuntimeError(f'Sync video returned None {video_id}')
+        if sync_chunk.shape[0] < self.sync_expected_length:
+            raise RuntimeError(f'Sync video too short {video_id}')
+
+        # truncate the video
+        clip_chunk = clip_chunk[:self.clip_expected_length]
+        if clip_chunk.shape[0] != self.clip_expected_length:
+            raise RuntimeError(f'CLIP video wrong length {video_id}, '
+                               f'expected {self.clip_expected_length}, '
+                               f'got {clip_chunk.shape[0]}')
+        clip_chunk = self.clip_augment(clip_chunk)
+
+        sync_chunk = sync_chunk[:self.sync_expected_length]
+        if sync_chunk.shape[0] != self.sync_expected_length:
+            raise RuntimeError(f'Sync video wrong length {video_id}, '
+                               f'expected {self.sync_expected_length}, '
+                               f'got {sync_chunk.shape[0]}')
+        sync_chunk = self.sync_augment(sync_chunk)
+
+        data = {
+            'name': video_id,
+            'caption': caption,
+            'clip_video': clip_chunk,
+            'sync_video': sync_chunk,
+        }
+
+        return data
+
+    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
+        return self.sample(idx)
+
+    def __len__(self):
+        return len(self.captions)

mmaudio/data/eval/video_dataset.py

@@ -0,0 +1,197 @@
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Union
+
+import pandas as pd
+import torch
+from torch.utils.data.dataset import Dataset
+from torchvision.transforms import v2
+from torio.io import StreamingMediaDecoder
+
+from mmaudio.utils.dist_utils import local_rank
+
+log = logging.getLogger()
+
+_CLIP_SIZE = 384
+_CLIP_FPS = 8.0
+
+_SYNC_SIZE = 224
+_SYNC_FPS = 25.0
+
+
+class VideoDataset(Dataset):
+
+    def __init__(
+        self,
+        video_root: Union[str, Path],
+        *,
+        duration_sec: float = 8.0,
+    ):
+        self.video_root = Path(video_root)
+
+        self.duration_sec = duration_sec
+
+        self.clip_expected_length = int(_CLIP_FPS * self.duration_sec)
+        self.sync_expected_length = int(_SYNC_FPS * self.duration_sec)
+
+        self.clip_transform = v2.Compose([
+            v2.Resize((_CLIP_SIZE, _CLIP_SIZE), interpolation=v2.InterpolationMode.BICUBIC),
+            v2.ToImage(),
+            v2.ToDtype(torch.float32, scale=True),
+        ])
+
+        self.sync_transform = v2.Compose([
+            v2.Resize(_SYNC_SIZE, interpolation=v2.InterpolationMode.BICUBIC),
+            v2.CenterCrop(_SYNC_SIZE),
+            v2.ToImage(),
+            v2.ToDtype(torch.float32, scale=True),
+            v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ])
+
+        # to be implemented by subclasses
+        self.captions = {}
+        self.videos = sorted(list(self.captions.keys()))
+
+    def sample(self, idx: int) -> dict[str, torch.Tensor]:
+        video_id = self.videos[idx]
+        caption = self.captions[video_id]
+
+        reader = StreamingMediaDecoder(self.video_root / (video_id + '.mp4'))
+        reader.add_basic_video_stream(
+            frames_per_chunk=int(_CLIP_FPS * self.duration_sec),
+            frame_rate=_CLIP_FPS,
+            format='rgb24',
+        )
+        reader.add_basic_video_stream(
+            frames_per_chunk=int(_SYNC_FPS * self.duration_sec),
+            frame_rate=_SYNC_FPS,
+            format='rgb24',
+        )
+
+        reader.fill_buffer()
+        data_chunk = reader.pop_chunks()
+
+        clip_chunk = data_chunk[0]
+        sync_chunk = data_chunk[1]
+        if clip_chunk is None:
+            raise RuntimeError(f'CLIP video returned None {video_id}')
+        if clip_chunk.shape[0] < self.clip_expected_length:
+            raise RuntimeError(
+                f'CLIP video too short {video_id}, expected {self.clip_expected_length}, got {clip_chunk.shape[0]}'
+            )
+
+        if sync_chunk is None:
+            raise RuntimeError(f'Sync video returned None {video_id}')
+        if sync_chunk.shape[0] < self.sync_expected_length:
+            raise RuntimeError(
+                f'Sync video too short {video_id}, expected {self.sync_expected_length}, got {sync_chunk.shape[0]}'
+            )
+
+        # truncate the video
+        clip_chunk = clip_chunk[:self.clip_expected_length]
+        if clip_chunk.shape[0] != self.clip_expected_length:
+            raise RuntimeError(f'CLIP video wrong length {video_id}, '
+                               f'expected {self.clip_expected_length}, '
+                               f'got {clip_chunk.shape[0]}')
+        clip_chunk = self.clip_transform(clip_chunk)
+
+        sync_chunk = sync_chunk[:self.sync_expected_length]
+        if sync_chunk.shape[0] != self.sync_expected_length:
+            raise RuntimeError(f'Sync video wrong length {video_id}, '
+                               f'expected {self.sync_expected_length}, '
+                               f'got {sync_chunk.shape[0]}')
+        sync_chunk = self.sync_transform(sync_chunk)
+
+        data = {
+            'name': video_id,
+            'caption': caption,
+            'clip_video': clip_chunk,
+            'sync_video': sync_chunk,
+        }
+
+        return data
+
+    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
+        try:
+            return self.sample(idx)
+        except Exception as e:
+            log.error(f'Error loading video {self.videos[idx]}: {e}')
+            return None
+
+    def __len__(self):
+        return len(self.captions)
+
+
+class VGGSound(VideoDataset):
+
+    def __init__(
+        self,
+        video_root: Union[str, Path],
+        csv_path: Union[str, Path],
+        *,
+        duration_sec: float = 8.0,
+    ):
+        super().__init__(video_root, duration_sec=duration_sec)
+        self.video_root = Path(video_root)
+        self.csv_path = Path(csv_path)
+
+        videos = sorted(os.listdir(self.video_root))
+        if local_rank == 0:
+            log.info(f'{len(videos)} videos found in {video_root}')
+        self.captions = {}
+
+        df = pd.read_csv(csv_path, header=None, names=['id', 'sec', 'caption',
+                                                       'split']).to_dict(orient='records')
+
+        videos_no_found = []
+        for row in df:
+            if row['split'] == 'test':
+                start_sec = int(row['sec'])
+                video_id = str(row['id'])
+                # this is how our videos are named
+                video_name = f'{video_id}_{start_sec:06d}'
+                if video_name + '.mp4' not in videos:
+                    videos_no_found.append(video_name)
+                    continue
+
+                self.captions[video_name] = row['caption']
+
+        if local_rank == 0:
+            log.info(f'{len(videos)} videos found in {video_root}')
+            log.info(f'{len(self.captions)} useable videos found')
+            if videos_no_found:
+                log.info(f'{len(videos_no_found)} found in {csv_path} but not in {video_root}')
+                log.info(
+                    'A small amount is expected, as not all videos are still available on YouTube')
+
+        self.videos = sorted(list(self.captions.keys()))
+
+
+class MovieGen(VideoDataset):
+
+    def __init__(
+        self,
+        video_root: Union[str, Path],
+        jsonl_root: Union[str, Path],
+        *,
+        duration_sec: float = 10.0,
+    ):
+        super().__init__(video_root, duration_sec=duration_sec)
+        self.video_root = Path(video_root)
+        self.jsonl_root = Path(jsonl_root)
+
+        videos = sorted(os.listdir(self.video_root))
+        videos = [v[:-4] for v in videos]  # remove extensions
+        self.captions = {}
+
+        for v in videos:
+            with open(self.jsonl_root / (v + '.jsonl')) as f:
+                data = json.load(f)
+                self.captions[v] = data['audio_prompt']
+
+        if local_rank == 0:
+            log.info(f'{len(videos)} videos found in {video_root}')
+
+        self.videos = videos

mmaudio/data/extracted_audio.py

@@ -0,0 +1,88 @@
+import logging
+from pathlib import Path
+from typing import Union
+
+import pandas as pd
+import torch
+from tensordict import TensorDict
+from torch.utils.data.dataset import Dataset
+
+from mmaudio.utils.dist_utils import local_rank
+
+log = logging.getLogger()
+
+
+class ExtractedAudio(Dataset):
+
+    def __init__(
+        self,
+        tsv_path: Union[str, Path],
+        *,
+        premade_mmap_dir: Union[str, Path],
+        data_dim: dict[str, int],
+    ):
+        super().__init__()
+
+        self.data_dim = data_dim
+        self.df_list = pd.read_csv(tsv_path, sep='\t').to_dict('records')
+        self.ids = [str(d['id']) for d in self.df_list]
+
+        log.info(f'Loading precomputed mmap from {premade_mmap_dir}')
+        # load precomputed memory mapped tensors
+        premade_mmap_dir = Path(premade_mmap_dir)
+        td = TensorDict.load_memmap(premade_mmap_dir)
+        log.info(f'Loaded precomputed mmap from {premade_mmap_dir}')
+        self.mean = td['mean']
+        self.std = td['std']
+        self.text_features = td['text_features']
+
+        log.info(f'Loaded {len(self)} samples from {premade_mmap_dir}.')
+        log.info(f'Loaded mean: {self.mean.shape}.')
+        log.info(f'Loaded std: {self.std.shape}.')
+        log.info(f'Loaded text features: {self.text_features.shape}.')
+
+        assert self.mean.shape[1] == self.data_dim['latent_seq_len'], \
+            f'{self.mean.shape[1]} != {self.data_dim["latent_seq_len"]}'
+        assert self.std.shape[1] == self.data_dim['latent_seq_len'], \
+            f'{self.std.shape[1]} != {self.data_dim["latent_seq_len"]}'
+
+        assert self.text_features.shape[1] == self.data_dim['text_seq_len'], \
+            f'{self.text_features.shape[1]} != {self.data_dim["text_seq_len"]}'
+        assert self.text_features.shape[-1] == self.data_dim['text_dim'], \
+            f'{self.text_features.shape[-1]} != {self.data_dim["text_dim"]}'
+
+        self.fake_clip_features = torch.zeros(self.data_dim['clip_seq_len'],
+                                              self.data_dim['clip_dim'])
+        self.fake_sync_features = torch.zeros(self.data_dim['sync_seq_len'],
+                                              self.data_dim['sync_dim'])
+        self.video_exist = torch.tensor(0, dtype=torch.bool)
+        self.text_exist = torch.tensor(1, dtype=torch.bool)
+
+    def compute_latent_stats(self) -> tuple[torch.Tensor, torch.Tensor]:
+        latents = self.mean
+        return latents.mean(dim=(0, 1)), latents.std(dim=(0, 1))
+
+    def get_memory_mapped_tensor(self) -> TensorDict:
+        td = TensorDict({
+            'mean': self.mean,
+            'std': self.std,
+            'text_features': self.text_features,
+        })
+        return td
+
+    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
+        data = {
+            'id': str(self.df_list[idx]['id']),
+            'a_mean': self.mean[idx],
+            'a_std': self.std[idx],
+            'clip_features': self.fake_clip_features,
+            'sync_features': self.fake_sync_features,
+            'text_features': self.text_features[idx],
+            'caption': self.df_list[idx]['caption'],
+            'video_exist': self.video_exist,
+            'text_exist': self.text_exist,
+        }
+        return data
+
+    def __len__(self):
+        return len(self.ids)

mmaudio/data/extracted_vgg.py

@@ -0,0 +1,101 @@
+import logging
+from pathlib import Path
+from typing import Union
+
+import pandas as pd
+import torch
+from tensordict import TensorDict
+from torch.utils.data.dataset import Dataset
+
+from mmaudio.utils.dist_utils import local_rank
+
+log = logging.getLogger()
+
+
+class ExtractedVGG(Dataset):
+
+    def __init__(
+        self,
+        tsv_path: Union[str, Path],
+        *,
+        premade_mmap_dir: Union[str, Path],
+        data_dim: dict[str, int],
+    ):
+        super().__init__()
+
+        self.data_dim = data_dim
+        self.df_list = pd.read_csv(tsv_path, sep='\t').to_dict('records')
+        self.ids = [d['id'] for d in self.df_list]
+
+        log.info(f'Loading precomputed mmap from {premade_mmap_dir}')
+        # load precomputed memory mapped tensors
+        premade_mmap_dir = Path(premade_mmap_dir)
+        td = TensorDict.load_memmap(premade_mmap_dir)
+        log.info(f'Loaded precomputed mmap from {premade_mmap_dir}')
+        self.mean = td['mean']
+        self.std = td['std']
+        self.clip_features = td['clip_features']
+        self.sync_features = td['sync_features']
+        self.text_features = td['text_features']
+
+        if local_rank == 0:
+            log.info(f'Loaded {len(self)} samples.')
+            log.info(f'Loaded mean: {self.mean.shape}.')
+            log.info(f'Loaded std: {self.std.shape}.')
+            log.info(f'Loaded clip_features: {self.clip_features.shape}.')
+            log.info(f'Loaded sync_features: {self.sync_features.shape}.')
+            log.info(f'Loaded text_features: {self.text_features.shape}.')
+
+        assert self.mean.shape[1] == self.data_dim['latent_seq_len'], \
+            f'{self.mean.shape[1]} != {self.data_dim["latent_seq_len"]}'
+        assert self.std.shape[1] == self.data_dim['latent_seq_len'], \
+            f'{self.std.shape[1]} != {self.data_dim["latent_seq_len"]}'
+
+        assert self.clip_features.shape[1] == self.data_dim['clip_seq_len'], \
+            f'{self.clip_features.shape[1]} != {self.data_dim["clip_seq_len"]}'
+        assert self.sync_features.shape[1] == self.data_dim['sync_seq_len'], \
+            f'{self.sync_features.shape[1]} != {self.data_dim["sync_seq_len"]}'
+        assert self.text_features.shape[1] == self.data_dim['text_seq_len'], \
+            f'{self.text_features.shape[1]} != {self.data_dim["text_seq_len"]}'
+
+        assert self.clip_features.shape[-1] == self.data_dim['clip_dim'], \
+            f'{self.clip_features.shape[-1]} != {self.data_dim["clip_dim"]}'
+        assert self.sync_features.shape[-1] == self.data_dim['sync_dim'], \
+            f'{self.sync_features.shape[-1]} != {self.data_dim["sync_dim"]}'
+        assert self.text_features.shape[-1] == self.data_dim['text_dim'], \
+            f'{self.text_features.shape[-1]} != {self.data_dim["text_dim"]}'
+
+        self.video_exist = torch.tensor(1, dtype=torch.bool)
+        self.text_exist = torch.tensor(1, dtype=torch.bool)
+
+    def compute_latent_stats(self) -> tuple[torch.Tensor, torch.Tensor]:
+        latents = self.mean
+        return latents.mean(dim=(0, 1)), latents.std(dim=(0, 1))
+
+    def get_memory_mapped_tensor(self) -> TensorDict:
+        td = TensorDict({
+            'mean': self.mean,
+            'std': self.std,
+            'clip_features': self.clip_features,
+            'sync_features': self.sync_features,
+            'text_features': self.text_features,
+        })
+        return td
+
+    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
+        data = {
+            'id': self.df_list[idx]['id'],
+            'a_mean': self.mean[idx],
+            'a_std': self.std[idx],
+            'clip_features': self.clip_features[idx],
+            'sync_features': self.sync_features[idx],
+            'text_features': self.text_features[idx],
+            'caption': self.df_list[idx]['label'],
+            'video_exist': self.video_exist,
+            'text_exist': self.text_exist,
+        }
+
+        return data
+
+    def __len__(self):
+        return len(self.ids)

mmaudio/data/extraction/vgg_sound.py

@@ -0,0 +1,193 @@
+import logging
+import os
+from pathlib import Path
+from typing import Optional, Union
+
+import pandas as pd
+import torch
+import torchaudio
+from torch.utils.data.dataset import Dataset
+from torchvision.transforms import v2
+from torio.io import StreamingMediaDecoder
+
+from mmaudio.utils.dist_utils import local_rank
+
+log = logging.getLogger()
+
+_CLIP_SIZE = 384
+_CLIP_FPS = 8.0
+
+_SYNC_SIZE = 224
+_SYNC_FPS = 25.0
+
+
+class VGGSound(Dataset):
+
+    def __init__(
+        self,
+        root: Union[str, Path],
+        *,
+        tsv_path: Union[str, Path] = 'sets/vgg3-train.tsv',
+        sample_rate: int = 16_000,
+        duration_sec: float = 8.0,
+        audio_samples: Optional[int] = None,
+        normalize_audio: bool = False,
+    ):
+        self.root = Path(root)
+        self.normalize_audio = normalize_audio
+        if audio_samples is None:
+            self.audio_samples = int(sample_rate * duration_sec)
+        else:
+            self.audio_samples = audio_samples
+            effective_duration = audio_samples / sample_rate
+            # make sure the duration is close enough, within 15ms
+            assert abs(effective_duration - duration_sec) < 0.015, \
+                f'audio_samples {audio_samples} does not match duration_sec {duration_sec}'
+
+        videos = sorted(os.listdir(self.root))
+        videos = set([Path(v).stem for v in videos])  # remove extensions
+        self.labels = {}
+        self.videos = []
+        missing_videos = []
+
+        # read the tsv for subset information
+        df_list = pd.read_csv(tsv_path, sep='\t', dtype={'id': str}).to_dict('records')
+        for record in df_list:
+            id = record['id']
+            label = record['label']
+            if id in videos:
+                self.labels[id] = label
+                self.videos.append(id)
+            else:
+                missing_videos.append(id)
+
+        if local_rank == 0:
+            log.info(f'{len(videos)} videos found in {root}')
+            log.info(f'{len(self.videos)} videos found in {tsv_path}')
+            log.info(f'{len(missing_videos)} videos missing in {root}')
+
+        self.sample_rate = sample_rate
+        self.duration_sec = duration_sec
+
+        self.expected_audio_length = audio_samples
+        self.clip_expected_length = int(_CLIP_FPS * self.duration_sec)
+        self.sync_expected_length = int(_SYNC_FPS * self.duration_sec)
+
+        self.clip_transform = v2.Compose([
+            v2.Resize((_CLIP_SIZE, _CLIP_SIZE), interpolation=v2.InterpolationMode.BICUBIC),
+            v2.ToImage(),
+            v2.ToDtype(torch.float32, scale=True),
+        ])
+
+        self.sync_transform = v2.Compose([
+            v2.Resize(_SYNC_SIZE, interpolation=v2.InterpolationMode.BICUBIC),
+            v2.CenterCrop(_SYNC_SIZE),
+            v2.ToImage(),
+            v2.ToDtype(torch.float32, scale=True),
+            v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ])
+
+        self.resampler = {}
+
+    def sample(self, idx: int) -> dict[str, torch.Tensor]:
+        video_id = self.videos[idx]
+        label = self.labels[video_id]
+
+        reader = StreamingMediaDecoder(self.root / (video_id + '.mp4'))
+        reader.add_basic_video_stream(
+            frames_per_chunk=int(_CLIP_FPS * self.duration_sec),
+            frame_rate=_CLIP_FPS,
+            format='rgb24',
+        )
+        reader.add_basic_video_stream(
+            frames_per_chunk=int(_SYNC_FPS * self.duration_sec),
+            frame_rate=_SYNC_FPS,
+            format='rgb24',
+        )
+        reader.add_basic_audio_stream(frames_per_chunk=2**30, )
+
+        reader.fill_buffer()
+        data_chunk = reader.pop_chunks()
+
+        clip_chunk = data_chunk[0]
+        sync_chunk = data_chunk[1]
+        audio_chunk = data_chunk[2]
+
+        if clip_chunk is None:
+            raise RuntimeError(f'CLIP video returned None {video_id}')
+        if clip_chunk.shape[0] < self.clip_expected_length:
+            raise RuntimeError(
+                f'CLIP video too short {video_id}, expected {self.clip_expected_length}, got {clip_chunk.shape[0]}'
+            )
+
+        if sync_chunk is None:
+            raise RuntimeError(f'Sync video returned None {video_id}')
+        if sync_chunk.shape[0] < self.sync_expected_length:
+            raise RuntimeError(
+                f'Sync video too short {video_id}, expected {self.sync_expected_length}, got {sync_chunk.shape[0]}'
+            )
+
+        # process audio
+        sample_rate = int(reader.get_out_stream_info(2).sample_rate)
+        audio_chunk = audio_chunk.transpose(0, 1)
+        audio_chunk = audio_chunk.mean(dim=0)  # mono
+        if self.normalize_audio:
+            abs_max = audio_chunk.abs().max()
+            audio_chunk = audio_chunk / abs_max * 0.95
+            if abs_max <= 1e-6:
+                raise RuntimeError(f'Audio is silent {video_id}')
+
+        # resample
+        if sample_rate == self.sample_rate:
+            audio_chunk = audio_chunk
+        else:
+            if sample_rate not in self.resampler:
+                # https://pytorch.org/audio/stable/tutorials/audio_resampling_tutorial.html#kaiser-best
+                self.resampler[sample_rate] = torchaudio.transforms.Resample(
+                    sample_rate,
+                    self.sample_rate,
+                    lowpass_filter_width=64,
+                    rolloff=0.9475937167399596,
+                    resampling_method='sinc_interp_kaiser',
+                    beta=14.769656459379492,
+                )
+            audio_chunk = self.resampler[sample_rate](audio_chunk)
+
+        if audio_chunk.shape[0] < self.expected_audio_length:
+            raise RuntimeError(f'Audio too short {video_id}')
+        audio_chunk = audio_chunk[:self.expected_audio_length]
+
+        # truncate the video
+        clip_chunk = clip_chunk[:self.clip_expected_length]
+        if clip_chunk.shape[0] != self.clip_expected_length:
+            raise RuntimeError(f'CLIP video wrong length {video_id}, '
+                               f'expected {self.clip_expected_length}, '
+                               f'got {clip_chunk.shape[0]}')
+        clip_chunk = self.clip_transform(clip_chunk)
+
+        sync_chunk = sync_chunk[:self.sync_expected_length]
+        if sync_chunk.shape[0] != self.sync_expected_length:
+            raise RuntimeError(f'Sync video wrong length {video_id}, '
+                               f'expected {self.sync_expected_length}, '
+                               f'got {sync_chunk.shape[0]}')
+        sync_chunk = self.sync_transform(sync_chunk)
+
+        data = {
+            'id': video_id,
+            'caption': label,
+            'audio': audio_chunk,
+            'clip_video': clip_chunk,
+            'sync_video': sync_chunk,
+        }
+
+        return data
+
+    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
+        try:
+            return self.sample(idx)
+        except Exception as e:
+            log.error(f'Error loading video {self.videos[idx]}: {e}')
+            return None
+
+    def __len__(self):
+        return len(self.labels)

mmaudio/data/extraction/wav_dataset.py

@@ -0,0 +1,132 @@
+import logging
+import os
+from pathlib import Path
+from typing import Union
+
+import open_clip
+import pandas as pd
+import torch
+import torchaudio
+from torch.utils.data.dataset import Dataset
+
+log = logging.getLogger()
+
+
+class WavTextClipsDataset(Dataset):
+
+    def __init__(
+        self,
+        root: Union[str, Path],
+        *,
+        captions_tsv: Union[str, Path],
+        clips_tsv: Union[str, Path],
+        sample_rate: int,
+        num_samples: int,
+        normalize_audio: bool = False,
+        reject_silent: bool = False,
+        tokenizer_id: str = 'ViT-H-14-378-quickgelu',
+    ):
+        self.root = Path(root)
+        self.sample_rate = sample_rate
+        self.num_samples = num_samples
+        self.normalize_audio = normalize_audio
+        self.reject_silent = reject_silent
+        self.tokenizer = open_clip.get_tokenizer(tokenizer_id)
+
+        audios = sorted(os.listdir(self.root))
+        audios = set([
+            Path(audio).stem for audio in audios
+            if audio.endswith('.wav') or audio.endswith('.flac')
+        ])
+        self.captions = {}
+
+        # read the caption tsv
+        df_list = pd.read_csv(captions_tsv, sep='\t', dtype={'id': str}).to_dict('records')
+        for record in df_list:
+            id = record['id']
+            caption = record['caption']
+            self.captions[id] = caption
+
+        # read the clip tsv
+        df_list = pd.read_csv(clips_tsv, sep='\t', dtype={
+            'id': str,
+            'name': str
+        }).to_dict('records')
+        self.clips = []
+        for record in df_list:
+            record['id'] = record['id']
+            record['name'] = record['name']
+            id = record['id']
+            name = record['name']
+            if name not in self.captions:
+                log.warning(f'Audio {name} not found in {captions_tsv}')
+                continue
+            record['caption'] = self.captions[name]
+            self.clips.append(record)
+
+        log.info(f'Found {len(self.clips)} audio files in {self.root}')
+
+        self.resampler = {}
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        try:
+            clip = self.clips[idx]
+            audio_name = clip['name']
+            audio_id = clip['id']
+            caption = clip['caption']
+            start_sample = clip['start_sample']
+            end_sample = clip['end_sample']
+
+            audio_path = self.root / f'{audio_name}.flac'
+            if not audio_path.exists():
+                audio_path = self.root / f'{audio_name}.wav'
+                assert audio_path.exists()
+
+            audio_chunk, sample_rate = torchaudio.load(audio_path)
+            audio_chunk = audio_chunk.mean(dim=0)  # mono
+            abs_max = audio_chunk.abs().max()
+            if self.normalize_audio:
+                audio_chunk = audio_chunk / abs_max * 0.95
+
+            if self.reject_silent and abs_max < 1e-6:
+                log.warning(f'Rejecting silent audio')
+                return None
+
+            audio_chunk = audio_chunk[start_sample:end_sample]
+
+            # resample
+            if sample_rate == self.sample_rate:
+                audio_chunk = audio_chunk
+            else:
+                if sample_rate not in self.resampler:
+                    # https://pytorch.org/audio/stable/tutorials/audio_resampling_tutorial.html#kaiser-best
+                    self.resampler[sample_rate] = torchaudio.transforms.Resample(
+                        sample_rate,
+                        self.sample_rate,
+                        lowpass_filter_width=64,
+                        rolloff=0.9475937167399596,
+                        resampling_method='sinc_interp_kaiser',
+                        beta=14.769656459379492,
+                    )
+                audio_chunk = self.resampler[sample_rate](audio_chunk)
+
+            if audio_chunk.shape[0] < self.num_samples:
+                raise ValueError('Audio is too short')
+            audio_chunk = audio_chunk[:self.num_samples]
+
+            tokens = self.tokenizer([caption])[0]
+
+            output = {
+                'waveform': audio_chunk,
+                'id': audio_id,
+                'caption': caption,
+                'tokens': tokens,
+            }
+
+            return output
+        except Exception as e:
+            log.error(f'Error reading {audio_path}: {e}')
+            return None
+
+    def __len__(self):
+        return len(self.clips)

mmaudio/data/mm_dataset.py

@@ -0,0 +1,45 @@
+import bisect
+
+import torch
+from torch.utils.data.dataset import Dataset
+
+
+# modified from https://pytorch.org/docs/stable/_modules/torch/utils/data/dataset.html#ConcatDataset
+class MultiModalDataset(Dataset):
+    datasets: list[Dataset]
+    cumulative_sizes: list[int]
+
+    @staticmethod
+    def cumsum(sequence):
+        r, s = [], 0
+        for e in sequence:
+            l = len(e)
+            r.append(l + s)
+            s += l
+        return r
+
+    def __init__(self, video_datasets: list[Dataset], audio_datasets: list[Dataset]):
+        super().__init__()
+        self.video_datasets = list(video_datasets)
+        self.audio_datasets = list(audio_datasets)
+        self.datasets = self.video_datasets + self.audio_datasets
+
+        self.cumulative_sizes = self.cumsum(self.datasets)
+
+    def __len__(self):
+        return self.cumulative_sizes[-1]
+
+    def __getitem__(self, idx):
+        if idx < 0:
+            if -idx > len(self):
+                raise ValueError("absolute value of index should not exceed dataset length")
+            idx = len(self) + idx
+        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
+        if dataset_idx == 0:
+            sample_idx = idx
+        else:
+            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
+        return self.datasets[dataset_idx][sample_idx]
+
+    def compute_latent_stats(self) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.video_datasets[0].compute_latent_stats()

mmaudio/data/utils.py

@@ -0,0 +1,148 @@
+import logging
+import os
+import random
+import tempfile
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import torch
+import torch.distributed as dist
+from tensordict import MemoryMappedTensor
+from torch.utils.data import DataLoader
+from torch.utils.data.dataset import Dataset
+from tqdm import tqdm
+
+from mmaudio.utils.dist_utils import local_rank, world_size
+
+scratch_path = Path(os.environ['SLURM_SCRATCH'] if 'SLURM_SCRATCH' in os.environ else '/dev/shm')
+shm_path = Path('/dev/shm')
+
+log = logging.getLogger()
+
+
+def reseed(seed):
+    random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def local_scatter_torch(obj: Optional[Any]):
+    if world_size == 1:
+        # Just one worker. Do nothing.
+        return obj
+
+    array = [obj] * world_size
+    target_array = [None]
+    if local_rank == 0:
+        dist.scatter_object_list(target_array, scatter_object_input_list=array, src=0)
+    else:
+        dist.scatter_object_list(target_array, scatter_object_input_list=None, src=0)
+    return target_array[0]
+
+
+class ShardDataset(Dataset):
+
+    def __init__(self, root):
+        self.root = root
+        self.shards = sorted(os.listdir(root))
+
+    def __len__(self):
+        return len(self.shards)
+
+    def __getitem__(self, idx):
+        return torch.load(os.path.join(self.root, self.shards[idx]), weights_only=True)
+
+
+def get_tmp_dir(in_memory: bool) -> Path:
+    return shm_path if in_memory else scratch_path
+
+
+def load_shards_and_share(data_path: Union[str, Path], ids: list[int],
+                          in_memory: bool) -> MemoryMappedTensor:
+    if local_rank == 0:
+        with tempfile.NamedTemporaryFile(prefix='shared-tensor-', dir=get_tmp_dir(in_memory)) as f:
+            log.info(f'Loading shards from {data_path} into {f.name}...')
+            data = load_shards(data_path, ids=ids, tmp_file_path=f.name)
+            data = share_tensor_to_all(data)
+            torch.distributed.barrier()
+            f.close()  # why does the context manager not close the file for me?
+    else:
+        log.info('Waiting for the data to be shared with me...')
+        data = share_tensor_to_all(None)
+        torch.distributed.barrier()
+
+    return data
+
+
+def load_shards(
+    data_path: Union[str, Path],
+    ids: list[int],
+    *,
+    tmp_file_path: str,
+) -> Union[torch.Tensor, dict[str, torch.Tensor]]:
+
+    id_set = set(ids)
+    shards = sorted(os.listdir(data_path))
+    log.info(f'Found {len(shards)} shards in {data_path}.')
+    first_shard = torch.load(os.path.join(data_path, shards[0]), weights_only=True)
+
+    log.info(f'Rank {local_rank} created file {tmp_file_path}')
+    first_item = next(iter(first_shard.values()))
+    log.info(f'First item shape: {first_item.shape}')
+    mm_tensor = MemoryMappedTensor.empty(shape=(len(ids), *first_item.shape),
+                                         dtype=torch.float32,
+                                         filename=tmp_file_path,
+                                         existsok=True)
+    total_count = 0
+    used_index = set()
+    id_indexing = {i: idx for idx, i in enumerate(ids)}
+    # faster with no workers; otherwise we need to set_sharing_strategy('file_system')
+    loader = DataLoader(ShardDataset(data_path), batch_size=1, num_workers=0)
+    for data in tqdm(loader, desc='Loading shards'):
+        for i, v in data.items():
+            if i not in id_set:
+                continue
+
+            # tensor_index = ids.index(i)
+            tensor_index = id_indexing[i]
+            if tensor_index in used_index:
+                raise ValueError(f'Duplicate id {i} found in {data_path}.')
+            used_index.add(tensor_index)
+            mm_tensor[tensor_index] = v
+            total_count += 1
+
+    assert total_count == len(ids), f'Expected {len(ids)} tensors, got {total_count}.'
+    log.info(f'Loaded {total_count} tensors from {data_path}.')
+
+    return mm_tensor
+
+
+def share_tensor_to_all(x: Optional[MemoryMappedTensor]) -> MemoryMappedTensor:
+    """
+    x: the tensor to be shared; None if local_rank != 0
+    return: the shared tensor
+    """
+
+    # there is no need to share your stuff with anyone if you are alone; must be in memory
+    if world_size == 1:
+        return x
+
+    if local_rank == 0:
+        assert x is not None, 'x must not be None if local_rank == 0'
+    else:
+        assert x is None, 'x must be None if local_rank != 0'
+
+    if local_rank == 0:
+        filename = x.filename
+        meta_information = (filename, x.shape, x.dtype)
+    else:
+        meta_information = None
+
+    filename, data_shape, data_type = local_scatter_torch(meta_information)
+    if local_rank == 0:
+        data = x
+    else:
+        data = MemoryMappedTensor.from_filename(filename=filename,
+                                                dtype=data_type,
+                                                shape=data_shape)
+
+    return data

mmaudio/eval_utils.py

@@ -0,0 +1,255 @@
+import dataclasses
+import logging
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import torch
+from colorlog import ColoredFormatter
+from PIL import Image
+from torchvision.transforms import v2
+
+from mmaudio.data.av_utils import ImageInfo, VideoInfo, read_frames, reencode_with_audio
+from mmaudio.model.flow_matching import FlowMatching
+from mmaudio.model.networks import MMAudio
+from mmaudio.model.sequence_config import CONFIG_16K, CONFIG_44K, SequenceConfig
+from mmaudio.model.utils.features_utils import FeaturesUtils
+from mmaudio.utils.download_utils import download_model_if_needed
+
+log = logging.getLogger()
+
+
+@dataclasses.dataclass
+class ModelConfig:
+    model_name: str
+    model_path: Path
+    vae_path: Path
+    bigvgan_16k_path: Optional[Path]
+    mode: str
+    synchformer_ckpt: Path = Path('./ext_weights/synchformer_state_dict.pth')
+
+    @property
+    def seq_cfg(self) -> SequenceConfig:
+        if self.mode == '16k':
+            return CONFIG_16K
+        elif self.mode == '44k':
+            return CONFIG_44K
+
+    def download_if_needed(self):
+        download_model_if_needed(self.model_path)
+        download_model_if_needed(self.vae_path)
+        if self.bigvgan_16k_path is not None:
+            download_model_if_needed(self.bigvgan_16k_path)
+        download_model_if_needed(self.synchformer_ckpt)
+
+
+small_16k = ModelConfig(model_name='small_16k',
+                        model_path=Path('./weights/mmaudio_small_16k.pth'),
+                        vae_path=Path('./ext_weights/v1-16.pth'),
+                        bigvgan_16k_path=Path('./ext_weights/best_netG.pt'),
+                        mode='16k')
+small_44k = ModelConfig(model_name='small_44k',
+                        model_path=Path('./weights/mmaudio_small_44k.pth'),
+                        vae_path=Path('./ext_weights/v1-44.pth'),
+                        bigvgan_16k_path=None,
+                        mode='44k')
+medium_44k = ModelConfig(model_name='medium_44k',
+                         model_path=Path('./weights/mmaudio_medium_44k.pth'),
+                         vae_path=Path('./ext_weights/v1-44.pth'),
+                         bigvgan_16k_path=None,
+                         mode='44k')
+large_44k = ModelConfig(model_name='large_44k',
+                        model_path=Path('./weights/mmaudio_large_44k.pth'),
+                        vae_path=Path('./ext_weights/v1-44.pth'),
+                        bigvgan_16k_path=None,
+                        mode='44k')
+large_44k_v2 = ModelConfig(model_name='large_44k_v2',
+                           model_path=Path('./weights/mmaudio_large_44k_v2.pth'),
+                           vae_path=Path('./ext_weights/v1-44.pth'),
+                           bigvgan_16k_path=None,
+                           mode='44k')
+all_model_cfg: dict[str, ModelConfig] = {
+    'small_16k': small_16k,
+    'small_44k': small_44k,
+    'medium_44k': medium_44k,
+    'large_44k': large_44k,
+    'large_44k_v2': large_44k_v2,
+}
+
+
+def generate(
+    clip_video: Optional[torch.Tensor],
+    sync_video: Optional[torch.Tensor],
+    text: Optional[list[str]],
+    *,
+    negative_text: Optional[list[str]] = None,
+    feature_utils: FeaturesUtils,
+    net: MMAudio,
+    fm: FlowMatching,
+    rng: torch.Generator,
+    cfg_strength: float,
+    clip_batch_size_multiplier: int = 40,
+    sync_batch_size_multiplier: int = 40,
+    image_input: bool = False,
+) -> torch.Tensor:
+    device = feature_utils.device
+    dtype = feature_utils.dtype
+
+    bs = len(text)
+    if clip_video is not None:
+        clip_video = clip_video.to(device, dtype, non_blocking=True)
+        clip_features = feature_utils.encode_video_with_clip(clip_video,
+                                                             batch_size=bs *
+                                                             clip_batch_size_multiplier)
+        if image_input:
+            clip_features = clip_features.expand(-1, net.clip_seq_len, -1)
+    else:
+        clip_features = net.get_empty_clip_sequence(bs)
+
+    if sync_video is not None and not image_input:
+        sync_video = sync_video.to(device, dtype, non_blocking=True)
+        sync_features = feature_utils.encode_video_with_sync(sync_video,
+                                                             batch_size=bs *
+                                                             sync_batch_size_multiplier)
+    else:
+        sync_features = net.get_empty_sync_sequence(bs)
+
+    if text is not None:
+        text_features = feature_utils.encode_text(text)
+    else:
+        text_features = net.get_empty_string_sequence(bs)
+
+    if negative_text is not None:
+        assert len(negative_text) == bs
+        negative_text_features = feature_utils.encode_text(negative_text)
+    else:
+        negative_text_features = net.get_empty_string_sequence(bs)
+
+    x0 = torch.randn(bs,
+                     net.latent_seq_len,
+                     net.latent_dim,
+                     device=device,
+                     dtype=dtype,
+                     generator=rng)
+    preprocessed_conditions = net.preprocess_conditions(clip_features, sync_features, text_features)
+    empty_conditions = net.get_empty_conditions(
+        bs, negative_text_features=negative_text_features if negative_text is not None else None)
+
+    cfg_ode_wrapper = lambda t, x: net.ode_wrapper(t, x, preprocessed_conditions, empty_conditions,
+                                                   cfg_strength)
+    x1 = fm.to_data(cfg_ode_wrapper, x0)
+    x1 = net.unnormalize(x1)
+    spec = feature_utils.decode(x1)
+    audio = feature_utils.vocode(spec)
+    return audio
+
+
+LOGFORMAT = "[%(log_color)s%(levelname)-8s%(reset)s]: %(log_color)s%(message)s%(reset)s"
+
+
+def setup_eval_logging(log_level: int = logging.INFO):
+    logging.root.setLevel(log_level)
+    formatter = ColoredFormatter(LOGFORMAT)
+    stream = logging.StreamHandler()
+    stream.setLevel(log_level)
+    stream.setFormatter(formatter)
+    log = logging.getLogger()
+    log.setLevel(log_level)
+    log.addHandler(stream)
+
+
+_CLIP_SIZE = 384
+_CLIP_FPS = 8.0
+
+_SYNC_SIZE = 224
+_SYNC_FPS = 25.0
+
+
+def load_video(video_path: Path, duration_sec: float, load_all_frames: bool = True) -> VideoInfo:
+
+    clip_transform = v2.Compose([
+        v2.Resize((_CLIP_SIZE, _CLIP_SIZE), interpolation=v2.InterpolationMode.BICUBIC),
+        v2.ToImage(),
+        v2.ToDtype(torch.float32, scale=True),
+    ])
+
+    sync_transform = v2.Compose([
+        v2.Resize(_SYNC_SIZE, interpolation=v2.InterpolationMode.BICUBIC),
+        v2.CenterCrop(_SYNC_SIZE),
+        v2.ToImage(),
+        v2.ToDtype(torch.float32, scale=True),
+        v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    ])
+
+    output_frames, all_frames, orig_fps = read_frames(video_path,
+                                                      list_of_fps=[_CLIP_FPS, _SYNC_FPS],
+                                                      start_sec=0,
+                                                      end_sec=duration_sec,
+                                                      need_all_frames=load_all_frames)
+
+    clip_chunk, sync_chunk = output_frames
+    clip_chunk = torch.from_numpy(clip_chunk).permute(0, 3, 1, 2)
+    sync_chunk = torch.from_numpy(sync_chunk).permute(0, 3, 1, 2)
+
+    clip_frames = clip_transform(clip_chunk)
+    sync_frames = sync_transform(sync_chunk)
+
+    clip_length_sec = clip_frames.shape[0] / _CLIP_FPS
+    sync_length_sec = sync_frames.shape[0] / _SYNC_FPS
+
+    if clip_length_sec < duration_sec:
+        log.warning(f'Clip video is too short: {clip_length_sec:.2f} < {duration_sec:.2f}')
+        log.warning(f'Truncating to {clip_length_sec:.2f} sec')
+        duration_sec = clip_length_sec
+
+    if sync_length_sec < duration_sec:
+        log.warning(f'Sync video is too short: {sync_length_sec:.2f} < {duration_sec:.2f}')
+        log.warning(f'Truncating to {sync_length_sec:.2f} sec')
+        duration_sec = sync_length_sec
+
+    clip_frames = clip_frames[:int(_CLIP_FPS * duration_sec)]
+    sync_frames = sync_frames[:int(_SYNC_FPS * duration_sec)]
+
+    video_info = VideoInfo(
+        duration_sec=duration_sec,
+        fps=orig_fps,
+        clip_frames=clip_frames,
+        sync_frames=sync_frames,
+        all_frames=all_frames if load_all_frames else None,
+    )
+    return video_info
+
+
+def load_image(image_path: Path) -> VideoInfo:
+    clip_transform = v2.Compose([
+        v2.Resize((_CLIP_SIZE, _CLIP_SIZE), interpolation=v2.InterpolationMode.BICUBIC),
+        v2.ToImage(),
+        v2.ToDtype(torch.float32, scale=True),
+    ])
+
+    sync_transform = v2.Compose([
+        v2.Resize(_SYNC_SIZE, interpolation=v2.InterpolationMode.BICUBIC),
+        v2.CenterCrop(_SYNC_SIZE),
+        v2.ToImage(),
+        v2.ToDtype(torch.float32, scale=True),
+        v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    ])
+
+    frame = np.array(Image.open(image_path))
+
+    clip_chunk = torch.from_numpy(frame).unsqueeze(0).permute(0, 3, 1, 2)
+    sync_chunk = torch.from_numpy(frame).unsqueeze(0).permute(0, 3, 1, 2)
+
+    clip_frames = clip_transform(clip_chunk)
+    sync_frames = sync_transform(sync_chunk)
+
+    video_info = ImageInfo(
+        clip_frames=clip_frames,
+        sync_frames=sync_frames,
+        original_frame=frame,
+    )
+    return video_info
+
+
+def make_video(video_info: VideoInfo, output_path: Path, audio: torch.Tensor, sampling_rate: int):
+    reencode_with_audio(video_info, output_path, audio, sampling_rate)

mmaudio/ext/__init__.py

@@ -0,0 +1 @@
+

mmaudio/ext/autoencoder/__init__.py

@@ -0,0 +1 @@
+from .autoencoder import AutoEncoderModule

mmaudio/ext/autoencoder/autoencoder.py

@@ -0,0 +1,52 @@
+from typing import Literal, Optional
+
+import torch
+import torch.nn as nn
+
+from mmaudio.ext.autoencoder.vae import VAE, get_my_vae
+from mmaudio.ext.bigvgan import BigVGAN
+from mmaudio.ext.bigvgan_v2.bigvgan import BigVGAN as BigVGANv2
+from mmaudio.model.utils.distributions import DiagonalGaussianDistribution
+
+
+class AutoEncoderModule(nn.Module):
+
+    def __init__(self,
+                 *,
+                 vae_ckpt_path,
+                 vocoder_ckpt_path: Optional[str] = None,
+                 mode: Literal['16k', '44k'],
+                 need_vae_encoder: bool = True):
+        super().__init__()
+        self.vae: VAE = get_my_vae(mode).eval()
+        vae_state_dict = torch.load(vae_ckpt_path, weights_only=True, map_location='cpu')
+        self.vae.load_state_dict(vae_state_dict)
+        self.vae.remove_weight_norm()
+
+        if mode == '16k':
+            assert vocoder_ckpt_path is not None
+            self.vocoder = BigVGAN(vocoder_ckpt_path).eval()
+        elif mode == '44k':
+            self.vocoder = BigVGANv2.from_pretrained('nvidia/bigvgan_v2_44khz_128band_512x',
+                                                     use_cuda_kernel=False)
+            self.vocoder.remove_weight_norm()
+        else:
+            raise ValueError(f'Unknown mode: {mode}')
+
+        for param in self.parameters():
+            param.requires_grad = False
+
+        if not need_vae_encoder:
+            del self.vae.encoder
+
+    @torch.inference_mode()
+    def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution:
+        return self.vae.encode(x)
+
+    @torch.inference_mode()
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        return self.vae.decode(z)
+
+    @torch.inference_mode()
+    def vocode(self, spec: torch.Tensor) -> torch.Tensor:
+        return self.vocoder(spec)

mmaudio/ext/autoencoder/edm2_utils.py

@@ -0,0 +1,168 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# This work is licensed under a Creative Commons
+# Attribution-NonCommercial-ShareAlike 4.0 International License.
+# You should have received a copy of the license along with this
+# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/
+"""Improved diffusion model architecture proposed in the paper
+"Analyzing and Improving the Training Dynamics of Diffusion Models"."""
+
+import numpy as np
+import torch
+
+#----------------------------------------------------------------------------
+# Variant of constant() that inherits dtype and device from the given
+# reference tensor by default.
+
+_constant_cache = dict()
+
+
+def constant(value, shape=None, dtype=None, device=None, memory_format=None):
+    value = np.asarray(value)
+    if shape is not None:
+        shape = tuple(shape)
+    if dtype is None:
+        dtype = torch.get_default_dtype()
+    if device is None:
+        device = torch.device('cpu')
+    if memory_format is None:
+        memory_format = torch.contiguous_format
+
+    key = (value.shape, value.dtype, value.tobytes(), shape, dtype, device, memory_format)
+    tensor = _constant_cache.get(key, None)
+    if tensor is None:
+        tensor = torch.as_tensor(value.copy(), dtype=dtype, device=device)
+        if shape is not None:
+            tensor, _ = torch.broadcast_tensors(tensor, torch.empty(shape))
+        tensor = tensor.contiguous(memory_format=memory_format)
+        _constant_cache[key] = tensor
+    return tensor
+
+
+def const_like(ref, value, shape=None, dtype=None, device=None, memory_format=None):
+    if dtype is None:
+        dtype = ref.dtype
+    if device is None:
+        device = ref.device
+    return constant(value, shape=shape, dtype=dtype, device=device, memory_format=memory_format)
+
+
+#----------------------------------------------------------------------------
+# Normalize given tensor to unit magnitude with respect to the given
+# dimensions. Default = all dimensions except the first.
+
+
+def normalize(x, dim=None, eps=1e-4):
+    if dim is None:
+        dim = list(range(1, x.ndim))
+    norm = torch.linalg.vector_norm(x, dim=dim, keepdim=True, dtype=torch.float32)
+    norm = torch.add(eps, norm, alpha=np.sqrt(norm.numel() / x.numel()))
+    return x / norm.to(x.dtype)
+
+
+class Normalize(torch.nn.Module):
+
+    def __init__(self, dim=None, eps=1e-4):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def forward(self, x):
+        return normalize(x, dim=self.dim, eps=self.eps)
+
+
+#----------------------------------------------------------------------------
+# Upsample or downsample the given tensor with the given filter,
+# or keep it as is.
+
+
+def resample(x, f=[1, 1], mode='keep'):
+    if mode == 'keep':
+        return x
+    f = np.float32(f)
+    assert f.ndim == 1 and len(f) % 2 == 0
+    pad = (len(f) - 1) // 2
+    f = f / f.sum()
+    f = np.outer(f, f)[np.newaxis, np.newaxis, :, :]
+    f = const_like(x, f)
+    c = x.shape[1]
+    if mode == 'down':
+        return torch.nn.functional.conv2d(x,
+                                          f.tile([c, 1, 1, 1]),
+                                          groups=c,
+                                          stride=2,
+                                          padding=(pad, ))
+    assert mode == 'up'
+    return torch.nn.functional.conv_transpose2d(x, (f * 4).tile([c, 1, 1, 1]),
+                                                groups=c,
+                                                stride=2,
+                                                padding=(pad, ))
+
+
+#----------------------------------------------------------------------------
+# Magnitude-preserving SiLU (Equation 81).
+
+
+def mp_silu(x):
+    return torch.nn.functional.silu(x) / 0.596
+
+
+class MPSiLU(torch.nn.Module):
+
+    def forward(self, x):
+        return mp_silu(x)
+
+
+#----------------------------------------------------------------------------
+# Magnitude-preserving sum (Equation 88).
+
+
+def mp_sum(a, b, t=0.5):
+    return a.lerp(b, t) / np.sqrt((1 - t)**2 + t**2)
+
+
+#----------------------------------------------------------------------------
+# Magnitude-preserving concatenation (Equation 103).
+
+
+def mp_cat(a, b, dim=1, t=0.5):
+    Na = a.shape[dim]
+    Nb = b.shape[dim]
+    C = np.sqrt((Na + Nb) / ((1 - t)**2 + t**2))
+    wa = C / np.sqrt(Na) * (1 - t)
+    wb = C / np.sqrt(Nb) * t
+    return torch.cat([wa * a, wb * b], dim=dim)
+
+
+#----------------------------------------------------------------------------
+# Magnitude-preserving convolution or fully-connected layer (Equation 47)
+# with force weight normalization (Equation 66).
+
+
+class MPConv1D(torch.nn.Module):
+
+    def __init__(self, in_channels, out_channels, kernel_size):
+        super().__init__()
+        self.out_channels = out_channels
+        self.weight = torch.nn.Parameter(torch.randn(out_channels, in_channels, kernel_size))
+
+        self.weight_norm_removed = False
+
+    def forward(self, x, gain=1):
+        assert self.weight_norm_removed, 'call remove_weight_norm() before inference'
+
+        w = self.weight * gain
+        if w.ndim == 2:
+            return x @ w.t()
+        assert w.ndim == 3
+        return torch.nn.functional.conv1d(x, w, padding=(w.shape[-1] // 2, ))
+
+    def remove_weight_norm(self):
+        w = self.weight.to(torch.float32)
+        w = normalize(w)  # traditional weight normalization
+        w = w / np.sqrt(w[0].numel())
+        w = w.to(self.weight.dtype)
+        self.weight.data.copy_(w)
+
+        self.weight_norm_removed = True
+        return self

mmaudio/ext/autoencoder/vae.py

@@ -0,0 +1,369 @@
+import logging
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from mmaudio.ext.autoencoder.edm2_utils import MPConv1D
+from mmaudio.ext.autoencoder.vae_modules import (AttnBlock1D, Downsample1D, ResnetBlock1D,
+                                                 Upsample1D, nonlinearity)
+from mmaudio.model.utils.distributions import DiagonalGaussianDistribution
+
+log = logging.getLogger()
+
+DATA_MEAN_80D = [
+    -1.6058, -1.3676, -1.2520, -1.2453, -1.2078, -1.2224, -1.2419, -1.2439, -1.2922, -1.2927,
+    -1.3170, -1.3543, -1.3401, -1.3836, -1.3907, -1.3912, -1.4313, -1.4152, -1.4527, -1.4728,
+    -1.4568, -1.5101, -1.5051, -1.5172, -1.5623, -1.5373, -1.5746, -1.5687, -1.6032, -1.6131,
+    -1.6081, -1.6331, -1.6489, -1.6489, -1.6700, -1.6738, -1.6953, -1.6969, -1.7048, -1.7280,
+    -1.7361, -1.7495, -1.7658, -1.7814, -1.7889, -1.8064, -1.8221, -1.8377, -1.8417, -1.8643,
+    -1.8857, -1.8929, -1.9173, -1.9379, -1.9531, -1.9673, -1.9824, -2.0042, -2.0215, -2.0436,
+    -2.0766, -2.1064, -2.1418, -2.1855, -2.2319, -2.2767, -2.3161, -2.3572, -2.3954, -2.4282,
+    -2.4659, -2.5072, -2.5552, -2.6074, -2.6584, -2.7107, -2.7634, -2.8266, -2.8981, -2.9673
+]
+
+DATA_STD_80D = [
+    1.0291, 1.0411, 1.0043, 0.9820, 0.9677, 0.9543, 0.9450, 0.9392, 0.9343, 0.9297, 0.9276, 0.9263,
+    0.9242, 0.9254, 0.9232, 0.9281, 0.9263, 0.9315, 0.9274, 0.9247, 0.9277, 0.9199, 0.9188, 0.9194,
+    0.9160, 0.9161, 0.9146, 0.9161, 0.9100, 0.9095, 0.9145, 0.9076, 0.9066, 0.9095, 0.9032, 0.9043,
+    0.9038, 0.9011, 0.9019, 0.9010, 0.8984, 0.8983, 0.8986, 0.8961, 0.8962, 0.8978, 0.8962, 0.8973,
+    0.8993, 0.8976, 0.8995, 0.9016, 0.8982, 0.8972, 0.8974, 0.8949, 0.8940, 0.8947, 0.8936, 0.8939,
+    0.8951, 0.8956, 0.9017, 0.9167, 0.9436, 0.9690, 1.0003, 1.0225, 1.0381, 1.0491, 1.0545, 1.0604,
+    1.0761, 1.0929, 1.1089, 1.1196, 1.1176, 1.1156, 1.1117, 1.1070
+]
+
+DATA_MEAN_128D = [
+    -3.3462, -2.6723, -2.4893, -2.3143, -2.2664, -2.3317, -2.1802, -2.4006, -2.2357, -2.4597,
+    -2.3717, -2.4690, -2.5142, -2.4919, -2.6610, -2.5047, -2.7483, -2.5926, -2.7462, -2.7033,
+    -2.7386, -2.8112, -2.7502, -2.9594, -2.7473, -3.0035, -2.8891, -2.9922, -2.9856, -3.0157,
+    -3.1191, -2.9893, -3.1718, -3.0745, -3.1879, -3.2310, -3.1424, -3.2296, -3.2791, -3.2782,
+    -3.2756, -3.3134, -3.3509, -3.3750, -3.3951, -3.3698, -3.4505, -3.4509, -3.5089, -3.4647,
+    -3.5536, -3.5788, -3.5867, -3.6036, -3.6400, -3.6747, -3.7072, -3.7279, -3.7283, -3.7795,
+    -3.8259, -3.8447, -3.8663, -3.9182, -3.9605, -3.9861, -4.0105, -4.0373, -4.0762, -4.1121,
+    -4.1488, -4.1874, -4.2461, -4.3170, -4.3639, -4.4452, -4.5282, -4.6297, -4.7019, -4.7960,
+    -4.8700, -4.9507, -5.0303, -5.0866, -5.1634, -5.2342, -5.3242, -5.4053, -5.4927, -5.5712,
+    -5.6464, -5.7052, -5.7619, -5.8410, -5.9188, -6.0103, -6.0955, -6.1673, -6.2362, -6.3120,
+    -6.3926, -6.4797, -6.5565, -6.6511, -6.8130, -6.9961, -7.1275, -7.2457, -7.3576, -7.4663,
+    -7.6136, -7.7469, -7.8815, -8.0132, -8.1515, -8.3071, -8.4722, -8.7418, -9.3975, -9.6628,
+    -9.7671, -9.8863, -9.9992, -10.0860, -10.1709, -10.5418, -11.2795, -11.3861
+]
+
+DATA_STD_128D = [
+    2.3804, 2.4368, 2.3772, 2.3145, 2.2803, 2.2510, 2.2316, 2.2083, 2.1996, 2.1835, 2.1769, 2.1659,
+    2.1631, 2.1618, 2.1540, 2.1606, 2.1571, 2.1567, 2.1612, 2.1579, 2.1679, 2.1683, 2.1634, 2.1557,
+    2.1668, 2.1518, 2.1415, 2.1449, 2.1406, 2.1350, 2.1313, 2.1415, 2.1281, 2.1352, 2.1219, 2.1182,
+    2.1327, 2.1195, 2.1137, 2.1080, 2.1179, 2.1036, 2.1087, 2.1036, 2.1015, 2.1068, 2.0975, 2.0991,
+    2.0902, 2.1015, 2.0857, 2.0920, 2.0893, 2.0897, 2.0910, 2.0881, 2.0925, 2.0873, 2.0960, 2.0900,
+    2.0957, 2.0958, 2.0978, 2.0936, 2.0886, 2.0905, 2.0845, 2.0855, 2.0796, 2.0840, 2.0813, 2.0817,
+    2.0838, 2.0840, 2.0917, 2.1061, 2.1431, 2.1976, 2.2482, 2.3055, 2.3700, 2.4088, 2.4372, 2.4609,
+    2.4731, 2.4847, 2.5072, 2.5451, 2.5772, 2.6147, 2.6529, 2.6596, 2.6645, 2.6726, 2.6803, 2.6812,
+    2.6899, 2.6916, 2.6931, 2.6998, 2.7062, 2.7262, 2.7222, 2.7158, 2.7041, 2.7485, 2.7491, 2.7451,
+    2.7485, 2.7233, 2.7297, 2.7233, 2.7145, 2.6958, 2.6788, 2.6439, 2.6007, 2.4786, 2.2469, 2.1877,
+    2.1392, 2.0717, 2.0107, 1.9676, 1.9140, 1.7102, 0.9101, 0.7164
+]
+
+
+class VAE(nn.Module):
+
+    def __init__(
+        self,
+        *,
+        data_dim: int,
+        embed_dim: int,
+        hidden_dim: int,
+    ):
+        super().__init__()
+
+        if data_dim == 80:
+            self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_80D, dtype=torch.float32))
+            self.data_std = nn.Buffer(torch.tensor(DATA_STD_80D, dtype=torch.float32))
+        elif data_dim == 128:
+            self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_128D, dtype=torch.float32))
+            self.data_std = nn.Buffer(torch.tensor(DATA_STD_128D, dtype=torch.float32))
+
+        self.data_mean = self.data_mean.view(1, -1, 1)
+        self.data_std = self.data_std.view(1, -1, 1)
+
+        self.encoder = Encoder1D(
+            dim=hidden_dim,
+            ch_mult=(1, 2, 4),
+            num_res_blocks=2,
+            attn_layers=[3],
+            down_layers=[0],
+            in_dim=data_dim,
+            embed_dim=embed_dim,
+        )
+        self.decoder = Decoder1D(
+            dim=hidden_dim,
+            ch_mult=(1, 2, 4),
+            num_res_blocks=2,
+            attn_layers=[3],
+            down_layers=[0],
+            in_dim=data_dim,
+            out_dim=data_dim,
+            embed_dim=embed_dim,
+        )
+
+        self.embed_dim = embed_dim
+        # self.quant_conv = nn.Conv1d(2 * embed_dim, 2 * embed_dim, 1)
+        # self.post_quant_conv = nn.Conv1d(embed_dim, embed_dim, 1)
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        pass
+
+    def encode(self, x: torch.Tensor, normalize: bool = True) -> DiagonalGaussianDistribution:
+        if normalize:
+            x = self.normalize(x)
+        moments = self.encoder(x)
+        posterior = DiagonalGaussianDistribution(moments)
+        return posterior
+
+    def decode(self, z: torch.Tensor, unnormalize: bool = True) -> torch.Tensor:
+        dec = self.decoder(z)
+        if unnormalize:
+            dec = self.unnormalize(dec)
+        return dec
+
+    def normalize(self, x: torch.Tensor) -> torch.Tensor:
+        return (x - self.data_mean) / self.data_std
+
+    def unnormalize(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.data_std + self.data_mean
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        sample_posterior: bool = True,
+        rng: Optional[torch.Generator] = None,
+        normalize: bool = True,
+        unnormalize: bool = True,
+    ) -> tuple[torch.Tensor, DiagonalGaussianDistribution]:
+
+        posterior = self.encode(x, normalize=normalize)
+        if sample_posterior:
+            z = posterior.sample(rng)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z, unnormalize=unnormalize)
+        return dec, posterior
+
+    def load_weights(self, src_dict) -> None:
+        self.load_state_dict(src_dict, strict=True)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def get_last_layer(self):
+        return self.decoder.conv_out.weight
+
+    def remove_weight_norm(self):
+        for name, m in self.named_modules():
+            if isinstance(m, MPConv1D):
+                m.remove_weight_norm()
+                log.debug(f"Removed weight norm from {name}")
+        return self
+
+
+class Encoder1D(nn.Module):
+
+    def __init__(self,
+                 *,
+                 dim: int,
+                 ch_mult: tuple[int] = (1, 2, 4, 8),
+                 num_res_blocks: int,
+                 attn_layers: list[int] = [],
+                 down_layers: list[int] = [],
+                 resamp_with_conv: bool = True,
+                 in_dim: int,
+                 embed_dim: int,
+                 double_z: bool = True,
+                 kernel_size: int = 3,
+                 clip_act: float = 256.0):
+        super().__init__()
+        self.dim = dim
+        self.num_layers = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.in_channels = in_dim
+        self.clip_act = clip_act
+        self.down_layers = down_layers
+        self.attn_layers = attn_layers
+        self.conv_in = MPConv1D(in_dim, self.dim, kernel_size=kernel_size)
+
+        in_ch_mult = (1, ) + tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        # downsampling
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_layers):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = dim * in_ch_mult[i_level]
+            block_out = dim * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(
+                    ResnetBlock1D(in_dim=block_in,
+                                  out_dim=block_out,
+                                  kernel_size=kernel_size,
+                                  use_norm=True))
+                block_in = block_out
+                if i_level in attn_layers:
+                    attn.append(AttnBlock1D(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level in down_layers:
+                down.downsample = Downsample1D(block_in, resamp_with_conv)
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock1D(in_dim=block_in,
+                                         out_dim=block_in,
+                                         kernel_size=kernel_size,
+                                         use_norm=True)
+        self.mid.attn_1 = AttnBlock1D(block_in)
+        self.mid.block_2 = ResnetBlock1D(in_dim=block_in,
+                                         out_dim=block_in,
+                                         kernel_size=kernel_size,
+                                         use_norm=True)
+
+        # end
+        self.conv_out = MPConv1D(block_in,
+                                 2 * embed_dim if double_z else embed_dim,
+                                 kernel_size=kernel_size)
+
+        self.learnable_gain = nn.Parameter(torch.zeros([]))
+
+    def forward(self, x):
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_layers):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                h = h.clamp(-self.clip_act, self.clip_act)
+                hs.append(h)
+            if i_level in self.down_layers:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        h = h.clamp(-self.clip_act, self.clip_act)
+
+        # end
+        h = nonlinearity(h)
+        h = self.conv_out(h, gain=(self.learnable_gain + 1))
+        return h
+
+
+class Decoder1D(nn.Module):
+
+    def __init__(self,
+                 *,
+                 dim: int,
+                 out_dim: int,
+                 ch_mult: tuple[int] = (1, 2, 4, 8),
+                 num_res_blocks: int,
+                 attn_layers: list[int] = [],
+                 down_layers: list[int] = [],
+                 kernel_size: int = 3,
+                 resamp_with_conv: bool = True,
+                 in_dim: int,
+                 embed_dim: int,
+                 clip_act: float = 256.0):
+        super().__init__()
+        self.ch = dim
+        self.num_layers = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.in_channels = in_dim
+        self.clip_act = clip_act
+        self.down_layers = [i + 1 for i in down_layers]  # each downlayer add one
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        block_in = dim * ch_mult[self.num_layers - 1]
+
+        # z to block_in
+        self.conv_in = MPConv1D(embed_dim, block_in, kernel_size=kernel_size)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True)
+        self.mid.attn_1 = AttnBlock1D(block_in)
+        self.mid.block_2 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_layers)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = dim * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock1D(in_dim=block_in, out_dim=block_out, use_norm=True))
+                block_in = block_out
+                if i_level in attn_layers:
+                    attn.append(AttnBlock1D(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level in self.down_layers:
+                up.upsample = Upsample1D(block_in, resamp_with_conv)
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.conv_out = MPConv1D(block_in, out_dim, kernel_size=kernel_size)
+        self.learnable_gain = nn.Parameter(torch.zeros([]))
+
+    def forward(self, z):
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        h = h.clamp(-self.clip_act, self.clip_act)
+
+        # upsampling
+        for i_level in reversed(range(self.num_layers)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+                h = h.clamp(-self.clip_act, self.clip_act)
+            if i_level in self.down_layers:
+                h = self.up[i_level].upsample(h)
+
+        h = nonlinearity(h)
+        h = self.conv_out(h, gain=(self.learnable_gain + 1))
+        return h
+
+
+def VAE_16k(**kwargs) -> VAE:
+    return VAE(data_dim=80, embed_dim=20, hidden_dim=384, **kwargs)
+
+
+def VAE_44k(**kwargs) -> VAE:
+    return VAE(data_dim=128, embed_dim=40, hidden_dim=512, **kwargs)
+
+
+def get_my_vae(name: str, **kwargs) -> VAE:
+    if name == '16k':
+        return VAE_16k(**kwargs)
+    if name == '44k':
+        return VAE_44k(**kwargs)
+    raise ValueError(f'Unknown model: {name}')
+
+
+if __name__ == '__main__':
+    network = get_my_vae('standard')
+
+    # print the number of parameters in terms of millions
+    num_params = sum(p.numel() for p in network.parameters()) / 1e6
+    print(f'Number of parameters: {num_params:.2f}M')