import spaces # imported before torch, as ZeroGPU requires import math import os import tempfile import soundfile as sf import torch from audiosr import build_model, super_resolution_long_audio from huggingface_hub import hf_hub_download from pyharp import ModelCard, build_endpoint, SampleRate, Channels from pyharp.tags import Subcategory import gradio as gr import random import numpy as np import librosa from scipy import signal import audiosr.pipeline as _audiosr_pipeline from audiosr.utils import _locate_cutoff_freq, lowpass as _audiosr_lowpass os.environ["TOKENIZERS_PARALLELISM"] = "true" torch.set_float32_matmul_precision("high") def _patched_lowpass_filtering_prepare_inference(dl_output): """ Patches a boundary bug in audiosr==0.0.7: cutoff_freq caps at exactly 24000 Hz (both via normal detection saturating there, and via the "give up" fallback for quiet/narrowband audio). sampling_rate for this step is hardcoded to 48000 upstream, so nyquist == 24000 exactly, producing Wn == cutoff_freq / nyquist == 1.0 for a wide range of real inputs. scipy's iirfilter requires 0 < Wn < 1 strictly. Clamped with a small safety margin below Nyquist instead of patching the pip package. """ waveform = dl_output["waveform"] sampling_rate = dl_output["sampling_rate"] cutoff_freq = (_locate_cutoff_freq(dl_output["stft"], percentile=0.985) / 1024) * 24000 if cutoff_freq < 1000: cutoff_freq = 24000 nyquist = 0.5 * sampling_rate if cutoff_freq >= nyquist: cutoff_freq = nyquist * 0.98 # safety margin order = 8 ftype = np.random.choice(["butter", "cheby1", "ellip", "bessel"]) filtered_audio = _audiosr_lowpass( waveform.numpy().squeeze(), highcut=cutoff_freq, fs=sampling_rate, order=order, _type=ftype, ) filtered_audio = torch.FloatTensor(filtered_audio.copy()).unsqueeze(0) if waveform.size(-1) <= filtered_audio.size(-1): filtered_audio = filtered_audio[..., : waveform.size(-1)] else: filtered_audio = torch.nn.functional.pad(filtered_audio, (0, waveform.size(-1) - filtered_audio.size(-1))) return {"waveform_lowpass": filtered_audio} _audiosr_pipeline.lowpass_filtering_prepare_inference = _patched_lowpass_filtering_prepare_inference DEVICE = "cuda" if torch.cuda.is_available() else "cpu" model = None model_card = ModelCard( name="AudioSR", description=( "Upsample any audio to 48kHz using audio super-resolution. Output is 48 " "kHz, which is the point of the model rather than a change of format. " "Super-resolution runs in mono, so a stereo input's result is copied " "across its channels and carries no stereo image. Repository: " "https://github.com/haoheliu/versatile_audio_super_resolution Paper: " "https://arxiv.org/abs/2309.07314" ), author="Haohe Liu, Ke Chen, Qiao Tian, Wenwu Wang, Mark D. Plumbley", tags=[Subcategory.BANDWIDTH_EXTENSION, SampleRate(48000), Channels(1), 'domain:music', 'domain:speech', 'domain:general-audio'], ) # Built in the worker, which runs process_fn, on its main thread, where ZeroGPU's emulation records # the CUDA placement, so that no request spends the caller's allowance on building it. The server # never runs process_fn. audiosr downloads its checkpoint unpinned, and ignores a path it is given, # so the pinned file, which the README preloads, is handed to it in place of its own download. if __name__ != "__main__": _checkpoint = hf_hub_download("haoheliu/audiosr_basic", "pytorch_model.bin", revision="74a47f49061a1e788e968cc43ad45c0b6243f37d") _audiosr_pipeline.download_checkpoint = lambda model_name: _checkpoint model = build_model(model_name="basic", device=DEVICE) def match_array_shapes(array_1: np.ndarray, array_2: np.ndarray): if (len(array_1.shape) == 1) & (len(array_2.shape) == 1): if array_1.shape[0] > array_2.shape[0]: array_1 = array_1[:array_2.shape[0]] elif array_1.shape[0] < array_2.shape[0]: array_1 = np.pad(array_1, ((array_2.shape[0] - array_1.shape[0], 0)), 'constant', constant_values=0) else: if array_1.shape[1] > array_2.shape[1]: array_1 = array_1[:,:array_2.shape[1]] elif array_1.shape[1] < array_2.shape[1]: padding = array_2.shape[1] - array_1.shape[1] array_1 = np.pad(array_1, ((0,0), (0,padding)), 'constant', constant_values=0) return array_1 def lr_filter(audio, cutoff, filter_type, order=12, sr=48000): audio = audio.T nyquist = 0.5 * sr normal_cutoff = cutoff / nyquist b, a = signal.butter(order//2, normal_cutoff, btype=filter_type, analog=False) sos = signal.tf2sos(b, a) filtered_audio = signal.sosfiltfilt(sos, audio) return filtered_audio.T # The audio is upsampled in 15-second chunks that start 13 seconds apart, each at the number of # DDIM steps asked for, so the cost grows with both # 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_path: str, ddim_steps: int, guidance_scale: float, seed: int) -> float: # Measured at about 34 s for 8 s of audio (one chunk) and 47 s for 30 s (three) at 50 steps, # after about 15 s of setup, padded chunks = math.ceil(sf.info(input_audio_path).duration / 13) return 25 + 12 * chunks * int(ddim_steps) / 50 def process_fn(input_audio_path: str, ddim_steps: int, guidance_scale: float, seed: str, multiband_ensemble: bool, input_cutoff: int) -> str: # Checked before the GPU is reserved, so a request the model cannot take spends none of the # caller's ZeroGPU allowance if not input_audio_path: raise gr.Error("Provide the audio to upsample.") try: info = sf.info(input_audio_path) except Exception as exc: raise gr.Error(f"Could not read the audio file: {exc}") from exc if info.frames == 0: raise gr.Error("The audio file is empty.") if gpu_seconds(input_audio_path, ddim_steps, guidance_scale, 0) > MAX_GPU_SECONDS: raise gr.Error( f"{int(ddim_steps)} DDIM steps on {info.duration:.0f} s of audio need more GPU time than " "a request can reserve. Lower the steps or trim the audio." ) try: seed_val = int(seed) if seed and seed.strip() not in ("0", "") else random.randint(1, 2**32 - 1) except (TypeError, ValueError): seed_val = random.randint(1, 2**32 - 1) # AudioSR's internal cutoff-frequency detection (lowpass_filtering_prepare_inference) # hardcodes a 24kHz-Nyquist assumption regardless of actual input sample rate, # which can produce an invalid filter Wn and crash on input below 48kHz. # Resample up front so that assumption always holds. orig_sr = info.samplerate original_channels = info.channels # A lossless input keeps its container. A lossy one comes back as WAV, since encoding the # result lossily again would expose the artifacts of the first encode and add a second set. input_ext = os.path.splitext(input_audio_path)[1].lower() output_ext = input_ext if input_ext in {".wav", ".flac"} else ".wav" if orig_sr < 48000: y, _ = librosa.load(input_audio_path, sr=None, mono=False) y = librosa.resample(y, orig_sr=orig_sr, target_sr=48000) resampled_path = tempfile.mktemp(suffix=".wav") sf.write(resampled_path, y.T if y.ndim > 1 else y, samplerate=48000) input_audio_path = resampled_path output = upsample(input_audio_path, int(ddim_steps), float(guidance_scale), seed_val) if multiband_ensemble: crossover_freq = int(input_cutoff) - 1000 low, _ = librosa.load(input_audio_path, sr=48000, mono=True) out = output.squeeze() if output.ndim > 1 else output out = match_array_shapes(out, low) low = lr_filter(low, crossover_freq, 'lowpass', order=10) high = lr_filter(out, crossover_freq, 'highpass', order=10) high = lr_filter(high, 23000, 'lowpass', order=2) output = low + high else: if output.shape[0] == 1: output = output.squeeze(0) else: output = output.T # Super-resolution runs in mono, so a stereo input comes back as one # channel; it is copied across the input's channels rather than left mono, # which a session on a stereo track would otherwise misread if original_channels > 1 and output.ndim == 1: output = np.tile(output[:, None], (1, original_channels)) output_path = tempfile.mktemp(suffix=output_ext) sf.write(output_path, output, samplerate=48000) return output_path @spaces.GPU(duration=gpu_seconds) @torch.inference_mode() def upsample(input_audio_path: str, ddim_steps: int, guidance_scale: float, seed: int): waveform = super_resolution_long_audio( model, input_audio_path, seed=seed, guidance_scale=guidance_scale, ddim_steps=ddim_steps, ) return waveform.cpu().numpy() if __name__ == "__main__": with gr.Blocks() as demo: input_components = [ gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info(( "Low-bandwidth audio to upsample. The result is written at 48 kHz " "whatever the input rate was." )), gr.Slider(minimum=10, maximum=500, step=10, value=50, label="DDIM Steps", info="More steps = better quality but slower"), gr.Slider(minimum=1.0, maximum=20.0, step=0.5, value=3.5, label="Guidance Scale", info="Higher values follow the conditioning more closely"), gr.Textbox(value="0", label="Seed", info="0 = random seed"), gr.Checkbox(value=False, label="Multiband Ensemble", info="Blend original low frequencies with upsampled highs"), gr.Slider(minimum=4000, maximum=20000, step=1000, value=12000, label="Input Cutoff (Hz)", info="Crossover frequency for multiband ensemble"), ] output_components = [ gr.Audio(type="filepath", label="Upsampled Audio").set_info("Audio upsampled to 48kHz."), ] app = build_endpoint( model_card=model_card, input_components=input_components, output_components=output_components, process_fn=process_fn, ) print("Launching Gradio...", flush=True) demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True, pwa=True, ssr_mode=False)