c2f prompts

.gitignore

@@ -175,14 +175,11 @@ 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,7 +7,6 @@ sdk: gradio
 sdk_version: 3.36.1
 app_file: app.py
 pinned: false
-python_version: 3.9
 ---
 
 # VampNet
@@ -19,15 +18,7 @@ you can try vampnet in a co-creative looper called unloop. see this link: https:
 
 # 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
@@ -100,7 +91,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",
+)
 
-interface = load_interface()
+# loader = AudioLoader()
+print(f"interface device is {interface.device}")
 
+# 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], mask=mask)
 
     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,25 +249,19 @@ with gr.Blocks() as demo:
                         "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,
+                        "periodic_p": 7,
+                        "onset_mask_width": 0,
                         "beat_mask_width": 0,
                         "beat_mask_downbeats": False,
                     },
                     "strong periodic variation": {
-                        "periodic_p": 17,
+                        "periodic_p": 13,
                         "onset_mask_width": 5,
                         "beat_mask_width": 0,
                         "beat_mask_downbeats": False,
                     },
                     "very strong periodic variation": {
-                        "periodic_p": 21,
+                        "periodic_p": 17,
                         "onset_mask_width": 5,
                         "beat_mask_width": 0,
                         "beat_mask_downbeats": False,
@@ -378,15 +269,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 +293,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 +318,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 +377,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,13 +404,6 @@ 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",
@@ -571,24 +428,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 +455,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 +468,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 +498,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/lora/lora.yml

@@ -4,16 +4,14 @@ $include:
 fine_tune: True
 
 train/AudioDataset.n_examples: 100000000
-val/AudioDataset.n_examples: 500
+val/AudioDataset.n_examples: 100
 
 
 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
+save_iters: [100000, 200000, 300000, 4000000, 500000]
 
 AdamW.lr: 0.0001
 

conf/vampnet.yml

@@ -32,7 +32,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

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

@@ -48,10 +48,11 @@ def fine_tune(audio_files_or_folders: List[str], name: str):
     }
 
     interface_conf = {
-        "Interface.coarse_ckpt": f"./runs/{name}/coarse/latest/vampnet/weights.pth",
+        "Interface.coarse_ckpt": f"./models/vampnet/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/vampnet/c2f.pth",
+        "Interface.coarse2fine_lora_ckpt": f"./runs/{name}/c2f/latest/lora.pth",
 
         "Interface.codec_ckpt": "./models/vampnet/codec.pth",
         "AudioLoader.sources": [audio_files_or_folders],

scripts/exp/train.py

@@ -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
@@ -29,9 +29,6 @@ from audiotools.ml.decorators import (
 
 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)
@@ -224,7 +221,7 @@ def train_loop(state: State, batch: dict, accel: Accelerator):
 
         dtype = torch.bfloat16 if accel.amp else None
         with accel.autocast(dtype=dtype):
-            z_hat = state.model(z_mask_latent)
+            z_hat = state.model(z_mask_latent, r)
 
         target = codebook_flatten(
             z[:, vn.n_conditioning_codebooks :, :],
@@ -289,7 +286,7 @@ def val_loop(state: State, batch: dict, accel: Accelerator):
 
     z_mask_latent = vn.embedding.from_codes(z_mask, state.codec)
 
-    z_hat = state.model(z_mask_latent)
+    z_hat = state.model(z_mask_latent, r)
 
     target = codebook_flatten(
         z[:, vn.n_conditioning_codebooks :, :],
@@ -408,19 +405,19 @@ def save_imputation(state, z, val_idx, writer):
 
     for i in range(len(val_idx)):
         imputed_noisy[i].cpu().write_audio_to_tb(
-            f"inpainted_prompt/{i}",
+            f"imputed_noisy/{i}",
             writer,
             step=state.tracker.step,
             plot_fn=None,
         )
         imputed[i].cpu().write_audio_to_tb(
-            f"inpainted_middle/{i}",
+            f"imputed/{i}",
             writer,
             step=state.tracker.step,
             plot_fn=None,
         )
         imputed_true[i].cpu().write_audio_to_tb(
-            f"reconstructed/{i}",
+            f"imputed_true/{i}",
             writer,
             step=state.tracker.step,
             plot_fn=None,
@@ -450,7 +447,7 @@ def save_samples(state: State, val_idx: int, writer: SummaryWriter):
 
     z_mask_latent = vn.embedding.from_codes(z_mask, state.codec)
 
-    z_hat = state.model(z_mask_latent)
+    z_hat = state.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)
@@ -469,7 +466,7 @@ def save_samples(state: State, val_idx: int, writer: SummaryWriter):
         }
         for k, v in audio_dict.items():
             v.cpu().write_audio_to_tb(
-                f"onestep/_{i}.r={r[i]:0.2f}/{k}",
+                f"samples/_{i}.r={r[i]:0.2f}/{k}",
                 writer,
                 step=state.tracker.step,
                 plot_fn=None,
@@ -488,6 +485,7 @@ def load(
     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:
@@ -500,7 +498,7 @@ 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():
@@ -513,14 +511,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 = torch.compile(VampNet()) if model is None else model
+    model = VampNet() if model is None else model
+
     model = accel.prepare_model(model)
 
     # assert accel.unwrap(model).n_codebooks == codec.quantizer.n_codebooks
@@ -604,7 +599,7 @@ def train(
         accel=accel, 
         tracker=tracker, 
         save_path=save_path)
-    print("initialized state.")
+
 
     train_dataloader = accel.prepare_dataloader(
         state.train_data,
@@ -619,15 +614,13 @@ def train(
         num_workers=num_workers,
         batch_size=batch_size,
         collate_fn=state.val_data.collate,
-        persistent_workers=num_workers > 0,
+        persistent_workers=True,
     )
-    print("initialized dataloader.")
 
     
 
     if fine_tune:
         lora.mark_only_lora_as_trainable(state.model)
-        print("marked only lora as trainable.")
 
     # Wrap the functions so that they neatly track in TensorBoard + progress bars
     # and only run when specific conditions are met.
@@ -642,7 +635,6 @@ def train(
     save_samples = when(lambda: accel.local_rank == 0)(save_samples)
     checkpoint = when(lambda: accel.local_rank == 0)(checkpoint)
 
-    print("starting training loop.")
     with tracker.live:
         for tracker.step, batch in enumerate(train_dataloader, start=tracker.step):
             train_loop(state, batch, accel)

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

scripts/utils/remove_quiet_files.py

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

scripts/utils/split_long_audio_file.py

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

scripts/utils/xeno-canto-dl.py

@@ -1,234 +0,0 @@
-from xenopy import Query
-
-
-SPECIES = [
-    "American Robin",
-    "Northern Cardinal",
-    "Mourning Dove",
-    "American Crow",
-    "Baltimore Oriole",
-    "Blue Jay",
-    "Eastern Bluebird",
-    "House Finch",
-    "American Goldfinch",
-    "House Sparrow",
-    "Song Sparrow",
-    "Tufted Titmouse",
-    "White-breasted Nuthatch",
-    "European Starling",
-    "American Redstart",
-    "Red-winged Blackbird",
-    "Brown-headed Cowbird",
-    "Common Grackle",
-    "Boat-tailed Grackle",
-    "Common Yellowthroat",
-    "Northern Mockingbird",
-    "Carolina Wren",
-    "Eastern Meadowlark",
-    "Chipping Sparrow",
-    "Tree Swallow",
-    "Barn Swallow",
-    "Cliff Swallow",
-    "Pine Siskin",
-    "Indigo Bunting",
-    "Eastern Towhee",
-    "Carolina Chickadee",
-    "Great Crested Flycatcher",
-    "Eastern Wood-Pewee",
-    "Ovenbird",
-    "Northern Flicker",
-    "Red-eyed Vireo",
-    "American Woodcock",
-    "Eastern Phoebe",
-    "Downy Woodpecker",
-    "Scarlet Tanager",
-    "Yellow Warbler",
-    "White-eyed Vireo",
-    "Common Loon",
-    "White-throated Sparrow",
-    "Yellow-throated Vireo",
-    "Great Blue Heron",
-    "Belted Kingfisher",
-    "Pied-billed Grebe",
-    "Wild Turkey",
-    "Wood Thrush",
-    "Rose-breasted Grosbeak",
-    "Field Sparrow",
-    "Hooded Warbler",
-    "Northern Parula",
-    "Chestnut-sided Warbler",
-    "Blue-winged Warbler",
-    "Red-bellied Woodpecker",
-    "Yellow-billed Cuckoo",
-    "Gray Catbird",
-    "Northern Saw-whet Owl",
-    "Osprey",
-    "Common Nighthawk",
-    "Broad-winged Hawk",
-    "Black-throated Green Warbler",
-    "Great Horned Owl",
-    "Common Raven",
-    "Barred Owl",
-    "Canada Warbler",
-    "Magnolia Warbler",
-    "Black-and-white Warbler",
-    "Eastern Kingbird",
-    "Swainson's Thrush",
-    "Worm-eating Warbler",
-    "Prairie Warbler",
-    "Baltimore Oriole",
-    "Black-throated Blue Warbler",
-    "Louisiana Waterthrush",
-    "Blackburnian Warbler",
-    "Black-capped Chickadee",
-    "Cerulean Warbler",
-    "Red-shouldered Hawk",
-    "Cooper's Hawk",
-    "Yellow-throated Warbler",
-    "Blue-headed Vireo",
-    "Blackpoll Warbler",
-    "Ruffed Grouse",
-    "Kentucky Warbler",
-    "Hermit Thrush",
-    "Cedar Waxwing",
-    "Eastern Screech-Owl",
-    "Northern Goshawk",
-    "Green Heron",
-    "Red-tailed Hawk",
-    "Black Vulture",
-    "Hairy Woodpecker",
-    "Golden-crowned Kinglet",
-    "Ruby-crowned Kinglet",
-    "Bicknell's Thrush",
-    "Blue-gray Gnatcatcher",
-    "Veery",
-    "Pileated Woodpecker",
-    "Purple Finch",
-    "White-crowned Sparrow",
-    "Snow Bunting",
-    "Pine Grosbeak",
-    "American Tree Sparrow",
-    "Dark-eyed Junco",
-    "Snowy Owl",
-    "White-winged Crossbill",
-    "Red Crossbill",
-    "Common Redpoll",
-    "Northern Shrike",
-    "Northern Harrier",
-    "Rough-legged Hawk",
-    "Long-eared Owl",
-    "Evening Grosbeak",
-    "Northern Pintail",
-    "American Black Duck",
-    "Mallard",
-    "Canvasback",
-    "Redhead",
-    "Ring-necked Duck",
-    "Greater Scaup",
-    "Lesser Scaup",
-    "Bufflehead",
-    "Common Goldeneye",
-    "Hooded Merganser",
-    "Common Merganser",
-    "Red-breasted Merganser",
-    "Ruddy Duck",
-    "Wood Duck",
-    "Gadwall",
-    "American Wigeon",
-    "Northern Shoveler",
-    "Green-winged Teal",
-    "Blue-winged Teal",
-    "Cinnamon Teal",
-    "Ringed Teal",
-    "Cape Teal",
-    "Northern Fulmar",
-    "Yellow-billed Loon",
-    "Red-throated Loon",
-    "Arctic Loon",
-    "Pacific Loon",
-    "Horned Grebe",
-    "Red-necked Grebe",
-    "Eared Grebe",
-    "Western Grebe",
-    "Clark's Grebe",
-    "Double-crested Cormorant",
-    "Pelagic Cormorant",
-    "Great Cormorant",
-    "American White Pelican",
-    "Brown Pelican",
-    "Brandt's Cormorant",
-    "Least Bittern",
-    "Great Egret",
-    "Snowy Egret",
-    "Little Blue Heron",
-    "Tricolored Heron",
-    "Reddish Egret",
-    "Black-crowned Night-Heron",
-    "Yellow-crowned Night-Heron",
-    "White Ibis",
-    "Glossy Ibis",
-    "Roseate Spoonbill",
-    "Wood Stork",
-    "Black-bellied Whistling-Duck",
-    "Fulvous Whistling-Duck",
-    "Greater White-fronted Goose",
-    "Snow Goose",
-    "Ross's Goose",
-    "Canada Goose",
-    "Brant",
-    "Mute Swan",
-    "Tundra Swan",
-    "Whooper Swan",
-    "Sandhill Crane",
-    "Black-necked Stilt",
-    "American Avocet",
-    "Northern Jacana",
-    "Greater Yellowlegs",
-    "Lesser Yellowlegs",
-    "Willet",
-    "Spotted Sandpiper",
-    "Upland Sandpiper",
-    "Whimbrel",
-    "Long-billed Curlew",
-    "Marbled Godwit",
-    "Ruddy Turnstone",
-    "Red Knot",
-    "Sanderling",
-    "Semipalmated Sandpiper",
-    "Western Sandpiper",
-    "Least Sandpiper",
-    "White-rumped Sandpiper",
-    "Baird's Sandpiper",
-    "Pectoral Sandpiper",
-    "Dunlin",
-    "Buff-breasted Sandpiper",
-    "Short-billed Dowitcher",
-    "Long-billed Dowitcher",
-    "Common Snipe",
-    "American Woodcock",
-    "Wilson's Phalarope",
-    "Red-necked Phalarope",
-    "Red Phalarope"
-]
-
-from pathlib import Path
-
-def remove_spaces(s):
-    return s.replace(" ", "")
-
-for species in SPECIES: 
-    if Path("/media/CHONK/hugo/xeno-canto-full/" + remove_spaces(species)).exists():
-        continue
-    try:
-        q = Query(
-            name=species, q="A", length="10-30", 
-            )
-
-        # retrieve metadata
-        metafiles = q.retrieve_meta(verbose=True)
-        # retrieve recordings
-        q.retrieve_recordings(multiprocess=True, nproc=10, attempts=10, outdir="/media/CHONK/hugo/xeno-canto-full/")
-
-    except:
-        print("Failed to download " + species)
-        continue

setup.py

@@ -28,13 +28,12 @@ setup(
     install_requires=[
         "torch",
         "argbind>=0.3.2",
-        "numpy==1.23",
+        "numpy==1.22",
         "wavebeat @ git+https://github.com/hugofloresgarcia/wavebeat",
         "lac @ git+https://github.com/hugofloresgarcia/lac.git",
         "descript-audiotools @ git+https://github.com/descriptinc/[email protected]",
         "gradio", 
+        "tensorboardX",
         "loralib",
-        "torch_pitch_shift",
-        "madmom",
     ],
 )

vampnet/interface.py

@@ -65,7 +65,7 @@ class Interface(torch.nn.Module):
     ):
         super().__init__()
         assert codec_ckpt is not None, "must provide a codec checkpoint"
-        self.codec = DAC.load(Path(codec_ckpt))
+        self.codec = DAC.load(codec_ckpt)
         self.codec.eval()
         self.codec.to(device)
 
@@ -120,16 +120,17 @@ class Interface(torch.nn.Module):
             if coarse_ckpt is not None:
                 self.coarse.to("cpu")
                 state_dict = torch.load(coarse_ckpt, map_location="cpu")
-                print(f"loading coarse from {coarse_ckpt}")
+
                 self.coarse.load_state_dict(state_dict, strict=False)
                 self.coarse.to(self.device)
             if c2f_ckpt is not None:
                 self.c2f.to("cpu")
                 state_dict = torch.load(c2f_ckpt, map_location="cpu")
-                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):
@@ -193,8 +194,8 @@ class Interface(torch.nn.Module):
 
     def make_beat_mask(self, 
             signal: AudioSignal, 
-            before_beat_s: float = 0.0,
-            after_beat_s: float = 0.02,
+            before_beat_s: float = 0.1,
+            after_beat_s: float = 0.1,
             mask_downbeats: bool = True,
             mask_upbeats: bool = True,
             downbeat_downsample_factor: int = None,

vampnet/mask.py

@@ -191,47 +191,29 @@ def onset_mask(
     width: int = 1
 ):
     import librosa
-    import madmom
-    from madmom.features.onsets import RNNOnsetProcessor, OnsetPeakPickingProcessor
-    import tempfile
-    import numpy as np 
-
-    with tempfile.NamedTemporaryFile(suffix='.wav') as f:
-        sig = sig.clone()
-        sig.write(f.name)
-
-        proc = RNNOnsetProcessor(online=False)
-        onsetproc = OnsetPeakPickingProcessor(threshold=0.3,
-                                              fps=sig.sample_rate/interface.codec.hop_length)
-        
-        act = proc(f.name)
-        onset_times = onsetproc(act)
-
-        # convert to indices for z array
-        onset_indices = librosa.time_to_frames(onset_times, sr=sig.sample_rate, hop_length=interface.codec.hop_length)
-
-        if onset_indices.shape[0] == 0:
-            mask = empty_mask(z)   
-            print(f"no onsets found, returning empty mask")
-        else: 
-            torch.set_printoptions(threshold=1000)
-            print("onset indices: ", onset_indices)
-            print("onset times: ", onset_times)
-
-            # create a mask, set onset 
-            mask = torch.ones_like(z)
-            n_timesteps = z.shape[-1]
-
-            for onset_index in onset_indices:
-                onset_index = min(onset_index, n_timesteps - 1)
-                onset_index = max(onset_index, 0)
-                mask[:, :, onset_index - width:onset_index + width] = 0.0
-
-            print(mask)
+
+    onset_indices = librosa.onset.onset_detect(
+        y=sig.clone().to_mono().samples.cpu().numpy()[0, 0], 
+        sr=sig.sample_rate,
+        hop_length=interface.codec.hop_length, 
+        backtrack=True,
+    )
+
+    # create a mask, set onset 
+    mask = torch.ones_like(z)
+    n_timesteps = z.shape[-1]
+
+    for onset_index in onset_indices:
+        onset_index = min(onset_index, n_timesteps - 1)
+        onset_index = max(onset_index, 0)
+        mask[:, :, onset_index - width:onset_index + width] = 0.0
+
+    print(mask)
     
     return mask
 
 
 
 if __name__ == "__main__":
-    pass
+    torch.set_printoptions(threshold=10000)
+

vampnet/modules/transformer.py

@@ -410,9 +410,7 @@ 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):
         """Computes a full transformer stack
         Parameters
         ----------
@@ -439,8 +437,6 @@ class TransformerStack(nn.Module):
         encoder_decoder_position_bias = None
 
         # Compute transformer layers
-        if return_activations:
-            activations = []
         for layer in self.layers:
             x, position_bias, encoder_decoder_position_bias = layer(
                 x=x,
@@ -451,15 +447,8 @@ class TransformerStack(nn.Module):
                 position_bias=position_bias,
                 encoder_decoder_position_bias=encoder_decoder_position_bias,
             )
-            if return_activations:
-                activations.append(x.detach())
 
-    
-        out = self.norm(x) if self.norm is not None else x
-        if return_activations:
-            return out, torch.stack(activations)
-        else:
-            return out
+        return self.norm(x) if self.norm is not None else x
 
 
 class VampNet(at.ml.BaseModel):
@@ -467,7 +456,7 @@ class VampNet(at.ml.BaseModel):
         self,
         n_heads: int = 20,
         n_layers: int = 16,
-        r_cond_dim: int = 0,
+        r_cond_dim: int = 64,
         n_codebooks: int = 9,
         n_conditioning_codebooks: int = 0,
         latent_dim: int = 8,
@@ -478,7 +467,6 @@ class VampNet(at.ml.BaseModel):
         dropout: float = 0.1
     ):
         super().__init__()
-        assert r_cond_dim == 0, f"r_cond_dim must be 0 (not supported), but got {r_cond_dim}"
         self.n_heads = n_heads
         self.n_layers = n_layers
         self.r_cond_dim = r_cond_dim
@@ -525,25 +513,21 @@ class VampNet(at.ml.BaseModel):
             ),
         )
 
-    def forward(self, x, return_activations: bool = False):
+    def forward(self, x, cond):
         x = self.embedding(x)
         x_mask = torch.ones_like(x, dtype=torch.bool)[:, :1, :].squeeze(1)
 
-        x = rearrange(x, "b d n -> b n d")
-        out = self.transformer(x=x, x_mask=x_mask, return_activations=return_activations)
-        if return_activations:
-            out, activations = out
+        cond = self.r_embed(cond)
 
+        x = rearrange(x, "b d n -> b n d")
+        out = self.transformer(x=x, x_mask=x_mask, cond=cond)
         out = rearrange(out, "b n d -> b d n")
 
-        out = self.classifier(out, None) # no cond here!
+        out = self.classifier(out, cond)
 
         out = rearrange(out, "b (p c) t -> b p (t c)", c=self.n_predict_codebooks)
 
-        if return_activations:
-            return out, activations
-        else:
-            return out
+        return out
     
     def r_embed(self, r, max_positions=10000):
         if self.r_cond_dim > 0:
@@ -594,23 +578,22 @@ class VampNet(at.ml.BaseModel):
         self,
         codec,
         time_steps: int = 300,
-        sampling_steps: int = 36,
+        sampling_steps: int = 24,
         start_tokens: Optional[torch.Tensor] = None,
-        sampling_temperature: float = 1.0,
         mask: Optional[torch.Tensor] = None,
-        mask_temperature: float = 10.5,
+        temperature: Union[float, Tuple[float, float]] = 2.5,
         typical_filtering=False,
         typical_mass=0.2,
         typical_min_tokens=1,
-        top_p=None,
         return_signal=True,
-        seed: int = None, 
-        sample_cutoff: float = 1.0,
     ):
-        if seed is not None:
-            at.util.seed(seed)
         logging.debug(f"beginning generation with {sampling_steps} steps")
 
+        #####################
+        # resolve temperature #
+        #####################
+        assert isinstance(temperature, float)
+        logging.debug(f"temperature: {temperature}")
 
 
         ##################### 
@@ -662,6 +645,9 @@ class VampNet(at.ml.BaseModel):
         for i in range(sampling_steps):
             logging.debug(f"step {i} of {sampling_steps}")
 
+            # our current temperature
+            logging.debug(f"temperature: {temperature}")
+
             # our current schedule step
             r = scalar_to_batch_tensor(
                 (i + 1) / sampling_steps, 
@@ -676,24 +662,41 @@ class VampNet(at.ml.BaseModel):
 
             # infer from latents
             # NOTE: this collapses the codebook dimension into the sequence dimension
-            logits = self.forward(latents) # b, prob, seq
+            logits = self.forward(latents, r) # b, prob, seq
             logits = logits.permute(0, 2, 1)  # b, seq, prob
-            b = logits.shape[0]
+            if typical_filtering:
+                typical_filter(logits, 
+                               typical_mass=typical_mass, 
+                               typical_min_tokens=typical_min_tokens
+                )
+
 
             logging.debug(f"permuted logits with shape: {logits.shape}")
 
-            sampled_z, selected_probs = sample_from_logits(
-                logits, sample=(
-                   (i / sampling_steps) <= sample_cutoff
-                ), 
-                temperature=sampling_temperature,
-                typical_filtering=typical_filtering, typical_mass=typical_mass,
-                typical_min_tokens=typical_min_tokens,
-                top_k=None, top_p=top_p, return_probs=True,
-            )
 
+            # logits2probs
+            probs = torch.softmax(logits, dim=-1)
+            logging.debug(f"computed probs with shape: {probs.shape}")
+
+
+            # sample from logits with multinomial sampling
+            b = probs.shape[0]
+            probs = rearrange(probs, "b seq prob -> (b seq) prob")
+
+            sampled_z =  torch.multinomial(probs, 1).squeeze(-1)
+
+            sampled_z = rearrange(sampled_z, "(b seq)-> b seq", b=b)
+            probs = rearrange(probs, "(b seq) prob -> b seq prob", b=b)
             logging.debug(f"sampled z with shape: {sampled_z.shape}")
 
+            # get the confidences: which tokens did we sample? 
+            selected_probs = (
+                torch.take_along_dim(
+                    probs, sampled_z.long().unsqueeze(-1), 
+                    dim=-1
+                ).squeeze(-1)
+            )
+
             # 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
@@ -730,7 +733,7 @@ class VampNet(at.ml.BaseModel):
 
             # get our new mask
             mask = mask_by_random_topk(
-                num_to_mask, selected_probs, mask_temperature * (1-r)
+                num_to_mask, selected_probs, temperature * (1-r)
             )  
 
             # update the mask
@@ -763,97 +766,8 @@ class VampNet(at.ml.BaseModel):
         else:
             return sampled_z
 
-def sample_from_logits(
-        logits, 
-        sample: bool = True,
-        temperature: float = 1.0,
-        top_k: int = None,
-        top_p: float = None,
-        typical_filtering: bool = False,
-        typical_mass: float = 0.2,
-        typical_min_tokens: int = 1,
-        return_probs: bool = False
-    ):
-    """Convenience function to sample from a categorial distribution with input as
-    unnormalized logits.
-
-    Parameters
-    ----------
-    logits : Tensor[..., vocab_size]
-    config: SamplingConfig
-        The set of hyperparameters to be used for sampling
-        sample : bool, optional
-            Whether to perform multinomial sampling, by default True
-        temperature : float, optional
-            Scaling parameter when multinomial samping, by default 1.0
-        top_k : int, optional
-            Restricts sampling to only `top_k` values acc. to probability,
-            by default None
-        top_p : float, optional
-            Restricts sampling to only those values with cumulative
-            probability = `top_p`, by default None
-
-    Returns
-    -------
-    Tensor[...]
-        Sampled tokens
-    """
-    shp = logits.shape[:-1]
-
-    if typical_filtering:
-        typical_filter(logits, 
-                        typical_mass=typical_mass, 
-                        typical_min_tokens=typical_min_tokens
-        )
-
-    # Apply top_k sampling
-    if top_k is not None:
-        v, _ = logits.topk(top_k)
-        logits[logits < v[..., [-1]]] = -float("inf")
-
-    # Apply top_p (nucleus) sampling
-    if top_p is not None and top_p < 1.0:
-        v, sorted_indices = logits.sort(descending=True)
-        cumulative_probs = v.softmax(dim=-1).cumsum(dim=-1)
-
-        sorted_indices_to_remove = cumulative_probs > top_p
-        # Right shift indices_to_remove to keep 1st token over threshold
-        sorted_indices_to_remove = F.pad(sorted_indices_to_remove, (1, 0), value=False)[
-            ..., :-1
-        ]
-
-        # Compute indices_to_remove in unsorted array
-        indices_to_remove = sorted_indices_to_remove.scatter(
-            -1, sorted_indices, sorted_indices_to_remove
-        )
-
-        logits[indices_to_remove] = -float("inf")
-
-    # Perform multinomial sampling after normalizing logits
-    probs = (
-        F.softmax(logits / temperature, dim=-1)
-        if temperature > 0
-        else logits.softmax(dim=-1)
-    )
-    token = (
-        probs.view(-1, probs.size(-1)).multinomial(1).squeeze(1).view(*shp)
-        if sample
-        else logits.argmax(-1)
-    )
-
-    if return_probs:
-        token_probs = probs.take_along_dim(token.unsqueeze(-1), dim=-1).squeeze(-1)
-        return token, token_probs
-    else:
-        return token
-    
-
 
-def mask_by_random_topk(
-        num_to_mask: int, 
-        probs: torch.Tensor, 
-        temperature: float = 1.0, 
-    ):
+def mask_by_random_topk(num_to_mask: int, probs: torch.Tensor, temperature: float = 1.0):
     """
     Args:
         num_to_mask (int): number of tokens to mask
@@ -866,8 +780,7 @@ def mask_by_random_topk(
     logging.debug(f"temperature: {temperature}")
     logging.debug("")
 
-    noise = gumbel_noise_like(probs)
-    confidence = torch.log(probs) + temperature * noise
+    confidence = torch.log(probs) + temperature * gumbel_noise_like(probs)
     logging.debug(f"confidence shape: {confidence.shape}")
 
     sorted_confidence, sorted_idx = confidence.sort(dim=-1)
@@ -937,7 +850,7 @@ if __name__ == "__main__":
         z_mask_latent = torch.rand(
             batch_size, model.latent_dim * model.n_codebooks, seq_len
         ).to(device)
-        z_hat = model(z_mask_latent)
+        z_hat = model(z_mask_latent, r)
 
         pred = z_hat.argmax(dim=1)
         pred = model.embedding.unflatten(pred, n_codebooks=model.n_predict_codebooks)