exps

scripts/utils/process_folder-c2f.py

@@ -15,57 +15,47 @@ def coarse2fine_infer(
         model, 
         vqvae, 
         device,
-        signal_window=3, 
-        signal_hop=1.5,
-        max_excerpts=20, 
     ):
-    output = defaultdict(list)
-
-    # split into 3 seconds
-    windows = [s for s in signal.clone().windows(signal_window, signal_hop)]
-    windows = windows[1:] # skip first window since it's half zero padded
-    random.shuffle(windows)
-    for w in windows[:max_excerpts]: 
-        # batch the signal into chunks of 3
-        with torch.no_grad():
-            # get codes
-            w = w.to(device)
-            z = vqvae.encode(w.audio_data, w.sample_rate)["codes"]
-
-            model.to(device)
-            output["reconstructed"] = model.to_signal(z, vqvae).cpu()
-
-            # make a full mask
-            mask = torch.ones_like(z)
-            mask[:, :model.n_conditioning_codebooks, :] = 0
-
-            output["sampled"].append(model.sample(
-                codec=vqvae, 
-                time_steps=z.shape[-1], 
-                sampling_steps=12, 
-                start_tokens=z, 
-                mask=mask, 
-                temperature=0.85, 
-                top_k=None, 
-                sample="gumbel", 
-                typical_filtering=True, 
-                return_signal=True
-            ).cpu())
-
-            output["argmax"].append(model.sample(
-                codec=vqvae, 
-                time_steps=z.shape[-1], 
-                sampling_steps=1, 
-                start_tokens=z, 
-                mask=mask, 
-                temperature=1.0, 
-                top_k=None, 
-                sample="argmax", 
-                typical_filtering=True, 
-                return_signal=True
-            ).cpu())
-
-    return output
+        output = {}
+        w = signal
+        w = w.to(device)
+        z = vqvae.encode(w.audio_data, w.sample_rate)["codes"]
+
+        model.to(device)
+        output["reconstructed"] = model.to_signal(z, vqvae).cpu()
+
+        # make a full mask
+        mask = torch.ones_like(z)
+        mask[:, :model.n_conditioning_codebooks, :] = 0
+
+        output["sampled"] = model.sample(
+            codec=vqvae, 
+            time_steps=z.shape[-1], 
+            sampling_steps=12, 
+            start_tokens=z, 
+            mask=mask, 
+            temperature=0.85, 
+            top_k=None, 
+            sample="gumbel", 
+            typical_filtering=True, 
+            return_signal=True
+        ).cpu()
+
+        output["argmax"] = model.sample(
+            codec=vqvae, 
+            time_steps=z.shape[-1], 
+            sampling_steps=1, 
+            start_tokens=z, 
+            mask=mask, 
+            temperature=1.0, 
+            top_k=None, 
+            sample="argmax", 
+            typical_filtering=True, 
+            return_signal=True
+        ).cpu()
+
+        return output
+
 
 
 @argbind.bind(without_prefix=True)
@@ -73,11 +63,10 @@ def main(
         sources=[
             "/data/spotdl/audio/val", "/data/spotdl/audio/test"
         ], 
-        audio_ext="mp3",
         exp_name="noise_mode",
         model_paths=[
-            "runs/c2f-exp-03.22.23/ckpt/mask/best/vampnet/weights.pth",
-            "runs/c2f-exp-03.22.23/ckpt/random/best/vampnet/weights.pth",
+            "runs/c2f-exp-03.22.23/ckpt/mask/epoch=400/vampnet/weights.pth",
+            "runs/c2f-exp-03.22.23/ckpt/random/epoch=400/vampnet/weights.pth",
         ],
         model_keys=[
             "mask",
@@ -86,10 +75,11 @@ def main(
         vqvae_path: str = "runs/codec-ckpt/codec.pth",
         device: str = "cuda",
         output_dir: str = ".",
+        max_excerpts: int = 5000,
+        duration: float = 3.0,
     ):
     from vampnet.modules.transformer import VampNet
     from lac.model.lac import LAC
-    from audiotools.post import audio_zip
 
     models = {
         k: VampNet.load(p) for k, p in zip(model_keys, model_paths)
@@ -105,26 +95,26 @@ def main(
 
     output_dir = Path(output_dir) / f"{exp_name}-samples"
 
-    for source in sources:
-        print(f"Processing {source}...")
-        source_files = list(Path(source).glob(f"**/*.{audio_ext}"))
-        random.shuffle(source_files)
-        for path in tqdm(source_files):
-            sig = AudioSignal(path)
-            sig.resample(vqvae.sample_rate).normalize(-24).ensure_max_of_audio(1.0)
-
-            out_dir = output_dir / path.stem 
+    from audiotools.data.datasets import AudioLoader, AudioDataset
+
+    loader = AudioLoader(sources=sources)
+    dataset = AudioDataset(loader, 
+        sample_rate=vqvae.sample_rate, 
+        duration=duration, 
+        n_examples=max_excerpts, 
+        without_replacement=True,
+    )
+    for i in tqdm(range(max_excerpts)):
+        sig = dataset[i]["signal"]
+        sig.resample(vqvae.sample_rate).normalize(-24).ensure_max_of_audio(1.0)
+
+        for model_key, model in models.items():
+            out = coarse2fine_infer(sig, model, vqvae, device)
+            out_dir = output_dir / model_key / Path(sig.path_to_file).stem
             out_dir.mkdir(parents=True, exist_ok=True)
-            if out_dir.exists():
-                print(f"Skipping {path.stem} since {out_dir} already exists.")
-                continue
-
-            for model_key, model in models.items():
-                out = coarse2fine_infer(sig, model, vqvae, device)
-                for k, sig_list in out.items():
-                    for i, s in enumerate(sig_list):
-                        s.write(out_dir / f"{model_key}-{k}-{i}.wav")
-            
+            for k, s in out.items():
+                s.write(out_dir / f"{k}.wav")
+        
 
 if __name__ == "__main__":
     args = argbind.parse_args()

scripts/utils/vamp_folder.py

@@ -0,0 +1,122 @@
+from pathlib import Path
+
+import argbind
+from tqdm import tqdm
+import torch
+
+from vampnet.interface import Interface
+
+Interface = argbind.bind(Interface, positional=True)
+
+def baseline(sig, interface):
+    return sig
+
+def reconstructed(sig, interface):
+    return interface.to_signal(
+        interface.encode(sig)
+    )
+
+def coarse2fine(sig, interface):
+    z = interface.encode(sig)
+    z = z[:, :interface.c2f.n_conditioning_codebooks, :]
+
+    z = interface.coarse_to_fine(z)
+    return interface.to_signal(z)
+
+def one_codebook(sig, interface):
+    z = interface.encode(sig)
+
+    mask = torch.zeros_like(z)
+    mask[:, 1:, :] = 1
+
+    zv = interface.coarse_vamp_v2(
+        sig, ext_mask=mask,
+    )
+    zv = interface.coarse_to_fine(zv)  
+
+    return interface.to_signal(zv)
+
+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)
+
+def two_codebooks_downsampled_4x(sig, interface):
+    z = interface.encode(sig)
+
+    mask = torch.zeros_like(z)
+    mask[:, 2:, :] = 1
+
+    zv = interface.coarse_vamp_v2(
+        sig, ext_mask=mask, downsample_factor=4
+    )
+    zv = interface.coarse_to_fine(zv)
+
+    return interface.to_signal(zv)
+
+def four_codebooks_downsampled_8x(sig, interface):
+    zv = interface.coarse_vamp_v2(
+        sig, downsample_factor=8
+    )
+    zv = interface.coarse_to_fine(zv)  
+    return interface.to_signal(zv)
+
+
+
+
+
+SAMPLE_CONDS ={
+    "baseline": baseline,
+    "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,
+}
+
+
+@argbind.bind(without_prefix=True)
+def main(
+        sources=[
+            "/data/spotdl/audio/val", "/data/spotdl/audio/test"
+        ], 
+        output_dir: str = "./samples",
+        max_excerpts: int = 5000,
+    ):
+    interface = Interface()
+
+    output_dir = Path(output_dir) 
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    from audiotools.data.datasets import AudioLoader, AudioDataset
+
+    loader = AudioLoader(sources=sources)
+    dataset = AudioDataset(loader, 
+        sample_rate=interface.codec.sample_rate, 
+        duration=interface.coarse.chunk_size_s, 
+        n_examples=max_excerpts, 
+        without_replacement=True,
+    )
+
+    for i in tqdm(range(max_excerpts)):
+        sig = dataset[i]["signal"]
+        
+        results = {
+            name: cond(sig, interface)
+            for name, cond in SAMPLE_CONDS.items()
+        }
+
+        for name, sig in results.items():
+            output_dir = Path(output_dir) / name
+            output_dir.mkdir(exist_ok=True, parents=True)
+
+            sig.write(output_dir / f"{i}.wav")
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+
+    with argbind.scope(args):
+        main()

vampnet/interface.py

@@ -196,6 +196,7 @@ class Interface:
                 time_steps=chunk_len,
                 start_tokens=chunk,
                 return_signal=False,
+                **kwargs
             )
             fine_z.append(chunk)
 

vampnet/modules/base.py

@@ -288,12 +288,12 @@ class VampBase(at.ml.BaseModel):
         self,
         codec,
         time_steps: int = 300,
-        sampling_steps: int = 24,
+        sampling_steps: int = 12,
         start_tokens: Optional[torch.Tensor] = None,
         mask: Optional[torch.Tensor] = None,
         temperature: Union[float, Tuple[float, float]] = 0.8,
         top_k: int = None,
-        sample: str = "multinomial",
+        sample: str = "gumbel",
         typical_filtering=False,
         typical_mass=0.2,
         typical_min_tokens=1,