Add WhAM Gradio app with ZeroGPU support

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Project CETI
+
+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

@@ -1,13 +1,211 @@
 ---
-title: Wham
-emoji: 📊
-colorFrom: red
-colorTo: pink
+title: WhAM
+emoji: 🐋
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
-sdk_version: 6.8.0
-python_version: '3.12'
 app_file: app.py
 pinned: false
+hardware: zero-a10g
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# WhAM: a Whale Acoustics Model
+[![arXiv](https://img.shields.io/badge/arXiv-2512.02206-b31b1b.svg)](https://arxiv.org/abs/2512.02206)
+[![Model Weights](https://img.shields.io/badge/Zenodo-Model%20Weights-blue.svg)](https://doi.org/10.5281/zenodo.17633708)
+[![Hugging Face Dataset](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-DSWP%20Dataset-yellow)](https://huggingface.co/datasets/orrp/DSWP)
+![WhAM](assets/inference.png "WhAM")
+WhAM is a transformer-based audio-to-audio model designed to synthesize and analyze sperm whale codas. Based on [VampNet](https://github.com/hugofloresgarcia/vampnet), WhAM uses masked acoustic token modeling to capture temporal and spectral features of whale communication. WhAM generates codas from a given audio context, enabling three core capabilities:
+
+ - Acoustic Translation: The ability to style-transfer arbitrary audio prompts (e.g., human speech, noise) into the acoustic texture of sperm whale codas.
+
+ - Synthesizing novel "pseudocodas".
+
+ - Providing audio embeddings for downstream tasks such as social unit and spectral feature ("vowel") classification.
+
+See our [NeurIPS 2025](https://openreview.net/pdf?id=IL1wvzOgqD) publication for more details.
+
+## Installation
+
+1.  **Clone the repository:**
+    ```bash
+    git clone https://github.com/Project-CETI/wham.git
+    cd wham
+    ```
+
+2.  **Set up the environment:**
+    ```bash
+    conda create -n wham python=3.9
+    conda activate wham
+    ```
+
+3.  **Install dependencies:**
+    ```bash
+    # Install the wham package
+    pip install -e .
+
+    # Install VampNet
+    pip install -e ./vampnet
+
+    # Install madmom
+    pip install --no-build-isolation madmom
+
+    # Install ffmpeg
+    conda install -c conda-forge ffmpeg
+    ```
+
+4.  **Download model weights:**
+    Download the [weights](https://zenodo.org/records/17633708) and extract to `vampnet/models/`.
+
+## Generation
+
+To run WhAM locally and prompt it in your browser:
+
+```bash
+python vampnet/app.py --args.load conf/interface.yml --Interface.device cuda
+```
+
+This will provide you with a Gradio link to test WhAM on inputs of your choice.
+
+## Training Data
+
+![Training](assets/training.png "Training")
+
+You only need to follow these to fine-tune your own version of WhAM. First, obtain the original VampNet weights by following the instructions in the ![original repo](https://github.com/hugofloresgarcia/vampnet/tree/ismir-2023). Download 
+c2f.pth and codec.pth and replace the weights you previously downloaded in `vampnet/models`.
+
+Second, obtain data:
+
+1.  **Domain adaptation data:**
+   
+    - Download audio samples from the [WMMS 'Best Of' Cut](https://whoicf2.whoi.edu/science/B/whalesounds/index.cfm). Save them under `vampnet/training_data/domain_adaptation`.
+    
+    - Download audio samples from the [BirdSet Dataset](https://huggingface.co/datasets/DBD-research-group/BirdSet). Save these under the same directory
+
+    - Finally, download all samples from the [AudioSet Dataset](https://research.google.com/audioset/ontology/index.html) with the label `Animal` and once again save these into the directory
+
+3.  **Species-specific finetuning:** Finetuning can be performed on the openly available **[Dominica Sperm Whale Project (DSWP)](https://huggingface.co/datasets/orrp/DSWP)** dataset, available on Hugging Face.
+
+
+With data in hand, navigate into `vampnet` and perform Domain Adaptation:
+```bash
+python vampnet/scripts/exp/fine_tune.py "training_data/domain_adaptation" domain_adapted && python vampnet/scripts/exp/train.py --args.load conf/generated/domain_adapted/coarse.yml && python vampnet/scripts/exp/train.py --args.load conf/generated/domain_adapted/c2f.yml
+```
+
+Then fine-tune the domain-adapted model. Create the config file with the command:
+
+```bash
+python vampnet/scripts/exp/fine_tune.py "training_data/species_specific_finetuning" fine-tuned
+```
+
+To select which weights you want to use as a checkpoint, change `fine_tune_checkpoint` in `conf/generated/fine-tuned/[c2f/coarse].yml` to `./runs/domain_adaptation/[coarse/c2f]/[checkpoint]/vampnets/weights.pth`. `[checkpoint]` can be `latest` in order to use the last saved checkpoint from the previous run, though it is recommended to manually verify the quality of generations over various checkpoints as overtraining can often cause degradation in audio quality, especially with smaller datasets. After making that change, run the command:
+
+```bash
+python vampnet/scripts/exp/train.py --args.load conf/generated/fine-tuned/coarse.yml && python vampnet/scripts/exp/train.py --args.load conf/generated/fine-tuned/c2f.yml
+```
+
+After following these steps, you should be able to generate audio via the browser by running:
+```bash
+python app.py --args.load vampnet/conf/generated/fine-tuned/interface.yml
+```
+
+**Note**: The coarse and fine weights can be trained separately if compute allows. In this case, you would call the two scripts:
+
+```bash
+python vampnet/scripts/exp/train.py --args.load conf/generated/[fine-tuned/domain_adaptated]/coarse.yml
+```
+
+```bash
+python vampnet/scripts/exp/train.py --args.load conf/generated/[fine-tuned/domain_adaptated]/c2f.yml
+```
+
+After both are finished running, ensure that both resulting weights are copied into the same copy of WhAM.
+
+
+
+## Testing Data
+
+1.  **Marine Mammel Data:**
+    Download audio samples from the [WMMS 'Best Of' Cut](https://whoicf2.whoi.edu/science/B/whalesounds/index.cfm). Save them under `data/testing_data/marine_mammals/data/[SPECIES_NAME]`.
+    * `[SPECIES_NAME]` must match the species names found in `wham/generation/prompt_configs.py`.
+      
+2.  **Sperm Whale Codas:**
+    To evaluate on sperm whale codas, you can use the openly available [DSWP](https://huggingface.co/datasets/orrp/DSWP) dataset.
+
+3. Generate artifical beeps for experiments. `data/generate_beeps.sh`
+
+
+## Reproducing Paper Results
+Note: Access to the DSWP+CETI annotated is required to reproduce all results; as of time of publication, only part of this data is publicly available. Still, we include the following code as it may be useful for researchers who may benefit from our evaluation pipeline.
+
+### 1. Downstream Classification Tasks
+To reproduce **Table 1** (Classification Accuracies) and **Figure 7** (Ablation Study):
+
+**Table 1 Results:**
+```bash
+cd wham/embedding
+./downstream_tasks.sh
+```
+* Runs all downstream classification tasks.
+* **Baselines:** Run once.
+* **Models (AVES, VampNet):** Run over 3 random seeds; reports mean and standard deviation.
+
+**Figure 7 Results (Ablation):**
+```bash
+cd wham/embedding
+./downstream_ablation.sh
+```
+* Outputs accuracy scores for ablation variants (averaged across 3 seeds with error bars).
+
+### 2. Generative Metrics
+
+**Figure 12: Frechet Audio Distance (FAD) Scores**
+Calculate the distance between WhAM's generated results and real codas:
+```bash
+# Calculate for all species
+bash wham/generation/eval/calculate_FAD.sh
+
+# Calculate for a single species
+bash wham/generation/eval/calculate_FAD.sh [species_name]
+```
+* *Runtime:* ~3 hours on an NVIDIA A10 GPU.
+
+**Figure 3: FAD with Custom/BirdNET Embeddings**
+To compare against other embeddings:
+1.  Convert your `.wav` files to `.npy` embeddings.
+2.  Place raw coda embeddings in: `data/testing_data/coda_embeddings`
+3.  Place comparison embeddings in subfolders within: `data/testing_data/comparison_embeddings`
+4.  Run:
+    ```bash
+    python wham/generation/eval/calculate_custom_fad.py
+    ```
+*For BirdNET embeddings, refer to the [official repo](https://github.com/BirdNET-Team/BirdNET-Analyzer).*
+
+**Table 2: Embedding Type Ablation**
+Calculate distances between raw codas, denoised versions, and noise profiles:
+```bash
+bash wham/generation/eval/FAD_ablation.sh
+```
+* *Prerequisites:* Ensure `data/testing_data/ablation/noise` and `data/testing_data/ablation/denoised` are populated.
+* *Runtime:* ~1.5 hours on an NVIDIA A10 GPU.
+
+**Figure 13: Tokenizer Reconstruction**
+Test the mean squared reconstruction error:
+```bash
+bash wham/generation/eval/evaluate_tokenizer.sh
+```
+
+---
+
+## Citation
+
+Please use the following citation if you use this code, model or data.
+
+```bibtex
+@inproceedings{wham2025,
+  title={Towards A Translative Model of Sperm Whale Vocalization},
+  author={Orr Paradise, Pranav Muralikrishnan, Liangyuan Chen, Hugo Flores Garcia, Bryan Pardo, Roee Diamant, David F. Gruber, Shane Gero, Shafi Goldwasser},
+  booktitle={Advances in Neural Information Processing Systems 39: Annual Conference
+                  on Neural Information Processing Systems 2025, NeurIPS 2025, San Diego, CA, USA},
+  year={2025}
+}
+```

app.py

@@ -0,0 +1,38 @@
+import os
+import sys
+
+from huggingface_hub import hf_hub_download
+
+REPO_ROOT = os.path.dirname(os.path.abspath(__file__))
+MODEL_DIR = os.path.join(REPO_ROOT, "vampnet", "models")
+os.makedirs(MODEL_DIR, exist_ok=True)
+
+MODEL_REPO = "anvitax/wham-weights"
+WEIGHT_FILES = ["coarse.pth", "c2f.pth", "codec.pth", "wavebeat.pth"]
+
+for fname in WEIGHT_FILES:
+    target = os.path.join(MODEL_DIR, fname)
+    if not os.path.exists(target):
+        print(f"Downloading {fname} from {MODEL_REPO}...")
+        hf_hub_download(repo_id=MODEL_REPO, filename=fname, local_dir=MODEL_DIR)
+    else:
+        print(f"Found {fname}")
+
+sys.path.insert(0, os.path.join(REPO_ROOT, "vampnet"))
+
+os.chdir(os.path.join(REPO_ROOT, "vampnet"))
+
+try:
+    import spaces
+    device = "cpu"
+except ImportError:
+    import torch
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+sys.argv = [
+    "app.py",
+    "--args.load", "conf/interface.yml",
+    "--Interface.device", device,
+]
+
+exec(open("app.py").read())

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -0,0 +1,16 @@
+torch
+gradio
+argbind>=0.3.2
+numpy<1.24
+pydantic==2.10.6
+huggingface_hub
+loralib
+torch_pitch_shift
+soundfile
+pydub
+tqdm
+Cython
+wavebeat @ git+https://github.com/hugofloresgarcia/wavebeat
+lac @ git+https://github.com/hugofloresgarcia/lac.git
+descript-audiotools @ git+https://github.com/hugofloresgarcia/audiotools.git
+pyharp

setup.py

@@ -0,0 +1,33 @@
+from setuptools import setup, find_packages
+
+with open("README.md") as f:
+    long_description = f.read()
+
+setup(
+    name="wham",
+    version="0.0.1",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/orrp/wam",
+    license="MIT",
+    packages=find_packages(),
+    package_dir={},
+    install_requires=[
+        "descript-audiotools @ git+https://github.com/hugofloresgarcia/audiotools.git",
+        "argbind",
+        "pandas",
+        "pathlib",
+        "pydub",
+        "ffmpeg-python",
+        "tqdm",
+        "scikit-learn",
+        "wandb",
+        "gdown",  # For fetching large files from Google Drive
+        "soundfile",
+        "transformers",
+        "torch",
+        "Cython",
+        "fadtk",
+        "urllib3==2.0"
+    ],
+)

vampnet/.gitignore

@@ -0,0 +1,189 @@
+# 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/env.sh
+venv/
+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/
+
+# Files created by experiments
+output/
+snapshot/
+*.m4a
+notebooks/scratch.ipynb
+notebooks/inspect.ipynb
+notebooks/effects.ipynb
+notebooks/*.ipynb
+notebooks/*.gif
+notebooks/*.wav
+notebooks/*.mp4
+*runs/
+boards/
+samples/
+*.ipynb
+
+results.json
+metrics.csv
+mprofile_*
+mem.png
+
+results/
+mprofile*
+*.png
+# do not ignore the test wav file
+!tests/audio/short_test_audio.wav
+!tests/audio/output.wav
+*/.DS_Store
+.DS_Store
+env.sh
+_codebraid/
+**/*.html
+**/*.exec.md
+flagged/
+log.txt
+ckpt/
+.syncthing*
+tests/assets/
+archived/
+
+scratch/
+
+runs-archive
+lyrebird-audiotools
+lyrebird-audio-codec
+samples-*/**
+
+gradio-outputs/
+models/
+samples*/
+models-all/
+models.zip
+.git-old
+conf/generated/*
+runs*/
+
+
+gtzan.zip
+.gtzan_emb_cache
+runs

vampnet/.pre-commit-config.yaml

@@ -0,0 +1,15 @@
+repos:
+- repo: https://github.com/asottile/reorder_python_imports
+  rev: v2.5.0
+  hooks:
+    - id: reorder-python-imports
+- repo: https://github.com/psf/black
+  rev: 23.1.0
+  hooks:
+    - id: black
+      language_version: python3
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.0.1
+  hooks:
+    - id: end-of-file-fixer
+    - id: trailing-whitespace

vampnet/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Hugo Flores García and Prem Seetharaman
+
+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.

vampnet/README.md

@@ -0,0 +1,94 @@
+# VampNet
+
+This repository contains recipes for training generative music models on top of the Descript Audio Codec.
+
+## try `unloop`
+you can try vampnet in a co-creative looper called unloop. see this link: https://github.com/hugofloresgarcia/unloop
+
+# Setting up
+
+**Requires Python 3.9**. 
+
+you'll need a Python 3.9 environment to run VampNet. This is due to a [known issue with madmom](https://github.com/hugofloresgarcia/vampnet/issues/15). 
+
+(for example, using conda)
+```bash
+conda create -n vampnet python=3.9
+conda activate vampnet
+```
+
+
+install VampNet
+
+```bash
+git clone https://github.com/hugofloresgarcia/vampnet.git
+pip install -e ./vampnet
+```
+
+## A note on argbind
+This repository relies on [argbind](https://github.com/pseeth/argbind) to manage CLIs and config files. 
+Config files are stored in the `conf/` folder. 
+
+## Getting the Pretrained Models
+
+### Licensing for Pretrained Models: 
+The weights for the models are licensed [`CC BY-NC-SA 4.0`](https://creativecommons.org/licenses/by-nc-sa/4.0/deed.ml). Likewise, any VampNet models fine-tuned on the pretrained models are also licensed [`CC BY-NC-SA 4.0`](https://creativecommons.org/licenses/by-nc-sa/4.0/deed.ml).
+
+Download the pretrained models from [this link](https://zenodo.org/record/8136629). Then, extract the models to the `models/` folder. 
+
+
+# Usage
+
+## Launching the Gradio Interface
+You can launch a gradio UI to play with vampnet. 
+
+```bash
+python app.py --args.load conf/interface.yml --Interface.device cuda
+```
+
+# Training / Fine-tuning 
+
+## Training a model
+
+To train a model, run the following script: 
+
+```bash
+python scripts/exp/train.py --args.load conf/vampnet.yml --save_path /path/to/checkpoints
+```
+
+You can edit `conf/vampnet.yml` to change the dataset paths or any training hyperparameters. 
+
+For coarse2fine models, you can use `conf/c2f.yml` as a starting configuration. 
+
+See `python scripts/exp/train.py -h` for a list of options.
+
+## Fine-tuning
+To fine-tune a model, use the script in `scripts/exp/fine_tune.py` to generate 3 configuration files: `c2f.yml`, `coarse.yml`, and `interface.yml`. 
+The first two are used to fine-tune the coarse and fine models, respectively. The last one is used to launch the gradio interface.
+
+```bash
+python scripts/exp/fine_tune.py "/path/to/audio1.mp3 /path/to/audio2/ /path/to/audio3.wav" <fine_tune_name>
+```
+
+This will create a folder under `conf/<fine_tune_name>/` with the 3 configuration files.
+
+The save_paths will be set to `runs/<fine_tune_name>/coarse` and `runs/<fine_tune_name>/c2f`. 
+
+launch the coarse job: 
+```bash
+python scripts/exp/train.py --args.load conf/<fine_tune_name>/coarse.yml 
+```
+
+this will save the coarse model to `runs/<fine_tune_name>/coarse/ckpt/best/`.
+
+launch the c2f job: 
+```bash
+python  scripts/exp/train.py --args.load conf/<fine_tune_name>/c2f.yml 
+```
+
+launch the interface: 
+```bash
+python  app.py --args.load conf/generated/<fine_tune_name>/interface.yml 
+```
+
+

vampnet/app.py

@@ -0,0 +1,677 @@
+from pathlib import Path
+from typing import Tuple
+import yaml
+import tempfile
+import uuid
+from dataclasses import dataclass, asdict
+
+import numpy as np
+import audiotools as at
+import argbind
+
+import gradio as gr
+from vampnet.interface import Interface
+from vampnet import mask as pmask
+
+try:
+    import spaces
+    ZERO_GPU = True
+except ImportError:
+    ZERO_GPU = False
+
+def gpu(fn):
+    if ZERO_GPU:
+        return spaces.GPU(fn)
+    return fn
+
+Interface = argbind.bind(Interface)
+# AudioLoader = argbind.bind(at.data.datasets.AudioLoader)
+
+conf = argbind.parse_args()
+
+
+from torch_pitch_shift import pitch_shift, get_fast_shifts
+def shift_pitch(signal, interval: int):
+    signal.samples = pitch_shift(
+        signal.samples, 
+        shift=interval, 
+        sample_rate=signal.sample_rate
+    )
+    return signal
+
+def load_interface():
+    with argbind.scope(conf):
+        interface = Interface()
+        # loader = AudioLoader()
+        print(f"interface device is {interface.device}")
+        return interface
+
+
+
+
+interface = load_interface()
+
+
+OUT_DIR = Path("gradio-outputs")
+OUT_DIR.mkdir(exist_ok=True, parents=True)
+
+
+def load_audio(file):
+    print(file)
+    filepath = file.name
+    sig = at.AudioSignal.salient_excerpt(
+        filepath, 
+        duration=interface.coarse.chunk_size_s
+    )
+    sig = interface.preprocess(sig)
+
+    out_dir = OUT_DIR / "tmp" / str(uuid.uuid4())
+    out_dir.mkdir(parents=True, exist_ok=True)
+    sig.write(out_dir / "input.wav")
+    return sig.path_to_file
+
+
+def load_example_audio():
+    return "./assets/example.wav"
+
+
+@gpu
+def _vamp(data, return_mask=False):
+    interface.to("cuda")
+
+    out_dir = OUT_DIR / str(uuid.uuid4())
+    out_dir.mkdir()
+    sig = at.AudioSignal(data[input_audio])
+    sig = interface.preprocess(sig)
+
+    loudness = sig.loudness()
+    print(f"input loudness is {loudness}")
+
+    if data[pitch_shift_amt] != 0:
+        sig = shift_pitch(sig, data[pitch_shift_amt])
+
+    z = interface.encode(sig)
+
+    ncc = data[n_conditioning_codebooks]
+
+    # build the mask
+    mask = pmask.linear_random(z, data[rand_mask_intensity])
+    mask = pmask.mask_and(
+        mask, pmask.inpaint(
+            z,
+            interface.s2t(data[prefix_s]),
+            interface.s2t(data[suffix_s])
+        )
+    )
+    mask = pmask.mask_and(
+        mask, pmask.periodic_mask(
+            z,
+            data[periodic_p],
+            data[periodic_w],
+            random_roll=True
+        )
+    )
+    if data[onset_mask_width] > 0:
+        mask = pmask.mask_or(
+            mask, pmask.onset_mask(sig, z, interface, width=data[onset_mask_width])
+        )
+    if data[beat_mask_width] > 0:
+        beat_mask = interface.make_beat_mask(
+            sig,
+            after_beat_s=(data[beat_mask_width]/1000), 
+            mask_upbeats=not data[beat_mask_downbeats],
+        )
+        mask = pmask.mask_and(mask, beat_mask)
+
+    # these should be the last two mask ops
+    mask = pmask.dropout(mask, data[dropout])
+    mask = pmask.codebook_unmask(mask, ncc)
+    mask = pmask.codebook_mask(mask, int(data[n_mask_codebooks]))
+
+
+
+    print(f"dropout {data[dropout]}")
+    print(f"masktemp {data[masktemp]}")
+    print(f"sampletemp {data[sampletemp]}")
+    print(f"top_p {data[top_p]}")
+    print(f"prefix_s {data[prefix_s]}")
+    print(f"suffix_s {data[suffix_s]}")
+    print(f"rand_mask_intensity {data[rand_mask_intensity]}")
+    print(f"num_steps {data[num_steps]}")
+    print(f"periodic_p {data[periodic_p]}")
+    print(f"periodic_w {data[periodic_w]}")
+    print(f"n_conditioning_codebooks {data[n_conditioning_codebooks]}")
+    print(f"use_coarse2fine {data[use_coarse2fine]}")
+    print(f"onset_mask_width {data[onset_mask_width]}")
+    print(f"beat_mask_width {data[beat_mask_width]}")
+    print(f"beat_mask_downbeats {data[beat_mask_downbeats]}")
+    print(f"stretch_factor {data[stretch_factor]}")
+    print(f"seed {data[seed]}")
+    print(f"pitch_shift_amt {data[pitch_shift_amt]}")
+    print(f"sample_cutoff {data[sample_cutoff]}")
+    
+    
+    _top_p = data[top_p] if data[top_p] > 0 else None
+    # save the mask as a txt file
+    np.savetxt(out_dir / "mask.txt", mask[:,0,:].long().cpu().numpy())
+
+    _seed = data[seed] if data[seed] > 0 else None
+    zv, mask_z = interface.coarse_vamp(
+        z, 
+        mask=mask,
+        sampling_steps=data[num_steps],
+        mask_temperature=data[masktemp]*10,
+        sampling_temperature=data[sampletemp],
+        return_mask=True, 
+        typical_filtering=data[typical_filtering], 
+        typical_mass=data[typical_mass], 
+        typical_min_tokens=data[typical_min_tokens], 
+        top_p=_top_p,
+        gen_fn=interface.coarse.generate,
+        seed=_seed,
+        sample_cutoff=data[sample_cutoff],
+    )
+
+    if use_coarse2fine: 
+        zv = interface.coarse_to_fine(
+            zv, 
+            mask_temperature=data[masktemp]*10, 
+            sampling_temperature=data[sampletemp],
+            mask=mask,
+            sampling_steps=data[num_steps],
+            sample_cutoff=data[sample_cutoff], 
+            seed=_seed,
+        )
+
+    sig = interface.to_signal(zv).cpu()
+    print("done")
+
+    print(f"output loudness is {sig.loudness()}")
+    sig = sig.normalize(loudness)    
+    print(f"normalized loudness is {sig.loudness()}")
+
+    sig.write(out_dir / "output.wav")
+
+    if return_mask:
+        mask = interface.to_signal(mask_z).cpu()
+        mask.write(out_dir / "mask.wav")
+        return sig.path_to_file, mask.path_to_file
+    else:
+        return sig.path_to_file
+
+def vamp(data):
+    return _vamp(data, return_mask=True)
+
+def api_vamp(data):
+    return _vamp(data, return_mask=False)
+        
+def save_vamp(data):
+    out_dir = OUT_DIR / "saved" / str(uuid.uuid4())
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    sig_in = at.AudioSignal(data[input_audio])
+    sig_out = at.AudioSignal(data[output_audio])
+
+    sig_in.write(out_dir / "input.wav")
+    sig_out.write(out_dir / "output.wav")
+    
+    _data = {
+        "masktemp": data[masktemp],
+        "sampletemp": data[sampletemp],
+        "top_p": data[top_p],
+        "prefix_s": data[prefix_s],
+        "suffix_s": data[suffix_s],
+        "rand_mask_intensity": data[rand_mask_intensity],
+        "num_steps": data[num_steps],
+        "notes": data[notes_text],
+        "periodic_period": data[periodic_p],
+        "periodic_width": data[periodic_w],
+        "n_conditioning_codebooks": data[n_conditioning_codebooks], 
+        "use_coarse2fine": data[use_coarse2fine],
+        "stretch_factor": data[stretch_factor],
+        "seed": data[seed],
+        "samplecutoff": data[sample_cutoff],
+    }
+
+    # save with yaml
+    with open(out_dir / "data.yaml", "w") as f:
+        yaml.dump(_data, f)
+
+    import zipfile
+    zip_path = str(out_dir.with_suffix(".zip"))
+    with zipfile.ZipFile(zip_path, "w") as zf:
+        for file in out_dir.iterdir():
+            zf.write(file, file.name)
+
+    return f"saved! your save code is {out_dir.stem}", zip_path
+
+
+@gpu
+def harp_vamp(_input_audio, _beat_mask_width, _sampletemp):
+    interface.to("cuda")
+
+    out_dir = OUT_DIR / str(uuid.uuid4())
+    out_dir.mkdir()
+    sig = at.AudioSignal(_input_audio)
+    sig = interface.preprocess(sig)
+
+    z = interface.encode(sig)
+
+    # build the mask
+    mask = pmask.linear_random(z, 1.0)
+    if _beat_mask_width > 0:
+        beat_mask = interface.make_beat_mask(
+            sig,
+            after_beat_s=(_beat_mask_width/1000), 
+        )
+        mask = pmask.mask_and(mask, beat_mask)
+
+    # save the mask as a txt file
+    zv, mask_z = interface.coarse_vamp(
+        z, 
+        mask=mask,
+        sampling_temperature=_sampletemp,
+        return_mask=True, 
+        gen_fn=interface.coarse.generate,
+    )
+
+
+    zv = interface.coarse_to_fine(
+        zv, 
+        sampling_temperature=_sampletemp,
+        mask=mask,
+    )
+
+    sig = interface.to_signal(zv).cpu()
+    print("done")
+
+    sig.write(out_dir / "output.wav")
+
+    return sig.path_to_file
+
+with gr.Blocks() as demo:
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("# VampNet Audio Vamping")
+            gr.Markdown("""## Description:
+            This is a demo of the VampNet, a generative audio model that transforms the input audio based on the chosen settings. 
+            You can control the extent and nature of variation with a set of manual controls and presets. 
+            Use this interface to experiment with different mask settings and explore the audio outputs.
+            """)
+
+            gr.Markdown("""
+            ## Instructions:
+            1. You can start by uploading some audio, or by loading the example audio. 
+            2. Choose a preset for the vamp operation, or manually adjust the controls to customize the mask settings. 
+            3. Click the "generate (vamp)!!!" button to apply the vamp operation. Listen to the output audio.
+            4. Optionally, you can add some notes and save the result. 
+            5. You can also use the output as the new input and continue experimenting!
+            """)
+    with gr.Row():
+        with gr.Column():
+
+
+            manual_audio_upload = gr.File(
+                label=f"upload some audio (will be randomly trimmed to max of {interface.coarse.chunk_size_s:.2f}s)",
+                file_types=["audio"]
+            )
+            load_example_audio_button = gr.Button("or load example audio")
+
+            input_audio = gr.Audio(
+                label="input audio",
+                interactive=False, 
+                type="filepath",
+            )
+
+            audio_mask = gr.Audio(
+                label="audio mask (listen to this to hear the mask hints)",
+                interactive=False, 
+                type="filepath",
+            )
+
+            # connect widgets
+            load_example_audio_button.click(
+                fn=load_example_audio,
+                inputs=[],
+                outputs=[ input_audio]
+            )
+
+            manual_audio_upload.change(
+                fn=load_audio,
+                inputs=[manual_audio_upload],
+                outputs=[ input_audio]
+            )
+                
+        # mask settings
+        with gr.Column():
+
+
+            presets = {
+                    "unconditional": {
+                        "periodic_p": 0,
+                        "onset_mask_width": 0,
+                        "beat_mask_width": 0,
+                        "beat_mask_downbeats": False,
+                    }, 
+                    "slight periodic variation": {
+                        "periodic_p": 5,
+                        "onset_mask_width": 5,
+                        "beat_mask_width": 0,
+                        "beat_mask_downbeats": False,
+                    },
+                    "moderate periodic variation": {
+                        "periodic_p": 13,
+                        "onset_mask_width": 5,
+                        "beat_mask_width": 0,
+                        "beat_mask_downbeats": False,
+                    },
+                    "strong periodic variation": {
+                        "periodic_p": 17,
+                        "onset_mask_width": 5,
+                        "beat_mask_width": 0,
+                        "beat_mask_downbeats": False,
+                    },
+                    "very strong periodic variation": {
+                        "periodic_p": 21,
+                        "onset_mask_width": 5,
+                        "beat_mask_width": 0,
+                        "beat_mask_downbeats": False,
+                    },
+                    "beat-driven variation": {
+                        "periodic_p": 0,
+                        "onset_mask_width": 0,
+                        "beat_mask_width": 50,
+                        "beat_mask_downbeats": False,
+                    },
+                    "beat-driven variation (downbeats only)": {
+                        "periodic_p": 0,
+                        "onset_mask_width": 0,
+                        "beat_mask_width": 50,
+                        "beat_mask_downbeats": True,
+                    },
+                    "beat-driven variation (downbeats only, strong)": {
+                        "periodic_p": 0,
+                        "onset_mask_width": 0,
+                        "beat_mask_width": 20,
+                        "beat_mask_downbeats": True,
+                    },
+                }
+
+            preset = gr.Dropdown(
+                label="preset", 
+                choices=list(presets.keys()),
+                value="strong periodic variation",
+            )
+            load_preset_button = gr.Button("load_preset")
+
+            with gr.Accordion("manual controls", open=True):
+                periodic_p = gr.Slider(
+                    label="periodic prompt  (0 - unconditional, 2 - lots of hints, 8 - a couple of hints, 16 - occasional hint, 32 - very occasional hint, etc)",
+                    minimum=0,
+                    maximum=128, 
+                    step=1,
+                    value=3, 
+                )
+
+
+                onset_mask_width = gr.Slider(
+                    label="onset mask width (multiplies with the periodic mask, 1 step ~= 10milliseconds) ",
+                    minimum=0,
+                    maximum=100,
+                    step=1,
+                    value=5,
+                )
+
+                beat_mask_width = gr.Slider(
+                    label="beat prompt (ms)",
+                    minimum=0,
+                    maximum=200,
+                    value=0,
+                )
+                beat_mask_downbeats = gr.Checkbox(
+                    label="beat mask downbeats only?", 
+                    value=False
+                )
+
+                n_mask_codebooks = gr.Number(
+                    label="first upper codebook level to mask",
+                    value=9,
+                )
+
+
+                with gr.Accordion("extras ", open=False):
+                    pitch_shift_amt = gr.Slider(
+                        label="pitch shift amount (semitones)",
+                        minimum=-12,
+                        maximum=12,
+                        step=1,
+                        value=0,
+                    )
+
+                    rand_mask_intensity = gr.Slider(
+                        label="random mask intensity. (If this is less than 1, scatters prompts throughout the audio, should be between 0.9 and 1.0)",
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=1.0
+                    )
+
+                    periodic_w = gr.Slider(
+                        label="periodic prompt width (steps, 1 step ~= 10milliseconds)",
+                        minimum=1,
+                        maximum=20,
+                        step=1,
+                        value=1,
+                    )
+                    n_conditioning_codebooks = gr.Number(
+                        label="number of conditioning codebooks. probably 0", 
+                        value=0,
+                        precision=0,
+                    )
+
+                    stretch_factor = gr.Slider(
+                        label="time stretch factor",
+                        minimum=0,
+                        maximum=64, 
+                        step=1,
+                        value=1, 
+                    )
+
+            preset_outputs = {
+                periodic_p, 
+                onset_mask_width, 
+                beat_mask_width,
+                beat_mask_downbeats,
+            }
+
+            def load_preset(_preset):
+                return tuple(presets[_preset].values())
+
+            load_preset_button.click(
+                fn=load_preset,
+                inputs=[preset],
+                outputs=preset_outputs
+            )
+
+
+            with gr.Accordion("prefix/suffix prompts", open=False):
+                prefix_s = gr.Slider(
+                    label="prefix hint length (seconds)",
+                    minimum=0.0,
+                    maximum=10.0,
+                    value=0.0
+                )
+                suffix_s = gr.Slider(
+                    label="suffix hint length (seconds)",
+                    minimum=0.0,
+                    maximum=10.0,
+                    value=0.0
+                )
+
+            masktemp = gr.Slider(
+                label="mask temperature",
+                minimum=0.0,
+                maximum=100.0,
+                value=1.5
+            )
+            sampletemp = gr.Slider(
+                label="sample temperature",
+                minimum=0.1,
+                maximum=10.0,
+                value=1.0, 
+                step=0.001
+            )
+        
+
+
+            with gr.Accordion("sampling settings", open=False):
+                top_p = gr.Slider(
+                    label="top p (0.0 = off)",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.0
+                )
+                typical_filtering = gr.Checkbox(
+                    label="typical filtering ",
+                    value=False
+                )
+                typical_mass = gr.Slider( 
+                    label="typical mass (should probably stay between 0.1 and 0.5)",
+                    minimum=0.01,
+                    maximum=0.99,
+                    value=0.15
+                )
+                typical_min_tokens = gr.Slider(
+                    label="typical min tokens (should probably stay between 1 and 256)",
+                    minimum=1,
+                    maximum=256,
+                    step=1,
+                    value=64
+                )
+                sample_cutoff = gr.Slider(
+                    label="sample cutoff",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.5, 
+                    step=0.01
+                )
+
+            use_coarse2fine = gr.Checkbox(
+                label="use coarse2fine",
+                value=True, 
+                visible=False
+            )
+
+            num_steps = gr.Slider(
+                label="number of steps (should normally be between 12 and 36)",
+                minimum=1,
+                maximum=128,
+                step=1,
+                value=36
+            )
+
+            dropout = gr.Slider(
+                label="mask dropout",
+                minimum=0.0,
+                maximum=1.0,
+                step=0.01,
+                value=0.0
+            )
+
+
+            seed = gr.Number(
+                label="seed (0 for random)",
+                value=0,
+                precision=0,
+            )
+
+
+
+        # mask settings
+        with gr.Column():
+
+            # lora_choice = gr.Dropdown(
+            #     label="lora choice", 
+            #     choices=list(loras.keys()),
+            #     value=LORA_NONE, 
+            #     visible=False
+            # )
+
+            vamp_button = gr.Button("generate (vamp)!!!")
+            output_audio = gr.Audio(
+                label="output audio",
+                interactive=False,
+                type="filepath"
+            )
+
+            notes_text = gr.Textbox(
+                label="type any notes about the generated audio here", 
+                value="",
+                interactive=True
+            )
+            save_button = gr.Button("save vamp")
+            download_file = gr.File(
+                label="vamp to download will appear here",
+                interactive=False
+            )
+            use_as_input_button = gr.Button("use output as input")
+            
+            thank_you = gr.Markdown("")
+
+
+    _inputs = {
+            input_audio, 
+            num_steps,
+            masktemp,
+            sampletemp,
+            top_p,
+            prefix_s, suffix_s, 
+            rand_mask_intensity, 
+            periodic_p, periodic_w,
+            n_conditioning_codebooks, 
+            dropout,
+            use_coarse2fine, 
+            stretch_factor, 
+            onset_mask_width, 
+            typical_filtering,
+            typical_mass,
+            typical_min_tokens,
+            beat_mask_width,
+            beat_mask_downbeats,
+            seed, 
+            # lora_choice,
+            n_mask_codebooks,
+            pitch_shift_amt, 
+            sample_cutoff
+        }
+  
+    # connect widgets
+    vamp_button.click(
+        fn=vamp,
+        inputs=_inputs,
+        outputs=[output_audio, audio_mask], 
+    )
+
+    api_vamp_button = gr.Button("api vamp", visible=False)
+    api_vamp_button.click(
+        fn=api_vamp,
+        inputs=_inputs, 
+        outputs=[output_audio], 
+        api_name="vamp"
+    )
+
+    use_as_input_button.click(
+        fn=lambda x: x,
+        inputs=[output_audio],
+        outputs=[input_audio]
+    )
+
+    save_button.click(
+        fn=save_vamp,
+        inputs=_inputs | {notes_text, output_audio},
+        outputs=[thank_you, download_file]
+    )
+
+
+# demo.launch(share=True, enable_queue=True, debug=True)
+demo.launch(share=True, debug=True)  # from wam: because enable_queue seems to not be supported?

vampnet/conf/c2f.yml

@@ -0,0 +1,14 @@
+$include:
+  - conf/vampnet.yml
+
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+
+VampNet.embedding_dim: 1280
+VampNet.n_layers: 16
+VampNet.n_heads: 20
+
+AudioDataset.duration: 3.0
+
+
+AudioDataset.loudness_cutoff: -40.0

vampnet/conf/interface.yml

@@ -0,0 +1,10 @@
+Interface.coarse_ckpt: ./models/coarse.pth
+Interface.coarse2fine_ckpt: ./models/c2f.pth
+Interface.codec_ckpt: ./models/codec.pth
+Interface.coarse_chunk_size_s: 10
+Interface.coarse2fine_chunk_size_s: 3
+Interface.wavebeat_ckpt: ./models/wavebeat.pth
+
+# AudioLoader.sources:
+#   - /media/CHONK/null
+

vampnet/conf/lora/lora.yml

@@ -0,0 +1,22 @@
+$include:
+  - conf/vampnet.yml
+
+fine_tune: True
+
+train/AudioDataset.n_examples: 100000000
+val/AudioDataset.n_examples: 500
+
+
+NoamScheduler.warmup: 500
+
+batch_size: 6
+num_workers: 7
+save_iters: [2000, 4000, 10000,20000, 40000, 100000]
+sample_freq: 2000
+val_freq: 1000
+
+AdamW.lr: 0.0001
+
+# let's us organize sound classes into folders and choose from those sound classes uniformly
+AudioDataset.without_replacement: False
+num_iters: 500000

vampnet/conf/vampnet.yml

@@ -0,0 +1,49 @@
+
+codec_ckpt: ./models/vampnet/codec.pth
+save_path: ckpt
+
+num_iters: 1000000000
+save_iters: [10000, 50000, 100000, 300000, 500000]
+val_idx: [0,1,2,3,4,5,6,7,8,9]
+sample_freq: 10000
+val_freq: 1000
+
+batch_size: 8
+num_workers: 10
+
+# Optimization
+amp: false
+
+CrossEntropyLoss.label_smoothing: 0.1
+
+AdamW.lr: 0.001
+
+NoamScheduler.factor: 2.0
+NoamScheduler.warmup: 10000
+
+VampNet.vocab_size: 1024
+VampNet.n_codebooks: 4
+VampNet.n_conditioning_codebooks: 0
+VampNet.r_cond_dim: 0
+VampNet.noise_mode: mask
+VampNet.embedding_dim: 1280
+VampNet.n_layers: 20
+VampNet.n_heads: 20
+VampNet.flash_attn: false
+VampNet.dropout: 0.1
+
+AudioLoader.relative_path: ""
+AudioDataset.loudness_cutoff: -30.0
+AudioDataset.without_replacement: true
+AudioLoader.shuffle: true
+
+AudioDataset.duration: 10.0
+
+train/AudioDataset.n_examples: 10000000
+train/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/audio-train
+
+val/AudioDataset.n_examples: 2000
+val/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/audio-val
+

vampnet/scripts/exp/eval.py

@@ -0,0 +1,110 @@
+from pathlib import Path
+import os
+from functools import partial
+
+from frechet_audio_distance import FrechetAudioDistance
+import pandas
+import argbind
+import torch
+from tqdm import tqdm
+
+import audiotools
+from audiotools import AudioSignal
+
+@argbind.bind(without_prefix=True)
+def eval(
+    exp_dir: str = None,
+    baseline_key: str = "baseline", 
+    audio_ext: str = ".wav",
+):
+    assert exp_dir is not None
+    exp_dir = Path(exp_dir)
+    assert exp_dir.exists(), f"exp_dir {exp_dir} does not exist"
+
+    # set up our metrics
+    # sisdr_loss = audiotools.metrics.distance.SISDRLoss()
+    # stft_loss = audiotools.metrics.spectral.MultiScaleSTFTLoss()
+    mel_loss = audiotools.metrics.spectral.MelSpectrogramLoss()
+    frechet = FrechetAudioDistance(
+        use_pca=False, 
+        use_activation=False,
+        verbose=True, 
+        audio_load_worker=4,
+    )
+    frechet.model.to("cuda" if torch.cuda.is_available() else "cpu")
+
+    # figure out what conditions we have
+    conditions = [d.name for d in exp_dir.iterdir() if d.is_dir()]
+
+    assert baseline_key in conditions, f"baseline_key {baseline_key} not found in {exp_dir}"
+    conditions.remove(baseline_key)
+
+    print(f"Found {len(conditions)} conditions in {exp_dir}")
+    print(f"conditions: {conditions}")
+
+    baseline_dir = exp_dir / baseline_key 
+    baseline_files = sorted(list(baseline_dir.glob(f"*{audio_ext}")), key=lambda x: int(x.stem))
+
+    metrics = []
+    for condition in tqdm(conditions):
+        cond_dir = exp_dir / condition
+        cond_files = sorted(list(cond_dir.glob(f"*{audio_ext}")), key=lambda x: int(x.stem))
+
+        print(f"computing fad for {baseline_dir} and {cond_dir}")
+        frechet_score = frechet.score(baseline_dir, cond_dir)
+
+        # make sure we have the same number of files
+        num_files = min(len(baseline_files), len(cond_files))
+        baseline_files = baseline_files[:num_files]
+        cond_files = cond_files[:num_files]
+        assert len(list(baseline_files)) == len(list(cond_files)), f"number of files in {baseline_dir} and {cond_dir} do not match. {len(list(baseline_files))} vs {len(list(cond_files))}"
+
+        def process(baseline_file, cond_file):
+            # make sure the files match (same name)
+            assert baseline_file.stem == cond_file.stem, f"baseline file {baseline_file} and cond file {cond_file} do not match"
+
+            # load the files
+            baseline_sig = AudioSignal(str(baseline_file))
+            cond_sig = AudioSignal(str(cond_file))
+
+            cond_sig.resample(baseline_sig.sample_rate)
+            cond_sig.truncate_samples(baseline_sig.length)
+
+            # if our condition is inpainting, we need to trim the conditioning off
+            if "inpaint" in condition:
+                ctx_amt = float(condition.split("_")[-1])
+                ctx_samples = int(ctx_amt * baseline_sig.sample_rate)
+                print(f"found inpainting condition. trimming off {ctx_samples} samples from {cond_file} and {baseline_file}")
+                cond_sig.trim(ctx_samples, ctx_samples)
+                baseline_sig.trim(ctx_samples, ctx_samples)
+
+            return {
+                # "sisdr": -sisdr_loss(baseline_sig, cond_sig).item(),
+                # "stft": stft_loss(baseline_sig, cond_sig).item(),
+                "mel": mel_loss(baseline_sig, cond_sig).item(),
+                "frechet": frechet_score,
+                # "visqol": vsq,
+                "condition": condition,
+                "file": baseline_file.stem,
+            }
+
+        print(f"processing {len(baseline_files)} files in {baseline_dir} and {cond_dir}")
+        metrics.extend(tqdm(map(process, baseline_files, cond_files), total=len(baseline_files)))
+
+    metric_keys = [k for k in metrics[0].keys() if k not in ("condition", "file")]
+
+
+    for mk in metric_keys:
+        stat = pandas.DataFrame(metrics)
+        stat = stat.groupby(['condition'])[mk].agg(['mean', 'count', 'std'])
+        stat.to_csv(exp_dir / f"stats-{mk}.csv")
+
+    df = pandas.DataFrame(metrics)
+    df.to_csv(exp_dir / "metrics-all.csv", index=False)
+
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+
+    with argbind.scope(args):
+        eval()

vampnet/scripts/exp/experiment.py

@@ -0,0 +1,254 @@
+from pathlib import Path
+import random
+from typing import List
+import tempfile
+import subprocess
+
+import argbind
+from tqdm import tqdm
+import torch
+
+from vampnet.interface import Interface
+from vampnet import mask as pmask
+import audiotools as at
+
+Interface: Interface = argbind.bind(Interface)
+
+
+
+def calculate_bitrate(
+        interface, num_codebooks, 
+        downsample_factor
+    ):
+    bit_width = 10
+    sr = interface.codec.sample_rate
+    hop = interface.codec.hop_size
+    rate = (sr / hop) * ((bit_width * num_codebooks) / downsample_factor)
+    return rate
+
+def baseline(sig, interface):
+    return interface.preprocess(sig)
+
+def reconstructed(sig, interface):
+    return interface.to_signal(
+        interface.encode(sig)
+    )
+
+def coarse2fine(sig, interface):
+    z = interface.encode(sig)
+    z = z[:, :interface.c2f.n_conditioning_codebooks, :]
+
+    z = interface.coarse_to_fine(z)
+    return interface.to_signal(z)
+
+class CoarseCond:
+
+    def __init__(self, num_conditioning_codebooks, downsample_factor):
+        self.num_conditioning_codebooks = num_conditioning_codebooks
+        self.downsample_factor = downsample_factor
+
+    def __call__(self, sig, interface):
+        z = interface.encode(sig)
+        mask = pmask.full_mask(z)
+        mask = pmask.codebook_unmask(mask, self.num_conditioning_codebooks)
+        mask = pmask.periodic_mask(mask, self.downsample_factor)
+
+        zv = interface.coarse_vamp(z, mask)
+        zv = interface.coarse_to_fine(zv)
+        return interface.to_signal(zv)
+
+def opus(sig, interface, bitrate=128):
+    sig = interface.preprocess(sig)
+    
+    with tempfile.NamedTemporaryFile(suffix=".wav") as f:
+        sig.write(f.name)
+
+        opus_name = Path(f.name).with_suffix(".opus")
+        # convert to opus
+        cmd = [
+            "ffmpeg", "-y", "-i", f.name, 
+            "-c:a", "libopus", 
+            "-b:a", f"{bitrate}", 
+           opus_name
+        ]
+        subprocess.run(cmd, check=True)
+
+        # convert back to wav
+        output_name = Path(f"{f.name}-opus").with_suffix(".wav")
+        cmd = [
+            "ffmpeg", "-y", "-i", opus_name, 
+            output_name
+        ]
+
+        subprocess.run(cmd, check=True)
+
+        sig = at.AudioSignal(
+            output_name, 
+            sample_rate=sig.sample_rate
+        )
+    return sig
+
+def mask_ratio_1_step(ratio=1.0):
+    def wrapper(sig, interface):
+        z = interface.encode(sig)
+        mask = pmask.linear_random(z, ratio)
+        zv = interface.coarse_vamp(
+            z, 
+            mask,
+            sampling_steps=1, 
+        )
+
+        return interface.to_signal(zv)
+    return wrapper
+
+def num_sampling_steps(num_steps=1):
+    def wrapper(sig, interface: Interface):
+        z = interface.encode(sig)
+        mask = pmask.periodic_mask(z, 16)
+        zv = interface.coarse_vamp(
+            z, 
+            mask,
+            sampling_steps=num_steps, 
+        )
+
+        zv = interface.coarse_to_fine(zv)
+        return interface.to_signal(zv)
+    return wrapper
+
+def beat_mask(ctx_time):
+    def wrapper(sig, interface):
+        beat_mask = interface.make_beat_mask(
+            sig,
+            before_beat_s=ctx_time/2,
+            after_beat_s=ctx_time/2,
+            invert=True
+        )
+
+        z = interface.encode(sig)
+
+        zv = interface.coarse_vamp(
+            z, beat_mask
+        )
+
+        zv = interface.coarse_to_fine(zv)
+        return interface.to_signal(zv)
+    return wrapper
+
+def inpaint(ctx_time):
+    def wrapper(sig, interface: Interface):
+        z = interface.encode(sig)
+        mask = pmask.inpaint(z, interface.s2t(ctx_time), interface.s2t(ctx_time))
+
+        zv = interface.coarse_vamp(z, mask)
+        zv = interface.coarse_to_fine(zv)
+        
+        return interface.to_signal(zv)
+    return wrapper
+
+def token_noise(noise_amt):
+    def wrapper(sig, interface: Interface):
+        z = interface.encode(sig)
+        mask = pmask.random(z, noise_amt)
+        z = torch.where(
+            mask, 
+            torch.randint_like(z, 0, interface.coarse.vocab_size), 
+            z
+        )
+        return interface.to_signal(z)
+    return wrapper
+
+EXP_REGISTRY = {}
+
+EXP_REGISTRY["gen-compression"] = {
+    "baseline": baseline,
+    "reconstructed": reconstructed,
+    "coarse2fine": coarse2fine,
+    **{
+        f"{n}_codebooks_downsampled_{x}x": CoarseCond(num_conditioning_codebooks=n, downsample_factor=x)
+            for (n, x) in (
+                (1, 1), # 1 codebook, no downsampling
+                (4, 4), # 4 codebooks, downsampled 4x
+                (4, 16), # 4 codebooks, downsampled 16x
+                (4, 32), # 4 codebooks, downsampled 16x
+            )
+    }, 
+    **{
+        f"token_noise_{x}": mask_ratio_1_step(ratio=x)
+            for x in [0.25, 0.5, 0.75]
+    },
+
+}
+
+
+EXP_REGISTRY["sampling-steps"] = {
+    # "codec": reconstructed,
+    **{f"steps_{n}": num_sampling_steps(n)  for n in [1, 4, 12, 36, 64, 72]},
+}
+
+
+EXP_REGISTRY["musical-sampling"] = {
+    **{f"beat_mask_{t}": beat_mask(t) for t in [0.075]}, 
+    **{f"inpaint_{t}": inpaint(t) for t in [0.5, 1.0,]}, # multiply these by 2 (they go left and right)
+}
+
+@argbind.bind(without_prefix=True)
+def main(
+        sources=[
+            "/media/CHONK/hugo/spotdl/val",
+        ], 
+        output_dir: str = "./samples",
+        max_excerpts: int = 2000,
+        exp_type: str = "gen-compression", 
+        seed: int = 0,
+        ext: str = [".mp3"],
+    ):
+    at.util.seed(seed)
+    interface = Interface()
+
+    output_dir = Path(output_dir) 
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    from audiotools.data.datasets import AudioLoader, AudioDataset
+
+    loader = AudioLoader(sources=sources, shuffle_state=seed, ext=ext)
+    dataset = AudioDataset(loader, 
+        sample_rate=interface.codec.sample_rate, 
+        duration=interface.coarse.chunk_size_s, 
+        n_examples=max_excerpts, 
+        without_replacement=True,
+    )
+
+    if exp_type in EXP_REGISTRY:
+        SAMPLE_CONDS = EXP_REGISTRY[exp_type]
+    else:
+        raise ValueError(f"Unknown exp_type {exp_type}")
+
+
+    indices = list(range(max_excerpts))
+    random.shuffle(indices)
+    for i in tqdm(indices):
+        # if all our files are already there, skip
+        done = []
+        for name in SAMPLE_CONDS:
+            o_dir = Path(output_dir) / name
+            done.append((o_dir / f"{i}.wav").exists())
+        if all(done):
+            continue
+
+        sig = dataset[i]["signal"]
+        results = {
+            name: cond(sig, interface).cpu()
+            for name, cond in SAMPLE_CONDS.items()
+        }
+
+        for name, sig in results.items():
+            o_dir = Path(output_dir) / name
+            o_dir.mkdir(exist_ok=True, parents=True)
+
+            sig.write(o_dir / f"{i}.wav")
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+
+    with argbind.scope(args):
+        main()

vampnet/scripts/exp/fine_tune.py

@@ -0,0 +1,81 @@
+import argbind
+from pathlib import Path
+import yaml
+from typing import List
+
+
+
+
+"""example output: (yaml)
+
+"""
+
+@argbind.bind(without_prefix=True, positional=True)
+def fine_tune(audio_files_or_folders: List[str], name: str):
+
+    conf_dir = Path("conf")
+    assert conf_dir.exists(), "conf directory not found. are you in the vampnet directory?"
+
+    conf_dir = conf_dir / "generated"
+    conf_dir.mkdir(exist_ok=True)
+
+    finetune_dir = conf_dir / name
+    finetune_dir.mkdir(exist_ok=True)
+
+    finetune_c2f_conf = {
+        "$include": ["conf/lora/lora.yml"],
+        "fine_tune": True,
+        "train/AudioLoader.sources": audio_files_or_folders,
+        "val/AudioLoader.sources": audio_files_or_folders,
+        "VampNet.n_codebooks": 14,
+        "VampNet.n_conditioning_codebooks": 4,
+        "VampNet.embedding_dim": 1280,
+        "VampNet.n_layers": 16,
+        "VampNet.n_heads": 20,
+        "AudioDataset.duration": 3.0,
+        "AudioDataset.loudness_cutoff": -40.0,
+        "save_path": f"./runs/{name}/c2f",
+        "fine_tune_checkpoint": "./models/vampnet/c2f.pth"
+    }
+
+    finetune_coarse_conf = {
+        "$include": ["conf/lora/lora.yml"],
+        "fine_tune": True,
+        "train/AudioLoader.sources": audio_files_or_folders,
+        "val/AudioLoader.sources": audio_files_or_folders,
+        "save_path": f"./runs/{name}/coarse",
+        "fine_tune_checkpoint": "./models/vampnet/coarse.pth"
+    }
+
+    interface_conf = {
+        "Interface.coarse_ckpt": f"./runs/{name}/coarse/latest/vampnet/weights.pth",
+
+        "Interface.coarse2fine_ckpt": f"./runs/{name}/c2f/latest/vampnet/weights.pth",
+        "Interface.wavebeat_ckpt": "./models/wavebeat.pth",
+
+        "Interface.codec_ckpt": "./models/vampnet/codec.pth",
+        "AudioLoader.sources": [audio_files_or_folders],
+    }
+
+    # save the confs
+    with open(finetune_dir / "c2f.yml", "w") as f:
+        yaml.dump(finetune_c2f_conf, f)
+
+    with open(finetune_dir / "coarse.yml", "w") as f:
+        yaml.dump(finetune_coarse_conf, f)
+    
+    with open(finetune_dir / "interface.yml", "w") as f: 
+        yaml.dump(interface_conf, f)
+
+
+    print(f"generated confs in {finetune_dir}. run training jobs with `python scripts/exp/train.py --args.load {finetune_dir}/<c2f/coarse>.yml` ")
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+
+    with argbind.scope(args):
+        fine_tune()
+
+
+
+    

vampnet/scripts/exp/train.py

@@ -0,0 +1,680 @@
+import os
+import sys
+import warnings
+from pathlib import Path
+from typing import Optional
+from dataclasses import dataclass
+
+import argbind
+import audiotools as at
+import torch
+import torch.nn as nn
+from audiotools import AudioSignal
+from audiotools.data import transforms
+from einops import rearrange
+from rich import pretty
+from rich.traceback import install
+from torch.utils.tensorboard import SummaryWriter
+
+import vampnet
+from vampnet.modules.transformer import VampNet
+from vampnet.util import codebook_unflatten, codebook_flatten
+from vampnet import mask as pmask
+# from dac.model.dac import DAC
+from lac.model.lac import LAC as DAC
+
+from audiotools.ml.decorators import (
+    timer, Tracker, when
+)
+
+import loralib as lora
+
+import torch._dynamo
+torch._dynamo.config.verbose=True
+
+
+# Enable cudnn autotuner to speed up training
+# (can be altered by the funcs.seed function)
+torch.backends.cudnn.benchmark = bool(int(os.getenv("CUDNN_BENCHMARK", 1)))
+# Uncomment to trade memory for speed.
+
+# Install to make things look nice
+warnings.filterwarnings("ignore", category=UserWarning)
+pretty.install()
+install()
+
+# optim
+Accelerator = argbind.bind(at.ml.Accelerator, without_prefix=True)
+CrossEntropyLoss = argbind.bind(nn.CrossEntropyLoss)
+AdamW = argbind.bind(torch.optim.AdamW)
+NoamScheduler = argbind.bind(vampnet.scheduler.NoamScheduler)
+
+# transforms
+filter_fn = lambda fn: hasattr(fn, "transform") and fn.__qualname__ not in [
+    "BaseTransform",
+    "Compose",
+    "Choose",
+]
+tfm = argbind.bind_module(transforms, "train", "val", filter_fn=filter_fn)
+
+# model
+VampNet = argbind.bind(VampNet)
+
+
+# data
+AudioLoader = argbind.bind(at.datasets.AudioLoader)
+AudioDataset = argbind.bind(at.datasets.AudioDataset, "train", "val")
+
+IGNORE_INDEX = -100
+
+
+@argbind.bind("train", "val", without_prefix=True)
+def build_transform():
+    transform = transforms.Compose(
+        tfm.VolumeNorm(("const", -24)),
+        # tfm.PitchShift(),
+        tfm.RescaleAudio(),
+    )
+    return transform
+
+
+@torch.no_grad()
+def apply_transform(transform_fn, batch):
+    sig: AudioSignal = batch["signal"]
+    kwargs = batch["transform_args"]
+
+    sig: AudioSignal = transform_fn(sig.clone(), **kwargs)
+    return sig
+
+
+def build_datasets(args, sample_rate: int):
+    with argbind.scope(args, "train"):
+        train_data = AudioDataset(
+            AudioLoader(), sample_rate, transform=build_transform()
+        )
+    with argbind.scope(args, "val"):
+        val_data = AudioDataset(AudioLoader(), sample_rate, transform=build_transform())
+    return train_data, val_data
+
+
+def rand_float(shape, low, high, rng):
+    return rng.draw(shape)[:, 0] * (high - low) + low
+
+
+def flip_coin(shape, p, rng):
+    return rng.draw(shape)[:, 0] < p
+
+
+def num_params_hook(o, p):
+    return o + f" {p/1e6:<.3f}M params."
+
+
+def add_num_params_repr_hook(model):
+    import numpy as np
+    from functools import partial
+
+    for n, m in model.named_modules():
+        o = m.extra_repr()
+        p = sum([np.prod(p.size()) for p in m.parameters()])
+
+        setattr(m, "extra_repr", partial(num_params_hook, o=o, p=p))
+
+
+def accuracy(
+    preds: torch.Tensor,
+    target: torch.Tensor,
+    top_k: int = 1,
+    ignore_index: Optional[int] = None,
+) -> torch.Tensor:
+    # Flatten the predictions and targets to be of shape (batch_size * sequence_length, n_class)
+    preds = rearrange(preds, "b p s -> (b s) p")
+    target = rearrange(target, "b s -> (b s)")
+
+    # return torchmetrics.functional.accuracy(preds, target, task='multiclass', top_k=topk, num_classes=preds.shape[-1], ignore_index=ignore_index)
+    if ignore_index is not None:
+        # Create a mask for the ignored index
+        mask = target != ignore_index
+        # Apply the mask to the target and predictions
+        preds = preds[mask]
+        target = target[mask]
+
+    # Get the top-k predicted classes and their indices
+    _, pred_indices = torch.topk(preds, k=top_k, dim=-1)
+
+    # Determine if the true target is in the top-k predicted classes
+    correct = torch.sum(torch.eq(pred_indices, target.unsqueeze(1)), dim=1)
+
+    # Calculate the accuracy
+    accuracy = torch.mean(correct.float())
+
+    return accuracy
+
+def _metrics(z_hat, r, target, flat_mask, output):
+    for r_range in [(0, 0.5), (0.5, 1.0)]:
+        unmasked_target = target.masked_fill(flat_mask.bool(), IGNORE_INDEX)
+        masked_target = target.masked_fill(~flat_mask.bool(), IGNORE_INDEX)
+
+        assert target.shape[0] == r.shape[0]
+        # grab the indices of the r values that are in the range
+        r_idx = (r >= r_range[0]) & (r < r_range[1])
+
+        # grab the target and z_hat values that are in the range
+        r_unmasked_target = unmasked_target[r_idx]
+        r_masked_target = masked_target[r_idx]
+        r_z_hat = z_hat[r_idx]
+
+        for topk in (1, 25):
+            s, e = r_range
+            tag = f"accuracy-{s}-{e}/top{topk}"
+
+            output[f"{tag}/unmasked"] = accuracy(
+                preds=r_z_hat,
+                target=r_unmasked_target,
+                ignore_index=IGNORE_INDEX,
+                top_k=topk,
+            )
+            output[f"{tag}/masked"] = accuracy(
+                preds=r_z_hat,
+                target=r_masked_target,
+                ignore_index=IGNORE_INDEX,
+                top_k=topk,
+            )
+
+
+@dataclass
+class State:
+    model: VampNet
+    codec: DAC
+
+    optimizer: AdamW
+    scheduler: NoamScheduler
+    criterion: CrossEntropyLoss
+    grad_clip_val: float
+
+    rng: torch.quasirandom.SobolEngine
+
+    train_data: AudioDataset
+    val_data: AudioDataset
+
+    tracker: Tracker
+
+
+@timer()
+def train_loop(state: State, batch: dict, accel: Accelerator):
+    state.model.train()
+    batch = at.util.prepare_batch(batch, accel.device)
+    signal = apply_transform(state.train_data.transform, batch)
+
+    output = {}
+    vn = accel.unwrap(state.model)
+    with accel.autocast():
+        with torch.inference_mode():
+            state.codec.to(accel.device)
+            z = state.codec.encode(signal.samples, signal.sample_rate)["codes"]
+            z = z[:, : vn.n_codebooks, :]
+
+        n_batch = z.shape[0]
+        r = state.rng.draw(n_batch)[:, 0].to(accel.device)
+
+        mask = pmask.random(z, r)
+        mask = pmask.codebook_unmask(mask, vn.n_conditioning_codebooks)
+        z_mask, mask = pmask.apply_mask(z, mask, vn.mask_token)
+        
+        z_mask_latent = vn.embedding.from_codes(z_mask, state.codec)
+
+        dtype = torch.bfloat16 if accel.amp else None
+        with accel.autocast(dtype=dtype):
+            z_hat = state.model(z_mask_latent)
+
+        target = codebook_flatten(
+            z[:, vn.n_conditioning_codebooks :, :],
+        )
+
+        flat_mask = codebook_flatten(
+            mask[:, vn.n_conditioning_codebooks :, :],
+        )
+
+        # replace target with ignore index for masked tokens
+        t_masked = target.masked_fill(~flat_mask.bool(), IGNORE_INDEX)
+        output["loss"] = state.criterion(z_hat, t_masked)
+
+        _metrics(
+            r=r,
+            z_hat=z_hat,
+            target=target,
+            flat_mask=flat_mask,
+            output=output,
+        )
+
+    
+    accel.backward(output["loss"])
+
+    output["other/learning_rate"] = state.optimizer.param_groups[0]["lr"]
+    output["other/batch_size"] = z.shape[0]
+
+
+    accel.scaler.unscale_(state.optimizer)
+    output["other/grad_norm"] = torch.nn.utils.clip_grad_norm_(
+        state.model.parameters(), state.grad_clip_val
+    )
+
+    accel.step(state.optimizer)
+    state.optimizer.zero_grad()
+
+    state.scheduler.step()
+    accel.update()
+
+
+    return {k: v for k, v in sorted(output.items())}
+
+
+@timer()
+@torch.no_grad()
+def val_loop(state: State, batch: dict, accel: Accelerator):
+    state.model.eval()
+    state.codec.eval()
+    batch = at.util.prepare_batch(batch, accel.device)
+    signal = apply_transform(state.val_data.transform, batch)
+
+    vn = accel.unwrap(state.model)
+    z = state.codec.encode(signal.samples, signal.sample_rate)["codes"]
+    z = z[:, : vn.n_codebooks, :]
+
+    n_batch = z.shape[0]
+    r = state.rng.draw(n_batch)[:, 0].to(accel.device)
+
+    mask = pmask.random(z, r)
+    mask = pmask.codebook_unmask(mask, vn.n_conditioning_codebooks)
+    z_mask, mask = pmask.apply_mask(z, mask, vn.mask_token)
+
+    z_mask_latent = vn.embedding.from_codes(z_mask, state.codec)
+
+    z_hat = state.model(z_mask_latent)
+
+    target = codebook_flatten(
+        z[:, vn.n_conditioning_codebooks :, :],
+    )
+
+    flat_mask = codebook_flatten(
+        mask[:, vn.n_conditioning_codebooks :, :]
+    )
+
+    output = {}
+    # replace target with ignore index for masked tokens
+    t_masked = target.masked_fill(~flat_mask.bool(), IGNORE_INDEX)
+    output["loss"] = state.criterion(z_hat, t_masked)
+
+    _metrics(
+        r=r,
+        z_hat=z_hat,
+        target=target,
+        flat_mask=flat_mask,
+        output=output,
+    )
+
+    return output
+
+
+def validate(state, val_dataloader, accel):
+    for batch in val_dataloader:
+        output = val_loop(state, batch, accel)
+    # Consolidate state dicts if using ZeroRedundancyOptimizer
+    if hasattr(state.optimizer, "consolidate_state_dict"):
+        state.optimizer.consolidate_state_dict()
+    return output
+
+
+def checkpoint(state, save_iters, save_path, fine_tune):
+    if accel.local_rank != 0:
+        state.tracker.print(f"ERROR:Skipping checkpoint on rank {accel.local_rank}")
+        return
+
+    metadata = {"logs": dict(state.tracker.history)}
+
+    tags = ["latest"]
+    state.tracker.print(f"Saving to {str(Path('.').absolute())}")
+
+    if state.tracker.step in save_iters:
+        tags.append(f"{state.tracker.step // 1000}k")
+
+    if state.tracker.is_best("val", "loss"):
+        state.tracker.print(f"Best model so far")
+        tags.append("best")
+
+    if fine_tune:
+        for tag in tags: 
+            # save the lora model 
+            (Path(save_path) / tag).mkdir(parents=True, exist_ok=True)
+            torch.save(
+                lora.lora_state_dict(accel.unwrap(state.model)), 
+                f"{save_path}/{tag}/lora.pth"
+            )
+
+    for tag in tags:
+        model_extra = {
+            "optimizer.pth": state.optimizer.state_dict(),
+            "scheduler.pth": state.scheduler.state_dict(),
+            "tracker.pth": state.tracker.state_dict(),
+            "metadata.pth": metadata,
+        }
+
+        accel.unwrap(state.model).metadata = metadata
+        accel.unwrap(state.model).save_to_folder(
+            f"{save_path}/{tag}", model_extra, package=False
+        )
+
+
+def save_sampled(state, z, writer):
+    num_samples = z.shape[0]
+
+    for i in range(num_samples):
+        sampled = accel.unwrap(state.model).generate(
+            codec=state.codec,
+            time_steps=z.shape[-1],
+            start_tokens=z[i : i + 1],
+        )
+        sampled.cpu().write_audio_to_tb(
+            f"sampled/{i}",
+            writer,
+            step=state.tracker.step,
+            plot_fn=None,
+        )
+
+
+def save_imputation(state, z, val_idx, writer):
+    n_prefix = int(z.shape[-1] * 0.25)
+    n_suffix = int(z.shape[-1] *  0.25)
+
+    vn = accel.unwrap(state.model)
+
+    mask = pmask.inpaint(z, n_prefix, n_suffix)
+    mask = pmask.codebook_unmask(mask, vn.n_conditioning_codebooks)
+    z_mask, mask = pmask.apply_mask(z, mask, vn.mask_token)
+
+    imputed_noisy = vn.to_signal(z_mask, state.codec)
+    imputed_true = vn.to_signal(z, state.codec)
+
+    imputed = []
+    for i in range(len(z)):
+        imputed.append(
+            vn.generate(
+                codec=state.codec,
+                time_steps=z.shape[-1],
+                start_tokens=z[i][None, ...],
+                mask=mask[i][None, ...],
+            )   
+        )   
+    imputed = AudioSignal.batch(imputed)
+
+    for i in range(len(val_idx)):
+        imputed_noisy[i].cpu().write_audio_to_tb(
+            f"inpainted_prompt/{i}",
+            writer,
+            step=state.tracker.step,
+            plot_fn=None,
+        )
+        imputed[i].cpu().write_audio_to_tb(
+            f"inpainted_middle/{i}",
+            writer,
+            step=state.tracker.step,
+            plot_fn=None,
+        )
+        imputed_true[i].cpu().write_audio_to_tb(
+            f"reconstructed/{i}",
+            writer,
+            step=state.tracker.step,
+            plot_fn=None,
+        )
+
+
+@torch.no_grad()
+def save_samples(state: State, val_idx: int, writer: SummaryWriter):
+    state.model.eval()
+    state.codec.eval()
+    vn = accel.unwrap(state.model)
+
+    batch = [state.val_data[i] for i in val_idx]
+    batch = at.util.prepare_batch(state.val_data.collate(batch), accel.device)
+
+    signal = apply_transform(state.val_data.transform, batch)
+
+    z = state.codec.encode(signal.samples, signal.sample_rate)["codes"]
+    z = z[:, : vn.n_codebooks, :]
+
+    r = torch.linspace(0.1, 0.95, len(val_idx)).to(accel.device)
+
+
+    mask = pmask.random(z, r)
+    mask = pmask.codebook_unmask(mask, vn.n_conditioning_codebooks)
+    z_mask, mask = pmask.apply_mask(z, mask, vn.mask_token)
+
+    z_mask_latent = vn.embedding.from_codes(z_mask, state.codec)
+
+    z_hat = state.model(z_mask_latent)
+
+    z_pred = torch.softmax(z_hat, dim=1).argmax(dim=1)
+    z_pred = codebook_unflatten(z_pred, n_c=vn.n_predict_codebooks)
+    z_pred = torch.cat([z[:, : vn.n_conditioning_codebooks, :], z_pred], dim=1)
+
+    generated = vn.to_signal(z_pred, state.codec)
+    reconstructed = vn.to_signal(z, state.codec)
+    masked = vn.to_signal(z_mask.squeeze(1), state.codec)
+
+    for i in range(generated.batch_size):
+        audio_dict = {
+            "original": signal[i],
+            "masked": masked[i],
+            "generated": generated[i],
+            "reconstructed": reconstructed[i],
+        }
+        for k, v in audio_dict.items():
+            v.cpu().write_audio_to_tb(
+                f"onestep/_{i}.r={r[i]:0.2f}/{k}",
+                writer,
+                step=state.tracker.step,
+                plot_fn=None,
+            )
+
+    save_sampled(state=state, z=z, writer=writer)
+    save_imputation(state=state, z=z, val_idx=val_idx, writer=writer)
+
+
+
+@argbind.bind(without_prefix=True)
+def load(
+    args,
+    accel: at.ml.Accelerator,
+    tracker: Tracker,
+    save_path: str,
+    resume: bool = False,
+    tag: str = "latest",
+    fine_tune_checkpoint: Optional[str] = None,
+    grad_clip_val: float = 5.0,
+) -> State:
+    codec = DAC.load(args["codec_ckpt"], map_location="cpu")
+    codec.eval()
+
+    model, v_extra = None, {}
+
+    if args["fine_tune"]:
+        assert fine_tune_checkpoint is not None, "Must provide a fine-tune checkpoint"
+        model = torch.compile(
+            VampNet.load(location=Path(fine_tune_checkpoint), 
+                         map_location="cpu", 
+            )
+        )
+        
+    if resume:
+        kwargs = {
+            "folder": f"{save_path}/{tag}",
+            "map_location": "cpu",
+            "package": False,
+        }
+        tracker.print(f"Loading checkpoint from {kwargs['folder']}")
+        if (Path(kwargs["folder"]) / "vampnet").exists():
+            model, v_extra = VampNet.load_from_folder(**kwargs)
+        else:
+            raise ValueError(
+                f"Could not find a VampNet checkpoint in {kwargs['folder']}"
+            )
+
+
+
+
+    model = torch.compile(VampNet()) if model is None else model
+    model = accel.prepare_model(model)
+
+    # assert accel.unwrap(model).n_codebooks == codec.quantizer.n_codebooks
+    assert (
+        accel.unwrap(model).vocab_size == codec.quantizer.quantizers[0].codebook_size
+    )
+
+    optimizer = AdamW(model.parameters(), use_zero=accel.use_ddp)
+    scheduler = NoamScheduler(optimizer, d_model=accel.unwrap(model).embedding_dim)
+    scheduler.step()
+
+    if "optimizer.pth" in v_extra:
+        optimizer.load_state_dict(v_extra["optimizer.pth"])
+        scheduler.load_state_dict(v_extra["scheduler.pth"])
+    if "tracker.pth" in v_extra:
+        tracker.load_state_dict(v_extra["tracker.pth"])
+    
+    criterion = CrossEntropyLoss()
+
+    sample_rate = codec.sample_rate
+
+    # a better rng for sampling from our schedule
+    rng = torch.quasirandom.SobolEngine(1, scramble=True, seed=args["seed"])  
+
+    # log a model summary w/ num params
+    if accel.local_rank == 0:
+        add_num_params_repr_hook(accel.unwrap(model))
+        with open(f"{save_path}/model.txt", "w") as f:
+            f.write(repr(accel.unwrap(model)))
+
+    # load the datasets
+    train_data, val_data = build_datasets(args, sample_rate)
+
+    return State(
+        tracker=tracker,
+        model=model,
+        codec=codec,
+        optimizer=optimizer,
+        scheduler=scheduler,
+        criterion=criterion,
+        rng=rng,
+        train_data=train_data,
+        val_data=val_data,
+        grad_clip_val=grad_clip_val,
+    )
+
+
+@argbind.bind(without_prefix=True)
+def train(
+    args,
+    accel: at.ml.Accelerator,
+    seed: int = 0,
+    codec_ckpt: str = None,
+    save_path: str = "ckpt",
+    num_iters: int = int(1000e6),
+    save_iters: list = [10000, 50000, 100000, 300000, 500000,],
+    sample_freq: int = 10000, 
+    val_freq: int = 1000,
+    batch_size: int = 12,
+    val_idx: list = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+    num_workers: int = 10,
+    fine_tune: bool = False, 
+):
+    assert codec_ckpt is not None, "codec_ckpt is required"
+
+    seed = seed + accel.local_rank
+    at.util.seed(seed)
+    writer = None
+
+    if accel.local_rank == 0:
+        writer = SummaryWriter(log_dir=f"{save_path}/logs/")
+        argbind.dump_args(args, f"{save_path}/args.yml")
+
+        tracker = Tracker(
+            writer=writer, log_file=f"{save_path}/log.txt", rank=accel.local_rank
+        )
+
+    # load the codec model
+    state: State = load(
+        args=args, 
+        accel=accel, 
+        tracker=tracker, 
+        save_path=save_path)
+    print("initialized state.")
+
+    train_dataloader = accel.prepare_dataloader(
+        state.train_data,
+        start_idx=state.tracker.step * batch_size,
+        num_workers=num_workers,
+        batch_size=batch_size,
+        collate_fn=state.train_data.collate,
+    )
+    val_dataloader = accel.prepare_dataloader(
+        state.val_data,
+        start_idx=0,
+        num_workers=num_workers,
+        batch_size=batch_size,
+        collate_fn=state.val_data.collate,
+        persistent_workers=num_workers > 0,
+    )
+    print("initialized dataloader.")
+
+    
+
+    if fine_tune:
+        lora.mark_only_lora_as_trainable(state.model)
+        print("marked only lora as trainable.")
+
+    # Wrap the functions so that they neatly track in TensorBoard + progress bars
+    # and only run when specific conditions are met.
+    global train_loop, val_loop, validate, save_samples, checkpoint
+
+    train_loop = tracker.log("train", "value", history=False)(
+        tracker.track("train", num_iters, completed=state.tracker.step)(train_loop)
+    )
+    val_loop = tracker.track("val", len(val_dataloader))(val_loop)
+    validate = tracker.log("val", "mean")(validate)
+
+    save_samples = when(lambda: accel.local_rank == 0)(save_samples)
+    checkpoint = when(lambda: accel.local_rank == 0)(checkpoint)
+
+    print("starting training loop.")
+    with tracker.live:
+        for tracker.step, batch in enumerate(train_dataloader, start=tracker.step):
+            train_loop(state, batch, accel)
+
+            last_iter = (
+                tracker.step == num_iters - 1 if num_iters is not None else False
+            )
+
+            if tracker.step % sample_freq == 0 or last_iter:
+                save_samples(state, val_idx, writer)
+
+            if tracker.step % val_freq == 0 or last_iter:
+                validate(state, val_dataloader, accel)
+                checkpoint(
+                    state=state, 
+                    save_iters=save_iters, 
+                    save_path=save_path, 
+                    fine_tune=fine_tune)
+
+                # Reset validation progress bar, print summary since last validation.
+                tracker.done("val", f"Iteration {tracker.step}")
+
+            if last_iter:
+                break
+
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+    args["args.debug"] = int(os.getenv("LOCAL_RANK", 0)) == 0
+    with argbind.scope(args):
+        with Accelerator() as accel:
+            if accel.local_rank != 0:
+                sys.tracebacklimit = 0
+            train(args, accel)

vampnet/scripts/utils/README.md

@@ -0,0 +1,28 @@
+# Scripts
+
+## process_zip.py
+
+Some requirements that may not be installed in the docker image:
+* argbind
+* wav2wav (pip install git+https://github.com/descriptinc/lyrebird-wav2wav.git or `pip install git+https://github.com/descriptinc/lyrebird-wav2wav.git@<branchname>`)
+
+### zip folder structure
+
+The zip folder should have the following internal structure:
+
+```
+base_folder/
+    test_case_1/
+        before.wav
+    test_case_2/
+        before.wav
+    ...
+    test_case_n/
+        before.wav
+```
+
+Note: There can be issues with the output zip if the input zip folder structure is too deep or too shallow. IF you want/need to use a zip file with a different folder structure, adjust this:
+https://github.com/descriptinc/lyrebird-wav2wav/blob/136c923ce19df03876a515ca0ed83854710cfa30/scripts/utils/process_zip.py#L28
+
+### Execution
+`python process_zip.py <path/to/zip> -tag <string>`

vampnet/scripts/utils/data/augment.py

@@ -0,0 +1,67 @@
+from pathlib import Path
+
+import audiotools as at
+from audiotools import AudioSignal
+
+import argbind
+import tqdm
+import torch
+
+
+from torch_pitch_shift import pitch_shift, get_fast_shifts
+from torch_time_stretch import time_stretch, get_fast_stretches
+
+from audiotools.core.util import sample_from_dist
+
+
+@argbind.bind(without_prefix=True)
+def augment(
+    audio_folder: Path = None,
+    dest_folder: Path = None,
+    n_augmentations: int = 10,
+):
+    """ 
+        Augment a folder of audio files by applying audiotools and pedalboard transforms. 
+
+        The dest foler will contain a folder for each of the clean dataset's files. 
+        Under each of these folders, there will be a clean file and many augmented files.
+    """
+    assert audio_folder is not None
+    assert dest_folder is not None
+    audio_files = at.util.find_audio(audio_folder)
+
+    for audio_file in tqdm.tqdm(audio_files):
+        subtree = dest_folder / audio_file.relative_to(audio_folder).parent
+        subdir = subtree / audio_file.stem
+        subdir.mkdir(parents=True, exist_ok=True)
+
+        src = AudioSignal(audio_file).to("cuda" if torch.cuda.is_available() else "cpu")
+
+        
+        for i, chunk in tqdm.tqdm(enumerate(src.windows(10, 10))):
+            # apply pedalboard transforms
+            for j in range(n_augmentations):
+                # pitch shift between -7 and 7 semitones
+                import random
+                dst = chunk.clone()
+                dst.samples = pitch_shift(
+                    dst.samples, 
+                    shift=random.choice(get_fast_shifts(src.sample_rate, 
+                            condition=lambda x: x >= 0.25 and x <= 1.0)), 
+                    sample_rate=src.sample_rate
+                )
+                dst.samples = time_stretch(
+                    dst.samples,
+                    stretch=random.choice(get_fast_stretches(src.sample_rate, 
+                                          condition=lambda x: x >= 0.667 and x <= 1.5, )),
+                    sample_rate=src.sample_rate, 
+                )
+
+                dst.cpu().write(subdir / f"{i}-{j}.wav")
+
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+
+    with argbind.scope(args):
+        augment()

vampnet/scripts/utils/data/maestro-reorg.py

@@ -0,0 +1,39 @@
+from pathlib import Path
+import json
+import os
+
+maestro_path = Path("/media/CHONK/hugo/maestro-v3.0.0")
+output_path = Path("/media/CHONK/hugo/maestro-v3.0.0-split")
+
+# split
+with open(maestro_path / "maestro-v3.0.0.json") as f:
+    maestro = json.load(f)
+
+breakpoint()
+train = []
+validation = []
+test = []
+for key, split in maestro["split"].items():
+    audio_filename = maestro['audio_filename'][key]
+    if split == "train":
+        train.append(audio_filename)
+    elif split == "test":
+        test.append(audio_filename)
+    elif split == "validation":
+        validation.append(audio_filename)
+    else:
+        raise ValueError(f"Unknown split {split}")
+
+# symlink all files
+for audio_filename in train:
+    p = output_path / "train" / audio_filename
+    p.parent.mkdir(parents=True, exist_ok=True)
+    os.symlink(maestro_path / audio_filename, p)
+for audio_filename in validation:
+    p = output_path / "validation" / audio_filename
+    p.parent.mkdir(parents=True, exist_ok=True)
+    os.symlink(maestro_path / audio_filename, p)
+for audio_filename in test:
+    p = output_path / "test" / audio_filename
+    p.parent.mkdir(parents=True, exist_ok=True)
+    os.symlink(maestro_path / audio_filename, p)

vampnet/scripts/utils/plots.py

@@ -0,0 +1,43 @@
+import matplotlib.pyplot as plt
+import seaborn as sns
+from pandas.api.types import CategoricalDtype
+
+def plot_metrics(metrics, condition_to_latex, title, color_palette):
+    # Add a new column to your dataframe with the latex representation
+    metrics['condition_latex'] = metrics['condition'].map(condition_to_latex)
+
+    # Order condition_latex as per the condition_to_latex dictionary
+    cat_type = CategoricalDtype(categories=condition_to_latex.values(), ordered=True)
+    metrics['condition_latex'] = metrics['condition_latex'].astype(cat_type)
+
+    # Compute mean and std for each condition for each metric
+    grouped = metrics.groupby('condition_latex')[['mel', 'frechet']].agg(['mean', 'std'])
+
+    fig, axs = plt.subplots(2, 1, figsize=(7, 5.25))
+
+    # Set the main title for the figure
+    fig.suptitle(title, fontsize=16)
+
+    # Get color for each bar in the plot
+    bar_colors = [color_palette[condition] for condition in grouped.index]
+
+    # Plot mel
+    sns.boxplot(x='condition_latex', y='mel', data=metrics, ax=axs[0], palette=color_palette, showfliers=False)
+    axs[0].set_ylabel('Mel Spectrogram Loss \u2190')
+    axs[0].set_xlabel('') # Remove x-axis label
+    axs[0].set_xticklabels(grouped.index, rotation=0, ha='center')
+
+    # Plot frechet
+    axs[1].bar(grouped.index, grouped['frechet']['mean'], yerr=grouped['frechet']['std'], color=bar_colors)
+    axs[1].set_ylabel('FAD \u2190')
+    axs[1].set_xlabel('') # Remove x-axis label
+    axs[1].set_xticklabels(grouped.index, rotation=0, ha='center')
+
+    # Adjust the space between plots
+    plt.subplots_adjust(hspace=0.1)
+
+    # Remove any unnecessary space around the plot
+    plt.tight_layout(rect=[0, 0, 1, 0.96])
+
+    # Reduce the space between suptitle and the plot
+    plt.subplots_adjust(top=0.92)

vampnet/scripts/utils/remove_quiet_files.py

@@ -0,0 +1,29 @@
+# removes files with loudness below 24db
+
+from pathlib import Path 
+import shutil
+import audiotools as at
+import argbind
+
+@argbind.bind(without_prefix=True)
+def remove_quiet_files(
+    src_dir: Path = None,
+    dest_dir: Path = None,
+    min_loudness: float = -30,
+):
+    # copy src to dest
+    dest_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copytree(src_dir, dest_dir, dirs_exist_ok=True)
+    
+    audio_files = at.util.find_audio(dest_dir)
+    for audio_file in audio_files:
+        sig = at.AudioSignal(audio_file)
+        if sig.loudness() < min_loudness:
+            audio_file.unlink()
+            print(f"removed {audio_file}")
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+
+    with argbind.scope(args):
+        remove_quiet_files()

vampnet/scripts/utils/split.py

@@ -0,0 +1,64 @@
+from pathlib import Path
+import random
+import shutil
+import os
+import json 
+
+import argbind
+from tqdm import tqdm
+from tqdm.contrib.concurrent import thread_map
+
+from audiotools.core import util
+
+
+@argbind.bind(without_prefix=True)
+def train_test_split(
+    audio_folder: str = ".", 
+    test_size: float = 0.2,
+    seed: int = 42,
+    pattern: str = "**/*.mp3",
+):
+    print(f"finding audio")
+
+    audio_folder = Path(audio_folder)
+    audio_files = list(tqdm(audio_folder.glob(pattern)))
+    print(f"found {len(audio_files)} audio files")
+    
+    # split according to test_size
+    n_test = int(len(audio_files) * test_size)
+    n_train = len(audio_files) - n_test
+
+    # shuffle
+    random.seed(seed)
+    random.shuffle(audio_files)
+
+    train_files = audio_files[:n_train]
+    test_files = audio_files[n_train:]
+
+
+    print(f"Train files: {len(train_files)}")
+    print(f"Test files: {len(test_files)}")
+    continue_ = input("Continue [yn]? ") or "n"
+
+    if continue_ != "y":
+        return
+    
+    for split, files in (
+        ("train", train_files), ("test", test_files)
+    ):
+        for file in tqdm(files):
+            out_file = audio_folder.parent / f"{audio_folder.name}-{split}" / Path(file).name
+            out_file.parent.mkdir(exist_ok=True, parents=True)
+            os.symlink(file, out_file)
+
+        # save split as json
+        with open(Path(audio_folder) / f"{split}.json", "w") as f:
+            json.dump([str(f) for f in files], f)
+    
+
+    
+if __name__ == "__main__":
+    args  = argbind.parse_args()
+
+    with argbind.scope(args):
+        train_test_split()

vampnet/scripts/utils/split_long_audio_file.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+import argbind
+
+import audiotools as at
+import tqdm
+
+
+@argbind.bind(without_prefix=True)
+def split_long_audio_file(
+    file: str = None, 
+    max_chunk_size_s: int = 60*10
+):
+    file = Path(file)
+    output_dir = file.parent / file.stem
+    output_dir.mkdir()
+    
+    sig = at.AudioSignal(file)
+
+    # split into chunks
+    for i, sig in tqdm.tqdm(enumerate(sig.windows(
+        window_duration=max_chunk_size_s, hop_duration=max_chunk_size_s/2, 
+        preprocess=True))
+    ):
+        sig.write(output_dir / f"{i}.wav")
+
+    print(f"wrote {len(list(output_dir.glob('*.wav')))} files to {output_dir}")
+    
+    return output_dir
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+
+    with argbind.scope(args):
+        split_long_audio_file()

vampnet/scripts/utils/stage.py

@@ -0,0 +1,30 @@
+import os
+import subprocess
+from pathlib import Path
+
+import argbind
+import rich
+from audiotools.ml import Experiment
+
+
+@argbind.bind(without_prefix=True)
+def run(
+    run_dir: str = os.getenv("PATH_TO_RUNS", "runs"),
+    name: str = None,
+    recent: bool = False,
+):
+    if recent:
+        paths = sorted(Path(run_dir).iterdir(), key=os.path.getmtime)
+        paths = [p.name for p in paths if p.is_dir()]
+        if paths:
+            name = paths[-1]
+
+    with Experiment(run_dir, name) as exp:
+        exp.snapshot()
+        rich.print(f"Created a snapshot of {exp.parent_directory} at {exp.exp_dir}")
+
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+    with argbind.scope(args):
+        run()

vampnet/scripts/utils/visualize_embeddings.py

@@ -0,0 +1,276 @@
+"""
+TODO: train a linear probe
+usage:
+   python gtzan_embeddings.py --args.load conf/interface.yml --Interface.device cuda --path_to_audio /path/to/audio/labels  --output_dir /path/to/output
+"""
+from pathlib import Path
+from typing import List
+
+import audiotools as at
+from audiotools import AudioSignal
+import argbind
+import torch
+import numpy as np
+import zipfile
+import json
+
+from vampnet.interface import Interface
+import tqdm
+
+# bind the Interface to argbind
+Interface = argbind.bind(Interface)
+
+PREFIX = "vampnet-embedding-"
+
+DEBUG = False
+
+
+def smart_plotly_export(fig, save_path: Path):
+    img_format = save_path.suffix[1:]
+    if img_format == "html":
+        fig.write_html(save_path)
+    elif img_format == "bytes":
+        return fig.to_image(format="png")
+    # TODO: come back and make this prettier
+    elif img_format == "numpy":
+        import io
+        from PIL import Image
+
+        def plotly_fig2array(fig):
+            # convert Plotly fig to  an array
+            fig_bytes = fig.to_image(format="png", width=1200, height=700)
+            buf = io.BytesIO(fig_bytes)
+            img = Image.open(buf)
+            return np.asarray(img)
+
+        return plotly_fig2array(fig)
+    elif img_format == "jpeg" or "png" or "webp":
+        fig.write_image(save_path)
+    else:
+        raise ValueError("invalid image format")
+
+
+def dim_reduce(annotated_embeddings, layer, output_dir, n_components=3, method="tsne"):
+    """
+    dimensionality reduction for visualization!
+    saves an html plotly figure to save_path
+    parameters:
+        emb (np.ndarray): the samples to be reduced with shape (samples, features)
+        labels (list): list of labels for embedding
+        save_path (str): path where u wanna save ur figure
+        method (str): umap, tsne, or pca
+        title (str): title for ur figure
+    returns:
+        proj (np.ndarray): projection vector with shape (samples, dimensions)
+    """
+    import pandas as pd
+    import plotly.express as px
+
+    fig_name = f"vampnet-embeddings-layer={layer}"
+    fig_title = f"{fig_name}_{method}"
+    save_path = (output_dir / fig_name).with_suffix(".html")
+
+    if method == "umap":
+        from umap import UMAP
+        reducer = UMAP(n_components=n_components)
+    elif method == "tsne":
+        from sklearn.manifold import TSNE
+
+        reducer = TSNE(n_components=n_components)
+    elif method == "pca":
+        from sklearn.decomposition import PCA
+
+        reducer = PCA(n_components=n_components)
+    else:
+        raise ValueError(f"invalid method: {method}")
+
+    labels = [emb.label for emb in annotated_embeddings]
+    names = [emb.filename for emb in annotated_embeddings]
+    embs = [emb.embedding for emb in annotated_embeddings]
+    embs_at_layer = np.stack(embs)[:, layer, :]
+    projs = reducer.fit_transform(embs_at_layer)
+
+    df = pd.DataFrame(
+        {
+            "label": labels,
+            "name": names,
+            "x": projs[:, 0],
+            "y": projs[:, 1],
+        }
+    )
+    if n_components == 2:
+        fig = px.scatter(
+            df, x="x", y="y", color="label", hover_name="name", title=fig_title,
+        )
+
+    elif n_components == 3:
+        df['z'] = projs[:, 2]
+        fig = px.scatter_3d(
+            df, x="x", y="y", z="z", color="label", hover_name="name", title=fig_title
+        )
+    else:
+        raise ValueError(f"can't plot {n_components} components")
+
+    fig.update_traces(
+        marker=dict(size=6, line=dict(width=1, color="DarkSlateGrey")),
+        selector=dict(mode="markers"),
+    )
+
+    return smart_plotly_export(fig, save_path)
+
+
+# per JukeMIR, we want the emebddings from the middle layer?
+def vampnet_embed(sig: AudioSignal, interface: Interface, layer=10):
+    with torch.inference_mode():
+        # preprocess the signal
+        sig = interface.preprocess(sig)
+
+        # get the coarse vampnet model
+        vampnet = interface.coarse
+
+        # get the tokens
+        z = interface.encode(sig)[:, : vampnet.n_codebooks, :]
+        z_latents = vampnet.embedding.from_codes(z, interface.codec)
+
+        # do a forward pass through the model, get the embeddings
+        _z, embeddings = vampnet(z_latents, return_activations=True)
+        # print(f"got embeddings with shape {embeddings.shape}")
+        # [layer, batch, time, n_dims]
+        # [20, 1, 600ish, 768]
+
+        # squeeze batch dim (1 bc layer should be dim 0)
+        assert (
+            embeddings.shape[1] == 1
+        ), f"expected batch dim to be 1, got {embeddings.shape[0]}"
+        embeddings = embeddings.squeeze(1)
+
+        num_layers = embeddings.shape[0]
+        assert (
+            layer < num_layers
+        ), f"layer {layer} is out of bounds for model with {num_layers} layers"
+
+        # do meanpooling over the time dimension
+        embeddings = embeddings.mean(dim=-2)
+        # [20, 768]
+
+        # return the embeddings
+        return embeddings
+
+
+from dataclasses import dataclass, fields
+
+
+@dataclass
+class AnnotatedEmbedding:
+    label: str
+    filename: str
+    embedding: np.ndarray
+
+    def save(self, path):
+        """Save the Embedding object to a given path as a zip file."""
+        with zipfile.ZipFile(path, "w") as archive:
+
+            # Save numpy array
+            with archive.open("embedding.npy", "w") as f:
+                np.save(f, self.embedding)
+
+            # Save non-numpy data as json
+            non_numpy_data = {
+                f.name: getattr(self, f.name)
+                for f in fields(self)
+                if f.name != "embedding"
+            }
+            with archive.open("data.json", "w") as f:
+                f.write(json.dumps(non_numpy_data).encode("utf-8"))
+
+    @classmethod
+    def load(cls, path):
+        """Load the Embedding object from a given zip path."""
+        with zipfile.ZipFile(path, "r") as archive:
+
+            # Load numpy array
+            with archive.open("embedding.npy") as f:
+                embedding = np.load(f)
+
+            # Load non-numpy data from json
+            with archive.open("data.json") as f:
+                data = json.loads(f.read().decode("utf-8"))
+
+        return cls(embedding=embedding, **data)
+
+
+@argbind.bind(without_prefix=True)
+def main(
+    path_to_audio: str = None,
+    cache_dir: str = "./.emb_cache",
+    output_dir: str = "./vampnet_embeddings",
+    layers: List[int] = [1, 3, 5, 7, 9, 11, 13, 15, 17, 19],
+    method: str = "tsne",
+    n_components: int = 2,
+):
+    path_to_audio = Path(path_to_audio)
+    assert path_to_audio.exists(), f"{path_to_audio} does not exist"
+
+    cache_dir = Path(cache_dir)
+    output_dir = Path(output_dir)
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    # load our interface
+    # argbind will automatically load the default config,
+    interface = Interface()
+
+    # we expect path_to_audio to consist of a folder for each label, so let's get the list of labels
+    labels = [Path(x).name for x in path_to_audio.iterdir() if x.is_dir()]
+    print(f"Found {len(labels)} labels")
+    print(f"labels: {labels}")
+
+    # collect audio files, labels, and embeddings
+    annotated_embeddings = []
+    for label in labels:
+        audio_files = list(at.util.find_audio(path_to_audio / label))
+        print(f"Found {len(audio_files)} audio files for label {label}")
+
+        for audio_file in tqdm.tqdm(audio_files, desc=f"embedding label {label}"):
+            # check if we have a cached embedding for this file
+            cached_path = cache_dir / f"{label}_{audio_file.stem}.emb"
+            if cached_path.exists():
+                # if so, load it
+                if DEBUG:
+                    print(f"loading cached embedding for {cached_path.stem}")
+                embedding = AnnotatedEmbedding.load(cached_path)
+            else:
+                try:
+                    sig = AudioSignal(audio_file)
+                except Exception as e:
+                    print(f"failed to load {audio_file.name} with error {e}")
+                    print(f"skipping {audio_file.name}")
+                    continue
+
+                # gets the embedding
+                emb = vampnet_embed(sig, interface).cpu().numpy()
+
+                # create an embedding we can save/load
+                embedding = AnnotatedEmbedding(
+                    label=label, filename=audio_file.name, embedding=emb
+                )
+
+                # cache the embeddings
+                cached_path.parent.mkdir(exist_ok=True, parents=True)
+                embedding.save(cached_path)
+            annotated_embeddings.append(embedding)
+
+    # now, let's do a dim reduction on the embeddings and visualize them.
+    for layer in tqdm.tqdm(layers, desc="dim reduction"):
+        dim_reduce(
+            annotated_embeddings,
+            layer,
+            output_dir=output_dir,
+            n_components=n_components,
+            method=method,
+        )
+
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+    with argbind.scope(args):
+        main()

vampnet/scripts/utils/xeno-canto-dl.py

@@ -0,0 +1,234 @@
+from xenopy import Query
+
+
+SPECIES = [
+    "American Robin",
+    "Northern Cardinal",
+    "Mourning Dove",
+    "American Crow",
+    "Baltimore Oriole",
+    "Blue Jay",
+    "Eastern Bluebird",
+    "House Finch",
+    "American Goldfinch",
+    "House Sparrow",
+    "Song Sparrow",
+    "Tufted Titmouse",
+    "White-breasted Nuthatch",
+    "European Starling",
+    "American Redstart",
+    "Red-winged Blackbird",
+    "Brown-headed Cowbird",
+    "Common Grackle",
+    "Boat-tailed Grackle",
+    "Common Yellowthroat",
+    "Northern Mockingbird",
+    "Carolina Wren",
+    "Eastern Meadowlark",
+    "Chipping Sparrow",
+    "Tree Swallow",
+    "Barn Swallow",
+    "Cliff Swallow",
+    "Pine Siskin",
+    "Indigo Bunting",
+    "Eastern Towhee",
+    "Carolina Chickadee",
+    "Great Crested Flycatcher",
+    "Eastern Wood-Pewee",
+    "Ovenbird",
+    "Northern Flicker",
+    "Red-eyed Vireo",
+    "American Woodcock",
+    "Eastern Phoebe",
+    "Downy Woodpecker",
+    "Scarlet Tanager",
+    "Yellow Warbler",
+    "White-eyed Vireo",
+    "Common Loon",
+    "White-throated Sparrow",
+    "Yellow-throated Vireo",
+    "Great Blue Heron",
+    "Belted Kingfisher",
+    "Pied-billed Grebe",
+    "Wild Turkey",
+    "Wood Thrush",
+    "Rose-breasted Grosbeak",
+    "Field Sparrow",
+    "Hooded Warbler",
+    "Northern Parula",
+    "Chestnut-sided Warbler",
+    "Blue-winged Warbler",
+    "Red-bellied Woodpecker",
+    "Yellow-billed Cuckoo",
+    "Gray Catbird",
+    "Northern Saw-whet Owl",
+    "Osprey",
+    "Common Nighthawk",
+    "Broad-winged Hawk",
+    "Black-throated Green Warbler",
+    "Great Horned Owl",
+    "Common Raven",
+    "Barred Owl",
+    "Canada Warbler",
+    "Magnolia Warbler",
+    "Black-and-white Warbler",
+    "Eastern Kingbird",
+    "Swainson's Thrush",
+    "Worm-eating Warbler",
+    "Prairie Warbler",
+    "Baltimore Oriole",
+    "Black-throated Blue Warbler",
+    "Louisiana Waterthrush",
+    "Blackburnian Warbler",
+    "Black-capped Chickadee",
+    "Cerulean Warbler",
+    "Red-shouldered Hawk",
+    "Cooper's Hawk",
+    "Yellow-throated Warbler",
+    "Blue-headed Vireo",
+    "Blackpoll Warbler",
+    "Ruffed Grouse",
+    "Kentucky Warbler",
+    "Hermit Thrush",
+    "Cedar Waxwing",
+    "Eastern Screech-Owl",
+    "Northern Goshawk",
+    "Green Heron",
+    "Red-tailed Hawk",
+    "Black Vulture",
+    "Hairy Woodpecker",
+    "Golden-crowned Kinglet",
+    "Ruby-crowned Kinglet",
+    "Bicknell's Thrush",
+    "Blue-gray Gnatcatcher",
+    "Veery",
+    "Pileated Woodpecker",
+    "Purple Finch",
+    "White-crowned Sparrow",
+    "Snow Bunting",
+    "Pine Grosbeak",
+    "American Tree Sparrow",
+    "Dark-eyed Junco",
+    "Snowy Owl",
+    "White-winged Crossbill",
+    "Red Crossbill",
+    "Common Redpoll",
+    "Northern Shrike",
+    "Northern Harrier",
+    "Rough-legged Hawk",
+    "Long-eared Owl",
+    "Evening Grosbeak",
+    "Northern Pintail",
+    "American Black Duck",
+    "Mallard",
+    "Canvasback",
+    "Redhead",
+    "Ring-necked Duck",
+    "Greater Scaup",
+    "Lesser Scaup",
+    "Bufflehead",
+    "Common Goldeneye",
+    "Hooded Merganser",
+    "Common Merganser",
+    "Red-breasted Merganser",
+    "Ruddy Duck",
+    "Wood Duck",
+    "Gadwall",
+    "American Wigeon",
+    "Northern Shoveler",
+    "Green-winged Teal",
+    "Blue-winged Teal",
+    "Cinnamon Teal",
+    "Ringed Teal",
+    "Cape Teal",
+    "Northern Fulmar",
+    "Yellow-billed Loon",
+    "Red-throated Loon",
+    "Arctic Loon",
+    "Pacific Loon",
+    "Horned Grebe",
+    "Red-necked Grebe",
+    "Eared Grebe",
+    "Western Grebe",
+    "Clark's Grebe",
+    "Double-crested Cormorant",
+    "Pelagic Cormorant",
+    "Great Cormorant",
+    "American White Pelican",
+    "Brown Pelican",
+    "Brandt's Cormorant",
+    "Least Bittern",
+    "Great Egret",
+    "Snowy Egret",
+    "Little Blue Heron",
+    "Tricolored Heron",
+    "Reddish Egret",
+    "Black-crowned Night-Heron",
+    "Yellow-crowned Night-Heron",
+    "White Ibis",
+    "Glossy Ibis",
+    "Roseate Spoonbill",
+    "Wood Stork",
+    "Black-bellied Whistling-Duck",
+    "Fulvous Whistling-Duck",
+    "Greater White-fronted Goose",
+    "Snow Goose",
+    "Ross's Goose",
+    "Canada Goose",
+    "Brant",
+    "Mute Swan",
+    "Tundra Swan",
+    "Whooper Swan",
+    "Sandhill Crane",
+    "Black-necked Stilt",
+    "American Avocet",
+    "Northern Jacana",
+    "Greater Yellowlegs",
+    "Lesser Yellowlegs",
+    "Willet",
+    "Spotted Sandpiper",
+    "Upland Sandpiper",
+    "Whimbrel",
+    "Long-billed Curlew",
+    "Marbled Godwit",
+    "Ruddy Turnstone",
+    "Red Knot",
+    "Sanderling",
+    "Semipalmated Sandpiper",
+    "Western Sandpiper",
+    "Least Sandpiper",
+    "White-rumped Sandpiper",
+    "Baird's Sandpiper",
+    "Pectoral Sandpiper",
+    "Dunlin",
+    "Buff-breasted Sandpiper",
+    "Short-billed Dowitcher",
+    "Long-billed Dowitcher",
+    "Common Snipe",
+    "American Woodcock",
+    "Wilson's Phalarope",
+    "Red-necked Phalarope",
+    "Red Phalarope"
+]
+
+from pathlib import Path
+
+def remove_spaces(s):
+    return s.replace(" ", "")
+
+for species in SPECIES: 
+    if Path("/media/CHONK/hugo/xeno-canto-full/" + remove_spaces(species)).exists():
+        continue
+    try:
+        q = Query(
+            name=species, q="A", length="10-30", 
+            )
+
+        # retrieve metadata
+        metafiles = q.retrieve_meta(verbose=True)
+        # retrieve recordings
+        q.retrieve_recordings(multiprocess=True, nproc=10, attempts=10, outdir="/media/CHONK/hugo/xeno-canto-full/")
+
+    except:
+        print("Failed to download " + species)
+        continue

vampnet/setup.py

@@ -0,0 +1,47 @@
+from setuptools import find_packages
+from setuptools import setup
+
+with open("README.md") as f:
+    long_description = f.read()
+
+setup(
+    name="vampnet",
+    version="0.0.1",
+    classifiers=[
+        "Intended Audience :: Developers",
+        "Natural Language :: English",
+        "Programming Language :: Python :: 3.7",
+        "Topic :: Artistic Software",
+        "Topic :: Multimedia",
+        "Topic :: Multimedia :: Sound/Audio",
+        "Topic :: Multimedia :: Sound/Audio :: Editors",
+        "Topic :: Software Development :: Libraries",
+    ],
+    description="Generative Music Modeling.",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    author="Hugo Flores García, Prem Seetharaman",
+    author_email="hfgacrcia@descript.com",
+    url="https://github.com/hugofloresgarcia/vampnet",
+    license="MIT",
+    packages=find_packages(),
+    setup_requires=[
+        "Cython",
+    ],
+    install_requires=[
+        "Cython",  # Added by WAM because it seems to be needed by this repo?
+        "torch",
+        "pydantic==2.10.6",
+        "argbind>=0.3.2",
+        "numpy<1.24",
+        "wavebeat @ git+https://github.com/hugofloresgarcia/wavebeat",
+        "lac @ git+https://github.com/hugofloresgarcia/lac.git",
+        "descript-audiotools @ git+https://github.com/hugofloresgarcia/audiotools.git",
+        "gradio",
+        "loralib",
+        "torch_pitch_shift",
+        "plotly",  # Added by WAM for clustering (see https://github.com/hugofloresgarcia/vampnet/issues/20)
+        "pyharp",
+        
+    ],
+)

vampnet/vampnet/__init__.py

@@ -0,0 +1,6 @@
+
+from . import modules
+from . import scheduler
+from .interface import Interface
+
+__version__ = "0.0.1"

vampnet/vampnet/beats.py

@@ -0,0 +1,250 @@
+import json
+import logging
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+from typing import List
+from typing import Tuple
+from typing import Union
+
+import librosa
+import torch
+import numpy as np
+from audiotools import AudioSignal
+
+
+logging.basicConfig(level=logging.INFO)
+
+###################
+# beat sync utils #
+###################
+
+AGGREGATOR_REGISTRY = {
+    "mean": np.mean,
+    "median": np.median,
+    "max": np.max,
+    "min": np.min,
+}
+
+
+def list_aggregators() -> list:
+    return list(AGGREGATOR_REGISTRY.keys())
+
+
+@dataclass
+class TimeSegment:
+    start: float
+    end: float
+
+    @property
+    def duration(self):
+        return self.end - self.start
+
+    def __str__(self) -> str:
+        return f"{self.start} - {self.end}"
+
+    def find_overlapping_segment(
+        self, segments: List["TimeSegment"]
+    ) -> Union["TimeSegment", None]:
+        """Find the first segment that overlaps with this segment, or None if no segment overlaps"""
+        for s in segments:
+            if s.start <= self.start and s.end >= self.end:
+                return s
+        return None
+
+
+def mkdir(path: Union[Path, str]) -> Path:
+    p = Path(path)
+    p.mkdir(parents=True, exist_ok=True)
+    return p
+
+
+
+###################
+#    beat data    #
+###################
+@dataclass
+class BeatSegment(TimeSegment):
+    downbeat: bool = False  # if there's a downbeat on the start_time
+
+
+class Beats:
+    def __init__(self, beat_times, downbeat_times):
+        if isinstance(beat_times, np.ndarray):
+            beat_times = beat_times.tolist()
+        if isinstance(downbeat_times, np.ndarray):
+            downbeat_times = downbeat_times.tolist()
+        self._beat_times = beat_times
+        self._downbeat_times = downbeat_times
+        self._use_downbeats = False
+
+    def use_downbeats(self, use_downbeats: bool = True):
+        """use downbeats instead of beats when calling beat_times"""
+        self._use_downbeats = use_downbeats
+
+    def beat_segments(self, signal: AudioSignal) -> List[BeatSegment]:
+        """
+        segments a song into time segments corresponding to beats.
+        the first segment starts at 0 and ends at the first beat time.
+        the last segment starts at the last beat time and ends at the end of the song.
+        """
+        beat_times = self._beat_times.copy()
+        downbeat_times = self._downbeat_times
+        beat_times.insert(0, 0)
+        beat_times.append(signal.signal_duration)
+
+        downbeat_ids = np.intersect1d(beat_times, downbeat_times, return_indices=True)[
+            1
+        ]
+        is_downbeat = [
+            True if i in downbeat_ids else False for i in range(len(beat_times))
+        ]
+        segments = [
+            BeatSegment(start_time, end_time, downbeat)
+            for start_time, end_time, downbeat in zip(
+                beat_times[:-1], beat_times[1:], is_downbeat
+            )
+        ]
+        return segments
+
+    def get_beats(self) -> np.ndarray:
+        """returns an array of beat times, in seconds
+        if downbeats is True, returns an array of downbeat times, in seconds
+        """
+        return np.array(
+            self._downbeat_times if self._use_downbeats else self._beat_times
+        )
+
+    @property
+    def beat_times(self) -> np.ndarray:
+        """return beat times"""
+        return np.array(self._beat_times)
+
+    @property
+    def downbeat_times(self) -> np.ndarray:
+        """return downbeat times"""
+        return np.array(self._downbeat_times)
+
+    def beat_times_to_feature_frames(
+        self, signal: AudioSignal, features: np.ndarray
+    ) -> np.ndarray:
+        """convert beat times to frames, given an array of time-varying features"""
+        beat_times = self.get_beats()
+        beat_frames = (
+            beat_times * signal.sample_rate / signal.signal_length * features.shape[-1]
+        ).astype(np.int64)
+        return beat_frames
+
+    def sync_features(
+        self, feature_frames: np.ndarray, features: np.ndarray, aggregate="median"
+    ) -> np.ndarray:
+        """sync features to beats"""
+        if aggregate not in AGGREGATOR_REGISTRY:
+            raise ValueError(f"unknown aggregation method {aggregate}")
+
+        return librosa.util.sync(
+            features, feature_frames, aggregate=AGGREGATOR_REGISTRY[aggregate]
+        )
+
+    def to_json(self) -> dict:
+        """return beats and downbeats as json"""
+        return {
+            "beats": self._beat_times,
+            "downbeats": self._downbeat_times,
+            "use_downbeats": self._use_downbeats,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        """load beats and downbeats from json"""
+        inst = cls(data["beats"], data["downbeats"])
+        inst.use_downbeats(data["use_downbeats"])
+        return inst
+
+    def save(self, output_dir: Path):
+        """save beats and downbeats to json"""
+        mkdir(output_dir)
+        with open(output_dir / "beats.json", "w") as f:
+            json.dump(self.to_json(), f)
+
+    @classmethod
+    def load(cls, input_dir: Path):
+        """load beats and downbeats from json"""
+        beats_file = Path(input_dir) / "beats.json"
+        with open(beats_file, "r") as f:
+            data = json.load(f)
+        return cls.from_dict(data)
+
+
+###################
+#  beat tracking  #
+###################
+
+
+class BeatTracker:
+    def extract_beats(self, signal: AudioSignal) -> Tuple[np.ndarray, np.ndarray]:
+        """extract beats from an audio signal"""
+        raise NotImplementedError
+
+    def __call__(self, signal: AudioSignal) -> Beats:
+        """extract beats from an audio signal
+        NOTE: if the first beat (and/or downbeat) is detected within the first 100ms of the audio,
+        it is discarded. This is to avoid empty bins with no beat synced features in the first beat.
+        Args:
+            signal (AudioSignal): signal to beat track
+        Returns:
+            Tuple[np.ndarray, np.ndarray]: beats and downbeats
+        """
+        beats, downbeats = self.extract_beats(signal)
+        return Beats(beats, downbeats)
+
+
+class WaveBeat(BeatTracker):
+    def __init__(self, ckpt_path: str = "checkpoints/wavebeat", device: str = "cpu"):
+        from wavebeat.dstcn import dsTCNModel
+
+        model = dsTCNModel.load_from_checkpoint(ckpt_path, map_location=torch.device(device))
+        model.eval()
+
+        self.device = device
+        self.model = model
+
+    def extract_beats(self, signal: AudioSignal) -> Tuple[np.ndarray, np.ndarray]:
+        """returns beat and downbeat times, in  seconds"""
+        # extract beats
+        beats, downbeats = self.model.predict_beats_from_array(
+            audio=signal.audio_data.squeeze(0),
+            sr=signal.sample_rate,
+            use_gpu=self.device != "cpu",
+        )
+
+        return beats, downbeats
+
+
+class MadmomBeats(BeatTracker):
+    def __init__(self):
+        raise NotImplementedError
+
+    def extract_beats(self, signal: AudioSignal) -> Tuple[np.ndarray, np.ndarray]:
+        """returns beat and downbeat times, in  seconds"""
+        pass
+
+
+BEAT_TRACKER_REGISTRY = {
+    "wavebeat": WaveBeat,
+    "madmom": MadmomBeats,
+}
+
+
+def list_beat_trackers() -> list:
+    return list(BEAT_TRACKER_REGISTRY.keys())
+
+
+def load_beat_tracker(beat_tracker: str, **kwargs) -> BeatTracker:
+    if beat_tracker not in BEAT_TRACKER_REGISTRY:
+        raise ValueError(
+            f"Unknown beat tracker {beat_tracker}. Available: {list_beat_trackers()}"
+        )
+
+    return BEAT_TRACKER_REGISTRY[beat_tracker](**kwargs)

vampnet/vampnet/interface.py

@@ -0,0 +1,422 @@
+import os
+from pathlib import Path
+import math
+
+import torch
+import numpy as np
+from audiotools import AudioSignal
+import tqdm
+
+from .modules.transformer import VampNet
+from .beats import WaveBeat
+from .mask import *
+
+# from dac.model.dac import DAC
+from lac.model.lac import LAC as DAC
+
+
+def signal_concat(
+    audio_signals: list,
+):
+    audio_data = torch.cat([x.audio_data for x in audio_signals], dim=-1)
+
+    return AudioSignal(audio_data, sample_rate=audio_signals[0].sample_rate)
+
+
+def _load_model(
+    ckpt: str, 
+    lora_ckpt: str = None,
+    device: str = "cpu",
+    chunk_size_s: int = 10,
+):
+    # we need to set strict to False if the model has lora weights to add later
+    model = VampNet.load(location=Path(ckpt), map_location="cpu", strict=False)
+
+    # load lora weights if needed
+    if lora_ckpt is not None:
+        if not Path(lora_ckpt).exists():
+            should_cont = input(
+                f"lora checkpoint {lora_ckpt} does not exist. continue? (y/n) "
+            )
+            if should_cont != "y":
+                raise Exception("aborting")
+        else:
+            model.load_state_dict(torch.load(lora_ckpt, map_location="cpu"), strict=False)
+
+    model.to(device)
+    model.eval()
+    model.chunk_size_s = chunk_size_s
+    return model
+
+
+
+class Interface(torch.nn.Module):
+    def __init__(
+        self,
+        coarse_ckpt: str = None,
+        coarse_lora_ckpt: str = None,
+        coarse2fine_ckpt: str = None,
+        coarse2fine_lora_ckpt: str = None,
+        codec_ckpt: str = None,
+        wavebeat_ckpt: str = None,
+        device: str = "cpu",
+        coarse_chunk_size_s: int =  10, 
+        coarse2fine_chunk_size_s: int =  3,
+    ):
+        super().__init__()
+        assert codec_ckpt is not None, "must provide a codec checkpoint"
+        self.codec = DAC.load(Path(codec_ckpt))
+        self.codec.eval()
+        self.codec.to(device)
+
+        assert coarse_ckpt is not None, "must provide a coarse checkpoint"
+        self.coarse = _load_model(
+            ckpt=coarse_ckpt,
+            lora_ckpt=coarse_lora_ckpt,
+            device=device,
+            chunk_size_s=coarse_chunk_size_s,
+        )
+
+        # check if we have a coarse2fine ckpt
+        if coarse2fine_ckpt is not None:
+            self.c2f = _load_model(
+                ckpt=coarse2fine_ckpt,
+                lora_ckpt=coarse2fine_lora_ckpt,
+                device=device,
+                chunk_size_s=coarse2fine_chunk_size_s,
+            )
+        else:
+            self.c2f = None
+
+        if wavebeat_ckpt is not None:
+            print(f"loading wavebeat from {wavebeat_ckpt}")
+            self.beat_tracker = WaveBeat(wavebeat_ckpt)
+            self.beat_tracker.model.to(device)
+        else:
+            self.beat_tracker = None
+
+        self.device = device
+
+    def lora_load(
+        self, 
+        coarse_ckpt: str = None,
+        c2f_ckpt: str = None,
+        full_ckpts: bool = False,
+    ):
+        if full_ckpts:
+            if coarse_ckpt is not None:
+                self.coarse = _load_model(
+                    ckpt=coarse_ckpt,  
+                    device=self.device,
+                    chunk_size_s=self.coarse.chunk_size_s,
+                )
+            if c2f_ckpt is not None:
+                self.c2f = _load_model(
+                    ckpt=c2f_ckpt,
+                    device=self.device,
+                    chunk_size_s=self.c2f.chunk_size_s,
+                )
+        else:
+            if coarse_ckpt is not None:
+                self.coarse.to("cpu")
+                state_dict = torch.load(coarse_ckpt, map_location="cpu")
+                print(f"loading coarse from {coarse_ckpt}")
+                self.coarse.load_state_dict(state_dict, strict=False)
+                self.coarse.to(self.device)
+            if c2f_ckpt is not None:
+                self.c2f.to("cpu")
+                state_dict = torch.load(c2f_ckpt, map_location="cpu")
+                print(f"loading c2f from {c2f_ckpt}")
+                self.c2f.load_state_dict(state_dict, strict=False)
+                self.c2f.to(self.device)
+        
+    def s2t(self, seconds: float):
+        """seconds to tokens"""
+        if isinstance(seconds, np.ndarray):
+            return np.ceil(seconds * self.codec.sample_rate / self.codec.hop_length)
+        else:
+            return math.ceil(seconds * self.codec.sample_rate / self.codec.hop_length)
+
+    def s2t2s(self, seconds: float):
+        """seconds to tokens to seconds"""
+        return self.t2s(self.s2t(seconds))
+    
+    def t2s(self, tokens: int):
+        """tokens to seconds"""
+        return tokens * self.codec.hop_length / self.codec.sample_rate
+
+    def to(self, device):
+        self.device = device
+        self.coarse.to(device)
+        self.codec.to(device)
+
+        if self.c2f is not None:
+            self.c2f.to(device)
+
+        if self.beat_tracker is not None:
+            self.beat_tracker.model.to(device)
+        return self
+
+    def to_signal(self, z: torch.Tensor):
+        return self.coarse.to_signal(z, self.codec)
+    
+    def preprocess(self, signal: AudioSignal):
+        signal = (
+            signal.clone()
+            .resample(self.codec.sample_rate)
+            .to_mono()
+            .normalize(-24)
+            .ensure_max_of_audio(1.0)
+        )
+        return signal
+    
+    @torch.inference_mode()
+    def encode(self, signal: AudioSignal):
+        signal = self.preprocess(signal).to(self.device)
+        z = self.codec.encode(signal.samples, signal.sample_rate)["codes"]
+        return z
+
+    def snap_to_beats(
+        self, 
+        signal: AudioSignal
+    ):
+        assert hasattr(self, "beat_tracker"), "No beat tracker loaded"
+        beats, downbeats = self.beat_tracker.extract_beats(signal)
+        
+        # trim the signa around the first beat time
+        samples_begin = int(beats[0] * signal.sample_rate )
+        samples_end = int(beats[-1] * signal.sample_rate)
+        print(beats[0])
+        signal = signal.clone().trim(samples_begin, signal.length - samples_end)
+
+        return signal
+
+    def make_beat_mask(self, 
+            signal: AudioSignal, 
+            before_beat_s: float = 0.0,
+            after_beat_s: float = 0.02,
+            mask_downbeats: bool = True,
+            mask_upbeats: bool = True,
+            downbeat_downsample_factor: int = None,
+            beat_downsample_factor: int = None,
+            dropout: float = 0.0,
+            invert: bool = True,
+    ):
+        """make a beat synced mask. that is, make a mask that 
+        places 1s at and around the beat, and 0s everywhere else. 
+        """
+        assert self.beat_tracker is not None, "No beat tracker loaded"
+
+        # get the beat times
+        beats, downbeats = self.beat_tracker.extract_beats(signal)
+
+        # get the beat indices in z
+        beats_z, downbeats_z = self.s2t(beats), self.s2t(downbeats)
+
+        # remove downbeats from beats
+        beats_z = torch.tensor(beats_z)[~torch.isin(torch.tensor(beats_z), torch.tensor(downbeats_z))]
+        beats_z = beats_z.tolist()
+        downbeats_z = downbeats_z.tolist()
+
+        # make the mask 
+        seq_len = self.s2t(signal.duration)
+        mask = torch.zeros(seq_len, device=self.device)
+        
+        mask_b4 = self.s2t(before_beat_s)
+        mask_after = self.s2t(after_beat_s)
+
+        if beat_downsample_factor is not None:
+            if beat_downsample_factor < 1:
+                raise ValueError("mask_beat_downsample_factor must be >= 1 or None")
+        else:
+            beat_downsample_factor = 1
+
+        if downbeat_downsample_factor is not None:
+            if downbeat_downsample_factor < 1:
+                raise ValueError("mask_beat_downsample_factor must be >= 1 or None")
+        else:
+            downbeat_downsample_factor = 1
+
+        beats_z = beats_z[::beat_downsample_factor]
+        downbeats_z = downbeats_z[::downbeat_downsample_factor]
+        print(f"beats_z: {len(beats_z)}")
+        print(f"downbeats_z: {len(downbeats_z)}")
+    
+        if mask_upbeats:
+            for beat_idx in beats_z:
+                _slice = int(beat_idx - mask_b4), int(beat_idx + mask_after)
+                num_steps = mask[_slice[0]:_slice[1]].shape[0]
+                _m = torch.ones(num_steps, device=self.device)
+                _m_mask = torch.bernoulli(_m * (1 - dropout))
+                _m = _m * _m_mask.long()
+                
+                mask[_slice[0]:_slice[1]] = _m
+
+        if mask_downbeats:
+            for downbeat_idx in downbeats_z:
+                _slice = int(downbeat_idx - mask_b4), int(downbeat_idx + mask_after)
+                num_steps = mask[_slice[0]:_slice[1]].shape[0]
+                _m = torch.ones(num_steps, device=self.device)
+                _m_mask = torch.bernoulli(_m * (1 - dropout))
+                _m = _m * _m_mask.long()
+                
+                mask[_slice[0]:_slice[1]] = _m
+        
+        mask = mask.clamp(0, 1)
+        if invert:
+            mask = 1 - mask
+        
+        mask = mask[None, None, :].bool().long()
+        if self.c2f is not None:
+            mask = mask.repeat(1, self.c2f.n_codebooks, 1)
+        else:
+            mask = mask.repeat(1, self.coarse.n_codebooks, 1)
+        return mask
+        
+    def coarse_to_fine(
+        self, 
+        z: torch.Tensor,
+        mask: torch.Tensor = None,
+        **kwargs
+    ):
+        assert self.c2f is not None, "No coarse2fine model loaded"
+        length = z.shape[-1]
+        chunk_len = self.s2t(self.c2f.chunk_size_s)
+        n_chunks = math.ceil(z.shape[-1] / chunk_len)
+
+        # zero pad to chunk_len
+        if length % chunk_len != 0:
+            pad_len = chunk_len - (length % chunk_len)
+            z = torch.nn.functional.pad(z, (0, pad_len))
+            mask = torch.nn.functional.pad(mask, (0, pad_len)) if mask is not None else None
+
+        n_codebooks_to_append = self.c2f.n_codebooks - z.shape[1]
+        if n_codebooks_to_append > 0:
+            z = torch.cat([
+                z,
+                torch.zeros(z.shape[0], n_codebooks_to_append, z.shape[-1]).long().to(self.device)
+            ], dim=1)
+
+        # set the mask to 0 for all conditioning codebooks
+        if mask is not None:
+            mask = mask.clone()
+            mask[:, :self.c2f.n_conditioning_codebooks, :] = 0
+
+        fine_z = []
+        for i in range(n_chunks):
+            chunk = z[:, :, i * chunk_len : (i + 1) * chunk_len]
+            mask_chunk = mask[:, :, i * chunk_len : (i + 1) * chunk_len] if mask is not None else None
+            
+            chunk = self.c2f.generate(
+                codec=self.codec,
+                time_steps=chunk_len,
+                start_tokens=chunk,
+                return_signal=False,
+                mask=mask_chunk,
+                **kwargs
+            )
+            fine_z.append(chunk)
+
+        fine_z = torch.cat(fine_z, dim=-1)
+        return fine_z[:, :, :length].clone()
+    
+    def coarse_vamp(
+        self, 
+        z, 
+        mask,
+        return_mask=False,
+        gen_fn=None,
+        **kwargs
+    ):
+        # coarse z
+        cz = z[:, : self.coarse.n_codebooks, :].clone()
+        assert cz.shape[-1] <= self.s2t(self.coarse.chunk_size_s), f"the sequence of tokens provided must match the one specified in the coarse chunk size, but got {cz.shape[-1]} and {self.s2t(self.coarse.chunk_size_s)}"
+
+        mask = mask[:, : self.coarse.n_codebooks, :]
+
+        cz_masked, mask = apply_mask(cz, mask, self.coarse.mask_token)
+        cz_masked = cz_masked[:, : self.coarse.n_codebooks, :]
+
+        gen_fn = gen_fn or self.coarse.generate
+        c_vamp = gen_fn(
+            codec=self.codec,
+            time_steps=cz.shape[-1],
+            start_tokens=cz,
+            mask=mask, 
+            return_signal=False,
+            **kwargs
+        )
+
+        # add the fine codes back in
+        c_vamp = torch.cat(
+            [c_vamp, z[:, self.coarse.n_codebooks :, :]], 
+            dim=1
+        )
+
+        if return_mask:
+            return c_vamp, cz_masked
+        
+        return c_vamp
+
+
+if __name__ == "__main__":
+    import audiotools as at
+    import logging
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    torch.set_printoptions(threshold=10000)
+    at.util.seed(42)
+
+    interface = Interface(
+        coarse_ckpt="./models/vampnet/coarse.pth", 
+        coarse2fine_ckpt="./models/vampnet/c2f.pth", 
+        codec_ckpt="./models/vampnet/codec.pth",
+        device="cuda", 
+        wavebeat_ckpt="./models/wavebeat.pth"
+    )
+
+
+    sig = at.AudioSignal('assets/example.wav')
+
+    z = interface.encode(sig)
+    breakpoint()
+
+    # mask = linear_random(z, 1.0)
+    # mask = mask_and(
+    #     mask, periodic_mask(
+    #         z,
+    #         32,
+    #         1,
+    #         random_roll=True
+    #     )
+    # )
+
+    # mask = interface.make_beat_mask(
+    #     sig, 0.0, 0.075
+    # )
+    # mask = dropout(mask, 0.0)
+    # mask = codebook_unmask(mask, 0)
+
+    mask = inpaint(z, n_prefix=100, n_suffix=100)
+    
+    zv, mask_z = interface.coarse_vamp(
+        z, 
+        mask=mask,
+        sampling_steps=36,
+        temperature=8.0,
+        return_mask=True, 
+        gen_fn=interface.coarse.generate
+    )
+    
+
+    use_coarse2fine = True
+    if use_coarse2fine: 
+        zv = interface.coarse_to_fine(zv, temperature=0.8, mask=mask)
+        breakpoint()
+
+    mask = interface.to_signal(mask_z).cpu()
+
+    sig = interface.to_signal(zv).cpu()
+    print("done")
+
+        

vampnet/vampnet/mask.py

@@ -0,0 +1,242 @@
+from typing import Optional
+
+import torch
+from audiotools import AudioSignal
+
+from .util import scalar_to_batch_tensor
+
+def _gamma(r):
+    return (r * torch.pi / 2).cos().clamp(1e-10, 1.0)
+
+def _invgamma(y):
+    if not torch.is_tensor(y):
+        y = torch.tensor(y)[None]
+    return 2 * y.acos() / torch.pi
+
+def full_mask(x: torch.Tensor):
+    assert x.ndim == 3, "x must be (batch, n_codebooks, seq)"
+    return torch.ones_like(x).long()
+
+def empty_mask(x: torch.Tensor):
+    assert x.ndim == 3, "x must be (batch, n_codebooks, seq)"
+    return torch.zeros_like(x).long()
+
+def apply_mask(
+        x: torch.Tensor, 
+        mask: torch.Tensor, 
+        mask_token: int
+    ):
+    assert mask.ndim == 3, "mask must be (batch, n_codebooks, seq), but got {mask.ndim}"
+    assert mask.shape == x.shape, f"mask must be same shape as x, but got {mask.shape} and {x.shape}" 
+    assert mask.dtype == torch.long, "mask must be long dtype, but got {mask.dtype}"
+    assert ~torch.any(mask > 1), "mask must be binary"
+    assert ~torch.any(mask < 0), "mask must be binary"
+
+    fill_x = torch.full_like(x, mask_token)
+    x = x * (1 - mask) + fill_x * mask
+
+    return x, mask
+
+def random(
+    x: torch.Tensor,
+    r: torch.Tensor
+):
+    assert x.ndim == 3, "x must be (batch, n_codebooks, seq)"
+    if not isinstance(r, torch.Tensor):
+        r = scalar_to_batch_tensor(r, x.shape[0]).to(x.device)
+
+    r = _gamma(r)[:, None, None]
+    probs = torch.ones_like(x) * r
+
+    mask = torch.bernoulli(probs)
+    mask = mask.round().long()
+
+    return mask
+
+def linear_random(
+    x: torch.Tensor,
+    r: torch.Tensor,
+):
+    assert x.ndim == 3, "x must be (batch, n_codebooks, seq)"
+    if not isinstance(r, torch.Tensor):
+        r = scalar_to_batch_tensor(r, x.shape[0]).to(x.device).float()
+
+    probs = torch.ones_like(x).to(x.device).float()
+    # expand to batch and codebook dims
+    probs = probs.expand(x.shape[0], x.shape[1], -1)
+    probs = probs * r
+
+    mask = torch.bernoulli(probs)
+    mask = mask.round().long()
+
+    return mask
+
+def inpaint(x: torch.Tensor, 
+    n_prefix,
+    n_suffix,
+):
+    assert n_prefix is not None
+    assert n_suffix is not None
+    
+    mask = full_mask(x)
+
+    # if we have a prefix or suffix, set their mask prob to 0
+    if n_prefix > 0:
+        if not isinstance(n_prefix, torch.Tensor):
+            n_prefix = scalar_to_batch_tensor(n_prefix, x.shape[0]).to(x.device) 
+        for i, n in enumerate(n_prefix):
+            if n > 0:
+                mask[i, :, :n] = 0.0
+    if n_suffix > 0:
+        if not isinstance(n_suffix, torch.Tensor):
+            n_suffix = scalar_to_batch_tensor(n_suffix, x.shape[0]).to(x.device)
+        for i, n in enumerate(n_suffix):
+            if n > 0:
+                mask[i, :, -n:] = 0.0
+
+    
+    return mask
+
+def periodic_mask(x: torch.Tensor, 
+                period: int, width: int = 1, 
+                random_roll=False,
+    ):
+    mask = full_mask(x)
+    if period == 0:
+        return mask
+
+    if not isinstance(period, torch.Tensor):
+        period = scalar_to_batch_tensor(period, x.shape[0])
+    for i, factor in enumerate(period):
+        if factor == 0:
+            continue
+        for j in range(mask.shape[-1]):
+            if j % factor == 0:
+                # figure out how wide the mask should be
+                j_start = max(0, j - width // 2  )
+                j_end = min(mask.shape[-1] - 1, j + width // 2 ) + 1 
+                # flip a coin for each position in the mask
+                j_mask = torch.bernoulli(torch.ones(j_end - j_start))
+                assert torch.all(j_mask == 1)
+                j_fill = torch.ones_like(j_mask) * (1 - j_mask)
+                assert torch.all(j_fill == 0)
+                # fill
+                mask[i, :, j_start:j_end] = j_fill
+    if random_roll:
+        # add a random offset to the mask
+        offset = torch.randint(0, period[0], (1,))
+        mask = torch.roll(mask, offset.item(), dims=-1)
+
+    return mask
+
+def codebook_unmask(
+    mask: torch.Tensor, 
+    n_conditioning_codebooks: int
+):
+    if n_conditioning_codebooks == None:
+        return mask
+    # if we have any conditioning codebooks, set their mask  to 0
+    mask = mask.clone()
+    mask[:, :n_conditioning_codebooks, :] = 0
+    return mask
+
+def codebook_mask(mask: torch.Tensor, start: int):
+    mask = mask.clone()
+    mask[:, start:, :] = 1
+    return mask
+
+def mask_and(
+    mask1: torch.Tensor, 
+    mask2: torch.Tensor
+):
+    assert mask1.shape == mask2.shape, "masks must be same shape"
+    return torch.min(mask1, mask2)
+
+def dropout(
+    mask: torch.Tensor,
+    p: float,
+):
+    assert 0 <= p <= 1, "p must be between 0 and 1"
+    assert mask.max() <= 1, "mask must be binary"
+    assert mask.min() >= 0, "mask must be binary"
+    mask = (~mask.bool()).float()
+    mask = torch.bernoulli(mask * (1 - p))
+    mask = ~mask.round().bool()
+    return mask.long()
+
+def mask_or(
+    mask1: torch.Tensor, 
+    mask2: torch.Tensor
+):
+    assert mask1.shape == mask2.shape, f"masks must be same shape, but got {mask1.shape} and {mask2.shape}"
+    assert mask1.max() <= 1, "mask1 must be binary"
+    assert mask2.max() <= 1, "mask2 must be binary"
+    assert mask1.min() >= 0, "mask1 must be binary"
+    assert mask2.min() >= 0, "mask2 must be binary"
+    return (mask1 + mask2).clamp(0, 1)
+
+def time_stretch_mask(
+    x: torch.Tensor, 
+    stretch_factor: int,
+):
+    assert stretch_factor >= 1, "stretch factor must be >= 1"
+    c_seq_len = x.shape[-1]
+    x = x.repeat_interleave(stretch_factor, dim=-1)
+
+    # trim cz to the original length
+    x = x[:, :, :c_seq_len]
+
+    mask = periodic_mask(x, stretch_factor, width=1)
+    return mask
+
+def onset_mask(
+    sig: AudioSignal, 
+    z: torch.Tensor,
+    interface,
+    width: int = 1
+):
+    import librosa
+    import madmom
+    from madmom.features.onsets import RNNOnsetProcessor, OnsetPeakPickingProcessor
+    import tempfile
+    import numpy as np 
+
+    with tempfile.NamedTemporaryFile(suffix='.wav') as f:
+        sig = sig.clone()
+        sig.write(f.name)
+
+        proc = RNNOnsetProcessor(online=False)
+        onsetproc = OnsetPeakPickingProcessor(threshold=0.3,
+                                              fps=sig.sample_rate/interface.codec.hop_length)
+        
+        act = proc(f.name)
+        onset_times = onsetproc(act)
+
+        # convert to indices for z array
+        onset_indices = librosa.time_to_frames(onset_times, sr=sig.sample_rate, hop_length=interface.codec.hop_length)
+
+        if onset_indices.shape[0] == 0:
+            mask = empty_mask(z)   
+            print(f"no onsets found, returning empty mask")
+        else: 
+            torch.set_printoptions(threshold=1000)
+            print("onset indices: ", onset_indices)
+            print("onset times: ", onset_times)
+
+            # create a mask, set onset 
+            mask = torch.ones_like(z)
+            n_timesteps = z.shape[-1]
+
+            for onset_index in onset_indices:
+                onset_index = min(onset_index, n_timesteps - 1)
+                onset_index = max(onset_index, 0)
+                mask[:, :, onset_index - width:onset_index + width] = 0.0
+
+            print(mask)
+    
+    return mask
+
+
+
+if __name__ == "__main__":
+    pass

vampnet/vampnet/modules/__init__.py

@@ -0,0 +1,6 @@
+import audiotools
+
+audiotools.ml.BaseModel.INTERN += ["vampnet.modules.**"]
+audiotools.ml.BaseModel.EXTERN += ["einops", "flash_attn.flash_attention", "loralib"]
+
+from .transformer import VampNet

vampnet/vampnet/modules/activations.py

@@ -0,0 +1,55 @@
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+
+class NewGELU(nn.Module):
+    """
+    Implementation of the GELU activation function currently in Google BERT repo
+    (identical to OpenAI GPT). Also see the Gaussian Error Linear Units
+    paper: https://arxiv.org/abs/1606.08415
+    """
+
+    def forward(self, x):
+        return (
+            0.5
+            * x
+            * (
+                1.0
+                + torch.tanh(
+                    math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))
+                )
+            )
+        )
+
+class GatedGELU(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.gelu = NewGELU()
+
+    def forward(self, x, dim: int = -1):
+        p1, p2 = x.chunk(2, dim=dim)
+        return p1 * self.gelu(p2)
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(channels))
+
+    def forward(self, x):
+        return x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x).pow(2)
+
+def get_activation(name: str = "relu"):
+    if name == "relu":
+        return nn.ReLU
+    elif name == "gelu":
+        return NewGELU
+    elif name == "geglu":
+        return GatedGELU
+    elif name == "snake":
+        return Snake1d
+    else:
+        raise ValueError(f"Unrecognized activation {name}")

vampnet/vampnet/modules/layers.py

@@ -0,0 +1,164 @@
+import time
+from typing import Optional
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+# Scripting this brings model speed up 1.4x
+@torch.jit.script
+def snake(x, alpha):
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    x = x.reshape(shape)
+    return x
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x):
+        return snake(x, self.alpha)
+
+
+def num_params(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+def recurse_children(module, fn):
+    for child in module.children():
+        if isinstance(child, nn.ModuleList):
+            for c in child:
+                yield recurse_children(c, fn)
+        if isinstance(child, nn.ModuleDict):
+            for c in child.values():
+                yield recurse_children(c, fn)
+
+        yield recurse_children(child, fn)
+        yield fn(child)
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+class SequentialWithFiLM(nn.Module):
+    """
+    handy wrapper for nn.Sequential that allows FiLM layers to be
+    inserted in between other layers.
+    """
+
+    def __init__(self, *layers):
+        super().__init__()
+        self.layers = nn.ModuleList(layers)
+
+    @staticmethod
+    def has_film(module):
+        mod_has_film = any(
+            [res for res in recurse_children(module, lambda c: isinstance(c, FiLM))]
+        )
+        return mod_has_film
+
+    def forward(self, x, cond):
+        for layer in self.layers:
+            if self.has_film(layer):
+                x = layer(x, cond)
+            else:
+                x = layer(x)
+        return x
+
+
+class FiLM(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+
+        if input_dim > 0:
+            self.beta = nn.Linear(input_dim, output_dim)
+            self.gamma = nn.Linear(input_dim, output_dim)
+
+    def forward(self, x, r):
+        if self.input_dim == 0:
+            return x
+        else:
+            beta, gamma = self.beta(r), self.gamma(r)
+            beta, gamma = (
+                beta.view(x.size(0), self.output_dim, 1),
+                gamma.view(x.size(0), self.output_dim, 1),
+            )
+            x = x * (gamma + 1) + beta
+        return x
+
+
+class CodebookEmbedding(nn.Module):
+    def __init__(
+        self,
+        vocab_size: int,
+        latent_dim: int,
+        n_codebooks: int,
+        emb_dim: int,
+        special_tokens: Optional[Tuple[str]] = None,
+    ):
+        super().__init__()
+        self.n_codebooks = n_codebooks
+        self.emb_dim = emb_dim
+        self.latent_dim = latent_dim
+        self.vocab_size = vocab_size
+
+        if special_tokens is not None:
+            for tkn in special_tokens:
+                self.special = nn.ParameterDict(
+                    {
+                        tkn: nn.Parameter(torch.randn(n_codebooks, self.latent_dim))
+                        for tkn in special_tokens
+                    }
+                )
+                self.special_idxs = {
+                    tkn: i + vocab_size for i, tkn in enumerate(special_tokens)
+                }
+
+        self.out_proj = nn.Conv1d(n_codebooks * self.latent_dim, self.emb_dim, 1)
+
+    def from_codes(self, codes: torch.Tensor, codec):
+        """ 
+        get a sequence of continuous embeddings from a sequence of discrete codes. 
+        unlike it's counterpart in the original VQ-VAE, this function adds for any special tokens
+        necessary for the language model, like <MASK>. 
+        """
+        n_codebooks = codes.shape[1]
+        latent = []
+        for i in range(n_codebooks):
+            c = codes[:, i, :]
+
+            lookup_table = codec.quantizer.quantizers[i].codebook.weight
+            if hasattr(self, "special"):
+                special_lookup = torch.cat(
+                    [self.special[tkn][i : i + 1] for tkn in self.special], dim=0
+                )
+                lookup_table = torch.cat([lookup_table, special_lookup], dim=0)
+
+            l = F.embedding(c, lookup_table).transpose(1, 2)
+            latent.append(l)
+
+        latent = torch.cat(latent, dim=1)
+        return latent
+
+    def forward(self, latents: torch.Tensor):
+        """
+        project a sequence of latents to a sequence of embeddings
+        """
+        x = self.out_proj(latents)
+        return x
+

vampnet/vampnet/modules/transformer.py

@@ -0,0 +1,953 @@
+import math
+import logging
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+import loralib as lora
+import audiotools as at
+
+from .activations import get_activation
+from .layers import CodebookEmbedding
+from .layers import FiLM
+from .layers import SequentialWithFiLM
+from .layers import WNConv1d
+from ..util import scalar_to_batch_tensor, codebook_flatten, codebook_unflatten
+from ..mask import _gamma
+
+LORA_R = 8
+
+# def log(t, eps=1e-20):
+#     return torch.log(t + eps)
+
+
+def gumbel_noise_like(t):
+    noise = torch.zeros_like(t).uniform_(1e-20, 1)
+    return -torch.log(-torch.log(noise))
+
+
+def gumbel_sample(t, temperature=1.0, dim=-1):
+    return ((t / max(temperature, 1e-10)) + gumbel_noise_like(t)).argmax(dim=dim)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.var_eps = eps
+
+    def forward(self, x):
+        """Returns root mean square normalized version of input `x`
+        # T5 uses a layer_norm which only scales and doesn't shift, which is also known
+        # as Root Mean Square Layer Normalization https://arxiv.org/abs/1910.07467
+        # thus varience is calculated w/o mean and there is no bias
+        Parameters
+        ----------
+        x : Tensor[B x T x D]
+        Returns
+        -------
+        Tensor[B x T x D]
+        """
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.var_eps)
+
+        return self.weight * x
+
+
+class FeedForward(nn.Module):
+    def __init__(
+        self, d_model: int = 512, dropout: float = 0.1, activation: str = "geglu"
+    ):
+        super().__init__()
+        factor = 2 if activation == "geglu" else 1
+        self.w_1 = lora.Linear(d_model, d_model * 4, bias=False, r=LORA_R)
+        self.w_2 = lora.Linear(d_model * 4 // factor, d_model, bias=False, r=LORA_R)
+        self.drop = nn.Dropout(dropout)
+        self.act = get_activation(activation)()
+
+    def forward(self, x):
+        """Computes position-wise feed-forward layer
+        Parameters
+        ----------
+        x : Tensor[B x T x D]
+        Returns
+        -------
+        Tensor[B x T x D]
+        """
+        x = self.w_1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.w_2(x)
+        return x
+
+
+class MultiHeadRelativeAttention(nn.Module):
+    def __init__(
+        self,
+        n_head: int = 8,
+        d_model: int = 512,
+        dropout: float = 0.1,
+        bidirectional: bool = True,
+        has_relative_attention_bias: bool = True,
+        attention_num_buckets: int = 32,
+        attention_max_distance: int = 128,
+    ):
+        super().__init__()
+        d_head = d_model // n_head
+        self.n_head = n_head
+        self.d_head = d_head
+        self.bidirectional = bidirectional
+        self.has_relative_attention_bias = has_relative_attention_bias
+        self.attention_num_buckets = attention_num_buckets
+        self.attention_max_distance = attention_max_distance
+
+        # Create linear query, key, value projections
+        self.w_qs = lora.Linear(d_model, d_model, bias=False, r=LORA_R)
+        self.w_ks = nn.Linear(d_model, d_model, bias=False)
+        self.w_vs = lora.Linear(d_model, d_model, bias=False, r=LORA_R)
+
+        # Create linear final output projection
+        self.fc = lora.Linear(d_model, d_model, bias=False, r=LORA_R)
+
+        # Dropout for attention output weights
+        self.dropout = nn.Dropout(dropout)
+
+        # Create relative positional embeddings (if turned on)
+        if has_relative_attention_bias:
+            self.relative_attention_bias = nn.Embedding(attention_num_buckets, n_head)
+
+    def _relative_position_bucket(self, relative_position):
+        """Converts unbounded relative position into bounded set of buckets
+        with half "exact" buckets (1 position = 1 bucket) and half "log-spaced"
+        buckets
+        Parameters
+        ----------
+        relative_position : Tensor[T_q x T_kv]
+            Relative positions between queries and key_value items
+        Returns
+        -------
+        Tensor[T_q x T_kv]
+            Input relative positions converted into buckets
+        """
+        relative_buckets = 0
+        num_buckets = self.attention_num_buckets
+        max_distance = self.attention_max_distance
+
+        # Convert relative position for (-inf, inf) to [0, inf]
+        # Negative relative positions correspond to past
+        # Positive relative positions correspond to future
+        if self.bidirectional:
+            # use half buckets for each side (past / future)
+            num_buckets //= 2
+
+            # Shift the position positions by `num_buckets` to wrap around
+            # negative positions
+            relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
+            relative_position = torch.abs(relative_position)
+        else:
+            # If not bidirectional, ignore positive positions and wrap
+            # negative positions to positive
+            relative_position = -torch.min(
+                relative_position, torch.zeros_like(relative_position)
+            )
+
+        # Allocate half of the buckets are for exact increments in positions
+        max_exact = num_buckets // 2
+        is_small = relative_position < max_exact
+
+        # The other half of the buckets are for logarithmically bigger bins in
+        # positions up to `max_distance`
+        relative_postion_if_large = max_exact + (
+            torch.log(relative_position.float() / max_exact)
+            / math.log(max_distance / max_exact)
+            * (num_buckets - max_exact)
+        ).to(torch.long)
+
+        # Clip the max relative position to `num_buckets - 1`
+        relative_postion_if_large = torch.min(
+            relative_postion_if_large,
+            torch.full_like(relative_postion_if_large, num_buckets - 1),
+        )
+
+        # Choose relative buckets based on small or large positions
+        relative_buckets += torch.where(
+            is_small, relative_position, relative_postion_if_large
+        )
+
+        return relative_buckets
+
+    def compute_bias(self, query_length, key_length):
+        """Computes a position bias scalar for each index in query_length x key_length
+        Parameters
+        ----------
+        query_length : int
+        key_length : int
+        Returns
+        -------
+        Tensor[heads x 1 x T_q x T_kv]
+            Position bias to be applied on attention logits
+        """
+
+        query_position = torch.arange(query_length, dtype=torch.long)[:, None]
+        key_position = torch.arange(key_length, dtype=torch.long)[None, :]
+        relative_position = key_position - query_position
+
+        # Convert relative position to buckets
+        relative_position_bucket = self._relative_position_bucket(relative_position)
+        relative_position_bucket = relative_position_bucket.to(
+            self.relative_attention_bias.weight.device
+        )
+
+        # Index attention bias values
+        values = self.relative_attention_bias(relative_position_bucket)
+        values = rearrange(values, "q k h -> h 1 q k")
+
+        return values
+
+    def forward(self, q, k, v, mask=None, position_bias=None):
+        """Computes attention over (keys, values) for every timestep in query
+        Parameters
+        ----------
+        q : Tensor[B x T_q x d_model]
+            Query vectors
+        k : Tensor[B x T_kv x d_model]
+            Key vectors to compute attention over
+        v : Tensor[B x T_kv x d_model]
+            Value vectors corresponding to the keys
+        mask : Tensor[B x T_q x T_kv], optional
+        position_bias: Tensor[head x 1 x T_q x T_kv]
+        Returns
+        -------
+        Tensor[B x T_q x d_model]
+            Outputs after attending (key, value) using queries
+        """
+        # Compute query, key, value projections
+        q = rearrange(self.w_qs(q), "b l (head k) -> head b l k", head=self.n_head)
+        k = rearrange(self.w_ks(k), "b t (head k) -> head b t k", head=self.n_head)
+        v = rearrange(self.w_vs(v), "b t (head k) -> head b t k", head=self.n_head)
+
+        # Compute attention matrix
+        attn = torch.einsum("hblk,hbtk->hblt", [q, k]) / np.sqrt(q.shape[-1])
+
+        # Add relative position bias to attention scores
+        if position_bias is None:
+            if self.has_relative_attention_bias:
+                position_bias = self.compute_bias(q.size(-2), k.size(-2))
+            else:
+                position_bias = torch.zeros_like(attn)
+        attn += position_bias
+
+        # Apply mask to attention scores to prevent looking up invalid locations
+        if mask is not None:
+            attn = attn.masked_fill(mask[None] == 0, -1e9)
+
+        # Normalize attention scores and add dropout
+        attn = torch.softmax(attn, dim=3)
+        attn = self.dropout(attn)
+
+        # Compute attended outputs (product of attention matrix and values)
+        output = torch.einsum("hblt,hbtv->hblv", [attn, v])
+        output = rearrange(output, "head b l v -> b l (head v)")
+        output = self.fc(output)
+
+        return output, position_bias
+
+
+class TransformerLayer(nn.Module):
+    def __init__(
+        self,
+        d_model: int = 512,
+        d_cond: int = 64,
+        n_heads: int = 8,
+        bidirectional: bool = True,
+        is_decoder: bool = False,
+        has_relative_attention_bias: bool = False,
+        flash_attn: bool = False,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        # Store args
+        self.is_decoder = is_decoder
+
+        # Create self-attention layer
+        self.norm_1 = RMSNorm(d_model)
+        self.film_1 = FiLM(d_cond, d_model)
+        self.flash_attn = flash_attn
+
+        if flash_attn:
+            from flash_attn.flash_attention import FlashMHA
+            self.self_attn = FlashMHA(
+                embed_dim=d_model,
+                num_heads=n_heads,
+                attention_dropout=dropout,
+                causal=False,
+            )
+        else:
+            self.self_attn = MultiHeadRelativeAttention(
+                n_heads, d_model, dropout, bidirectional, has_relative_attention_bias
+            )
+
+        # (Optional) Create cross-attention layer
+        if is_decoder:
+            self.norm_2 = RMSNorm(d_model)
+            self.film_2 = FiLM(d_cond, d_model)
+            self.cross_attn = MultiHeadRelativeAttention(
+                n_heads,
+                d_model,
+                dropout,
+                bidirectional=True,
+                has_relative_attention_bias=False,
+            )
+
+        # Create last feed-forward layer
+        self.norm_3 = RMSNorm(d_model)
+        self.film_3 = FiLM(d_cond, d_model)
+        self.feed_forward = FeedForward(d_model=d_model, dropout=dropout)
+
+        # Create dropout
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        x,
+        x_mask,
+        cond,
+        src=None,
+        src_mask=None,
+        position_bias=None,
+        encoder_decoder_position_bias=None,
+    ):
+        """Computes one transformer layer consisting of self attention, (op) cross attention
+        and feedforward layer
+        Parameters
+        ----------
+        x : Tensor[B x T_q x D]
+        x_mask : Tensor[B x T_q]
+        src : Tensor[B x T_kv x D], optional
+        src_mask : Tensor[B x T_kv x D], optional
+        position_bias : Tensor[heads x B x T_q x T_q], optional
+            Relative position bias for self attention layer
+        encoder_decoder_position_bias : Tensor[heads x B x T_q x T_kv], optional
+            Relative position bias for cross attention layer
+        Returns
+        -------
+        Tensor[B x T_q x D]
+        """
+        y = self.norm_1(x)
+        y = self.film_1(y.permute(0, 2, 1), cond).permute(0, 2, 1)
+        if self.flash_attn:
+            with torch.autocast(y.device.type, dtype=torch.bfloat16):
+                y = self.self_attn(y)[0]
+        else:
+            y, position_bias = self.self_attn(y, y, y, x_mask, position_bias)
+        x = x + self.dropout(y)
+
+        if self.is_decoder:
+            y = self.norm_2(x)
+            y = self.film_2(y.permute(0, 2, 1), cond).permute(0, 2, 1)
+            y, encoder_decoder_position_bias = self.cross_attn(
+                y, src, src, src_mask, encoder_decoder_position_bias
+            )
+            x = x + self.dropout(y)
+
+        y = self.norm_3(x)
+        y = self.film_3(
+            y.permute(
+                0,
+                2,
+                1,
+            ),
+            cond,
+        ).permute(0, 2, 1)
+        y = self.feed_forward(y)
+        x = x + self.dropout(y)
+
+        return x, position_bias, encoder_decoder_position_bias
+
+
+class TransformerStack(nn.Module):
+    def __init__(
+        self,
+        d_model: int = 512,
+        d_cond: int = 64,
+        n_heads: int = 8,
+        n_layers: int = 8,
+        last_layer: bool = True,
+        bidirectional: bool = True,
+        flash_attn: bool = False,
+        is_decoder: bool = False,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        # Store args
+        self.bidirectional = bidirectional
+        self.is_decoder = is_decoder
+
+        # Create transformer layers
+        # In T5, relative attention bias is shared by all layers in the stack
+        self.layers = nn.ModuleList(
+            [
+                TransformerLayer(
+                    d_model,
+                    d_cond,
+                    n_heads,
+                    bidirectional,
+                    is_decoder,
+                    has_relative_attention_bias=True if (i == 0) else False,
+                    flash_attn=flash_attn,
+                    dropout=dropout,
+                )
+                for i in range(n_layers)
+            ]
+        )
+
+        # Perform last normalization
+        self.norm = RMSNorm(d_model) if last_layer else None
+
+    def subsequent_mask(self, size):
+        return torch.ones(1, size, size).tril().bool()
+
+    def forward(self, x, x_mask, cond=None, src=None, src_mask=None,
+                return_activations: bool = False
+        ):
+        """Computes a full transformer stack
+        Parameters
+        ----------
+        x : Tensor[B x T_q x D]
+        x_mask : Tensor[B x T_q]
+        src : Tensor[B x T_kv x D], optional
+        src_mask : Tensor[B x T_kv], optional
+        Returns
+        -------
+        Tensor[B x T_q x D]
+        """
+
+        # Convert `src_mask` to (B x T_q x T_kv) shape for cross attention masking
+        if self.is_decoder:
+            src_mask = x_mask.unsqueeze(-1) * src_mask.unsqueeze(-2)
+
+        # Convert `x_mask` to (B x T_q x T_q) shape for self attention masking
+        x_mask = x_mask.unsqueeze(-2)
+        if not self.bidirectional:
+            x_mask = x_mask * self.subsequent_mask(x.size(1)).to(x_mask.device)
+
+        # Initialize position biases
+        position_bias = None
+        encoder_decoder_position_bias = None
+
+        # Compute transformer layers
+        if return_activations:
+            activations = []
+        for layer in self.layers:
+            x, position_bias, encoder_decoder_position_bias = layer(
+                x=x,
+                x_mask=x_mask,
+                cond=cond,
+                src=src,
+                src_mask=src_mask,
+                position_bias=position_bias,
+                encoder_decoder_position_bias=encoder_decoder_position_bias,
+            )
+            if return_activations:
+                activations.append(x.detach())
+
+    
+        out = self.norm(x) if self.norm is not None else x
+        if return_activations:
+            return out, torch.stack(activations)
+        else:
+            return out
+
+
+class VampNet(at.ml.BaseModel):
+    def __init__(
+        self,
+        n_heads: int = 20,
+        n_layers: int = 16,
+        r_cond_dim: int = 0,
+        n_codebooks: int = 9,
+        n_conditioning_codebooks: int = 0,
+        latent_dim: int = 8,
+        embedding_dim: int = 1280,
+        vocab_size: int = 1024,
+        flash_attn: bool = True,
+        noise_mode: str = "mask",
+        dropout: float = 0.1
+    ):
+        super().__init__()
+        assert r_cond_dim == 0, f"r_cond_dim must be 0 (not supported), but got {r_cond_dim}"
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.r_cond_dim = r_cond_dim
+        self.n_codebooks = n_codebooks
+        self.n_conditioning_codebooks = n_conditioning_codebooks
+        self.embedding_dim = embedding_dim
+        self.vocab_size = vocab_size
+        self.latent_dim = latent_dim
+        self.flash_attn = flash_attn
+        self.noise_mode = noise_mode
+
+        assert self.noise_mode == "mask", "deprecated"
+
+        self.embedding = CodebookEmbedding(
+            latent_dim=latent_dim,
+            n_codebooks=n_codebooks,
+            vocab_size=vocab_size,
+            emb_dim=embedding_dim,
+            special_tokens=["MASK"],
+        )
+        self.mask_token = self.embedding.special_idxs["MASK"]
+
+        self.transformer = TransformerStack(
+            d_model=embedding_dim,
+            d_cond=r_cond_dim,
+            n_heads=n_heads,
+            n_layers=n_layers,
+            last_layer=True,
+            bidirectional=True,
+            flash_attn=flash_attn,
+            is_decoder=False,
+            dropout=dropout,
+        )
+
+        # Add final conv layer
+        self.n_predict_codebooks = n_codebooks - n_conditioning_codebooks
+        self.classifier = SequentialWithFiLM(
+            WNConv1d(
+                embedding_dim,
+                vocab_size * self.n_predict_codebooks,
+                kernel_size=1,
+                padding="same",
+                # groups=self.n_predict_codebooks,
+            ),
+        )
+
+    def forward(self, x, return_activations: bool = False):
+        x = self.embedding(x)
+        x_mask = torch.ones_like(x, dtype=torch.bool)[:, :1, :].squeeze(1)
+
+        x = rearrange(x, "b d n -> b n d")
+        out = self.transformer(x=x, x_mask=x_mask, return_activations=return_activations)
+        if return_activations:
+            out, activations = out
+
+        out = rearrange(out, "b n d -> b d n")
+
+        out = self.classifier(out, None) # no cond here!
+
+        out = rearrange(out, "b (p c) t -> b p (t c)", c=self.n_predict_codebooks)
+
+        if return_activations:
+            return out, activations
+        else:
+            return out
+    
+    def r_embed(self, r, max_positions=10000):
+        if self.r_cond_dim > 0:
+            dtype = r.dtype
+
+            r = _gamma(r) * max_positions
+            half_dim = self.r_cond_dim // 2
+
+            emb = math.log(max_positions) / (half_dim - 1)
+            emb = torch.arange(half_dim, device=r.device).float().mul(-emb).exp()
+
+            emb = r[:, None] * emb[None, :]
+            emb = torch.cat([emb.sin(), emb.cos()], dim=1)
+
+            if self.r_cond_dim % 2 == 1:  # zero pad
+                emb = nn.functional.pad(emb, (0, 1), mode="constant")
+
+            return emb.to(dtype)
+        else:
+            return r
+    
+    @torch.no_grad()
+    def to_signal(self, z, codec):
+        """
+        convert a sequence of latents to a signal. 
+        """
+        assert z.ndim == 3
+
+        signal = at.AudioSignal(
+            codec.decode(
+                codec.quantizer.from_latents(self.embedding.from_codes(z, codec))[0]
+            )["audio"],
+            codec.sample_rate,
+        )
+
+        # find where the mask token is and replace it with silence in the audio
+        for tstep in range(z.shape[-1]):
+            if torch.any(z[:, :, tstep] == self.mask_token):
+                sample_idx_0 = tstep * codec.hop_length
+                sample_idx_1 = sample_idx_0 + codec.hop_length
+                signal.samples[:, :, sample_idx_0:sample_idx_1] = 0.0
+
+        return signal
+
+
+    @torch.no_grad()
+    def generate(
+        self,
+        codec,
+        time_steps: int = 300,
+        sampling_steps: int = 36,
+        start_tokens: Optional[torch.Tensor] = None,
+        sampling_temperature: float = 1.0,
+        mask: Optional[torch.Tensor] = None,
+        mask_temperature: float = 10.5,
+        typical_filtering=False,
+        typical_mass=0.2,
+        typical_min_tokens=1,
+        top_p=None,
+        return_signal=True,
+        seed: int = None, 
+        sample_cutoff: float = 1.0,
+    ):
+        if seed is not None:
+            at.util.seed(seed)
+        logging.debug(f"beginning generation with {sampling_steps} steps")
+
+
+
+        ##################### 
+        # resolve initial z #
+        #####################
+        z = start_tokens
+
+        if z is None:
+            z = torch.full((1, self.n_codebooks, time_steps), self.mask_token).to(
+                self.device
+            )
+
+        logging.debug(f"created z with shape {z.shape}")
+
+
+        #################
+        # resolve mask #
+        #################
+
+        if mask is None:
+            mask = torch.ones_like(z).to(self.device).int()
+            mask[:, : self.n_conditioning_codebooks, :] = 0.0
+        if mask.ndim == 2:
+            mask = mask[:, None, :].repeat(1, z.shape[1], 1)
+        # init_mask = mask.clone()
+        
+        logging.debug(f"created mask with shape {mask.shape}")
+
+
+        ###########
+        # set up #
+        ##########
+        # apply the mask to z
+        z_masked = z.masked_fill(mask.bool(), self.mask_token)
+        # logging.debug(f"z_masked: {z_masked}")
+
+        # how many mask tokens to begin with?
+        num_mask_tokens_at_start = (z_masked == self.mask_token).sum()
+        logging.debug(f"num mask tokens at start: {num_mask_tokens_at_start}")
+
+        # how many codebooks are we inferring vs conditioning on?
+        n_infer_codebooks = self.n_codebooks - self.n_conditioning_codebooks
+        logging.debug(f"n infer codebooks: {n_infer_codebooks}")
+
+        #################
+        # begin sampling #
+        #################
+
+        for i in range(sampling_steps):
+            logging.debug(f"step {i} of {sampling_steps}")
+
+            # our current schedule step
+            r = scalar_to_batch_tensor(
+                (i + 1) / sampling_steps, 
+                z.shape[0]
+            ).to(z.device)
+            logging.debug(f"r: {r}")
+
+            # get latents
+            latents = self.embedding.from_codes(z_masked, codec)
+            logging.debug(f"computed latents with shape: {latents.shape}")
+
+
+            # infer from latents
+            # NOTE: this collapses the codebook dimension into the sequence dimension
+            logits = self.forward(latents) # b, prob, seq
+            logits = logits.permute(0, 2, 1)  # b, seq, prob
+            b = logits.shape[0]
+
+            logging.debug(f"permuted logits with shape: {logits.shape}")
+
+            sampled_z, selected_probs = sample_from_logits(
+                logits, sample=(
+                   (i / sampling_steps) <= sample_cutoff
+                ), 
+                temperature=sampling_temperature,
+                typical_filtering=typical_filtering, typical_mass=typical_mass,
+                typical_min_tokens=typical_min_tokens,
+                top_k=None, top_p=top_p, return_probs=True,
+            )
+
+            logging.debug(f"sampled z with shape: {sampled_z.shape}")
+
+            # flatten z_masked and mask, so we can deal with the sampling logic
+            # we'll unflatten them at the end of the loop for the next forward pass
+            # remove conditioning codebooks, we'll add them back at the end
+            z_masked = codebook_flatten(z_masked[:, self.n_conditioning_codebooks:, :])           
+
+            mask = (z_masked == self.mask_token).int()
+            
+            # update the mask, remove conditioning codebooks from the mask
+            logging.debug(f"updated mask with shape: {mask.shape}")
+            # add z back into sampled z where the mask was false
+            sampled_z = torch.where(
+                mask.bool(), sampled_z, z_masked
+            )
+            logging.debug(f"added z back into sampled z with shape: {sampled_z.shape}")
+
+            # ignore any tokens that weren't masked
+            selected_probs = torch.where(
+                mask.bool(), selected_probs, torch.inf
+            )
+
+            # get the num tokens to mask, according to the schedule
+            num_to_mask = torch.floor(_gamma(r) * num_mask_tokens_at_start).unsqueeze(1).long()
+            logging.debug(f"num to mask: {num_to_mask}")
+
+            if i != (sampling_steps - 1):
+                num_to_mask = torch.maximum(
+                    torch.tensor(1),
+                    torch.minimum(
+                        mask.sum(dim=-1, keepdim=True) - 1,
+                        num_to_mask
+                    )
+                )
+
+
+            # get our new mask
+            mask = mask_by_random_topk(
+                num_to_mask, selected_probs, mask_temperature * (1-r)
+            )  
+
+            # update the mask
+            z_masked = torch.where(
+                mask.bool(), self.mask_token, sampled_z
+            )
+            logging.debug(f"updated z_masked with shape: {z_masked.shape}")
+
+            z_masked = codebook_unflatten(z_masked, n_infer_codebooks)
+            mask = codebook_unflatten(mask, n_infer_codebooks)
+            logging.debug(f"unflattened z_masked with shape: {z_masked.shape}")
+
+            # add conditioning codebooks back to z_masked
+            z_masked = torch.cat(
+                (z[:, :self.n_conditioning_codebooks, :], z_masked), dim=1
+            )
+            logging.debug(f"added conditioning codebooks back to z_masked with shape: {z_masked.shape}")
+
+
+        # add conditioning codebooks back to sampled_z
+        sampled_z = codebook_unflatten(sampled_z, n_infer_codebooks)
+        sampled_z = torch.cat(
+            (z[:, :self.n_conditioning_codebooks, :], sampled_z), dim=1
+        )
+
+        logging.debug(f"finished sampling")
+
+        if return_signal:
+            return self.to_signal(sampled_z, codec)
+        else:
+            return sampled_z
+
+def sample_from_logits(
+        logits, 
+        sample: bool = True,
+        temperature: float = 1.0,
+        top_k: int = None,
+        top_p: float = None,
+        typical_filtering: bool = False,
+        typical_mass: float = 0.2,
+        typical_min_tokens: int = 1,
+        return_probs: bool = False
+    ):
+    """Convenience function to sample from a categorial distribution with input as
+    unnormalized logits.
+
+    Parameters
+    ----------
+    logits : Tensor[..., vocab_size]
+    config: SamplingConfig
+        The set of hyperparameters to be used for sampling
+        sample : bool, optional
+            Whether to perform multinomial sampling, by default True
+        temperature : float, optional
+            Scaling parameter when multinomial samping, by default 1.0
+        top_k : int, optional
+            Restricts sampling to only `top_k` values acc. to probability,
+            by default None
+        top_p : float, optional
+            Restricts sampling to only those values with cumulative
+            probability = `top_p`, by default None
+
+    Returns
+    -------
+    Tensor[...]
+        Sampled tokens
+    """
+    shp = logits.shape[:-1]
+
+    if typical_filtering:
+        typical_filter(logits, 
+                        typical_mass=typical_mass, 
+                        typical_min_tokens=typical_min_tokens
+        )
+
+    # Apply top_k sampling
+    if top_k is not None:
+        v, _ = logits.topk(top_k)
+        logits[logits < v[..., [-1]]] = -float("inf")
+
+    # Apply top_p (nucleus) sampling
+    if top_p is not None and top_p < 1.0:
+        v, sorted_indices = logits.sort(descending=True)
+        cumulative_probs = v.softmax(dim=-1).cumsum(dim=-1)
+
+        sorted_indices_to_remove = cumulative_probs > top_p
+        # Right shift indices_to_remove to keep 1st token over threshold
+        sorted_indices_to_remove = F.pad(sorted_indices_to_remove, (1, 0), value=False)[
+            ..., :-1
+        ]
+
+        # Compute indices_to_remove in unsorted array
+        indices_to_remove = sorted_indices_to_remove.scatter(
+            -1, sorted_indices, sorted_indices_to_remove
+        )
+
+        logits[indices_to_remove] = -float("inf")
+
+    # Perform multinomial sampling after normalizing logits
+    probs = (
+        F.softmax(logits / temperature, dim=-1)
+        if temperature > 0
+        else logits.softmax(dim=-1)
+    )
+    token = (
+        probs.view(-1, probs.size(-1)).multinomial(1).squeeze(1).view(*shp)
+        if sample
+        else logits.argmax(-1)
+    )
+
+    if return_probs:
+        token_probs = probs.take_along_dim(token.unsqueeze(-1), dim=-1).squeeze(-1)
+        return token, token_probs
+    else:
+        return token
+    
+
+
+def mask_by_random_topk(
+        num_to_mask: int, 
+        probs: torch.Tensor, 
+        temperature: float = 1.0, 
+    ):
+    """
+    Args:
+        num_to_mask (int): number of tokens to mask
+        probs (torch.Tensor): probabilities for each sampled event, shape (batch, seq)
+        temperature (float, optional): temperature. Defaults to 1.0.
+    """
+    logging.debug(f"masking by random topk")
+    logging.debug(f"num to mask: {num_to_mask}")
+    logging.debug(f"probs shape: {probs.shape}")
+    logging.debug(f"temperature: {temperature}")
+    logging.debug("")
+
+    noise = gumbel_noise_like(probs)
+    confidence = torch.log(probs) + temperature * noise
+    logging.debug(f"confidence shape: {confidence.shape}")
+
+    sorted_confidence, sorted_idx = confidence.sort(dim=-1)
+    logging.debug(f"sorted confidence shape: {sorted_confidence.shape}")
+    logging.debug(f"sorted idx shape: {sorted_idx.shape}")
+
+    # get the cut off threshold, given the mask length
+    cut_off = torch.take_along_dim(
+        sorted_confidence, num_to_mask, axis=-1
+    )
+    logging.debug(f"cut off shape: {cut_off.shape}")
+
+    # mask out the tokens
+    mask = confidence < cut_off
+    logging.debug(f"mask shape: {mask.shape}")
+
+    return mask
+
+def typical_filter(
+        logits, 
+        typical_mass: float = 0.95,
+        typical_min_tokens: int = 1,):
+    nb, nt, _ = logits.shape
+    x_flat = rearrange(logits, "b t l -> (b t ) l")
+    x_flat_norm = torch.nn.functional.log_softmax(x_flat, dim=-1)
+    x_flat_norm_p = torch.exp(x_flat_norm)
+    entropy = -(x_flat_norm * x_flat_norm_p).nansum(-1, keepdim=True)
+
+    c_flat_shifted = torch.abs((-x_flat_norm) - entropy)
+    c_flat_sorted, x_flat_indices = torch.sort(c_flat_shifted, descending=False)
+    x_flat_cumsum = (
+        x_flat.gather(-1, x_flat_indices).softmax(dim=-1).cumsum(dim=-1)
+    )
+
+    last_ind = (x_flat_cumsum < typical_mass).sum(dim=-1)
+    sorted_indices_to_remove = c_flat_sorted > c_flat_sorted.gather(
+        1, last_ind.view(-1, 1)
+    )
+    if typical_min_tokens > 1:
+        sorted_indices_to_remove[..., :typical_min_tokens] = 0
+    indices_to_remove = sorted_indices_to_remove.scatter(
+        1, x_flat_indices, sorted_indices_to_remove
+    )
+    x_flat = x_flat.masked_fill(indices_to_remove, -float("Inf"))
+    logits = rearrange(x_flat, "(b t) l -> b t l", t=nt)
+    return logits
+
+
+if __name__ == "__main__":
+    # import argbind
+    from .layers import num_params
+
+    VampNet = argbind.bind(VampNet)
+
+    @argbind.bind(without_prefix=True)
+    def try_model(device: str = "cuda", batch_size: int = 2, seq_len_s: float = 10.0):
+        seq_len = int(32000 / 512 * seq_len_s)
+
+        model = VampNet().to(device)
+
+        z = torch.randint(
+            0, model.vocab_size, size=(batch_size, model.n_codebooks, seq_len)
+        ).to(device)
+
+        r = torch.zeros(batch_size).to(device)
+        
+        z_mask_latent = torch.rand(
+            batch_size, model.latent_dim * model.n_codebooks, seq_len
+        ).to(device)
+        z_hat = model(z_mask_latent)
+
+        pred = z_hat.argmax(dim=1)
+        pred = model.embedding.unflatten(pred, n_codebooks=model.n_predict_codebooks)
+
+        print(f"model has {num_params(model)/1e6:<.3f}M parameters")
+        print(f"prediction has shape {pred.shape}")
+        breakpoint()
+
+    args = argbind.parse_args()
+    with argbind.scope(args):
+        try_model()
+
+

vampnet/vampnet/scheduler.py

@@ -0,0 +1,47 @@
+import copy
+from typing import List
+
+import torch
+
+class NoamScheduler:
+    """OG scheduler from transformer paper: https://arxiv.org/pdf/1706.03762.pdf
+    Implementation from Annotated Transformer: https://nlp.seas.harvard.edu/2018/04/03/attention.html
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        d_model: int = 512,
+        factor: float = 1.0,
+        warmup: int = 4000,
+    ):
+        # Store hparams
+        self.warmup = warmup
+        self.factor = factor
+        self.d_model = d_model
+
+        # Initialize variables `lr` and `steps`
+        self.lr = None
+        self.steps = 0
+
+        # Store the optimizer
+        self.optimizer = optimizer
+
+    def state_dict(self):
+        return {
+            key: value for key, value in self.__dict__.items() if key != "optimizer"
+        }
+
+    def load_state_dict(self, state_dict):
+        self.__dict__.update(state_dict)
+
+    def step(self):
+        self.steps += 1
+        self.lr = self.factor * (
+            self.d_model ** (-0.5)
+            * min(self.steps ** (-0.5), self.steps * self.warmup ** (-1.5))
+        )
+
+        for p in self.optimizer.param_groups:
+            p["lr"] = self.lr
+

vampnet/vampnet/util.py

@@ -0,0 +1,46 @@
+import tqdm
+
+import torch
+from einops import rearrange
+
+def scalar_to_batch_tensor(x, batch_size):
+    return torch.tensor(x).repeat(batch_size)
+
+
+def parallelize(
+        fn, 
+        *iterables,
+        parallel: str = "thread_map",
+        **kwargs
+    ):
+    if parallel == "thread_map":
+        from tqdm.contrib.concurrent import thread_map
+        return thread_map(
+            fn, 
+            *iterables, 
+            **kwargs
+        )
+    elif parallel == "process_map":
+        from tqdm.contrib.concurrent import process_map
+        return process_map(
+            fn, 
+            *iterables, 
+            **kwargs
+        )
+    elif parallel == "single":
+        return [fn(x) for x in tqdm.tqdm(*iterables)]
+    else:
+        raise ValueError(f"parallel must be one of 'thread_map', 'process_map', 'single', but got {parallel}")
+    
+def codebook_flatten(tokens: torch.Tensor):
+    """ 
+    flatten a sequence of tokens from (batch, codebook, time) to (batch, codebook * time)
+    """
+    return rearrange(tokens, "b c t -> b (t c)")
+
+def codebook_unflatten(flat_tokens: torch.Tensor, n_c: int = None):
+    """
+    unflatten a sequence of tokens from (batch, codebook * time) to (batch, codebook, time)
+    """
+    tokens = rearrange(flat_tokens, "b (t c) -> b c t", c=n_c)
+    return tokens