changes

conf/{interface.yml → interface-jazzpop-exp.yml}

@@ -2,4 +2,8 @@ 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
+Interface.coarse2fine_chunk_size_s: 3
+
+AudioLoader.sources:
+  - /data/spotdl/audio/val
+  - /data/spotdl/audio/test

demo.py

@@ -0,0 +1,220 @@
+from pathlib import Path
+from typing import Tuple
+import yaml
+
+import numpy as np
+import audiotools as at
+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)
+with argbind.scope(conf):
+    interface = Interface()
+    loader = AudioLoader()
+
+dataset = at.data.datasets.AudioDataset(
+    loader,
+    sample_rate=interface.codec.sample_rate,
+    duration=interface.coarse.chunk_size_s,
+    n_examples=5000,
+    without_replacement=True,
+)
+
+
+def load_audio(file):
+    print(file)
+    filepath = file.name
+    sig = at.AudioSignal.salient_excerpt(
+        filepath, 
+        duration=interface.coarse.chunk_size_s
+    )
+    sig = interface.preprocess(sig)
+
+    audio = sig.samples.numpy()[0]
+    sr = sig.sample_rate
+    return sr, audio.T
+
+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
+
+def mask_audio(
+        prefix_s, suffix_s, rand_mask_intensity, 
+        mask_periodic_amt, beat_unmask_dur, 
+        mask_dwn_chk, dwn_factor, 
+        mask_up_chk, up_factor
+    ):
+    pass
+
+def vamp(
+    input_audio, prefix_s, suffix_s, rand_mask_intensity,
+    mask_periodic_amt, beat_unmask_dur,
+    mask_dwn_chk, dwn_factor,
+    mask_up_chk, up_factor
+):
+    print(input_audio)
+
+
+with gr.Blocks() as demo:
+
+    gr.Markdown('# Vampnet')
+    
+    with gr.Row():
+        # input audio
+        with gr.Column():
+            gr.Markdown("## Input Audio")
+
+            manual_audio_upload = gr.File(
+                label=f"upload some audio (will be randomly trimmed to max of {interface.coarse.chunk_size_s:.2f}s)",
+                file_types=["audio"]
+            )
+            load_random_audio_button = gr.Button("or load random audio")
+
+            input_audio = gr.Audio(
+                label="input audio",
+                interactive=False, 
+            )
+            input_audio_viz = gr.HTML(
+                label="input audio",
+            )
+
+            # connect widgets
+            load_random_audio_button.click(
+                fn=load_random_audio,
+                inputs=[],
+                outputs=[ input_audio]
+            )
+
+            manual_audio_upload.change(
+                fn=load_audio,
+                inputs=[manual_audio_upload],
+                outputs=[ input_audio]
+            )
+                
+
+        # mask settings
+        with gr.Column():
+            gr.Markdown("## Mask Settings")
+            prefix_s = gr.Slider(
+                label="prefix length (seconds)",
+                minimum=0.0,
+                maximum=10.0,
+                value=0.0
+            )
+            suffix_s = gr.Slider(
+                label="suffix length (seconds)",
+                minimum=0.0,
+                maximum=10.0,
+                value=0.0
+            )
+
+            rand_mask_intensity = gr.Slider(
+                label="random mask intensity (lower means more freedom)",
+                minimum=0.0,
+                maximum=1.0,
+                value=1.0
+            )
+
+            mask_periodic_amt = gr.Slider(
+                label="periodic unmasking factor (higher means more freedom)",
+                minimum=0,
+                maximum=32, 
+                step=1,
+                value=2, 
+            )
+            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"):
+                beat_unmask_dur = gr.Slider(
+                    label="duration", 
+                    minimum=0.0,
+                    maximum=3.0,
+                    value=0.1
+                )
+                with gr.Accordion("downbeat settings"):
+                    mask_dwn_chk = gr.Checkbox(
+                        label="unmask downbeats",
+                        value=True
+                    )
+                    dwn_factor = gr.Slider(
+                        label="downbeat downsample factor (unmask every Nth downbeat)",
+                        value=1, 
+                        minimum=1,
+                        maximum=16, 
+                        step=1
+                    )
+                with gr.Accordion("upbeat settings"):
+                    mask_up_chk = gr.Checkbox(
+                        label="unmask upbeats",
+                        value=True
+                    )
+                    up_factor = gr.Slider(
+                        label="upbeat downsample factor (unmask every Nth upbeat)",
+                        value=1,
+                        minimum=1,
+                        maximum=16,
+                        step=1
+                    )
+            
+    # process and output
+    with gr.Row():
+        with gr.Column():
+            vamp_button = gr.Button("vamp")
+
+            output_audio = gr.Audio(
+                label="output audio",
+                interactive=False,
+                visible=False
+            )
+            output_audio_viz = gr.Video(
+                label="output audio",
+                interactive=False
+            )
+
+    # connect widgets
+    compute_mask_button.click(
+        fn=mask_audio,
+        inputs=[
+            prefix_s, suffix_s, rand_mask_intensity, 
+            mask_periodic_amt, beat_unmask_dur, 
+            mask_dwn_chk, dwn_factor, 
+            mask_up_chk, up_factor
+        ],
+        outputs=[mask_output, mask_output_viz]
+    )
+
+    # connect widgets
+    vamp_button.click(
+        fn=vamp,
+        inputs=[input_audio,
+            prefix_s, suffix_s, rand_mask_intensity, 
+            mask_periodic_amt, beat_unmask_dur, 
+            mask_dwn_chk, dwn_factor, 
+            mask_up_chk, up_factor
+        ],
+        outputs=[output_audio, output_audio_viz]
+    )
+
+
+demo.launch(share=True)

scripts/exp/eval.py

@@ -57,30 +57,31 @@ def eval(
         cond_files = cond_files[:num_files]
         assert len(list(baseline_files)) == len(list(cond_files)), f"number of files in {baseline_dir} and {cond_dir} do not match. {len(list(baseline_files))} vs {len(list(cond_files))}"
 
-        pbar = tqdm(zip(baseline_files, cond_files), total=len(baseline_files))
-        for baseline_file, cond_file in pbar:
+        def process(baseline_file, cond_file):
             # make sure the files match (same name)
             assert baseline_file.stem == cond_file.stem, f"baseline file {baseline_file} and cond file {cond_file} do not match"
-            pbar.set_description(baseline_file.stem)
 
             # load the files
             baseline_sig = AudioSignal(str(baseline_file))
             cond_sig = AudioSignal(str(cond_file))
 
             # compute the metrics
-            try:
-                vsq = visqol(baseline_sig, cond_sig)
-            except:
-                vsq = 0.0
-            metrics.append({
+            # try:
+            #     vsq = visqol(baseline_sig, cond_sig)
+            # except:
+            #     vsq = 0.0
+            return {
                 "sisdr": -sisdr_loss(baseline_sig, cond_sig).item(),
                 "stft": stft_loss(baseline_sig, cond_sig).item(),
                 "mel": mel_loss(baseline_sig, cond_sig).item(),
                 "frechet": frechet_score,
-                "visqol": vsq,
+                # "visqol": vsq,
                 "condition": condition,
                 "file": baseline_file.stem,
-            })
+            }
+
+        print(f"processing {len(baseline_files)} files in {baseline_dir} and {cond_dir}")
+        metrics.extend(tqdm(map(process, baseline_files, cond_files), total=len(baseline_files)))
 
     metric_keys = [k for k in metrics[0].keys() if k not in ("condition", "file")]
 

scripts/utils/vamp_folder.py

@@ -1,8 +1,10 @@
 from pathlib import Path
+import random
+from typing import List
 
 import argbind
 from tqdm import tqdm
-import torch
+import argbind
 
 from vampnet.interface import Interface
 import audiotools as at
@@ -12,9 +14,9 @@ Interface = argbind.bind(Interface)
 # condition wrapper for printing
 def condition(cond):
     def wrapper(sig, interface):
-        print(f"Condition: {cond.__name__}")
+        # print(f"Condition: {cond.__name__}")
         sig = cond(sig, interface)
-        print(f"Condition: {cond.__name__} (done)\n")
+        # print(f"Condition: {cond.__name__} (done)\n")
         return sig
     return wrapper
 
@@ -49,48 +51,27 @@ def coarse2fine_argmax(sig, interface):
 
 @condition
 def one_codebook(sig, interface):
-    z = interface.encode(sig)
-
-    nb, _, nt = z.shape 
-    nc = interface.coarse.n_codebooks
-    mask = torch.zeros(nb, nc, nt).to(interface.device)
-    mask[:, 1:, :] = 1
-
     zv = interface.coarse_vamp_v2(
-        sig, ext_mask=mask,
+        sig, n_conditioning_codebooks=1
     )
     zv = interface.coarse_to_fine(zv)  
 
     return interface.to_signal(zv)
 
-@condition
-def four_codebooks_downsampled_4x(sig, interface):
-    zv = interface.coarse_vamp_v2(
-        sig, downsample_factor=4
-    )
-    zv = interface.coarse_to_fine(zv)  
-    return interface.to_signal(zv)
-
 @condition
 def two_codebooks_downsampled_4x(sig, interface):
-    z = interface.encode(sig)
-
-    nb, _, nt = z.shape 
-    nc = interface.coarse.n_codebooks
-    mask = torch.zeros(nb, nc, nt).to(interface.device)
-    mask[:, 2:, :] = 1
-
     zv = interface.coarse_vamp_v2(
-        sig, ext_mask=mask, downsample_factor=4
+        sig, n_conditioning_codebooks=2,
+        downsample_factor=4
     )
     zv = interface.coarse_to_fine(zv)
 
     return interface.to_signal(zv)
 
-@condition
-def four_codebooks_downsampled_8x(sig, interface):
+
+def four_codebooks_downsampled(sig, interface, x=12):
     zv = interface.coarse_vamp_v2(
-        sig, downsample_factor=8
+        sig, downsample_factor=12
     )
     zv = interface.coarse_to_fine(zv)  
     return interface.to_signal(zv)
@@ -101,9 +82,13 @@ COARSE_SAMPLE_CONDS ={
     "reconstructed": reconstructed,
     "coarse2fine": coarse2fine,
     "one_codebook": one_codebook,
-    "four_codebooks_downsampled_4x": four_codebooks_downsampled_4x,
     "two_codebooks_downsampled_4x": two_codebooks_downsampled_4x,
-    "four_codebooks_downsampled_8x": four_codebooks_downsampled_8x,
+    # four codebooks at different downsample factors
+    **{
+        f"four_codebooks_downsampled_{x}x": lambda sig, interface: four_codebooks_downsampled(sig, interface, x=x)
+        for x in [4, 8, 12, 16, 20, 24]
+    }
+
 }
 
 C2F_SAMPLE_CONDS = {
@@ -131,7 +116,7 @@ def main(
 
     from audiotools.data.datasets import AudioLoader, AudioDataset
 
-    loader = AudioLoader(sources=sources)
+    loader = AudioLoader(sources=sources, shuffle_state=seed)
     dataset = AudioDataset(loader, 
         sample_rate=interface.codec.sample_rate, 
         duration=interface.coarse.chunk_size_s, 
@@ -141,7 +126,18 @@ def main(
 
     SAMPLE_CONDS = COARSE_SAMPLE_CONDS if exp_type == "coarse" else C2F_SAMPLE_CONDS
 
-    for i in tqdm(range(max_excerpts)):
+
+    indices = list(range(max_excerpts))
+    random.shuffle(indices)
+    for i in tqdm(indices):
+        # if all our files are already there, skip
+        # done = []
+        # for name in SAMPLE_CONDS:
+        #     o_dir = Path(output_dir) / name
+        #     done.append((o_dir / f"{i}.wav").exists())
+        # if all(done):
+        #     continue
+
         sig = dataset[i]["signal"]
         
         results = {

setup.py

@@ -26,16 +26,15 @@ setup(
     license="MIT",
     packages=find_packages(),
     install_requires=[
-        "torch<=1.11.0",
+        "torch",
         "argbind>=0.3.2",
         "pytorch-ignite",
         "rich",
-        "audiotools @ git+https://github.com/descriptinc/lyrebird-audiotools.git@0.6.3",
-        "lac @ git+https://github.com/descriptinc/lyrebird-audio-codec.git@main",
+        "audiotools @ git+https://github.com/descriptinc/lyrebird-audiotools.git@hf/backup-info",
+        "lac @ git+https://github.com/descriptinc/lyrebird-audio-codec.git",
         "tqdm",
         "tensorboard",
         "google-cloud-logging==2.2.0",
-        "torchmetrics>=0.7.3",
         "einops",
         "frechet_audio_distance"
     ],

vampnet/gradio.py

@@ -1,4 +0,0 @@
-
-import gradio as gr
-
-

vampnet/interface.py

@@ -315,6 +315,7 @@ class Interface:
         debug=False,
         swap_prefix_suffix=False, 
         ext_mask=None,
+        n_conditioning_codebooks=None,
         verbose=False,
         **kwargs
     ):
@@ -351,7 +352,8 @@ class Interface:
                 n_suffix=n_suffix, 
                 downsample_factor=downsample_factor,
                 mask=cz_mask, 
-                ext_mask=ext_mask
+                ext_mask=ext_mask, 
+                n_conditioning_codebooks=n_conditioning_codebooks
             )
             if debug:
                 print("tokens to infer")

vampnet/modules/base.py

@@ -41,6 +41,7 @@ class VampBase(at.ml.BaseModel):
         n_prefix: Optional[torch.Tensor] = None,
         n_suffix: Optional[torch.Tensor] = None,
         downsample_factor: Optional[int] = None, 
+        n_conditioning_codebooks: Optional[int] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         assert x.ndim == 3, "x must be (batch, n_codebooks, seq)"
 
@@ -79,7 +80,8 @@ class VampBase(at.ml.BaseModel):
             mask = mask.round().long()
             
             # if we have any conditioning codebooks, set their mask  to 0
-            mask[:, : self.n_conditioning_codebooks, :] = 0
+            n_conditioning_codebooks = n_conditioning_codebooks or self.n_conditioning_codebooks
+            mask[:, :n_conditioning_codebooks, :] = 0
         else:
             assert mask.ndim == 3, "mask must be (batch, n_codebooks, seq)"
             assert mask.shape == x.shape, "mask must be same shape as x"

vampnet/util.py

@@ -0,0 +1,63 @@
+import tqdm
+# import pathos
+
+def process_map(fn, *iterables, **tqdm_kwargs):
+    """
+    Equivalent of `list(map(fn, *iterables))`
+    driven by `concurrent.futures.ProcessPoolExecutor`.
+
+    Parameters
+    ----------
+    tqdm_class  : optional
+        `tqdm` class to use for bars [default: tqdm.auto.tqdm].
+    max_workers  : int, optional
+        Maximum number of workers to spawn; passed to
+        `concurrent.futures.ProcessPoolExecutor.__init__`.
+        [default: min(32, cpu_count() + 4)].
+    chunksize  : int, optional
+        Size of chunks sent to worker processes; passed to
+        `concurrent.futures.ProcessPoolExecutor.map`. [default: 1].
+    lock_name  : str, optional
+        Member of `tqdm_class.get_lock()` to use [default: mp_lock].
+    """
+    from concurrent.futures import ProcessPoolExecutor
+    if iterables and "chunksize" not in tqdm_kwargs:
+        # default `chunksize=1` has poor performance for large iterables
+        # (most time spent dispatching items to workers).
+        longest_iterable_len = max(map(length_hint, iterables))
+        if longest_iterable_len > 1000:
+            from warnings import warn
+            warn("Iterable length %d > 1000 but `chunksize` is not set."
+                 " This may seriously degrade multiprocess performance."
+                 " Set `chunksize=1` or more." % longest_iterable_len,
+                 TqdmWarning, stacklevel=2)
+    if "lock_name" not in tqdm_kwargs:
+        tqdm_kwargs = tqdm_kwargs.copy()
+        tqdm_kwargs["lock_name"] = "mp_lock"
+    return _executor_map(ProcessPoolExecutor, fn, *iterables, **tqdm_kwargs)
+
+
+def parallelize(
+        fn, 
+        *iterables,
+        parallel: str = "thread_map",
+        **kwargs
+    ):
+    if parallel == "thread_map":
+        from tqdm.contrib.concurrent import thread_map
+        return thread_map(
+            fn, 
+            *iterables, 
+            **kwargs
+        )
+    elif parallel == "process_map":
+        from tqdm.contrib.concurrent import process_map
+        return process_map(
+            fn, 
+            *iterables, 
+            **kwargs
+        )
+    elif parallel == "single":
+        return [fn(x) for x in tqdm.tqdm(*iterables)]
+    else:
+        raise ValueError(f"parallel must be one of 'thread_map', 'process_map', 'single', but got {parallel}")