first commit

.dockerignore

@@ -0,0 +1,2 @@
+*.wav
+runs/

.gitattributes

@@ -0,0 +1,4 @@
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+tests/assets/cpu_test_file.pt filter=lfs diff=lfs merge=lfs -text
+tests/assets/quick.pth filter=lfs diff=lfs merge=lfs -text
+tests/assets/slow.pth filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,173 @@
+# 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
+*.wav
+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

.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

Dockerfile

@@ -0,0 +1,36 @@
+FROM us.gcr.io/lyrebird-research/research-image/audio:beta
+
+COPY requirements.txt requirements.txt
+ARG GITHUB_TOKEN
+RUN echo machine github.com login ${GITHUB_TOKEN} > ~/.netrc
+
+COPY env/alias.sh /alias.sh
+COPY env/entry_script.sh /entry_script.sh
+RUN cat /alias.sh >> ~/.zshrc
+
+# USER researcher
+RUN pip install --upgrade -r requirements.txt
+RUN pip install --upgrade tensorflow
+RUN pip install --upgrade librosa
+RUN pip install --upgrade numba
+ENV PYTHONPATH "$PYTHONPATH:/u/home/src"
+ENV NUMBA_CACHE_DIR=/tmp/
+
+USER root
+RUN wget https://github.com/jgm/pandoc/releases/download/2.18/pandoc-2.18-1-amd64.deb
+RUN dpkg -i pandoc-2.18-1-amd64.deb
+RUN apt-get update && apt-get install task-spooler
+
+RUN head -n -1 /entry_script.sh > /entry_script_jupyter.sh
+RUN head -n -1 /entry_script.sh > /entry_script_tensorboard.sh
+RUN head -n -1 /entry_script.sh > /entry_script_gradio.sh
+
+RUN echo \
+    'su -p ${USER} -c "source ~/.zshrc && jupyter lab --ip=0.0.0.0"' >> \
+    /entry_script_jupyter.sh
+RUN echo \
+    'su -p ${USER} -c "source ~/.zshrc && tensorboard --logdir=$TENSORBOARD_PATH --samples_per_plugin audio=500 --bind_all"' >> \
+    /entry_script_tensorboard.sh
+RUN echo \
+    'su -p ${USER} -c "source ~/.zshrc && python app.py --args.load=conf/app.yml"' >> \
+    /entry_script_gradio.sh

README.md

@@ -0,0 +1,287 @@
+# Lyrebird Wav2Wav
+
+This repository contains recipes for training Wav2Wav models.
+
+## Install hooks
+
+First install the pre-commit util:
+
+https://pre-commit.com/#install
+
+    pip install pre-commit  # with pip
+    brew install pre-commit  # on Mac
+
+Then install the git hooks
+
+    pre-commit install
+    # check .pre-commit-config.yaml for details of hooks
+
+Upon `git commit`, the pre-commit hooks will be run automatically on the stage files (i.e. added by `git add`)
+
+**N.B. By default, pre-commit checks only run on staged files**
+
+If you need to run it on all files:
+
+    pre-commit run --all-files
+
+## Usage & model zoo
+
+To download the model, one must be authenticated to the `lyrebird-research` project on Google Cloud.
+To see all available models, run
+
+```bash
+python -m wav2wav.list_models
+```
+
+which outputs something like this:
+
+```
+gs://research-models/wav2wav
+└── prod
+    └── v3
+        └── ckpt
+            ├── best
+            │   └── generator
+            │       ├── ❌ model.onnx
+            │       ├── ❌ nvidia_geforce_rtx_2080_ti_11_7.trt
+            │       ├── ✅ package.pth
+            │       ├── ❌ tesla_t4_11_7.trt
+            │       └── ✅ weights.pth
+            └── latest
+                └── generator
+                    ├── ❌ package.pth
+                    └── ❌ weights.pth
+    └── v2
+        ...
+└── dev
+    ...
+```
+
+This will show all the models that are available on GCP. Models that are available locally are marked with a ✅, while those not available locally
+are marked with ❌. `.onnx` indicates a model that must be run with
+the `ONNX` runtime, while `.trt` indicate models that have been optimized
+with TensorRT. Note that TensorRT models are specific to GPU and CUDA
+runtime, and their file names indicate what to use to run them.
+
+`package.pth` is a version of the model that is saved using `torch.package`,
+and contains a copy of the model code within it, which allow it to work
+even if the model code in `wav2wav/modules/generator.py` changes. `weights.pth`
+contains the model weights, and the code must match the code used
+to create the model.
+
+To use a model from this list, simply write its path and give it to the `enhance` script,
+like so:
+
+```
+python -m wav2wav.interface \
+  [input_path]
+  --model_path=prod/v3/ckpt/best/generator/weights.pth
+  --output_path [output_path]
+```
+
+Models are downloaded to the location set by the environment variable `MODEL_LOCAL_PATH`, and defaults to `~/.wav2wav/models`. Similarly,
+The model bucket is determined by `MODEL_GCS_PATH` and defaults to
+`gs://research-models/wav2wav/`.
+
+## Development
+### Setting everything up
+
+Run the setup script to set up your environment via:
+
+```bash
+python env/setup.py
+```
+
+The setup script does not require any dependencies beyond just Python.
+Once run, follow the instructions it prints out to create your
+environment file, which will be at `env/env.sh`.
+
+Note that if this is a new machine, and
+the data is not downloaded somewhere on it already, it will ask you
+for a directory to download the data to.
+
+For Github setup, if you don't have a .netrc token, create one by going to your Github profile -> Developer settings -> Personal access tokens -> Generate new token. Copy the token and [keep it secret, keep it safe](https://www.youtube.com/watch?v=iThtELZvfPs).
+
+When complete, run:
+
+```bash
+source env/env.sh
+```
+
+Now build and launch the Docker containers:
+
+```bash
+docker compose up -d
+```
+
+This builds and runs a Jupyter notebook and Tensorboard
+in the background, which points to your `TENSORBOARD_PATH`
+env. variable.
+
+Now, launch your development environment via:
+
+```bash
+docker compose run dev
+```
+
+To tear down your development environment, just do
+
+```bash
+docker compose down
+```
+
+### Downloading data and pre-processing
+Next, from within the Docker environment (or an appropriately configured Conda environment with environment variables set as above), do the following:
+
+```
+python -m wav2wav.preprocess.download
+```
+
+This will download all the necessary data, which are referenced by
+the CSV files in `conf/audio/*`. These CSVs were generated via
+`python -m wav2wav.preprocess.organize`.
+
+### Launching an experiment
+
+Experiments are first _staged_ by running the `stage` command (which corresponds to the script `scripts/exp/stage.py`).
+
+`stage` creates a directory with a copy of all of the Git-tracked files in the root repository.`stage` launches a shell into said directory, so all commands are run on the
+copy of the original repository code. This is useful for rewinding to an old experiment
+and resuming it, for example. Even if the repository code changes, the snapshot in the experiment directory is unchanged from the original run, so it can be re-used.
+
+Then, the experiment can be run via:
+
+```bash
+torchrun --nproc_per_node gpu \
+  scripts/exp/train.py \
+  --args.load=conf/args.yml \
+```
+
+The full settings are in [conf/daps/train.yml](conf/daps/train.yml).
+
+### Evaluating an experiment
+
+There are two ways to evaluate an experiment: quantitative and qualitative.
+For the first, we can use the `scripts/exp/evaluate.py` script. This script evaluates the model over the `val_data` and `test_data`, defined in your
+`train` script, and takes as input an experiment directory. The metrics
+computed by this script are saved to the same folder.
+
+The other way is via a preference test. Let's say we want to compare
+the v3 prod model against the v2 prod model. to do this, we use the
+`scripts/exp/qa.py` script. This script creates a zip file containing all
+the samples and an HTML page for easy viewing. It also creates a Papaya
+preference test. Use it like this:
+
+```bash
+WAV2WAV_MODELS=a,b python scripts/exp/qa.py \
+  --a/model_path prod/v3/ckpt/best/generator/package.pth \
+  --b/model_path prod/v2/ckpt/best/generator/package.pth \
+  --a/name "v3" --b/name "v2" \
+  --device cuda:0 \
+  --n_samples 20 \
+  --zip_path "samples/out.zip"
+```
+
+### Useful commands
+
+#### Monitoring the machine
+
+There's a useful `tmux` workspace that you can launch via:
+
+```bash
+tmuxp load ./workspace.yml
+```
+
+which will have a split pane with a shell to launch commands on the left,
+and GPU monitoring, `htop`, and a script that watches for changes in your
+directory on the right, in three split panes.
+
+#### Cleaning up after a run
+
+Sometimes DDP runs fail to clear themselves out of the machine. To fix this, run
+
+```bash
+cleanup
+```
+
+### Deploying a new model to production
+
+Okay, so you ran a model and it seems promising and you want to upload it
+to GCS so it can be QA'd fully, and then shipped. First, upload
+your experiment to the `dev` bucket on GCS via:
+
+```bash
+gsutil cp -r /path/to/{exp_name} gs://research-models/wav2wav/dev/{exp_name}
+```
+
+Once uploaded, QA can access the models by specifying
+`model_path=dev/{exp_name}/ckpt/{best,latest}/generator/package.pth` when using the
+`wav2wav.interface.enhance` function. If it passes QA, and is scheduled to
+ship to production, then next we have to generate the TensorRT model file,
+which requires us to have a machine that matches that of a production machine.
+
+There is a script that automates this procedure, that does not require any
+fiddling from our end. Navigate to the repository root and run:
+
+```
+python scripts/utils/convert_on_gcp.py dev/{exp_name}/ckpt/{best,latest}//generator/weights.pth
+```
+
+This will provision the machine, download the relevant model from GCS, optimize it on
+the production GPU with the correct CUDA runtime, and then upload the generated `.trt`
+and `.onnx` models back to the bucket.
+
+Finally, copy the model to the `prod` bucket, incrementing the version number by one:
+
+```bash
+gsutil cp -r gs://research-models/wav2wav/dev/{exp_name} gs://research-models/wav2wav/prod/v{N}
+```
+
+where `N` is the next version (e.g. if v3 is the latest, the new one is v4). Then, update
+the model table in [Notion](https://www.notion.so/descript/fc04de4b46e6417eba1d06bdc8de6c75?v=e56db4e6b37c4d9b9eca8d9be15c826a) with the new model.
+
+Once the above is all done, we update the code in two places:
+
+1. In `interface.py`, we update `PROD_MODEL_PATH` to point to the `weights.pth`
+   for whichever tag ended up shipping (either `best` or `latest`).
+2. In `interface.py`, we update `PROD_TRT_PATH` to point the generated
+   TensorRT checkpoint generated by the script above.
+
+After merging to master, a new Docker image will be created, and one can update the relevant lines
+in descript-workflows like in this [PR](https://github.com/descriptinc/descript-workflows/pull/477/files).
+
+We have Github action workflows in [.github/workflows/deploy.yml](.github/workflows/deploy.yml) to build and deploy new docker images. Two images are built - one for staging and another for production.
+To deploy a new release version, follow the instructions in [this coda doc](https://coda.io/d/Research-Engineering_dOABAWL46p-/Deploying-Services_su1am#_lu7E8).
+
+Coda doc with informations about deploying speech-enhance worker is [here](https://coda.io/d/Research-Engineering_dOABAWL46p-/Deploying-Services_su1am#_lu7E8).
+
+And that's it! Once the new staging is built, you're done.
+
+## Testing
+
+### Profiling and Regression testing
+
+- The [profiling script](tests/profile_inference.py) profiles the `wav2wav.interface.enhance` function.
+- NOTE: ALWAYS run the profiler on a T4 GPU. ALWAYS run the profiling in isolation i.e kill all other processes on the GPU. Recommended vm size on GCP is `n1-standard-32` as the stress test of six hours of audio requires ~35GB of system memory.
+- To run profiling use the [profiling script](tests/profile_inference.py) via command `python3 -m tests.profile_inference`. Results will be printed after `1` run.
+- Use the [test_regression.py](tests/test_regression.py) script to run tests that
+  - compare performance stats of current model with known best model
+  - test for output deviation from the last model
+- Run `git lfs checkout` to checkout input file and model weights required for testing the model.
+- To launch these tests, run `python3 -m pytest tests/test_regression.py -v`.
+- As a side effect, this will update the `tests/stat.csv` file if the current model performs better than last best known model as per `tests/stat.csv`.
+- NOTE: In case of architecture change, purge the weights files : `tests/assets/{quick|slow}.pth` and reference stat file : `tests/assets/baseline.json` file. Running the [test_regression.py](tests/test_regression.py) script in absence of reference stat file, will generate new baseline referece stats as well as append new performance stats to stats file. In the absence of saved weights, new weights are generated and saved on disk. Make sure to commit these files (stat.csv, baseline.json, *.pth) when the model architecture changes.
+
+### Unit tests
+Regular unit tests that test functionality such as training resume etc. These are run on CPU. Update them when new features are added.
+
+### Profiling tests
+These tests profile the model's resource consumption. They are run on T4 GPU with 32 cores and >35GB memory. Their usage is reported in the above sections.
+
+### Functional tests
+These tests detect deviation from known baseline model. A category of these tests ensure that a new pytorch model doesn't deviate from the previous one. Another category ensures that the TensorRT version of the current pytorch model doens't deviate from it. These tests are marked with the marker `output_qa` and can be run via the command line `python3 -m pytest -v -m output_qa`. Some of these tests require a GPU.
+
+### CI tests
+- The tests are divided into two categories depending on the platform requirement - CPU tests and GPU tests.
+- The CPU tests contains unit tests.
+- The GPU tests contain a subset of functional tests. These tests can be run by the command `python3 -m pytest -v -m gpu_ci_test`.

changelog.md

@@ -0,0 +1,11 @@
+# v0.3.4
+- Upgraded nussl, removing unneeded code from interface, and improving interaction.
+
+# v0.3.1
+- Rich interaction with the model.
+
+# v0.2.0
+- Release with script to process zipfiles for QA + gsutil sync for model weights.
+
+# v0.1.0
+- Initial release.

conf/vampnet-c2f.yml

@@ -0,0 +1,9 @@
+$include:
+  - conf/vampnet.yml
+
+VampNet.n_codebooks: 9
+VampNet.n_conditioning_codebooks: 3
+
+train/AudioDataset.duration: 3
+val/AudioDataset.duration: 3
+test/AudioDataset.duration: 3

conf/vampnet.yml

@@ -0,0 +1,56 @@
+
+wav2wav_ckpt: /u/home/src/runs/codec-ckpt/codec.pth
+save_path: ckpt
+max_epochs: 1000000
+epoch_length: 1000
+save_audio_epochs: 2
+val_idx: [0,1,2,3,4,5,6,7,8,9]
+
+prefix_amt: 0.0
+suffix_amt: 0.0
+prefix_dropout: 0.1
+suffix_dropout: 0.1
+
+batch_size: 120
+num_workers: 80
+
+# Optimization
+detect_anomaly: false
+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: 3
+VampNet.n_conditioning_codebooks: 0
+VampNet.r_cond_dim: 64
+VampNet.embedding_dim: 1280
+VampNet.n_layers: 16
+VampNet.n_heads: 20
+VampNet.flash_attn: false
+VampNet.dropout: 0.05
+
+AudioLoader.relative_path: /data/
+AudioDataset.loudness_cutoff: -30.0
+AudioDataset.without_replacement: true
+AudioLoader.shuffle: true
+
+train/AudioDataset.duration: 5.0
+train/AudioDataset.n_examples: 10000000
+train/AudioLoader.sources:
+  - /data/spotdl/audio/train
+
+val/AudioDataset.duration: 5.0
+val/AudioDataset.n_examples: 2000
+val/AudioLoader.sources:
+  - /data/spotdl/audio/val
+
+test/AudioDataset.duration: 5.0
+test/AudioDataset.n_examples: 1000
+test/AudioLoader.sources:
+  - /data/spotdl/audio/test

docker-compose.yml

@@ -0,0 +1,90 @@
+
+version: "3.5"
+services:
+  tensorrt:
+    build:
+      context: .
+      dockerfile: ./deployment_build/dockerfile
+      args:
+        GITHUB_TOKEN: ${GITHUB_TOKEN}
+    profiles:
+      - tensorrt
+    volumes:
+      - ./:/u/home/src
+      - ~/.config/gcloud:/root/.config/gcloud
+    deploy:
+      resources:
+        limits:
+          # match production limits
+          cpus: '7'
+          memory: 25000M
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]
+    working_dir: /u/home/src
+    entrypoint:
+      - python
+      - -m
+      - wav2wav.converter
+  base:
+    build:
+      context: .
+      dockerfile: ./Dockerfile
+      args:
+        GITHUB_TOKEN: ${GITHUB_TOKEN}
+    volumes:
+      - .:/u/home/src
+      - ~/.wav2wav:/u/home/.wav2wav
+      - ${PATH_TO_DATA}:/data
+      - ${PATH_TO_RUNS}:/runs
+      - ~/.config/gcloud:/u/home/.config/gcloud
+      - ~/.zsh_history:/u/home/.zsh_history
+    environment:
+      - GITHUB_TOKEN
+      - DISCOURSE_API_USERNAME
+      - DISCOURSE_SERVER
+      - DISCOURSE_API_KEY
+      - HOST_USER_ID
+      - HOST_USER_GID
+      - JUPYTER_TOKEN
+      - PATH_TO_DATA=/data
+      - PATH_TO_RUNS=/runs
+      - TENSORBOARD_PATH
+      - MPLCONFIGDIR=/u/home/.mplconfig
+    shm_size: 32G
+    working_dir: /u/home/src
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              capabilities: [gpu]
+  dev:
+    extends: base
+    profiles:
+      - interactive
+    stdin_open: true
+    tty: true
+  jupyter:
+    extends: base
+    ports:
+      - ${JUPYTER_PORT}:8888
+    entrypoint:
+      - /bin/bash
+      - /entry_script_jupyter.sh
+  tensorboard:
+    extends: base
+    ports:
+      - ${TENSORBOARD_PORT}:6006
+    entrypoint:
+      - /bin/bash
+      - /entry_script_tensorboard.sh
+  gradio:
+    extends: base
+    ports:
+      - 7860:7860
+    entrypoint:
+      - /bin/bash
+      - /entry_script_gradio.sh

env/alias.sh

@@ -0,0 +1,3 @@
+alias cleanup="pkill python && echo -en '\e[?25h'"
+alias stage="python ./scripts/utils/stage.py"
+alias fix_cursor="echo -en '\e[?25h'"

env/data.sh

@@ -0,0 +1,36 @@
+export PATH_TO_DATA=~/data
+
+if [[ $(hostname) == "oon17" ]]; then
+    export PATH_TO_DATA=/home/prem/shared/data/
+fi
+
+if [[ $(hostname) == "oon19" ]]; then
+    export PATH_TO_DATA=/home/prem/shared/data/
+fi
+
+if [[ $(hostname) == "lucas-ssound-trt-vm" ]]; then
+    export PATH_TO_DATA=~/data
+fi
+
+if [[ $(hostname) == "a100-ssound" ]]; then
+    export PATH_TO_DATA=~/data
+fi
+
+if [[ $(hostname) == "oon25" ]]; then
+    export PATH_TO_DATA=/data
+fi
+
+if [[ $(hostname) == "macbook-pro-2.lan" ]]; then
+    export PATH_TO_DATA=~/data
+fi
+
+if [[ $(hostname) == "oon11" ]]; then
+    export PATH_TO_DATA=/data2/syncthing_lucas/data
+fi
+
+if [[ $(hostname) == "oon12" ]]; then
+    export PATH_TO_DATA=/data
+fi
+if [[ $(hostname) == "oon26" ]]; then
+    export PATH_TO_DATA=/data
+fi

env/entry_script.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+set -e
+
+if [ -z "${USER}" ]; then
+  echo "We need USER to be set!"; exit 100
+fi
+
+# check if host uid and gid are set
+if [ -z "${HOST_USER_ID}" ]; then
+    echo "Please set HOST_USER_ID env. variables to continue." ; exit 0
+fi
+
+if [ -z "${HOST_USER_GID}" ]; then
+    echo "Please set HOST_USER_GID env. variables to continue." ; exit 0
+fi
+
+USER_ID=$HOST_USER_ID
+USER_GID=$HOST_USER_GID
+USER_HOME=/u/home
+
+# modify uid and gid to match host
+sed -i -e "s/^${USER}:\([^:]*\):[0-9]*:[0-9]*/${USER}:\1:${USER_ID}:${USER_GID}/"  /etc/passwd
+
+# create a group for host gid
+groupadd -f --gid "${USER_GID}" "host_group"
+
+chown $USER_ID $USER_HOME
+chown $USER_ID /u/home/.zshrc
+chown $USER_ID /u/home/.oh-my-zsh
+
+mkdir -p /u/home/.cache
+chown -R $USER_ID:$USER_GID /u/home/.cache/
+
+_term() {
+  echo "Caught SIGTERM signal!"
+  kill -TERM "$child" 2>/dev/null
+}
+
+trap _term SIGTERM
+
+su -p "${USER}"

env/setup.py

@@ -0,0 +1,123 @@
+# This script guides the user through setting up their env.sh
+# if env.sh does not exist. Should have no dependencies other
+# than Python standard library.
+import shlex
+import socket
+import subprocess
+import textwrap
+
+
+def run(cmd):
+    return subprocess.check_output(shlex.split(cmd)).decode("utf-8")
+
+
+print("1. Setting up Google Cloud access")
+print("---------------------------------")
+gcloud_authorized = "gs://research-data-raw" in run("gsutil ls")
+if not gcloud_authorized:
+    run("gcloud auth login")
+
+run("gcloud config set project lyrebird-research")
+run("gcloud auth configure-docker")
+
+print()
+print("2. Setting up Github access")
+print("---------------------------")
+
+lines = textwrap.wrap(
+    "First, let's get your Github token, so all "
+    "packages can be installed. Create one by going to your "
+    "Github profile -> Developer settings -> Personal access tokens -> "
+    "Generate new token. Copy the token below."
+)
+[print(l) for l in lines]
+
+GITHUB_TOKEN = input("\nGithub token: ") or "undefined"
+
+print()
+print("3. Setting up Jupyter and Tensorboard")
+print("-------------------------------------")
+
+JUPYTER_TOKEN = input("Password for Jupyter server (default:password): ") or "password"
+JUPYTER_PORT = input("Jupyter port to run on (default:8888): ") or "8888"
+TENSORBOARD_PORT = input("Tensorboard port to run on (default:6006): ") or "6006"
+
+print()
+print("4. Setting up paths.")
+print("--------------------")
+
+PATH_TO_RUNS = input("Where runs should go (default:./runs/): ") or "./runs/"
+TENSORBOARD_PATH = (
+    input("Bucket/dir for tensorboard logs (default=PATH_TO_RUNS): ") or PATH_TO_RUNS
+)
+
+with open("env/data.sh") as f:
+    data_script = f.read()
+
+write_to_data_sh = False
+if socket.gethostname() not in data_script:
+    print("Looks like the data path for this machine is not setup.")
+    PATH_TO_DATA = input(f"Path to data on {socket.gethostname()}: ") or "~/data"
+
+    data_command = f"""
+if [[ $(hostname) == "{socket.gethostname()}" ]]; then
+    export PATH_TO_DATA={PATH_TO_DATA}
+fi
+    """
+    write_to_data_sh = True
+
+
+print()
+print("5. Setting up Papaya")
+print("-----------------------------------------")
+
+PAPAYA_USER_TOKEN = input("Papaya user token: ") or "undefined"
+
+env_script = f"""
+source env/alias.sh
+source env/data.sh
+export GITHUB_TOKEN={GITHUB_TOKEN}
+
+export PAPAYA_USER_TOKEN={PAPAYA_USER_TOKEN}
+
+export HOST_USER_ID=$(id -u)
+export HOST_USER_GID=$(id -g)
+
+export JUPYTER_TOKEN={JUPYTER_TOKEN}
+export JUPYTER_PORT={JUPYTER_PORT}
+export TENSORBOARD_PORT={TENSORBOARD_PORT}
+
+export PATH_TO_RUNS={PATH_TO_RUNS}
+export TENSORBOARD_PATH={TENSORBOARD_PATH}
+"""
+
+print()
+print("6. Potential file contents.")
+print("---------------------------")
+
+print("env/env.sh: \n")
+print("##################")
+print(env_script)
+print("##################")
+
+if write_to_data_sh:
+    data_script += data_command
+
+print("env/data.sh:")
+print("##################")
+print(data_script)
+print("##################")
+
+print()
+write_to_files = input("Write to file [yn]? ") or "n"
+if write_to_files == "y":
+    with open("env/env.sh", "w") as f:
+        f.write(env_script.strip())
+    with open("env/data.sh", "w") as f:
+        f.write(data_script.strip())
+
+print()
+print("8. Finalize setup.")
+print("------------------")
+print("Run the following command to complete setup.")
+print("source env/env.sh")

requirements.txt

@@ -0,0 +1,29 @@
+argbind>=0.3.1
+pytorch-ignite
+rich
+audiotools @ git+https://github.com/descriptinc/lyrebird-audiotools.git@0.6.3
+tqdm
+tensorboard
+google-cloud-logging==2.2.0
+pytest
+pytest-cov
+papaya_client @ git+https://github.com/descriptinc/lyrebird-papaya.git@master
+pynvml
+psutil
+pandas
+onnx
+onnx-simplifier
+seaborn
+jupyterlab
+jupyterlab-link-share
+pandas
+watchdog
+pesq
+tabulate
+torchmetrics
+codebraid==0.5.0
+jupyter-client==6.1.12
+tensorboardX
+gradio
+einops
+flash-attn

scripts/generative/eval.py

@@ -0,0 +1,124 @@
+import glob
+import imp
+import os
+from pathlib import Path
+
+import argbind
+import audiotools
+import numpy as np
+import pandas as pd
+import torch
+from flatten_dict import flatten
+from rich.progress import track
+from torch.utils.tensorboard import SummaryWriter
+
+import wav2wav
+
+train = imp.load_source("train", str(Path(__file__).absolute().parent / "train.py"))
+
+
+@argbind.bind(without_prefix=True)
+def evaluate(
+    args,
+    model_tag: str = "ckpt/best",
+    device: str = "cuda",
+    exp: str = None,
+    overwrite: bool = False,
+):
+    assert exp is not None
+
+    sisdr_loss = audiotools.metrics.distance.SISDRLoss()
+    stft_loss = audiotools.metrics.spectral.MultiScaleSTFTLoss()
+    mel_loss = audiotools.metrics.spectral.MelSpectrogramLoss()
+
+    with audiotools.util.chdir(exp):
+        vampnet = wav2wav.modules.vampnet.transformer.VampNet.load(
+            f"{model_tag}/vampnet/package.pth"
+        )
+        vampnet = vampnet.to(device)
+        if vampnet.cond_dim > 0:
+            condnet = wav2wav.modules.condnet.transformer.CondNet.load(
+                f"{model_tag}/condnet/package.pth"
+            )
+            condnet = condnet.to(device)
+        else:
+            condnet = None
+
+        vqvae = wav2wav.modules.generator.Generator.load(
+            f"{model_tag}/vqvae/package.pth"
+        )
+
+    _, _, test_data = train.build_datasets(args, vqvae.sample_rate)
+
+    with audiotools.util.chdir(exp):
+        datasets = {
+            "test": test_data,
+        }
+
+        metrics_path = Path(f"{model_tag}/metrics")
+        metrics_path.mkdir(parents=True, exist_ok=True)
+
+        for key, dataset in datasets.items():
+            csv_path = metrics_path / f"{key}.csv"
+            if csv_path.exists() and not overwrite:
+                break
+            metrics = []
+            for i in track(range(len(dataset))):
+                # TODO: for coarse2fine
+                # grab the signal
+                # mask all the codebooks except the conditioning ones
+                # and infer
+                # then compute metrics
+                # for a baseline, just use the coarsest codebook
+
+                try:
+                    visqol = audiotools.metrics.quality.visqol(
+                        enhanced, clean, "audio"
+                    ).item()
+                except:
+                    visqol = None
+
+                sisdr = sisdr_loss(enhanced, clean)
+                stft = stft_loss(enhanced, clean)
+                mel = mel_loss(enhanced, clean)
+
+                metrics.append(
+                    {
+                        "visqol": visqol,
+                        "sisdr": sisdr.item(),
+                        "stft": stft.item(),
+                        "mel": mel.item(),
+                        "dataset": key,
+                        "condition": exp,
+                    }
+                )
+                print(metrics[-1])
+
+                transform_args = flatten(item["transform_args"], "dot")
+                for k, v in transform_args.items():
+                    if torch.is_tensor(v):
+                        if len(v.shape) == 0:
+                            metrics[-1][k] = v.item()
+
+            metrics = pd.DataFrame.from_dict(metrics)
+            with open(csv_path, "w") as f:
+                metrics.to_csv(f)
+
+        data = summary(model_tag).to_dict()
+        metrics = {}
+        for k1, v1 in data.items():
+            for k2, v2 in v1.items():
+                metrics[f"metrics/{k2}/{k1}"] = v2
+
+        # Number of steps to record
+        writer = SummaryWriter(log_dir=metrics_path)
+        num_steps = 10
+        for k, v in metrics.items():
+            for i in range(num_steps):
+                writer.add_scalar(k, v, i)
+
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+    with argbind.scope(args):
+        evaluate(args)

scripts/generative/train.py

@@ -0,0 +1,662 @@
+import os
+import subprocess
+import time
+import warnings
+from pathlib import Path
+from typing import Optional
+
+import argbind
+import audiotools as at
+import torch
+import torch.nn as nn
+from audiotools import AudioSignal
+from audiotools.data import transforms
+from einops import rearrange
+from rich import pretty
+from rich.traceback import install
+from tensorboardX import SummaryWriter
+
+import vampnet
+from vampnet.modules.transformer import VampNet
+from lac.model.lac import LAC
+
+
+# Enable cudnn autotuner to speed up training
+# (can be altered by the funcs.seed function)
+torch.backends.cudnn.benchmark = bool(int(os.getenv("CUDNN_BENCHMARK", 1)))
+# Uncomment to trade memory for speed.
+
+# Install to make things look nice
+warnings.filterwarnings("ignore", category=UserWarning)
+pretty.install()
+install()
+
+# optim
+Accelerator = argbind.bind(at.ml.Accelerator, without_prefix=True)
+CrossEntropyLoss = argbind.bind(nn.CrossEntropyLoss)
+AdamW = argbind.bind(torch.optim.AdamW)
+NoamScheduler = argbind.bind(vampnet.scheduler.NoamScheduler)
+
+# transforms
+filter_fn = lambda fn: hasattr(fn, "transform") and fn.__qualname__ not in [
+    "BaseTransform",
+    "Compose",
+    "Choose",
+]
+tfm = argbind.bind_module(transforms, "train", "val", filter_fn=filter_fn)
+
+# model
+VampNet = argbind.bind(VampNet)
+
+
+# data
+AudioLoader = argbind.bind(at.datasets.AudioLoader)
+AudioDataset = argbind.bind(at.datasets.AudioDataset, "train", "val")
+
+IGNORE_INDEX = -100
+
+
+@argbind.bind("train", "val", without_prefix=True)
+def build_transform():
+    transform = transforms.Compose(
+        tfm.VolumeNorm(("uniform", -32, -14)),
+        tfm.VolumeChange(("uniform", -6, 3)),
+        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())
+    with argbind.scope(args, "test"):
+        test_data = AudioDataset(
+            AudioLoader(), sample_rate, transform=build_transform()
+        )
+    return train_data, val_data, test_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
+
+
+@argbind.bind(without_prefix=True)
+def load(
+    args,
+    accel: at.ml.Accelerator,
+    save_path: str,
+    resume: bool = False,
+    tag: str = "latest",
+    load_weights: bool = False,
+):
+    model, v_extra = None, {}
+
+    if resume:
+        kwargs = {
+            "folder": f"{save_path}/{tag}",
+            "map_location": "cpu",
+            "package": not load_weights,
+        }
+        if (Path(kwargs["folder"]) / "model").exists():
+            model, v_extra = model.load_from_folder(**kwargs)
+
+    codec = LAC.load(args["codec_ckpt"], map_location="cpu")
+    codec.eval()
+    model = VampNet() if model is None else model
+
+    model = accel.prepare_model(model)
+
+    # assert accel.unwrap(model).n_codebooks == codec.quantizer.n_codebooks
+    assert (
+        accel.unwrap(model).vocab_size == codec.quantizer.quantizers[0].codebook_size
+    )
+
+    optimizer = AdamW(model.parameters(), use_zero=accel.use_ddp)
+    scheduler = NoamScheduler(optimizer, d_model=accel.unwrap(model).embedding_dim)
+    scheduler.step()
+
+    trainer_state = {"state_dict": None, "start_idx": 0}
+
+    if "optimizer.pth" in v_extra:
+        optimizer.load_state_dict(v_extra["optimizer.pth"])
+    if "scheduler.pth" in v_extra:
+        scheduler.load_state_dict(v_extra["scheduler.pth"])
+    if "trainer.pth" in v_extra:
+        trainer_state = v_extra["trainer.pth"]
+
+    return {
+        "model": model,
+        "codec": codec,
+        "optimizer": optimizer,
+        "scheduler": scheduler,
+        "trainer_state": trainer_state,
+    }
+
+
+def get_gpu_memory_map():
+    """Get the current gpu usage.
+
+    Returns
+    -------
+    usage: dict
+        Keys are device ids as integers.
+        Values are memory usage as integers in MB.
+    """
+    result = subprocess.check_output(
+        ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,nounits,noheader"],
+        encoding="utf-8",
+    )
+    # Convert lines into a dictionary
+    gpu_memory = [int(x) for x in result.strip().split("\n")]
+    gpu_memory_map = dict(zip(range(len(gpu_memory)), gpu_memory))
+    gpu_memory_map = {f"gpu/{k}": v / 1024 for k, v in gpu_memory_map.items()}
+    return gpu_memory_map
+
+
+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,
+    **kwargs,
+) -> 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
+
+
+@argbind.bind(without_prefix=True)
+def train(
+    args,
+    accel: at.ml.Accelerator,
+    codec_ckpt: str = None,
+    seed: int = 0,
+    save_path: str = "ckpt",
+    max_epochs: int = int(100e3),
+    epoch_length: int = 1000,
+    save_audio_epochs: int = 10,
+    batch_size: int = 48,
+    grad_acc_steps: int = 1,
+    val_idx: list = [0, 1, 2, 3, 4],
+    num_workers: int = 20,
+    detect_anomaly: bool = False,
+    grad_clip_val: float = 5.0,
+    prefix_amt: float = 0.0,
+    suffix_amt: float = 0.0,
+    prefix_dropout: float = 0.1,
+    suffix_dropout: float = 0.1,
+    quiet: bool = False,
+):
+    assert codec_ckpt is not None, "codec_ckpt is required"
+
+    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")
+
+    # load the codec model
+    loaded = load(args, accel, save_path)
+    model = loaded["model"]
+    codec = loaded["codec"]
+    optimizer = loaded["optimizer"]
+    scheduler = loaded["scheduler"]
+    trainer_state = loaded["trainer_state"]
+
+    sample_rate = codec.sample_rate
+
+    # a better rng for sampling from our schedule
+    rng = torch.quasirandom.SobolEngine(1, scramble=True)  
+
+    # 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)
+    train_dataloader = accel.prepare_dataloader(
+        train_data,
+        start_idx=trainer_state["start_idx"],
+        num_workers=num_workers,
+        batch_size=batch_size,
+        collate_fn=train_data.collate,
+    )
+    val_dataloader = accel.prepare_dataloader(
+        val_data,
+        start_idx=0,
+        num_workers=num_workers,
+        batch_size=batch_size,
+        collate_fn=val_data.collate,
+    )
+
+    criterion = CrossEntropyLoss()
+
+    class Trainer(at.ml.BaseTrainer):
+        _last_grad_norm = 0.0
+
+        def metrics(self, vn, 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,
+                        task="multiclass",
+                        num_classes=vn.vocab_size,
+                    )
+                    output[f"{tag}/masked"] = accuracy(
+                        preds=r_z_hat,
+                        target=r_masked_target,
+                        ignore_index=IGNORE_INDEX,
+                        top_k=topk,
+                        task="multiclass",
+                        num_classes=vn.vocab_size,
+                    )
+
+        def train_loop(self, engine, batch):
+
+            model.train()
+            batch = at.util.prepare_batch(batch, accel.device)
+            signal = apply_transform(train_data.transform, batch)
+
+            output = {}
+            vn = accel.unwrap(model)
+            with accel.autocast():
+                with torch.inference_mode():
+                    z = codec.encode(signal.samples, signal.sample_rate)["codes"]
+                    z = z[:, : vn.n_codebooks, :]
+
+                n_batch = z.shape[0]
+                r = rng.draw(n_batch)[:, 0].to(accel.device)
+
+                if prefix_amt > 0.0:
+                    prefix_mask = flip_coin(n_batch, 1 - prefix_dropout, rng)
+                    n_prefix = int(prefix_amt * z.shape[-1]) * prefix_mask
+                else:
+                    n_prefix = None
+                if suffix_amt > 0.0:
+                    suffix_mask = flip_coin(n_batch, 1 - suffix_dropout, rng)
+                    n_suffix = int(suffix_amt * z.shape[-1]) * suffix_mask
+                else:
+                    n_suffix = None
+
+                z_mask, mask = vn.add_noise(
+                    z, r, n_prefix=n_prefix, n_suffix=n_suffix
+                )
+                z_mask_latent = vn.embedding.from_codes(z_mask, codec)
+
+                dtype = torch.bfloat16 if accel.amp else None
+                with accel.autocast(dtype=dtype):
+                    z_hat = model(z_mask_latent, r)
+                    # for mask mode
+                    z_hat = vn.add_truth_to_logits(z, z_hat, mask)
+
+                target = vn.embedding.flatten(
+                    z[:, vn.n_conditioning_codebooks :, :],
+                    n_codebooks=vn.n_predict_codebooks,
+                )
+
+                flat_mask = vn.embedding.flatten(
+                    mask[:, vn.n_conditioning_codebooks :, :],
+                    n_codebooks=vn.n_predict_codebooks,
+                )
+
+                if vn.noise_mode == "mask":
+                    # replace target with ignore index for masked tokens
+                    t_masked = target.masked_fill(~flat_mask.bool(), IGNORE_INDEX)
+                    output["loss"] = criterion(z_hat, t_masked)
+                else:
+                    output["loss"] = criterion(z_hat, target)
+
+                self.metrics(
+                    vn=vn,
+                    r=r,
+                    z_hat=z_hat,
+                    target=target,
+                    flat_mask=flat_mask,
+                    output=output,
+                )
+
+            
+            accel.backward(output["loss"] / grad_acc_steps)
+
+            output["other/learning_rate"] = optimizer.param_groups[0]["lr"]
+            output["other/batch_size"] = z.shape[0]
+
+            output.update(get_gpu_memory_map())
+
+            if (
+                (engine.state.iteration % grad_acc_steps == 0)
+                or (engine.state.iteration % epoch_length == 0)
+                or (engine.state.iteration % epoch_length == 1)
+            ):  # (or we reached the end of the epoch)
+                accel.scaler.unscale_(optimizer)
+                output["other/grad_norm"] = torch.nn.utils.clip_grad_norm_(
+                    model.parameters(), grad_clip_val
+                )
+                self._last_grad_norm = output["other/grad_norm"]
+
+                accel.step(optimizer)
+                optimizer.zero_grad()
+
+                scheduler.step()
+                accel.update()
+            else:
+                output["other/grad_norm"] = self._last_grad_norm
+
+            return {k: v for k, v in sorted(output.items())}
+
+        @torch.no_grad()
+        def val_loop(self, engine, batch):
+            model.eval()
+            codec.eval()
+            batch = at.util.prepare_batch(batch, accel.device)
+            signal = apply_transform(val_data.transform, batch)
+
+            vn = accel.unwrap(model)
+            z = codec.encode(signal.samples, signal.sample_rate)["codes"]
+            z = z[:, : vn.n_codebooks, :]
+
+            n_batch = z.shape[0]
+            r = rng.draw(n_batch)[:, 0].to(accel.device)
+
+            if prefix_amt > 0.0:
+                prefix_mask = flip_coin(n_batch, 1 - prefix_dropout, rng)
+                n_prefix = int(prefix_amt * z.shape[-1]) * prefix_mask
+            else:
+                n_prefix = None
+            if suffix_amt > 0.0:
+                suffix_mask = flip_coin(n_batch, 1 - suffix_dropout, rng)
+                n_suffix = int(suffix_amt * z.shape[-1]) * suffix_mask
+            else:
+                n_suffix = None
+
+            z_mask, mask = vn.add_noise(z, r, n_prefix=n_prefix, n_suffix=n_suffix)
+            z_mask_latent = vn.embedding.from_codes(z_mask, codec)
+
+            z_hat = model(z_mask_latent, r)
+            # for mask mode
+            z_hat = vn.add_truth_to_logits(z, z_hat, mask)
+
+            target = vn.embedding.flatten(
+                z[:, vn.n_conditioning_codebooks :, :],
+                n_codebooks=vn.n_predict_codebooks,
+            )
+
+            flat_mask = vn.embedding.flatten(
+                mask[:, vn.n_conditioning_codebooks :, :],
+                n_codebooks=vn.n_predict_codebooks,
+            )
+
+            output = {}
+            if vn.noise_mode == "mask":
+                # replace target with ignore index for masked tokens
+                t_masked = target.masked_fill(~flat_mask.bool(), IGNORE_INDEX)
+                output["loss"] = criterion(z_hat, t_masked)
+            else:
+                output["loss"] = criterion(z_hat, target)
+
+            self.metrics(
+                vn=vn,
+                r=r,
+                z_hat=z_hat,
+                target=target,
+                flat_mask=flat_mask,
+                output=output,
+            )
+
+            return output
+
+        def checkpoint(self, engine):
+            if accel.local_rank != 0:
+                print(f"ERROR:Skipping checkpoint on rank {accel.local_rank}")
+                return
+
+            metadata = {"logs": dict(engine.state.logs["epoch"])}
+
+            if self.state.epoch % save_audio_epochs == 0:
+                self.save_samples()
+
+            tags = ["latest"]
+            loss_key = "loss/val" if "loss/val" in metadata["logs"] else "loss/train"
+            self.print(f"Saving to {str(Path('.').absolute())}")
+
+            if self.is_best(engine, loss_key):
+                self.print(f"Best model so far")
+                tags.append("best")
+
+            for tag in tags:
+                model_extra = {
+                    "optimizer.pth": optimizer.state_dict(),
+                    "scheduler.pth": scheduler.state_dict(),
+                    "trainer.pth": {
+                        "start_idx": self.state.iteration * batch_size,
+                        "state_dict": self.state_dict(),
+                    },
+                    "metadata.pth": metadata,
+                }
+
+                accel.unwrap(model).metadata = metadata
+                accel.unwrap(model).save_to_folder(
+                    f"{save_path}/{tag}", model_extra
+                )
+
+        def save_sampled(self, z):
+            num_samples = z.shape[0]
+
+            for i in range(num_samples):
+                sampled = accel.unwrap(model).sample(
+                    codec,
+                    time_steps=z.shape[-1],
+                    start_tokens=z[i : i + 1],
+                )
+                sampled.cpu().write_audio_to_tb(
+                    f"sampled/{i}",
+                    self.writer,
+                    step=self.state.epoch,
+                    plot_fn=None,
+                )
+
+        def save_imputation(self, z: torch.Tensor):
+            # imputations
+            mask_begin = z.shape[-1] // 4
+            mask_end = (z.shape[-1] * 3) // 4
+
+            imp_mask = torch.zeros(z.shape[0], z.shape[-1]).to(accel.device).int()
+            imp_mask[:, mask_begin:mask_end] = 1
+
+            imp_noisy = (
+                z * (1 - imp_mask[:, None, :])
+                + torch.randint_like(z, 0, accel.unwrap(model).vocab_size)
+                * imp_mask[:, None, :]
+            )
+            imputed_noisy = accel.unwrap(model).to_signal(imp_noisy, codec)
+            imputed_true = accel.unwrap(model).to_signal(z, codec)
+
+            imputed = []
+            for i in range(len(z)):
+                imputed.append(
+                    accel.unwrap(model).sample(
+                        codec,
+                        time_steps=z.shape[-1],
+                        start_tokens=z[i][None, ...],
+                        mask=imp_mask[i][None, ...],
+                    )
+                )
+            imputed = AudioSignal.batch(imputed)
+
+            for i in range(len(val_idx)):
+                imputed_noisy[i].cpu().write_audio_to_tb(
+                    f"imputed_noisy/{i}",
+                    self.writer,
+                    step=self.state.epoch,
+                    plot_fn=None,
+                )
+                imputed[i].cpu().write_audio_to_tb(
+                    f"imputed/{i}",
+                    self.writer,
+                    step=self.state.epoch,
+                    plot_fn=None,
+                )
+                imputed_true[i].cpu().write_audio_to_tb(
+                    f"imputed_true/{i}",
+                    self.writer,
+                    step=self.state.epoch,
+                    plot_fn=None,
+                )
+
+        @torch.no_grad()
+        def save_samples(self):
+            model.eval()
+            codec.eval()
+            vn = accel.unwrap(model)
+
+            batch = [val_data[i] for i in val_idx]
+            batch = at.util.prepare_batch(val_data.collate(batch), accel.device)
+
+            signal = apply_transform(val_data.transform, batch)
+
+            z = 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)
+
+            n_batch = z.shape[0]
+
+            if prefix_amt > 0.0:
+                prefix_mask = flip_coin(n_batch, 1 - prefix_dropout, rng)
+                n_prefix = int(prefix_amt * z.shape[-1]) * prefix_mask
+            else:
+                n_prefix = None
+            if suffix_amt > 0.0:
+                suffix_mask = flip_coin(n_batch, 1 - suffix_dropout, rng)
+                n_suffix = int(suffix_amt * z.shape[-1]) * suffix_mask
+            else:
+                n_suffix = None
+
+            z_mask, mask = vn.add_noise(z, r, n_prefix=n_prefix, n_suffix=n_suffix)
+            z_mask_latent = vn.embedding.from_codes(z_mask, codec)
+
+            z_hat = model(z_mask_latent, r)
+            # for mask mode
+            z_hat = vn.add_truth_to_logits(z, z_hat, mask)
+
+            z_pred = torch.softmax(z_hat, dim=1).argmax(dim=1)
+            z_pred = vn.embedding.unflatten(z_pred, n_codebooks=vn.n_predict_codebooks)
+            z_pred = torch.cat([z[:, : vn.n_conditioning_codebooks, :], z_pred], dim=1)
+
+            print("z_mask", z_mask.shape)
+            generated = vn.to_signal(z_pred, codec)
+            reconstructed = vn.to_signal(z, codec)
+            masked = vn.to_signal(z_mask.squeeze(1), 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"samples/_{i}.r={r[i]:0.2f}/{k}",
+                        self.writer,
+                        step=self.state.epoch,
+                        plot_fn=None,
+                    )
+
+            self.save_sampled(z)
+            self.save_imputation(z)
+
+    trainer = Trainer(writer=writer, quiet=quiet)
+
+    if trainer_state["state_dict"] is not None:
+        trainer.load_state_dict(trainer_state["state_dict"])
+    if hasattr(train_dataloader.sampler, "set_epoch"):
+        train_dataloader.sampler.set_epoch(trainer.trainer.state.epoch)
+
+    trainer.run(
+        train_dataloader,
+        val_dataloader,
+        num_epochs=max_epochs,
+        epoch_length=epoch_length,
+        detect_anomaly=detect_anomaly,
+    )
+
+
+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:
+            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/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()

setup.py

@@ -0,0 +1,40 @@
+from setuptools import find_packages
+from setuptools import setup
+
+with open("README.md") as f:
+    long_description = f.read()
+
+setup(
+    name="vampnet",
+    version="0.0.1",
+    classifiers=[
+        "Intended Audience :: Developers",
+        "Natural Language :: English",
+        "Programming Language :: Python :: 3.7",
+        "Topic :: Artistic Software",
+        "Topic :: Multimedia",
+        "Topic :: Multimedia :: Sound/Audio",
+        "Topic :: Multimedia :: Sound/Audio :: Editors",
+        "Topic :: Software Development :: Libraries",
+    ],
+    description="Generative Music Modeling.",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    author="Hugo Flores García",
+    author_email="hfgacrcia@descript.com",
+    url="https://github.com/descriptinc/lyrebird-vampnet",
+    license="MIT",
+    packages=find_packages(),
+    install_requires=[
+        "torch<=1.11.0",
+        "argbind>=0.3.2",
+        "pytorch-ignite",
+        "rich",
+        "audiotools @ git+https://github.com/descriptinc/lyrebird-audiotools.git@0.6.0",
+        "tqdm",
+        "tensorboard",
+        "google-cloud-logging==2.2.0",
+        "torchmetrics>=0.7.3",
+        "einops",
+    ],
+)

vampnet/__init__.py

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

vampnet/enchilada.py

@@ -0,0 +1,179 @@
+import os
+from pathlib import Path
+
+import torch
+from audiotools import AudioSignal
+
+from .modules.transformer import VampNet
+from lac.model.lac import LAC
+
+
+class TheWholeEnchilada:
+    def __init__(
+        self,
+        coarse_ckpt: str,
+        coarse2fine_ckpt: str,
+        codec_ckpt: str,
+        device: str = "cpu",
+    ):
+        self.codec = LAC.load(Path(codec_ckpt))
+        self.codec.eval()
+        self.codec.to(device)
+
+        self.coarse = VampNet.load(location=Path(coarse_ckpt), map_location="cpu")
+        self.coarse.to(device)
+        self.coarse.eval()
+
+        self.coarse2fine = VampNet.load(
+            location=Path(coarse2fine_ckpt), map_location="cpu"
+        )
+        # FIXME
+        print(
+            f"WARNING: PATCHING coarse2fine seq_len to 288, for backwards compatibility with a specific jazzpop model. it used to be {self.coarse2fine.seq_len}"
+        )
+        self.coarse2fine.seq_len = 288
+
+        self.coarse2fine.to(device)
+        self.coarse2fine.eval()
+
+        self.device = device
+
+    def seconds_to_tokens(self, seconds: float):
+        return int(seconds * self.codec.sample_rate / self.codec.hop_length)
+
+    def to(self, device):
+        self.device = device
+        self.coarse.to(device)
+        self.coarse2fine.to(device)
+        self.codec.to(device)
+        return self
+
+    def encode(self, signal: AudioSignal):
+        with torch.inference_mode():
+            # coarse z
+            cz = self.codec.encode(signal.samples, signal.sample_rate)["codes"]
+
+        return cz
+
+    def vamp(
+        self,
+        signal,
+        prefix_dur_s: float = 1.25,
+        suffix_dur_s: float = 1.25,
+        downsample_hint: bool = True,
+        downsample_factor: int = 4,
+        num_loops: int = 3,
+        **kwargs,
+    ):
+        """
+        Loop imputation of a signal.
+        """
+        signal.to(self.device).resample(self.codec.sample_rate).to_mono()
+
+        z = self.encode(signal)
+
+        cz = z[:, : self.coarse.n_codebooks, :].clone()
+        original_cz = cz.clone()
+        seq_len = original_cz.shape[-1]
+        assert (
+            seq_len == self.coarse.seq_len
+        ), f"expected seq_len {self.coarse.seq_len}, got {seq_len} for token sequence length. Is your signal the same duration as the model was trained with? "
+
+        vamp_hop_s = prefix_dur_s
+        vamp_hop = self.seconds_to_tokens(vamp_hop_s)
+
+        cmask = torch.ones_like(cz)
+
+        if downsample_hint:
+            # downsample by factor of 4
+            for i in range(cmask.shape[-1]):
+                if i % downsample_factor == 0:
+                    cmask[:, :, i] = 0
+
+        if prefix_dur_s > 0:
+            prefix_len = self.seconds_to_tokens(prefix_dur_s)
+            cmask[:, :, :prefix_len] = 0
+            print(f"prefix_len: {prefix_len}")
+        else:
+            prefix_len = 0
+
+        if suffix_dur_s > 0:
+            suffix_len = self.seconds_to_tokens(suffix_dur_s)
+            cmask[:, :, -suffix_len:] = 0
+            print(f"suffix_len: {suffix_len}")
+        else:
+            suffix_len = 0
+
+        prefix_z = cz[:, :, :prefix_len]
+
+        coarse_vamp = [prefix_z.clone()]
+        for i in range(num_loops):
+            sampled_cz = self.coarse.sample(
+                codec=self.codec,
+                time_steps=seq_len,
+                mask=cmask,
+                start_tokens=cz,
+                return_signal=False,
+                **kwargs,
+            )
+
+            new_prefix = sampled_cz[:, :, prefix_len : prefix_len + vamp_hop]
+            coarse_vamp.append(new_prefix.clone())
+
+            # replace the prefix in cz with the new prefix
+            # don't worry about a copy of the prefix still being
+            # in the mask area, since that will be masked out
+            cz[:, :, :vamp_hop] = new_prefix.clone()
+            print("to append and to prefix")
+
+        # we're done, so add the suffix
+        coarse_vamp.append(sampled_cz[:, :, prefix_len + vamp_hop :])
+
+        # concatenate the vamps
+        coarse_vamp = torch.cat(coarse_vamp, dim=-1)
+
+        # add a layer of
+        fine_prefix = z[:, self.coarse.n_codebooks :, :prefix_len]
+        fine_suffix = z[:, self.coarse.n_codebooks :, -suffix_len:]
+        fine_vamp = torch.randint(
+            0,
+            self.coarse2fine.vocab_size,
+            (
+                coarse_vamp.shape[0],
+                self.coarse2fine.n_predict_codebooks,
+                coarse_vamp.shape[-1],
+            ),
+        ).to(self.device)
+        fine_vamp[:, :, :prefix_len] = fine_prefix
+        fine_vamp[:, :, -suffix_len:] = fine_suffix
+
+        vamp_z = torch.cat([coarse_vamp, fine_vamp], dim=1)
+
+        # now we sample from the coarse2fine model
+        # to get the fine details
+        start_pos = 0
+
+        c2f_vamp = []
+        while start_pos < vamp_z.shape[-1]:
+            end_pos = min(start_pos + self.coarse2fine.seq_len, vamp_z.shape[-1])
+
+            c2fz = vamp_z[:, :, start_pos:end_pos]
+            self.coarse2fine: VampNet
+            sampled_c2fz = self.coarse2fine.sample(
+                codec=self.codec,
+                start_tokens=c2fz,
+                return_signal=False,
+                mask=None,
+            )
+            c2f_vamp.append(sampled_c2fz)
+            start_pos += self.coarse2fine.seq_len
+
+        c2f_vamp = torch.cat(c2f_vamp, dim=-1)
+
+        # make it a signal
+        vamp_signal = self.coarse2fine.to_signal(c2f_vamp, self.codec)
+
+        return {
+            "full": vamp_signal,
+            "coarse": self.coarse.to_signal(coarse_vamp, self.codec),
+        }

vampnet/modules/__init__.py

@@ -0,0 +1,4 @@
+import audiotools
+
+audiotools.ml.BaseModel.INTERN += ["vampnet.modules.**"]
+audiotools.ml.BaseModel.EXTERN += ["einops", "flash_attn.flash_attention"]

vampnet/modules/base.py

@@ -0,0 +1,461 @@
+import math
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+import audiotools as at
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from tqdm import tqdm
+
+
+def log(t, eps=1e-20):
+    return torch.log(t + eps)
+
+
+def gumbel_noise(t):
+    noise = torch.zeros_like(t).uniform_(0, 1)
+    return -log(-log(noise))
+
+
+def gumbel_sample(t, temperature=1.0, dim=-1):
+    return ((t / max(temperature, 1e-10)) + gumbel_noise(t)).argmax(dim=dim)
+
+
+class VampBase(at.ml.BaseModel):
+    def forward(self, x: torch.Tensor, r: torch.Tensor):
+        raise NotImplementedError
+
+    def add_noise(
+        self,
+        x: torch.Tensor,
+        r: torch.Tensor,
+        random_x: Optional[torch.Tensor] = None,
+        mask: Optional[torch.Tensor] = None,
+        n_prefix: Optional[torch.Tensor] = None,
+        n_suffix: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        assert x.ndim == 3, "x must be (batch, n_codebooks, seq)"
+
+        if mask is None:
+            r = self.gamma(r)[:, None, None]
+            probs = torch.ones_like(x) * r
+
+            # if we have a prefix or suffix, set their mask prob to 0
+            if n_prefix is not None:
+                for i, n in enumerate(n_prefix):
+                    probs[i, :, :n] = 0.0
+            if n_suffix is not None:
+                for i, n in enumerate(n_suffix):
+                    probs[i, :, -n:] = 0.0
+
+            mask = torch.bernoulli(probs)
+            mask = mask.round().long()
+            
+            # if we have any conditioning codebooks, set their mask  to 0
+            mask[:, : self.n_conditioning_codebooks, :] = 0
+        else:
+            assert mask.ndim == 3, "mask must be (batch, n_codebooks, seq)"
+            assert mask.shape == x.shape, "mask must be same shape as x"
+
+        if random_x is None:
+            random_x = torch.randint_like(x, 0, self.vocab_size)
+
+        if self.noise_mode == "mask":
+            random_x = torch.full_like(x, self.mask_token)
+        elif self.noise_mode == "random":
+            if random_x is None:
+                random_x = torch.randint_like(x, 0, self.vocab_size)
+        else:
+            raise ValueError(f"invalid noise mode {self.noise_mode}")
+
+        x = x * (1 - mask) + random_x * mask
+        return x, mask
+
+    def add_truth_to_logits(
+        self,
+        z_true,
+        z_hat,
+        mask,
+    ):
+        if self.noise_mode == "mask":
+            z_true = z_true[:, self.n_conditioning_codebooks :, :]
+            mask = mask[:, self.n_conditioning_codebooks :, :]
+
+            truth = F.one_hot(z_true, self.vocab_size)
+            print(truth.shape)
+            # truth = rearrange(truth, "b c t p -> b p (t c)")
+            mask = mask[:, :, :, None].expand(-1, -1, -1, self.vocab_size)
+            z_hat = rearrange(
+                z_hat,
+                "b p (t c) -> b c t p",
+                c=self.n_codebooks - self.n_conditioning_codebooks,
+            )
+
+            z_hat = z_hat * mask + truth * (1 - mask)
+
+            z_hat = rearrange(z_hat, "b c t p -> b p (t c)")
+
+        return z_hat
+
+    def gamma(self, r):
+        return (r * torch.pi / 2).cos()
+
+    def r_embed(self, r, max_positions=10000):
+        """ """
+        assert hasattr(self, "r_cond_dim"), "must set r_cond_dim before calling r_embed"
+
+        if self.r_cond_dim > 0:
+            dtype = r.dtype
+
+            r = self.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, vqvae):
+        if z.ndim == 2:
+            z = self.embedding.unflatten(z)
+        assert z.ndim == 3
+
+        signal = at.AudioSignal(
+            vqvae.decode(
+                vqvae.quantizer.from_latents(self.embedding.from_codes(z, vqvae))[0]
+            )["audio"],
+            vqvae.sample_rate,
+        )
+
+        return signal
+
+    @torch.no_grad()
+    def sample(self, **kwargs):
+        if self.noise_mode == "mask":
+            return self.maskgit_sample(**kwargs)
+        else:
+            return self.paella_sample(**kwargs)
+
+    def paella_sample(
+        self,
+        vqvae,
+        time_steps: int = 400,
+        sampling_steps: int = 12,
+        start_tokens: Optional[torch.Tensor] = None,
+        mask: Optional[torch.Tensor] = None,
+        device: str = "cpu",
+        temperature: Union[float, Tuple[float, float]] = 1.0,
+        top_k: int = None,
+        sample: str = "gumbel",
+        renoise_mode: str = "start",
+        renoise_steps=None,
+        typical_filtering=True,
+        typical_mass=0.2,
+        typical_min_tokens=1,
+        return_signal=True,
+    ):
+        r = torch.linspace(0, 1, sampling_steps + 1)[:-1][:, None].to(device)
+        if renoise_steps == None:
+            renoise_steps = sampling_steps - 1
+
+        if isinstance(temperature, float):
+            temperature = torch.tensor(temperature).repeat(sampling_steps)
+        elif isinstance(temperature, tuple):
+            assert len(temperature) == 2
+            l, h = temperature
+            temperature = torch.linspace(l, h, sampling_steps)
+        else:
+            raise TypeError(f"invalid type for temperature")
+
+        if self.n_conditioning_codebooks > 0:
+            assert (
+                start_tokens is not None
+            ), "must provide start_tokens if n_conditioning_codebooks > 0"
+
+        if start_tokens is None:
+            if self.noise_mode == "noise":
+                z = torch.randint(
+                    0, self.vocab_size, size=(1, self.n_codebooks, time_steps)
+                ).to(device)
+            elif self.noise_mode == "mask":
+                z = torch.full((1, self.n_codebooks, time_steps), self.mask_token)
+        else:
+            z = start_tokens
+            assert (
+                z.ndim == 3
+            ), f"start_tokens must be shape (batch, n_codebooks, seq_len), got {z.shape}"
+            assert z.shape[0] == 1, f"batch size must be 1"
+
+        if mask is None:
+            mask = torch.ones(z.shape[0], z.shape[-1]).to(device).int()
+
+        # apply mask
+        assert mask.shape == (
+            z.shape[0],
+            z.shape[-1],
+        ), f"mask must be shape (batch, seq_len), got {mask.shape}"
+        mask = mask[:, None, :]
+        mask = mask.repeat(1, z.shape[1], 1)
+        mask[:, : self.n_conditioning_codebooks, :] = 0.0
+
+        if self.noise_mode == "mask":
+            z_true = z.clone()
+
+        z, mask = self.add_noise(z, r=r[0], random_x=None, mask=mask)
+        z_init = z.clone()
+        for i, tmpt in enumerate(temperature):
+            if renoise_mode == "prev":
+                z_prev = z.clone()
+
+            latents = self.embedding.from_codes(z, vqvae)
+            logits = self.forward(latents, r[i])
+
+            # for mask mode
+            logits = self.add_truth_to_logits(z_true, logits, mask)
+
+            # Apply topk sampling
+            logits = logits.permute(0, 2, 1)
+
+            z = self.sample_from_logits(
+                logits,
+                tmpt,
+                top_k,
+                sample=sample,
+                typical_filtering=typical_filtering,
+                typical_mass=typical_mass,
+                typical_min_tokens=typical_min_tokens,
+            )
+
+            # add back in conditioning codebooks
+            z = self.embedding.unflatten(z, n_codebooks=self.n_predict_codebooks)
+            z = torch.cat(
+                [z_init[:, : self.n_conditioning_codebooks, :], z], dim=1
+            ).int()
+
+            if i < renoise_steps:
+                if renoise_mode == "prev":
+                    z, _ = self.add_noise(z, r[i + 1], random_x=z_prev)
+                elif renoise_mode == "start":
+                    z, _ = self.add_noise(z, r[i + 1], random_x=z_init)
+                elif renoise_mode == "rand":
+                    z, _ = self.add_noise(z, r[i + 1])
+                else:
+                    raise ValueError(f"Invalid renoise_mode: {renoise_mode}")
+
+            if mask is not None:
+                z = start_tokens * (1 - mask) + z * mask
+
+        if return_signal:
+            return self.to_signal(z, vqvae)
+        else:
+            return z
+
+    def maskgit_sample(
+        self,
+        vqvae,
+        time_steps: int = 300,
+        sampling_steps: int = 24,
+        start_tokens: Optional[torch.Tensor] = None,
+        mask: Optional[torch.Tensor] = None,
+        temperature: Union[float, Tuple[float, float]] = 1.0,
+        top_k: int = None,
+        sample: str = "multinomial",
+        typical_filtering=False,
+        typical_mass=0.2,
+        typical_min_tokens=1,
+        return_signal=True,
+    ):
+        if isinstance(temperature, float):
+            temperature = torch.tensor(temperature).repeat(sampling_steps)
+        elif isinstance(temperature, tuple):
+            assert len(temperature) == 2
+            l, h = temperature
+            temperature = torch.linspace(l, h, sampling_steps)
+        else:
+            raise TypeError(f"invalid type for temperature")
+
+        def flatten(codes):
+            return rearrange(codes, "b c t -> b (t c)")
+
+        def unflatten(codes, c):
+            return rearrange(codes, "b (t c) -> b c t", c=c)
+
+        z = start_tokens
+
+        if z is None:
+            z = torch.full((1, self.n_codebooks, time_steps), self.mask_token).to(
+                self.device
+            )
+
+        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)
+
+        # figure out which timesteps we're keeping
+        keep_mask = 1 - mask
+
+        # any conditioning codebook levels need to be in the keep mask
+        # if self.n_conditioning_codebooks > 0:
+        #     cond_mask = torch.ones(z.shape[0], self.n_conditioning_codebooks, z.shape[-1]).to(z.device)
+        #     keep_mask = torch.cat([cond_mask, keep_mask], dim=1)
+
+        # flatten
+        keep_mask = flatten(keep_mask)
+
+        # our r steps
+        r_steps = torch.linspace(0, 1, sampling_steps + 1)[1:].to(self.device)
+
+        # how many tokens did we keep on init?
+        num_kept_on_init = keep_mask.sum()
+
+        # how many codebooks are we inferring vs conditioning on?
+        n_infer_codebooks = self.n_codebooks - self.n_conditioning_codebooks
+
+        for i in tqdm(range(sampling_steps)):
+            # our current temperature
+            tmpt = temperature[i]
+
+            # our current schedule step
+            r = r_steps[i : i + 1]
+
+            with torch.inference_mode():
+                # mask our z
+                keep_mask_unflat = unflatten(keep_mask, c=self.n_codebooks)
+                z_masked = z.masked_fill(~keep_mask_unflat.bool(), self.mask_token)
+
+                # get latents
+                latents = self.embedding.from_codes(z_masked, vqvae)
+
+                # infer from latents
+                logits = self.forward(latents, r)
+                logits = logits.permute(0, 2, 1)  # b, seq, prob
+
+                # the schedule determines how many samples to keep
+                num_tokens_to_infer = (z.shape[-1] * z.shape[-2]) - num_kept_on_init
+                num_to_keep = num_kept_on_init + int(
+                    num_tokens_to_infer * (self.gamma(1 - r))
+                )
+
+                # figure out which logits we wanna keep
+                if num_to_keep > 0:
+                    probs = logits.softmax(dim=-1)
+
+                    keep_probs = F.one_hot(z, self.vocab_size)[:, :, :]
+
+                    probs = rearrange(
+                        probs, "b (t c) p -> b c t p", c=n_infer_codebooks
+                    )
+                    probs = torch.cat(
+                        [keep_probs[:, : self.n_conditioning_codebooks, ...], probs],
+                        dim=1,
+                    )
+
+                    keep_probs = rearrange(
+                        keep_probs, "b c t p -> b (t c) p", c=self.n_codebooks
+                    )
+                    probs = rearrange(probs, "b c t p -> b (t c) p", c=self.n_codebooks)
+
+                    keep_prob_mask = keep_mask.unsqueeze(-1).repeat(
+                        1, 1, self.vocab_size
+                    )
+                    probs = (keep_prob_mask.long() * keep_probs) + (
+                        1 - keep_prob_mask.long()
+                    ) * probs
+
+                    highest_probs = probs.max(dim=-1, keepdim=False)[0]
+                    v, _ = highest_probs.topk(num_to_keep, dim=-1)
+
+                    keep_mask = torch.ones_like(keep_mask).bool().clone()
+                    keep_mask[highest_probs < v[..., [-1]]] = 0
+
+                logits = torch.log(probs)
+
+                z_inferred = self.sample_from_logits(
+                    logits=logits,
+                    top_k=top_k,
+                    temperature=tmpt,
+                    sample=sample,
+                    typical_filtering=typical_filtering,
+                    typical_mass=typical_mass,
+                    typical_min_tokens=typical_min_tokens,
+                )
+
+                z = rearrange(z_inferred, "b (t c) -> b c t", c=self.n_codebooks)
+
+                # add conditioning codebooks back
+                # z = torch.cat([z[:, :self.n_conditioning_codebooks, :], z_inferred], dim=1)
+
+        if return_signal:
+            return self.to_signal(z, vqvae)
+        else:
+            return z
+
+    def sample_from_logits(
+        self,
+        logits,
+        top_k: int = None,
+        temperature: float = 1.0,
+        sample: str = "multinomial",
+        typical_filtering=False,
+        typical_mass=0.2,
+        typical_min_tokens=1,
+    ):
+        # add temperature
+        logits = logits / temperature
+
+        # add topk
+        if top_k is not None:
+            v, topk_idx = logits.topk(top_k)
+            logits[logits < v[..., [-1]]] = -float("inf")
+
+        if typical_filtering:
+            assert top_k is None
+            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)
+
+        if sample == "multinomial":
+            probs = torch.softmax(logits, dim=-1)
+            inferred = torch.stack([pr.multinomial(1).squeeze(-1) for pr in probs])
+        elif sample == "argmax":
+            inferred = torch.softmax(probs, dim=-1).argmax(dim=-1)
+        elif sample == "gumbel":
+            inferred = gumbel_sample(logits, dim=-1)
+        else:
+            raise ValueError(f"invalid sampling method: {sample}")
+
+        return inferred

vampnet/modules/modules.py

@@ -0,0 +1,168 @@
+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
+
+
+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)
+
+
+# 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 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, vqvae):
+        n_codebooks = codes.shape[1]
+        latent = []
+        for i in range(n_codebooks):
+            c = codes[:, i, :]
+
+            lookup_table = vqvae.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):
+        x = self.out_proj(latents)
+        return x
+
+    def flatten(self, tokens: torch.Tensor, n_codebooks: int = None):
+        n_c = n_codebooks if n_codebooks is not None else self.n_codebooks
+        return rearrange(tokens, "b c t -> b (t c)", c=n_c)
+
+    def unflatten(self, flat_tokens: torch.Tensor, n_codebooks: int = None):
+        nb, nt = flat_tokens.shape
+
+        n_c = n_codebooks if n_codebooks is not None else self.n_codebooks
+        tokens = rearrange(flat_tokens, "b (t c) -> b c t", c=n_c)
+
+        return tokens

vampnet/modules/transformer.py

@@ -0,0 +1,606 @@
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from .base import VampBase
+from .modules import CodebookEmbedding
+from .modules import FiLM
+from .modules import SequentialWithFiLM
+from .modules import WNConv1d
+
+
+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
+
+
+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}")
+
+
+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)
+
+
+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 = nn.Linear(d_model, d_model * 4, bias=False)
+        self.w_2 = nn.Linear(d_model * 4 // factor, d_model, bias=False)
+        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 = nn.Linear(d_model, d_model, bias=False)
+        self.w_ks = nn.Linear(d_model, d_model, bias=False)
+        self.w_vs = nn.Linear(d_model, d_model, bias=False)
+
+        # Create linear final output projection
+        self.fc = nn.Linear(d_model, d_model, bias=False)
+
+        # 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=(i == 0),
+                    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):
+        """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
+        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,
+            )
+
+        return self.norm(x) if self.norm is not None else x
+
+
+class VampNet(VampBase):
+    def __init__(
+        self,
+        n_heads: int = 20,
+        n_layers: int = 16,
+        r_cond_dim: int = 64,
+        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",
+        seq_len: int = 313,
+    ):
+        super().__init__()
+        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
+        self.seq_len = seq_len
+
+        if noise_mode == "mask":
+            special_tokens = ["MASK"]
+        elif noise_mode == "random":
+            special_tokens = None
+        else:
+            raise ValueError(f"Unknown noise mode: {noise_mode}")
+
+        self.embedding = CodebookEmbedding(
+            latent_dim=latent_dim,
+            n_codebooks=n_codebooks,
+            vocab_size=vocab_size,
+            emb_dim=embedding_dim,
+            special_tokens=special_tokens,
+        )
+
+        if noise_mode == "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=0.1,
+        )
+
+        # 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, cond):
+        x = self.embedding(x)
+        x_mask = torch.ones_like(x, dtype=torch.bool)[:, :1, :].squeeze(1)
+
+        cond = self.r_embed(cond)
+
+        x = rearrange(x, "b d n -> b n d")
+        out = self.transformer(x=x, x_mask=x_mask, cond=cond)
+        out = rearrange(out, "b n d -> b d n")
+
+        out = self.classifier(out, cond)
+
+        out = rearrange(out, "b (p c) t -> b p (t c)", c=self.n_predict_codebooks)
+
+        return out
+
+
+if __name__ == "__main__":
+    # import argbind
+    from .modules 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, mask = model.add_noise(z, r)
+
+        z_mask_latent = torch.rand(
+            batch_size, model.latent_dim * model.n_codebooks, seq_len
+        ).to(device)
+        z_hat = model(z_mask_latent, r)
+
+        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
+