app

app.py

@@ -1,31 +1,20 @@
-# huggingface space exclusive
-import os
-
-# print("installing pyharp")
-# os.system('pip install "pyharp@git+https://github.com/audacitorch/pyharp.git"')
-# print("installing madmom")
-# os.system('pip install cython')
-# os.system('pip install madmom')
-
 from pathlib import Path
-from typing import Tuple
 import yaml
-import tempfile
 import uuid
-import shutil
-from dataclasses import dataclass, asdict
 
 import numpy as np
 import audiotools as at
 import argbind
+import shutil
 import torch
+from datetime import datetime
 
 import gradio as gr
-from vampnet.interface import Interface
+from vampnet.interface import Interface, signal_concat
 from vampnet import mask as pmask
 
-from pyharp import ModelCard, build_endpoint
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 
 interface = Interface(
     device=device,
@@ -46,6 +35,16 @@ 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
 
     
@@ -53,15 +52,15 @@ for conf_file in generated_confs.glob("*/interface.yml"):
 OUT_DIR = Path("gradio-outputs")
 OUT_DIR.mkdir(exist_ok=True, parents=True)
 
-
+MAX_DURATION_S = 60
 def load_audio(file):
     print(file)
     filepath = file.name
     sig = at.AudioSignal.salient_excerpt(
-        filepath, 
-        duration=interface.coarse.chunk_size_s
+        filepath, duration=MAX_DURATION_S
     )
-    sig = interface.preprocess(sig)
+    # sig = interface.preprocess(sig)
+    sig = at.AudioSignal(filepath)
 
     out_dir = OUT_DIR / "tmp" / str(uuid.uuid4())
     out_dir.mkdir(parents=True, exist_ok=True)
@@ -72,92 +71,234 @@ def load_audio(file):
 def load_example_audio():
     return "./assets/example.wav"
 
+from torch_pitch_shift import pitch_shift, get_fast_shifts
+def shift_pitch(signal, interval: int):
+    signal.samples = pitch_shift(
+        signal.samples, 
+        shift=interval, 
+        sample_rate=signal.sample_rate
+    )
+    return signal
+
+def _vamp(seed, input_audio, model_choice, pitch_shift_amt, periodic_p, p2, n_mask_codebooks, n_mask_codebooks_2, rand_mask_intensity, prefix_s, suffix_s, periodic_w, onset_mask_width, dropout, masktemp, sampletemp, typical_filtering, typical_mass, typical_min_tokens, top_p, sample_cutoff, win_dur, num_feedback_steps, stretch_factor, api=False):
+    _seed = seed if seed > 0 else None
+    if _seed is None:
+        _seed = int(torch.randint(0, 2**32, (1,)).item())
+    at.util.seed(_seed)
+
+    datentime = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+    out_dir = OUT_DIR / f"{Path(input_audio).stem}-{datentime}-seed-{_seed}-model-{model_choice}"
+    out_dir.mkdir(parents=True)
+    sig = at.AudioSignal(input_audio)
+    sig.write(out_dir / "input.wav")
 
-def _vamp(sig, data):
-
-    sig = interface.preprocess(sig)
+    # 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"],
+    )
 
     loudness = sig.loudness()
     print(f"input loudness is {loudness}")
 
-    z = interface.encode(sig)
-
-    # build the mask
-    mask = pmask.full_mask(z)
-    mask = pmask.mask_and(
-        mask, pmask.periodic_mask(
-            z,
-            data[periodic_p],
-            random_roll=True
-        )
+    if pitch_shift_amt != 0:
+        sig = shift_pitch(sig, pitch_shift_amt)
+
+    _p2 = periodic_p if p2 == 0 else p2
+    _n_codebooks_2 = n_mask_codebooks if n_mask_codebooks_2 == 0 else n_mask_codebooks_2
+
+    build_mask_kwargs = dict(
+        rand_mask_intensity=rand_mask_intensity,
+        prefix_s=prefix_s,
+        suffix_s=suffix_s,
+        periodic_prompt=int(periodic_p),
+        periodic_prompt2=int(_p2),
+        periodic_prompt_width=periodic_w,
+        onset_mask_width=onset_mask_width,
+        _dropout=dropout,
+        upper_codebook_mask=int(n_mask_codebooks), 
+        upper_codebook_mask_2=int(_n_codebooks_2),
     )
 
-    # these should be the last two mask ops
-    mask = pmask.codebook_mask(mask, int(data[n_mask_codebooks]))
-
-    print(f"sampletemp {data[sampletemp]}")
-    print(f"num_steps {data[num_steps]}")
-    print(f"periodic_p {data[periodic_p]}")
-
-    print(f"processing coarse...")
-    zv, mask_z = interface.coarse_vamp(
-        z, 
-        mask=mask,
-        sampling_steps=data[num_steps],
-        mask_temperature=1.5*10,
-        sampling_temperature=data[sampletemp],
-        return_mask=True, 
-        top_p=0.85,
-        gen_fn=interface.coarse.generate,
-        sample_cutoff=1.0,
+    vamp_kwargs = dict(
+        mask_temperature=masktemp*10,
+        sampling_temperature=sampletemp,
+        typical_filtering=typical_filtering, 
+        typical_mass=typical_mass, 
+        typical_min_tokens=typical_min_tokens, 
+        top_p=top_p if top_p > 0 else None,
+        seed=_seed,
+        sample_cutoff=sample_cutoff,
     )
 
-    print(f"processing coarse to fine...")
-    zv = interface.coarse_to_fine(
-        zv, 
-        mask_temperature=1.5*10, 
-        sampling_temperature=data[sampletemp],
-        mask=mask,
-        sampling_steps=data[num_steps] // 2,
-        sample_cutoff=1.0, 
+    # save the mask as a txt file
+    interface.set_chunk_size(win_dur)
+    sig, mask, codes = interface.ez_vamp(
+        sig, 
+        batch_size=4 if not api else 1,
+        feedback_steps=num_feedback_steps,
+        time_stretch_factor=stretch_factor,
+        build_mask_kwargs=build_mask_kwargs,
+        vamp_kwargs=vamp_kwargs,
+        return_mask=True,
     )
 
-    sig = interface.to_signal(zv).cpu()
-    print("done")
-    return sig
-
-
-def process_fn(data):
-    # remove any old files in the output directory (from previous runs)
-    shutil.rmtree(OUT_DIR)
-    OUT_DIR.mkdir()
-
-    out_dir = OUT_DIR / str(uuid.uuid4())
-    out_dir.mkdir()
-    sig = at.AudioSignal(data[input_audio])
-
-    for _pass in range(data[num_passes]):
-        sig = _vamp(sig, data)
+    if api:
+        sig.write(out_dir / "out.wav")
+
+        return sig.path_to_file
+
+    if not api: 
+        # write codes to numpy file
+        np.save(out_dir / "codes.npy", codes.cpu().numpy())
+        metadata = {}
+        metadata["seed"] = _seed
+        metadata["model_choice"] = model_choice
+        metadata["mask_kwargs"] = build_mask_kwargs
+        metadata["vamp_kwargs"] = vamp_kwargs
+        metadata["loudness"] = loudness
+        # save the metadata
+        with open(out_dir / "metadata.yml", "w") as f:
+            yaml.dump(metadata, f)
+
+        sig0 = sig[0].write(out_dir / "out1.wav")
+        sig1 = sig[1].write(out_dir / "out2.wav")
+        sig2 = sig[2].write(out_dir / "out3.wav")
+        sig3 = sig[3].write(out_dir / "out4.wav")
+
+        # write the mask to txt
+        with open(out_dir / "mask.txt", "w") as f:
+            m = mask[0].cpu().numpy()
+            # write to txt, each time step on a new line
+            for i in range(m.shape[-1]):
+                f.write(f"{m[:, i]}\n")
+
+
+        import matplotlib.pyplot as plt
+        plt.clf()
+        interface.visualize_codes(mask)
+        plt.savefig(out_dir / "mask.png")
+        plt.clf()
+        interface.visualize_codes(codes)
+        plt.savefig(out_dir / "codes.png")
+        plt.close()
+        
+        # zip out dir, and return the path to the zip
+        shutil.make_archive(out_dir, 'zip', out_dir)
+
+        # chunk in groups of 1024 timesteps
+        _mask_sigs = []
+        for i in range(0, mask.shape[-1], 1024):
+            _mask_sigs.append(interface.to_signal(mask[:, :, i:i+1024].to(interface.device)).cpu())
+        mask = signal_concat(_mask_sigs)
+        mask.write(out_dir / "mask.wav")
+
+        return (
+            sig0.path_to_file, sig1.path_to_file, 
+            sig2.path_to_file, sig3.path_to_file,
+            mask.path_to_file, str(out_dir.with_suffix(".zip")), out_dir / "mask.png"
+        )
 
+def 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],
+        p2=data[p2],
+        n_mask_codebooks=data[n_mask_codebooks],
+        n_mask_codebooks_2=data[n_mask_codebooks_2],
+        rand_mask_intensity=data[rand_mask_intensity],
+        prefix_s=data[prefix_s],
+        suffix_s=data[suffix_s],
+        periodic_w=data[periodic_w],
+        onset_mask_width=data[onset_mask_width],
+        dropout=data[dropout],
+        masktemp=data[masktemp],
+        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],
+        win_dur=data[win_dur],
+        num_feedback_steps=data[num_feedback_steps],
+        stretch_factor=data[stretch_factor],
+        api=False, 
+    )
 
-    sig.write(out_dir / "output.wav")
+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],
+        p2=data[p2],
+        n_mask_codebooks=data[n_mask_codebooks],
+        n_mask_codebooks_2=data[n_mask_codebooks_2],
+        rand_mask_intensity=data[rand_mask_intensity],
+        prefix_s=data[prefix_s],
+        suffix_s=data[suffix_s],
+        periodic_w=data[periodic_w],
+        onset_mask_width=data[onset_mask_width],
+        dropout=data[dropout],
+        masktemp=data[masktemp],
+        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],
+        win_dur=data[win_dur],
+        num_feedback_steps=data[num_feedback_steps],
+        stretch_factor=data[stretch_factor],
+        api=True, 
+    )
 
-    return sig.path_to_file
 
+def harp_vamp(input_audio,
+            periodic_p, 
+            n_mask_codebooks, 
+            pitch_shift_amt, 
+            win_dur, 
+            num_feedback_steps):
+    return _vamp(
+        seed=0, 
+        input_audio=input_audio,
+        model_choice="default",
+        pitch_shift_amt=pitch_shift_amt,
+        periodic_p=periodic_p,
+        p2=0,
+        n_mask_codebooks=n_mask_codebooks,
+        n_mask_codebooks_2=0,
+        rand_mask_intensity=1.0,
+        prefix_s=0.0,
+        suffix_s=0.0,
+        periodic_w=1,
+        onset_mask_width=0,
+        dropout=0.0,
+        masktemp=1.5,
+        sampletemp=1.0,
+        typical_filtering=True,
+        typical_mass=0.15,
+        typical_min_tokens=64,
+        top_p=0.9,
+        sample_cutoff=1.0,
+        win_dur=win_dur,
+        num_feedback_steps=num_feedback_steps,
+        stretch_factor=1.0,
+        api=True,
+    )
+    
+    
 
 with gr.Blocks() as demo:
-
     with gr.Row():
         with gr.Column():
-            gr.Markdown("# nesquik 🌰🐿️👾 ")
-            gr.Markdown(" the ultimate bitcrusher! will do its best to convert your instrumental music into an 8-bit chiptune.")
-
-    with gr.Row():
-        with gr.Column():
-
-
             manual_audio_upload = gr.File(
-                label=f"upload some audio (will be randomly trimmed to max of {interface.coarse.chunk_size_s:.2f}s)",
+                label=f"upload some audio (will be randomly trimmed to max of 100s)",
                 file_types=["audio"]
             )
             load_example_audio_button = gr.Button("or load example audio")
@@ -168,6 +309,11 @@ with gr.Blocks() as demo:
                 type="filepath",
             )
 
+            audio_mask = gr.Audio(
+                label="audio mask (listen to this to hear the mask hints)",
+                interactive=False, 
+                type="filepath",
+            )
 
             # connect widgets
             load_example_audio_button.click(
@@ -182,85 +328,282 @@ with gr.Blocks() as demo:
                 outputs=[ input_audio]
             )
                 
+
+
         # mask settings
         with gr.Column():
-            with gr.Accordion("controls", open=False):
+            with gr.Accordion("manual controls", open=True):
                 periodic_p = gr.Slider(
                     label="periodic prompt",
-                    minimum=1,
-                    maximum=3, 
+                    minimum=0,
+                    maximum=128, 
+                    step=1,
+                    value=3, 
+                )
+                p2 = gr.Slider(
+                    label="periodic prompt 2 (0 - same as p1, 2 - lots of hints, 8 - a couple of hints, 16 - occasional hint, 32 - very occasional hint, etc)",
+                    minimum=0,
+                    maximum=128,
                     step=1,
-                    value=2, 
+                    value=0,
+                )
+
+                onset_mask_width = gr.Slider(
+                    label="onset mask width (multiplies with the periodic mask, 1 step ~= 10milliseconds) ",
+                    minimum=0,
+                    maximum=100,
+                    step=1,
+                    value=0,
                 )
 
                 n_mask_codebooks = gr.Slider(
-                    label="first upper codebook level to mask",
+                    label="compression prompt ",
+                    value=3,
                     minimum=0, 
-                    maximum=9, 
-                    value=2,
+                    maximum=14,
                     step=1,
                 )
+                n_mask_codebooks_2 = gr.Number(
+                    label="compression prompt 2 via linear interpolation (0 == constant)",
+                    value=0,
+                )
 
-                sampletemp = gr.Slider(
-                    label="sample temperature",
-                    minimum=0.8,
-                    maximum=1.5,
-                    value=1.0, 
-                    step=0.001
+            with gr.Accordion("extras ", open=False):
+                pitch_shift_amt = gr.Slider(
+                    label="pitch shift amount (semitones)",
+                    minimum=-12,
+                    maximum=12,
+                    step=1,
+                    value=0,
                 )
-        
-                num_steps = gr.Slider(
-                    label="number of steps (should normally be between 12 and 36)",
+
+                stretch_factor = gr.Slider(
+                    label="time stretch factor",
+                    minimum=0,
+                    maximum=64, 
+                    step=1,
+                    value=1, 
+                )
+
+                rand_mask_intensity = gr.Slider(
+                    label="random mask intensity. (If this is less than 1, scatters prompts throughout the audio, should be between 0.9 and 1.0)",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=1.0
+                )
+
+                periodic_w = gr.Slider(
+                    label="periodic prompt width (steps, 1 step ~= 10milliseconds)",
                     minimum=1,
-                    maximum=36,
-                    step=6,
-                    value=24
+                    maximum=20,
+                    step=1,
+                    value=1,
+                )
+
+            with gr.Accordion("prefix/suffix prompts", open=True):
+                prefix_s = gr.Slider(
+                    label="prefix hint length (seconds)",
+                    minimum=0.0,
+                    maximum=10.0,
+                    value=0.0
+                )
+                suffix_s = gr.Slider(
+                    label="suffix hint length (seconds)",
+                    minimum=0.0,
+                    maximum=10.0,
+                    value=0.0
                 )
 
-                num_passes = gr.Slider(
-                    label="number of passes (more passes = more time, but better results)",
-                    minimum=2,
-                    maximum=6,
+            masktemp = gr.Slider(
+                label="mask temperature",
+                minimum=0.0,
+                maximum=100.0,
+                value=1.5
+            )
+            sampletemp = gr.Slider(
+                label="sample temperature",
+                minimum=0.1,
+                maximum=10.0,
+                value=1.0, 
+                step=0.001
+            )
+        
+
+
+            with gr.Accordion("sampling settings", open=False):
+                top_p = gr.Slider(
+                    label="top p (0.0 = off)",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.9
+                )
+                typical_filtering = gr.Checkbox(
+                    label="typical filtering ",
+                    value=True
+                )
+                typical_mass = gr.Slider( 
+                    label="typical mass (should probably stay between 0.1 and 0.5)",
+                    minimum=0.01,
+                    maximum=0.99,
+                    value=0.15
+                )
+                typical_min_tokens = gr.Slider(
+                    label="typical min tokens (should probably stay between 1 and 256)",
+                    minimum=1,
+                    maximum=256,
                     step=1,
-                    value=4
+                    value=64
                 )
+                sample_cutoff = gr.Slider(
+                    label="sample cutoff",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=1.0, 
+                    step=0.01
+                )
+
+            dropout = gr.Slider(
+                label="mask dropout",
+                minimum=0.0,
+                maximum=1.0,
+                step=0.01,
+                value=0.0
+            )
 
 
-            vamp_button = gr.Button("nes, quick!!!!!")
-            output_audio = gr.Audio(
-                label="output audio",
+            seed = gr.Number(
+                label="seed (0 for random)",
+                value=0,
+                precision=0,
+            )
+
+
+
+        # mask settings
+        with gr.Column():
+
+            model_choice = gr.Dropdown(
+                label="model choice", 
+                choices=list(MODEL_CHOICES.keys()),
+                value="default", 
+                visible=True
+            )
+
+            num_feedback_steps = gr.Slider(
+                label="number of feedback steps (each one takes a while)",
+                minimum=1,
+                maximum=16,
+                step=1,
+                value=1
+            )
+
+            win_dur= gr.Slider(
+                label="window duration (seconds)",
+                minimum=2,
+                maximum=10,
+                value=6)
+            
+
+            vamp_button = gr.Button("generate (vamp)!!!")
+            maskimg = gr.Image(
+                label="mask image",
                 interactive=False,
                 type="filepath"
             )
+            out1 = gr.Audio(
+                label="output audio 1",
+                interactive=False,
+                type="filepath"
+            )
+            out2 = gr.Audio(
+                label="output audio 2",
+                interactive=False,
+                type="filepath"
+            )
+            out3 = gr.Audio(
+                label="output audio 3",
+                interactive=False,
+                type="filepath"
+            )
+            out4 = gr.Audio(
+                label="output audio 4",
+                interactive=False,
+                type="filepath"
+            )
+            
+            thank_you = gr.Markdown("")
+
+            # download all the outputs
+            download = gr.File(type="file", label="download outputs")
+
 
     _inputs = {
             input_audio, 
-            num_steps,
+            masktemp,
             sampletemp,
-            periodic_p,
+            top_p,
+            prefix_s, suffix_s, 
+            rand_mask_intensity, 
+            periodic_p, periodic_w,
+            dropout,
+            stretch_factor, 
+            onset_mask_width, 
+            typical_filtering,
+            typical_mass,
+            typical_min_tokens,
+            seed, 
+            model_choice,
             n_mask_codebooks,
-            num_passes
+            pitch_shift_amt, 
+            sample_cutoff, 
+            num_feedback_steps, 
+            p2, 
+            n_mask_codebooks_2, 
+            win_dur
         }
   
     # connect widgets
     vamp_button.click(
-        fn=process_fn,
+        fn=vamp,
         inputs=_inputs,
-        outputs=[output_audio], 
+        outputs=[out1, out2, out3, out4, audio_mask, download, maskimg], 
+    )
+
+    api_vamp_button = gr.Button("api vamp", visible=False)
+    api_vamp_button.click(
+        fn=api_vamp,
+        inputs=_inputs, 
+        outputs=[out1], 
+        api_name="vamp"
     )
 
+    from pyharp import ModelCard, build_endpoint
+
+    model_card = ModelCard(
+        name="salad bowl", 
+        description="sounds", 
+        author="hugo flores garcía", 
+        tags=["generative","sound"], 
+    )
 
     build_endpoint(
-        inputs=list(_inputs),  
-        output=output_audio, 
-        process_fn=process_fn,
-        card=ModelCard(
-            name="nesquik 🌰🐿️👾", 
-            description="turn your music into NES music!! quick!! NOTE: vampnet's has a maximum context length of 10 seconds. Please split all audio clips into 10 second chunks, or processing will result in an error. ", 
-            author="Hugo Flores García", 
-            tags=["music", "generative"]
-        ), 
-        visible=False
+        inputs=[
+            input_audio,
+            periodic_p, 
+            n_mask_codebooks, 
+            pitch_shift_amt, 
+            win_dur,
+            num_feedback_steps
+        ], 
+        output=out1,
+        process_fn=harp_vamp,
+        card=model_card
     )
 
-demo.queue().launch()
+
+try:
+    demo.queue()
+    demo.launch(share=True)
+except KeyboardInterrupt:
+    shutil.rmtree("gradio-outputs", ignore_errors=True)
+    raise