from pathlib import Path import yaml import uuid 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, signal_concat from vampnet import mask as pmask device = "cuda" if torch.cuda.is_available() else "cpu" interface = Interface( device=device, coarse_ckpt="models/vampnet/coarse.pth", coarse2fine_ckpt="models/vampnet/c2f.pth", codec_ckpt="models/vampnet/codec.pth", ) # populate the model choices with any interface.yml files in the generated confs MODEL_CHOICES = { "default": { "Interface.coarse_ckpt": str(interface.coarse_path), "Interface.coarse2fine_ckpt": str(interface.c2f_path), "Interface.codec_ckpt": str(interface.codec_path), } } 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 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=MAX_DURATION_S ) # sig = interface.preprocess(sig) sig = at.AudioSignal(filepath) 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" 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") # 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}") 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), ) 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, ) # 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, ) 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, ) 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, ) def harp_vamp(input_audio, periodic_p, n_mask_codebooks, pitch_shift_amt, win_dur): 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=1, stretch_factor=1.0, api=True, ) 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 100s)", 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(): with gr.Accordion("manual controls", open=True): periodic_p = gr.Slider( label="periodic prompt", 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=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="compression prompt ", value=3, minimum=0, maximum=14, step=1, ) n_mask_codebooks_2 = gr.Number( label="compression prompt 2 via linear interpolation (0 == constant)", value=0, ) with gr.Accordion("extras ", open=False): pitch_shift_amt = gr.Slider( label="pitch shift amount (semitones)", minimum=-12, maximum=12, step=1, value=0, ) 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=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 ) 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=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 ) 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, masktemp, sampletemp, 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, pitch_shift_amt, sample_cutoff, num_feedback_steps, p2, n_mask_codebooks_2, win_dur } # connect widgets vamp_button.click( fn=vamp, inputs=_inputs, 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=[ input_audio, periodic_p, n_mask_codebooks, pitch_shift_amt, win_dur, ], output=out1, process_fn=harp_vamp, card=model_card ) try: demo.queue() demo.launch(share=True) except KeyboardInterrupt: shutil.rmtree("gradio-outputs", ignore_errors=True) raise