beat tracker fix

.gitignore

@@ -175,14 +175,10 @@ lyrebird-audio-codec
 samples-*/**
 
 gradio-outputs/
-models/
 samples*/
 models-all/
 models.zip
+audiotools/ 
+descript-audio-codec/
+# *.pth
 .git-old
-conf/generated/*
-runs*/
-
-
-gtzan.zip
-.gtzan_emb_cache

README.md

@@ -7,27 +7,16 @@ sdk: gradio
 sdk_version: 3.36.1
 app_file: app.py
 pinned: false
-python_version: 3.9
+duplicated_from: hugggof/vampnet
 ---
 
 # VampNet
 
-This repository contains recipes for training generative music models on top of the Descript Audio Codec.
-
-## try `unloop`
-you can try vampnet in a co-creative looper called unloop. see this link: https://github.com/hugofloresgarcia/unloop
+This repository contains recipes for training generative music models on top of the Lyrebird Audio Codec.
 
 # Setting up
 
-**Requires Python 3.9**. 
-
-you'll need a Python 3.9 environment to run VampNet. This is due to a [known issue with madmom](https://github.com/hugofloresgarcia/vampnet/issues/15). 
-
-(for example, using conda)
-```bash
-conda create -n vampnet python=3.9
-conda activate vampnet
-```
+Requires Python 3.9 or later. 
 
 
 install VampNet
@@ -46,7 +35,7 @@ Config files are stored in the `conf/` folder.
 ### Licensing for Pretrained Models: 
 The weights for the models are licensed [`CC BY-NC-SA 4.0`](https://creativecommons.org/licenses/by-nc-sa/4.0/deed.ml). Likewise, any VampNet models fine-tuned on the pretrained models are also licensed [`CC BY-NC-SA 4.0`](https://creativecommons.org/licenses/by-nc-sa/4.0/deed.ml).
 
-Download the pretrained models from [this link](https://zenodo.org/record/8136629). Then, extract the models to the `models/` folder. 
+Download the pretrained models from [this link](https://zenodo.org/record/8136545). Then, extract the models to the `models/` folder. 
 
 
 # Usage
@@ -100,7 +89,7 @@ python  scripts/exp/train.py --args.load conf/<fine_tune_name>/c2f.yml
 
 launch the interface: 
 ```bash
-python  app.py --args.load conf/generated/<fine_tune_name>/interface.yml 
+python  demo.py --args.load conf/generated/<fine_tune_name>/interface.yml 
 ```
 
 

app.py

@@ -1,12 +1,3 @@
-# huggingface space exclusive
-import os
-
-# print("installing pyharp")
-# os.system('pip install "pyharp@git+https://github.com/audacitorch/pyharp.git"')
-# print("installing madmom")
-os.system('pip install cython')
-os.system('pip install madmom')
-
 from pathlib import Path
 from typing import Tuple
 import yaml
@@ -24,38 +15,27 @@ import gradio as gr
 from vampnet.interface import Interface
 from vampnet import mask as pmask
 
-from pyharp import ModelCard, build_endpoint
-
-
-
-# loader = AudioLoader()
+# Interface = argbind.bind(Interface)
 # AudioLoader = argbind.bind(at.data.datasets.AudioLoader)
 
-conf = argbind.parse_args()
-
-
-from torch_pitch_shift import pitch_shift, get_fast_shifts
-def shift_pitch(signal, interval: int):
-    signal.samples = pitch_shift(
-        signal.samples, 
-        shift=interval, 
-        sample_rate=signal.sample_rate
-    )
-    return signal
-
-def load_interface():
-    interface = Interface(
-        coarse_ckpt="./models/vampnet/coarse.pth",
-        coarse2fine_ckpt="./models/vampnet/c2f.pth",
-        codec_ckpt="./models/vampnet/codec.pth",
-        wavebeat_ckpt="./models/wavebeat.pth",
-        device="cuda" if torch.cuda.is_available() else "cpu",
-    )
-    return interface
+interface = Interface(
+    coarse_ckpt="./models/vampnet/coarse.pth",
+    coarse2fine_ckpt="./models/vampnet/c2f.pth",
+    codec_ckpt="./models/vampnet/codec.pth",
+    wavebeat_ckpt="./models/wavebeat.pth",
+    device="cuda" if torch.cuda.is_available() else "cpu",
+)
 
+# loader = AudioLoader()
+print(f"interface device is {interface.device}")
 
-interface = load_interface()
-
+# dataset = at.data.datasets.AudioDataset(
+#     loader,
+#     sample_rate=interface.codec.sample_rate,
+#     duration=interface.coarse.chunk_size_s,
+#     n_examples=5000,
+#     without_replacement=True,
+# )
 
 OUT_DIR = Path("gradio-outputs")
 OUT_DIR.mkdir(exist_ok=True, parents=True)
@@ -70,7 +50,7 @@ def load_audio(file):
     )
     sig = interface.preprocess(sig)
 
-    out_dir = OUT_DIR / "tmp" / str(uuid.uuid4())
+    out_dir = OUT_DIR / str(uuid.uuid4())
     out_dir.mkdir(parents=True, exist_ok=True)
     sig.write(out_dir / "input.wav")
     return sig.path_to_file
@@ -88,10 +68,6 @@ def _vamp(data, return_mask=False):
     out_dir = OUT_DIR / str(uuid.uuid4())
     out_dir.mkdir()
     sig = at.AudioSignal(data[input_audio])
-    sig = interface.preprocess(sig)
-
-    if data[pitch_shift_amt] != 0:
-        sig = shift_pitch(sig, data[pitch_shift_amt])
 
     z = interface.encode(sig)
 
@@ -131,58 +107,24 @@ def _vamp(data, return_mask=False):
     mask = pmask.codebook_unmask(mask, ncc)
 
 
-    print(f"dropout {data[dropout]}")
-    print(f"masktemp {data[masktemp]}")
-    print(f"sampletemp {data[sampletemp]}")
-    print(f"top_p {data[top_p]}")
-    print(f"prefix_s {data[prefix_s]}")
-    print(f"suffix_s {data[suffix_s]}")
-    print(f"rand_mask_intensity {data[rand_mask_intensity]}")
-    print(f"num_steps {data[num_steps]}")
-    print(f"periodic_p {data[periodic_p]}")
-    print(f"periodic_w {data[periodic_w]}")
-    print(f"n_conditioning_codebooks {data[n_conditioning_codebooks]}")
-    print(f"use_coarse2fine {data[use_coarse2fine]}")
-    print(f"onset_mask_width {data[onset_mask_width]}")
-    print(f"beat_mask_width {data[beat_mask_width]}")
-    print(f"beat_mask_downbeats {data[beat_mask_downbeats]}")
-    print(f"stretch_factor {data[stretch_factor]}")
-    print(f"seed {data[seed]}")
-    print(f"pitch_shift_amt {data[pitch_shift_amt]}")
-    print(f"sample_cutoff {data[sample_cutoff]}")
-    
-    
-    _top_p = data[top_p] if data[top_p] > 0 else None
+    print(f"created mask with: linear random {data[rand_mask_intensity]}, inpaint {data[prefix_s]}:{data[suffix_s]}, periodic {data[periodic_p]}:{data[periodic_w]}, dropout {data[dropout]}, codebook unmask {ncc}, onset mask {data[onset_mask_width]}, num steps {data[num_steps]}, init temp {data[temp]},  use coarse2fine {data[use_coarse2fine]}")
     # save the mask as a txt file
     np.savetxt(out_dir / "mask.txt", mask[:,0,:].long().cpu().numpy())
 
-    _seed = data[seed] if data[seed] > 0 else None
     zv, mask_z = interface.coarse_vamp(
         z, 
         mask=mask,
         sampling_steps=data[num_steps],
-        mask_temperature=data[masktemp]*10,
-        sampling_temperature=data[sampletemp],
+        temperature=float(data[temp]*10),
         return_mask=True, 
         typical_filtering=data[typical_filtering], 
         typical_mass=data[typical_mass], 
         typical_min_tokens=data[typical_min_tokens], 
-        top_p=_top_p,
         gen_fn=interface.coarse.generate,
-        seed=_seed,
-        sample_cutoff=data[sample_cutoff],
     )
 
     if use_coarse2fine: 
-        zv = interface.coarse_to_fine(
-            zv, 
-            mask_temperature=data[masktemp]*10, 
-            sampling_temperature=data[sampletemp],
-            mask=mask,
-            sampling_steps=data[num_steps] // 2,
-            sample_cutoff=data[sample_cutoff], 
-            seed=_seed,
-        )
+        zv = interface.coarse_to_fine(zv, temperature=data[temp])
 
     sig = interface.to_signal(zv).cpu()
     print("done")
@@ -215,9 +157,7 @@ def save_vamp(data):
     sig_out.write(out_dir / "output.wav")
     
     _data = {
-        "masktemp": data[masktemp],
-        "sampletemp": data[sampletemp],
-        "top_p": data[top_p],
+        "temp": data[temp],
         "prefix_s": data[prefix_s],
         "suffix_s": data[suffix_s],
         "rand_mask_intensity": data[rand_mask_intensity],
@@ -228,8 +168,6 @@ def save_vamp(data):
         "n_conditioning_codebooks": data[n_conditioning_codebooks], 
         "use_coarse2fine": data[use_coarse2fine],
         "stretch_factor": data[stretch_factor],
-        "seed": data[seed],
-        "samplecutoff": data[sample_cutoff],
     }
 
     # save with yaml
@@ -245,54 +183,13 @@ def save_vamp(data):
     return f"saved! your save code is {out_dir.stem}", zip_path
 
 
-def harp_vamp(_input_audio, _beat_mask_width, _sampletemp):
-
-    out_dir = OUT_DIR / str(uuid.uuid4())
-    out_dir.mkdir()
-    sig = at.AudioSignal(_input_audio)
-    sig = interface.preprocess(sig)
-
-    z = interface.encode(sig)
-
-    # build the mask
-    mask = pmask.linear_random(z, 1.0)
-    if _beat_mask_width > 0:
-        beat_mask = interface.make_beat_mask(
-            sig,
-            after_beat_s=(_beat_mask_width/1000), 
-        )
-        mask = pmask.mask_and(mask, beat_mask)
-
-    # save the mask as a txt file
-    zv, mask_z = interface.coarse_vamp(
-        z, 
-        mask=mask,
-        sampling_temperature=_sampletemp,
-        return_mask=True, 
-        gen_fn=interface.coarse.generate,
-    )
-
-
-    zv = interface.coarse_to_fine(
-        zv, 
-        sampling_temperature=_sampletemp,
-        mask=mask,
-    )
-
-    sig = interface.to_signal(zv).cpu()
-    print("done")
-
-    sig.write(out_dir / "output.wav")
-
-    return sig.path_to_file
-
 with gr.Blocks() as demo:
 
     with gr.Row():
         with gr.Column():
-            gr.Markdown("# VampNet Audio Vamping")
+            gr.Markdown("# VampNet")
             gr.Markdown("""## Description:
-            This is a demo of the VampNet, a generative audio model that transforms the input audio based on the chosen settings. 
+            This is a demo of VampNet, a masked generative music model capable of doing music variations. 
             You can control the extent and nature of variation with a set of manual controls and presets. 
             Use this interface to experiment with different mask settings and explore the audio outputs.
             """)
@@ -300,8 +197,8 @@ with gr.Blocks() as demo:
             gr.Markdown("""
             ## Instructions:
             1. You can start by uploading some audio, or by loading the example audio. 
-            2. Choose a preset for the vamp operation, or manually adjust the controls to customize the mask settings. 
-            3. Click the "generate (vamp)!!!" button to apply the vamp operation. Listen to the output audio.
+            2. Choose a preset for the vamp operation, or manually adjust the controls to customize the mask settings. Click the load preset button. 
+            3. Click the "generate (vamp)!!!" button to generate audio. Listen to the output audio, and the masked audio to hear the mask hints.
             4. Optionally, you can add some notes and save the result. 
             5. You can also use the output as the new input and continue experimenting!
             """)
@@ -352,12 +249,15 @@ with gr.Blocks() as demo:
                         "beat_mask_downbeats": False,
                     }, 
                     "slight periodic variation": {
+<<<<<<< HEAD
                         "periodic_p": 5,
                         "onset_mask_width": 5,
                         "beat_mask_width": 0,
                         "beat_mask_downbeats": False,
                     },
                     "moderate periodic variation": {
+=======
+>>>>>>> main
                         "periodic_p": 13,
                         "onset_mask_width": 5,
                         "beat_mask_width": 0,
@@ -378,15 +278,9 @@ with gr.Blocks() as demo:
                     "beat-driven variation": {
                         "periodic_p": 0,
                         "onset_mask_width": 0,
-                        "beat_mask_width": 50,
+                        "beat_mask_width": 20,
                         "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,
@@ -408,20 +302,20 @@ with gr.Blocks() as demo:
                     minimum=0,
                     maximum=128, 
                     step=1,
-                    value=3, 
+                    value=13, 
                 )
 
 
                 onset_mask_width = gr.Slider(
                     label="onset mask width (multiplies with the periodic mask, 1 step ~= 10milliseconds) ",
                     minimum=0,
-                    maximum=100,
+                    maximum=20,
                     step=1,
                     value=5,
                 )
 
                 beat_mask_width = gr.Slider(
-                    label="beat prompt (ms)",
+                    label="beat mask width (in milliseconds)",
                     minimum=0,
                     maximum=200,
                     value=0,
@@ -433,14 +327,6 @@ with gr.Blocks() as demo:
 
 
                 with gr.Accordion("extras ", open=False):
-                    pitch_shift_amt = gr.Slider(
-                        label="pitch shift amount (semitones)",
-                        minimum=-12,
-                        maximum=12,
-                        step=1,
-                        value=0,
-                    )
-
                     rand_mask_intensity = gr.Slider(
                         label="random mask intensity. (If this is less than 1, scatters prompts throughout the audio, should be between 0.9 and 1.0)",
                         minimum=0.0,
@@ -500,34 +386,21 @@ with gr.Blocks() as demo:
                     value=0.0
                 )
 
-            masktemp = gr.Slider(
-                label="mask temperature",
+            temp = gr.Slider(
+                label="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
+                value=1.8
             )
-        
+
 
 
             with gr.Accordion("sampling settings", open=False):
-                top_p = gr.Slider(
-                    label="top p (0.0 = off)",
-                    minimum=0.0,
-                    maximum=1.0,
-                    value=0.9
-                )
                 typical_filtering = gr.Checkbox(
                     label="typical filtering ",
                     value=False
                 )
-                typical_mass = gr.Slider( 
+                typical_mass = gr.Slider(
                     label="typical mass (should probably stay between 0.1 and 0.5)",
                     minimum=0.01,
                     maximum=0.99,
@@ -540,18 +413,10 @@ with gr.Blocks() as demo:
                     step=1,
                     value=64
                 )
-                sample_cutoff = gr.Slider(
-                    label="sample cutoff",
-                    minimum=0.0,
-                    maximum=1.0,
-                    value=0.5, 
-                    step=0.01
-                )
 
             use_coarse2fine = gr.Checkbox(
                 label="use coarse2fine",
-                value=True, 
-                visible=False
+                value=True
             )
 
             num_steps = gr.Slider(
@@ -571,24 +436,8 @@ with gr.Blocks() as demo:
             )
 
 
-            seed = gr.Number(
-                label="seed (0 for random)",
-                value=0,
-                precision=0,
-            )
-
-
-
         # mask settings
         with gr.Column():
-
-            # lora_choice = gr.Dropdown(
-            #     label="lora choice", 
-            #     choices=list(loras.keys()),
-            #     value=LORA_NONE, 
-            #     visible=False
-            # )
-
             vamp_button = gr.Button("generate (vamp)!!!")
             output_audio = gr.Audio(
                 label="output audio",
@@ -614,9 +463,7 @@ with gr.Blocks() as demo:
     _inputs = {
             input_audio, 
             num_steps,
-            masktemp,
-            sampletemp,
-            top_p,
+            temp,
             prefix_s, suffix_s, 
             rand_mask_intensity, 
             periodic_p, periodic_w,
@@ -629,11 +476,7 @@ with gr.Blocks() as demo:
             typical_mass,
             typical_min_tokens,
             beat_mask_width,
-            beat_mask_downbeats,
-            seed, 
-            # lora_choice,
-            pitch_shift_amt, 
-            sample_cutoff
+            beat_mask_downbeats
         }
   
     # connect widgets
@@ -663,24 +506,4 @@ with gr.Blocks() as demo:
         outputs=[thank_you, download_file]
     )
 
-    # harp stuff
-    harp_inputs = [
-        input_audio, 
-        beat_mask_width,
-        sampletemp,
-    ]
-    
-    build_endpoint(
-        inputs=harp_inputs,  
-        output=output_audio, 
-        process_fn=harp_vamp,
-        card=ModelCard(
-            name="vampnet", 
-            description="Generate variations on music input, based on small prompts around the beat. NOTE: vampnet's has a maximum context length of 10 seconds. Please split all audio clips into 10 second chunks, or processing will result in an error. ", 
-            author="Hugo Flores García", 
-            tags=["music", "generative"]
-        ), 
-        visible=False
-    )
-
-demo.launch()
+demo.queue().launch()

conf/generated-v0/berta-goldman-speech/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+save_path: ./runs/berta-goldman-speech/c2f
+train/AudioLoader.sources:
+- /media/CHONK/hugo/Berta-Caceres-2015-Goldman-Speech.mp3
+val/AudioLoader.sources:
+- /media/CHONK/hugo/Berta-Caceres-2015-Goldman-Speech.mp3

conf/generated-v0/berta-goldman-speech/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+save_path: ./runs/berta-goldman-speech/coarse
+train/AudioLoader.sources:
+- /media/CHONK/hugo/Berta-Caceres-2015-Goldman-Speech.mp3
+val/AudioLoader.sources:
+- /media/CHONK/hugo/Berta-Caceres-2015-Goldman-Speech.mp3

conf/generated-v0/berta-goldman-speech/interface.yml

@@ -0,0 +1,5 @@
+AudioLoader.sources:
+- /media/CHONK/hugo/Berta-Caceres-2015-Goldman-Speech.mp3
+Interface.coarse2fine_ckpt: ./runs/berta-goldman-speech/c2f/best/vampnet/weights.pth
+Interface.coarse_ckpt: ./runs/berta-goldman-speech/coarse/best/vampnet/weights.pth
+Interface.codec_ckpt: ./models/spotdl/codec.pth

conf/generated-v0/gamelan-xeno-canto/c2f.yml

@@ -0,0 +1,17 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+save_path: ./runs/gamelan-xeno-canto/c2f
+train/AudioLoader.sources:
+- /media/CHONK/hugo/loras/Sound Tracker - Gamelan (Indonesia) [UEWCCSuHsuQ].mp3
+- /media/CHONK/hugo/loras/xeno-canto-2
+val/AudioLoader.sources:
+- /media/CHONK/hugo/loras/Sound Tracker - Gamelan (Indonesia) [UEWCCSuHsuQ].mp3
+- /media/CHONK/hugo/loras/xeno-canto-2

conf/generated-v0/gamelan-xeno-canto/coarse.yml

@@ -0,0 +1,10 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+save_path: ./runs/gamelan-xeno-canto/coarse
+train/AudioLoader.sources:
+- /media/CHONK/hugo/loras/Sound Tracker - Gamelan (Indonesia) [UEWCCSuHsuQ].mp3
+- /media/CHONK/hugo/loras/xeno-canto-2
+val/AudioLoader.sources:
+- /media/CHONK/hugo/loras/Sound Tracker - Gamelan (Indonesia) [UEWCCSuHsuQ].mp3
+- /media/CHONK/hugo/loras/xeno-canto-2

conf/generated-v0/gamelan-xeno-canto/interface.yml

@@ -0,0 +1,6 @@
+AudioLoader.sources:
+- /media/CHONK/hugo/loras/Sound Tracker - Gamelan (Indonesia) [UEWCCSuHsuQ].mp3
+- /media/CHONK/hugo/loras/xeno-canto-2
+Interface.coarse2fine_ckpt: ./runs/gamelan-xeno-canto/c2f/best/vampnet/weights.pth
+Interface.coarse_ckpt: ./runs/gamelan-xeno-canto/coarse/best/vampnet/weights.pth
+Interface.codec_ckpt: ./models/spotdl/codec.pth

conf/generated-v0/nasralla/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+save_path: ./runs/nasralla/c2f
+train/AudioLoader.sources:
+- /media/CHONK/hugo/nasralla
+val/AudioLoader.sources:
+- /media/CHONK/hugo/nasralla

conf/generated-v0/nasralla/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+save_path: ./runs/nasralla/coarse
+train/AudioLoader.sources:
+- /media/CHONK/hugo/nasralla
+val/AudioLoader.sources:
+- /media/CHONK/hugo/nasralla

conf/generated-v0/nasralla/interface.yml

@@ -0,0 +1,5 @@
+AudioLoader.sources:
+- /media/CHONK/hugo/nasralla
+Interface.coarse2fine_ckpt: ./runs/nasralla/c2f/best/vampnet/weights.pth
+Interface.coarse_ckpt: ./runs/nasralla/coarse/best/vampnet/weights.pth
+Interface.codec_ckpt: ./models/spotdl/codec.pth

conf/generated/breaks-steps/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/c2f.pth
+save_path: ./runs/breaks-steps/c2f
+train/AudioLoader.sources: &id001
+- /media/CHONK/hugo/breaks-steps
+val/AudioLoader.sources: *id001

conf/generated/breaks-steps/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/coarse.pth
+save_path: ./runs/breaks-steps/coarse
+train/AudioLoader.sources: &id001
+- /media/CHONK/hugo/breaks-steps
+val/AudioLoader.sources: *id001

conf/generated/breaks-steps/interface.yml

@@ -0,0 +1,7 @@
+AudioLoader.sources:
+- - /media/CHONK/hugo/breaks-steps
+Interface.coarse2fine_ckpt: ./models/spotdl/c2f.pth
+Interface.coarse2fine_lora_ckpt: ./runs/breaks-steps/c2f/latest/lora.pth
+Interface.coarse_ckpt: ./models/spotdl/coarse.pth
+Interface.coarse_lora_ckpt: ./runs/breaks-steps/coarse/latest/lora.pth
+Interface.codec_ckpt: ./models/spotdl/codec.pth

conf/generated/bulgarian-tv-choir/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/c2f.pth
+save_path: ./runs/bulgarian-tv-choir/c2f
+train/AudioLoader.sources: &id001
+- /media/CHONK/hugo/loras/bulgarian-female-tv-choir/
+val/AudioLoader.sources: *id001

conf/generated/bulgarian-tv-choir/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/coarse.pth
+save_path: ./runs/bulgarian-tv-choir/coarse
+train/AudioLoader.sources: &id001
+- /media/CHONK/hugo/loras/bulgarian-female-tv-choir/
+val/AudioLoader.sources: *id001

conf/generated/bulgarian-tv-choir/interface.yml

@@ -0,0 +1,7 @@
+AudioLoader.sources:
+- - /media/CHONK/hugo/loras/bulgarian-female-tv-choir/
+Interface.coarse2fine_ckpt: ./models/spotdl/c2f.pth
+Interface.coarse2fine_lora_ckpt: ./runs/bulgarian-tv-choir/c2f/latest/lora.pth
+Interface.coarse_ckpt: ./models/spotdl/coarse.pth
+Interface.coarse_lora_ckpt: ./runs/bulgarian-tv-choir/coarse/latest/lora.pth
+Interface.codec_ckpt: ./models/spotdl/codec.pth

conf/generated/dariacore/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/c2f.pth
+save_path: ./runs/dariacore/c2f
+train/AudioLoader.sources: &id001
+- /media/CHONK/hugo/loras/dariacore
+val/AudioLoader.sources: *id001

conf/generated/dariacore/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/coarse.pth
+save_path: ./runs/dariacore/coarse
+train/AudioLoader.sources: &id001
+- /media/CHONK/hugo/loras/dariacore
+val/AudioLoader.sources: *id001

conf/generated/dariacore/interface.yml

@@ -0,0 +1,7 @@
+AudioLoader.sources:
+- - /media/CHONK/hugo/loras/dariacore
+Interface.coarse2fine_ckpt: ./models/spotdl/c2f.pth
+Interface.coarse2fine_lora_ckpt: ./runs/dariacore/c2f/latest/lora.pth
+Interface.coarse_ckpt: ./models/spotdl/coarse.pth
+Interface.coarse_lora_ckpt: ./runs/dariacore/coarse/latest/lora.pth
+Interface.codec_ckpt: ./models/spotdl/codec.pth

conf/generated/musica-bolero-marimba/c2f.yml

@@ -0,0 +1,18 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/c2f.pth
+save_path: ./runs/musica-bolero-marimba/c2f
+train/AudioLoader.sources:
+- /media/CHONK/hugo/loras/boleros
+- /media/CHONK/hugo/loras/marimba-honduras
+val/AudioLoader.sources:
+- /media/CHONK/hugo/loras/boleros
+- /media/CHONK/hugo/loras/marimba-honduras

conf/generated/musica-bolero-marimba/coarse.yml

@@ -0,0 +1,11 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/coarse.pth
+save_path: ./runs/musica-bolero-marimba/coarse
+train/AudioLoader.sources:
+- /media/CHONK/hugo/loras/boleros
+- /media/CHONK/hugo/loras/marimba-honduras
+val/AudioLoader.sources:
+- /media/CHONK/hugo/loras/boleros
+- /media/CHONK/hugo/loras/marimba-honduras

conf/generated/musica-bolero-marimba/interface.yml

@@ -0,0 +1,8 @@
+AudioLoader.sources:
+- /media/CHONK/hugo/loras/boleros
+- /media/CHONK/hugo/loras/marimba-honduras
+Interface.coarse2fine_ckpt: ./models/spotdl/c2f.pth
+Interface.coarse2fine_lora_ckpt: ./runs/musica-bolero-marimba/c2f/latest/lora.pth
+Interface.coarse_ckpt: ./models/spotdl/coarse.pth
+Interface.coarse_lora_ckpt: ./runs/musica-bolero-marimba/coarse/latest/lora.pth
+Interface.codec_ckpt: ./models/spotdl/codec.pth

conf/generated/panchos/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/c2f.pth
+save_path: ./runs/panchos/c2f
+train/AudioLoader.sources: &id001
+- /media/CHONK/hugo/loras/panchos/
+val/AudioLoader.sources: *id001

conf/generated/panchos/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/coarse.pth
+save_path: ./runs/panchos/coarse
+train/AudioLoader.sources: &id001
+- /media/CHONK/hugo/loras/panchos/
+val/AudioLoader.sources: *id001

conf/generated/panchos/interface.yml

@@ -0,0 +1,7 @@
+AudioLoader.sources:
+- - /media/CHONK/hugo/loras/panchos/
+Interface.coarse2fine_ckpt: ./models/spotdl/c2f.pth
+Interface.coarse2fine_lora_ckpt: ./runs/panchos/c2f/latest/lora.pth
+Interface.coarse_ckpt: ./models/spotdl/coarse.pth
+Interface.coarse_lora_ckpt: ./runs/panchos/coarse/latest/lora.pth
+Interface.codec_ckpt: ./models/spotdl/codec.pth

conf/generated/titi-monkey/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/c2f.pth
+save_path: ./runs/titi-monkey/c2f
+train/AudioLoader.sources: &id001
+- /media/CHONK/hugo/loras/titi-monkey.mp3
+val/AudioLoader.sources: *id001

conf/generated/titi-monkey/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/coarse.pth
+save_path: ./runs/titi-monkey/coarse
+train/AudioLoader.sources: &id001
+- /media/CHONK/hugo/loras/titi-monkey.mp3
+val/AudioLoader.sources: *id001

conf/generated/titi-monkey/interface.yml

@@ -0,0 +1,7 @@
+AudioLoader.sources:
+- - /media/CHONK/hugo/loras/titi-monkey.mp3
+Interface.coarse2fine_ckpt: ./models/spotdl/c2f.pth
+Interface.coarse2fine_lora_ckpt: ./runs/titi-monkey/c2f/latest/lora.pth
+Interface.coarse_ckpt: ./models/spotdl/coarse.pth
+Interface.coarse_lora_ckpt: ./runs/titi-monkey/coarse/latest/lora.pth
+Interface.codec_ckpt: ./models/spotdl/codec.pth

conf/generated/xeno-canto/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/c2f.pth
+save_path: ./runs/xeno-canto/c2f
+train/AudioLoader.sources: &id001
+- /media/CHONK/hugo/loras/xeno-canto-2/
+val/AudioLoader.sources: *id001

conf/generated/xeno-canto/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+fine_tune_checkpoint: ./models/spotdl/coarse.pth
+save_path: ./runs/xeno-canto/coarse
+train/AudioLoader.sources: &id001
+- /media/CHONK/hugo/loras/xeno-canto-2/
+val/AudioLoader.sources: *id001

conf/generated/xeno-canto/interface.yml

@@ -0,0 +1,7 @@
+AudioLoader.sources:
+- - /media/CHONK/hugo/loras/xeno-canto-2/
+Interface.coarse2fine_ckpt: ./mod els/spotdl/c2f.pth
+Interface.coarse2fine_lora_ckpt: ./runs/xeno-canto/c2f/latest/lora.pth
+Interface.coarse_ckpt: ./models/spotdl/coarse.pth
+Interface.coarse_lora_ckpt: ./runs/xeno-canto/coarse/latest/lora.pth
+Interface.codec_ckpt: ./models/spotdl/codec.pth

conf/lora/birds.yml

@@ -0,0 +1,10 @@
+$include:
+  - conf/lora/lora.yml
+
+fine_tune: True
+
+train/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/birds
+
+val/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/birds

conf/lora/birdss.yml

@@ -0,0 +1,12 @@
+$include:
+  - conf/lora/lora.yml
+
+fine_tune: True
+
+train/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/birds
+  - /media/CHONK/hugo/spotdl/subsets/this-is-charlie-parker/
+
+val/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/birds
+  - /media/CHONK/hugo/spotdl/subsets/this-is-charlie-parker/

conf/lora/constructions.yml

@@ -0,0 +1,10 @@
+$include:
+  - conf/lora/lora.yml
+
+fine_tune: True
+
+train/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/constructions/third.mp3
+
+val/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/constructions/third.mp3

conf/lora/ella-baila-sola.yml

@@ -0,0 +1,10 @@
+$include:
+  - conf/lora/lora.yml
+
+fine_tune: True
+
+train/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/ella-baila-sola.mp3
+
+val/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/ella-baila-sola.mp3

conf/lora/gas-station.yml

@@ -0,0 +1,10 @@
+$include:
+  - conf/lora/lora.yml
+
+fine_tune: True
+
+train/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/gas-station-sushi.mp3
+
+val/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/gas-station-sushi.mp3

conf/lora/lora-is-this-charlie-parker.yml

@@ -0,0 +1,10 @@
+$include:
+  - conf/lora/lora.yml
+
+fine_tune: True
+
+train/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/this-is-charlie-parker/Charlie Parker - Donna Lee.mp3
+
+val/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/this-is-charlie-parker/Charlie Parker - Donna Lee.mp3

conf/lora/lora.yml

@@ -3,20 +3,20 @@ $include:
 
 fine_tune: True
 
-train/AudioDataset.n_examples: 100000000
-val/AudioDataset.n_examples: 500
+train/AudioDataset.n_examples: 10000000
+
+val/AudioDataset.n_examples: 10
 
 
 NoamScheduler.warmup: 500
 
-batch_size: 6
+batch_size: 7
 num_workers: 7
-save_iters: [10000, 20000, 30000, 40000, 50000, 100000]
-sample_freq: 1000
-val_freq: 500
+epoch_length: 100
+save_audio_epochs: 10
 
 AdamW.lr: 0.0001
 
 # let's us organize sound classes into folders and choose from those sound classes uniformly
 AudioDataset.without_replacement: False
-num_iters: 500000
+max_epochs: 500

conf/lora/underworld.yml

@@ -0,0 +1,10 @@
+$include:
+  - conf/lora/lora.yml
+
+fine_tune: True
+
+train/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/underworld.mp3
+
+val/AudioLoader.sources:
+  - /media/CHONK/hugo/spotdl/subsets/underworld.mp3

conf/lora/xeno-canto/c2f.yml

@@ -0,0 +1,21 @@
+$include:
+  - conf/lora/lora.yml
+
+fine_tune: True
+
+train/AudioLoader.sources:
+  - /media/CHONK/hugo/xeno-canto-2
+
+val/AudioLoader.sources:
+  - /media/CHONK/hugo/xeno-canto-2
+
+
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+
+VampNet.embedding_dim: 1280
+VampNet.n_layers: 16
+VampNet.n_heads: 20
+
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0

conf/lora/xeno-canto/coarse.yml

@@ -0,0 +1,10 @@
+$include:
+  - conf/lora/lora.yml
+
+fine_tune: True
+
+train/AudioLoader.sources:
+  - /media/CHONK/hugo/xeno-canto-2
+
+val/AudioLoader.sources:
+  - /media/CHONK/hugo/xeno-canto-2

conf/vampnet-musdb-drums.yml

@@ -0,0 +1,22 @@
+$include:
+  - conf/vampnet.yml
+
+VampNet.embedding_dim: 512
+VampNet.n_layers: 12
+VampNet.n_heads: 8
+
+AudioDataset.duration: 12.0
+
+train/AudioDataset.n_examples: 10000000
+train/AudioLoader.sources:
+  - /data/musdb18hq/train/**/*drums.wav
+
+
+val/AudioDataset.n_examples: 500
+val/AudioLoader.sources:
+  - /data/musdb18hq/test/**/*drums.wav
+
+
+test/AudioDataset.n_examples: 1000
+test/AudioLoader.sources:
+  - /data/musdb18hq/test/**/*drums.wav

conf/vampnet.yml

@@ -1,17 +1,21 @@
 
-codec_ckpt: ./models/vampnet/codec.pth
+codec_ckpt: ./models/spotdl/codec.pth
 save_path: ckpt
-
-num_iters: 1000000000
-save_iters: [10000, 50000, 100000, 300000, 500000]
+max_epochs: 1000
+epoch_length: 1000
+save_audio_epochs: 2
 val_idx: [0,1,2,3,4,5,6,7,8,9]
-sample_freq: 10000
-val_freq: 1000
+
+prefix_amt: 0.0
+suffix_amt: 0.0
+prefix_dropout: 0.1
+suffix_dropout: 0.1
 
 batch_size: 8
 num_workers: 10
 
 # Optimization
+detect_anomaly: false
 amp: false
 
 CrossEntropyLoss.label_smoothing: 0.1
@@ -21,6 +25,9 @@ AdamW.lr: 0.001
 NoamScheduler.factor: 2.0
 NoamScheduler.warmup: 10000
 
+PitchShift.shift_amount: [const, 0]
+PitchShift.prob: 0.0
+
 VampNet.vocab_size: 1024
 VampNet.n_codebooks: 4
 VampNet.n_conditioning_codebooks: 0
@@ -32,7 +39,7 @@ VampNet.n_heads: 20
 VampNet.flash_attn: false
 VampNet.dropout: 0.1
 
-AudioLoader.relative_path: ""
+AudioLoader.relative_path: /data/
 AudioDataset.loudness_cutoff: -30.0
 AudioDataset.without_replacement: true
 AudioLoader.shuffle: true
@@ -41,9 +48,12 @@ AudioDataset.duration: 10.0
 
 train/AudioDataset.n_examples: 10000000
 train/AudioLoader.sources:
-  - /media/CHONK/hugo/spotdl/audio-train
+  - /data/spotdl/audio/train
 
 val/AudioDataset.n_examples: 2000
 val/AudioLoader.sources:
-  - /media/CHONK/hugo/spotdl/audio-val
+  - /data/spotdl/audio/val
 
+test/AudioDataset.n_examples: 1000
+test/AudioLoader.sources:
+  - /data/spotdl/audio/test

requirements.txt

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

scripts/exp/fine_tune.py

@@ -35,7 +35,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/spotdl/c2f.pth"
     }
 
     finetune_coarse_conf = {
@@ -44,16 +44,17 @@ 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/spotdl/coarse.pth"
     }
 
     interface_conf = {
-        "Interface.coarse_ckpt": f"./runs/{name}/coarse/latest/vampnet/weights.pth",
+        "Interface.coarse_ckpt": f"./models/spotdl/coarse.pth",
+        "Interface.coarse_lora_ckpt": f"./runs/{name}/coarse/latest/lora.pth",
 
-        "Interface.coarse2fine_ckpt": f"./runs/{name}/c2f/latest/vampnet/weights.pth",
-        "Interface.wavebeat_ckpt": "./models/wavebeat.pth",
+        "Interface.coarse2fine_ckpt": f"./models/spotdl/c2f.pth",
+        "Interface.coarse2fine_lora_ckpt": f"./runs/{name}/c2f/latest/lora.pth",
 
-        "Interface.codec_ckpt": "./models/vampnet/codec.pth",
+        "Interface.codec_ckpt": "./models/spotdl/codec.pth",
         "AudioLoader.sources": [audio_files_or_folders],
     }
 

scripts/exp/train.py

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

scripts/utils/{data/augment.py → augment.py}

@@ -5,19 +5,34 @@ from audiotools import AudioSignal
 
 import argbind
 import tqdm
-import torch
 
 
-from torch_pitch_shift import pitch_shift, get_fast_shifts
-from torch_time_stretch import time_stretch, get_fast_stretches
+from pedalboard import (
+   Compressor, Gain, Chorus, LadderFilter, Phaser, Convolution, Reverb, Pedalboard 
+)
+from pedalboard.io import AudioFile 
 
-from audiotools.core.util import sample_from_dist
+# Read in a whole file, resampling to our desired sample rate:
+samplerate = 44100.0
+with AudioFile('guitar-input.wav').resampled_to(samplerate) as f:
+  audio = f.read(f.frames)
+
+# Make a pretty interesting sounding guitar pedalboard:
+board = Pedalboard([
+    Compressor(threshold_db=-50, ratio=25),
+    Gain(gain_db=30),
+    Chorus(),
+    LadderFilter(mode=LadderFilter.Mode.HPF12, cutoff_hz=900),
+    Phaser(),
+    Convolution("./guitar_amp.wav", 1.0),
+    Reverb(room_size=0.25),
+])
 
 
 @argbind.bind(without_prefix=True)
 def augment(
-    audio_folder: Path = None,
-    dest_folder: Path = None,
+    audio_folder: Path,
+    dest_folder: Path,
     n_augmentations: int = 10,
 ):
     """ 
@@ -26,8 +41,7 @@ def augment(
         The dest foler will contain a folder for each of the clean dataset's files. 
         Under each of these folders, there will be a clean file and many augmented files.
     """
-    assert audio_folder is not None
-    assert dest_folder is not None
+
     audio_files = at.util.find_audio(audio_folder)
 
     for audio_file in tqdm.tqdm(audio_files):
@@ -35,33 +49,5 @@ def augment(
         subdir = subtree / audio_file.stem
         subdir.mkdir(parents=True, exist_ok=True)
 
-        src = AudioSignal(audio_file).to("cuda" if torch.cuda.is_available() else "cpu")
-
-        
-        for i, chunk in tqdm.tqdm(enumerate(src.windows(10, 10))):
-            # apply pedalboard transforms
-            for j in range(n_augmentations):
-                # pitch shift between -7 and 7 semitones
-                import random
-                dst = chunk.clone()
-                dst.samples = pitch_shift(
-                    dst.samples, 
-                    shift=random.choice(get_fast_shifts(src.sample_rate, 
-                            condition=lambda x: x >= 0.25 and x <= 1.0)), 
-                    sample_rate=src.sample_rate
-                )
-                dst.samples = time_stretch(
-                    dst.samples,
-                    stretch=random.choice(get_fast_stretches(src.sample_rate, 
-                                          condition=lambda x: x >= 0.667 and x <= 1.5, )),
-                    sample_rate=src.sample_rate, 
-                )
-
-                dst.cpu().write(subdir / f"{i}-{j}.wav")
-
-
-if __name__ == "__main__":
-    args = argbind.parse_args()
-
-    with argbind.scope(args):
-        augment()
+        # apply pedalboard transforms
+        for i in range(n_augmentations):

scripts/utils/gtzan_embeddings.py

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