Refactoring, linting, switching to pyproject.toml and Docker

.dockerignore

@@ -0,0 +1,6 @@
+.venv
+.git
+__pycache__
+*.pyc
+.ruff_cache
+wham/vampnet/models/*

Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.10-slim
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv
+
+# Install ffmpeg for audio processing
+RUN apt-get update && apt-get install -y ffmpeg git build-essential && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+# Install dependencies using uv
+COPY pyproject.toml .
+RUN uv pip install --system .
+
+# Copy your code and run the app
+COPY . .
+EXPOSE 7860
+ENV GRADIO_SERVER_NAME="0.0.0.0"
+CMD ["python", "vampnet/app.py"]

README.md

@@ -3,11 +3,9 @@ title: WhAM
 emoji: 🐋
 colorFrom: blue
 colorTo: indigo
-sdk: gradio
-app_file: vampnet/app.py
+sdk: docker
 pinned: false
 hardware: a10g-small
-python_version: "3.10"
 ---
 
 # WhAM: a Whale Acoustics Model

pyproject.toml

@@ -11,7 +11,7 @@ authors = [
     { name = "Project CETI" }
 ]
 license = { text = "MIT" }
-requires-python = ">=3.9"
+requires-python = ">=3.10,<3.11"
 dependencies = [
     "torch",
     "gradio",
@@ -33,10 +33,10 @@ dependencies = [
     "gdown",
     "transformers",
     "fadtk",
-    "urllib3==2.0",
+    "urllib3>=2.0.2",
     "plotly",
     "pyharp",
-    # Git-based dependencies
+    "ruff",
     "wavebeat @ git+https://github.com/hugofloresgarcia/wavebeat.git",
     "lac @ git+https://github.com/hugofloresgarcia/lac.git",
     "descript-audiotools @ git+https://github.com/hugofloresgarcia/audiotools.git"
@@ -47,15 +47,13 @@ where = ["."]
 include = ["wham*", "vampnet*"]
 
 [tool.ruff]
-# Target Python 3.9+
-target-version = "py39"
+target-version = "py310"
 line-length = 88
 
 [tool.ruff.lint]
-# Enable Pyflakes (F), pycodestyle (E, W), and isort (I)
 select = ["E", "F", "W", "I"]
-ignore = []
+fixable = ["ALL"]
 
-[tool.ruff.format]
-quote-style = "double"
-indent-style = "space"
+[tool.ruff.lint.isort]
+known-first-party = ["wham", "vampnet"]
+section-order = ["future", "standard-library", "third-party", "first-party", "local-folder"]

vampnet/app.py

@@ -1,43 +1,40 @@
 import os
 import sys
+import uuid
+from pathlib import Path
 
+import argbind
+import audiotools as at
+import gradio as gr
+import numpy as np
+import torch
+import yaml
+
+from vampnet import mask as pmask
+from vampnet.interface import Interface
 
 SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
 os.chdir(SCRIPT_DIR)
 
-import torch
-device = "cuda" if torch.cuda.is_available()
+device = "cuda" if torch.cuda.is_available() else "cpu"
 sys.argv = ["app.py", "--args.load", "conf/interface.yml", "--Interface.device", device]
 
-from pathlib import Path
-from typing import Tuple
-import yaml
-import tempfile
-import uuid
-from dataclasses import dataclass, asdict
-
-import numpy as np
-import audiotools as at
-import argbind
-
-import gradio as gr
-from vampnet.interface import Interface
-from vampnet import mask as pmask
 
 Interface = argbind.bind(Interface)
 
 conf = argbind.parse_args()
 
 
-from torch_pitch_shift import pitch_shift, get_fast_shifts
+from torch_pitch_shift import pitch_shift
+
+
 def shift_pitch(signal, interval: int):
     signal.samples = pitch_shift(
-        signal.samples,
-        shift=interval,
-        sample_rate=signal.sample_rate
+        signal.samples, shift=interval, sample_rate=signal.sample_rate
     )
     return signal
 
+
 def load_interface():
     with argbind.scope(conf):
         interface = Interface()
@@ -46,8 +43,6 @@ def load_interface():
         return interface
 
 
-
-
 interface = load_interface()
 
 
@@ -59,8 +54,7 @@ def load_audio(file):
     print(file)
     filepath = file.name
     sig = at.AudioSignal.salient_excerpt(
-        filepath,
-        duration=interface.coarse.chunk_size_s
+        filepath, duration=interface.coarse.chunk_size_s
     )
     sig = interface.preprocess(sig)
 
@@ -121,19 +115,10 @@ def _vamp(
     # build the mask
     mask = pmask.linear_random(z, _rand_mask_intensity)
     mask = pmask.mask_and(
-        mask, pmask.inpaint(
-            z,
-            interface.s2t(_prefix_s),
-            interface.s2t(_suffix_s)
-        )
+        mask, pmask.inpaint(z, interface.s2t(_prefix_s), interface.s2t(_suffix_s))
     )
     mask = pmask.mask_and(
-        mask, pmask.periodic_mask(
-            z,
-            _periodic_p,
-            _periodic_w,
-            random_roll=True
-        )
+        mask, pmask.periodic_mask(z, _periodic_p, _periodic_w, random_roll=True)
     )
     if _onset_mask_width > 0:
         mask = pmask.mask_or(
@@ -142,7 +127,7 @@ def _vamp(
     if _beat_mask_width > 0:
         beat_mask = interface.make_beat_mask(
             sig,
-            after_beat_s=(_beat_mask_width/1000),
+            after_beat_s=(_beat_mask_width / 1000),
             mask_upbeats=not _beat_mask_downbeats,
         )
         mask = pmask.mask_and(mask, beat_mask)
@@ -174,14 +159,14 @@ def _vamp(
 
     _top_p_val = _top_p if _top_p > 0 else None
     # save the mask as a txt file
-    np.savetxt(out_dir / "mask.txt", mask[:,0,:].long().cpu().numpy())
+    np.savetxt(out_dir / "mask.txt", mask[:, 0, :].long().cpu().numpy())
 
     _seed_val = _seed if _seed > 0 else None
     zv, mask_z = interface.coarse_vamp(
         z,
         mask=mask,
         sampling_steps=_num_steps,
-        mask_temperature=_masktemp*10,
+        mask_temperature=_masktemp * 10,
         sampling_temperature=_sampletemp,
         return_mask=True,
         typical_filtering=_typical_filtering,
@@ -196,7 +181,7 @@ def _vamp(
     if _use_coarse2fine:
         zv = interface.coarse_to_fine(
             zv,
-            mask_temperature=_masktemp*10,
+            mask_temperature=_masktemp * 10,
             sampling_temperature=_sampletemp,
             mask=mask,
             sampling_steps=_num_steps,
@@ -220,6 +205,7 @@ def _vamp(
     else:
         return sig.path_to_file
 
+
 def _extract_and_call_vamp(data, return_mask):
     """Extract plain values from Gradio data dict so only picklable args cross the ZeroGPU boundary."""
     return _vamp(
@@ -250,12 +236,15 @@ def _extract_and_call_vamp(data, return_mask):
         return_mask=return_mask,
     )
 
+
 def vamp(data):
     return _extract_and_call_vamp(data, return_mask=True)
 
+
 def api_vamp(data):
     return _extract_and_call_vamp(data, return_mask=False)
 
+
 def save_vamp(data):
     out_dir = OUT_DIR / "saved" / str(uuid.uuid4())
     out_dir.mkdir(parents=True, exist_ok=True)
@@ -289,6 +278,7 @@ def save_vamp(data):
         yaml.dump(_data, f)
 
     import zipfile
+
     zip_path = str(out_dir.with_suffix(".zip"))
     with zipfile.ZipFile(zip_path, "w") as zf:
         for file in out_dir.iterdir():
@@ -312,7 +302,7 @@ def harp_vamp(_input_audio, _beat_mask_width, _sampletemp):
     if _beat_mask_width > 0:
         beat_mask = interface.make_beat_mask(
             sig,
-            after_beat_s=(_beat_mask_width/1000),
+            after_beat_s=(_beat_mask_width / 1000),
         )
         mask = pmask.mask_and(mask, beat_mask)
 
@@ -325,7 +315,6 @@ def harp_vamp(_input_audio, _beat_mask_width, _sampletemp):
         gen_fn=interface.coarse.generate,
     )
 
-
     zv = interface.coarse_to_fine(
         zv,
         sampling_temperature=_sampletemp,
@@ -339,8 +328,8 @@ def harp_vamp(_input_audio, _beat_mask_width, _sampletemp):
 
     return sig.path_to_file
 
-with gr.Blocks() as demo:
 
+with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             gr.Markdown("# VampNet Audio Vamping")
@@ -360,11 +349,9 @@ with gr.Blocks() as demo:
             """)
     with gr.Row():
         with gr.Column():
-
-
             manual_audio_upload = gr.File(
                 label=f"upload some audio (will be randomly trimmed to max of {interface.coarse.chunk_size_s:.2f}s)",
-                file_types=["audio"]
+                file_types=["audio"],
             )
             load_example_audio_button = gr.Button("or load example audio")
 
@@ -382,71 +369,65 @@ with gr.Blocks() as demo:
 
             # connect widgets
             load_example_audio_button.click(
-                fn=load_example_audio,
-                inputs=[],
-                outputs=[ input_audio]
+                fn=load_example_audio, inputs=[], outputs=[input_audio]
             )
 
             manual_audio_upload.change(
-                fn=load_audio,
-                inputs=[manual_audio_upload],
-                outputs=[ input_audio]
+                fn=load_audio, inputs=[manual_audio_upload], outputs=[input_audio]
             )
 
         # mask settings
         with gr.Column():
-
-
             presets = {
-                    "unconditional": {
-                        "periodic_p": 0,
-                        "onset_mask_width": 0,
-                        "beat_mask_width": 0,
-                        "beat_mask_downbeats": False,
-                    },
-                    "slight periodic variation": {
-                        "periodic_p": 5,
-                        "onset_mask_width": 5,
-                        "beat_mask_width": 0,
-                        "beat_mask_downbeats": False,
-                    },
-                    "moderate periodic variation": {
-                        "periodic_p": 13,
-                        "onset_mask_width": 5,
-                        "beat_mask_width": 0,
-                        "beat_mask_downbeats": False,
-                    },
-                    "strong periodic variation": {
-                        "periodic_p": 17,
-                        "onset_mask_width": 5,
-                        "beat_mask_width": 0,
-                        "beat_mask_downbeats": False,
-                    },
-                    "very strong periodic variation": {
-                        "periodic_p": 21,
-                        "onset_mask_width": 5,
-                        "beat_mask_width": 0,
-                        "beat_mask_downbeats": False,
-                    },
-                    "beat-driven variation": {
-                        "periodic_p": 0,
-                        "onset_mask_width": 0,
-                        "beat_mask_width": 50,
-                        "beat_mask_downbeats": False,
-                    },
-                    "beat-driven variation (downbeats only)": {
-                        "periodic_p": 0,
-                        "onset_mask_width": 0,
-                        "beat_mask_width": 50,
-                        "beat_mask_downbeats": True,
-                    },
-                    "beat-driven variation (downbeats only, strong)": {
-                        "periodic_p": 0,
-                        "onset_mask_width": 0,
-                        "beat_mask_width": 20,
-                        "beat_mask_downbeats": True,
-                    },
-                }
+                "unconditional": {
+                    "periodic_p": 0,
+                    "onset_mask_width": 0,
+                    "beat_mask_width": 0,
+                    "beat_mask_downbeats": False,
+                },
+                "slight periodic variation": {
+                    "periodic_p": 5,
+                    "onset_mask_width": 5,
+                    "beat_mask_width": 0,
+                    "beat_mask_downbeats": False,
+                },
+                "moderate periodic variation": {
+                    "periodic_p": 13,
+                    "onset_mask_width": 5,
+                    "beat_mask_width": 0,
+                    "beat_mask_downbeats": False,
+                },
+                "strong periodic variation": {
+                    "periodic_p": 17,
+                    "onset_mask_width": 5,
+                    "beat_mask_width": 0,
+                    "beat_mask_downbeats": False,
+                },
+                "very strong periodic variation": {
+                    "periodic_p": 21,
+                    "onset_mask_width": 5,
+                    "beat_mask_width": 0,
+                    "beat_mask_downbeats": False,
+                },
+                "beat-driven variation": {
+                    "periodic_p": 0,
+                    "onset_mask_width": 0,
+                    "beat_mask_width": 50,
+                    "beat_mask_downbeats": False,
+                },
+                "beat-driven variation (downbeats only)": {
+                    "periodic_p": 0,
+                    "onset_mask_width": 0,
+                    "beat_mask_width": 50,
+                    "beat_mask_downbeats": True,
+                },
+                "beat-driven variation (downbeats only, strong)": {
+                    "periodic_p": 0,
+                    "onset_mask_width": 0,
+                    "beat_mask_width": 20,
+                    "beat_mask_downbeats": True,
+                },
+            }
 
             preset = gr.Dropdown(
                 label="preset",
@@ -464,7 +445,6 @@ with gr.Blocks() as demo:
                     value=3,
                 )
 
-
                 onset_mask_width = gr.Slider(
                     label="onset mask width (multiplies with the periodic mask, 1 step ~= 10milliseconds) ",
                     minimum=0,
@@ -480,8 +460,7 @@ with gr.Blocks() as demo:
                     value=0,
                 )
                 beat_mask_downbeats = gr.Checkbox(
-                    label="beat mask downbeats only?",
-                    value=False
+                    label="beat mask downbeats only?", value=False
                 )
 
                 n_mask_codebooks = gr.Number(
@@ -489,7 +468,6 @@ with gr.Blocks() as demo:
                     value=9,
                 )
 
-
                 with gr.Accordion("extras ", open=False):
                     pitch_shift_amt = gr.Slider(
                         label="pitch shift amount (semitones)",
@@ -503,7 +481,7 @@ with gr.Blocks() as demo:
                         label="random mask intensity. (If this is less than 1, scatters prompts throughout the audio, should be between 0.9 and 1.0)",
                         minimum=0.0,
                         maximum=1.0,
-                        value=1.0
+                        value=1.0,
                     )
 
                     periodic_w = gr.Slider(
@@ -538,78 +516,62 @@ with gr.Blocks() as demo:
                 return tuple(presets[_preset].values())
 
             load_preset_button.click(
-                fn=load_preset,
-                inputs=[preset],
-                outputs=preset_outputs
+                fn=load_preset, inputs=[preset], outputs=preset_outputs
             )
 
-
             with gr.Accordion("prefix/suffix prompts", open=False):
                 prefix_s = gr.Slider(
                     label="prefix hint length (seconds)",
                     minimum=0.0,
                     maximum=10.0,
-                    value=0.0
+                    value=0.0,
                 )
                 suffix_s = gr.Slider(
                     label="suffix hint length (seconds)",
                     minimum=0.0,
                     maximum=10.0,
-                    value=0.0
+                    value=0.0,
                 )
 
             masktemp = gr.Slider(
-                label="mask temperature",
-                minimum=0.0,
-                maximum=100.0,
-                value=1.5
+                label="mask temperature", minimum=0.0, maximum=100.0, value=1.5
             )
             sampletemp = gr.Slider(
                 label="sample temperature",
                 minimum=0.1,
                 maximum=10.0,
                 value=1.0,
-                step=0.001
+                step=0.001,
             )
 
-
-
             with gr.Accordion("sampling settings", open=False):
                 top_p = gr.Slider(
-                    label="top p (0.0 = off)",
-                    minimum=0.0,
-                    maximum=1.0,
-                    value=0.0
-                )
-                typical_filtering = gr.Checkbox(
-                    label="typical filtering ",
-                    value=False
+                    label="top p (0.0 = off)", minimum=0.0, maximum=1.0, value=0.0
                 )
+                typical_filtering = gr.Checkbox(label="typical filtering ", value=False)
                 typical_mass = gr.Slider(
                     label="typical mass (should probably stay between 0.1 and 0.5)",
                     minimum=0.01,
                     maximum=0.99,
-                    value=0.15
+                    value=0.15,
                 )
                 typical_min_tokens = gr.Slider(
                     label="typical min tokens (should probably stay between 1 and 256)",
                     minimum=1,
                     maximum=256,
                     step=1,
-                    value=64
+                    value=64,
                 )
                 sample_cutoff = gr.Slider(
                     label="sample cutoff",
                     minimum=0.0,
                     maximum=1.0,
                     value=0.5,
-                    step=0.01
+                    step=0.01,
                 )
 
             use_coarse2fine = gr.Checkbox(
-                label="use coarse2fine",
-                value=True,
-                visible=False
+                label="use coarse2fine", value=True, visible=False
             )
 
             num_steps = gr.Slider(
@@ -617,29 +579,21 @@ with gr.Blocks() as demo:
                 minimum=1,
                 maximum=128,
                 step=1,
-                value=36
+                value=36,
             )
 
             dropout = gr.Slider(
-                label="mask dropout",
-                minimum=0.0,
-                maximum=1.0,
-                step=0.01,
-                value=0.0
+                label="mask dropout", minimum=0.0, maximum=1.0, step=0.01, value=0.0
             )
 
-
             seed = gr.Number(
                 label="seed (0 for random)",
                 value=0,
                 precision=0,
             )
 
-
-
         # mask settings
         with gr.Column():
-
             # lora_choice = gr.Dropdown(
             #     label="lora choice",
             #     choices=list(loras.keys()),
@@ -649,51 +603,49 @@ with gr.Blocks() as demo:
 
             vamp_button = gr.Button("generate (vamp)!!!")
             output_audio = gr.Audio(
-                label="output audio",
-                interactive=False,
-                type="filepath"
+                label="output audio", interactive=False, type="filepath"
             )
 
             notes_text = gr.Textbox(
                 label="type any notes about the generated audio here",
                 value="",
-                interactive=True
+                interactive=True,
             )
             save_button = gr.Button("save vamp")
             download_file = gr.File(
-                label="vamp to download will appear here",
-                interactive=False
+                label="vamp to download will appear here", interactive=False
             )
             use_as_input_button = gr.Button("use output as input")
 
             thank_you = gr.Markdown("")
 
-
     _inputs = {
-            input_audio,
-            num_steps,
-            masktemp,
-            sampletemp,
-            top_p,
-            prefix_s, suffix_s,
-            rand_mask_intensity,
-            periodic_p, periodic_w,
-            n_conditioning_codebooks,
-            dropout,
-            use_coarse2fine,
-            stretch_factor,
-            onset_mask_width,
-            typical_filtering,
-            typical_mass,
-            typical_min_tokens,
-            beat_mask_width,
-            beat_mask_downbeats,
-            seed,
-            # lora_choice,
-            n_mask_codebooks,
-            pitch_shift_amt,
-            sample_cutoff
-        }
+        input_audio,
+        num_steps,
+        masktemp,
+        sampletemp,
+        top_p,
+        prefix_s,
+        suffix_s,
+        rand_mask_intensity,
+        periodic_p,
+        periodic_w,
+        n_conditioning_codebooks,
+        dropout,
+        use_coarse2fine,
+        stretch_factor,
+        onset_mask_width,
+        typical_filtering,
+        typical_mass,
+        typical_min_tokens,
+        beat_mask_width,
+        beat_mask_downbeats,
+        seed,
+        # lora_choice,
+        n_mask_codebooks,
+        pitch_shift_amt,
+        sample_cutoff,
+    }
 
     # connect widgets
     vamp_button.click(
@@ -704,22 +656,17 @@ with gr.Blocks() as demo:
 
     api_vamp_button = gr.Button("api vamp", visible=False)
     api_vamp_button.click(
-        fn=api_vamp,
-        inputs=_inputs,
-        outputs=[output_audio],
-        api_name="vamp"
+        fn=api_vamp, inputs=_inputs, outputs=[output_audio], api_name="vamp"
     )
 
     use_as_input_button.click(
-        fn=lambda x: x,
-        inputs=[output_audio],
-        outputs=[input_audio]
+        fn=lambda x: x, inputs=[output_audio], outputs=[input_audio]
     )
 
     save_button.click(
         fn=save_vamp,
         inputs=_inputs | {notes_text, output_audio},
-        outputs=[thank_you, download_file]
+        outputs=[thank_you, download_file],
     )
 
 

vampnet/scripts/exp/eval.py

@@ -1,20 +1,18 @@
 from pathlib import Path
-import os
-from functools import partial
 
-from frechet_audio_distance import FrechetAudioDistance
-import pandas
 import argbind
+import audiotools
+import pandas
 import torch
+from audiotools import AudioSignal
+from frechet_audio_distance import FrechetAudioDistance
 from tqdm import tqdm
 
-import audiotools
-from audiotools import AudioSignal
 
 @argbind.bind(without_prefix=True)
 def eval(
     exp_dir: str = None,
-    baseline_key: str = "baseline", 
+    baseline_key: str = "baseline",
     audio_ext: str = ".wav",
 ):
     assert exp_dir is not None
@@ -26,9 +24,9 @@ def eval(
     # stft_loss = audiotools.metrics.spectral.MultiScaleSTFTLoss()
     mel_loss = audiotools.metrics.spectral.MelSpectrogramLoss()
     frechet = FrechetAudioDistance(
-        use_pca=False, 
+        use_pca=False,
         use_activation=False,
-        verbose=True, 
+        verbose=True,
         audio_load_worker=4,
     )
     frechet.model.to("cuda" if torch.cuda.is_available() else "cpu")
@@ -36,19 +34,25 @@ def eval(
     # figure out what conditions we have
     conditions = [d.name for d in exp_dir.iterdir() if d.is_dir()]
 
-    assert baseline_key in conditions, f"baseline_key {baseline_key} not found in {exp_dir}"
+    assert baseline_key in conditions, (
+        f"baseline_key {baseline_key} not found in {exp_dir}"
+    )
     conditions.remove(baseline_key)
 
     print(f"Found {len(conditions)} conditions in {exp_dir}")
     print(f"conditions: {conditions}")
 
-    baseline_dir = exp_dir / baseline_key 
-    baseline_files = sorted(list(baseline_dir.glob(f"*{audio_ext}")), key=lambda x: int(x.stem))
+    baseline_dir = exp_dir / baseline_key
+    baseline_files = sorted(
+        list(baseline_dir.glob(f"*{audio_ext}")), key=lambda x: int(x.stem)
+    )
 
     metrics = []
     for condition in tqdm(conditions):
         cond_dir = exp_dir / condition
-        cond_files = sorted(list(cond_dir.glob(f"*{audio_ext}")), key=lambda x: int(x.stem))
+        cond_files = sorted(
+            list(cond_dir.glob(f"*{audio_ext}")), key=lambda x: int(x.stem)
+        )
 
         print(f"computing fad for {baseline_dir} and {cond_dir}")
         frechet_score = frechet.score(baseline_dir, cond_dir)
@@ -57,11 +61,15 @@ def eval(
         num_files = min(len(baseline_files), len(cond_files))
         baseline_files = baseline_files[:num_files]
         cond_files = cond_files[:num_files]
-        assert len(list(baseline_files)) == len(list(cond_files)), f"number of files in {baseline_dir} and {cond_dir} do not match. {len(list(baseline_files))} vs {len(list(cond_files))}"
+        assert len(list(baseline_files)) == len(list(cond_files)), (
+            f"number of files in {baseline_dir} and {cond_dir} do not match. {len(list(baseline_files))} vs {len(list(cond_files))}"
+        )
 
         def process(baseline_file, cond_file):
             # make sure the files match (same name)
-            assert baseline_file.stem == cond_file.stem, f"baseline file {baseline_file} and cond file {cond_file} do not match"
+            assert baseline_file.stem == cond_file.stem, (
+                f"baseline file {baseline_file} and cond file {cond_file} do not match"
+            )
 
             # load the files
             baseline_sig = AudioSignal(str(baseline_file))
@@ -74,7 +82,9 @@ def eval(
             if "inpaint" in condition:
                 ctx_amt = float(condition.split("_")[-1])
                 ctx_samples = int(ctx_amt * baseline_sig.sample_rate)
-                print(f"found inpainting condition. trimming off {ctx_samples} samples from {cond_file} and {baseline_file}")
+                print(
+                    f"found inpainting condition. trimming off {ctx_samples} samples from {cond_file} and {baseline_file}"
+                )
                 cond_sig.trim(ctx_samples, ctx_samples)
                 baseline_sig.trim(ctx_samples, ctx_samples)
 
@@ -88,15 +98,18 @@ def eval(
                 "file": baseline_file.stem,
             }
 
-        print(f"processing {len(baseline_files)} files in {baseline_dir} and {cond_dir}")
-        metrics.extend(tqdm(map(process, baseline_files, cond_files), total=len(baseline_files)))
+        print(
+            f"processing {len(baseline_files)} files in {baseline_dir} and {cond_dir}"
+        )
+        metrics.extend(
+            tqdm(map(process, baseline_files, cond_files), total=len(baseline_files))
+        )
 
     metric_keys = [k for k in metrics[0].keys() if k not in ("condition", "file")]
 
-
     for mk in metric_keys:
         stat = pandas.DataFrame(metrics)
-        stat = stat.groupby(['condition'])[mk].agg(['mean', 'count', 'std'])
+        stat = stat.groupby(["condition"])[mk].agg(["mean", "count", "std"])
         stat.to_csv(exp_dir / f"stats-{mk}.csv")
 
     df = pandas.DataFrame(metrics)
@@ -107,4 +120,4 @@ if __name__ == "__main__":
     args = argbind.parse_args()
 
     with argbind.scope(args):
-        eval()
+        eval()

vampnet/scripts/exp/experiment.py

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

vampnet/scripts/exp/fine_tune.py

@@ -1,20 +1,21 @@
-import argbind
 from pathlib import Path
-import yaml
 from typing import List
 
-
-
+import argbind
+import yaml
 
 """example output: (yaml)
 
 """
 
+
 @argbind.bind(without_prefix=True, positional=True)
 def fine_tune(audio_files_or_folders: List[str], name: str):
 
     conf_dir = Path("conf")
-    assert conf_dir.exists(), "conf directory not found. are you in the vampnet directory?"
+    assert conf_dir.exists(), (
+        "conf directory not found. are you in the vampnet directory?"
+    )
 
     conf_dir = conf_dir / "generated"
     conf_dir.mkdir(exist_ok=True)
@@ -35,7 +36,7 @@ def fine_tune(audio_files_or_folders: List[str], name: str):
         "AudioDataset.duration": 3.0,
         "AudioDataset.loudness_cutoff": -40.0,
         "save_path": f"./runs/{name}/c2f",
-        "fine_tune_checkpoint": "./models/vampnet/c2f.pth"
+        "fine_tune_checkpoint": "./models/vampnet/c2f.pth",
     }
 
     finetune_coarse_conf = {
@@ -44,15 +45,13 @@ def fine_tune(audio_files_or_folders: List[str], name: str):
         "train/AudioLoader.sources": audio_files_or_folders,
         "val/AudioLoader.sources": audio_files_or_folders,
         "save_path": f"./runs/{name}/coarse",
-        "fine_tune_checkpoint": "./models/vampnet/coarse.pth"
+        "fine_tune_checkpoint": "./models/vampnet/coarse.pth",
     }
 
     interface_conf = {
         "Interface.coarse_ckpt": f"./runs/{name}/coarse/latest/vampnet/weights.pth",
-
         "Interface.coarse2fine_ckpt": f"./runs/{name}/c2f/latest/vampnet/weights.pth",
         "Interface.wavebeat_ckpt": "./models/wavebeat.pth",
-
         "Interface.codec_ckpt": "./models/vampnet/codec.pth",
         "AudioLoader.sources": [audio_files_or_folders],
     }
@@ -63,19 +62,17 @@ def fine_tune(audio_files_or_folders: List[str], name: str):
 
     with open(finetune_dir / "coarse.yml", "w") as f:
         yaml.dump(finetune_coarse_conf, f)
-    
-    with open(finetune_dir / "interface.yml", "w") as f: 
+
+    with open(finetune_dir / "interface.yml", "w") as f:
         yaml.dump(interface_conf, f)
 
+    print(
+        f"generated confs in {finetune_dir}. run training jobs with `python scripts/exp/train.py --args.load {finetune_dir}/<c2f/coarse>.yml` "
+    )
 
-    print(f"generated confs in {finetune_dir}. run training jobs with `python scripts/exp/train.py --args.load {finetune_dir}/<c2f/coarse>.yml` ")
 
 if __name__ == "__main__":
     args = argbind.parse_args()
 
     with argbind.scope(args):
         fine_tune()
-
-
-
-    

vampnet/scripts/exp/train.py

@@ -1,36 +1,33 @@
 import os
 import sys
 import warnings
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
-from dataclasses import dataclass
 
 import argbind
 import audiotools as at
+import loralib as lora
 import torch
+import torch._dynamo
 import torch.nn as nn
 from audiotools import AudioSignal
 from audiotools.data import transforms
+from audiotools.ml.decorators import Tracker, timer, when
 from einops import rearrange
+
+# from dac.model.dac import DAC
+from lac.model.lac import LAC as DAC
 from rich import pretty
 from rich.traceback import install
 from torch.utils.tensorboard import SummaryWriter
 
 import vampnet
-from vampnet.modules.transformer import VampNet
-from vampnet.util import codebook_unflatten, codebook_flatten
 from vampnet import mask as pmask
-# from dac.model.dac import DAC
-from lac.model.lac import LAC as DAC
-
-from audiotools.ml.decorators import (
-    timer, Tracker, when
-)
-
-import loralib as lora
+from vampnet.modules.transformer import VampNet
+from vampnet.util import codebook_flatten, codebook_unflatten
 
-import torch._dynamo
-torch._dynamo.config.verbose=True
+torch._dynamo.config.verbose = True
 
 
 # Enable cudnn autotuner to speed up training
@@ -50,11 +47,15 @@ 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",
-]
+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
@@ -106,13 +107,14 @@ def flip_coin(shape, p, rng):
 
 
 def num_params_hook(o, p):
-    return o + f" {p/1e6:<.3f}M params."
+    return o + f" {p / 1e6:<.3f}M params."
 
 
 def add_num_params_repr_hook(model):
-    import numpy as np
     from functools import partial
 
+    import numpy as np
+
     for n, m in model.named_modules():
         o = m.extra_repr()
         p = sum([np.prod(p.size()) for p in m.parameters()])
@@ -149,6 +151,7 @@ def accuracy(
 
     return accuracy
 
+
 def _metrics(z_hat, r, target, flat_mask, output):
     for r_range in [(0, 0.5), (0.5, 1.0)]:
         unmasked_target = target.masked_fill(flat_mask.bool(), IGNORE_INDEX)
@@ -219,7 +222,7 @@ def train_loop(state: State, batch: dict, accel: Accelerator):
         mask = pmask.random(z, r)
         mask = pmask.codebook_unmask(mask, vn.n_conditioning_codebooks)
         z_mask, mask = pmask.apply_mask(z, mask, vn.mask_token)
-        
+
         z_mask_latent = vn.embedding.from_codes(z_mask, state.codec)
 
         dtype = torch.bfloat16 if accel.amp else None
@@ -246,13 +249,11 @@ def train_loop(state: State, batch: dict, accel: Accelerator):
             output=output,
         )
 
-    
     accel.backward(output["loss"])
 
     output["other/learning_rate"] = state.optimizer.param_groups[0]["lr"]
     output["other/batch_size"] = z.shape[0]
 
-
     accel.scaler.unscale_(state.optimizer)
     output["other/grad_norm"] = torch.nn.utils.clip_grad_norm_(
         state.model.parameters(), state.grad_clip_val
@@ -264,7 +265,6 @@ def train_loop(state: State, batch: dict, accel: Accelerator):
     state.scheduler.step()
     accel.update()
 
-
     return {k: v for k, v in sorted(output.items())}
 
 
@@ -295,9 +295,7 @@ def val_loop(state: State, batch: dict, accel: Accelerator):
         z[:, vn.n_conditioning_codebooks :, :],
     )
 
-    flat_mask = codebook_flatten(
-        mask[:, vn.n_conditioning_codebooks :, :]
-    )
+    flat_mask = codebook_flatten(mask[:, vn.n_conditioning_codebooks :, :])
 
     output = {}
     # replace target with ignore index for masked tokens
@@ -338,16 +336,16 @@ def checkpoint(state, save_iters, save_path, fine_tune):
         tags.append(f"{state.tracker.step // 1000}k")
 
     if state.tracker.is_best("val", "loss"):
-        state.tracker.print(f"Best model so far")
+        state.tracker.print("Best model so far")
         tags.append("best")
 
     if fine_tune:
-        for tag in tags: 
-            # save the lora model 
+        for tag in tags:
+            # save the lora model
             (Path(save_path) / tag).mkdir(parents=True, exist_ok=True)
             torch.save(
-                lora.lora_state_dict(accel.unwrap(state.model)), 
-                f"{save_path}/{tag}/lora.pth"
+                lora.lora_state_dict(accel.unwrap(state.model)),
+                f"{save_path}/{tag}/lora.pth",
             )
 
     for tag in tags:
@@ -383,7 +381,7 @@ def save_sampled(state, z, writer):
 
 def save_imputation(state, z, val_idx, writer):
     n_prefix = int(z.shape[-1] * 0.25)
-    n_suffix = int(z.shape[-1] *  0.25)
+    n_suffix = int(z.shape[-1] * 0.25)
 
     vn = accel.unwrap(state.model)
 
@@ -402,8 +400,8 @@ def save_imputation(state, z, val_idx, writer):
                 time_steps=z.shape[-1],
                 start_tokens=z[i][None, ...],
                 mask=mask[i][None, ...],
-            )   
-        )   
+            )
+        )
     imputed = AudioSignal.batch(imputed)
 
     for i in range(len(val_idx)):
@@ -443,7 +441,6 @@ def save_samples(state: State, val_idx: int, writer: SummaryWriter):
 
     r = torch.linspace(0.1, 0.95, len(val_idx)).to(accel.device)
 
-
     mask = pmask.random(z, r)
     mask = pmask.codebook_unmask(mask, vn.n_conditioning_codebooks)
     z_mask, mask = pmask.apply_mask(z, mask, vn.mask_token)
@@ -479,7 +476,6 @@ def save_samples(state: State, val_idx: int, writer: SummaryWriter):
     save_imputation(state=state, z=z, val_idx=val_idx, writer=writer)
 
 
-
 @argbind.bind(without_prefix=True)
 def load(
     args,
@@ -499,11 +495,12 @@ def load(
     if args["fine_tune"]:
         assert fine_tune_checkpoint is not None, "Must provide a fine-tune checkpoint"
         model = torch.compile(
-            VampNet.load(location=Path(fine_tune_checkpoint), 
-                         map_location="cpu", 
+            VampNet.load(
+                location=Path(fine_tune_checkpoint),
+                map_location="cpu",
             )
         )
-        
+
     if resume:
         kwargs = {
             "folder": f"{save_path}/{tag}",
@@ -518,16 +515,11 @@ def load(
                 f"Could not find a VampNet checkpoint in {kwargs['folder']}"
             )
 
-
-
-
     model = torch.compile(VampNet()) if model is None else model
     model = accel.prepare_model(model)
 
     # assert accel.unwrap(model).n_codebooks == codec.quantizer.n_codebooks
-    assert (
-        accel.unwrap(model).vocab_size == codec.quantizer.quantizers[0].codebook_size
-    )
+    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)
@@ -538,13 +530,13 @@ def load(
         scheduler.load_state_dict(v_extra["scheduler.pth"])
     if "tracker.pth" in v_extra:
         tracker.load_state_dict(v_extra["tracker.pth"])
-    
+
     criterion = CrossEntropyLoss()
 
     sample_rate = codec.sample_rate
 
     # a better rng for sampling from our schedule
-    rng = torch.quasirandom.SobolEngine(1, scramble=True, seed=args["seed"])  
+    rng = torch.quasirandom.SobolEngine(1, scramble=True, seed=args["seed"])
 
     # log a model summary w/ num params
     if accel.local_rank == 0:
@@ -577,13 +569,19 @@ def train(
     codec_ckpt: str = None,
     save_path: str = "ckpt",
     num_iters: int = int(1000e6),
-    save_iters: list = [10000, 50000, 100000, 300000, 500000,],
-    sample_freq: int = 10000, 
+    save_iters: list = [
+        10000,
+        50000,
+        100000,
+        300000,
+        500000,
+    ],
+    sample_freq: int = 10000,
     val_freq: int = 1000,
     batch_size: int = 12,
     val_idx: list = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
     num_workers: int = 10,
-    fine_tune: bool = False, 
+    fine_tune: bool = False,
 ):
     assert codec_ckpt is not None, "codec_ckpt is required"
 
@@ -600,11 +598,7 @@ def train(
         )
 
     # load the codec model
-    state: State = load(
-        args=args, 
-        accel=accel, 
-        tracker=tracker, 
-        save_path=save_path)
+    state: State = load(args=args, accel=accel, tracker=tracker, save_path=save_path)
     print("initialized state.")
 
     train_dataloader = accel.prepare_dataloader(
@@ -624,8 +618,6 @@ def train(
     )
     print("initialized dataloader.")
 
-    
-
     if fine_tune:
         lora.mark_only_lora_as_trainable(state.model)
         print("marked only lora as trainable.")
@@ -658,10 +650,11 @@ def train(
             if tracker.step % val_freq == 0 or last_iter:
                 validate(state, val_dataloader, accel)
                 checkpoint(
-                    state=state, 
-                    save_iters=save_iters, 
-                    save_path=save_path, 
-                    fine_tune=fine_tune)
+                    state=state,
+                    save_iters=save_iters,
+                    save_path=save_path,
+                    fine_tune=fine_tune,
+                )
 
                 # Reset validation progress bar, print summary since last validation.
                 tracker.done("val", f"Iteration {tracker.step}")

vampnet/scripts/utils/data/augment.py

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

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

@@ -1,6 +1,6 @@
-from pathlib import Path
 import json
 import os
+from pathlib import Path
 
 maestro_path = Path("/media/CHONK/hugo/maestro-v3.0.0")
 output_path = Path("/media/CHONK/hugo/maestro-v3.0.0-split")
@@ -14,7 +14,7 @@ train = []
 validation = []
 test = []
 for key, split in maestro["split"].items():
-    audio_filename = maestro['audio_filename'][key]
+    audio_filename = maestro["audio_filename"][key]
     if split == "train":
         train.append(audio_filename)
     elif split == "test":
@@ -36,4 +36,4 @@ for audio_filename in validation:
 for audio_filename in test:
     p = output_path / "test" / audio_filename
     p.parent.mkdir(parents=True, exist_ok=True)
-    os.symlink(maestro_path / audio_filename, p)
+    os.symlink(maestro_path / audio_filename, p)

vampnet/scripts/utils/plots.py

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

vampnet/scripts/utils/remove_quiet_files.py

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

vampnet/scripts/utils/split.py

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

vampnet/scripts/utils/split_long_audio_file.py

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

vampnet/scripts/utils/stage.py

@@ -1,5 +1,4 @@
 import os
-import subprocess
 from pathlib import Path
 
 import argbind

vampnet/scripts/utils/visualize_embeddings.py

@@ -3,19 +3,20 @@ TODO: train a linear probe
 usage:
    python gtzan_embeddings.py --args.load conf/interface.yml --Interface.device cuda --path_to_audio /path/to/audio/labels  --output_dir /path/to/output
 """
+
+import json
+import zipfile
 from pathlib import Path
 from typing import List
 
-import audiotools as at
-from audiotools import AudioSignal
 import argbind
-import torch
+import audiotools as at
 import numpy as np
-import zipfile
-import json
+import torch
+import tqdm
+from audiotools import AudioSignal
 
 from vampnet.interface import Interface
-import tqdm
 
 # bind the Interface to argbind
 Interface = argbind.bind(Interface)
@@ -34,6 +35,7 @@ def smart_plotly_export(fig, save_path: Path):
     # TODO: come back and make this prettier
     elif img_format == "numpy":
         import io
+
         from PIL import Image
 
         def plotly_fig2array(fig):
@@ -72,6 +74,7 @@ def dim_reduce(annotated_embeddings, layer, output_dir, n_components=3, method="
 
     if method == "umap":
         from umap import UMAP
+
         reducer = UMAP(n_components=n_components)
     elif method == "tsne":
         from sklearn.manifold import TSNE
@@ -100,11 +103,16 @@ def dim_reduce(annotated_embeddings, layer, output_dir, n_components=3, method="
     )
     if n_components == 2:
         fig = px.scatter(
-            df, x="x", y="y", color="label", hover_name="name", title=fig_title,
+            df,
+            x="x",
+            y="y",
+            color="label",
+            hover_name="name",
+            title=fig_title,
         )
 
     elif n_components == 3:
-        df['z'] = projs[:, 2]
+        df["z"] = projs[:, 2]
         fig = px.scatter_3d(
             df, x="x", y="y", z="z", color="label", hover_name="name", title=fig_title
         )
@@ -139,15 +147,15 @@ def vampnet_embed(sig: AudioSignal, interface: Interface, layer=10):
         # [20, 1, 600ish, 768]
 
         # squeeze batch dim (1 bc layer should be dim 0)
-        assert (
-            embeddings.shape[1] == 1
-        ), f"expected batch dim to be 1, got {embeddings.shape[0]}"
+        assert embeddings.shape[1] == 1, (
+            f"expected batch dim to be 1, got {embeddings.shape[0]}"
+        )
         embeddings = embeddings.squeeze(1)
 
         num_layers = embeddings.shape[0]
-        assert (
-            layer < num_layers
-        ), f"layer {layer} is out of bounds for model with {num_layers} layers"
+        assert layer < num_layers, (
+            f"layer {layer} is out of bounds for model with {num_layers} layers"
+        )
 
         # do meanpooling over the time dimension
         embeddings = embeddings.mean(dim=-2)
@@ -169,7 +177,6 @@ class AnnotatedEmbedding:
     def save(self, path):
         """Save the Embedding object to a given path as a zip file."""
         with zipfile.ZipFile(path, "w") as archive:
-
             # Save numpy array
             with archive.open("embedding.npy", "w") as f:
                 np.save(f, self.embedding)
@@ -187,7 +194,6 @@ class AnnotatedEmbedding:
     def load(cls, path):
         """Load the Embedding object from a given zip path."""
         with zipfile.ZipFile(path, "r") as archive:
-
             # Load numpy array
             with archive.open("embedding.npy") as f:
                 embedding = np.load(f)

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

@@ -1,6 +1,5 @@
 from xenopy import Query
 
-
 SPECIES = [
     "American Robin",
     "Northern Cardinal",
@@ -208,27 +207,36 @@ SPECIES = [
     "American Woodcock",
     "Wilson's Phalarope",
     "Red-necked Phalarope",
-    "Red Phalarope"
+    "Red Phalarope",
 ]
 
 from pathlib import Path
 
+
 def remove_spaces(s):
     return s.replace(" ", "")
 
-for species in SPECIES: 
+
+for species in SPECIES:
     if Path("/media/CHONK/hugo/xeno-canto-full/" + remove_spaces(species)).exists():
         continue
     try:
         q = Query(
-            name=species, q="A", length="10-30", 
-            )
+            name=species,
+            q="A",
+            length="10-30",
+        )
 
         # retrieve metadata
         metafiles = q.retrieve_meta(verbose=True)
         # retrieve recordings
-        q.retrieve_recordings(multiprocess=True, nproc=10, attempts=10, outdir="/media/CHONK/hugo/xeno-canto-full/")
+        q.retrieve_recordings(
+            multiprocess=True,
+            nproc=10,
+            attempts=10,
+            outdir="/media/CHONK/hugo/xeno-canto-full/",
+        )
 
     except:
         print("Failed to download " + species)
-        continue
+        continue

vampnet/setup.py

@@ -1,5 +1,4 @@
-from setuptools import find_packages
-from setuptools import setup
+from setuptools import find_packages, setup
 
 with open("README.md") as f:
     long_description = f.read()
@@ -29,7 +28,7 @@ setup(
         "Cython",
     ],
     install_requires=[
-        "Cython",  # Added by WAM because it seems to be needed by this repo?
+        "Cython",  # Added by WhAM because it seems to be needed by this repo?
         "torch",
         "pydantic==2.10.6",
         "argbind>=0.3.2",
@@ -40,8 +39,6 @@ setup(
         "gradio",
         "loralib",
         "torch_pitch_shift",
-        "plotly",  # Added by WAM for clustering (see https://github.com/hugofloresgarcia/vampnet/issues/20)
         "pyharp",
-        
     ],
 )

vampnet/vampnet/__init__.py

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

vampnet/vampnet/beats.py

@@ -1,19 +1,14 @@
 import json
 import logging
-import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any
-from typing import List
-from typing import Tuple
-from typing import Union
+from typing import List, Tuple, Union
 
 import librosa
-import torch
 import numpy as np
+import torch
 from audiotools import AudioSignal
 
-
 logging.basicConfig(level=logging.INFO)
 
 ###################
@@ -60,7 +55,6 @@ def mkdir(path: Union[Path, str]) -> Path:
     return p
 
 
-
 ###################
 #    beat data    #
 ###################
@@ -204,7 +198,9 @@ class WaveBeat(BeatTracker):
     def __init__(self, ckpt_path: str = "checkpoints/wavebeat", device: str = "cpu"):
         from wavebeat.dstcn import dsTCNModel
 
-        model = dsTCNModel.load_from_checkpoint(ckpt_path, map_location=torch.device(device), weights_only=False)
+        model = dsTCNModel.load_from_checkpoint(
+            ckpt_path, map_location=torch.device(device), weights_only=False
+        )
         model.eval()
 
         self.device = device
@@ -247,4 +243,4 @@ def load_beat_tracker(beat_tracker: str, **kwargs) -> BeatTracker:
             f"Unknown beat tracker {beat_tracker}. Available: {list_beat_trackers()}"
         )
 
-    return BEAT_TRACKER_REGISTRY[beat_tracker](**kwargs)
+    return BEAT_TRACKER_REGISTRY[beat_tracker](**kwargs)

vampnet/vampnet/interface.py

@@ -1,19 +1,17 @@
-import os
-from pathlib import Path
 import math
+from pathlib import Path
 
-import torch
 import numpy as np
+import torch
 from audiotools import AudioSignal
-import tqdm
-
-from .modules.transformer import VampNet
-from .beats import WaveBeat
-from .mask import *
 
 # from dac.model.dac import DAC
 from lac.model.lac import LAC as DAC
 
+from .beats import WaveBeat
+from .mask import *
+from .modules.transformer import VampNet
+
 
 def signal_concat(
     audio_signals: list,
@@ -24,7 +22,7 @@ def signal_concat(
 
 
 def _load_model(
-    ckpt: str, 
+    ckpt: str,
     lora_ckpt: str = None,
     device: str = "cpu",
     chunk_size_s: int = 10,
@@ -41,7 +39,9 @@ def _load_model(
             if should_cont != "y":
                 raise Exception("aborting")
         else:
-            model.load_state_dict(torch.load(lora_ckpt, map_location="cpu"), strict=False)
+            model.load_state_dict(
+                torch.load(lora_ckpt, map_location="cpu"), strict=False
+            )
 
     model.to(device)
     model.eval()
@@ -49,7 +49,6 @@ def _load_model(
     return model
 
 
-
 class Interface(torch.nn.Module):
     def __init__(
         self,
@@ -60,8 +59,8 @@ class Interface(torch.nn.Module):
         codec_ckpt: str = None,
         wavebeat_ckpt: str = None,
         device: str = "cpu",
-        coarse_chunk_size_s: int =  10, 
-        coarse2fine_chunk_size_s: int =  3,
+        coarse_chunk_size_s: int = 10,
+        coarse2fine_chunk_size_s: int = 3,
     ):
         super().__init__()
         assert codec_ckpt is not None, "must provide a codec checkpoint"
@@ -98,7 +97,7 @@ class Interface(torch.nn.Module):
         self.device = device
 
     def lora_load(
-        self, 
+        self,
         coarse_ckpt: str = None,
         c2f_ckpt: str = None,
         full_ckpts: bool = False,
@@ -106,7 +105,7 @@ class Interface(torch.nn.Module):
         if full_ckpts:
             if coarse_ckpt is not None:
                 self.coarse = _load_model(
-                    ckpt=coarse_ckpt,  
+                    ckpt=coarse_ckpt,
                     device=self.device,
                     chunk_size_s=self.coarse.chunk_size_s,
                 )
@@ -129,7 +128,7 @@ class Interface(torch.nn.Module):
                 print(f"loading c2f from {c2f_ckpt}")
                 self.c2f.load_state_dict(state_dict, strict=False)
                 self.c2f.to(self.device)
-        
+
     def s2t(self, seconds: float):
         """seconds to tokens"""
         if isinstance(seconds, np.ndarray):
@@ -140,7 +139,7 @@ class Interface(torch.nn.Module):
     def s2t2s(self, seconds: float):
         """seconds to tokens to seconds"""
         return self.t2s(self.s2t(seconds))
-    
+
     def t2s(self, tokens: int):
         """tokens to seconds"""
         return tokens * self.codec.hop_length / self.codec.sample_rate
@@ -159,7 +158,7 @@ class Interface(torch.nn.Module):
 
     def to_signal(self, z: torch.Tensor):
         return self.coarse.to_signal(z, self.codec)
-    
+
     def preprocess(self, signal: AudioSignal):
         signal = (
             signal.clone()
@@ -169,41 +168,39 @@ class Interface(torch.nn.Module):
             .ensure_max_of_audio(1.0)
         )
         return signal
-    
+
     @torch.inference_mode()
     def encode(self, signal: AudioSignal):
         signal = self.preprocess(signal).to(self.device)
         z = self.codec.encode(signal.samples, signal.sample_rate)["codes"]
         return z
 
-    def snap_to_beats(
-        self, 
-        signal: AudioSignal
-    ):
+    def snap_to_beats(self, signal: AudioSignal):
         assert hasattr(self, "beat_tracker"), "No beat tracker loaded"
         beats, downbeats = self.beat_tracker.extract_beats(signal)
-        
+
         # trim the signa around the first beat time
-        samples_begin = int(beats[0] * signal.sample_rate )
+        samples_begin = int(beats[0] * signal.sample_rate)
         samples_end = int(beats[-1] * signal.sample_rate)
         print(beats[0])
         signal = signal.clone().trim(samples_begin, signal.length - samples_end)
 
         return signal
 
-    def make_beat_mask(self, 
-            signal: AudioSignal, 
-            before_beat_s: float = 0.0,
-            after_beat_s: float = 0.02,
-            mask_downbeats: bool = True,
-            mask_upbeats: bool = True,
-            downbeat_downsample_factor: int = None,
-            beat_downsample_factor: int = None,
-            dropout: float = 0.0,
-            invert: bool = True,
+    def make_beat_mask(
+        self,
+        signal: AudioSignal,
+        before_beat_s: float = 0.0,
+        after_beat_s: float = 0.02,
+        mask_downbeats: bool = True,
+        mask_upbeats: bool = True,
+        downbeat_downsample_factor: int = None,
+        beat_downsample_factor: int = None,
+        dropout: float = 0.0,
+        invert: bool = True,
     ):
-        """make a beat synced mask. that is, make a mask that 
-        places 1s at and around the beat, and 0s everywhere else. 
+        """make a beat synced mask. that is, make a mask that
+        places 1s at and around the beat, and 0s everywhere else.
         """
         assert self.beat_tracker is not None, "No beat tracker loaded"
 
@@ -214,14 +211,16 @@ class Interface(torch.nn.Module):
         beats_z, downbeats_z = self.s2t(beats), self.s2t(downbeats)
 
         # remove downbeats from beats
-        beats_z = torch.tensor(beats_z)[~torch.isin(torch.tensor(beats_z), torch.tensor(downbeats_z))]
+        beats_z = torch.tensor(beats_z)[
+            ~torch.isin(torch.tensor(beats_z), torch.tensor(downbeats_z))
+        ]
         beats_z = beats_z.tolist()
         downbeats_z = downbeats_z.tolist()
 
-        # make the mask 
+        # make the mask
         seq_len = self.s2t(signal.duration)
         mask = torch.zeros(seq_len, device=self.device)
-        
+
         mask_b4 = self.s2t(before_beat_s)
         mask_after = self.s2t(after_beat_s)
 
@@ -241,44 +240,39 @@ class Interface(torch.nn.Module):
         downbeats_z = downbeats_z[::downbeat_downsample_factor]
         print(f"beats_z: {len(beats_z)}")
         print(f"downbeats_z: {len(downbeats_z)}")
-    
+
         if mask_upbeats:
             for beat_idx in beats_z:
                 _slice = int(beat_idx - mask_b4), int(beat_idx + mask_after)
-                num_steps = mask[_slice[0]:_slice[1]].shape[0]
+                num_steps = mask[_slice[0] : _slice[1]].shape[0]
                 _m = torch.ones(num_steps, device=self.device)
                 _m_mask = torch.bernoulli(_m * (1 - dropout))
                 _m = _m * _m_mask.long()
-                
-                mask[_slice[0]:_slice[1]] = _m
+
+                mask[_slice[0] : _slice[1]] = _m
 
         if mask_downbeats:
             for downbeat_idx in downbeats_z:
                 _slice = int(downbeat_idx - mask_b4), int(downbeat_idx + mask_after)
-                num_steps = mask[_slice[0]:_slice[1]].shape[0]
+                num_steps = mask[_slice[0] : _slice[1]].shape[0]
                 _m = torch.ones(num_steps, device=self.device)
                 _m_mask = torch.bernoulli(_m * (1 - dropout))
                 _m = _m * _m_mask.long()
-                
-                mask[_slice[0]:_slice[1]] = _m
-        
+
+                mask[_slice[0] : _slice[1]] = _m
+
         mask = mask.clamp(0, 1)
         if invert:
             mask = 1 - mask
-        
+
         mask = mask[None, None, :].bool().long()
         if self.c2f is not None:
             mask = mask.repeat(1, self.c2f.n_codebooks, 1)
         else:
             mask = mask.repeat(1, self.coarse.n_codebooks, 1)
         return mask
-        
-    def coarse_to_fine(
-        self, 
-        z: torch.Tensor,
-        mask: torch.Tensor = None,
-        **kwargs
-    ):
+
+    def coarse_to_fine(self, z: torch.Tensor, mask: torch.Tensor = None, **kwargs):
         assert self.c2f is not None, "No coarse2fine model loaded"
         length = z.shape[-1]
         chunk_len = self.s2t(self.c2f.chunk_size_s)
@@ -288,49 +282,57 @@ class Interface(torch.nn.Module):
         if length % chunk_len != 0:
             pad_len = chunk_len - (length % chunk_len)
             z = torch.nn.functional.pad(z, (0, pad_len))
-            mask = torch.nn.functional.pad(mask, (0, pad_len)) if mask is not None else None
+            mask = (
+                torch.nn.functional.pad(mask, (0, pad_len))
+                if mask is not None
+                else None
+            )
 
         n_codebooks_to_append = self.c2f.n_codebooks - z.shape[1]
         if n_codebooks_to_append > 0:
-            z = torch.cat([
-                z,
-                torch.zeros(z.shape[0], n_codebooks_to_append, z.shape[-1]).long().to(self.device)
-            ], dim=1)
+            z = torch.cat(
+                [
+                    z,
+                    torch.zeros(z.shape[0], n_codebooks_to_append, z.shape[-1])
+                    .long()
+                    .to(self.device),
+                ],
+                dim=1,
+            )
 
         # set the mask to 0 for all conditioning codebooks
         if mask is not None:
             mask = mask.clone()
-            mask[:, :self.c2f.n_conditioning_codebooks, :] = 0
+            mask[:, : self.c2f.n_conditioning_codebooks, :] = 0
 
         fine_z = []
         for i in range(n_chunks):
             chunk = z[:, :, i * chunk_len : (i + 1) * chunk_len]
-            mask_chunk = mask[:, :, i * chunk_len : (i + 1) * chunk_len] if mask is not None else None
-            
+            mask_chunk = (
+                mask[:, :, i * chunk_len : (i + 1) * chunk_len]
+                if mask is not None
+                else None
+            )
+
             chunk = self.c2f.generate(
                 codec=self.codec,
                 time_steps=chunk_len,
                 start_tokens=chunk,
                 return_signal=False,
                 mask=mask_chunk,
-                **kwargs
+                **kwargs,
             )
             fine_z.append(chunk)
 
         fine_z = torch.cat(fine_z, dim=-1)
         return fine_z[:, :, :length].clone()
-    
-    def coarse_vamp(
-        self, 
-        z, 
-        mask,
-        return_mask=False,
-        gen_fn=None,
-        **kwargs
-    ):
+
+    def coarse_vamp(self, z, mask, return_mask=False, gen_fn=None, **kwargs):
         # coarse z
         cz = z[:, : self.coarse.n_codebooks, :].clone()
-        assert cz.shape[-1] <= self.s2t(self.coarse.chunk_size_s), f"the sequence of tokens provided must match the one specified in the coarse chunk size, but got {cz.shape[-1]} and {self.s2t(self.coarse.chunk_size_s)}"
+        assert cz.shape[-1] <= self.s2t(self.coarse.chunk_size_s), (
+            f"the sequence of tokens provided must match the one specified in the coarse chunk size, but got {cz.shape[-1]} and {self.s2t(self.coarse.chunk_size_s)}"
+        )
 
         mask = mask[:, : self.coarse.n_codebooks, :]
 
@@ -342,41 +344,39 @@ class Interface(torch.nn.Module):
             codec=self.codec,
             time_steps=cz.shape[-1],
             start_tokens=cz,
-            mask=mask, 
+            mask=mask,
             return_signal=False,
-            **kwargs
+            **kwargs,
         )
 
         # add the fine codes back in
-        c_vamp = torch.cat(
-            [c_vamp, z[:, self.coarse.n_codebooks :, :]], 
-            dim=1
-        )
+        c_vamp = torch.cat([c_vamp, z[:, self.coarse.n_codebooks :, :]], dim=1)
 
         if return_mask:
             return c_vamp, cz_masked
-        
+
         return c_vamp
 
 
 if __name__ == "__main__":
-    import audiotools as at
     import logging
+
+    import audiotools as at
+
     logger = logging.getLogger()
     logger.setLevel(logging.INFO)
     torch.set_printoptions(threshold=10000)
     at.util.seed(42)
 
     interface = Interface(
-        coarse_ckpt="./models/vampnet/coarse.pth", 
-        coarse2fine_ckpt="./models/vampnet/c2f.pth", 
+        coarse_ckpt="./models/vampnet/coarse.pth",
+        coarse2fine_ckpt="./models/vampnet/c2f.pth",
         codec_ckpt="./models/vampnet/codec.pth",
-        device="cuda", 
-        wavebeat_ckpt="./models/wavebeat.pth"
+        device="cuda",
+        wavebeat_ckpt="./models/wavebeat.pth",
     )
 
-
-    sig = at.AudioSignal('assets/example.wav')
+    sig = at.AudioSignal("assets/example.wav")
 
     z = interface.encode(sig)
     breakpoint()
@@ -398,19 +398,18 @@ if __name__ == "__main__":
     # mask = codebook_unmask(mask, 0)
 
     mask = inpaint(z, n_prefix=100, n_suffix=100)
-    
+
     zv, mask_z = interface.coarse_vamp(
-        z, 
+        z,
         mask=mask,
         sampling_steps=36,
         temperature=8.0,
-        return_mask=True, 
-        gen_fn=interface.coarse.generate
+        return_mask=True,
+        gen_fn=interface.coarse.generate,
     )
-    
 
     use_coarse2fine = True
-    if use_coarse2fine: 
+    if use_coarse2fine:
         zv = interface.coarse_to_fine(zv, temperature=0.8, mask=mask)
         breakpoint()
 
@@ -418,5 +417,3 @@ if __name__ == "__main__":
 
     sig = interface.to_signal(zv).cpu()
     print("done")
-
-        

vampnet/vampnet/mask.py

@@ -1,33 +1,34 @@
-from typing import Optional
-
 import torch
 from audiotools import AudioSignal
 
 from .util import scalar_to_batch_tensor
 
+
 def _gamma(r):
     return (r * torch.pi / 2).cos().clamp(1e-10, 1.0)
 
+
 def _invgamma(y):
     if not torch.is_tensor(y):
         y = torch.tensor(y)[None]
     return 2 * y.acos() / torch.pi
 
+
 def full_mask(x: torch.Tensor):
     assert x.ndim == 3, "x must be (batch, n_codebooks, seq)"
     return torch.ones_like(x).long()
 
+
 def empty_mask(x: torch.Tensor):
     assert x.ndim == 3, "x must be (batch, n_codebooks, seq)"
     return torch.zeros_like(x).long()
 
-def apply_mask(
-        x: torch.Tensor, 
-        mask: torch.Tensor, 
-        mask_token: int
-    ):
+
+def apply_mask(x: torch.Tensor, mask: torch.Tensor, mask_token: int):
     assert mask.ndim == 3, "mask must be (batch, n_codebooks, seq), but got {mask.ndim}"
-    assert mask.shape == x.shape, f"mask must be same shape as x, but got {mask.shape} and {x.shape}" 
+    assert mask.shape == x.shape, (
+        f"mask must be same shape as x, but got {mask.shape} and {x.shape}"
+    )
     assert mask.dtype == torch.long, "mask must be long dtype, but got {mask.dtype}"
     assert ~torch.any(mask > 1), "mask must be binary"
     assert ~torch.any(mask < 0), "mask must be binary"
@@ -37,10 +38,8 @@ def apply_mask(
 
     return x, mask
 
-def random(
-    x: torch.Tensor,
-    r: torch.Tensor
-):
+
+def random(x: torch.Tensor, r: torch.Tensor):
     assert x.ndim == 3, "x must be (batch, n_codebooks, seq)"
     if not isinstance(r, torch.Tensor):
         r = scalar_to_batch_tensor(r, x.shape[0]).to(x.device)
@@ -53,6 +52,7 @@ def random(
 
     return mask
 
+
 def linear_random(
     x: torch.Tensor,
     r: torch.Tensor,
@@ -71,19 +71,21 @@ def linear_random(
 
     return mask
 
-def inpaint(x: torch.Tensor, 
+
+def inpaint(
+    x: torch.Tensor,
     n_prefix,
     n_suffix,
 ):
     assert n_prefix is not None
     assert n_suffix is not None
-    
+
     mask = full_mask(x)
 
     # if we have a prefix or suffix, set their mask prob to 0
     if n_prefix > 0:
         if not isinstance(n_prefix, torch.Tensor):
-            n_prefix = scalar_to_batch_tensor(n_prefix, x.shape[0]).to(x.device) 
+            n_prefix = scalar_to_batch_tensor(n_prefix, x.shape[0]).to(x.device)
         for i, n in enumerate(n_prefix):
             if n > 0:
                 mask[i, :, :n] = 0.0
@@ -94,13 +96,15 @@ def inpaint(x: torch.Tensor,
             if n > 0:
                 mask[i, :, -n:] = 0.0
 
-    
     return mask
 
-def periodic_mask(x: torch.Tensor, 
-                period: int, width: int = 1, 
-                random_roll=False,
-    ):
+
+def periodic_mask(
+    x: torch.Tensor,
+    period: int,
+    width: int = 1,
+    random_roll=False,
+):
     mask = full_mask(x)
     if period == 0:
         return mask
@@ -113,8 +117,8 @@ def periodic_mask(x: torch.Tensor,
         for j in range(mask.shape[-1]):
             if j % factor == 0:
                 # figure out how wide the mask should be
-                j_start = max(0, j - width // 2  )
-                j_end = min(mask.shape[-1] - 1, j + width // 2 ) + 1 
+                j_start = max(0, j - width // 2)
+                j_end = min(mask.shape[-1] - 1, j + width // 2) + 1
                 # flip a coin for each position in the mask
                 j_mask = torch.bernoulli(torch.ones(j_end - j_start))
                 assert torch.all(j_mask == 1)
@@ -129,10 +133,8 @@ def periodic_mask(x: torch.Tensor,
 
     return mask
 
-def codebook_unmask(
-    mask: torch.Tensor, 
-    n_conditioning_codebooks: int
-):
+
+def codebook_unmask(mask: torch.Tensor, n_conditioning_codebooks: int):
     if n_conditioning_codebooks == None:
         return mask
     # if we have any conditioning codebooks, set their mask  to 0
@@ -140,18 +142,18 @@ def codebook_unmask(
     mask[:, :n_conditioning_codebooks, :] = 0
     return mask
 
+
 def codebook_mask(mask: torch.Tensor, start: int):
     mask = mask.clone()
     mask[:, start:, :] = 1
     return mask
 
-def mask_and(
-    mask1: torch.Tensor, 
-    mask2: torch.Tensor
-):
+
+def mask_and(mask1: torch.Tensor, mask2: torch.Tensor):
     assert mask1.shape == mask2.shape, "masks must be same shape"
     return torch.min(mask1, mask2)
 
+
 def dropout(
     mask: torch.Tensor,
     p: float,
@@ -164,19 +166,20 @@ def dropout(
     mask = ~mask.round().bool()
     return mask.long()
 
-def mask_or(
-    mask1: torch.Tensor, 
-    mask2: torch.Tensor
-):
-    assert mask1.shape == mask2.shape, f"masks must be same shape, but got {mask1.shape} and {mask2.shape}"
+
+def mask_or(mask1: torch.Tensor, mask2: torch.Tensor):
+    assert mask1.shape == mask2.shape, (
+        f"masks must be same shape, but got {mask1.shape} and {mask2.shape}"
+    )
     assert mask1.max() <= 1, "mask1 must be binary"
     assert mask2.max() <= 1, "mask2 must be binary"
     assert mask1.min() >= 0, "mask1 must be binary"
     assert mask2.min() >= 0, "mask2 must be binary"
     return (mask1 + mask2).clamp(0, 1)
 
+
 def time_stretch_mask(
-    x: torch.Tensor, 
+    x: torch.Tensor,
     stretch_factor: int,
 ):
     assert stretch_factor >= 1, "stretch factor must be >= 1"
@@ -189,18 +192,19 @@ def time_stretch_mask(
     mask = periodic_mask(x, stretch_factor, width=1)
     return mask
 
+
 def _onset_times_madmom(wav_path, sample_rate, hop_length):
-    from madmom.features.onsets import RNNOnsetProcessor, OnsetPeakPickingProcessor
+    from madmom.features.onsets import OnsetPeakPickingProcessor, RNNOnsetProcessor
+
     proc = RNNOnsetProcessor(online=False)
-    onsetproc = OnsetPeakPickingProcessor(
-        threshold=0.3, fps=sample_rate / hop_length
-    )
+    onsetproc = OnsetPeakPickingProcessor(threshold=0.3, fps=sample_rate / hop_length)
     act = proc(wav_path)
     return onsetproc(act)
 
 
 def _onset_times_librosa(wav_path, sample_rate, hop_length):
     import librosa
+
     y, sr = librosa.load(wav_path, sr=sample_rate)
     onset_env = librosa.onset.onset_strength(y=y, sr=sr, hop_length=hop_length)
     onset_frames = librosa.onset.onset_detect(
@@ -209,18 +213,14 @@ def _onset_times_librosa(wav_path, sample_rate, hop_length):
     return librosa.frames_to_time(onset_frames, sr=sr, hop_length=hop_length)
 
 
-def onset_mask(
-    sig: AudioSignal, 
-    z: torch.Tensor,
-    interface,
-    width: int = 1
-):
-    import librosa
+def onset_mask(sig: AudioSignal, z: torch.Tensor, interface, width: int = 1):
     import tempfile
-    import numpy as np 
+
+    import librosa
 
     try:
         import madmom  # noqa: F401
+
         _get_onset_times = _onset_times_madmom
     except ImportError:
         print("madmom not installed, falling back to librosa for onset detection")
@@ -228,7 +228,7 @@ def onset_mask(
 
     hop_length = interface.codec.hop_length
 
-    with tempfile.NamedTemporaryFile(suffix='.wav') as f:
+    with tempfile.NamedTemporaryFile(suffix=".wav") as f:
         sig = sig.clone()
         sig.write(f.name)
 
@@ -238,9 +238,9 @@ def onset_mask(
         )
 
         if onset_indices.shape[0] == 0:
-            mask = empty_mask(z)   
-            print(f"no onsets found, returning empty mask")
-        else: 
+            mask = empty_mask(z)
+            print("no onsets found, returning empty mask")
+        else:
             torch.set_printoptions(threshold=1000)
             print("onset indices: ", onset_indices)
             print("onset times: ", onset_times)
@@ -251,12 +251,11 @@ def onset_mask(
             for onset_index in onset_indices:
                 onset_index = min(onset_index, n_timesteps - 1)
                 onset_index = max(onset_index, 0)
-                mask[:, :, onset_index - width:onset_index + width] = 0.0
+                mask[:, :, onset_index - width : onset_index + width] = 0.0
 
             print(mask)
-    
-    return mask
 
+    return mask
 
 
 if __name__ == "__main__":

vampnet/vampnet/modules/__init__.py

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

vampnet/vampnet/modules/activations.py

@@ -1,9 +1,7 @@
 import math
-import numpy as np
+
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
 
 
 class NewGELU(nn.Module):
@@ -25,6 +23,7 @@ class NewGELU(nn.Module):
             )
         )
 
+
 class GatedGELU(nn.Module):
     def __init__(self):
         super().__init__()
@@ -34,6 +33,7 @@ class GatedGELU(nn.Module):
         p1, p2 = x.chunk(2, dim=dim)
         return p1 * self.gelu(p2)
 
+
 class Snake1d(nn.Module):
     def __init__(self, channels):
         super().__init__()
@@ -42,6 +42,7 @@ class Snake1d(nn.Module):
     def forward(self, x):
         return x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x).pow(2)
 
+
 def get_activation(name: str = "relu"):
     if name == "relu":
         return nn.ReLU
@@ -52,4 +53,4 @@ def get_activation(name: str = "relu"):
     elif name == "snake":
         return Snake1d
     else:
-        raise ValueError(f"Unrecognized activation {name}")
+        raise ValueError(f"Unrecognized activation {name}")

vampnet/vampnet/modules/layers.py

@@ -1,13 +1,11 @@
-import time
-from typing import Optional
-from typing import Tuple
+from typing import Optional, Tuple
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from einops import rearrange
 from torch.nn.utils import weight_norm
 
+
 # Scripting this brings model speed up 1.4x
 @torch.jit.script
 def snake(x, alpha):
@@ -132,10 +130,10 @@ class CodebookEmbedding(nn.Module):
         self.out_proj = nn.Conv1d(n_codebooks * self.latent_dim, self.emb_dim, 1)
 
     def from_codes(self, codes: torch.Tensor, codec):
-        """ 
-        get a sequence of continuous embeddings from a sequence of discrete codes. 
+        """
+        get a sequence of continuous embeddings from a sequence of discrete codes.
         unlike it's counterpart in the original VQ-VAE, this function adds for any special tokens
-        necessary for the language model, like <MASK>. 
+        necessary for the language model, like <MASK>.
         """
         n_codebooks = codes.shape[1]
         latent = []
@@ -161,4 +159,3 @@ class CodebookEmbedding(nn.Module):
         """
         x = self.out_proj(latents)
         return x
-

vampnet/vampnet/modules/transformer.py

@@ -1,22 +1,19 @@
-import math
 import logging
-from typing import Optional, Tuple, Union
+import math
+from typing import Optional
 
+import audiotools as at
+import loralib as lora
 import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
-import loralib as lora
-import audiotools as at
 
-from .activations import get_activation
-from .layers import CodebookEmbedding
-from .layers import FiLM
-from .layers import SequentialWithFiLM
-from .layers import WNConv1d
-from ..util import scalar_to_batch_tensor, codebook_flatten, codebook_unflatten
 from ..mask import _gamma
+from ..util import codebook_flatten, codebook_unflatten, scalar_to_batch_tensor
+from .activations import get_activation
+from .layers import CodebookEmbedding, FiLM, SequentialWithFiLM, WNConv1d
 
 LORA_R = 8
 
@@ -279,6 +276,7 @@ class TransformerLayer(nn.Module):
 
         if flash_attn:
             from flash_attn.flash_attention import FlashMHA
+
             self.self_attn = FlashMHA(
                 embed_dim=d_model,
                 num_heads=n_heads,
@@ -410,9 +408,15 @@ class TransformerStack(nn.Module):
     def subsequent_mask(self, size):
         return torch.ones(1, size, size).tril().bool()
 
-    def forward(self, x, x_mask, cond=None, src=None, src_mask=None,
-                return_activations: bool = False
-        ):
+    def forward(
+        self,
+        x,
+        x_mask,
+        cond=None,
+        src=None,
+        src_mask=None,
+        return_activations: bool = False,
+    ):
         """Computes a full transformer stack
         Parameters
         ----------
@@ -454,7 +458,6 @@ class TransformerStack(nn.Module):
             if return_activations:
                 activations.append(x.detach())
 
-    
         out = self.norm(x) if self.norm is not None else x
         if return_activations:
             return out, torch.stack(activations)
@@ -475,10 +478,12 @@ class VampNet(at.ml.BaseModel):
         vocab_size: int = 1024,
         flash_attn: bool = True,
         noise_mode: str = "mask",
-        dropout: float = 0.1
+        dropout: float = 0.1,
     ):
         super().__init__()
-        assert r_cond_dim == 0, f"r_cond_dim must be 0 (not supported), but got {r_cond_dim}"
+        assert r_cond_dim == 0, (
+            f"r_cond_dim must be 0 (not supported), but got {r_cond_dim}"
+        )
         self.n_heads = n_heads
         self.n_layers = n_layers
         self.r_cond_dim = r_cond_dim
@@ -530,13 +535,15 @@ class VampNet(at.ml.BaseModel):
         x_mask = torch.ones_like(x, dtype=torch.bool)[:, :1, :].squeeze(1)
 
         x = rearrange(x, "b d n -> b n d")
-        out = self.transformer(x=x, x_mask=x_mask, return_activations=return_activations)
+        out = self.transformer(
+            x=x, x_mask=x_mask, return_activations=return_activations
+        )
         if return_activations:
             out, activations = out
 
         out = rearrange(out, "b n d -> b d n")
 
-        out = self.classifier(out, None) # no cond here!
+        out = self.classifier(out, None)  # no cond here!
 
         out = rearrange(out, "b (p c) t -> b p (t c)", c=self.n_predict_codebooks)
 
@@ -544,7 +551,7 @@ class VampNet(at.ml.BaseModel):
             return out, activations
         else:
             return out
-    
+
     def r_embed(self, r, max_positions=10000):
         if self.r_cond_dim > 0:
             dtype = r.dtype
@@ -564,11 +571,11 @@ class VampNet(at.ml.BaseModel):
             return emb.to(dtype)
         else:
             return r
-    
+
     @torch.no_grad()
     def to_signal(self, z, codec):
         """
-        convert a sequence of latents to a signal. 
+        convert a sequence of latents to a signal.
         """
         assert z.ndim == 3
 
@@ -588,7 +595,6 @@ class VampNet(at.ml.BaseModel):
 
         return signal
 
-
     @torch.no_grad()
     def generate(
         self,
@@ -604,16 +610,14 @@ class VampNet(at.ml.BaseModel):
         typical_min_tokens=1,
         top_p=None,
         return_signal=True,
-        seed: int = None, 
+        seed: int = None,
         sample_cutoff: float = 1.0,
     ):
         if seed is not None:
             at.util.seed(seed)
         logging.debug(f"beginning generation with {sampling_steps} steps")
 
-
-
-        ##################### 
+        #####################
         # resolve initial z #
         #####################
         z = start_tokens
@@ -625,7 +629,6 @@ class VampNet(at.ml.BaseModel):
 
         logging.debug(f"created z with shape {z.shape}")
 
-
         #################
         # resolve mask #
         #################
@@ -636,9 +639,8 @@ class VampNet(at.ml.BaseModel):
         if mask.ndim == 2:
             mask = mask[:, None, :].repeat(1, z.shape[1], 1)
         # init_mask = mask.clone()
-        
-        logging.debug(f"created mask with shape {mask.shape}")
 
+        logging.debug(f"created mask with shape {mask.shape}")
 
         ###########
         # set up #
@@ -663,33 +665,33 @@ class VampNet(at.ml.BaseModel):
             logging.debug(f"step {i} of {sampling_steps}")
 
             # our current schedule step
-            r = scalar_to_batch_tensor(
-                (i + 1) / sampling_steps, 
-                z.shape[0]
-            ).to(z.device)
+            r = scalar_to_batch_tensor((i + 1) / sampling_steps, z.shape[0]).to(
+                z.device
+            )
             logging.debug(f"r: {r}")
 
             # get latents
             latents = self.embedding.from_codes(z_masked, codec)
             logging.debug(f"computed latents with shape: {latents.shape}")
 
-
             # infer from latents
             # NOTE: this collapses the codebook dimension into the sequence dimension
-            logits = self.forward(latents) # b, prob, seq
+            logits = self.forward(latents)  # b, prob, seq
             logits = logits.permute(0, 2, 1)  # b, seq, prob
             b = logits.shape[0]
 
             logging.debug(f"permuted logits with shape: {logits.shape}")
 
             sampled_z, selected_probs = sample_from_logits(
-                logits, sample=(
-                   (i / sampling_steps) <= sample_cutoff
-                ), 
+                logits,
+                sample=((i / sampling_steps) <= sample_cutoff),
                 temperature=sampling_temperature,
-                typical_filtering=typical_filtering, typical_mass=typical_mass,
+                typical_filtering=typical_filtering,
+                typical_mass=typical_mass,
                 typical_min_tokens=typical_min_tokens,
-                top_k=None, top_p=top_p, return_probs=True,
+                top_k=None,
+                top_p=top_p,
+                return_probs=True,
             )
 
             logging.debug(f"sampled z with shape: {sampled_z.shape}")
@@ -697,46 +699,38 @@ class VampNet(at.ml.BaseModel):
             # flatten z_masked and mask, so we can deal with the sampling logic
             # we'll unflatten them at the end of the loop for the next forward pass
             # remove conditioning codebooks, we'll add them back at the end
-            z_masked = codebook_flatten(z_masked[:, self.n_conditioning_codebooks:, :])           
+            z_masked = codebook_flatten(z_masked[:, self.n_conditioning_codebooks :, :])
 
             mask = (z_masked == self.mask_token).int()
-            
+
             # update the mask, remove conditioning codebooks from the mask
             logging.debug(f"updated mask with shape: {mask.shape}")
             # add z back into sampled z where the mask was false
-            sampled_z = torch.where(
-                mask.bool(), sampled_z, z_masked
-            )
+            sampled_z = torch.where(mask.bool(), sampled_z, z_masked)
             logging.debug(f"added z back into sampled z with shape: {sampled_z.shape}")
 
             # ignore any tokens that weren't masked
-            selected_probs = torch.where(
-                mask.bool(), selected_probs, torch.inf
-            )
+            selected_probs = torch.where(mask.bool(), selected_probs, torch.inf)
 
             # get the num tokens to mask, according to the schedule
-            num_to_mask = torch.floor(_gamma(r) * num_mask_tokens_at_start).unsqueeze(1).long()
+            num_to_mask = (
+                torch.floor(_gamma(r) * num_mask_tokens_at_start).unsqueeze(1).long()
+            )
             logging.debug(f"num to mask: {num_to_mask}")
 
             if i != (sampling_steps - 1):
                 num_to_mask = torch.maximum(
                     torch.tensor(1),
-                    torch.minimum(
-                        mask.sum(dim=-1, keepdim=True) - 1,
-                        num_to_mask
-                    )
+                    torch.minimum(mask.sum(dim=-1, keepdim=True) - 1, num_to_mask),
                 )
 
-
             # get our new mask
             mask = mask_by_random_topk(
-                num_to_mask, selected_probs, mask_temperature * (1-r)
-            )  
+                num_to_mask, selected_probs, mask_temperature * (1 - r)
+            )
 
             # update the mask
-            z_masked = torch.where(
-                mask.bool(), self.mask_token, sampled_z
-            )
+            z_masked = torch.where(mask.bool(), self.mask_token, sampled_z)
             logging.debug(f"updated z_masked with shape: {z_masked.shape}")
 
             z_masked = codebook_unflatten(z_masked, n_infer_codebooks)
@@ -745,35 +739,37 @@ class VampNet(at.ml.BaseModel):
 
             # add conditioning codebooks back to z_masked
             z_masked = torch.cat(
-                (z[:, :self.n_conditioning_codebooks, :], z_masked), dim=1
+                (z[:, : self.n_conditioning_codebooks, :], z_masked), dim=1
+            )
+            logging.debug(
+                f"added conditioning codebooks back to z_masked with shape: {z_masked.shape}"
             )
-            logging.debug(f"added conditioning codebooks back to z_masked with shape: {z_masked.shape}")
-
 
         # add conditioning codebooks back to sampled_z
         sampled_z = codebook_unflatten(sampled_z, n_infer_codebooks)
         sampled_z = torch.cat(
-            (z[:, :self.n_conditioning_codebooks, :], sampled_z), dim=1
+            (z[:, : self.n_conditioning_codebooks, :], sampled_z), dim=1
         )
 
-        logging.debug(f"finished sampling")
+        logging.debug("finished sampling")
 
         if return_signal:
             return self.to_signal(sampled_z, codec)
         else:
             return sampled_z
 
+
 def sample_from_logits(
-        logits, 
-        sample: bool = True,
-        temperature: float = 1.0,
-        top_k: int = None,
-        top_p: float = None,
-        typical_filtering: bool = False,
-        typical_mass: float = 0.2,
-        typical_min_tokens: int = 1,
-        return_probs: bool = False
-    ):
+    logits,
+    sample: bool = True,
+    temperature: float = 1.0,
+    top_k: int = None,
+    top_p: float = None,
+    typical_filtering: bool = False,
+    typical_mass: float = 0.2,
+    typical_min_tokens: int = 1,
+    return_probs: bool = False,
+):
     """Convenience function to sample from a categorial distribution with input as
     unnormalized logits.
 
@@ -801,9 +797,8 @@ def sample_from_logits(
     shp = logits.shape[:-1]
 
     if typical_filtering:
-        typical_filter(logits, 
-                        typical_mass=typical_mass, 
-                        typical_min_tokens=typical_min_tokens
+        typical_filter(
+            logits, typical_mass=typical_mass, typical_min_tokens=typical_min_tokens
         )
 
     # Apply top_k sampling
@@ -846,21 +841,20 @@ def sample_from_logits(
         return token, token_probs
     else:
         return token
-    
 
 
 def mask_by_random_topk(
-        num_to_mask: int, 
-        probs: torch.Tensor, 
-        temperature: float = 1.0, 
-    ):
+    num_to_mask: int,
+    probs: torch.Tensor,
+    temperature: float = 1.0,
+):
     """
     Args:
         num_to_mask (int): number of tokens to mask
         probs (torch.Tensor): probabilities for each sampled event, shape (batch, seq)
         temperature (float, optional): temperature. Defaults to 1.0.
     """
-    logging.debug(f"masking by random topk")
+    logging.debug("masking by random topk")
     logging.debug(f"num to mask: {num_to_mask}")
     logging.debug(f"probs shape: {probs.shape}")
     logging.debug(f"temperature: {temperature}")
@@ -875,9 +869,7 @@ def mask_by_random_topk(
     logging.debug(f"sorted idx shape: {sorted_idx.shape}")
 
     # get the cut off threshold, given the mask length
-    cut_off = torch.take_along_dim(
-        sorted_confidence, num_to_mask, axis=-1
-    )
+    cut_off = torch.take_along_dim(sorted_confidence, num_to_mask, axis=-1)
     logging.debug(f"cut off shape: {cut_off.shape}")
 
     # mask out the tokens
@@ -886,10 +878,12 @@ def mask_by_random_topk(
 
     return mask
 
+
 def typical_filter(
-        logits, 
-        typical_mass: float = 0.95,
-        typical_min_tokens: int = 1,):
+    logits,
+    typical_mass: float = 0.95,
+    typical_min_tokens: int = 1,
+):
     nb, nt, _ = logits.shape
     x_flat = rearrange(logits, "b t l -> (b t ) l")
     x_flat_norm = torch.nn.functional.log_softmax(x_flat, dim=-1)
@@ -898,9 +892,7 @@ def typical_filter(
 
     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)
-    )
+    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(
@@ -933,7 +925,7 @@ if __name__ == "__main__":
         ).to(device)
 
         r = torch.zeros(batch_size).to(device)
-        
+
         z_mask_latent = torch.rand(
             batch_size, model.latent_dim * model.n_codebooks, seq_len
         ).to(device)
@@ -942,12 +934,10 @@ if __name__ == "__main__":
         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"model has {num_params(model) / 1e6:<.3f}M parameters")
         print(f"prediction has shape {pred.shape}")
         breakpoint()
 
     args = argbind.parse_args()
     with argbind.scope(args):
         try_model()
-
-

vampnet/vampnet/scheduler.py

@@ -1,8 +1,6 @@
-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
@@ -44,4 +42,3 @@ class NoamScheduler:
 
         for p in self.optimizer.param_groups:
             p["lr"] = self.lr
-

vampnet/vampnet/util.py

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

wham.egg-info/PKG-INFO

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

wham.egg-info/SOURCES.txt

@@ -0,0 +1,33 @@
+LICENSE
+README.md
+pyproject.toml
+vampnet/app.py
+vampnet/setup.py
+vampnet/scripts/exp/eval.py
+vampnet/scripts/exp/experiment.py
+vampnet/scripts/exp/fine_tune.py
+vampnet/scripts/exp/train.py
+vampnet/scripts/utils/plots.py
+vampnet/scripts/utils/remove_quiet_files.py
+vampnet/scripts/utils/split.py
+vampnet/scripts/utils/split_long_audio_file.py
+vampnet/scripts/utils/stage.py
+vampnet/scripts/utils/visualize_embeddings.py
+vampnet/scripts/utils/xeno-canto-dl.py
+vampnet/scripts/utils/data/augment.py
+vampnet/scripts/utils/data/maestro-reorg.py
+vampnet/vampnet/__init__.py
+vampnet/vampnet/beats.py
+vampnet/vampnet/interface.py
+vampnet/vampnet/mask.py
+vampnet/vampnet/scheduler.py
+vampnet/vampnet/util.py
+vampnet/vampnet/modules/__init__.py
+vampnet/vampnet/modules/activations.py
+vampnet/vampnet/modules/layers.py
+vampnet/vampnet/modules/transformer.py
+wham.egg-info/PKG-INFO
+wham.egg-info/SOURCES.txt
+wham.egg-info/dependency_links.txt
+wham.egg-info/requires.txt
+wham.egg-info/top_level.txt

wham.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

wham.egg-info/requires.txt

@@ -0,0 +1,27 @@
+torch
+gradio
+argbind>=0.3.2
+numpy<1.24
+pydantic<3,>=2.0
+huggingface_hub
+loralib
+torch_pitch_shift
+soundfile
+pydub
+tqdm
+Cython
+pandas
+pathlib
+ffmpeg-python
+scikit-learn
+wandb
+gdown
+transformers
+fadtk
+urllib3>=2.0.2
+plotly
+pyharp
+ruff
+wavebeat @ git+https://github.com/hugofloresgarcia/wavebeat.git
+lac @ git+https://github.com/hugofloresgarcia/lac.git
+descript-audiotools @ git+https://github.com/hugofloresgarcia/audiotools.git

wham.egg-info/top_level.txt

@@ -0,0 +1 @@
+vampnet