lora interface

app.py

@@ -18,10 +18,45 @@ Interface = argbind.bind(Interface)
 
 conf = argbind.parse_args()
 
-with argbind.scope(conf):
-    interface = Interface()
-    # loader = AudioLoader()
-    print(f"interface device is {interface.device}")
+def load_interface():
+    with argbind.scope(conf):
+        interface = Interface()
+        # loader = AudioLoader()
+        print(f"interface device is {interface.device}")
+        return interface
+
+
+LORA_NONE = "None"
+def load_loras():
+    loras = {}
+    # find confs under conf/generated
+    for conf_file in Path("conf/generated").glob("**/interface.yml"):
+        name = conf_file.parent.name
+        with open(conf_file) as f:
+            loras[name] = yaml.safe_load(f)
+    loras[LORA_NONE] = None
+    return loras
+
+interface = load_interface()
+loras = load_loras()
+cur_lora = LORA_NONE
+
+def load_lora(name):
+    global interface
+    global cur_lora
+    if name == cur_lora:
+        return
+    if name != LORA_NONE:
+        interface.lora_load(
+            coarse_ckpt=loras[name]["Interface.coarse_lora_ckpt"],
+            c2f_ckpt=loras[name]["Interface.coarse2fine_lora_ckpt"],
+            full_ckpts=False
+        )
+        cur_lora = name
+
+    else: 
+        interface = load_interface()
+        cur_lora = LORA_NONE
 
 # dataset = at.data.datasets.AudioDataset(
 #     loader,
@@ -55,6 +90,8 @@ def load_example_audio():
 
 
 def _vamp(data, return_mask=False):
+    load_lora(data[lora_choice])
+
     out_dir = OUT_DIR / str(uuid.uuid4())
     out_dir.mkdir()
     sig = at.AudioSignal(data[input_audio])
@@ -173,6 +210,7 @@ def save_vamp(data):
         "use_coarse2fine": data[use_coarse2fine],
         "stretch_factor": data[stretch_factor],
         "seed": data[seed],
+        "lora": data[lora_choice],
     }
 
     # save with yaml
@@ -472,6 +510,13 @@ with gr.Blocks() as demo:
 
         # mask settings
         with gr.Column():
+
+            lora_choice = gr.Dropdown(
+                label="lora choice", 
+                choices=list(loras.keys()),
+                value=LORA_NONE,
+            )
+
             vamp_button = gr.Button("generate (vamp)!!!")
             output_audio = gr.Audio(
                 label="output audio",
@@ -514,7 +559,7 @@ with gr.Blocks() as demo:
             beat_mask_width,
             beat_mask_downbeats,
             seed, 
-            seed
+            lora_choice,
         }
   
     # connect widgets

vampnet/interface.py

@@ -120,17 +120,16 @@ class Interface(torch.nn.Module):
             if coarse_ckpt is not None:
                 self.coarse.to("cpu")
                 state_dict = torch.load(coarse_ckpt, map_location="cpu")
-
+                print(f"loading coarse from {coarse_ckpt}")
                 self.coarse.load_state_dict(state_dict, strict=False)
                 self.coarse.to(self.device)
             if c2f_ckpt is not None:
                 self.c2f.to("cpu")
                 state_dict = torch.load(c2f_ckpt, map_location="cpu")
-
+                print(f"loading c2f from {c2f_ckpt}")
                 self.c2f.load_state_dict(state_dict, strict=False)
                 self.c2f.to(self.device)
         
-
     def s2t(self, seconds: float):
         """seconds to tokens"""
         if isinstance(seconds, np.ndarray):