fix dropout bug for masks, refactor interfaces, add finetune setup script

conf/generated/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/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/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/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/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/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/interface/spotdl.yml

@@ -7,5 +7,6 @@ Interface.coarse2fine_chunk_size_s: 3
 
 
 AudioLoader.sources:
-  - /data/spotdl/audio/val
-  - /data/spotdl/audio/test
+  # - /media/CHONK/hugo/spotdl/subsets/jazz-blues
+  - /media/CHONK/null
+

demo.py

@@ -63,9 +63,11 @@ def load_random_audio():
 
 
 def _vamp(data, return_mask=False):
-    print(data)
-    print(data[input_audio])
+    out_dir = OUT_DIR / str(uuid.uuid4())
+    out_dir.mkdir()
     sig = at.AudioSignal(data[input_audio])
+    #pitch shift input
+    sig = sig.shift_pitch(data[input_pitch_shift])
 
     # TODO: random pitch shift of segments in the signal to prompt! window size should be a parameter, pitch shift width should be a parameter
 
@@ -98,7 +100,9 @@ def _vamp(data, return_mask=False):
     mask = pmask.dropout(mask, data[dropout])
     mask = pmask.codebook_unmask(mask, ncc)
 
-    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]}")
+    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[init_temp]}, final temp {data[final_temp]}, use coarse2fine {data[use_coarse2fine]}")
+    # save the mask as a txt file
+    np.savetxt(out_dir / "mask.txt", mask[:,0,:].long().cpu().numpy())
 
     zv, mask_z = interface.coarse_vamp(
         z, 
@@ -114,8 +118,7 @@ def _vamp(data, return_mask=False):
     sig = interface.to_signal(zv).cpu()
     print("done")
 
-    out_dir = OUT_DIR / str(uuid.uuid4())
-    out_dir.mkdir()
+    
 
     sig.write(out_dir / "output.wav")
 
@@ -136,13 +139,13 @@ def save_vamp(data):
     out_dir = OUT_DIR / "saved" / str(uuid.uuid4())
     out_dir.mkdir(parents=True, exist_ok=True)
 
-    sig_in = at.AudioSignal(input_audio)
-    sig_out = at.AudioSignal(output_audio)
+    sig_in = at.AudioSignal(data[input_audio])
+    sig_out = at.AudioSignal(data[output_audio])
 
     sig_in.write(out_dir / "input.wav")
     sig_out.write(out_dir / "output.wav")
     
-    data = {
+    _data = {
         "init_temp": data[init_temp],
         "final_temp": data[final_temp],
         "prefix_s": data[prefix_s],
@@ -159,7 +162,7 @@ def save_vamp(data):
 
     # save with yaml
     with open(out_dir / "data.yaml", "w") as f:
-        yaml.dump(data, f)
+        yaml.dump(_data, f)
 
     import zipfile
     zip_path = out_dir.with_suffix(".zip")
@@ -321,6 +324,8 @@ with gr.Blocks() as demo:
                 type="filepath"
             )
 
+            use_as_input_button = gr.Button("use as input")
+
         
         # with gr.Column():
         #     with gr.Accordion(label="beat unmask (how much time around the beat should be hinted?)"):
@@ -386,9 +391,15 @@ with gr.Blocks() as demo:
         api_name="vamp"
     )
 
+    use_as_input_button.click(
+        fn=lambda x: x,
+        inputs=[output_audio],
+        outputs=[input_audio]
+    )
+
     save_button.click(
         fn=save_vamp,
-        inputs=_inputs | {notes_text},
+        inputs=_inputs | {notes_text, output_audio},
         outputs=[thank_you, download_file]
     )
 

scripts/exp/fine_tune.py

@@ -0,0 +1,85 @@
+import argbind
+from pathlib import Path
+import yaml
+
+
+
+
+"""example output: (yaml)
+
+"""
+
+@argbind.bind(without_prefix=True, positional=True)
+def fine_tune(audio_file_or_folder: str, name: str):
+
+    conf_dir = Path("conf")
+    assert conf_dir.exists(), "conf directory not found. are you in the vampnet directory?"
+
+    conf_dir = conf_dir / "generated"
+    conf_dir.mkdir(exist_ok=True)
+
+    finetune_dir = conf_dir / name
+    finetune_dir.mkdir(exist_ok=True)
+
+    finetune_c2f_conf = {
+        "$include": ["conf/lora/lora.yml"],
+        "fine_tune": True,
+        "train/AudioLoader.sources": [audio_file_or_folder],
+        "val/AudioLoader.sources": [audio_file_or_folder],
+        "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,
+        "save_path": f"./runs/{name}/c2f",
+    }
+
+    finetune_coarse_conf = {
+        "$include": ["conf/lora/lora.yml"],
+        "fine_tune": True,
+        "train/AudioLoader.sources": [audio_file_or_folder],
+        "val/AudioLoader.sources": [audio_file_or_folder],
+        "save_path": f"./runs/{name}/coarse",
+    }
+
+    interface_conf = {
+        "Interface.coarse_ckpt": f"./runs/{name}/coarse/best/vampnet/weights.pth",
+        "Interface.coarse2fine_ckpt": f"./runs/{name}/c2f/best/vampnet/weights.pth",
+        "Interface.codec_ckpt": "./models/spotdl/codec.pth",
+        "AudioLoader.sources": [audio_file_or_folder],
+    }
+
+    # save the confs
+    with open(finetune_dir / "c2f.yml", "w") as f:
+        yaml.dump(finetune_c2f_conf, f)
+
+    with open(finetune_dir / "coarse.yml", "w") as f:
+        yaml.dump(finetune_coarse_conf, f)
+    
+    with open(finetune_dir / "interface.yml", "w") as f: 
+        yaml.dump(interface_conf, f)
+
+    # copy the starter weights to the save paths
+    import shutil
+
+    def pmkdir(path):
+        Path(path).parent.mkdir(exist_ok=True, parents=True)
+        return path
+
+    shutil.copy("./models/spotdl/c2f.pth", pmkdir(f"./runs/{name}/c2f/starter/vampnet/weights.pth"))
+    shutil.copy("./models/spotdl/coarse.pth", pmkdir(f"./runs/{name}/coarse/starter/vampnet/weights.pth"))
+    
+
+    print(f"generated confs in {finetune_dir}. run training jobs with `python scripts/exp/train.py --args.load {finetune_dir}/<c2f/coarse>.yml --resume --load_weights --tag starter` ")
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+
+    with argbind.scope(args):
+        fine_tune()
+
+
+
+    

vampnet/mask.py

@@ -151,9 +151,13 @@ def dropout(
     mask: torch.Tensor,
     p: float,
 ):
-    # 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()
+    assert 0 <= p <= 1, "p must be between 0 and 1"
+    assert mask.max() <= 1, "mask must be binary"
+    assert mask.min() >= 0, "mask must be binary"
+    mask = (~mask.bool()).float()
+    mask = torch.bernoulli(mask * (1 - p))
+    mask = ~mask.round().bool()
+    return mask.long()
 
 def mask_or(
     mask1: torch.Tensor, 
@@ -191,7 +195,8 @@ def onset_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
+        hop_length=interface.codec.hop_length, 
+        backtrack=True,
     )
 
     # create a mask, set onset