from pathlib import Path from typing import Tuple import yaml import tempfile import uuid from dataclasses import dataclass, asdict import numpy as np import audiotools as at import argbind import gradio as gr from vampnet.interface import Interface from vampnet import mask as pmask Interface = argbind.bind(Interface) # AudioLoader = argbind.bind(at.data.datasets.AudioLoader) conf = argbind.parse_args() with argbind.scope(conf): interface = Interface() # loader = AudioLoader() print(f"interface device is {interface.device}") # dataset = at.data.datasets.AudioDataset( # loader, # sample_rate=interface.codec.sample_rate, # duration=interface.coarse.chunk_size_s, # n_examples=5000, # without_replacement=True, # ) OUT_DIR = Path("gradio-outputs") OUT_DIR.mkdir(exist_ok=True, parents=True) def load_audio(file): print(file) filepath = file.name sig = at.AudioSignal.salient_excerpt( filepath, duration=interface.coarse.chunk_size_s ) sig = interface.preprocess(sig) out_dir = OUT_DIR / "tmp" / str(uuid.uuid4()) out_dir.mkdir(parents=True, exist_ok=True) sig.write(out_dir / "input.wav") return sig.path_to_file def load_example_audio(): return "./assets/example.wav" def _vamp(data, return_mask=False): out_dir = OUT_DIR / str(uuid.uuid4()) out_dir.mkdir() sig = at.AudioSignal(data[input_audio]) z = interface.encode(sig) ncc = data[n_conditioning_codebooks] # build the mask mask = pmask.linear_random(z, data[rand_mask_intensity]) mask = pmask.mask_and( mask, pmask.inpaint( z, interface.s2t(data[prefix_s]), interface.s2t(data[suffix_s]) ) ) mask = pmask.mask_and( mask, pmask.periodic_mask( z, data[periodic_p], data[periodic_w], random_roll=True ) ) if data[onset_mask_width] > 0: mask = pmask.mask_or( mask, pmask.onset_mask(sig, z, interface, width=data[onset_mask_width]) ) if data[beat_mask_width] > 0: beat_mask = interface.make_beat_mask( sig, after_beat_s=(data[beat_mask_width]/1000), mask_upbeats=not data[beat_mask_downbeats], ) mask = pmask.mask_and(mask, beat_mask) # these should be the last two mask ops mask = pmask.dropout(mask, data[dropout]) mask = pmask.codebook_unmask(mask, ncc) print(f"created mask with: linear random {data[rand_mask_intensity]}, inpaint {data[prefix_s]}:{data[suffix_s]}, periodic {data[periodic_p]}:{data[periodic_w]}, dropout {data[dropout]}, codebook unmask {ncc}, onset mask {data[onset_mask_width]}, num steps {data[num_steps]}, init temp {data[temp]}, use coarse2fine {data[use_coarse2fine]}") # save the mask as a txt file np.savetxt(out_dir / "mask.txt", mask[:,0,:].long().cpu().numpy()) zv, mask_z = interface.coarse_vamp( z, mask=mask, sampling_steps=data[num_steps], temperature=data[temp]*10, return_mask=True, typical_filtering=data[typical_filtering], typical_mass=data[typical_mass], typical_min_tokens=data[typical_min_tokens], gen_fn=interface.coarse.generate, ) if use_coarse2fine: zv = interface.coarse_to_fine(zv, temperature=data[temp]) sig = interface.to_signal(zv).cpu() print("done") sig.write(out_dir / "output.wav") if return_mask: mask = interface.to_signal(mask_z).cpu() mask.write(out_dir / "mask.wav") return sig.path_to_file, mask.path_to_file else: return sig.path_to_file def vamp(data): return _vamp(data, return_mask=True) def api_vamp(data): return _vamp(data, return_mask=False) def save_vamp(data): out_dir = OUT_DIR / "saved" / str(uuid.uuid4()) out_dir.mkdir(parents=True, exist_ok=True) sig_in = at.AudioSignal(data[input_audio]) sig_out = at.AudioSignal(data[output_audio]) sig_in.write(out_dir / "input.wav") sig_out.write(out_dir / "output.wav") _data = { "temp": data[temp], "prefix_s": data[prefix_s], "suffix_s": data[suffix_s], "rand_mask_intensity": data[rand_mask_intensity], "num_steps": data[num_steps], "notes": data[notes_text], "periodic_period": data[periodic_p], "periodic_width": data[periodic_w], "n_conditioning_codebooks": data[n_conditioning_codebooks], "use_coarse2fine": data[use_coarse2fine], "stretch_factor": data[stretch_factor], } # save with yaml with open(out_dir / "data.yaml", "w") as f: yaml.dump(_data, f) import zipfile zip_path = out_dir.with_suffix(".zip") with zipfile.ZipFile(zip_path, "w") as zf: for file in out_dir.iterdir(): zf.write(file, file.name) return f"saved! your save code is {out_dir.stem}", zip_path with gr.Blocks() as demo: 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)", file_types=["audio"] ) load_example_audio_button = gr.Button("or load example audio") input_audio = gr.Audio( label="input audio", interactive=False, 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( fn=load_example_audio, inputs=[], outputs=[ input_audio] ) manual_audio_upload.change( fn=load_audio, inputs=[manual_audio_upload], outputs=[ input_audio] ) # mask settings with gr.Column(): periodic_p = gr.Slider( label="periodic prompt (0.0 means no prompt, 2 - lots of hints, 8 - a couple of hints, 16 - occasional hint, 32 - very occasional hint, etc)", minimum=0, maximum=128, step=1, value=3, ) onset_mask_width = gr.Slider( label="onset mask width (multiplies with the periodic mask, 1 step ~= 10milliseconds) ", minimum=0, maximum=20, step=1, value=5, ) beat_mask_width = gr.Slider( label="beat mask width (in milliseconds)", minimum=0, maximum=200, value=0, ) beat_mask_downbeats = gr.Checkbox( label="beat mask downbeats only?", value=False ) with gr.Accordion("extras ", open=False): 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=20, step=1, value=1, ) n_conditioning_codebooks = gr.Number( label="number of conditioning codebooks. probably 0", value=0, precision=0, ) stretch_factor = gr.Slider( label="time stretch factor", minimum=0, maximum=64, step=1, value=1, ) with gr.Accordion("prefix/suffix hints", open=False): 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 ) temp = gr.Slider( label="temperature", minimum=0.0, maximum=3.0, value=0.8 ) with gr.Accordion("sampling settings", open=False): typical_filtering = gr.Checkbox( label="typical filtering ", value=False ) 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=64 ) use_coarse2fine = gr.Checkbox( label="use coarse2fine", value=True ) num_steps = gr.Slider( label="number of steps (should normally be between 12 and 36)", minimum=1, maximum=128, step=1, value=36 ) dropout = gr.Slider( label="mask dropout", minimum=0.0, maximum=1.0, step=0.01, value=0.0 ) # mask settings with gr.Column(): vamp_button = gr.Button("generate (vamp)!!!") output_audio = gr.Audio( label="output audio", interactive=False, type="filepath" ) notes_text = gr.Textbox( label="type any notes about the generated audio here", value="", interactive=True ) save_button = gr.Button("save vamp") download_file = gr.File( label="vamp to download will appear here", interactive=False ) use_as_input_button = gr.Button("use output as input") thank_you = gr.Markdown("") _inputs = { input_audio, num_steps, temp, prefix_s, suffix_s, rand_mask_intensity, periodic_p, periodic_w, n_conditioning_codebooks, dropout, use_coarse2fine, stretch_factor, onset_mask_width, typical_filtering, typical_mass, typical_min_tokens, beat_mask_width, beat_mask_downbeats } # connect widgets vamp_button.click( fn=vamp, inputs=_inputs, outputs=[output_audio, audio_mask], ) api_vamp_button = gr.Button("api vamp", visible=False) api_vamp_button.click( fn=api_vamp, inputs=_inputs, outputs=[output_audio], api_name="vamp" ) use_as_input_button.click( fn=lambda x: x, inputs=[output_audio], outputs=[input_audio] ) save_button.click( fn=save_vamp, inputs=_inputs | {notes_text, output_audio}, outputs=[thank_you, download_file] ) demo.launch(share=True, enable_queue=False, debug=True)