gitattributes

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* 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
+coarse.pth filter=lfs diff=lfs merge=lfs -text
+c2f.pth filter=lfs diff=lfs merge=lfs -text
+wavebeat.pth filter=lfs diff=lfs merge=lfs -text
+codec.pth filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,191 @@
+# 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
+conf/generated/*
+runs*/
+
+
+gtzan.zip
+.gtzan_emb_cache
+runs
+
+data/
+src/

.pre-commit-config.yaml

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

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.

app.py

@@ -0,0 +1,704 @@
+# huggingface space exclusive
+import os
+
+# print("installing pyharp")
+# os.system('pip install "pyharp@git+https://github.com/audacitorch/pyharp.git"')
+# print("installing madmom")
+os.system('pip install cython')
+os.system('pip install madmom')
+
+from pathlib import Path
+from typing import Tuple
+import yaml
+import tempfile
+import uuid
+import shutil
+from dataclasses import dataclass, asdict
+
+import numpy as np
+import audiotools as at
+import argbind
+import torch
+
+import gradio as gr
+from vampnet.interface import Interface
+from vampnet import mask as pmask
+
+from pyharp import ModelCard, build_endpoint
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+interface = Interface(
+    device=device,
+    coarse_ckpt="models/vampvat/coarse.pth", 
+    coarse2fine_ckpt="models/vampvat/c2f.pth",
+    codec_ckpt="models/vampvat/codec.pth",
+)
+
+# 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)
+        MODEL_CHOICES[conf_file.parent.name] = _conf
+
+    
+
+OUT_DIR = Path("gradio-outputs")
+OUT_DIR.mkdir(exist_ok=True, parents=True)
+
+
+def load_audio(file):
+    print(file)
+    filepath = file.name
+    sig = at.AudioSignal.salient_excerpt(
+        filepath, 
+        duration=interface.coarse.chunk_size_s
+    )
+    sig = interface.preprocess(sig)
+
+    out_dir = OUT_DIR / "tmp" / str(uuid.uuid4())
+    out_dir.mkdir(parents=True, exist_ok=True)
+    sig.write(out_dir / "input.wav")
+    return sig.path_to_file
+
+
+def load_example_audio():
+    return "./assets/example.wav"
+
+from torch_pitch_shift import pitch_shift, get_fast_shifts
+def shift_pitch(signal, interval: int):
+    signal.samples = pitch_shift(
+        signal.samples, 
+        shift=interval, 
+        sample_rate=signal.sample_rate
+    )
+    return signal
+
+def _vamp(data, return_mask=False):
+    # remove any old files in the output directory (from previous runs)
+    shutil.rmtree(OUT_DIR)
+    OUT_DIR.mkdir()
+
+    out_dir = OUT_DIR / str(uuid.uuid4())
+    out_dir.mkdir()
+    sig = at.AudioSignal(data[input_audio])
+    sig = interface.preprocess(sig)
+
+    # reload the model if necessary
+    interface.reload(
+        coarse_ckpt=MODEL_CHOICES[data[model_choice]]["Interface.coarse_ckpt"],
+        c2f_ckpt=MODEL_CHOICES[data[model_choice]]["Interface.coarse2fine_ckpt"],
+    )
+
+    loudness = sig.loudness()
+    print(f"input loudness is {loudness}")
+
+    if data[pitch_shift_amt] != 0:
+        sig = shift_pitch(sig, data[pitch_shift_amt])
+
+    z = interface.encode(sig)
+
+    ncc = data[n_conditioning_codebooks]
+
+    # build the mask
+    mask = pmask.linear_random(z, data[rand_mask_intensity])
+    mask = pmask.mask_and(
+        mask, pmask.inpaint(
+            z,
+            interface.s2t(data[prefix_s]),
+            interface.s2t(data[suffix_s])
+        )
+    )
+    mask = pmask.mask_and(
+        mask, pmask.periodic_mask(
+            z,
+            data[periodic_p],
+            data[periodic_w],
+            random_roll=True
+        )
+    )
+    if data[onset_mask_width] > 0:
+        mask = pmask.mask_or(
+            mask, pmask.onset_mask(sig, z, interface, width=data[onset_mask_width])
+        )
+    if data[beat_mask_width] > 0:
+        beat_mask = interface.make_beat_mask(
+            sig,
+            after_beat_s=(data[beat_mask_width]/1000), 
+            mask_upbeats=not data[beat_mask_downbeats],
+        )
+        mask = pmask.mask_and(mask, beat_mask)
+
+    # these should be the last two mask ops
+    mask = pmask.dropout(mask, data[dropout])
+    mask = pmask.codebook_unmask(mask, ncc)
+    mask = pmask.codebook_mask(mask, int(data[n_mask_codebooks]))
+
+    print(f"dropout {data[dropout]}")
+    print(f"masktemp {data[masktemp]}")
+    print(f"sampletemp {data[sampletemp]}")
+    print(f"top_p {data[top_p]}")
+    print(f"prefix_s {data[prefix_s]}")
+    print(f"suffix_s {data[suffix_s]}")
+    print(f"rand_mask_intensity {data[rand_mask_intensity]}")
+    print(f"num_steps {data[num_steps]}")
+    print(f"periodic_p {data[periodic_p]}")
+    print(f"periodic_w {data[periodic_w]}")
+    print(f"n_conditioning_codebooks {data[n_conditioning_codebooks]}")
+    print(f"use_coarse2fine {data[use_coarse2fine]}")
+    print(f"onset_mask_width {data[onset_mask_width]}")
+    print(f"beat_mask_width {data[beat_mask_width]}")
+    print(f"beat_mask_downbeats {data[beat_mask_downbeats]}")
+    print(f"stretch_factor {data[stretch_factor]}")
+    print(f"seed {data[seed]}")
+    print(f"pitch_shift_amt {data[pitch_shift_amt]}")
+    print(f"sample_cutoff {data[sample_cutoff]}")
+    
+    
+    _top_p = data[top_p] if data[top_p] > 0 else None
+    # save the mask as a txt file
+    np.savetxt(out_dir / "mask.txt", mask[:,0,:].long().cpu().numpy())
+
+    _seed = data[seed] if data[seed] > 0 else None
+    print(f"processing coarse...")
+    zv, mask_z = interface.coarse_vamp(
+        z, 
+        mask=mask,
+        sampling_steps=data[num_steps],
+        mask_temperature=data[masktemp]*10,
+        sampling_temperature=data[sampletemp],
+        return_mask=True, 
+        typical_filtering=data[typical_filtering], 
+        typical_mass=data[typical_mass], 
+        typical_min_tokens=data[typical_min_tokens], 
+        top_p=_top_p,
+        gen_fn=interface.coarse.generate,
+        seed=_seed,
+        sample_cutoff=data[sample_cutoff],
+    )
+
+    if use_coarse2fine: 
+        print(f"processing coarse to fine...")
+        zv = interface.coarse_to_fine(
+            zv, 
+            mask_temperature=data[masktemp]*10, 
+            sampling_temperature=data[sampletemp],
+            mask=mask,
+            sampling_steps=data[num_steps] // 2,
+            sample_cutoff=data[sample_cutoff], 
+            seed=_seed,
+        )
+
+    sig = interface.to_signal(zv).cpu()
+    print("done")
+
+    
+
+    sig.write(out_dir / "output.wav")
+
+    if return_mask:
+        mask = interface.to_signal(mask_z).cpu()
+        mask.write(out_dir / "mask.wav")
+        return sig.path_to_file, mask.path_to_file
+    else:
+        return sig.path_to_file
+
+def vamp(data):
+    return _vamp(data, return_mask=True)
+
+def api_vamp(data):
+    return _vamp(data, return_mask=False)
+        
+def save_vamp(data):
+    out_dir = OUT_DIR / "saved" / str(uuid.uuid4())
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    sig_in = at.AudioSignal(data[input_audio])
+    sig_out = at.AudioSignal(data[output_audio])
+
+    sig_in.write(out_dir / "input.wav")
+    sig_out.write(out_dir / "output.wav")
+    
+    _data = {
+        "masktemp": data[masktemp],
+        "sampletemp": data[sampletemp],
+        "top_p": data[top_p],
+        "prefix_s": data[prefix_s],
+        "suffix_s": data[suffix_s],
+        "rand_mask_intensity": data[rand_mask_intensity],
+        "num_steps": data[num_steps],
+        "notes": data[notes_text],
+        "periodic_period": data[periodic_p],
+        "periodic_width": data[periodic_w],
+        "n_conditioning_codebooks": data[n_conditioning_codebooks], 
+        "use_coarse2fine": data[use_coarse2fine],
+        "stretch_factor": data[stretch_factor],
+        "seed": data[seed],
+        "samplecutoff": data[sample_cutoff],
+    }
+
+    # save with yaml
+    with open(out_dir / "data.yaml", "w") as f:
+        yaml.dump(_data, f)
+
+    import zipfile
+    zip_path = out_dir.with_suffix(".zip")
+    with zipfile.ZipFile(zip_path, "w") as zf:
+        for file in out_dir.iterdir():
+            zf.write(file, file.name)
+
+    return f"saved! your save code is {out_dir.stem}", zip_path
+
+
+def harp_vamp(_input_audio, _beat_mask_width, _sampletemp):
+
+    out_dir = OUT_DIR / str(uuid.uuid4())
+    out_dir.mkdir()
+    sig = at.AudioSignal(_input_audio)
+    sig = interface.preprocess(sig)
+
+    z = interface.encode(sig)
+
+    # build the mask
+    mask = pmask.linear_random(z, 1.0)
+    if _beat_mask_width > 0:
+        beat_mask = interface.make_beat_mask(
+            sig,
+            after_beat_s=(_beat_mask_width/1000), 
+        )
+        mask = pmask.mask_and(mask, beat_mask)
+
+    # save the mask as a txt file
+    zv, mask_z = interface.coarse_vamp(
+        z, 
+        mask=mask,
+        sampling_temperature=_sampletemp,
+        return_mask=True, 
+        gen_fn=interface.coarse.generate,
+    )
+
+
+    zv = interface.coarse_to_fine(
+        zv, 
+        sampling_temperature=_sampletemp,
+        mask=mask,
+    )
+
+    sig = interface.to_signal(zv).cpu()
+    print("done")
+
+    sig.write(out_dir / "output.wav")
+
+    return sig.path_to_file
+
+with gr.Blocks() as demo:
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("# VampNet Audio Vamping")
+            gr.Markdown("""## Description:
+            This is a demo of the VampNet, a generative audio model that transforms the input audio based on the chosen settings. 
+            You can control the extent and nature of variation with a set of manual controls and presets. 
+            Use this interface to experiment with different mask settings and explore the audio outputs.
+            """)
+
+            gr.Markdown("""
+            ## Instructions:
+            1. You can start by uploading some audio, or by loading the example audio. 
+            2. Choose a preset for the vamp operation, or manually adjust the controls to customize the mask settings. 
+            3. Click the "generate (vamp)!!!" button to apply the vamp operation. Listen to the output audio.
+            4. Optionally, you can add some notes and save the result. 
+            5. You can also use the output as the new input and continue experimenting!
+            """)
+    with gr.Row():
+        with gr.Column():
+
+
+            manual_audio_upload = gr.File(
+                label=f"upload some audio (will be randomly trimmed to max of {interface.coarse.chunk_size_s:.2f}s)",
+                file_types=["audio"]
+            )
+            load_example_audio_button = gr.Button("or load example audio")
+
+            input_audio = gr.Audio(
+                label="input audio",
+                interactive=False, 
+                type="filepath",
+            )
+
+            audio_mask = gr.Audio(
+                label="audio mask (listen to this to hear the mask hints)",
+                interactive=False, 
+                type="filepath",
+            )
+
+            # connect widgets
+            load_example_audio_button.click(
+                fn=load_example_audio,
+                inputs=[],
+                outputs=[ input_audio]
+            )
+
+            manual_audio_upload.change(
+                fn=load_audio,
+                inputs=[manual_audio_upload],
+                outputs=[ input_audio]
+            )
+                
+        # mask settings
+        with gr.Column():
+
+
+            presets = {
+                    "unconditional": {
+                        "periodic_p": 0,
+                        "onset_mask_width": 0,
+                        "beat_mask_width": 0,
+                        "beat_mask_downbeats": False,
+                    }, 
+                    "slight periodic variation": {
+                        "periodic_p": 5,
+                        "onset_mask_width": 5,
+                        "beat_mask_width": 0,
+                        "beat_mask_downbeats": False,
+                    },
+                    "moderate periodic variation": {
+                        "periodic_p": 13,
+                        "onset_mask_width": 5,
+                        "beat_mask_width": 0,
+                        "beat_mask_downbeats": False,
+                    },
+                    "strong periodic variation": {
+                        "periodic_p": 17,
+                        "onset_mask_width": 5,
+                        "beat_mask_width": 0,
+                        "beat_mask_downbeats": False,
+                    },
+                    "very strong periodic variation": {
+                        "periodic_p": 21,
+                        "onset_mask_width": 5,
+                        "beat_mask_width": 0,
+                        "beat_mask_downbeats": False,
+                    },
+                    "beat-driven variation": {
+                        "periodic_p": 0,
+                        "onset_mask_width": 0,
+                        "beat_mask_width": 50,
+                        "beat_mask_downbeats": False,
+                    },
+                    "beat-driven variation (downbeats only)": {
+                        "periodic_p": 0,
+                        "onset_mask_width": 0,
+                        "beat_mask_width": 50,
+                        "beat_mask_downbeats": True,
+                    },
+                    "beat-driven variation (downbeats only, strong)": {
+                        "periodic_p": 0,
+                        "onset_mask_width": 0,
+                        "beat_mask_width": 20,
+                        "beat_mask_downbeats": True,
+                    },
+                }
+
+            preset = gr.Dropdown(
+                label="preset", 
+                choices=list(presets.keys()),
+                value="strong periodic variation",
+            )
+            load_preset_button = gr.Button("load_preset")
+
+            with gr.Accordion("manual controls", open=True):
+                periodic_p = gr.Slider(
+                    label="periodic prompt  (0 - unconditional, 2 - lots of hints, 8 - a couple of hints, 16 - occasional hint, 32 - very occasional hint, etc)",
+                    minimum=0,
+                    maximum=128, 
+                    step=1,
+                    value=3, 
+                )
+
+
+                onset_mask_width = gr.Slider(
+                    label="onset mask width (multiplies with the periodic mask, 1 step ~= 10milliseconds) ",
+                    minimum=0,
+                    maximum=100,
+                    step=1,
+                    value=5,
+                )
+
+                beat_mask_width = gr.Slider(
+                    label="beat prompt (ms)",
+                    minimum=0,
+                    maximum=200,
+                    value=0,
+                )
+                beat_mask_downbeats = gr.Checkbox(
+                    label="beat mask downbeats only?", 
+                    value=False
+                )
+
+                n_mask_codebooks = gr.Number(
+                    label="first upper codebook level to mask",
+                    value=9,
+                )
+
+
+                with gr.Accordion("extras ", open=False):
+                    pitch_shift_amt = gr.Slider(
+                        label="pitch shift amount (semitones)",
+                        minimum=-12,
+                        maximum=12,
+                        step=1,
+                        value=0,
+                    )
+
+                    rand_mask_intensity = gr.Slider(
+                        label="random mask intensity. (If this is less than 1, scatters prompts throughout the audio, should be between 0.9 and 1.0)",
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=1.0
+                    )
+
+                    periodic_w = gr.Slider(
+                        label="periodic prompt width (steps, 1 step ~= 10milliseconds)",
+                        minimum=1,
+                        maximum=20,
+                        step=1,
+                        value=1,
+                    )
+                    n_conditioning_codebooks = gr.Number(
+                        label="number of conditioning codebooks. probably 0", 
+                        value=0,
+                        precision=0,
+                    )
+
+                    stretch_factor = gr.Slider(
+                        label="time stretch factor",
+                        minimum=0,
+                        maximum=64, 
+                        step=1,
+                        value=1, 
+                    )
+
+            preset_outputs = {
+                periodic_p, 
+                onset_mask_width, 
+                beat_mask_width,
+                beat_mask_downbeats,
+            }
+
+            def load_preset(_preset):
+                return tuple(presets[_preset].values())
+
+            load_preset_button.click(
+                fn=load_preset,
+                inputs=[preset],
+                outputs=preset_outputs
+            )
+
+
+            with gr.Accordion("prefix/suffix prompts", open=False):
+                prefix_s = gr.Slider(
+                    label="prefix hint length (seconds)",
+                    minimum=0.0,
+                    maximum=10.0,
+                    value=0.0
+                )
+                suffix_s = gr.Slider(
+                    label="suffix hint length (seconds)",
+                    minimum=0.0,
+                    maximum=10.0,
+                    value=0.0
+                )
+
+            masktemp = gr.Slider(
+                label="mask temperature",
+                minimum=0.0,
+                maximum=100.0,
+                value=1.5
+            )
+            sampletemp = gr.Slider(
+                label="sample temperature",
+                minimum=0.1,
+                maximum=10.0,
+                value=1.0, 
+                step=0.001
+            )
+        
+
+
+            with gr.Accordion("sampling settings", open=False):
+                top_p = gr.Slider(
+                    label="top p (0.0 = off)",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.9
+                )
+                typical_filtering = gr.Checkbox(
+                    label="typical filtering ",
+                    value=False
+                )
+                typical_mass = gr.Slider( 
+                    label="typical mass (should probably stay between 0.1 and 0.5)",
+                    minimum=0.01,
+                    maximum=0.99,
+                    value=0.15
+                )
+                typical_min_tokens = gr.Slider(
+                    label="typical min tokens (should probably stay between 1 and 256)",
+                    minimum=1,
+                    maximum=256,
+                    step=1,
+                    value=64
+                )
+                sample_cutoff = gr.Slider(
+                    label="sample cutoff",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.5, 
+                    step=0.01
+                )
+
+            use_coarse2fine = gr.Checkbox(
+                label="use coarse2fine",
+                value=True, 
+                visible=False
+            )
+
+            num_steps = gr.Slider(
+                label="number of steps (should normally be between 12 and 36)",
+                minimum=1,
+                maximum=128,
+                step=1,
+                value=36
+            )
+
+            dropout = gr.Slider(
+                label="mask dropout",
+                minimum=0.0,
+                maximum=1.0,
+                step=0.01,
+                value=0.0
+            )
+
+
+            seed = gr.Number(
+                label="seed (0 for random)",
+                value=0,
+                precision=0,
+            )
+
+
+
+        # mask settings
+        with gr.Column():
+
+            model_choice = gr.Dropdown(
+                label="model choice", 
+                choices=list(MODEL_CHOICES.keys()),
+                value="default", 
+                visible=True
+            )
+
+            vamp_button = gr.Button("generate (vamp)!!!")
+            output_audio = gr.Audio(
+                label="output audio",
+                interactive=False,
+                type="filepath"
+            )
+
+            notes_text = gr.Textbox(
+                label="type any notes about the generated audio here", 
+                value="",
+                interactive=True
+            )
+            save_button = gr.Button("save vamp")
+            download_file = gr.File(
+                label="vamp to download will appear here",
+                interactive=False
+            )
+            use_as_input_button = gr.Button("use output as input")
+            
+            thank_you = gr.Markdown("")
+
+
+    _inputs = {
+            input_audio, 
+            num_steps,
+            masktemp,
+            sampletemp,
+            top_p,
+            prefix_s, suffix_s, 
+            rand_mask_intensity, 
+            periodic_p, periodic_w,
+            n_conditioning_codebooks, 
+            dropout,
+            use_coarse2fine, 
+            stretch_factor, 
+            onset_mask_width, 
+            typical_filtering,
+            typical_mass,
+            typical_min_tokens,
+            beat_mask_width,
+            beat_mask_downbeats,
+            seed, 
+            model_choice,
+            n_mask_codebooks,
+            pitch_shift_amt, 
+            sample_cutoff
+        }
+  
+    # connect widgets
+    vamp_button.click(
+        fn=vamp,
+        inputs=_inputs,
+        outputs=[output_audio, audio_mask], 
+    )
+
+    api_vamp_button = gr.Button("api vamp", visible=False)
+    api_vamp_button.click(
+        fn=api_vamp,
+        inputs=_inputs, 
+        outputs=[output_audio], 
+        api_name="vamp"
+    )
+
+    use_as_input_button.click(
+        fn=lambda x: x,
+        inputs=[output_audio],
+        outputs=[input_audio]
+    )
+
+    save_button.click(
+        fn=save_vamp,
+        inputs=_inputs | {notes_text, output_audio},
+        outputs=[thank_you, download_file]
+    )
+
+    # harp stuff
+    harp_inputs = [
+        input_audio, 
+        beat_mask_width,
+        sampletemp,
+    ]
+    
+    build_endpoint(
+        inputs=harp_inputs,  
+        output=output_audio, 
+        process_fn=harp_vamp,
+        card=ModelCard(
+            name="vampnet", 
+            description="Generate variations on music input, based on small prompts around the beat. NOTE: vampnet's has a maximum context length of 10 seconds. Please split all audio clips into 10 second chunks, or processing will result in an error. ", 
+            author="Hugo Flores García", 
+            tags=["music", "generative"]
+        ), 
+        visible=False
+    )
+
+demo.queue().launch()

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: 6
+num_workers: 7
+save_iters: [2000, 4000, 10000,20000, 40000, 100000]
+sample_freq: 2000
+val_freq: 1000
+
+AdamW.lr: 0.0001
+
+# let's us organize sound classes into folders and choose from those sound classes uniformly
+AudioDataset.without_replacement: False
+num_iters: 500000

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
+

requirements.txt

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

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

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": str(finetune_dir / "ckpt/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": str(finetune_dir / "ckpt/coarse"),
+        "fine_tune_checkpoint": "./models/vampnet/coarse.pth"
+    }
+
+    interface_conf = {
+        "Interface.coarse_ckpt": f"{finetune_dir}/ckpt/coarse/latest/vampnet/weights.pth",
+
+        "Interface.coarse2fine_ckpt": f"{finetune_dir}/ckpt/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.to_signal(z_mask, state.codec)
+    imputed_true = vn.to_signal(z, state.codec)
+
+    imputed = []
+    for i in range(len(z)):
+        imputed.append(
+            vn.generate(
+                codec=state.codec,
+                time_steps=z.shape[-1],
+                start_tokens=z[i][None, ...],
+                mask=mask[i][None, ...],
+            )   
+        )   
+    imputed = AudioSignal.batch(imputed)
+
+    for i in range(len(val_idx)):
+        imputed_noisy[i].cpu().write_audio_to_tb(
+            f"inpainted_prompt/{i}",
+            writer,
+            step=state.tracker.step,
+            plot_fn=None,
+        )
+        imputed[i].cpu().write_audio_to_tb(
+            f"inpainted_middle/{i}",
+            writer,
+            step=state.tracker.step,
+            plot_fn=None,
+        )
+        imputed_true[i].cpu().write_audio_to_tb(
+            f"reconstructed/{i}",
+            writer,
+            step=state.tracker.step,
+            plot_fn=None,
+        )
+
+
+@torch.no_grad()
+def save_samples(state: State, val_idx: int, writer: SummaryWriter):
+    state.model.eval()
+    state.codec.eval()
+    vn = accel.unwrap(state.model)
+
+    batch = [state.val_data[i] for i in val_idx]
+    batch = at.util.prepare_batch(state.val_data.collate(batch), accel.device)
+
+    signal = apply_transform(state.val_data.transform, batch)
+
+    z = state.codec.encode(signal.samples, signal.sample_rate)["codes"]
+    z = z[:, : vn.n_codebooks, :]
+
+    r = torch.linspace(0.1, 0.95, len(val_idx)).to(accel.device)
+
+
+    mask = pmask.random(z, r)
+    mask = pmask.codebook_unmask(mask, vn.n_conditioning_codebooks)
+    z_mask, mask = pmask.apply_mask(z, mask, vn.mask_token)
+
+    z_mask_latent = vn.embedding.from_codes(z_mask, state.codec)
+
+    z_hat = state.model(z_mask_latent)
+
+    z_pred = torch.softmax(z_hat, dim=1).argmax(dim=1)
+    z_pred = codebook_unflatten(z_pred, n_c=vn.n_predict_codebooks)
+    z_pred = torch.cat([z[:, : vn.n_conditioning_codebooks, :], z_pred], dim=1)
+
+    generated = vn.to_signal(z_pred, state.codec)
+    reconstructed = vn.to_signal(z, state.codec)
+    masked = vn.to_signal(z_mask.squeeze(1), state.codec)
+
+    for i in range(generated.batch_size):
+        audio_dict = {
+            "original": signal[i],
+            "masked": masked[i],
+            "generated": generated[i],
+            "reconstructed": reconstructed[i],
+        }
+        for k, v in audio_dict.items():
+            v.cpu().write_audio_to_tb(
+                f"onestep/_{i}.r={r[i]:0.2f}/{k}",
+                writer,
+                step=state.tracker.step,
+                plot_fn=None,
+            )
+
+    save_sampled(state=state, z=z, writer=writer)
+    save_imputation(state=state, z=z, val_idx=val_idx, writer=writer)
+
+
+
+@argbind.bind(without_prefix=True)
+def load(
+    args,
+    accel: at.ml.Accelerator,
+    tracker: Tracker,
+    save_path: str,
+    resume: bool = False,
+    tag: str = "latest",
+    fine_tune_checkpoint: Optional[str] = None,
+    grad_clip_val: float = 5.0,
+) -> State:
+    codec = DAC.load(args["codec_ckpt"], map_location="cpu")
+    codec.eval()
+
+    model, v_extra = None, {}
+
+    if args["fine_tune"]:
+        assert fine_tune_checkpoint is not None, "Must provide a fine-tune checkpoint"
+        model = torch.compile(
+            VampNet.load(location=Path(fine_tune_checkpoint), 
+                         map_location="cpu", 
+            )
+        )
+        
+    if resume:
+        kwargs = {
+            "folder": f"{save_path}/{tag}",
+            "map_location": "cpu",
+            "package": False,
+        }
+        tracker.print(f"Loading checkpoint from {kwargs['folder']}")
+        if (Path(kwargs["folder"]) / "vampnet").exists():
+            model, v_extra = VampNet.load_from_folder(**kwargs)
+        else:
+            raise ValueError(
+                f"Could not find a VampNet checkpoint in {kwargs['folder']}"
+            )
+
+
+
+
+    model = torch.compile(VampNet()) if model is None else model
+    model = accel.prepare_model(model)
+
+    # assert accel.unwrap(model).n_codebooks == codec.quantizer.n_codebooks
+    assert (
+        accel.unwrap(model).vocab_size == codec.quantizer.quantizers[0].codebook_size
+    )
+
+
+    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,64 @@
+from pathlib import Path
+import random
+import shutil
+import os
+import json 
+
+import argbind
+from tqdm import tqdm
+from tqdm.contrib.concurrent import thread_map
+
+from audiotools.core import util
+
+
+@argbind.bind(without_prefix=True)
+def train_test_split(
+    audio_folder: str = ".", 
+    test_size: float = 0.2,
+    seed: int = 42,
+    pattern: str = "**/*.mp3",
+):
+    print(f"finding audio")
+
+    audio_folder = Path(audio_folder)
+    audio_files = list(tqdm(audio_folder.glob(pattern)))
+    print(f"found {len(audio_files)} audio files")
+    
+    # split according to test_size
+    n_test = int(len(audio_files) * test_size)
+    n_train = len(audio_files) - n_test
+
+    # shuffle
+    random.seed(seed)
+    random.shuffle(audio_files)
+
+    train_files = audio_files[:n_train]
+    test_files = audio_files[n_train:]
+
+
+    print(f"Train files: {len(train_files)}")
+    print(f"Test files: {len(test_files)}")
+    continue_ = input("Continue [yn]? ") or "n"
+
+    if continue_ != "y":
+        return
+    
+    for split, files in (
+        ("train", train_files), ("test", test_files)
+    ):
+        for file in tqdm(files):
+            out_file = audio_folder.parent / f"{audio_folder.name}-{split}" / Path(file).name
+            out_file.parent.mkdir(exist_ok=True, parents=True)
+            os.symlink(file, out_file)
+
+        # save split as json
+        with open(Path(audio_folder) / f"{split}.json", "w") as f:
+            json.dump([str(f) for f in files], f)
+    
+
+    
+if __name__ == "__main__":
+    args  = argbind.parse_args()
+
+    with argbind.scope(args):
+        train_test_split()

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,43 @@
+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/descriptinc/audiotools.git@0.7.2",
+        "gradio", 
+        "loralib",
+        "torch_pitch_shift",
+        "madmom",
+        "pyharp @ git+https://github.com/audacitorch/pyharp.git",
+        "plotly",
+        "umap_learn",
+    ],
+)

vampnet/__init__.py

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

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

vampnet/mask.py

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

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

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