interface for max

conf/interface-spotdl.yml

@@ -3,7 +3,7 @@ Interface.coarse2fine_ckpt: ./models/spotdl/c2f.pth
 Interface.codec_ckpt: ./models/spotdl/codec.pth
 Interface.coarse_chunk_size_s: 10
 Interface.coarse2fine_chunk_size_s: 3
-Interface.wavebeat_ckpt: ./models/wavebeat.pth
+# Interface.wavebeat_ckpt: ./models/wavebeat.pth
 
 
 AudioLoader.sources:

demo.py

@@ -59,6 +59,43 @@ def load_random_audio():
     return sig.path_to_file
 
 
+def ez_vamp(
+    input_audio, init_temp, final_temp,  
+    mask_periodic_amt, mask_periodic_width, num_steps,  
+):
+    print(input_audio)
+    sig = at.AudioSignal(input_audio)
+
+    print(f"running standard vampnet with {num_vamps} vamps")
+    zv = interface.coarse_vamp_v2(
+        sig, 
+        sampling_steps=num_steps,
+        temperature=(init_temp, final_temp),
+        prefix_dur_s=0.0,
+        suffix_dur_s=0.0,
+        num_vamps=1,
+        downsample_factor=mask_periodic_amt,
+        periodic_width=mask_periodic_width,
+        periodic_dropout=0.0,
+        periodic_width_dropout=0.0,
+        n_conditioning_codebooks=None,
+        intensity=1.0,
+        ext_mask=None, 
+    )
+
+    zv = interface.coarse_to_fine(zv)
+
+    sig = interface.to_signal(zv).cpu()
+    print("done")
+
+    out_dir = OUT_DIR / str(uuid.uuid4())
+    out_dir.mkdir()
+    sig.write(out_dir / "output.wav")
+    # mask.write(out_dir / "mask.wav")
+    # return sig.path_to_file, mask.path_to_file
+    return sig.path_to_file
+
+
 def vamp(
     input_audio, init_temp, final_temp, 
     prefix_s, suffix_s, rand_mask_intensity,
@@ -68,7 +105,7 @@ def vamp(
     num_vamps, mode, use_beats, num_steps, snap_to_beats, 
     beat_unmask_drop,  mask_periodic_width, 
     mask_periodic_dropout, mask_periodic_width_dropout, 
-    n_conditioning_codebooks
+    n_conditioning_codebooks, use_coarse2fine
 ):
     # try:
         print(input_audio)
@@ -119,36 +156,19 @@ def vamp(
                 return_mask=True
             )
     
-            zv = interface.coarse_to_fine(zv)
-            mask = interface.to_signal(mask_z).cpu()
+            if use_coarse2fine: 
+                zv = interface.coarse_to_fine(zv)
+            # mask = interface.to_signal(mask_z).cpu()
 
             sig = interface.to_signal(zv).cpu()
             print("done")
-        elif mode == "loop":
-            print(f"running loop vampnet with {num_vamps} vamps")
-            sig, mask = interface.loop(
-                sig, 
-                temperature=(init_temp, final_temp),
-                prefix_dur_s=prefix_s, 
-                suffix_dur_s=prefix_s, # suffix should be same length as prefix 
-                num_loops=num_vamps,
-                downsample_factor=mask_periodic_amt,
-                periodic_width=mask_periodic_width,
-                intensity=rand_mask_intensity,
-                ext_mask=beat_mask, 
-                verbose=True,
-                return_mask=True
-            )
-            sig = sig.cpu()
-            mask = mask.cpu()
-            print("done")
-
 
         out_dir = OUT_DIR / str(uuid.uuid4())
         out_dir.mkdir()
         sig.write(out_dir / "output.wav")
-        mask.write(out_dir / "mask.wav")
-        return sig.path_to_file, mask.path_to_file
+        # mask.write(out_dir / "mask.wav")
+        # return sig.path_to_file, mask.path_to_file
+        return sig.path_to_file, None
     # except Exception as e:
     #     raise gr.Error(f"failed with error: {e}")
         
@@ -160,7 +180,7 @@ def save_vamp(
     mask_up_chk, up_factor, 
     num_vamps, mode, output_audio, notes, use_beats, num_steps, snap_to_beats,
     beat_unmask_drop, mask_periodic_width, mask_periodic_dropout, mask_periodic_width_dropout, 
-    n_conditioning_codebooks
+    n_conditioning_codebooks, use_coarse2fine
 ):
     out_dir = OUT_DIR / "saved" / str(uuid.uuid4())
     out_dir.mkdir(parents=True, exist_ok=True)
@@ -193,7 +213,8 @@ def save_vamp(
         "mask_periodic_width": mask_periodic_width, 
         "mask_periodic_dropout": mask_periodic_dropout,
         "mask_periodic_width_dropout": mask_periodic_width_dropout, 
-        "n_conditioning_codebooks": n_conditioning_codebooks
+        "n_conditioning_codebooks": n_conditioning_codebooks, 
+        "use_coarse2fine": use_coarse2fine,
     }
 
     # save with yaml
@@ -260,6 +281,10 @@ with gr.Blocks() as demo:
                 choices=["standard",],
                 value="standard"
             )
+            use_coarse2fine = gr.Checkbox(
+                label="use coarse2fine",
+                value=True
+            )
             num_vamps = gr.Number(
                 label="number of vamps. more vamps = longer generated audio",
                 value=1,
@@ -476,9 +501,10 @@ with gr.Blocks() as demo:
             num_vamps, mode, use_beats, num_steps, snap_to_beats, 
             beat_unmask_drop, mask_periodic_width, 
             mask_periodic_dropout, mask_periodic_width_dropout, 
-            n_conditioning_codebooks
+            n_conditioning_codebooks, use_coarse2fine
         ],
-        outputs=[output_audio, audio_mask]
+        outputs=[output_audio, audio_mask], 
+        api_name="vamp"
     )
 
     save_button.click(
@@ -494,9 +520,18 @@ with gr.Blocks() as demo:
             notes_text, use_beats, num_steps, snap_to_beats, 
             beat_unmask_drop, mask_periodic_width, 
             mask_periodic_dropout, mask_periodic_width_dropout, 
-            n_conditioning_codebooks
+            n_conditioning_codebooks, use_coarse2fine
         ],
         outputs=[thank_you, download_file]
     )
 
-demo.launch(share=True, enable_queue=True)
+    ez_vamp_button = gr.Button("ez vamp")
+    ez_vamp_button.click(
+        fn=ez_vamp,
+        inputs=[input_audio, init_temp, final_temp, mask_periodic_amt,
+                mask_periodic_width, num_steps ],
+        outputs=[output_audio],
+        api_name="ez_vamp"
+    )
+
+demo.launch(share=True, enable_queue=False, debug=True)

vampnet/interface.py

@@ -20,7 +20,7 @@ def signal_concat(
     return AudioSignal(audio_data, sample_rate=audio_signals[0].sample_rate)
 
 
-class Interface:
+class Interface(torch.nn.Module):
     def __init__(
         self,
         coarse_ckpt: str = None,
@@ -31,6 +31,7 @@ class Interface:
         coarse_chunk_size_s: int =  5, 
         coarse2fine_chunk_size_s: int =  3,
     ):
+        super().__init__()
         assert codec_ckpt is not None, "must provide a codec checkpoint"
         self.codec = LAC.load(Path(codec_ckpt))
         self.codec.eval()
@@ -240,103 +241,7 @@ class Interface:
         fine_z = torch.cat(fine_z, dim=-1)
         return fine_z[:, :, :length].clone()
     
-    def coarse_vamp(
-        self, 
-        signal, 
-        prefix_dur_s: float = 1.25, 
-        suffix_dur_s: float = 1.25,
-        num_loops: int = 3,
-        mode="impute",
-        downsample_factor: int = None,
-        debug=False,
-        **kwargs
-    ):
-        z = self.encode(signal)
-
-        assert signal.duration == self.coarse.chunk_size_s, "signal duration must match coarse chunk size for now"
-
-        # coarse z
-        cz = z[:, : self.coarse.n_codebooks, :].clone()
-        c_seq_len = cz.shape[-1]
-        n_prefix = self.s2t(prefix_dur_s)
-        n_suffix = self.s2t(suffix_dur_s)
-        
-        # we'll keep the final codes sequence here
-        c_vamp = {
-            'prefix': [cz[:, :, :n_prefix].clone()],
-            'suffix': [cz[:, :, c_seq_len-n_suffix:].clone()]
-        }
-
-        _cz = cz.clone()
-        for _ in range(num_loops):
-            # add noise
-            cz_masked, cz_mask = self.coarse.add_noise(
-                _cz, r=0.0,
-                n_prefix=n_prefix,
-                n_suffix=n_suffix, 
-                downsample_factor=downsample_factor
-            )
-            if debug:
-                print("tokens to infer")
-                self.to_signal(cz_masked).cpu().widget()
-
-            # sample!
-            cz_sampled = self.coarse.sample(
-                codec=self.codec,
-                time_steps=self.s2t(self.coarse.chunk_size_s),
-                start_tokens=_cz,
-                mask=cz_mask, 
-                return_signal=False,
-                **kwargs
-            )
-
-            if debug:
-                print("tokens sampled")
-                self.to_signal(cz_sampled).cpu().widget()
-            
-            cz_imputed = cz_sampled[:, :, n_prefix:c_seq_len-n_suffix].clone()
-            
-            if mode == "impute":
-                 # split the imputed codes into two halves
-                cz_imputed_a = cz_imputed[:, :, : cz_imputed.shape[-1] // 2].clone()
-                cz_imputed_b = cz_imputed[:, :, cz_imputed.shape[-1] // 2 :].clone()
-            elif mode == "continue":
-                cz_imputed_a = cz_imputed[:, :, : cz_imputed.shape[-1]].clone()
-                cz_imputed_b = _cz[:, :, :0].clone() # empty 
-            elif mode == "reverse-continue":
-                cz_imputed_a = _cz[:, :, :0].clone() # empty
-                cz_imputed_b = cz_imputed[:, :, : cz_imputed.shape[-1]].clone()
-            else:
-                raise ValueError(f"mode {mode} not supported")
-
-            if debug:
-                # add to our c_vamp
-                if cz_imputed_a.shape[-1] > 0:
-                    print("new_prefix added")
-                    self.to_signal(cz_imputed_a).cpu().widget()
-                if cz_imputed_b.shape[-1] >  0:
-                    print("new_suffix added")   
-                    self.to_signal(cz_imputed_b).cpu().widget()
-
-            c_vamp['prefix'].append(cz_imputed_a.clone())
-            c_vamp['suffix'].insert(0, cz_imputed_b.clone())
-
-            n_to_insert = c_seq_len - (cz_imputed_a.shape[-1] + cz_imputed_b.shape[-1])
-            to_insert = torch.zeros(cz_imputed_a.shape[0], cz_imputed_a.shape[1], n_to_insert).long().to(self.device)
-            _cz = torch.cat([cz_imputed_a, to_insert, cz_imputed_b], dim=-1)
-
-            if debug:
-                print("tokens to infer next round (area to insert in the middle)")
-                self.to_signal(_cz).cpu().widget()
-
-
-
-
-        prefix_codes = torch.cat(c_vamp['prefix'], dim=-1)
-        suffix_codes = torch.cat(c_vamp['suffix'], dim=-1)
-        c_vamp = torch.cat([prefix_codes, suffix_codes], dim=-1)
-        return c_vamp
-
+ 
     def coarse_vamp_v2(
         self, 
         signal, 
@@ -390,6 +295,7 @@ class Interface:
                 downsample_factor=downsample_factor,
                 periodic_width=periodic_width,
                 periodic_dropout=periodic_dropout,
+                add_random_periodic_offset=True,
                 periodic_width_dropout=periodic_width_dropout,
                 mask=cz_mask, 
                 ext_mask=ext_mask, 

vampnet/modules/base.py

@@ -44,6 +44,7 @@ class VampBase(at.ml.BaseModel):
         periodic_width: int = 1,
         periodic_width_dropout: float = 0.0,
         periodic_dropout: float = 0.0,
+        add_random_periodic_offset: bool = False,  # TODO: should be always false lol this is hacky
         n_conditioning_codebooks: Optional[int] = None,
         noise_mode: str = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -84,6 +85,7 @@ class VampBase(at.ml.BaseModel):
                                 if torch.bernoulli(torch.tensor(periodic_dropout)).item() == 1:
                                     # if we win, skip
                                     continue
+
                             # figure out how wide the mask should be
                             j_start = max(0, j - periodic_width // 2)
                             j_end = min(probs.shape[-1] - 1, j + periodic_width // 2) + 1
@@ -92,6 +94,10 @@ class VampBase(at.ml.BaseModel):
                             j_fill = torch.ones_like(j_mask) * (1 - j_mask)
                             # fill
                             probs[i, :, j_start:j_end] = 1 - j_fill
+                if add_random_periodic_offset:
+                    # add a random offset to the mask
+                    offset = torch.randint(0, downsample_factor[0], (1,))
+                    probs = torch.roll(probs, offset.item(), dims=-1)
 
             mask = torch.bernoulli(probs)
             mask = mask.round().long()