Spaces:
Running on Zero
Running on Zero
| import spaces # imported first, as ZeroGPU requires | |
| try: # torch>=2.6 flipped torch.load(weights_only) to True; legacy ckpts need False | |
| import torch as _torch | |
| if getattr(_torch.load, "__harp_compat__", False) is False: | |
| _torch_load_orig = _torch.load | |
| def _torch_load_compat(*args, **kwargs): | |
| # Lightning passes weights_only=None, which torch reads as True | |
| if kwargs.get("weights_only") is None: | |
| kwargs["weights_only"] = False | |
| return _torch_load_orig(*args, **kwargs) | |
| _torch_load_compat.__harp_compat__ = True | |
| _torch.load = _torch_load_compat | |
| except Exception: # torch not installed / unexpected API -- nothing to patch | |
| pass | |
| import time | |
| import numpy as np | |
| import soundfile as sf | |
| import audiotools as at | |
| import torch | |
| from pyharp.core import build_endpoint, ModelCard | |
| from pyharp.media.audio import save_audio | |
| from pyharp.tags import Subcategory, SampleRate, Channels | |
| import gradio as gr | |
| from vampnet.interface import Interface | |
| from vampnet import mask as pmask | |
| model_card = ModelCard( | |
| name="VampNet", | |
| description=( | |
| "Generates new audio by remixing and extending your input, filling in " | |
| "the blanks to create a coherent continuation in the same style. " | |
| "Besides the default, three fine-tuned models are available. cat10 was " | |
| "fine-tuned on expressive cat sounds, lazaro-ros-sep on Afro-Cuban " | |
| "percussion extracted from the Lázaro Ros discography, and supermetroid " | |
| "on Super Metroid (SNES) gameplay audio, for retro, game-style " | |
| "soundscapes. Generation runs in mono at the model's own rate. The " | |
| "result is resampled back to the input's rate and copied across its " | |
| "channels, so the output has the format it came in with but no stereo " | |
| "image. The code is MIT, but the " | |
| "pretrained weights are CC-BY-NC-SA-4.0, so results cannot be used " | |
| "commercially and must be shared alike. Repository: " | |
| "https://github.com/hugofloresgarcia/vampnet Paper: " | |
| "https://arxiv.org/abs/2307.04686" | |
| ), | |
| author="Hugo Flores Garcia, Prem Seetharaman, Rithesh Kumar, Bryan Pardo", | |
| tags=[Subcategory.INFILLING, SampleRate(44100), Channels(1), 'sound-effects', 'domain:music'] | |
| ) | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| init_model_choice = open("DEFAULT_MODEL").read().strip() | |
| MODEL_CHOICES = Interface.available_models() | |
| # pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the | |
| # server, so the default model is built and placed there alone. Each weight set is about 2.4 GB, | |
| # too much to hold all four together, so the others are only fetched here, and built inside the | |
| # request that picks one. Their paths are kept, since load_finetuned would list the model | |
| # repository over the network from inside that request's GPU reservation. | |
| if __name__ != "__main__": | |
| from vampnet import download_default, download_finetuned | |
| interface = Interface.default() | |
| interface.load_finetuned(init_model_choice) | |
| interface.to(DEVICE) | |
| CHECKPOINTS = { | |
| name: download_default() if name == "default" else download_finetuned(name) | |
| for name in MODEL_CHOICES | |
| } | |
| print(f"Loaded initial model: {init_model_choice}") | |
| 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_internal( | |
| seed, input_audio, model_choice, | |
| pitch_shift_amt, periodic_p, | |
| n_mask_codebooks, onset_mask_width, | |
| dropout, sampletemp, typical_filtering, | |
| typical_mass, typical_min_tokens, top_p, | |
| sample_cutoff, stretch_factor, sampling_steps, beat_mask_ms, num_feedback_steps | |
| ): | |
| t0 = time.time() | |
| _seed = seed if seed > 0 else None | |
| if _seed is None: | |
| _seed = int(torch.randint(0, 2**32, (1,)).item()) | |
| at.util.seed(_seed) | |
| if input_audio is None: | |
| raise gr.Error("no input audio received!") | |
| sr, input_audio = input_audio | |
| input_audio = input_audio / np.iinfo(input_audio.dtype).max | |
| sig = at.AudioSignal(input_audio, sr).to_mono() | |
| loudness = sig.loudness() | |
| sig = interface._preprocess(sig) | |
| # Builds the chosen model's checkpoints, unless they are the ones already built | |
| interface.reload(*CHECKPOINTS[model_choice]) | |
| if pitch_shift_amt != 0: | |
| sig = shift_pitch(sig, pitch_shift_amt) | |
| codes = interface.encode(sig) | |
| mask = interface.build_mask( | |
| codes, | |
| sig=sig, | |
| periodic_prompt=periodic_p, | |
| onset_mask_width=onset_mask_width, | |
| _dropout=dropout, | |
| upper_codebook_mask=n_mask_codebooks, | |
| ) | |
| if beat_mask_ms > 0: | |
| 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) | |
| 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_z = interface.vamp( | |
| codes, mask, | |
| 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=top_p, | |
| seed=_seed, | |
| sample_cutoff=sample_cutoff, | |
| ) | |
| print(f"vamp took {time.time() - t0} seconds") | |
| sig = interface.decode(codes) | |
| sig = sig.normalize(loudness) | |
| return sig | |
| # Generation covers the whole input, stretched, in 10 s chunks, so the cost grows with the | |
| # output's length | |
| # PyHARP stops a job that runs past its timeout_s, 15 minutes by default, so a request whose | |
| # reservation would exceed that is refused rather than started. A shorter one is left to | |
| # ZeroGPU, which admits it when the caller's own allowance covers it. | |
| MAX_GPU_SECONDS = 900 | |
| def gpu_seconds(input_audio, sampletemp, periodic_p, dropout, n_mask_codebooks, model_choice, | |
| stretch_factor): | |
| # Measured at about 15 s for 30 s of output and 19 s for 64 s, padded, with room for loading | |
| # a fine-tuned model | |
| return 15 + 0.4 * sf.info(input_audio).duration * max(1, int(stretch_factor)) | |
| def process_fn(input_audio, sampletemp, periodic_p, dropout, n_mask_codebooks, model_choice, | |
| stretch_factor): | |
| # Checked before the GPU is reserved, so a request the model cannot take spends none of the | |
| # caller's ZeroGPU allowance | |
| if model_choice not in MODEL_CHOICES: | |
| raise gr.Error(f"Unknown model: {model_choice}") | |
| if not input_audio: | |
| raise gr.Error("Provide the audio to rework.") | |
| try: | |
| duration = sf.info(input_audio).duration | |
| except Exception as exc: | |
| raise gr.Error(f"Could not read the audio file: {exc}") from exc | |
| if duration == 0: | |
| raise gr.Error("The audio file is empty.") | |
| output_seconds = duration * max(1, int(stretch_factor)) | |
| if gpu_seconds(input_audio, sampletemp, periodic_p, dropout, n_mask_codebooks, model_choice, | |
| stretch_factor) > MAX_GPU_SECONDS: | |
| raise gr.Error( | |
| f"That would generate {output_seconds:.0f} s of audio, more than one request can in " | |
| f"the {MAX_GPU_SECONDS // 60} minutes it may run. Use a shorter input or a smaller " | |
| "stretch factor." | |
| ) | |
| return harp_vamp(input_audio, sampletemp, periodic_p, dropout, n_mask_codebooks, | |
| model_choice, stretch_factor) | |
| def harp_vamp(input_audio, sampletemp, periodic_p, dropout, n_mask_codebooks, model_choice, stretch_factor): | |
| sig = at.AudioSignal(input_audio) | |
| # Captured before the downmix, because generation runs in mono at VampNet's | |
| # own rate and both have to be put back afterwards | |
| original_sr = sig.sample_rate | |
| original_channels = sig.num_channels | |
| sig = sig.to_mono() | |
| input_audio = sig.cpu().detach().numpy()[0][0] | |
| input_audio = input_audio * np.iinfo(np.int16).max | |
| input_audio = input_audio.astype(np.int16) | |
| input_audio = input_audio.reshape(1, -1) | |
| input_audio = (sig.sample_rate, input_audio) | |
| sig = _vamp_internal( | |
| seed=0, | |
| input_audio=input_audio, | |
| model_choice=model_choice, | |
| pitch_shift_amt=0, | |
| periodic_p=int(periodic_p), | |
| n_mask_codebooks=int(n_mask_codebooks), | |
| onset_mask_width=0, | |
| dropout=dropout, | |
| sampletemp=sampletemp, | |
| typical_filtering=False, | |
| typical_mass=0.15, | |
| typical_min_tokens=1, | |
| top_p=None, | |
| sample_cutoff=1.0, | |
| stretch_factor=stretch_factor, | |
| sampling_steps=36, | |
| beat_mask_ms=int(0), | |
| num_feedback_steps=1, | |
| ) | |
| sig = sig.detach().cpu() | |
| if sig.sample_rate != original_sr: | |
| sig = sig.resample(original_sr) | |
| # Copies of one generated channel, not a stereo image, which the card says | |
| if original_channels > 1 and sig.num_channels == 1: | |
| sig.audio_data = sig.audio_data.repeat(1, original_channels, 1) | |
| return save_audio(sig) | |
| if __name__ == "__main__": | |
| with gr.Blocks() as demo: | |
| harp_input_components = [ | |
| gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info(( | |
| "The audio to rework. The model masks out parts of this recording " | |
| "and fills them back in, so the result keeps its character." | |
| )), | |
| gr.Slider(label="Randomness", minimum=0.1, maximum=10.0, value=1.0, step=0.001, | |
| info="Controls sampling diversity- higher values yield more novel variations."), | |
| gr.Slider(label="Rhythmic Reference", minimum=0, maximum=13, step=1, value=7, | |
| info="Determines how frequently rhythmic cues are repeated in generation."), | |
| gr.Slider(label="Variation Amount", minimum=0.0, maximum=1.0, step=0.01, value=0.0, | |
| info="How much of the input the model changes. 0 = close copy, higher = greater remix."), | |
| gr.Slider(label="Timbral Reference", value=3, minimum=1, maximum=14, step=1, | |
| info="Controls the richness of timbre retained from the source audio. Lower = simpler, higher = richer timbre."), | |
| gr.Dropdown(label="Model Choice", choices=list(MODEL_CHOICES), value=init_model_choice, | |
| info="Select which fine-tuned VampNet model to use."), | |
| gr.Slider(label="Time Stretch Factor", minimum=1, maximum=8, step=1, value=1, | |
| info="Scales the output's length relative to the input. 1 = the same length, 2 = twice as long."), | |
| ] | |
| harp_output_components = [ | |
| gr.Audio(type="filepath", label="Generated Audio").set_info(( | |
| "The reworked audio, at the same sample rate as the input." | |
| )), | |
| ] | |
| build_endpoint( | |
| model_card=model_card, | |
| input_components=harp_input_components, | |
| output_components=harp_output_components, | |
| process_fn=process_fn, | |
| ) | |
| demo.queue() | |
| demo.launch(pwa=True, show_error=True, ssr_mode=False) | |