app/interface fixes

app.py

@@ -19,38 +19,11 @@ from vampnet import mask as pmask
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 interface = Interface.default()
-
-# populate the model choices with any interface.yml files in the generated confs
-MODEL_CHOICES = {
-    "default": {
-        "Interface.coarse_ckpt": str(interface.coarse_path), 
-        "Interface.coarse2fine_ckpt": str(interface.c2f_path),
-        "Interface.codec_ckpt": str(interface.codec_path),
-    }
-}
-generated_confs = Path("conf/generated")
-for conf_file in generated_confs.glob("*/interface.yml"):
-    with open(conf_file) as f:
-        _conf = yaml.safe_load(f)
-
-        # check if the coarse, c2f, and codec ckpts exist
-        # otherwise, dont' add this model choice
-        if not (
-            Path(_conf["Interface.coarse_ckpt"]).exists() and 
-            Path(_conf["Interface.coarse2fine_ckpt"]).exists() and 
-            Path(_conf["Interface.codec_ckpt"]).exists()
-        ):
-            continue
-
-        MODEL_CHOICES[conf_file.parent.name] = _conf
-
     
 def to_output(sig):
     return sig.sample_rate, sig.cpu().detach().numpy()[0][0]
 
-
-
-MAX_DURATION_S = 5
+MAX_DURATION_S = 10
 def load_audio(file):
     print(file)
     if isinstance(file, str):
@@ -91,6 +64,7 @@ def _vamp(
         typical_mass, typical_min_tokens, top_p, 
         sample_cutoff, stretch_factor, api=False
     ):
+
     t0 = time.time()
     interface.to("cuda" if torch.cuda.is_available() else "cpu")
     print(f"using device {interface.device}")
@@ -105,15 +79,15 @@ def _vamp(
     sig = at.AudioSignal(input_audio, sr)
 
     # reload the model if necessary
-    interface.reload(
-        coarse_ckpt=MODEL_CHOICES[model_choice]["Interface.coarse_ckpt"],
-        c2f_ckpt=MODEL_CHOICES[model_choice]["Interface.coarse2fine_ckpt"],
-    )
+    interface.load_finetuned(model_choice)
 
     if pitch_shift_amt != 0:
         sig = shift_pitch(sig, pitch_shift_amt)
 
-    build_mask_kwargs = dict(
+    codes = interface.encode(sig)
+
+    mask = interface.build_mask(
+        codes, sig,
         rand_mask_intensity=1.0,
         prefix_s=0.0,
         suffix_s=0.0,
@@ -124,29 +98,26 @@ def _vamp(
         upper_codebook_mask=int(n_mask_codebooks), 
     )
 
-    vamp_kwargs = dict(
-        temperature=sampletemp,
-        typical_filtering=typical_filtering, 
-        typical_mass=typical_mass, 
-        typical_min_tokens=typical_min_tokens, 
-        top_p=None,
-        seed=_seed,
-        sample_cutoff=1.0,
-    )
 
     # save the mask as a txt file
     interface.set_chunk_size(10.0)
-    sig, mask, codes = interface.ez_vamp(
-        sig, 
+    codes, mask = interface.vamp(
+        codes, mask,
         batch_size=1 if api else 1,
         feedback_steps=1,
         time_stretch_factor=stretch_factor,
-        build_mask_kwargs=build_mask_kwargs,
-        vamp_kwargs=vamp_kwargs,
         return_mask=True,
+        temperature=sampletemp,
+        typical_filtering=typical_filtering, 
+        typical_mass=typical_mass, 
+        typical_min_tokens=typical_min_tokens, 
+        top_p=None,
+        seed=_seed,
+        sample_cutoff=1.0,
     )
     print(f"vamp took {time.time() - t0} seconds")
 
+    sig = interface.decode(codes)
 
     return to_output(sig)
 
@@ -352,7 +323,7 @@ with gr.Blocks() as demo:
 
             model_choice = gr.Dropdown(
                 label="model choice", 
-                choices=list(MODEL_CHOICES.keys()),
+                choices=list(interface.available_models()),
                 value="default", 
                 visible=True
             )

token_telephone/vamp_helper.py

@@ -136,7 +136,7 @@ def ez_variation(
 
     # save the mask as a txt file
     interface.set_chunk_size(10.0)
-    sig, mask, codes = interface.ez_vamp(
+    sig, mask, codes = interface.vamp(
         sig, 
         batch_size=1,
         feedback_steps=1,

vampnet/interface.py

@@ -128,13 +128,16 @@ class Interface(torch.nn.Module):
     @classmethod
     def available_models(cls):
         from . import list_finetuned
-        return list_finetuned()
+        return list_finetuned() + ["default"]
 
 
     def load_finetuned(self, name: str):
         assert name in self.available_models(), f"{name} is not a valid model name"
-        from . import download_finetuned
-        coarse_path, c2f_path = download_finetuned(name)
+        from . import download_finetuned, download_default
+        if name == "default":
+            coarse_path, c2f_path = download_default()
+        else:
+            coarse_path, c2f_path = download_finetuned(name)
         self.reload(
             coarse_ckpt=coarse_path,
             c2f_ckpt=c2f_path,