gooood outputs

.gitignore

@@ -174,3 +174,6 @@ runs-archive
 lyrebird-audiotools
 lyrebird-audio-codec
 samples-*/**
+
+gradio-outputs/
+models/

conf/interface-jazzpop-exp.yml

@@ -1,9 +0,0 @@
-Interface.coarse_ckpt: /runs/jazzpop-coarse-1m-steps.pth
-Interface.coarse2fine_ckpt: /runs/jazzpop-c2f.pth
-Interface.codec_ckpt: /runs/codec-ckpt/codec.pth
-Interface.coarse_chunk_size_s: 5
-Interface.coarse2fine_chunk_size_s: 3
-
-AudioLoader.sources:
-  - /data/spotdl/audio/val
-  - /data/spotdl/audio/test

conf/interface-jazzpop.yml

@@ -0,0 +1,10 @@
+Interface.coarse_ckpt: ./models/jazzpop/coarse.pth
+Interface.coarse2fine_ckpt: ./models/jazzpop/c2f.pth
+Interface.codec_ckpt: ./models/jazzpop/codec.pth
+Interface.coarse_chunk_size_s: 5
+Interface.coarse2fine_chunk_size_s: 3
+Interface.wavebeat_ckpt: ./models/wavebeat.pth
+
+AudioLoader.sources:
+  - /data/spotdl-jazzpop/audio/val
+  - /data/spotdl-jazzpop/audio/test

conf/interface-spotdl.yml

@@ -0,0 +1,11 @@
+Interface.coarse_ckpt: ./models/spotdl/coarse.pth
+Interface.coarse2fine_ckpt: ./models/spotdl/c2f.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:
+  - /data/spotdl/audio/val
+  - /data/spotdl/audio/test

demo.py

@@ -1,6 +1,8 @@
 from pathlib import Path
 from typing import Tuple
 import yaml
+import tempfile
+import uuid
 
 import numpy as np
 import audiotools as at
@@ -9,13 +11,15 @@ import argbind
 import gradio as gr
 from vampnet.interface import Interface
 
-conf = yaml.safe_load(Path("conf/interface-jazzpop-exp.yml").read_text())
-
 Interface = argbind.bind(Interface)
 AudioLoader = argbind.bind(at.data.datasets.AudioLoader)
+
+conf = argbind.parse_args()
+
 with argbind.scope(conf):
     interface = Interface()
     loader = AudioLoader()
+    print(f"interface device is {interface.device}")
 
 dataset = at.data.datasets.AudioDataset(
     loader,
@@ -26,6 +30,10 @@ dataset = at.data.datasets.AudioDataset(
 )
 
 
+OUT_DIR = Path("gradio-outputs")
+OUT_DIR.mkdir(exist_ok=True, parents=True)
+
+
 def load_audio(file):
     print(file)
     filepath = file.name
@@ -35,87 +43,207 @@ def load_audio(file):
     )
     sig = interface.preprocess(sig)
 
-    audio = sig.samples.numpy()[0]
-    sr = sig.sample_rate
-    return sr, audio.T
+    out_dir = OUT_DIR / "tmp" / str(uuid.uuid4())
+    out_dir.mkdir(parents=True, exist_ok=True)
+    sig.write(out_dir / "input.wav")
+    return sig.path_to_file
 
 def load_random_audio():
     index = np.random.randint(0, len(dataset))
     sig = dataset[index]["signal"]
     sig = interface.preprocess(sig)
 
-    audio = sig.samples.numpy()[0]
-    sr = sig.sample_rate
-    return sr, audio.T
+    out_dir = OUT_DIR / "tmp" / str(uuid.uuid4())
+    out_dir.mkdir(parents=True, exist_ok=True)
+    sig.write(out_dir / "input.wav")
+    return sig.path_to_file
 
 
 def vamp(
-    input_audio, prefix_s, suffix_s, rand_mask_intensity,
+    input_audio, init_temp, final_temp, 
+    prefix_s, suffix_s, rand_mask_intensity,
     mask_periodic_amt, beat_unmask_dur,
     mask_dwn_chk, dwn_factor,
     mask_up_chk, up_factor, 
-    num_vamps, mode
+    num_vamps, mode, use_beats, num_steps
 ):
-    try:
+    # try:
         print(input_audio)
 
-        sig = at.AudioSignal(
-            input_audio[1], 
-            sample_rate=input_audio[0]
-        )
+        sig = at.AudioSignal(input_audio.name)
         
-        if beat_unmask_dur > 0.0:
+        if beat_unmask_dur > 0.0 and use_beats:
             beat_mask = interface.make_beat_mask(
-                sig, 
-                before_beat_s=0.01,
+                sig,
+                before_beat_s=0.0,
                 after_beat_s=beat_unmask_dur,
                 mask_downbeats=mask_dwn_chk,
                 mask_upbeats=mask_up_chk,
-                downbeat_downsample_factor=dwn_factor, 
-                beat_downsample_factor=up_factor,
+                downbeat_downsample_factor=dwn_factor if dwn_factor > 0 else None, 
+                beat_downsample_factor=up_factor if up_factor > 0 else None,
                 dropout=0.7, 
                 invert=True
             )
+            print(beat_mask)
         else:
             beat_mask = None
 
         if mode == "standard": 
-            zv = interface.coarse_vamp_v2(
+            print(f"running standard vampnet with {num_vamps} vamps")
+            zv, mask_z = interface.coarse_vamp_v2(
                 sig, 
+                sampling_steps=num_steps,
+                temperature=(init_temp, final_temp),
                 prefix_dur_s=prefix_s,
                 suffix_dur_s=suffix_s,
                 num_vamps=num_vamps,
                 downsample_factor=mask_periodic_amt,
                 intensity=rand_mask_intensity,
-                ext_mask=beat_mask
+                ext_mask=beat_mask, 
+                verbose=True,
+                return_mask=True
             )
+    
+            zv = interface.coarse_to_fine(zv)
+            mask = interface.to_signal(mask_z).cpu()
+
+            sig = interface.to_signal(zv).cpu()
+            print("done")
         elif mode == "loop":
-            zv = interface.loop(
-                zv, 
+            print(f"running loop vampnet with {num_vamps} vamps")
+            sig, mask = interface.loop(
+                sig, 
+                temperature=(init_temp, final_temp),
                 prefix_dur_s=prefix_s, 
                 suffix_dur_s=suffix_s, 
                 num_loops=num_vamps,
                 downsample_factor=mask_periodic_amt,
                 intensity=rand_mask_intensity,
-                ext_mask=beat_mask
+                ext_mask=beat_mask, 
+                verbose=True,
+                return_mask=True
             )
+            sig = sig.cpu()
+            mask = mask.cpu()
+            print("done")
 
-        zv = interface.coarse_to_fine(zv)
-        sig = interface.to_signal(zv)
-        return sig.sample_rate, sig.samples[0].T
-    except Exception as e:
-        raise gr.Error(f"failed with error: {e}")
+
+        out_dir = OUT_DIR / str(uuid.uuid4())
+        out_dir.mkdir()
+        sig.write(out_dir / "output.wav")
+        mask.write(out_dir / "mask.wav")
+        return sig.path_to_file, mask.path_to_file
+    # except Exception as e:
+    #     raise gr.Error(f"failed with error: {e}")
         
+def save_vamp(
+    input_audio, init_temp, final_temp, 
+    prefix_s, suffix_s, rand_mask_intensity,
+    mask_periodic_amt, beat_unmask_dur,
+    mask_dwn_chk, dwn_factor,
+    mask_up_chk, up_factor, 
+    num_vamps, mode, output_audio, notes, use_beats, num_steps
+):
+    out_dir = OUT_DIR / "saved" / str(uuid.uuid4())
+    out_dir.mkdir(parents=True, exist_ok=True)
 
+    sig_in = at.AudioSignal(input_audio.name)
+    sig_out = at.AudioSignal(output_audio.name)
+
+    sig_in.write(out_dir / "input.wav")
+    sig_out.write(out_dir / "output.wav")
+    
+    data = {
+        "init_temp": init_temp,
+        "final_temp": final_temp,
+        "prefix_s": prefix_s,
+        "suffix_s": suffix_s,
+        "rand_mask_intensity": rand_mask_intensity,
+        "mask_periodic_amt": mask_periodic_amt,
+        "use_beats": use_beats,
+        "beat_unmask_dur": beat_unmask_dur,
+        "mask_dwn_chk": mask_dwn_chk,
+        "dwn_factor": dwn_factor,
+        "mask_up_chk": mask_up_chk,
+        "up_factor": up_factor,
+        "num_vamps": num_vamps,
+        "num_steps": num_steps,
+        "mode": mode,
+        "notes": notes,
+    }
+
+    # save with yaml
+    with open(out_dir / "data.yaml", "w") as f:
+        yaml.dump(data, f)
+
+    import zipfile
+    zip_path = out_dir.with_suffix(".zip")
+    with zipfile.ZipFile(zip_path, "w") as zf:
+        for file in out_dir.iterdir():
+            zf.write(file, file.name)
+
+    return f"saved! your save code is {out_dir.stem}", zip_path
 
 with gr.Blocks() as demo:
 
-    gr.Markdown('# Vampnet')
-    
     with gr.Row():
         # input audio
         with gr.Column():
+            gr.Markdown("""
+            # Vampnet
+            **Instructions**:
+            1. Upload some audio (or click the load random audio button)
+            2. Adjust the mask hints. The more hints, the more the generated music will follow the input music
+            3. Adjust the vampnet parameters. The more vamps, the longer the generated music will be
+            4. Click the "vamp" button
+            5. Listen to the generated audio
+            6. If you noticed something you liked, write some notes, click the "save vamp" button, and copy the save code
+
+            
+
+
+            """)
             gr.Markdown("## Input Audio")
+        with gr.Column():
+            gr.Markdown("""
+            ## Mask Hints
+            - most of the original audio will be masked and replaced with audio generated by vampnet
+            - mask hints are used to guide vampnet to generate audio that sounds like the original
+            - the more hints you give, the more the generated audio will sound like the original
+
+
+            
+
+            
+
+
+            """)
+        with gr.Column():
+            gr.Markdown("""
+            ### Tips
+            - use the beat sync button so the output audio has the same beat structure as the input audio
+            - if you want the generated audio to sound like the original, but with a different beat structure:
+                - uncheck the beat sync button
+                - decrease the periodic unmasking to anywhere from 2 to 8
+            - if you want a more "random" generation:
+                - uncheck the beat sync button (or reduce the beat unmask duration)
+                - increase the periodic unmasking to 16 or more
+
+            """)
+
+
+    with gr.Row():
+        with gr.Column():
+            mode = gr.Radio(
+                label="**mode**. note that loop mode requires a prefix and suffix longer than 0",
+                choices=["standard", "loop"],
+                value="standard"
+            )
+            num_vamps = gr.Number(
+                label="number of vamps (or loops). more vamps = longer generated audio",
+                value=1,
+                precision=0
+            )
 
             manual_audio_upload = gr.File(
                 label=f"upload some audio (will be randomly trimmed to max of {interface.coarse.chunk_size_s:.2f}s)",
@@ -126,9 +254,13 @@ with gr.Blocks() as demo:
             input_audio = gr.Audio(
                 label="input audio",
                 interactive=False, 
+                type="file",
             )
-            input_audio_viz = gr.HTML(
-                label="input audio",
+
+            audio_mask = gr.Audio(
+                label="audio mask (listen to this to hear the mask hints)",
+                interactive=False, 
+                type="file",
             )
 
             # connect widgets
@@ -147,113 +279,160 @@ with gr.Blocks() as demo:
 
         # mask settings
         with gr.Column():
-            gr.Markdown("## Mask Settings")
+
+            mask_periodic_amt = gr.Slider(
+                label="periodic unmasking factor (provides a rhythmic, periodic hint). 0.0 means no hint, 2 means one hint every 2 timesteps, etc, 4 means one hint every 4 timesteps, etc.",
+                minimum=0,
+                maximum=32, 
+                step=1,
+                value=16, 
+            )
+
+
+            rand_mask_intensity = gr.Slider(
+                label="random mask intensity. (If this is less than 1, scatters tiny hints throughout the audio, should be between 0.9 and 1.0)",
+                minimum=0.0,
+                maximum=1.0,
+                value=1.0
+            )
+
             prefix_s = gr.Slider(
-                label="prefix length (seconds)",
+                label="prefix hint length (seconds)",
                 minimum=0.0,
                 maximum=10.0,
                 value=0.0
             )
             suffix_s = gr.Slider(
-                label="suffix length (seconds)",
+                label="suffix hint length (seconds)",
                 minimum=0.0,
                 maximum=10.0,
                 value=0.0
             )
 
-            rand_mask_intensity = gr.Slider(
-                label="random mask intensity (lower means more freedom)",
+
+            init_temp = gr.Slider(
+                label="initial temperature (should probably stay between 0.6 and 1)",
                 minimum=0.0,
-                maximum=1.0,
-                value=1.0
+                maximum=1.5,
+                value=0.8
+            )
+            final_temp = gr.Slider(
+                label="final temperature (should probably stay between 0.7 and 2)",
+                minimum=0.0,
+                maximum=2.0,
+                value=0.9
             )
 
-            mask_periodic_amt = gr.Slider(
-                label="periodic unmasking factor (higher means more freedom)",
-                minimum=0,
-                maximum=32, 
+            use_beats = gr.Checkbox(
+                label="use beat hints",
+                value=True
+            )
+
+            num_steps = gr.Slider(
+                label="number of steps (should normally be between 12 and 36)",
+                minimum=4,
+                maximum=128,
                 step=1,
-                value=2, 
+                value=24
             )
-            compute_mask_button = gr.Button("compute mask")
-            mask_output = gr.Audio(
-                label="masked audio",
+
+
+            vamp_button = gr.Button("vamp!!!")
+
+            output_audio = gr.Audio(
+                label="output audio",
                 interactive=False,
-                visible=False
-            )
-            mask_output_viz = gr.Video(
-                label="masked audio",
-                interactive=False
+                type="file"
             )
+
+
+            # gr.Markdown("**NOTE**: for loop mode, both prefix and suffix must be greater than 0.")
+            # compute_mask_button = gr.Button("compute mask")
+            # mask_output = gr.Audio(
+            #     label="masked audio",
+            #     interactive=False,
+            #     visible=False
+            # )
+            # mask_output_viz = gr.Video(
+            #     label="masked audio",
+            #     interactive=False
+            # )
         
         with gr.Column():
-            gr.Markdown("## Beat Unmasking")
-            with gr.Accordion(label="beat unmask"):
+            with gr.Accordion(label="beat unmask (how much time around the beat should be hinted?)"):
+                
                 beat_unmask_dur = gr.Slider(
                     label="duration", 
                     minimum=0.0,
                     maximum=3.0,
                     value=0.1
                 )
-                with gr.Accordion("downbeat settings"):
+                with gr.Accordion("downbeat settings", open=False):
                     mask_dwn_chk = gr.Checkbox(
-                        label="unmask downbeats",
+                        label="hint downbeats",
                         value=True
                     )
                     dwn_factor = gr.Slider(
-                        label="downbeat downsample factor (unmask every Nth downbeat)",
-                        value=1, 
-                        minimum=1,
+                        label="downbeat downsample factor (hint only every Nth downbeat)",
+                        value=0, 
+                        minimum=0,
                         maximum=16, 
                         step=1
                     )
-                with gr.Accordion("upbeat settings"):
+                with gr.Accordion("upbeat settings", open=False):
                     mask_up_chk = gr.Checkbox(
-                        label="unmask upbeats",
+                        label="hint upbeats",
                         value=True
                     )
                     up_factor = gr.Slider(
-                        label="upbeat downsample factor (unmask every Nth upbeat)",
-                        value=1,
-                        minimum=1,
+                        label="upbeat downsample factor (hint only every Nth upbeat)",
+                        value=0,
+                        minimum=0,
                         maximum=16,
                         step=1
                     )
-            
-    # process and output
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("**NOTE**: for loop mode, both prefix and suffix must be greater than 0.")
-            mode = gr.Radio(
-                label="mode",
-                choices=["standard", "loop"],
-                value="standard"
-            )
-            num_vamps = gr.Number(
-                label="number of vamps",
-                value=1,
-                precision=0
-            )
-            vamp_button = gr.Button("vamp")
 
-            output_audio = gr.Audio(
-                label="output audio",
-                interactive=False,
-                visible=False
-            )
+                notes_text = gr.Textbox(
+                    label="type any notes about the generated audio here", 
+                    value="",
+                    interactive=True
+                )
+                save_button = gr.Button("download vamp")
+                download_file = gr.File(
+                    label="vamp to download will appear here",
+                    interactive=False
+                )
+                
+
+                thank_you = gr.Markdown("")
 
+            
     # connect widgets
     vamp_button.click(
         fn=vamp,
-        inputs=[input_audio,
+        inputs=[input_audio, init_temp,final_temp,
             prefix_s, suffix_s, rand_mask_intensity, 
             mask_periodic_amt, beat_unmask_dur, 
             mask_dwn_chk, dwn_factor, 
             mask_up_chk, up_factor, 
-            num_vamps, mode
+            num_vamps, mode, use_beats, num_steps
         ],
-        outputs=[output_audio]
+        outputs=[output_audio, audio_mask]
     )
 
+    save_button.click(
+        fn=save_vamp,
+        inputs=[
+            input_audio, init_temp, final_temp,
+            prefix_s, suffix_s, rand_mask_intensity,
+            mask_periodic_amt, beat_unmask_dur,
+            mask_dwn_chk, dwn_factor,
+            mask_up_chk, up_factor,
+            num_vamps, mode,
+            output_audio,
+            notes_text, use_beats, num_steps
+        ],
+        outputs=[thank_you, download_file]
+    )
 
-demo.launch(share=True, server_name="0.0.0.0")
+demo.launch(share=True, enable_queue=True)

vampnet/beats.py

@@ -215,7 +215,7 @@ class WaveBeat(BeatTracker):
         beats, downbeats = self.model.predict_beats_from_array(
             audio=signal.audio_data.squeeze(0),
             sr=signal.sample_rate,
-            use_gpu=self.device is not "cpu",
+            use_gpu=self.device != "cpu",
         )
 
         return beats, downbeats

vampnet/interface.py

@@ -26,6 +26,7 @@ class Interface:
         coarse_ckpt: str = None,
         coarse2fine_ckpt: str = None,
         codec_ckpt: str = None,
+        wavebeat_ckpt: str = None,
         device: str = "cpu",
         coarse_chunk_size_s: int =  5, 
         coarse2fine_chunk_size_s: int =  3,
@@ -51,6 +52,13 @@ class Interface:
         else:
             self.c2f = None
 
+        if wavebeat_ckpt is not None:
+            print(f"loading wavebeat from {wavebeat_ckpt}")
+            self.beat_tracker = WaveBeat(wavebeat_ckpt)
+            self.beat_tracker.model.to(device)
+        else:
+            self.beat_tracker = None
+
         self.device = device
 
     def s2t(self, seconds: float):
@@ -71,8 +79,13 @@ class Interface:
     def to(self, device):
         self.device = device
         self.coarse.to(device)
-        self.c2f.to(device)
         self.codec.to(device)
+
+        if self.c2f is not None:
+            self.c2f.to(device)
+
+        if self.beat_tracker is not None:
+            self.beat_tracker.model.to(device)
         return self
 
     def to_signal(self, z: torch.Tensor):
@@ -106,7 +119,7 @@ class Interface:
             mask_upbeats: bool = True,
             downbeat_downsample_factor: int = None,
             beat_downsample_factor: int = None,
-            dropout: float = 0.7,
+            dropout: float = 0.3,
             invert: bool = True,
     ):
         """make a beat synced mask. that is, make a mask that 
@@ -146,6 +159,8 @@ class Interface:
 
         beats_z = beats_z[::beat_downsample_factor]
         downbeats_z = downbeats_z[::downbeat_downsample_factor]
+        print(f"beats_z: {len(beats_z)}")
+        print(f"downbeats_z: {len(downbeats_z)}")
     
         if mask_upbeats:
             for beat_idx in beats_z:
@@ -153,8 +168,10 @@ class Interface:
                 num_steps = mask[_slice[0]:_slice[1]].shape[0]
                 _m = torch.ones(num_steps, device=self.device)
                 _m = torch.nn.functional.dropout(_m, p=dropout)
+                print(_m)
                 
                 mask[_slice[0]:_slice[1]] = _m
+                print(mask)
 
         if mask_downbeats:
             for downbeat_idx in downbeats_z:
@@ -165,6 +182,7 @@ class Interface:
                 
                 mask[_slice[0]:_slice[1]] = _m
         
+        mask = mask.clamp(0, 1)
         if invert:
             mask = 1 - mask
         
@@ -317,6 +335,7 @@ class Interface:
         ext_mask=None,
         n_conditioning_codebooks=None,
         verbose=False,
+        return_mask=False,
         **kwargs
     ):
         z = self.encode(signal)
@@ -448,6 +467,9 @@ class Interface:
         prefix_codes = torch.cat(c_vamp['prefix'], dim=-1)
         suffix_codes = torch.cat(c_vamp['suffix'], dim=-1)
         c_vamp = torch.cat([prefix_codes, suffix_codes], dim=-1)
+
+        if return_mask:
+            return c_vamp, cz_masked
         return c_vamp
 
     # create a variation of an audio signal
@@ -527,6 +549,7 @@ class Interface:
         num_loops: int = 4, 
         # overlap_hop_ratio: float = 1.0, # TODO: should this be fixed to 1.0?  or should we overlap and replace instead of overlap add
         verbose: bool = False,
+        return_mask: bool = False,
         **kwargs, 
     ):
         assert prefix_dur_s >= 0.0, "prefix duration must be >= 0"
@@ -549,8 +572,12 @@ class Interface:
                         prefix_dur_s=prefix_dur_s,
                         suffix_dur_s=suffix_dur_s,
                         swap_prefix_suffix=is_flipped,
+                        return_mask=return_mask,
                         **kwargs
                 )
+            if return_mask:
+                vamped, mask = vamped
+            
             # if we're flipped, we trim the prefix off of the end
             # otherwise we trim the suffix off of the end
             trim_len = prefix_len_tokens if is_flipped else suffix_len_tokens
@@ -568,6 +595,9 @@ class Interface:
             loops = [self.coarse_to_fine(l) for l in loops]
 
         loops = [self.to_signal(l) for l in loops]
-            
+
+        if return_mask:
+            return signal_concat(loops), self.to_signal(mask)
+        
         return signal_concat(loops)