moving out

demo.py

@@ -65,7 +65,10 @@ def vamp(
     mask_periodic_amt, beat_unmask_dur,
     mask_dwn_chk, dwn_factor,
     mask_up_chk, up_factor, 
-    num_vamps, mode, use_beats, num_steps, snap_to_beats
+    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
 ):
     # try:
         print(input_audio)
@@ -89,7 +92,7 @@ def vamp(
                 mask_upbeats=mask_up_chk,
                 downbeat_downsample_factor=dwn_factor if dwn_factor > 0 else None, 
                 beat_downsample_factor=up_factor if up_factor > 0 else None,
-                dropout=0.7, 
+                dropout=beat_unmask_drop, 
                 invert=True
             )
             print(beat_mask)
@@ -106,6 +109,10 @@ def vamp(
                 suffix_dur_s=suffix_s,
                 num_vamps=num_vamps,
                 downsample_factor=mask_periodic_amt,
+                periodic_width=mask_periodic_width,
+                periodic_dropout=mask_periodic_dropout,
+                periodic_width_dropout=mask_periodic_width_dropout,
+                n_conditioning_codebooks=n_conditioning_codebooks if n_conditioning_codebooks > 0 else None,
                 intensity=rand_mask_intensity,
                 ext_mask=beat_mask, 
                 verbose=True,
@@ -126,6 +133,7 @@ def vamp(
                 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,
@@ -150,7 +158,9 @@ def save_vamp(
     mask_periodic_amt, beat_unmask_dur,
     mask_dwn_chk, dwn_factor,
     mask_up_chk, up_factor, 
-    num_vamps, mode, output_audio, notes, use_beats, num_steps, snap_to_beats
+    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
 ):
     out_dir = OUT_DIR / "saved" / str(uuid.uuid4())
     out_dir.mkdir(parents=True, exist_ok=True)
@@ -179,6 +189,11 @@ def save_vamp(
         "snap_to_beats": snap_to_beats,
         "mode": mode,
         "notes": notes,
+        "beat_unmask_drop": beat_unmask_drop,
+        "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
     }
 
     # save with yaml
@@ -287,6 +302,12 @@ with gr.Blocks() as demo:
         # mask settings
         with gr.Column():
 
+            n_conditioning_codebooks = gr.Number(
+                label="number of conditioning codebooks. probably 0", 
+                value=0,
+                precision=0,
+            )
+
             mask_periodic_amt = gr.Slider(
                 label="periodic hint  (0.0 means no hint, 2 - lots of hints, 8 - a couple of hints, 16 - occasional hint, 32 - very occasional hint, etc)",
                 minimum=0,
@@ -294,11 +315,29 @@ with gr.Blocks() as demo:
                 step=1,
                 value=9, 
             )
-
+            mask_periodic_width = gr.Slider(
+                label="periodic hint width (steps, 1 step ~= 10milliseconds",
+                minimum=1,
+                maximum=100,
+                step=1,
+                value=1,
+            )
+            mask_periodic_dropout = gr.Slider(
+                label="periodic hint dropout (0.0 means no dropout, 1.0 means all dropout)",
+                minimum=0.0,
+                maximum=1.0,
+                value=0.0,
+            )
+            mask_periodic_width_dropout = gr.Slider(    
+                label="periodic hint width dropout (0.0 means no dropout, 1.0 means all dropout)",
+                minimum=0.0,
+                maximum=1.0,
+                value=0.0,
+            )
 
             rand_mask_intensity = gr.Slider(
                 label="random mask intensity. (If this is less than 1, scatters tiny hints throughout the audio, should be between 0.9 and 1.0)",
-                minimum=0.0,
+                minimum=0.8,
                 maximum=1.0,
                 value=1.0
             )
@@ -343,7 +382,7 @@ with gr.Blocks() as demo:
 
             num_steps = gr.Slider(
                 label="number of steps (should normally be between 12 and 36)",
-                minimum=4,
+                minimum=1,
                 maximum=128,
                 step=1,
                 value=36
@@ -379,6 +418,13 @@ with gr.Blocks() as demo:
                     maximum=3.0,
                     value=0.07
                 )
+                beat_unmask_drop = gr.Slider(
+                    label="dropout (within beat)", 
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.0
+                )
+
                 with gr.Accordion("downbeat settings", open=False):
                     mask_dwn_chk = gr.Checkbox(
                         label="hint downbeats",
@@ -427,7 +473,10 @@ with gr.Blocks() as demo:
             mask_periodic_amt, beat_unmask_dur, 
             mask_dwn_chk, dwn_factor, 
             mask_up_chk, up_factor, 
-            num_vamps, mode, use_beats, num_steps, snap_to_beats
+            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
         ],
         outputs=[output_audio, audio_mask]
     )
@@ -442,7 +491,10 @@ with gr.Blocks() as demo:
             mask_up_chk, up_factor,
             num_vamps, mode,
             output_audio,
-            notes_text, use_beats, num_steps, snap_to_beats
+            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
         ],
         outputs=[thank_you, download_file]
     )

scripts/utils/vamp_folder.py

@@ -220,7 +220,7 @@ EXP_REGISTRY["mask-ratio"] = {
 
 EXP_REGISTRY["sampling-steps"] = {
     "codec": reconstructed,
-    **{f"steps_{n}": num_sampling_steps(n)  for n in [1, 4, 12, 24, 36, 64, 72, 128]},
+    **{f"steps_{n}": num_sampling_steps(n)  for n in [1, 4, 12, 24, 36, 64, 72]},
 }
 
 EXP_REGISTRY["baseline"] = {

vampnet/interface.py

@@ -134,7 +134,7 @@ class Interface:
             mask_upbeats: bool = True,
             downbeat_downsample_factor: int = None,
             beat_downsample_factor: int = None,
-            dropout: float = 0.3,
+            dropout: float = 0.0,
             invert: bool = True,
     ):
         """make a beat synced mask. that is, make a mask that 
@@ -182,7 +182,8 @@ class Interface:
                 _slice = int(beat_idx - mask_b4), int(beat_idx + mask_after)
                 num_steps = mask[_slice[0]:_slice[1]].shape[0]
                 _m = torch.ones(num_steps, device=self.device)
-                _m = torch.nn.functional.dropout(_m, p=dropout)
+                _m_mask = torch.bernoulli(_m * (1 - dropout))
+                _m = _m * _m_mask.long()
                 
                 mask[_slice[0]:_slice[1]] = _m
 
@@ -191,7 +192,8 @@ class Interface:
                 _slice = int(downbeat_idx - mask_b4), int(downbeat_idx + mask_after)
                 num_steps = mask[_slice[0]:_slice[1]].shape[0]
                 _m = torch.ones(num_steps, device=self.device)
-                _m = torch.nn.functional.dropout(_m, p=dropout)
+                _m_mask = torch.bernoulli(_m * (1 - dropout))
+                _m = _m * _m_mask.long()
                 
                 mask[_slice[0]:_slice[1]] = _m
         
@@ -342,6 +344,9 @@ class Interface:
         suffix_dur_s: float = 0.0,
         num_vamps: int = 1,
         downsample_factor: int = None,
+        periodic_width: int = 1,
+        periodic_dropout=0.0,
+        periodic_width_dropout=0.0, 
         intensity: float = 1.0, 
         debug=False,
         swap_prefix_suffix=False, 
@@ -383,6 +388,9 @@ class Interface:
                 n_prefix=n_prefix,
                 n_suffix=n_suffix, 
                 downsample_factor=downsample_factor,
+                periodic_width=periodic_width,
+                periodic_dropout=periodic_dropout,
+                periodic_width_dropout=periodic_width_dropout,
                 mask=cz_mask, 
                 ext_mask=ext_mask, 
                 n_conditioning_codebooks=n_conditioning_codebooks
@@ -481,8 +489,11 @@ class Interface:
         suffix_codes = torch.cat(c_vamp['suffix'], dim=-1)
         c_vamp = torch.cat([prefix_codes, suffix_codes], dim=-1)
 
+        # replace the mask token in cz_masked with random tokens
+        # so that we can decode it
         if return_mask:
             return c_vamp, cz_masked
+        
         return c_vamp
 
     # create a variation of an audio signal

vampnet/modules/base.py

@@ -41,6 +41,9 @@ class VampBase(at.ml.BaseModel):
         n_prefix: Optional[torch.Tensor] = None,
         n_suffix: Optional[torch.Tensor] = None,
         downsample_factor: Optional[int] = None, 
+        periodic_width: int = 1,
+        periodic_width_dropout: float = 0.0,
+        periodic_dropout: float = 0.0,
         n_conditioning_codebooks: Optional[int] = None,
         noise_mode: str = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -75,7 +78,20 @@ class VampBase(at.ml.BaseModel):
                         continue
                     for j in range(probs.shape[-1]):
                         if j % factor == 0:
-                            probs[i, :, j] = 0.0
+                            # if we have periodic dropout
+                            if periodic_dropout > 0:
+                                # flip a coin
+                                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
+                            # flip a coin for each position in the mask
+                            j_mask = torch.bernoulli(torch.ones(j_end - j_start) * periodic_width_dropout)
+                            j_fill = torch.ones_like(j_mask) * (1 - j_mask)
+                            # fill
+                            probs[i, :, j_start:j_end] = 1 - j_fill
 
             mask = torch.bernoulli(probs)
             mask = mask.round().long()
@@ -175,6 +191,14 @@ class VampBase(at.ml.BaseModel):
             codec.sample_rate,
         )
 
+        # find where the mask token is and replace it with silence in the audio
+        for tstep in range(z.shape[-1]):
+            if torch.any(z[:, :, tstep] == self.mask_token):
+                print("mask token found at step", tstep)
+                sample_idx_0 = tstep * codec.hop_length
+                sample_idx_1 = sample_idx_0 + codec.hop_length
+                signal.samples[:, :, sample_idx_0:sample_idx_1] = 0.0
+
         return signal
 
     @torch.no_grad()