recovering from a gittastrophe

.gitattributes

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

.gitignore

@@ -0,0 +1,194 @@
+# 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/
+samples*/
+models-all/
+models.zip
+.git-old
+
+
+
+gtzan.zip
+.gtzan_emb_cache
+
+
+data/
+data
+pyharp
+
+models/vampnet/*
+models/*

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.

README.md

@@ -0,0 +1,113 @@
+---
+title: salad bowl (vampnet)
+emoji: 🥗
+colorFrom: yellow
+colorTo: green
+sdk: gradio
+sdk_version: 4.37.2
+python_version: 3.9.17
+app_file: app.py
+pinned: false
+license: cc-by-nc-4.0
+---
+
+# VampNet
+
+This repository contains recipes for training generative music models on top of the Descript Audio Codec.
+
+# 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
+```
+
+# 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
+```
+
+for multi-gpu training, use torchrun:
+
+```bash
+torchrun --nproc_per_node gpu scripts/exp/train.py --args.load conf/vampnet.yml --save_path path/to/ckpt
+```
+
+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.
+
+## Debugging training
+
+To debug training, it's easier to debug with 1 gpu and 0 workers
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python -m pdb scripts/exp/train.py --args.load conf/vampnet.yml --save_path /path/to/checkpoints --num_workers 0
+```
+
+## 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/generated/<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/generated/<fine_tune_name>/c2f.yml 
+```
+
+## 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. 
+
+### 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. 
+
+
+
+

app.py

@@ -0,0 +1,428 @@
+import spaces
+from pathlib import Path
+import yaml
+import time
+import uuid
+
+import numpy as np
+import audiotools as at
+import argbind
+import shutil
+import torch
+from datetime import datetime
+
+import gradio as gr
+from vampnet.interface import Interface, signal_concat
+from vampnet import mask as pmask
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+interface = Interface.default()
+
+# populate the model choices with any interface.yml files in the generated confs
+MODEL_CHOICES = {
+    "default": {
+        "Interface.coarse_ckpt": str(interface.coarse_path), 
+        "Interface.coarse2fine_ckpt": str(interface.c2f_path),
+        "Interface.codec_ckpt": str(interface.codec_path),
+    }
+}
+generated_confs = Path("conf/generated")
+for conf_file in generated_confs.glob("*/interface.yml"):
+    with open(conf_file) as f:
+        _conf = yaml.safe_load(f)
+
+        # check if the coarse, c2f, and codec ckpts exist
+        # otherwise, dont' add this model choice
+        if not (
+            Path(_conf["Interface.coarse_ckpt"]).exists() and 
+            Path(_conf["Interface.coarse2fine_ckpt"]).exists() and 
+            Path(_conf["Interface.codec_ckpt"]).exists()
+        ):
+            continue
+
+        MODEL_CHOICES[conf_file.parent.name] = _conf
+
+    
+def to_output(sig):
+    return sig.sample_rate, sig.cpu().detach().numpy()[0][0]
+
+
+
+MAX_DURATION_S = 5
+def load_audio(file):
+    print(file)
+    if isinstance(file, str):
+        filepath = file
+    elif isinstance(file, tuple):
+        # not a file
+        sr, samples = file
+        samples = samples / np.iinfo(samples.dtype).max
+        return sr, samples
+    else:
+        filepath = file.name
+    sig = at.AudioSignal.salient_excerpt(
+        filepath, duration=MAX_DURATION_S
+    )
+    sig = at.AudioSignal(filepath)
+    return to_output(sig)
+
+
+def load_example_audio():
+    return load_audio("./assets/example.wav")
+
+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
+
+
+@spaces.GPU
+def _vamp(
+        seed, input_audio, model_choice, 
+        pitch_shift_amt, periodic_p, 
+        n_mask_codebooks, periodic_w, onset_mask_width, 
+        dropout, sampletemp, typical_filtering, 
+        typical_mass, typical_min_tokens, top_p, 
+        sample_cutoff, stretch_factor, api=False
+    ):
+    t0 = time.time()
+    interface.to("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"using device {interface.device}")
+    _seed = seed if seed > 0 else None
+    if _seed is None:
+        _seed = int(torch.randint(0, 2**32, (1,)).item())
+    at.util.seed(_seed)
+
+    sr, input_audio = input_audio
+    input_audio = input_audio / np.iinfo(input_audio.dtype).max
+    
+    sig = at.AudioSignal(input_audio, sr)
+
+    # reload the model if necessary
+    interface.reload(
+        coarse_ckpt=MODEL_CHOICES[model_choice]["Interface.coarse_ckpt"],
+        c2f_ckpt=MODEL_CHOICES[model_choice]["Interface.coarse2fine_ckpt"],
+    )
+
+    if pitch_shift_amt != 0:
+        sig = shift_pitch(sig, pitch_shift_amt)
+
+    build_mask_kwargs = dict(
+        rand_mask_intensity=1.0,
+        prefix_s=0.0,
+        suffix_s=0.0,
+        periodic_prompt=int(periodic_p),
+        periodic_prompt_width=periodic_w,
+        onset_mask_width=onset_mask_width,
+        _dropout=dropout,
+        upper_codebook_mask=int(n_mask_codebooks), 
+    )
+
+    vamp_kwargs = dict(
+        temperature=sampletemp,
+        typical_filtering=typical_filtering, 
+        typical_mass=typical_mass, 
+        typical_min_tokens=typical_min_tokens, 
+        top_p=None,
+        seed=_seed,
+        sample_cutoff=1.0,
+    )
+
+    # save the mask as a txt file
+    interface.set_chunk_size(10.0)
+    sig, mask, codes = interface.ez_vamp(
+        sig, 
+        batch_size=1 if api else 1,
+        feedback_steps=1,
+        time_stretch_factor=stretch_factor,
+        build_mask_kwargs=build_mask_kwargs,
+        vamp_kwargs=vamp_kwargs,
+        return_mask=True,
+    )
+    print(f"vamp took {time.time() - t0} seconds")
+
+
+    return to_output(sig)
+
+def vamp(data):
+    return _vamp(
+        seed=data[seed], 
+        input_audio=data[input_audio],
+        model_choice=data[model_choice],
+        pitch_shift_amt=data[pitch_shift_amt],
+        periodic_p=data[periodic_p],
+        n_mask_codebooks=data[n_mask_codebooks],
+        periodic_w=data[periodic_w],
+        onset_mask_width=data[onset_mask_width],
+        dropout=data[dropout],
+        sampletemp=data[sampletemp],
+        typical_filtering=data[typical_filtering],
+        typical_mass=data[typical_mass],
+        typical_min_tokens=data[typical_min_tokens],
+        top_p=data[top_p],
+        sample_cutoff=data[sample_cutoff],
+        stretch_factor=data[stretch_factor],
+        api=False, 
+    )
+
+def api_vamp(data):
+    return _vamp(
+        seed=data[seed], 
+        input_audio=data[input_audio],
+        model_choice=data[model_choice],
+        pitch_shift_amt=data[pitch_shift_amt],
+        periodic_p=data[periodic_p],
+        n_mask_codebooks=data[n_mask_codebooks],
+        periodic_w=data[periodic_w],
+        onset_mask_width=data[onset_mask_width],
+        dropout=data[dropout],
+        sampletemp=data[sampletemp],
+        typical_filtering=data[typical_filtering],
+        typical_mass=data[typical_mass],
+        typical_min_tokens=data[typical_min_tokens],
+        top_p=data[top_p],
+        sample_cutoff=data[sample_cutoff],
+        stretch_factor=data[stretch_factor],
+        api=True, 
+    )
+
+
+    
+    
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            manual_audio_upload = gr.File(
+                label=f"upload some audio (will be randomly trimmed to max of 100s)",
+                file_types=["audio"]
+            )
+            load_example_audio_button = gr.Button("or load example audio")
+
+            input_audio = gr.Audio(
+                label="input audio",
+                interactive=False, 
+                type="numpy",
+            )
+
+            audio_mask = gr.Audio(
+                label="audio mask (listen to this to hear the mask hints)",
+                interactive=False, 
+                type="numpy",
+            )
+
+            # 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():
+            with gr.Accordion("manual controls", open=True):
+                periodic_p = gr.Slider(
+                    label="periodic prompt",
+                    minimum=0,
+                    maximum=13, 
+                    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=0, visible=False
+                )
+
+                n_mask_codebooks = gr.Slider(
+                    label="compression prompt ",
+                    value=3,
+                    minimum=1, 
+                    maximum=14,
+                    step=1,
+                )
+            
+            maskimg = gr.Image(
+                label="mask image",
+                interactive=False,
+                type="filepath"
+            )
+
+            with gr.Accordion("extras ", open=False):
+                pitch_shift_amt = gr.Slider(
+                    label="pitch shift amount (semitones)",
+                    minimum=-12,
+                    maximum=12,
+                    step=1,
+                    value=0,
+                )
+
+                stretch_factor = gr.Slider(
+                    label="time stretch factor",
+                    minimum=0,
+                    maximum=8, 
+                    step=1,
+                    value=1, 
+                )
+
+                periodic_w = gr.Slider(
+                    label="periodic prompt width (steps, 1 step ~= 10milliseconds)",
+                    minimum=1,
+                    maximum=20,
+                    step=1,
+                    value=1,
+                )
+
+
+            with gr.Accordion("sampling settings", open=False):
+                sampletemp = gr.Slider(
+                    label="sample temperature",
+                    minimum=0.1,
+                    maximum=10.0,
+                    value=1.0, 
+                    step=0.001
+                )
+            
+                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=True
+                )
+                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=0.9,
+                    value=1.0, 
+                    step=0.01
+                )
+
+
+            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():
+
+            model_choice = gr.Dropdown(
+                label="model choice", 
+                choices=list(MODEL_CHOICES.keys()),
+                value="default", 
+                visible=True
+            )
+
+
+            vamp_button = gr.Button("generate (vamp)!!!")
+
+
+            audio_outs = []
+            use_as_input_btns = []
+            for i in range(1):
+                with gr.Column():
+                    audio_outs.append(gr.Audio(
+                        label=f"output audio {i+1}",
+                        interactive=False,
+                        type="numpy"
+                    ))
+                    use_as_input_btns.append(
+                        gr.Button(f"use as input (feedback)")
+                    )
+
+            thank_you = gr.Markdown("")
+
+            # download all the outputs
+            # download = gr.File(type="filepath", label="download outputs")
+
+
+    _inputs = {
+            input_audio, 
+            sampletemp,
+            top_p,
+            periodic_p, periodic_w,
+            dropout,
+            stretch_factor, 
+            onset_mask_width, 
+            typical_filtering,
+            typical_mass,
+            typical_min_tokens,
+            seed, 
+            model_choice,
+            n_mask_codebooks,
+            pitch_shift_amt, 
+            sample_cutoff, 
+        }
+  
+    # connect widgets
+    vamp_button.click(
+        fn=vamp,
+        inputs=_inputs,
+        outputs=[audio_outs[0]], 
+    )
+
+    api_vamp_button = gr.Button("api vamp", visible=True)
+    api_vamp_button.click(
+        fn=api_vamp,
+        inputs=_inputs, 
+        outputs=[audio_outs[0]], 
+        api_name="vamp"
+    )
+
+    for i, btn in enumerate(use_as_input_btns):
+        btn.click(
+            fn=load_audio,
+            inputs=[audio_outs[i]],
+            outputs=[input_audio]
+        )
+
+try:
+    demo.queue()
+    demo.launch(share=True)
+except KeyboardInterrupt:
+    shutil.rmtree("gradio-outputs", ignore_errors=True)
+    raise

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

conf/interface.yml

@@ -0,0 +1,10 @@
+Interface.coarse_ckpt: ./models/vampnet/coarse.pth
+Interface.coarse2fine_ckpt: ./models/vampnet/c2f.pth
+Interface.codec_ckpt: ./models/vampnet/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
+

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: 7
+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

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
+

hello.py

@@ -0,0 +1,43 @@
+import random
+import vampnet
+import audiotools as at
+
+# load the default vampnet model
+interface = vampnet.interface.Interface.default()
+
+# list available finetuned models
+finetuned_model_choices = interface.available_models()
+print(f"available finetuned models: {finetuned_model_choices}")
+
+# pick a random finetuned model
+model_choice = random.choice(finetuned_model_choices)
+print(f"choosing model: {model_choice}")
+
+# load a finetuned model
+interface.load_finetuned(model_choice)
+
+# load an example audio file
+signal = at.AudioSignal("assets/example.wav")
+
+# get the tokens for the audio
+codes = interface.encode(signal)
+
+# build a mask for the audio
+mask = interface.build_mask(
+    codes, signal,
+    periodic_prompt=7, 
+    upper_codebook_mask=3,
+)
+
+# generate the output tokens
+output_tokens = interface.vamp(
+    codes, mask, return_mask=False,
+    temperature=1.0, 
+    typical_filtering=True, 
+)
+
+# convert them to a signal
+output_signal = interface.decode(output_tokens)
+
+# save the output signal
+output_signal.write("scratch/output.wav")

requirements.txt

@@ -0,0 +1,10 @@
+torch==2.1.0
+argbind>=0.3.2
+numpy==1.23
+loralib
+wavebeat @ git+https://github.com/hugofloresgarcia/wavebeat
+lac @ git+https://github.com/hugofloresgarcia/lac.git
+descript-audiotools @ git+https://github.com/hugofloresgarcia/audiotools.git
+-e git+https://github.com/audacitorch/pyharp.git#egg=pyharp
+torch_pitch_shift
+gradio==4.37.2

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()

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.decode(
+        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.decode(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.decode(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.decode(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.decode(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.decode(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.decode(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.decode(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()

scripts/exp/export.py

@@ -0,0 +1,22 @@
+from pathlib import Path
+
+run_dir = Path("runs/sample-instrument")
+name = run_dir.name
+
+repo_dir = Path("models/vampnet")
+
+
+for part in ("coarse", "c2f"):
+    outdir = repo_dir / "loras" / name 
+    outdir.mkdir(parents=True, exist_ok=True)
+    outpath = outdir / f"{part}.pth"
+    path = run_dir / part / "latest" / "vampnet" / "weights.pth"
+    path.rename(outpath)
+    print(f"moved {path} to {outpath}")
+
+# now, push to hub
+from huggingface_hub import Repository
+repo = Repository(repo_dir, git_user="hugofloresgarcia", git_email="huferflo@gmail.com")
+repo.push_to_hub(
+    commit_message=f"add {name}"
+)

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()
+
+
+
+    

scripts/exp/train.py

@@ -0,0 +1,686 @@
+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 as tfm
+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",
+]
+
+# 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 = tfm.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.decode(z_mask, state.codec)
+    imputed_true = vn.decode(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.decode(z_pred, state.codec)
+    reconstructed = vn.decode(z, state.codec)
+    masked = vn.decode(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
+    )
+
+
+    if accel.world_size > 1:
+        from torch.distributed.optim import ZeroRedundancyOptimizer
+        optimizer = ZeroRedundancyOptimizer(model.parameters(), AdamW)
+        print(f"OPTIMIZER LR is {optimizer.param_groups[0]['lr']}")
+    else:
+        optimizer = AdamW(model.parameters())
+
+    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)

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>`

scripts/utils/gtzan_embeddings.py

@@ -0,0 +1,264 @@
+"""
+TODO: train a linear probe
+usage:
+   python gtzan_embeddings.py --args.load conf/interface.yml --Interface.device cuda --path_to_gtzan /path/to/gtzan/genres_original  --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)
+
+DEBUG = False
+
+def smart_plotly_export(fig, save_path):
+    img_format = save_path.split('.')[-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(emb, labels, save_path, n_components=3, method='tsne', title=''):
+    """
+    dimensionality reduction for visualization!
+    saves an html plotly figure to save_path
+    parameters:
+        emb (np.ndarray): the samples to be reduces 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
+    if method == 'umap':
+        from umap import UMAP
+        reducer = umap.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
+ 
+    proj = reducer.fit_transform(emb)
+
+    if n_components == 2:
+        df = pd.DataFrame(dict(
+            x=proj[:, 0],
+            y=proj[:, 1],
+            instrument=labels
+        ))
+        fig = px.scatter(df, x='x', y='y', color='instrument',
+                        title=title+f"_{method}")
+
+    elif n_components == 3:
+        df = pd.DataFrame(dict(
+            x=proj[:, 0],
+            y=proj[:, 1],
+            z=proj[:, 2],
+            instrument=labels
+        ))
+        fig = px.scatter_3d(df, x='x', y='y', z='z',
+                        color='instrument',
+                        title=title)
+    else:
+        raise ValueError("cant plot more than 3 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 Embedding:
+    genre: 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_gtzan: str = None, 
+    cache_dir: str = "./.gtzan_emb_cache",
+    output_dir: str = "./gtzan_vampnet_embeddings",
+    layers: List[int] = [1, 3, 5, 7, 9, 11, 13, 15, 17, 19]
+):
+    path_to_gtzan = Path(path_to_gtzan)
+    assert path_to_gtzan.exists(), f"{path_to_gtzan} 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()
+
+    # gtzan should have a folder for each genre, so let's get the list of genres
+    genres = [Path(x).name for x in path_to_gtzan.iterdir() if x.is_dir()]
+    print(f"Found {len(genres)} genres")
+    print(f"genres: {genres}")
+
+    # collect audio files, genres, and embeddings
+    data = []
+    for genre in genres:
+        audio_files = list(at.util.find_audio(path_to_gtzan / genre))
+        print(f"Found {len(audio_files)} audio files for genre {genre}")
+
+        for audio_file in tqdm.tqdm(audio_files, desc=f"embedding genre {genre}"):
+            # check if we have a cached embedding for this file
+            cached_path = (cache_dir / f"{genre}_{audio_file.stem}.emb")
+            if cached_path.exists():
+                # if so, load it
+                if DEBUG:
+                    print(f"loading cached embedding for {cached_path.stem}")
+                embedding = Embedding.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 = Embedding(
+                    genre=genre,
+                    filename=audio_file.name,
+                    embedding=emb
+                )
+
+                # cache the embeddings
+                cached_path.parent.mkdir(exist_ok=True, parents=True)
+                embedding.save(cached_path)
+            data.append(embedding)
+
+    # now, let's do a dim reduction on the embeddings
+    # and visualize them. 
+
+    # collect a list of embeddings and labels
+    embeddings = [d.embedding for d in data]
+    labels = [d.genre for d in data]
+
+    # convert the embeddings to a numpy array
+    embeddings = np.stack(embeddings)
+
+    # do dimensionality reduction for each layer we're given
+    for layer in tqdm.tqdm(layers, desc="dim reduction"):
+        dim_reduce(
+            embeddings[:, layer, :], labels, 
+            save_path=str(output_dir / f'vampnet-gtzan-layer={layer}.html'), 
+            n_components=2, method='tsne', 
+            title=f'vampnet-gtzan-layer={layer}'
+        )
+        
+
+
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+    with argbind.scope(args):
+        main()

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)

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()

scripts/utils/split.py

@@ -0,0 +1,66 @@
+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,
+):
+    print(f"finding audio")
+
+    audio_folder = Path(audio_folder)
+    audio_files = util.find_audio(audio_folder)
+    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)
+            try:
+                os.symlink(file, out_file)
+            except FileExistsError:
+                print(f"File {out_file} already exists, skipping")
+
+        # 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()

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()

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()

scripts/utils/visualize_embeddings.py

@@ -0,0 +1,265 @@
+"""
+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)
+
+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:
+        annotated_embeddings (list): the annotated enmbeddings to be reduced; embeddings have 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.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()

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

setup.py

@@ -0,0 +1,40 @@
+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(),
+    install_requires=[
+        "torch",
+        "argbind>=0.3.2",
+        "numpy==1.23",
+        "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",
+    ],
+)

token_telephone/tt.py

@@ -0,0 +1,616 @@
+from ttutil import hsv_to_rgb, dbg, log, set_debug, pow2db, db2pow
+from dataclasses import dataclass, field
+import os
+from pathlib import Path
+import random
+import time
+from threading import Thread
+import gc
+gc.disable()
+
+import sounddevice as sd
+
+from blessed import Terminal
+
+import numpy as np
+import torch
+from einops import rearrange
+
+PROFILE = False
+DEBUG = False
+DEBUG_NO_VAMPNET = False
+set_debug(DEBUG)
+# if DEBUG:
+#     import gc
+#     # log when gc start and stops
+#     gc.set_debug(gc.DEBUG_STATS)
+
+@dataclass 
+class LoadState:
+    t0: float = None
+    loaded: bool = False
+
+load_state = LoadState()    
+
+def on_random_color():
+    def random_rgb_bg():
+        return np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255)
+    return term.on_color_rgb(*random_rgb_bg())
+
+# draw the intro screen before slow imports
+def color_tokenize_txt(text: str):
+    # apply a random bg color to each letter
+    return "".join(on_random_color()(letter) for letter in text)
+
+def color_tokenize_words(text: str):
+    return " ".join(on_random_color()(word) for word in text.split(" "))
+
+def draw_intro_screen():
+    global load_state
+    load_state.t0 = time.time()
+    avg_time = 20 # average loading time
+
+    while not load_state.loaded: 
+        print(term.clear)
+        print(term.move_xy(0, 1) + term.center(color_tokenize_words("hugo flores garcía")))
+        print(term.move_xy(0, 3) + term.center(color_tokenize_words("and")))
+        print(term.move_xy(0, 5) + term.center(color_tokenize_words("stephan moore")))
+        print(term.move_xy(0, 7) + term.center(color_tokenize_words("present")))
+        print(term.move_xy(0, 9) + term.center(term.bold(color_tokenize_txt("token telephone"))))
+
+        # print(term.move_xy(0, 10) + term.center(color_tokenize_txt("loading ")), end="")
+        # make a little loading bar
+        elapsed = time.time() - load_state.t0
+        num_dots = int((elapsed / avg_time) * 20)
+        num_spaces = 20 - num_dots
+        print(term.move_xy(0, 12) + term.center(color_tokenize_words("loading")))
+        print(term.move_xy(0, 13) + term.center(color_tokenize_txt(f"[{'.' * num_dots}") + f"{' ' * num_spaces}]"))
+        time.sleep(0.3)
+    
+    log(f"loading took {time.time() - load_state.t0} seconds")
+    return
+    
+# the program
+term = Terminal()
+
+# draw the intro screen on a background thread
+Thread(target=draw_intro_screen).start()
+
+# disable garbage collection
+from audiotools import AudioSignal
+from vamp_helper import load_interface, ez_variation
+
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 
+# ~~~~~~  configs!     ~~~~~~~~
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+MAX_LOUDNESS = -20
+MIN_LOUDNESS = -40
+COLS = 40
+ROWS = 13
+
+device = 'Scarlett 4i4 4th Gen'
+sample_rate = 48000
+num_channels = 4
+blocksize = 16384
+
+
+# TODO: 
+# still some quirks to work around recording time: 
+# do we wanna stop recording and wait a full cycle before letting people record again? 
+# how do we wanna balance the volume of a new input vs what's currently gonig on?
+# should people have to take turns in between new loops? 
+# otherwise, we're doing great i think
+# we also need to add a crossfade. This means maybe cutting off the last 0.1 seconds of the loop, and the beginning 0.1
+# and use that to crossfade.
+
+# TODO: do I wanna train a diff model to swap every 2hrs or something? 
+# how lond does model swapping take?  how can I make it faster?
+
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+# ~~~~~~  looper      ~~~~~~~~
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+@dataclass 
+class State:
+    # looper state
+    feedback: float = 0.25
+    duration: float = 5.0
+    record_channel: int = 0
+
+    loopbuf: np.ndarray = None # the main loop buffer. the token telephone audio is here
+    looper_in: np.ndarray = None # a buffer that stores the audio that's being recorded
+
+    buf_in: np.ndarray = None # the input block with audio samples in the audio callbac
+    lookback_buf: np.ndarray = None # stores some lookback audio for when the threshold is passed, to propery capture transients
+
+    recording: bool = False
+    playing: bool = False
+
+    # ramps
+    record_ramp_in: bool = False
+    record_ramp_out: bool = False
+
+    # n_record_layers: int = 2 # number of times we'll record over before clearing
+    # cur_rec_layer: int = 0
+    recording_locked: bool = False
+
+    rec_time: float = 0
+    cur_hold_time: float = None
+    pos: int = 0
+    rms_db: float = float("-inf")
+
+    trig_threshold_db = -25 # a more sane default is -20
+    hold_seconds = 1.0
+    rel_threshold_db = -40 # a more sane default is -30
+
+    status: str = field(default=None)
+
+    # token telephone configs
+    z_buf: torch.Tensor = None
+    input_ready = False
+    input_channel = 0
+    token_telephone_processing: bool = False
+    num_telephone_chans = 4
+    tt_cur_ch = 0
+
+    def __post_init__(self):
+        self.loopbuf = np.zeros((num_channels, int(self.duration * sample_rate)))
+        self.looper_in = np.zeros((1, int(self.duration * sample_rate)))
+
+        # hold 200ms of lookback to account for rising attacks. 
+        num_lookback_samples = max(int(sample_rate * 0.2), int(blocksize))
+        log(f"num_lookback_samples {num_lookback_samples} ({num_lookback_samples / sample_rate} seconds)")
+        self.lookback_buf = np.zeros((1, num_lookback_samples)) 
+        
+        self.buf_in = np.zeros((num_channels, blocksize))
+
+
+
+def check_if_record(st: State, ain: np.ndarray, on_release_callback=None):
+    # get our rms value
+    rms = pow2db(np.sqrt(np.mean(ain**2)))
+    st.rms_db = rms
+
+    # determine if we should ater the looper state
+    # if we werent recording and we cross the trigger threshold
+    # start recording
+    # if not st.recording and rms > st.trig_threshold_db and not st.recording_locked:
+    if not st.recording and rms > st.trig_threshold_db and not st.recording_locked:
+        st.recording = True
+        st.record_ramp_in = True
+
+    # if we were recording and we cross the release threshold
+    # begin the hold period
+    if (st.recording and rms < st.rel_threshold_db) or st.rec_time > (st.duration-st.hold_seconds):
+        # if we dont have a hold time, set it
+        if st.cur_hold_time is None:
+            st.cur_hold_time = time.time()
+
+        # release if we have a hold time and we've held for the required time, 
+        if (time.time() - st.cur_hold_time) > st.hold_seconds:
+            st.record_ramp_out = True
+            st.rec_time = 0
+            if on_release_callback is not None:
+                st.input_ready = True
+                on_release_callback(st)
+            st.cur_hold_time = None
+        else:
+            pass
+    else:
+        st.cur_hold_time = None
+
+
+def launch_token_telephone(st: State):
+    if interface is None:
+        log("no interface loaded, can't do token telephone!")
+        time.sleep(10)
+        return
+
+    # if we're already processing, do nothing
+    if st.token_telephone_processing:
+        return 
+    else:
+        log("starting token telephone!")
+        Thread(target=do_token_telephone, args=(st,)).start()
+
+
+def do_token_telephone(st: State,):
+    st.token_telephone_processing = True
+    while True:
+        lrc = st.record_channel
+        t0 = time.time()
+        cur_ch = st.tt_cur_ch
+
+        # if there was input ready, start back from the top. 
+        if st.input_ready:
+            log(f"there was input ready, processing!")
+            # NOTE: hugo, trying something new here. what happens if
+            # we don't reset the channel when input is ready, 
+            # and instead let it come in anywhere in the cycle?
+            # st.tt_cur_ch = 0 # uncomment to go back to reality
+
+            # clear the lrc, reset for next record. 
+            st.input_ready = False
+
+            # reocrd the channel that we'll be processing in and lock recording
+            st.input_channel = cur_ch
+            st.recording_locked = True
+
+            # first, let's preprocess looper in
+            sig_looper_in = AudioSignal(
+                torch.from_numpy(st.looper_in).unsqueeze(0),
+                sample_rate=sample_rate
+            )
+            sig_loopbuf_curch = AudioSignal(
+                torch.from_numpy(st.loopbuf[cur_ch:cur_ch+1]).unsqueeze(0),
+                sample_rate=sample_rate
+            )
+            # make sure looperin matches the midpoint in loudness
+            ldns_mid = max(sig_loopbuf_curch.loudness(), sig_looper_in.loudness())
+            sig_looper_in = sig_looper_in.normalize(ldns_mid)
+            st.looper_in = sig_looper_in.samples.cpu().numpy().squeeze(0)
+
+            st.loopbuf[cur_ch:cur_ch + 1] = (
+                st.looper_in + st.loopbuf[cur_ch:cur_ch+1] * st.feedback
+            )
+            # also lower the volumes of the other channels
+            for i in range(4):
+                if i != cur_ch:
+                    st.loopbuf[i:i+1] = st.loopbuf[i:i+1] * 0.5 # -3dB
+
+            st.looper_in = np.zeros_like(st.looper_in)
+
+        loop_input = st.loopbuf[cur_ch:cur_ch+1]
+
+        # ~~~ VAMPNET STUFF ~~~~
+        sig = AudioSignal(
+            torch.from_numpy(loop_input).unsqueeze(0),
+            sample_rate=sample_rate
+        )
+        input_loudness = sig.loudness()
+        log(f"INPUT loudness {input_loudness}")
+        if input_loudness > MAX_LOUDNESS:
+            log(f"input loudness {input_loudness} is over {MAX_LOUDNESS}!")
+            sig = sig.normalize(MAX_LOUDNESS)
+        elif input_loudness < MIN_LOUDNESS:
+            log(f"input loudness {input_loudness} is under {MIN_LOUDNESS}!")
+            sig = sig.normalize(MIN_LOUDNESS)
+
+        sig = ez_variation(interface, sig)
+        sig = sig.resample(sample_rate)
+
+        # notify if we've gone over the loudness
+        sig = sig.normalize(input_loudness)
+        outloudness = sig.loudness()
+        if outloudness > MAX_LOUDNESS:
+            log(f"out loudness {sig.loudness()} is over {MAX_LOUDNESS}!")
+            sig = sig.normalize(MAX_LOUDNESS)
+        elif outloudness < MIN_LOUDNESS:
+            log(f"out loudness {sig.loudness()} is under {MIN_LOUDNESS}!")
+            sig = sig.normalize(MIN_LOUDNESS)
+
+        # put it back in the loopbuf
+        # write to the next channel
+        # (TODO: instead of trimming to loopbuf.shape[1], maybe we can just have the loopbuf be the right size from init time.)
+        cur_ch = (cur_ch + 1) % st.num_telephone_chans
+        st.tt_cur_ch = cur_ch
+        if False: # HUGO: is there a time where we want feedback? 
+            st.loopbuf[cur_ch:cur_ch+1] = (
+                sig.samples.cpu().numpy().squeeze(0)[:, :st.loopbuf.shape[1]]
+                + st.feedback * st.loopbuf[cur_ch:cur_ch+1]
+            )
+        else:
+            st.loopbuf[cur_ch:cur_ch+1] = (
+                sig.samples.cpu().numpy().squeeze(0)[:, :st.loopbuf.shape[1]]
+            )
+
+        log(f"output loudness {sig.loudness()}")
+        log(f"telephone loop took {time.time() - t0} seconds... next channel {cur_ch}\n\n")
+
+        # if we've made it back to the input channel, we can unlock the recording
+        log(f"cur_ch {cur_ch} input_channel {st.input_channel}")
+        if cur_ch == st.input_channel:
+            st.recording_locked = False
+            log(f"recording unlocked!")
+
+
+        # unlock the recording if we've successfully written to all channels
+        # if st.recording_locked and cur_ch == 0:
+        #     st.recording_locked = False
+        #     log(f"recording locked {st.recording_locked}")
+
+    st.token_telephone_processing = False
+    return
+
+# TODO: since we're using this really high threshold
+# we always need to record about 100ms in advance, to catch the beginning of the attacks. 
+
+def looper_process_block(st, block: np.ndarray):
+    lrc = st.record_channel
+
+    # treat the lookback buffer as a circular buffer
+    st.lookback_buf = np.roll(st.lookback_buf, block.shape[1], axis=1)
+    st.lookback_buf[:, -block.shape[1]:] = block[lrc:lrc+1, :]
+
+
+    # check if we need to record.
+    if st.recording:
+        start_i = (st.pos + block.shape[1]) - st.lookback_buf.shape[1]
+        end_i = st.pos + st.lookback_buf.shape[1]
+
+        indices = np.take(
+            np.arange(st.loopbuf.shape[1]),
+            np.arange(start_i, end_i),
+            mode="wrap"
+        )
+        _audio_in = st.lookback_buf[:, :]
+        # ramp in if we need to
+        if st.record_ramp_in:
+            _audio_in = _audio_in * np.linspace(0, 1, _audio_in.shape[1])
+            st.record_ramp_in=False
+        
+        if st.record_ramp_out:
+            _audio_in = _audio_in * np.linspace(1, 0, _audio_in.shape[1])
+            st.record_ramp_out=False
+            st.recording = False
+
+        st.looper_in[:, indices] = (
+            0.9 * st.looper_in[:, indices] + _audio_in
+        )
+
+        # incremement the recording time
+        st.rec_time += st.lookback_buf.shape[1] / sample_rate
+    
+    # check if we need to play
+    crossfade_samples = int(0.1 * sample_rate)
+    if st.playing:
+        play_pos = (st.pos + block.shape[1]) % st.loopbuf.shape[1] # read one buffer ahead
+        indices = np.arange(play_pos, play_pos + block.shape[1]) 
+        block = st.loopbuf.take(indices, axis=1, mode="wrap")[:, :] # this doesn't have any crossfading. # TODO: this is still not working! 
+
+    # if we've recorded more than the loop size
+    if st.rec_time > st.duration and st.recording:
+        # play the loop
+        play_pos = st.pos + block.shape[1] # read one buffer ahead
+        indices = np.arange(play_pos, play_pos + block.shape[1])
+
+        block[lrc:lrc] = st.looper_in.take(indices, axis=1, mode="wrap")[:, :]
+    
+    # advance looper state
+    st.pos = (st.pos + block.shape[1]) % st.loopbuf.shape[1]    
+    
+    return block
+
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+# ~~~~~~  drawing      ~~~~~~~~
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+def draw_rms_bar(st, x, y, width, height):
+    rms_min = -50
+    rms_max = -10
+    rms = st.rms_db
+    rms = max(rms, rms_min)
+    threshold = st.trig_threshold_db
+    rel_threshold = st.rel_threshold_db
+
+    rms_block = int((rms - rms_min) / (rms_max - rms_min) * height)
+    threshold_block = (threshold - rms_min) / (rms_max - rms_min) * height
+    rel_threshold_block = (rel_threshold - rms_min) / (rms_max - rms_min) * height
+
+    # draw the rms curve
+    for i in range(rms_block, height+4):
+        with term.location(x+4, y+height-i):
+            print(term.clear_bol)
+    for i in range(rms_block):
+        rms_val = i * (rms_max - rms_min) / height + rms_min
+        with term.location(x, y+height-2-i):
+            if i < threshold_block:
+                print("   " + term.on_green(f"*"))
+            else:
+                print("   " + term.on_red(f"*"))
+
+    # at the very bottom of the bar, draw the rms value
+    with term.location(x, y+height-1):
+        print(f"{rms:.1f}dB")
+        # print(f"  rms")
+
+
+def draw_looper(st):
+    x = 0
+    y = 0
+    width = COLS
+    height = ROWS
+
+    tt_refresh_every = 0.3
+    if not hasattr(draw_looper, "last_draw"):
+        draw_looper.last_draw = 0
+        should_draw = True
+    else:
+        should_draw = (time.time() - draw_looper.last_draw) > tt_refresh_every
+        if should_draw:
+            draw_looper.last_draw = time.time()
+
+
+    draw_rms_bar(st, x, y, width - 10, height)
+        
+    if should_draw:
+        with term.location(width // 2-4, 1):
+            for i, letter in enumerate("token telephone"):
+                print(on_random_color()(letter), end="")
+
+    # with term.location(ROWS-2, COLS // 2):
+        # print(f"status {st.status}!!!")
+
+
+    # if we're recording, draw a red unlderlined "rec" sign on the bottom right
+    # with term.location(width-8, height-1):
+    #     if st.recording:
+    #         print(term.on_red("rec"))
+    #     else:
+    #         print(term.on_gray50("rec"))
+    
+    # # if we're playing draw a green underline "play" sign on the bottom right
+    # with term.location(width-4, height-1):
+    #     if st.playing:
+    #         print(term.on_green("play"))
+    #     else:
+    #         print(term.on_gray50("play"))
+
+
+    # draw the timeline at the bottom using ---
+    with term.location(6, height):
+        timeline = ["-"] * (width - 12)
+        playhead = int((st.pos / st.loopbuf.shape[1]) * (width - 12))
+        timeline[playhead] = "v"
+        print("|"+"".join(timeline) + "|")
+
+    
+    # draw the main message at the very center:
+    msg_loc = (width // 2, height // 2+1)
+    _x, _y = msg_loc
+    if not st.recording:
+        if not st.recording_locked:
+            print(term.move_xy(0, _y-1) + term.center("make a sound", width=width+5))
+            print(term.move_xy(0, _y+0) + term.center("to", width=width+5))
+            print(term.move_xy(0, _y+1) + term.center("record", width=width+5))
+        else:
+            # how many seconds left until we can record again?
+            # how many more chs do we need to go through before we can record again?
+            if st.tt_cur_ch < st.input_channel:
+                chs_remaining = st.input_channel - st.tt_cur_ch
+            else:
+                chs_remaining = 4-st.tt_cur_ch + st.input_channel
+            locked_time_remaining = chs_remaining * st.duration + st.duration - (st.pos / sample_rate)
+            print(term.move_xy(0, _y-1) + term.center("please wait", width=width+5))
+            print(term.move_xy(0, _y+0) + term.center(term.on_green(f"{locked_time_remaining:.1f}s"), width=width+5))
+            print(term.move_xy(0, _y+1) + term.center("for your turn :)", width=width+5))
+    else: 
+        print(term.move_xy(0, _y-1) + term.center(term.on_red("recording"), width=width+5))
+        print(term.move_xy(0, _y+0) + term.center(f"{(st.duration) - st.rec_time:.1f}s left", width=width+5))
+        print(term.move_xy(0, _y+1) + term.center("", width=width+5))
+
+
+    # we'll draw channel 0 (1) on the bottom right corner
+    # channel 1 (2) on the top right corner
+    # channel 2 (3) on the top left corner
+    # channel 3 (4) on the bottom left corner
+    my = 3 # margin
+    mx = 10
+    locations = {
+        1: (width - mx, height - my),
+        2: (width - mx, 1+my),
+        3: (mx, 1+my),
+        4: (mx, height - my),
+    }
+    for i in range(1, 5):
+        if should_draw:
+            if st.tt_cur_ch == i - 1 and st.token_telephone_processing:
+                x, y = locations[i]
+                on_random_colors = lambda n: "".join(on_random_color()(" ") for _ in range(n))
+                print(term.move_xy(x, y-1) + on_random_colors(5))
+                print(term.move_xy(x, y) + on_random_color()(" ") + f" {i} " + on_random_color()(" "))
+                print(term.move_xy(x, y+1) + on_random_colors(5))
+            else:
+                # same thing, but a gray instead of random colors
+                x, y = locations[i]
+                on_gray_colors = lambda n: "".join(term.on_gray50(" ") for _ in range(n))
+                print(term.move_xy(x, y-1) + on_gray_colors(5))
+                print(term.move_xy(x, y) + term.on_gray50(" ") + f" {i} " + term.on_gray50(" "))
+                print(term.move_xy(x, y+1) + on_gray_colors(5))
+
+
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+# ~~~~~~  live audio   ~~~~~~~~
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+def audio_init():
+    sd.default.samplerate = sample_rate
+    sd.default.device = device
+            
+# ~~~~~~ the main audio callback ~~~~~~~~~
+def callback(st, indata, outdata, frames, _time, status):
+    t0 = time.time()
+    lrc = st.record_channel
+
+    if status:
+        log(f"status is {status}")
+        st.status = status
+        
+    # log dtype, status, frames, time, max min
+    # log(f"indata {indata.dtype} max {indata.max()} min {indata.min()} {status} {frames} {_time}")
+
+
+    ain = rearrange(indata, 't n -> n t', n=num_channels)
+    
+    # convert audio to from int32 to float32
+    ain = ain.astype(np.float32) / np.iinfo(np.int16).max
+    buf_in = ain
+
+    # if it's all zeros, we're not recording
+    # so we can just pass it through
+    if np.all(buf_in == 0):
+        st.status = st.status + "no input"
+        return 
+
+    st.buf_in = buf_in
+    check_if_record(
+            st, buf_in, 
+            on_release_callback=launch_token_telephone
+    )
+    buf_in = looper_process_block(st, buf_in)
+
+    # pass our st.loopbuf to the output
+    ain = buf_in
+
+    # convert back to int32
+    ain =  (ain * np.iinfo(np.int16).max).astype(np.int16)
+
+    outdata[:] = rearrange(ain, 'n t -> t n')
+
+    # log(f"outdata {outdata.dtype} max {outdata.max()} min {outdata.min()} --- took {time.time() - t0} seconds")
+
+
+
+if DEBUG_NO_VAMPNET:
+    interface=None
+else:
+    interface = load_interface(model_choice="opera")
+
+load_state.loaded = True
+
+def main():
+    if PROFILE:
+        import yappi
+        yappi.start()
+
+    try:
+        audio_init()
+        st = State()
+        st.playing = True
+
+        from functools import partial
+        cb = partial(callback, st)
+
+        with term.fullscreen(), term.cbreak():
+            with sd.Stream(channels=num_channels, callback=cb, blocksize=blocksize, prime_output_buffers_using_stream_callback=True, dtype=np.int16):
+                while True:
+                    with term.hidden_cursor():
+                        if DEBUG: 
+                            time.sleep(100)
+                        else:
+                            draw_looper(st)
+
+    except KeyboardInterrupt:
+        print(term.clear)
+        if PROFILE: 
+            yappi.stop()
+
+            # retrieve thread stats by their thread id (given by yappi)
+            threads = yappi.get_thread_stats()
+            for thread in threads:
+                print(
+                    "Function stats for (%s) (%d)" % (thread.name, thread.id)
+                )  # it is the Thread.__class__.__name__
+                yappi.get_func_stats(ctx_id=thread.id).print_all()
+
+main()

token_telephone/ttutil.py

@@ -0,0 +1,65 @@
+import logging
+from pathlib import Path
+ROOT = Path(__file__).parent
+
+import numpy as np
+from queue import Queue
+
+# make a log file!!
+logfile= ROOT / "log.txt"
+if logfile.exists():
+    logfile.unlink()
+logging.basicConfig(filename=logfile, level=logging.INFO, datefmt="%Y-%m-%d %H:%M:%S", format="%(asctime)s | %(levelname)s | %(message)s")
+
+
+def hsv_to_rgb(h, s, v):
+    # from https://en.wikipedia.org/wiki/HSL_and_HSV#From_HSV
+    c = v * s
+    h_ = h / 60
+    x = c * (1 - abs(h_ % 2 - 1))
+    m = v - c
+
+    if h_ < 1:
+        r, g, b = c, x, 0
+    elif h_ < 2:
+        r, g, b = x, c, 0
+    elif h_ < 3:
+        r, g, b = 0, c, x
+    elif h_ < 4:
+        r, g, b = 0, x, c
+    elif h_ < 5:
+        r, g, b = x, 0, c
+    else:
+        r, g, b = c, 0, x
+
+    return r + m, g + m, b + m
+
+
+def dbg(*args):
+    print(" ".join(map(str, args)))
+
+
+# we'll want to log on a separate thread 
+# so that we can log without blocking the main thread
+
+# make a queue for logging
+log_queue = Queue()
+
+# log to a file instead of the console
+def log(msg):
+    # log_queue.put(msg)
+    logging.info(msg)
+    pass
+
+def set_debug(debug):
+    if debug:
+        # print log to console
+        logging.getLogger().addHandler(logging.StreamHandler())
+
+
+def pow2db(x):
+    return 10 * np.log10(x + 1e-6)
+
+
+def db2pow(x):
+    return 10 ** (x / 10)

token_telephone/vamp_helper.py

@@ -0,0 +1,172 @@
+from pathlib import Path
+import time
+import os
+from contextlib import contextmanager
+import random
+
+import numpy as np
+import audiotools as at
+from audiotools import AudioSignal
+import argbind
+import shutil
+import torch
+import yaml
+
+
+from vampnet.interface import Interface, signal_concat
+from vampnet import mask as pmask
+
+from ttutil import log
+
+# TODO: incorporate discord bot (if mem allows)
+# in a separate thread, send audio samples for listening
+# and send back the results
+# as well as the params for sampling
+# also a command that lets you clear the current signal 
+# if you want to start over
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+VAMPNET_DIR = Path(".").resolve()
+
+@contextmanager
+def chdir(path):
+    old_dir = os.getcwd()
+    os.chdir(path)
+    try:
+        yield
+    finally:
+        os.chdir(old_dir)
+
+def load_interface(model_choice="default") -> Interface:
+    with chdir(VAMPNET_DIR):
+
+
+        # populate the model choices with any interface.yml files in the generated confs
+        MODEL_CHOICES = {
+            "default": {
+                "Interface.coarse_ckpt": "models/vampnet/coarse.pth", 
+                "Interface.coarse2fine_ckpt": "models/vampnet/c2f.pth",
+                "Interface.codec_ckpt": "models/vampnet/codec.pth",
+            }
+        }
+        generated_confs = Path("conf/generated")
+        for conf_file in generated_confs.glob("*/interface.yml"):
+            with open(conf_file) as f:
+                _conf = yaml.safe_load(f)
+
+                # check if the coarse, c2f, and codec ckpts exist
+                # otherwise, dont' add this model choice
+                if not (
+                    Path(_conf["Interface.coarse_ckpt"]).exists() and 
+                    Path(_conf["Interface.coarse2fine_ckpt"]).exists() and 
+                    Path(_conf["Interface.codec_ckpt"]).exists()
+                ):
+                    continue
+
+                MODEL_CHOICES[conf_file.parent.name] = _conf
+
+    interface = Interface(
+        device=device, 
+        coarse_ckpt=MODEL_CHOICES[model_choice]["Interface.coarse_ckpt"],
+        coarse2fine_ckpt=MODEL_CHOICES[model_choice]["Interface.coarse2fine_ckpt"],
+        codec_ckpt=MODEL_CHOICES[model_choice]["Interface.codec_ckpt"],
+    )
+
+    interface.model_choices = MODEL_CHOICES
+    interface.to("cuda" if torch.cuda.is_available() else "cpu")
+    return interface
+
+def load_model(interface: Interface, model_choice: str):
+    interface.reload(
+        interface.model_choices[model_choice]["Interface.coarse_ckpt"],
+        interface.model_choices[model_choice]["Interface.coarse2fine_ckpt"],
+    )
+
+def ez_variation(
+        interface,
+        sig: AudioSignal,
+        seed: int = None, 
+        model_choice: str = None,  
+    ):
+    t0 = time.time()
+    
+    if seed is None:
+        seed = int(torch.randint(0, 2**32, (1,)).item())
+    at.util.seed(seed)
+
+    # reload the model if necessary
+    if model_choice is not None:
+        load_model(interface, model_choice)
+
+    # SAMPLING MASK PARAMS, hard code for now, we'll prob want a more preset-ey thing for the actual thin
+    # we probably honestly just want to oscillate between the same 4 presets
+    # in a predictable order such that they have a predictable outcome
+    periodic_p = random.choice([3])
+    n_mask_codebooks = 3
+    sampletemp = random.choice([1.0,])
+    dropout = random.choice([0.0, 0.0])
+
+    top_p = None # NOTE: top p may be the culprit behind the collapse into single pitches. 
+
+    # parameters for the build_mask function
+    build_mask_kwargs = dict(
+        rand_mask_intensity=1.0,
+        prefix_s=0.0,
+        suffix_s=0.0,
+        periodic_prompt=int(periodic_p),
+        periodic_prompt2=int(periodic_p),
+        periodic_prompt_width=1,
+        _dropout=dropout,
+        upper_codebook_mask=int(n_mask_codebooks), 
+        upper_codebook_mask_2=int(n_mask_codebooks),
+    )
+
+    # parameters for the vamp function
+    vamp_kwargs = dict(
+        temperature=sampletemp,
+        typical_filtering=True, 
+        typical_mass=0.15, 
+        typical_min_tokens=64, 
+        top_p=top_p,
+        seed=seed,
+        sample_cutoff=1.0,
+    )
+
+    # save the mask as a txt file
+    interface.set_chunk_size(10.0)
+    sig, mask, codes = interface.ez_vamp(
+        sig, 
+        batch_size=1,
+        feedback_steps=1,
+        time_stretch_factor=1,
+        build_mask_kwargs=build_mask_kwargs,
+        vamp_kwargs=vamp_kwargs,
+        return_mask=True,
+    )
+
+    log(f"vamp took {time.time() - t0} seconds")
+    return sig
+
+
+
+def main():
+    import tqdm
+
+    interface = load_interface()
+    sig = AudioSignal.excerpt("assets/example.wav", duration=7.0)
+    sig = interface.preprocess(sig)
+    sig.write('ttout/in.wav')
+    insig = sig.clone()
+
+    fdbk_every = 4
+    fdbk = 0.5
+
+    for i in tqdm.tqdm(range(1000)): 
+        sig = ez_variation(interface, sig, model_choice="orchestral")
+        sig.write(f'ttout/out{i}.wav')
+    
+
+if __name__ == "__main__":
+    main()

vampnet/__init__.py

@@ -0,0 +1,90 @@
+
+from . import modules
+from pathlib import Path
+from . import scheduler
+from .interface import Interface
+from .modules.transformer import VampNet
+
+
+__version__ = "0.0.1"
+
+ROOT = Path(__file__).parent.parent
+MODELS_DIR = ROOT / "models" / "vampnet"
+
+from huggingface_hub import hf_hub_download, HfFileSystem
+DEFAULT_HF_MODEL_REPO = "hugggof/vampnet"
+FS = HfFileSystem()
+
+def download_codec():
+    # from dac.model.dac import DAC
+    from lac.model.lac import LAC as DAC
+    repo_id = DEFAULT_HF_MODEL_REPO
+    filename = "codec.pth"
+    codec_path = hf_hub_download(
+        repo_id=repo_id,
+        filename=filename,
+        subfolder=None, 
+        local_dir=MODELS_DIR
+    )
+    return codec_path
+    
+
+def download_default():
+    filenames = ["coarse.pth", "c2f.pth"]
+    repo_id = DEFAULT_HF_MODEL_REPO
+    paths = []
+    for filename in filenames:
+        path = f"{MODELS_DIR}/{filename}"
+        if not Path(path).exists():
+            path = hf_hub_download(
+                repo_id=repo_id,
+                filename=filename,
+                subfolder=None, 
+                local_dir=MODELS_DIR,
+                local_dir_use_symlinks=False,
+                local_files_only=False
+            )
+        paths.append(path)
+    
+    # load the models
+    return paths[0], paths[1]
+
+
+def download_finetuned(name):
+    repo_id = f"{DEFAULT_HF_MODEL_REPO}"
+    filenames = ["coarse.pth", "c2f.pth"]
+    paths = []
+    for filename in filenames:
+        path = f"{MODELS_DIR}/{name}/loras/{filename}"
+        if not Path(path).exists():
+            path = hf_hub_download(
+                repo_id=repo_id,
+                filename=filename,
+                subfolder=f"loras/{name}", 
+                local_dir=MODELS_DIR, 
+                local_dir_use_symlinks=False,
+                local_files_only=False
+            )
+        paths.append(path)
+    
+    # load the models
+    return paths[0], paths[1]
+    
+def list_finetuned():
+    diritems = FS.listdir(f"{DEFAULT_HF_MODEL_REPO}/loras")
+    # iterate through all the names
+    valid_diritems = []
+    for item in diritems:
+        model_file_items = FS.listdir(item["name"])
+        item_names = [item["name"].split("/")[-1] for item in model_file_items]
+        # check that theres a "c2f.pth" and "coarse.pth" in the items
+        c2f_exists = "c2f.pth" in item_names
+        coarse_exists = "coarse.pth" in item_names
+        if c2f_exists and coarse_exists:
+            valid_diritems.append(item)
+
+    # get the names of the valid items
+    names = [item["name"].split("/")[-1] for item in valid_diritems]
+    return names
+
+

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/interface.py

@@ -0,0 +1,623 @@
+import os
+from pathlib import Path
+import math
+import logging
+
+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,
+        compile=True,
+    ):
+        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)
+        self.codec_path = Path(codec_ckpt)
+
+        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,
+        )
+        self.coarse_path = Path(coarse_ckpt)
+
+        # check if we have a coarse2fine ckpt
+        if coarse2fine_ckpt is not None:
+            self.c2f_path = Path(coarse2fine_ckpt)
+            self.c2f = _load_model(
+                ckpt=coarse2fine_ckpt,
+                lora_ckpt=coarse2fine_lora_ckpt,
+                device=device,
+                chunk_size_s=coarse2fine_chunk_size_s,
+            )
+        else:
+            self.c2f_path = None
+            self.c2f = None
+
+        if wavebeat_ckpt is not None:
+            logging.debug(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
+        self.loudness = -24.0
+
+        if compile:
+            logging.debug(f"compiling models")
+            self.coarse = torch.compile(self.coarse)
+            if self.c2f is not None:
+                self.c2f = torch.compile(self.c2f)
+            self.codec = torch.compile(self.codec)
+
+
+    @classmethod
+    def default(cls):
+        from . import download_codec, download_default
+        print(f"loading default vampnet")
+        codec_path = download_codec()
+        coarse_path, c2f_path = download_default()
+    
+        return Interface(
+            coarse_ckpt=coarse_path,
+            coarse2fine_ckpt=c2f_path,
+            codec_ckpt=codec_path,
+        )
+
+    @classmethod
+    def available_models(cls):
+        from . import list_finetuned
+        return list_finetuned()
+
+
+    def load_finetuned(self, name: str):
+        assert name in self.available_models(), f"{name} is not a valid model name"
+        from . import download_finetuned
+        coarse_path, c2f_path = download_finetuned(name)
+        self.reload(
+            coarse_ckpt=coarse_path,
+            c2f_ckpt=c2f_path,
+        )
+
+    def reload(
+        self, 
+        coarse_ckpt: str = None,
+        c2f_ckpt: str = None,
+    ):
+        if coarse_ckpt is not None:
+            # check if we already loaded, if so, don't reload
+            if self.coarse_path == Path(coarse_ckpt):
+                logging.debug(f"already loaded {coarse_ckpt}")
+            else:
+                self.coarse = _load_model(
+                    ckpt=coarse_ckpt,  
+                    device=self.device,
+                    chunk_size_s=self.coarse.chunk_size_s,
+                )
+                self.coarse_path = Path(coarse_ckpt)
+                logging.debug(f"loaded {coarse_ckpt}")
+
+        if c2f_ckpt is not None:
+            if self.c2f_path == Path(c2f_ckpt):
+                logging.debug(f"already loaded {c2f_ckpt}")
+            else:
+                self.c2f = _load_model(
+                    ckpt=c2f_ckpt,
+                    device=self.device,
+                    chunk_size_s=self.c2f.chunk_size_s,
+                )
+                self.c2f_path = Path(c2f_ckpt)
+                logging.debug(f"loaded {c2f_ckpt}")
+        
+    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 decode(self, z: torch.Tensor):
+        return self.coarse.decode(z, self.codec)
+    
+    def _preprocess(self, signal: AudioSignal):
+        signal = (
+            signal.clone()
+            .resample(self.codec.sample_rate)
+            .to_mono()
+            .normalize(self.loudness)
+            .ensure_max_of_audio(1.0)
+        )
+        logging.debug(f"length before codec preproc: {signal.samples.shape}")
+        signal.samples, length = self.codec.preprocess(signal.samples, signal.sample_rate)
+        logging.debug(f"length after codec preproc: {signal.samples.shape}")
+        return signal
+    
+    @torch.inference_mode()
+    def encode(self, signal: AudioSignal):
+        signal = signal.to(self.device)
+        signal = self._preprocess(signal)
+        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)
+        logging.debug(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]
+        logging.debug(f"beats_z: {len(beats_z)}")
+        logging.debug(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 set_chunk_size(self, chunk_size_s: float):
+        self.coarse.chunk_size_s = chunk_size_s
+        
+    @torch.inference_mode()
+    def coarse_to_fine(
+        self, 
+        z: torch.Tensor,
+        mask: torch.Tensor = None,
+        return_mask: bool = False,
+        **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), value=1) 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)
+            logging.debug(f"appended {n_codebooks_to_append} codebooks to z")
+
+        # 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
+            
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                chunk = self.c2f.generate(
+                    codec=self.codec,
+                    time_steps=chunk_len,
+                    start_tokens=chunk,
+                    return_signal=False,
+                    mask=mask_chunk,
+                    cfg_guidance=None,
+                    **kwargs
+                )
+                fine_z.append(chunk)
+
+        fine_z = torch.cat(fine_z, dim=-1)
+        if return_mask:
+            return fine_z[:, :, :length].clone(), apply_mask(fine_z, mask, self.c2f.mask_token)[0][:, :, :length].clone()
+        
+        return fine_z[:, :, :length].clone()
+    
+    @torch.inference_mode()
+    def coarse_vamp(
+        self, 
+        z, 
+        mask,
+        return_mask=False,
+        gen_fn=None,
+        **kwargs
+    ):
+        # coarse z
+        cz = z[:, : self.coarse.n_codebooks, :].clone()
+        mask = mask[:, : self.coarse.n_codebooks, :]
+        # 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)}"
+
+        # cut into chunks, keep the last chunk separate if it's too small
+        chunk_len = self.s2t(self.coarse.chunk_size_s)
+        n_chunks = math.ceil(cz.shape[-1] / chunk_len)
+        last_chunk_len = cz.shape[-1] % chunk_len
+
+        cz_chunks = []
+        mask_chunks = []
+        for i in range(n_chunks):
+            chunk = cz[:, :, i * chunk_len : (i + 1) * chunk_len]
+            mask_chunk = mask[:, :, i * chunk_len : (i + 1) * chunk_len]
+
+            # make sure that the very first and last timestep of each chunk is 0 so that we don't get a weird
+            # discontinuity when we stitch the chunks back together
+            # only if there's already a 0 somewhere in the chunk
+            if torch.any(mask_chunk == 0):
+                mask_chunk[:, :, 0] = 0
+                mask_chunk[:, :, -1] = 0
+
+            cz_chunks.append(chunk)
+            mask_chunks.append(mask_chunk)
+
+        # now vamp each chunk
+        cz_masked_chunks = []
+        cz_vamped_chunks = []
+        for chunk, mask_chunk in zip(cz_chunks, mask_chunks):
+            cz_masked_chunk, mask_chunk = apply_mask(chunk, mask_chunk, self.coarse.mask_token)
+            cz_masked_chunk = cz_masked_chunk[:, : self.coarse.n_codebooks, :]
+            cz_masked_chunks.append(cz_masked_chunk)
+            
+
+            gen_fn = gen_fn or self.coarse.generate
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                c_vamp_chunk = gen_fn(
+                    codec=self.codec,
+                    time_steps=chunk_len,
+                    start_tokens=cz_masked_chunk,
+                    return_signal=False,
+                    mask=mask_chunk,
+                    **kwargs
+                )
+                cz_vamped_chunks.append(c_vamp_chunk)
+        
+        # stitch the chunks back together
+        cz_masked = torch.cat(cz_masked_chunks, dim=-1)
+        c_vamp = torch.cat(cz_vamped_chunks, dim=-1)
+
+        # 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
+    
+    def build_mask(self, 
+        z: torch.Tensor,
+        sig: AudioSignal = None,
+        rand_mask_intensity: float = 1.0,
+        prefix_s: float = 0.0,
+        suffix_s: float = 0.0,
+        periodic_prompt: int = 7,
+        periodic_prompt_width: int = 1,
+        onset_mask_width: int = 0, 
+        _dropout: float = 0.0,
+        upper_codebook_mask: int = 3,
+        ncc: int = 0,
+    ):
+        mask = linear_random(z, rand_mask_intensity)
+        mask = mask_and(
+            mask,
+            inpaint(z, self.s2t(prefix_s), self.s2t(suffix_s)),
+        )
+
+        pmask = periodic_mask(z, periodic_prompt, periodic_prompt_width, random_roll=True)
+        mask = mask_and(mask, pmask)
+
+        if onset_mask_width > 0:
+            assert sig is not None, f"must provide a signal to use onset mask"
+            mask = mask_and(
+                mask, onset_mask(
+                    sig, z, self, 
+                    width=onset_mask_width
+                )
+            )
+
+        mask = dropout(mask, _dropout)
+        mask = codebook_unmask(mask, ncc)
+
+        mask = codebook_mask(mask, int(upper_codebook_mask), None)
+        return mask
+
+    def vamp(
+        self, 
+        codes: torch.Tensor,
+        mask: torch.Tensor,
+        batch_size: int = 1,
+        feedback_steps: int = 1,
+        time_stretch_factor: int = 1,
+        return_mask: bool = False,
+        **kwargs,
+    ):
+        z = codes
+
+        # expand z to batch size
+        z = z.expand(batch_size, -1, -1)
+        mask = mask.expand(batch_size, -1, -1)
+
+        # stretch mask and z to match the time stretch factor
+        # we'll add (stretch_factor - 1) mask tokens in between each timestep of z
+        # and we'll make the mask 1 in all the new slots we added
+        if time_stretch_factor > 1:
+            z = z.repeat_interleave(time_stretch_factor, dim=-1)
+            mask = mask.repeat_interleave(time_stretch_factor, dim=-1)
+            added_mask = torch.ones_like(mask)
+            added_mask[:, :, ::time_stretch_factor] = 0
+            mask = mask.bool() | added_mask.bool()
+            mask = mask.long()
+            
+        # the forward pass
+        logging.debug(z.shape)
+        logging.debug("coarse!")
+        zv, mask_z = self.coarse_vamp(
+            z,
+            mask=mask,
+            return_mask=True, 
+            **kwargs
+        )
+
+        # add the top codebooks back in
+        if zv.shape[1] < z.shape[1]:
+            logging.debug(f"adding {z.shape[1] - zv.shape[1]} codebooks back in")
+            zv = torch.cat(
+                [zv, z[:, self.coarse.n_codebooks :, :]], 
+                dim=1
+            )
+
+        # now, coarse2fine
+        logging.debug(f"coarse2fine!")
+        zv, fine_zv_mask = self.coarse_to_fine(
+            zv, 
+            mask=mask,
+            typical_filtering=True,
+            _sampling_steps=[2],
+            return_mask=True
+        )
+        mask_z = torch.cat(
+            [mask_z[:, :self.coarse.n_codebooks, :], fine_zv_mask[:, self.coarse.n_codebooks:, :]], 
+            dim=1
+        )
+
+        z = zv
+
+        if return_mask:
+            return z, mask_z.cpu(),
+        else:
+            return z
+
+    def visualize_codes(self, z: torch.Tensor):
+        import matplotlib.pyplot as plt
+        # make sure the figsize is square when imshow is called
+        fig = plt.figure(figsize=(10, 7))
+        # in subplots, plot z[0] and the mask
+        # set title to "codes" and "mask"
+        fig.add_subplot(2, 1, 1)
+        plt.imshow(z[0].cpu().numpy(), aspect='auto', origin='lower', cmap="tab20")
+        plt.title("codes")
+        plt.ylabel("codebook index")
+        # set the xticks to seconds
+
+
+if __name__ == "__main__":
+    import audiotools as at
+    import logging
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    torch.set_logging.debugoptions(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)
+
+
+    mask = interface.build_mask(
+        z=z,
+        sig=sig,
+        rand_mask_intensity=1.0,
+        prefix_s=0.0,
+        suffix_s=0.0,
+        periodic_prompt=7,
+        periodic_prompt2=7,
+        periodic_prompt_width=1,
+        onset_mask_width=5, 
+        _dropout=0.0,
+        upper_codebook_mask=3,
+        upper_codebook_mask_2=None,
+        ncc=0,
+    )
+
+    zv, mask_z = interface.coarse_vamp(
+        z, 
+        mask=mask,
+        return_mask=True, 
+        gen_fn=interface.coarse.generate
+    )
+    
+
+    use_coarse2fine = True
+    if use_coarse2fine: 
+        zv = interface.coarse_to_fine(zv, mask=mask)
+        breakpoint()
+
+    mask = interface.decode(mask_z).cpu()
+
+    sig = interface.decode(zv).cpu()
+
+
+    logging.debug("done")
+
+        

vampnet/mask.py

@@ -0,0 +1,226 @@
+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()
+        r = r[:, None, None]
+
+    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, val1: int, val2: int = None):
+    mask = mask.clone()
+    mask[:, val1:, :] = 1
+    # val2 = val2 or val1
+    # vs = torch.linspace(val1, val2, mask.shape[1])
+    # for t, v in enumerate(vs):
+    #     v = int(v)
+    #     mask[:, v:, t] = 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
+
+    onset_frame_idxs = librosa.onset.onset_detect(
+        y=sig.samples[0][0].detach().cpu().numpy(), sr=sig.sample_rate, 
+        hop_length=interface.codec.hop_length,
+        backtrack=True,
+    )
+    if len(onset_frame_idxs) == 0:
+        print("no onsets detected")
+    print("onset_frame_idxs", onset_frame_idxs)
+    print("mask shape", z.shape)
+
+    mask = torch.ones_like(z)
+    for idx in onset_frame_idxs:
+        mask[:, :, idx-width:idx+width] = 0
+
+    return mask
+
+
+
+if __name__ == "__main__":
+    sig = AudioSignal("assets/example.wav")

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/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/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/modules/transformer.py

@@ -0,0 +1,965 @@
+import math
+import logging
+from typing import Optional, Tuple, Union, List
+
+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 decode(self, z, codec):
+        """
+        convert a sequence of latents to a signal. 
+        """
+        assert z.ndim == 3
+
+        # remove mask token
+        z = z.masked_fill(z == self.mask_token, 0)
+        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.all(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.inference_mode()
+    def generate(
+        self, 
+        codec,
+        time_steps: int = 300,
+        _sampling_steps: List[int] = [12],
+        start_tokens: Optional[torch.Tensor] = None,
+        temperature: float = 1.0,
+        mask: Optional[torch.Tensor] = None,
+        mask_temperature: float = 10.5,
+        typical_filtering=True,
+        typical_mass=0.2,
+        typical_min_tokens=64,
+        top_p=None,
+        seed: int = None,
+        sample_cutoff: float = 1.0,
+        return_signal=True,
+        debug=False,
+        causal_weight: float = 0.0,
+        cfg_guidance: float = None,
+    ):
+        if seed is not None:
+            at.util.seed(seed)
+        sampling_steps = sum(_sampling_steps)
+        logging.debug(f"beginning generation with {sampling_steps} steps")
+
+        ##################### 
+        # resolve initial z #
+        #####################
+        z = start_tokens
+        nb = z.shape[0]
+
+        if z is None:
+            z = torch.full((1, self.n_codebooks, time_steps), self.mask_token).to(
+                self.device
+            )
+
+
+
+        #################
+        # 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()
+        
+
+
+        ###########
+        # 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()
+
+        # how many codebooks are we inferring vs conditioning on?
+        n_infer_codebooks = self.n_codebooks - self.n_conditioning_codebooks
+
+        if cfg_guidance is not None:
+            # we need to repeat our tensors
+            z_uncond = torch.full_like(z, self.mask_token)
+
+            z_masked = torch.cat(
+                (z_masked, z_uncond), dim=0
+            )
+            z = torch.cat(
+                (z, z_uncond), dim=0
+            )
+            mask = torch.cat(
+                (mask, torch.full_like(mask, 1)), dim=0
+            )
+
+        #################
+        # begin sampling #
+        #################
+        from tqdm import tqdm
+        for i in range(sampling_steps):
+
+            # our current schedule step
+            r = scalar_to_batch_tensor(
+                (i + 1) / sampling_steps, 
+                z.shape[0]
+            ).to(z.device)
+
+            # get latents
+            latents = self.embedding.from_codes(z_masked, codec)
+
+
+            # infer from latents
+            # NOTE: this collapses the codebook dimension into the sequence dimension
+            logits = self.forward(latents) # b, prob, seq
+
+            if cfg_guidance is not None:
+                logits_cond, logits_uncond = logits[:nb], logits[nb:]
+                logits_cond = cfg_guidance * logits_cond + cfg_guidance * (1 - logits_uncond)
+
+            logits = logits.permute(0, 2, 1)  # b, seq, prob
+            b = logits.shape[0]
+
+            sampled_z, selected_probs = sample_from_logits(
+                logits, sample=(
+                   (i / sampling_steps) <= sample_cutoff
+                ), 
+                temperature=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,
+            )
+
+
+            # 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
+            # add z back into sampled z where the mask was false
+            sampled_z = torch.where(
+                mask.bool(), sampled_z, z_masked
+            )
+
+            # 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
+            )
+
+            z_masked = codebook_unflatten(z_masked, n_infer_codebooks)
+            mask = codebook_unflatten(mask, n_infer_codebooks)
+
+            # add conditioning codebooks back to z_masked
+            z_masked = torch.cat(
+                (z[:, :self.n_conditioning_codebooks, :], z_masked), dim=1
+            )
+
+        # 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
+        )
+
+        if cfg_guidance is not None:
+            sampled_z = sampled_z[:nb]
+
+        if return_signal:
+            return self.decode(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)
+    temperature = temperature.unsqueeze(-1)
+    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)
+
+        logging.debug(f"model has {num_params(model)/1e6:<.3f}M parameters")
+        logging.debug(f"prediction has shape {pred.shape}")
+
+    args = argbind.parse_args()
+    with argbind.scope(args):
+        try_model()
+
+

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/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