coarse eval

conf/interface.yml

@@ -0,0 +1,5 @@
+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

scripts/utils/vamp_folder.py

@@ -6,7 +6,7 @@ import torch
 
 from vampnet.interface import Interface
 
-Interface = argbind.bind(Interface, positional=True)
+Interface = argbind.bind(Interface)
 
 def baseline(sig, interface):
     return sig
@@ -26,7 +26,9 @@ def coarse2fine(sig, interface):
 def one_codebook(sig, interface):
     z = interface.encode(sig)
 
-    mask = torch.zeros_like(z)
+    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(
@@ -46,7 +48,9 @@ def four_codebooks_downsampled_4x(sig, interface):
 def two_codebooks_downsampled_4x(sig, interface):
     z = interface.encode(sig)
 
-    mask = torch.zeros_like(z)
+    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(
@@ -64,9 +68,6 @@ def four_codebooks_downsampled_8x(sig, interface):
     return interface.to_signal(zv)
 
 
-
-
-
 SAMPLE_CONDS ={
     "baseline": baseline,
     "reconstructed": reconstructed,
@@ -105,15 +106,15 @@ def main(
         sig = dataset[i]["signal"]
         
         results = {
-            name: cond(sig, interface)
+            name: cond(sig, interface).cpu()
             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)
+            o_dir = Path(output_dir) / name
+            o_dir.mkdir(exist_ok=True, parents=True)
 
-            sig.write(output_dir / f"{i}.wav")
+            sig.write(o_dir / f"{i}.wav")
 
 if __name__ == "__main__":
     args = argbind.parse_args()

vampnet/interface.py

@@ -23,17 +23,19 @@ def signal_concat(
 class Interface:
     def __init__(
         self,
-        coarse_ckpt: str,
-        coarse2fine_ckpt: str,
-        codec_ckpt: str,
+        coarse_ckpt: str = None,
+        coarse2fine_ckpt: str = None,
+        codec_ckpt: str = None,
         device: str = "cpu",
         coarse_chunk_size_s: int =  5, 
         coarse2fine_chunk_size_s: int =  3,
     ):
+        assert codec_ckpt is not None, "must provide a codec checkpoint"
         self.codec = LAC.load(Path(codec_ckpt))
         self.codec.eval()
         self.codec.to(device)
 
+        assert coarse_ckpt is not None, "must provide a coarse checkpoint"
         self.coarse = VampNet.load(location=Path(coarse_ckpt), map_location="cpu")
         self.coarse.to(device)
         self.coarse.eval()