several mods

TODOS

@@ -0,0 +1 @@
+[ ] add sketch2sound finetuning

app.py

@@ -21,6 +21,7 @@ interface = Interface.default()
 init_model_choice = open("DEFAULT_MODEL").read().strip()
 # load the init model
 interface.load_finetuned(init_model_choice)
+interface.to(device)
     
 def to_output(sig):
     return sig.sample_rate, sig.cpu().detach().numpy()[0][0]
@@ -105,9 +106,33 @@ def _vamp(
         n_mask_codebooks, periodic_w, onset_mask_width, 
         dropout, sampletemp, typical_filtering, 
         typical_mass, typical_min_tokens, top_p, 
-        sample_cutoff, stretch_factor, api=False
+        sample_cutoff, stretch_factor, sampling_steps, beat_mask_ms, num_feedback_steps, api=False
     ):
 
+    print("args!")
+    print(f"seed: {seed}")
+    print(f"input_audio: {input_audio}")
+    print(f"model_choice: {model_choice}")
+    print(f"pitch_shift_amt: {pitch_shift_amt}")
+    print(f"periodic_p: {periodic_p}")
+    print(f"n_mask_codebooks: {n_mask_codebooks}")
+    print(f"periodic_w: {periodic_w}")
+    print(f"onset_mask_width: {onset_mask_width}")
+    print(f"dropout: {dropout}")
+    print(f"sampletemp: {sampletemp}")
+    print(f"typical_filtering: {typical_filtering}")
+    print(f"typical_mass: {typical_mass}")
+    print(f"typical_min_tokens: {typical_min_tokens}")
+    print(f"top_p: {top_p}")
+    print(f"sample_cutoff: {sample_cutoff}")
+    print(f"stretch_factor: {stretch_factor}")
+    print(f"sampling_steps: {sampling_steps}")
+    print(f"api: {api}")
+    print(f"beat_mask_ms: {beat_mask_ms}")
+    print(f"using device {interface.device}")
+    print(f"num feedback steps: {num_feedback_steps}")
+
+
     t0 = time.time()
     interface.to("cuda" if torch.cuda.is_available() else "cpu")
     print(f"using device {interface.device}")
@@ -121,6 +146,9 @@ def _vamp(
     
     sig = at.AudioSignal(input_audio, sr).to_mono()
 
+    loudness = sig.loudness()
+    sig = interface._preprocess(sig)
+
     # reload the model if necessary
     interface.load_finetuned(model_choice)
 
@@ -129,38 +157,70 @@ def _vamp(
 
     codes = interface.encode(sig)
 
-    mask = new_vampnet_mask(
-        interface, 
-        codes, 
-        onset_idxs=onsets(sig, hop_length=interface.codec.hop_length),
-        width=onset_mask_width,
+    # mask = new_vampnet_mask(
+    #     interface, 
+    #     codes, 
+    #     onset_idxs=onsets(sig, hop_length=interface.codec.hop_length),
+    #     width=onset_mask_width,
+    #     periodic_prompt=periodic_p,
+    #     upper_codebook_mask=n_mask_codebooks,
+    #     drop_amt=dropout
+    # ).long()
+
+    
+    mask = interface.build_mask(
+        codes,
+        sig=sig, 
         periodic_prompt=periodic_p,
+        periodic_prompt_width=periodic_w,
+        onset_mask_width=onset_mask_width,
+        _dropout=dropout,
         upper_codebook_mask=n_mask_codebooks,
-        drop_amt=dropout
-    ).long()
+    )
+    if beat_mask_ms > 0:
+        # bm = pmask.mask_or(
+        #     pmask.periodic_mask(
+        #         codes, periodic_p, periodic_w, random_roll=False
+        #     ),
+        # )
+        mask = pmask.mask_and(
+            mask, interface.make_beat_mask(
+                sig, after_beat_s=beat_mask_ms/1000.,
+            )
+        )
+        mask = pmask.codebook_mask(mask, n_mask_codebooks)
+    np.savetxt("scratch/rms_mask.txt", mask[0].cpu().numpy(), fmt='%d')
 
-    # save the mask as a txt file
     interface.set_chunk_size(10.0)
+
+    # lord help me
+    if top_p is not None:
+        if top_p > 0:
+            pass
+        else:
+            top_p = None
+
     codes, mask = interface.vamp(
         codes, mask,
-        batch_size=1 if api else 1,
-        feedback_steps=1,
-        _sampling_steps=12 if sig.duration <6.0 else 24,
+        batch_size=2,
+        feedback_steps=num_feedback_steps,
+        _sampling_steps=sampling_steps,
         time_stretch_factor=stretch_factor,
         return_mask=True,
         temperature=sampletemp,
         typical_filtering=typical_filtering, 
         typical_mass=typical_mass, 
         typical_min_tokens=typical_min_tokens, 
-        top_p=None,
+        top_p=top_p,
         seed=_seed,
-        sample_cutoff=1.0,
+        sample_cutoff=sample_cutoff,
     )
     print(f"vamp took {time.time() - t0} seconds")
 
     sig = interface.decode(codes)
+    sig = sig.normalize(loudness)
 
-    return to_output(sig)
+    return to_output(sig[0]), to_output(sig[1])
 
 def vamp(data):
     return _vamp(
@@ -180,31 +240,29 @@ def vamp(data):
         top_p=data[top_p],
         sample_cutoff=data[sample_cutoff],
         stretch_factor=data[stretch_factor],
+        sampling_steps=data[sampling_steps],
+        beat_mask_ms=data[beat_mask_ms],
+        num_feedback_steps=data[num_feedback_steps],
         api=False, 
     )
 
-# def api_vamp(data):
-#     return _vamp(
-#         seed=data[seed], 
-#         input_audio=data[input_audio],
-#         model_choice=data[model_choice],
-#         pitch_shift_amt=data[pitch_shift_amt],
-#         periodic_p=data[periodic_p],
-#         n_mask_codebooks=data[n_mask_codebooks],
-#         periodic_w=data[periodic_w],
-#         onset_mask_width=data[onset_mask_width],
-#         dropout=data[dropout],
-#         sampletemp=data[sampletemp],
-#         typical_filtering=data[typical_filtering],
-#         typical_mass=data[typical_mass],
-#         typical_min_tokens=data[typical_min_tokens],
-#         top_p=data[top_p],
-#         sample_cutoff=data[sample_cutoff],
-#         stretch_factor=data[stretch_factor],
-#         api=True, 
-#     )
-
-def api_vamp(input_audio, sampletemp, top_p, periodic_p, periodic_w, dropout, stretch_factor, onset_mask_width, typical_filtering, typical_mass, typical_min_tokens, seed, model_choice, n_mask_codebooks, pitch_shift_amt, sample_cutoff):
+
+def api_vamp(input_audio, 
+                sampletemp, top_p, 
+                periodic_p, periodic_w, 
+                dropout, 
+                stretch_factor,
+                onset_mask_width,
+                typical_filtering,
+                typical_mass,
+                typical_min_tokens,
+                seed,
+                model_choice,
+                n_mask_codebooks,
+                pitch_shift_amt,
+                sample_cutoff, 
+                sampling_steps, 
+                beat_mask_ms, num_feedback_steps):
     return _vamp(
         seed=seed, 
         input_audio=input_audio,
@@ -222,50 +280,12 @@ def api_vamp(input_audio, sampletemp, top_p, periodic_p, periodic_w, dropout, st
         top_p=top_p,
         sample_cutoff=sample_cutoff,
         stretch_factor=stretch_factor,
+        sampling_steps=sampling_steps,
+        beat_mask_ms=beat_mask_ms,
+        num_feedback_steps=num_feedback_steps,
         api=True, 
     )
 
-OUT_DIR = Path("gradio-outputs")
-OUT_DIR.mkdir(exist_ok=True)
-def harp_vamp(input_audio_file, periodic_p, n_mask_codebooks):
-    sig = at.AudioSignal(input_audio_file)
-    sr, samples = sig.sample_rate, sig.samples[0][0].detach().cpu().numpy()
-    # convert to int32
-    samples = (samples * np.iinfo(np.int32).max).astype(np.int32)
-    sr, samples =  _vamp(
-        seed=0,
-        input_audio=(sr, samples),
-        model_choice=init_model_choice,
-        pitch_shift_amt=0,
-        periodic_p=periodic_p,
-        n_mask_codebooks=n_mask_codebooks,
-        periodic_w=1,
-        onset_mask_width=0,
-        dropout=0.0,
-        sampletemp=1.0,
-        typical_filtering=True,
-        typical_mass=0.15,  
-        typical_min_tokens=64,
-        top_p=0.0,
-        sample_cutoff=1.0,
-        stretch_factor=1,
-    )
-    
-    sig = at.AudioSignal(samples, sr)
-    # write to file
-    # clear the outdir
-    for p in OUT_DIR.glob("*"):
-        p.unlink()
-    OUT_DIR.mkdir(exist_ok=True)
-    # outpath = OUT_DIR / f"{uuid.uuid4()}.wav"
-    from pyharp import AudioLabel, LabelList, save_audio
-    outpath = save_audio(sig)
-    sig.write(outpath)
-    output_labels = LabelList()
-    output_labels.append(AudioLabel(label='~', t=0.0, amplitude=0.5, description='generated audio'))
-    return outpath, output_labels
-    
-
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
@@ -359,6 +379,11 @@ with gr.Blocks() as demo:
                     value=1,
                 )
 
+                beat_mask_ms = gr.Number(
+                    label="beat mask width (milliseconds)",
+                    value=0,
+                )
+
 
             with gr.Accordion("sampling settings", open=False):
                 sampletemp = gr.Slider(
@@ -399,6 +424,22 @@ with gr.Blocks() as demo:
                     value=1.0, 
                     step=0.01
                 )
+                sampling_steps = gr.Slider(
+                    label="sampling steps",
+                    minimum=1,
+                    maximum=128,
+                    step=1,
+                    value=36
+                )
+                num_feedback_steps = gr.Slider(
+                    label="feedback steps",
+                    minimum=1,
+                    maximum=16,
+                    step=1,
+                    value=1
+                )
+
+
 
 
             dropout = gr.Slider(
@@ -433,7 +474,7 @@ with gr.Blocks() as demo:
 
             audio_outs = []
             use_as_input_btns = []
-            for i in range(1):
+            for i in range(2):
                 with gr.Column():
                     audio_outs.append(gr.Audio(
                         label=f"output audio {i+1}",
@@ -466,13 +507,16 @@ with gr.Blocks() as demo:
             n_mask_codebooks,
             pitch_shift_amt, 
             sample_cutoff, 
+            sampling_steps, 
+            beat_mask_ms,
+            num_feedback_steps
         }
   
     # connect widgets
     vamp_button.click(
         fn=vamp,
         inputs=_inputs,
-        outputs=[audio_outs[0]], 
+        outputs=[audio_outs[0], audio_outs[1]], 
     )
 
     api_vamp_button = gr.Button("api vamp", visible=True)
@@ -491,31 +535,15 @@ with gr.Blocks() as demo:
                 model_choice,
                 n_mask_codebooks,
                 pitch_shift_amt,
-                sample_cutoff
+                sample_cutoff, 
+                sampling_steps, 
+                beat_mask_ms,
+                num_feedback_steps
         ], 
-        outputs=[audio_outs[0]],
+        outputs=[audio_outs[0], audio_outs[1]],
         api_name="vamp"
     )
 
-    from pyharp import ModelCard, build_endpoint
-    card = ModelCard(
-        name="vampnet", 
-        description="vampnet! is a model for generating audio from audio",
-        author="hugo flores garcía", 
-        tags=["music generation"], 
-        midi_in=False, 
-        midi_out=False
-    )
-        
-    # Build a HARP-compatible endpoint
-    app = build_endpoint(model_card=card,
-                         components=[
-                            periodic_p, 
-                            n_mask_codebooks,
-                        ],
-                         process_fn=harp_vamp)
-
-
 
 try:
     demo.queue()

conf/generated/ivo/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+fine_tune_checkpoint: ./models/vampnet/c2f.pth
+save_path: ./runs/ivo/c2f
+train/AudioLoader.sources: &id001
+- ./scratch/miguel/ivo/separated
+val/AudioLoader.sources: *id001

conf/generated/ivo/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+fine_tune_checkpoint: ./models/vampnet/coarse.pth
+save_path: ./runs/ivo/coarse
+train/AudioLoader.sources: &id001
+- ./scratch/miguel/ivo/separated
+val/AudioLoader.sources: *id001

conf/generated/ivo/interface.yml

@@ -0,0 +1,6 @@
+AudioLoader.sources:
+- - ./scratch/miguel/ivo/separated
+Interface.coarse2fine_ckpt: ./runs/ivo/c2f/latest/vampnet/weights.pth
+Interface.coarse_ckpt: ./runs/ivo/coarse/latest/vampnet/weights.pth
+Interface.codec_ckpt: ./models/vampnet/codec.pth
+Interface.wavebeat_ckpt: ./models/wavebeat.pth

conf/generated/lazaro-ros-sep/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+fine_tune_checkpoint: ./models/vampnet/c2f.pth
+save_path: ./runs/lazaro-ros-sep/c2f
+train/AudioLoader.sources: &id001
+- ./scratch/miguel/lazaro-ros/separated
+val/AudioLoader.sources: *id001

conf/generated/lazaro-ros-sep/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+fine_tune_checkpoint: ./models/vampnet/coarse.pth
+save_path: ./runs/lazaro-ros-sep/coarse
+train/AudioLoader.sources: &id001
+- ./scratch/miguel/lazaro-ros/separated
+val/AudioLoader.sources: *id001

conf/generated/lazaro-ros-sep/interface.yml

@@ -0,0 +1,6 @@
+AudioLoader.sources:
+- - ./scratch/miguel/lazaro-ros/separated
+Interface.coarse2fine_ckpt: ./runs/lazaro-ros-sep/c2f/latest/vampnet/weights.pth
+Interface.coarse_ckpt: ./runs/lazaro-ros-sep/coarse/latest/vampnet/weights.pth
+Interface.codec_ckpt: ./models/vampnet/codec.pth
+Interface.wavebeat_ckpt: ./models/wavebeat.pth

conf/generated/lazaro-ros/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+fine_tune_checkpoint: ./models/vampnet/c2f.pth
+save_path: ./runs/lazaro-ros/c2f
+train/AudioLoader.sources: &id001
+- ./scratch/miguel/lazaro-ros
+val/AudioLoader.sources: *id001

conf/generated/lazaro-ros/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+fine_tune_checkpoint: ./models/vampnet/coarse.pth
+save_path: ./runs/lazaro-ros/coarse
+train/AudioLoader.sources: &id001
+- ./scratch/miguel/lazaro-ros
+val/AudioLoader.sources: *id001

conf/generated/lazaro-ros/interface.yml

@@ -0,0 +1,6 @@
+AudioLoader.sources:
+- - ./scratch/miguel/lazaro-ros
+Interface.coarse2fine_ckpt: ./runs/lazaro-ros/c2f/latest/vampnet/weights.pth
+Interface.coarse_ckpt: ./runs/lazaro-ros/coarse/latest/vampnet/weights.pth
+Interface.codec_ckpt: ./models/vampnet/codec.pth
+Interface.wavebeat_ckpt: ./models/wavebeat.pth

conf/generated/march-31/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+fine_tune_checkpoint: ./models/vampnet/c2f.pth
+save_path: ./runs/march-31/c2f
+train/AudioLoader.sources: &id001
+- sound-journal-march-31
+val/AudioLoader.sources: *id001

conf/generated/march-31/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+fine_tune_checkpoint: ./models/vampnet/coarse.pth
+save_path: ./runs/march-31/coarse
+train/AudioLoader.sources: &id001
+- sound-journal-march-31
+val/AudioLoader.sources: *id001

conf/generated/march-31/interface.yml

@@ -0,0 +1,6 @@
+AudioLoader.sources:
+- - sound-journal-march-31
+Interface.coarse2fine_ckpt: ./runs/march-31/c2f/latest/vampnet/weights.pth
+Interface.coarse_ckpt: ./runs/march-31/coarse/latest/vampnet/weights.pth
+Interface.codec_ckpt: ./models/vampnet/codec.pth
+Interface.wavebeat_ckpt: ./models/wavebeat.pth

conf/generated/sax-new/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+fine_tune_checkpoint: ./models/vampnet/c2f.pth
+save_path: ./runs/sax-new/c2f
+train/AudioLoader.sources: &id001
+- ./scratch/miguel/saxophone-new/
+val/AudioLoader.sources: *id001

conf/generated/sax-new/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+fine_tune_checkpoint: ./models/vampnet/coarse.pth
+save_path: ./runs/sax-new/coarse
+train/AudioLoader.sources: &id001
+- ./scratch/miguel/saxophone-new/
+val/AudioLoader.sources: *id001

conf/generated/sax-new/interface.yml

@@ -0,0 +1,6 @@
+AudioLoader.sources:
+- - ./scratch/miguel/saxophone-new/
+Interface.coarse2fine_ckpt: ./runs/sax-new/c2f/latest/vampnet/weights.pth
+Interface.coarse_ckpt: ./runs/sax-new/coarse/latest/vampnet/weights.pth
+Interface.codec_ckpt: ./models/vampnet/codec.pth
+Interface.wavebeat_ckpt: ./models/wavebeat.pth

conf/generated/saxophone/c2f.yml

@@ -0,0 +1,15 @@
+$include:
+- conf/lora/lora.yml
+AudioDataset.duration: 3.0
+AudioDataset.loudness_cutoff: -40.0
+VampNet.embedding_dim: 1280
+VampNet.n_codebooks: 14
+VampNet.n_conditioning_codebooks: 4
+VampNet.n_heads: 20
+VampNet.n_layers: 16
+fine_tune: true
+fine_tune_checkpoint: ./models/vampnet/c2f.pth
+save_path: ./runs/saxophone/c2f
+train/AudioLoader.sources: &id001
+- scratch/sounds
+val/AudioLoader.sources: *id001

conf/generated/saxophone/coarse.yml

@@ -0,0 +1,8 @@
+$include:
+- conf/lora/lora.yml
+fine_tune: true
+fine_tune_checkpoint: ./models/vampnet/coarse.pth
+save_path: ./runs/saxophone/coarse
+train/AudioLoader.sources: &id001
+- scratch/sounds
+val/AudioLoader.sources: *id001

conf/generated/saxophone/interface.yml

@@ -0,0 +1,6 @@
+AudioLoader.sources:
+- - scratch/sounds
+Interface.coarse2fine_ckpt: ./runs/saxophone/c2f/latest/vampnet/weights.pth
+Interface.coarse_ckpt: ./runs/saxophone/coarse/latest/vampnet/weights.pth
+Interface.codec_ckpt: ./models/vampnet/codec.pth
+Interface.wavebeat_ckpt: ./models/wavebeat.pth

conf/lora/lora-s2s.yml

@@ -0,0 +1,27 @@
+$include:
+  - conf/vampnet.yml
+
+fine_tune: True
+
+train/AudioDataset.n_examples: 100000000
+val/AudioDataset.n_examples: 500
+
+
+NoamScheduler.warmup: 500
+
+batch_size: 7
+num_workers: 7
+save_iters: [2000, 4000, 10000,20000, 40000, 100000]
+sample_freq: 2000
+val_freq: 1000
+
+AdamW.lr: 0.0001
+
+# let's us organize sound classes into folders and choose from those sound classes uniformly
+AudioDataset.without_replacement: False
+num_iters: 500000
+
+
+# control signals to use as conditioning. 
+Sketch2SoundController.ctrl_keys: ['rmsq16',]
+

conf/lora/lora.yml

@@ -19,4 +19,4 @@ AdamW.lr: 0.0001
 
 # let's us organize sound classes into folders and choose from those sound classes uniformly
 AudioDataset.without_replacement: False
-num_iters: 500000
+num_iters: 500000

scripts/exp/export.py

@@ -1,11 +1,10 @@
 from pathlib import Path
 
-run_dir = Path("runs/sample-instrument")
+run_dir = Path("runs/lazaro-ros-sep")
 name = run_dir.name
 
 repo_dir = Path("models/vampnet")
 
-
 for part in ("coarse", "c2f"):
     outdir = repo_dir / "loras" / name 
     outdir.mkdir(parents=True, exist_ok=True)
@@ -16,7 +15,7 @@ for part in ("coarse", "c2f"):
 
 # now, push to hub
 from huggingface_hub import Repository
-repo = Repository(repo_dir, git_user="hugofloresgarcia", git_email="[email protected]")
+repo = Repository(str(repo_dir),  git_user="hugofloresgarcia", git_email="[email protected]")
 repo.push_to_hub(
     commit_message=f"add {name}"
 )

scripts/exp/train.py

@@ -18,6 +18,7 @@ from torch.utils.tensorboard import SummaryWriter
 
 import vampnet
 from vampnet.modules.transformer import VampNet
+# from vampnet.control import Sketch2SoundController
 from vampnet.util import codebook_unflatten, codebook_flatten
 from vampnet import mask as pmask
 # from dac.model.dac import DAC
@@ -66,6 +67,8 @@ AudioDataset = argbind.bind(at.datasets.AudioDataset, "train", "val")
 
 IGNORE_INDEX = -100
 
+# Sketch2SoundController = argbind.bind(Sketch2SoundController)
+
 
 @argbind.bind("train", "val", without_prefix=True)
 def build_transform():
@@ -118,6 +121,36 @@ def add_num_params_repr_hook(model):
 
         setattr(m, "extra_repr", partial(num_params_hook, o=o, p=p))
 
+def get_controls(state, sig: at.AudioSignal):
+    # get controls
+    n_batch = sig.samples.shape[0]  
+    if state.controller is not None:
+        ctrls = state.controller.extract(sig)
+        # draw control masks
+        ctrl_masks = state.controller.random_mask(
+            ctrls, 
+            r=state.rng.draw(n_batch)[:, 0].to(state.device)
+        )
+    else:
+        ctrls = None
+        ctrl_masks = None
+    
+    return ctrls, ctrl_masks
+
+
+def generate_z_mask(state, z, vn, n_batch, ctrl_masks=None):
+    r = state.rng.draw(n_batch)[:, 0].to(state.device)
+
+    mask, ii = state.model.random_mask(z, r)
+    mask = pmask.codebook_unmask(mask, vn.n_conditioning_codebooks)
+
+    # outpaint? 
+    # if state.outpaint_prob > 0:
+    #     if flip_coin(state.outpaint_prob):
+    #         mask, ctrl_masks = state.build_tria_mask(mask, ctrl_masks)
+    z_mask = pmask.apply_mask(z, mask, vn.mask_token)
+    
+    return z_mask, mask, ii, r, ctrl_masks
 
 def accuracy(
     preds: torch.Tensor,
@@ -184,6 +217,8 @@ def _metrics(z_hat, r, target, flat_mask, output):
 class State:
     model: VampNet
     codec: DAC
+    # controller: Sketch2SoundController
+    controller: Optional[object]
 
     optimizer: AdamW
     scheduler: NoamScheduler
@@ -218,6 +253,11 @@ def train_loop(state: State, batch: dict, accel: Accelerator):
         mask = pmask.random(z, r)
         mask = pmask.codebook_unmask(mask, vn.n_conditioning_codebooks)
         z_mask, mask = pmask.apply_mask(z, mask, vn.mask_token)
+
+        # get controls
+        ctrls, ctrl_masks = get_controls(state, signal)
+
+        # TODO: KEEP INCORPORATING ZMASK CODE
         
         z_mask_latent = vn.embedding.from_codes(z_mask, state.codec)
 
@@ -266,6 +306,22 @@ def train_loop(state: State, batch: dict, accel: Accelerator):
 
     return {k: v for k, v in sorted(output.items())}
 
+# def get_controls(self, sig: sn.Signal, controller):
+#     # get controls
+#     n_batch = sig.wav.shape[0]  
+#     if self.controller is not None:
+#         ctrls = self.controller.extract(sig)
+#         # draw control masks
+#         ctrl_masks = self.controller.random_mask(
+#             ctrls, 
+#             r=self.rng.draw(n_batch)[:, 0].to(self.device)
+#         )
+#     else:
+#         ctrls = None
+#         ctrl_masks = None
+    
+#     return ctrls, ctrl_masks
+
 
 @timer()
 @torch.no_grad()
@@ -561,6 +617,8 @@ def load(
     # load the datasets
     train_data, val_data = build_datasets(args, sample_rate)
 
+    # controller = Sketch2SoundController(sample_rate=sample_rate, hop_length=codec.hop_length)
+
     return State(
         tracker=tracker,
         model=model,
@@ -572,6 +630,7 @@ def load(
         train_data=train_data,
         val_data=val_data,
         grad_clip_val=grad_clip_val,
+        controller=None,
     )
 
 
@@ -612,6 +671,7 @@ def train(
         tracker=tracker, 
         save_path=save_path)
     print("initialized state.")
+    state.device = accel.device
 
     train_dataloader = accel.prepare_dataloader(
         state.train_data,

token_telephone/tt.py

@@ -16,10 +16,25 @@ import numpy as np
 import torch
 from einops import rearrange
 
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 
+# ~~~~~~  configs!     ~~~~~~~~
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+MAX_LOUDNESS = -20
+MIN_LOUDNESS = -40
+COLS = 40
+ROWS = 13
+
+device = 'Scarlett 4i4 4th Gen'
+sample_rate = 48000
+num_channels = 4
+blocksize = 16384
+
 PROFILE = False
 DEBUG = False
 DEBUG_NO_VAMPNET = False
 set_debug(DEBUG)
+
 # if DEBUG:
 #     import gc
 #     # log when gc start and stops
@@ -80,19 +95,6 @@ Thread(target=draw_intro_screen).start()
 from audiotools import AudioSignal
 from vamp_helper import load_interface, ez_variation
 
-# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 
-# ~~~~~~  configs!     ~~~~~~~~
-# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-MAX_LOUDNESS = -20
-MIN_LOUDNESS = -40
-COLS = 40
-ROWS = 13
-
-device = 'Scarlett 4i4 4th Gen'
-sample_rate = 48000
-num_channels = 4
-blocksize = 16384
 
 
 # TODO: 

vampnet/beats.py

@@ -213,10 +213,11 @@ class WaveBeat(BeatTracker):
     def extract_beats(self, signal: AudioSignal) -> Tuple[np.ndarray, np.ndarray]:
         """returns beat and downbeat times, in  seconds"""
         # extract beats
+        self.model.to('cuda' if torch.cuda.is_available() else 'cpu')
         beats, downbeats = self.model.predict_beats_from_array(
             audio=signal.audio_data.squeeze(0),
             sr=signal.sample_rate,
-            use_gpu=self.device != "cpu",
+            use_gpu=torch.cuda.is_available(),
         )
 
         return beats, downbeats

vampnet/control.py

@@ -0,0 +1,277 @@
+from dataclasses import dataclass
+from functools import partial
+from typing import Optional
+
+from torch import nn
+
+import vampnet.dsp.signal as sn
+from vampnet.dsp.signal import Signal
+from vampnet.mask import random_along_time
+from torch import Tensor
+import torch
+
+
+class MedianFilterAugment(nn.Module):
+
+    def __init__(self, 
+        kernel_size: int, 
+        train_min: int = 1, 
+        train_max: int = 20,
+    ):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.train_min = train_min
+        self.train_max = train_max
+
+    def forward(self, x: Tensor) -> Tensor:
+        if self.training:
+            sizes = torch.randint(
+                self.train_min, 
+                self.train_max, 
+                size=(x.shape[0],)
+            )
+        else:
+            sizes = self.kernel_size
+        # print(f"median filter sizes: {sizes}")
+        return sn.median_filter_1d(x, sizes)
+
+class RMS(nn.Module):
+
+    def __init__(self, 
+        hop_length, 
+        window_length=2048, 
+        n_quantize=None, 
+        sample_rate=44100, 
+        median_filter_size: Optional[int] = None,
+        train_median_filter_min=1, 
+        train_median_filter_max=15,
+    ):
+        super().__init__()
+
+        self.hop_length = hop_length
+        self.window_length = window_length
+        self.n_quantize = n_quantize
+        self.sample_rate = sample_rate
+
+        self.mf = MedianFilterAugment(
+            kernel_size=median_filter_size, 
+            train_min=train_median_filter_min, 
+            train_max=train_median_filter_max
+        ) if median_filter_size is not None else None
+    
+    @property
+    def dim(self):
+        return 1
+
+    def extract(self, sig: Signal) -> Tensor:
+        rmsd = sn.rms(sig, 
+            window_length=self.window_length, 
+            hop_length=self.hop_length, 
+        )[:, :, :-1] # TODO: cutting the last frame to match DAC tokens but why :'(
+        nb, _, _ = rmsd.shape
+
+        if self.n_quantize is not None:
+            # standardize to 0-1
+            rmsd = (rmsd - rmsd.min()) / (rmsd.max() - rmsd.min())
+
+            # quantize to 128 steps
+            rmsd = torch.round(rmsd * self.n_quantize)
+            rmsd =  rmsd / self.n_quantize
+
+        if self.mf is not None:
+            rmsd = self.mf(rmsd)
+        
+        return rmsd
+
+
+
+class HarmonicChroma(nn.Module):
+
+    def __init__(self, 
+        hop_length: int, window_length: int = 4096,
+        n_chroma: int = 48, sample_rate: int = 44100,
+        top_n: int = 0
+    ):
+        super().__init__()
+        from torchaudio.prototype.transforms import ChromaScale
+        self.hop_length = hop_length
+        self.window_length = window_length
+        self.n_chroma = n_chroma
+        self.sample_rate = sample_rate
+        self.top_n = top_n
+
+        # HUGO: this representation, as is,
+        # encodes timbre information in the chroma
+        # which is not what we want!!!
+        # would a median filter help perhaps? 
+        self.chroma = ChromaScale(
+            sample_rate=self.sample_rate,
+            n_freqs=self.window_length // 2 + 1,
+            n_chroma=self.n_chroma,
+            octwidth=5.0,
+        )
+
+    @property
+    def dim(self):
+        return self.n_chroma
+
+    def extract(self, sig: Signal) -> Tensor:
+        from vampnet.dsp.hpss import hpss
+        self.chroma.to(sig.wav.device)
+
+        # spectrogram 
+        spec = sn.stft(sig, 
+            window_length=self.window_length, 
+            hop_length=self.hop_length
+        )
+        # magnitude
+        spec = torch.abs(spec)
+
+        # hpss
+        spec = hpss(spec, kernel_size=51, hard=True)[0]
+
+        # chroma
+        chroma = self.chroma(spec)
+
+        # get the rms of this spec
+        rms_d = sn.rms_from_spec(
+            spec, window_length=self.window_length
+        )
+
+        # convert the rms to db
+        rms_d = 10 * torch.log10(rms_d + 1e-7)
+
+        # make a mask based on the rms < -40
+        mask = torch.where(rms_d < -40, torch.zeros_like(rms_d), torch.ones_like(rms_d))
+
+        # remove anything below 80 (where the fuck did I get this number from?)
+        chroma = torch.where(chroma < 100, torch.zeros_like(chroma), chroma)
+
+        # Get top 2 values and indices along the -2 dimension
+        if self.top_n:
+            _, topk_indices = torch.topk(chroma, self.top_n, dim=-2)
+
+            # Create a mask for the top 2 values
+            topk_mask = torch.zeros_like(chroma).scatter_(-2, topk_indices, 1.0)
+
+            # Retain only the top 2 values
+            chroma = chroma * topk_mask
+
+        # apply the mask
+        chroma = chroma * mask.unsqueeze(-2)
+
+        # Apply softmax along dim=-2
+        if self.top_n > 0:
+            chroma = torch.nn.functional.softmax(chroma, dim=-2)
+
+            # mask out any timesteps whose chroma have all equal values (all 0s before softmax)
+            # TODO: i did this with chatgpt, there's gott a be a better way
+            chroma_mean = chroma.mean(dim=-2, keepdim=True)
+            chroma_diff = torch.abs(chroma - chroma_mean)
+            equal_mask = torch.all(chroma_diff < 1e-6, dim=-2, keepdim=True)
+            
+            # Set chroma values to zero for timesteps with all equal values
+            chroma = torch.where(equal_mask, torch.zeros_like(chroma), chroma)
+
+
+        return chroma[:, 0, :, :-1] # mono only :(  FIX ME!
+
+
+# TODO: try harmonic mel? 
+
+CONTROLLERS = {
+    "rms": RMS, 
+    "rmsq128": partial(RMS, n_quantize=128),
+    "rmsq16": partial(RMS, n_quantize=16),
+    "rms-median": partial(RMS, median_filter_size=5),
+    "rmsq16-median": partial(RMS, n_quantize=16, median_filter_size=3),
+    "hchroma": HarmonicChroma,
+    "hchroma-12c-top2": partial(HarmonicChroma, n_chroma=12,  top_n=2), # TODO: refactor me. If this works, this should just be named hchroma. 
+    "hchroma-36c-top3": partial(HarmonicChroma, n_chroma=36,  top_n=3) # TODO: refactor me. If this works, this should just be named hchroma.
+}
+ 
+class Sketch2SoundController(nn.Module):
+
+    def __init__(
+        self,
+        ctrl_keys: list[str], 
+        hop_length: str, 
+        sample_rate: int,
+    ):
+        super().__init__()
+
+        assert all([k in CONTROLLERS for k in ctrl_keys]), f"got an unsupported control key in {ctrl_keys}!\n  supported: {CONTROLLERS.keys()}"
+
+        self.hop_length = hop_length
+        self.ctrl_keys = ctrl_keys
+        self.sample_rate = sample_rate
+
+        self.controllers = {
+            k: CONTROLLERS[k](hop_length=hop_length, sample_rate=sample_rate)
+            for k in self.ctrl_keys
+        }
+
+    @property
+    def ctrl_dims(self, ) -> dict[str, int]:
+        return {
+            k: controller.dim for k, controller in self.controllers.items()
+        }
+
+    def extract(self, sig: Signal) -> dict[str, Tensor]:
+        ctrls = {
+            k: controller.extract(sig) for k, controller in self.controllers.items()
+        }
+        return ctrls
+
+    def random_mask(self, ctrls: dict[str, Tensor], r: float):
+        masks = {}
+        for k, ctrl in ctrls.items():
+            masks[k] = 1-random_along_time(ctrl, r)
+        return masks
+
+    def empty_mask(self, ctrls: dict[str, Tensor]):
+        first_key = next(iter(ctrls))
+        mask = torch.zeros_like(ctrls[first_key])
+        return {k: mask for k in ctrls}
+        
+
+def test_controller():
+    controller = Sketch2SoundController(
+        ctrl_keys=["rms-median", "rms", "rmsq128"], 
+        hop_length=512, 
+        sample_rate=44100
+    )
+    controller.train()
+    # sig = sn.read_from_file("assets/example.wav")
+    # sig = sn.read_from_file("/Users/hugo/Downloads/DCS_SE_FullChoir_ScaleUpDown06_A2_DYN.wav")
+    # sig = sn.excerpt('/Users/hugo/Downloads/(guitarra - hugo mix) bubararu - tambor negro.wav', offset=0, duration=10)
+    sig = sn.read_from_file("assets/voice-prompt.wav")
+    ctrls = controller.extract(sig)
+    print(f"given sig of shape {sig.wav.shape}, extracted controls: {ctrls}")
+
+    # print the whole thing
+    # torch.set_printoptions(profile="full")
+    # print(ctrls["hchroma"][0][0][:, 200:210])
+
+    # imshow the chroma
+    import matplotlib.pyplot as plt
+
+    # Define relative heights for the subplots
+    fig, (ax1, ax2, ax3, ax4) = plt.subplots(
+        4, 1, 
+        sharex=True, 
+    )
+
+    # Display the spectrogram on the top
+    ax1.imshow(sn.stft(sig, hop_length=512, window_length=2048).abs()[0][0].cpu().log().numpy(), aspect='auto', origin='lower')
+    # display rms on the bottom
+    ax2.plot(ctrls["rms-median"][0][0])
+    ax3.plot(ctrls["rms"][0][0])
+    ax4.plot(ctrls["rmsq128"][0][0])
+
+    plt.tight_layout()  # Ensure proper spacing
+    plt.savefig("img.png")
+
+
+if __name__ == "__main__":
+    test_controller()

vampnet/interface.py

@@ -59,7 +59,7 @@ class Interface(torch.nn.Module):
         coarse2fine_ckpt: str = None,
         coarse2fine_lora_ckpt: str = None,
         codec_ckpt: str = None,
-        wavebeat_ckpt: str = None,
+        wavebeat_ckpt: str = "./models/vampnet/wavebeat.pth",
         device: str = "cpu",
         coarse_chunk_size_s: int =  10, 
         coarse2fine_chunk_size_s: int =  3,
@@ -96,7 +96,7 @@ class Interface(torch.nn.Module):
 
         if wavebeat_ckpt is not None:
             logging.debug(f"loading wavebeat from {wavebeat_ckpt}")
-            self.beat_tracker = WaveBeat(wavebeat_ckpt)
+            self.beat_tracker = WaveBeat(wavebeat_ckpt, device=device)
             self.beat_tracker.model.to(device)
         else:
             self.beat_tracker = None
@@ -254,6 +254,7 @@ class Interface(torch.nn.Module):
         """
         assert self.beat_tracker is not None, "No beat tracker loaded"
 
+
         # get the beat times
         beats, downbeats = self.beat_tracker.extract_beats(signal)
 
@@ -516,12 +517,19 @@ class Interface(torch.nn.Module):
         # the forward pass
         logging.debug(z.shape)
         logging.debug("coarse!")
-        zv, mask_z = self.coarse_vamp(
-            z,
-            mask=mask,
-            return_mask=True, 
-            **kwargs
-        )
+        zv = z
+        for i in range(feedback_steps):
+            zv, mask_z = self.coarse_vamp(
+                zv, 
+                mask=mask,
+                return_mask=True, 
+                **kwargs)
+            # roll the mask around a random amount
+            mask_z = mask_z.roll(
+                shifts=(i + 1) % feedback_steps, 
+                dims=-1
+            )
+
 
         # add the top codebooks back in
         if zv.shape[1] < z.shape[1]:

vampnet/mask.py

@@ -163,14 +163,18 @@ def dropout(
     mask: torch.Tensor,
     p: float,
 ):
-    assert 0 <= p <= 1, "p must be between 0 and 1"
-    assert mask.max() <= 1, "mask must be binary"
-    assert mask.min() >= 0, "mask must be binary"
-    mask = (~mask.bool()).float()
-    mask = torch.bernoulli(mask * (1 - p))
-    mask = ~mask.round().bool()
+    # instead of the above, mask along the last dimensions
+    tsteps = mask.shape[-1]
+    tsteps_to_drop = int(tsteps * p)
+    tsteps_to_keep = tsteps - tsteps_to_drop
+    idxs_to_drop = torch.randint(0, tsteps, (tsteps_to_drop,))
+    mask = mask.clone()
+    mask[:, :, idxs_to_drop] = 1
     return mask.long()
 
+
+
+
 def mask_or(
     mask1: torch.Tensor, 
     mask2: torch.Tensor

vampnet/modules/transformer.py

@@ -6,6 +6,7 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch import Tensor
 from einops import rearrange
 import loralib as lora
 import audiotools as at
@@ -405,7 +406,7 @@ class TransformerStack(nn.Module):
         )
 
         # Perform last normalization
-        self.norm = RMSNorm(d_model) if last_layer else None
+        self.norm = RMSNorm(d_model) if last_layer else None   
 
     def subsequent_mask(self, size):
         return torch.ones(1, size, size).tril().bool()
@@ -461,6 +462,75 @@ class TransformerStack(nn.Module):
         else:
             return out
 
+class CFGDropout(nn.Module):
+    
+    def __init__(self, p: float = 0.2):
+        super().__init__()
+        self.p = p
+
+    def forward(self, x: Tensor):
+        # dropout along the batch dim
+        if self.training:
+            mask = torch.rand(x.shape[0], 1, 1, device=x.device) > self.p
+        else:
+            mask = torch.ones(x.shape[0], 1, 1, device=x.device)
+        return x * mask
+
+class ControlEncoder(nn.Module):
+
+    def __init__(self, 
+        ctrl_dims: dict[str, int], 
+        embedding_dim: int,
+        cfg_dropout_prob: float
+    ):
+        super().__init__()
+        self.ctrl_encoders = nn.ModuleDict({
+            key: nn.Linear(dim, embedding_dim)
+                for key, dim in ctrl_dims.items()
+        })
+
+        self.cfg_dropout = CFGDropout(p=cfg_dropout_prob)
+        self.all_dropout = CFGDropout(p=cfg_dropout_prob / 2)
+
+    def forward(self, 
+        embedding: Tensor, # embedding to which we will add ctrls
+        ctrls: dict[str, Tensor], 
+        ctrl_masks: dict[str, Tensor]
+    ):
+        # INPUT: ctrl tensor should be shape (b d n)
+
+        # assert that we got all the right ctrls and ctrl_masks according to the encoders that we have
+        assert list(sorted(ctrls.keys())) == list(sorted(self.ctrl_encoders.keys())), "ctrls and ctrl_encoders keys do not match"
+        assert list(sorted(ctrl_masks.keys())) == list(sorted(self.ctrl_encoders.keys())), "ctrl_masks and ctrl_encoders keys do not match"
+
+        out_emb = torch.zeros_like(embedding)
+        for ck in ctrls:
+            ctrld = ctrls[ck]
+            ctrlmask = ctrl_masks[ck]
+
+            assert ctrld.shape[-1] == embedding.shape[-1], "ctrls should match x along time dimension"
+            assert ctrlmask.ndim == 2, "ctrlmask should be 2d"
+            assert ctrlmask.shape[-1] == ctrld.shape[-1], "ctrlmask should match ctrld along time dimension"
+
+            # project ctrl with encoder
+            ctrld = rearrange(ctrld, "b d n -> b n d")
+            ctrl_emb = self.ctrl_encoders[ck](ctrld)
+            ctrld = rearrange(ctrld, "b n d -> b d n")
+            ctrl_emb = rearrange(ctrl_emb, "b n d -> b d n")
+
+            # apply ctrl mask
+            ctrl_emb = ctrl_emb * ctrlmask[:, None, :]
+
+            # apply cfg dropout
+            ctrl_emb = self.cfg_dropout(ctrl_emb)
+
+            # add to the out_emb
+            out_emb = out_emb + ctrl_emb
+
+        # randomly dropout all ctrls
+        out_emb = self.all_dropout(out_emb)
+
+        return out_emb
 
 class VampNet(at.ml.BaseModel):
     def __init__(
@@ -475,7 +545,10 @@ class VampNet(at.ml.BaseModel):
         vocab_size: int = 1024,
         flash_attn: bool = True,
         noise_mode: str = "mask",
-        dropout: float = 0.1
+        dropout: float = 0.1,
+        ctrl_dims: Optional[dict[str, int]] = None, 
+        cfg_dropout_prob: float = 0.2, 
+        cond_dim: int = 0,
     ):
         super().__init__()
         assert r_cond_dim == 0, f"r_cond_dim must be 0 (not supported), but got {r_cond_dim}"
@@ -489,6 +562,11 @@ class VampNet(at.ml.BaseModel):
         self.latent_dim = latent_dim
         self.flash_attn = flash_attn
         self.noise_mode = noise_mode
+        self.cond_dim = cond_dim
+        self.r_cond_dim = r_cond_dim
+        self.dropout = dropout
+        self.cfg_dropout_prob = cfg_dropout_prob
+        # self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
         assert self.noise_mode == "mask", "deprecated"
 
@@ -525,10 +603,25 @@ class VampNet(at.ml.BaseModel):
             ),
         )
 
-    def forward(self, x, return_activations: bool = False):
+        if self.cond_dim > 0:
+            self.cfg_dropout = CFGDropout(p=cfg_dropout_prob)
+
+        self.ctrl_dims = ctrl_dims
+        if self.ctrl_dims is not None:
+            self.ctrl_encoder = ControlEncoder(
+                ctrl_dims, 
+                embedding_dim=embedding_dim, 
+                cfg_dropout_prob=cfg_dropout_prob
+            )
+
+    def forward(self, x, ctrls=None, ctrl_masks=None, return_activations: bool = False):
         x = self.embedding(x)
         x_mask = torch.ones_like(x, dtype=torch.bool)[:, :1, :].squeeze(1)
 
+        if self.ctrl_dims is not None:
+            # apply controls
+            x = x + self.ctrl_encoder(x, ctrls, ctrl_masks)
+            
         x = rearrange(x, "b d n -> b n d")
         out = self.transformer(x=x, x_mask=x_mask, return_activations=return_activations)
         if return_activations:
@@ -600,6 +693,8 @@ class VampNet(at.ml.BaseModel):
         temperature: float = 1.0,
         mask: Optional[torch.Tensor] = None,
         mask_temperature: float = 10.5,
+        ctrls:dict = None,
+        ctrl_masks:dict = None,
         typical_filtering=True,
         typical_mass=0.15,
         typical_min_tokens=64,
@@ -609,7 +704,9 @@ class VampNet(at.ml.BaseModel):
         return_signal=True,
         debug=False,
         causal_weight: float = 0.0,
+        cfg_scale: float = 3.0,
         cfg_guidance: float = None,
+        cond = None # unused
     ):
         if seed is not None:
             at.util.seed(seed)
@@ -622,6 +719,22 @@ class VampNet(at.ml.BaseModel):
         z = start_tokens
         nb = z.shape[0]
 
+        use_cfg = ctrls is not None
+        tocfg = lambda x: x.repeat(2, 1, 1) if use_cfg else x
+        tocfgblank = lambda x: torch.cat([x, torch.zeros_like(x)], dim=0) if use_cfg else x
+        def fromcfg(x):
+            if use_cfg:
+                xcond, xuncond = x.chunk(2)
+                return xuncond + cfg_scale * (xcond - xuncond)
+            return x
+
+        z = tocfg(z)
+        if ctrls is not None:
+            ctrls = {k: tocfg(v) for k, v in ctrls.items()}
+            ctrl_masks = {k: tocfgblank(v) for k, v in ctrl_masks.items()}
+        if cond is not None:
+            cond = tocfg(cond)
+
         if z is None:
             z = torch.full((1, self.n_codebooks, time_steps), self.mask_token).to(
                 self.device
@@ -727,6 +840,7 @@ class VampNet(at.ml.BaseModel):
             # infer from latents
             # NOTE: this collapses the codebook dimension into the sequence dimension
             logits = self.forward(latents) # b, prob, seq
+            logits = fromcfg(logits)
 
             if cfg_guidance is not None:
                 logits_cond, logits_uncond = logits[:nb], logits[nb:]
@@ -774,9 +888,6 @@ class VampNet(at.ml.BaseModel):
                 plt.imshow(_mask[0].cpu().numpy())
                 plt.savefig(f"{STEP_FOLDER}/mask.png")
 
-
-
-            
             # update the mask, remove conditioning codebooks from the mask
             # add z back into sampled z where the mask was false
             sampled_z = torch.where(