demo cleanup, onset masks, pitch shifting

conf/interface/xeno-canto.yml

@@ -0,0 +1,11 @@
+Interface.coarse_ckpt: ./runs/xeno-canto-2/coarse/best/vampnet/weights.pth
+Interface.coarse2fine_ckpt: ./runs/xeno-canto-2/c2f/best/vampnet/weights.pth
+Interface.codec_ckpt: ./models/spotdl/codec.pth
+Interface.coarse_chunk_size_s: 10
+Interface.coarse2fine_chunk_size_s: 3
+# Interface.wavebeat_ckpt: ./models/wavebeat.pth
+
+
+AudioLoader.sources:
+  - /media/CHONK/hugo/xeno-canto-2
+  - /media/CHONK/hugo/xeno-canto-2

conf/lora/lora.yml

@@ -8,9 +8,11 @@ train/AudioDataset.n_examples: 10000000
 val/AudioDataset.n_examples: 10
 
 
-NoamScheduler.warmup: 400
+NoamScheduler.warmup: 500
 
+batch_size: 7
+num_workers: 7
 epoch_length: 100
-save_audio_epochs: 2
+save_audio_epochs: 4
 
 AdamW.lr: 0.0001

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.yml

@@ -25,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

demo.py

@@ -62,10 +62,13 @@ def load_random_audio():
     return sig.path_to_file
 
 
-def vamp(data):
+def _vamp(data, return_mask=False):
+    print(data)
     print(data[input_audio])
     sig = at.AudioSignal(data[input_audio])
 
+    # TODO: random pitch shift of segments in the signal to prompt! window size should be a parameter, pitch shift width should be a parameter
+
     z = interface.encode(sig)
 
     ncc = data[n_conditioning_codebooks]
@@ -87,6 +90,11 @@ def vamp(data):
             random_roll=True
         )
     )
+    if data[onset_mask_width] > 0:
+        mask = pmask.mask_or(
+            mask, pmask.onset_mask(sig, z, interface, width=data[onset_mask_width])
+        )
+    # these should be the last two mask ops
     mask = pmask.dropout(mask, data[dropout])
     mask = pmask.codebook_unmask(mask, ncc)
 
@@ -103,9 +111,6 @@ def vamp(data):
     if use_coarse2fine: 
         zv = interface.coarse_to_fine(zv)
 
-
-    mask = interface.to_signal(mask_z).cpu()
-
     sig = interface.to_signal(zv).cpu()
     print("done")
 
@@ -113,8 +118,19 @@ def vamp(data):
     out_dir.mkdir()
 
     sig.write(out_dir / "output.wav")
-    mask.write(out_dir / "mask.wav")
-    return sig.path_to_file, mask.path_to_file
+
+    if return_mask:
+        mask = interface.to_signal(mask_z).cpu()
+        mask.write(out_dir / "mask.wav")
+        return sig.path_to_file, mask.path_to_file
+    else:
+        return sig.path_to_file
+
+def vamp(data):
+    return _vamp(data, return_mask=True)
+
+def api_vamp(data):
+    return _vamp(data, return_mask=False)
         
 def save_vamp(data):
     out_dir = OUT_DIR / "saved" / str(uuid.uuid4())
@@ -198,6 +214,14 @@ with gr.Blocks() as demo:
         # mask settings
         with gr.Column():
 
+            input_pitch_shift = gr.Slider(
+                label="input pitch shift (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,
@@ -220,6 +244,14 @@ with gr.Blocks() as demo:
                 value=1,
             )
 
+            onset_mask_width = gr.Slider(
+                label="onset mask width (steps, 1 step ~= 10milliseconds)",
+                minimum=0,
+                maximum=20,
+                step=1,
+                value=0,
+            )
+
             with gr.Accordion("extras ", open=False):
                 n_conditioning_codebooks = gr.Number(
                     label="number of conditioning codebooks. probably 0", 
@@ -322,11 +354,9 @@ with gr.Blocks() as demo:
             )
             
             thank_you = gr.Markdown("")
-  
-    # connect widgets
-    vamp_button.click(
-        fn=vamp,
-        inputs={
+
+
+    _inputs = {
             input_audio, 
             num_steps,
             init_temp, final_temp,
@@ -336,27 +366,29 @@ with gr.Blocks() as demo:
             n_conditioning_codebooks, 
             dropout,
             use_coarse2fine, 
-            stretch_factor
-        },
+            stretch_factor, 
+            onset_mask_width, 
+            input_pitch_shift
+        }
+  
+    # connect widgets
+    vamp_button.click(
+        fn=vamp,
+        inputs=_inputs,
         outputs=[output_audio, audio_mask], 
+    )
+
+    api_vamp_button = gr.Button("api vamp")
+    api_vamp_button.click(
+        fn=api_vamp,
+        inputs=_inputs, 
+        outputs=[output_audio], 
         api_name="vamp"
     )
 
     save_button.click(
         fn=save_vamp,
-        inputs={
-            input_audio, 
-            num_steps,
-            init_temp, final_temp,
-            prefix_s, suffix_s, 
-            rand_mask_intensity, 
-            periodic_p, periodic_w,
-            n_conditioning_codebooks,
-            dropout,
-            use_coarse2fine, 
-            stretch_factor, 
-            notes_text
-        },
+        inputs=_inputs | {notes_text},
         outputs=[thank_you, download_file]
     )
 

scripts/exp/train.py

@@ -62,6 +62,7 @@ IGNORE_INDEX = -100
 def build_transform():
     transform = transforms.Compose(
         tfm.VolumeNorm(("const", -24)),
+        # tfm.PitchShift(),
         tfm.RescaleAudio(),
     )
     return transform

vampnet/mask.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 import torch
+from audiotools import AudioSignal
 
 from .util import scalar_to_batch_tensor
 
@@ -150,7 +151,9 @@ def dropout(
     mask: torch.Tensor,
     p: float,
 ):
-    return torch.bernoulli((torch.ones_like(mask) * (1-p)).float()).long() * mask
+    # negate the mask (we want the 0s to be 1s, since we want to drop the prompt, not the mask)
+    mask = (~(mask.bool())).long()
+    return torch.nn.functional.dropout(mask.float(), p=p, training=True).long().bool().long()
 
 def mask_or(
     mask1: torch.Tensor, 
@@ -166,7 +169,6 @@ def mask_or(
 def time_stretch_mask(
     x: torch.Tensor, 
     stretch_factor: int,
-    mask_token: int
 ):
     assert stretch_factor >= 1, "stretch factor must be >= 1"
     c_seq_len = x.shape[-1]
@@ -176,7 +178,35 @@ def time_stretch_mask(
     x = x[:, :, :c_seq_len]
 
     mask = periodic_mask(x, stretch_factor, width=1)
-    return apply_mask(x, mask, mask_token)
+    return mask
+
+def onset_mask(
+    sig: AudioSignal, 
+    z: torch.Tensor,
+    interface,
+    width: int = 1
+):
+    import librosa
+
+    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
+    )
+
+    # 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__":

vampnet/modules/transformer.py

@@ -62,8 +62,8 @@ class FeedForward(nn.Module):
     ):
         super().__init__()
         factor = 2 if activation == "geglu" else 1
-        self.w_1 = nn.Linear(d_model, d_model * 4, bias=False)
-        self.w_2 = nn.Linear(d_model * 4 // factor, d_model, bias=False)
+        self.w_1 = lora.Linear(d_model, d_model * 4, bias=False, r=LORA_R)
+        self.w_2 = lora.Linear(d_model * 4 // factor, d_model, bias=False, r=LORA_R)
         self.drop = nn.Dropout(dropout)
         self.act = get_activation(activation)()
 
@@ -109,7 +109,7 @@ class MultiHeadRelativeAttention(nn.Module):
         self.w_vs = lora.Linear(d_model, d_model, bias=False, r=LORA_R)
 
         # Create linear final output projection
-        self.fc = nn.Linear(d_model, d_model, bias=False)
+        self.fc = lora.Linear(d_model, d_model, bias=False, r=LORA_R)
 
         # Dropout for attention output weights
         self.dropout = nn.Dropout(dropout)