import spaces # imported first, as ZeroGPU requires import tempfile from math import ceil from pathlib import Path import gradio as gr import numpy as np import soundfile as sf import torch import torch.nn.functional as F import torchaudio from flowhigh import FlowHighSR from huggingface_hub import snapshot_download from pyharp import ModelCard, build_endpoint from pyharp.tags import SampleRate, Subcategory SAMPLE_RATE = 48000 model_card = ModelCard( name="FLowHigh", description=( "Extend the bandwidth of low-resolution audio with flow matching. " "Repository: https://github.com/resemble-ai/flowhigh Paper: " "https://arxiv.org/abs/2501.04926" ), author="Jun-Hak Yun, Seung-Bin Kim, Seong-Whan Lee", tags=[Subcategory.BANDWIDTH_EXTENSION, SampleRate(SAMPLE_RATE), 'domain:speech'], ) # pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the # server, so the model is downloaded and placed there alone if __name__ != "__main__": checkpoint_dir = snapshot_download( repo_id="ResembleAI/FlowHigh", revision="4281fe4119e5f3d209a9a893218fed85af2e5bbe", allow_patterns=["*.json", "*.pt"], ) model = FlowHighSR.from_local(checkpoint_dir, device="cuda") # 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. LOSSLESS = {".wav", ".flac"} # The model runs over the whole recording at once, a pass per channel, so the cost grows with # its length and its channels. Its attention spans the whole recording, so memory grows with the # square of the length: 120 s measured at about 6 s, while 300 s ran out of GPU memory MAX_INPUT_SECONDS = 120 def gpu_seconds(input_audio_path: str, output_sample_rate: str) -> float: # Measured at about 6 s for 120 s of mono speech, padded info = sf.info(input_audio_path) return 10 + 0.1 * info.duration * info.channels def process_fn(input_audio_path: str, output_sample_rate: str) -> str: # Checked before the GPU is reserved, so a request the model cannot take spends none of the # caller's ZeroGPU allowance # Read with soundfile rather than torchaudio, whose ffmpeg decoding fails outside a GPU call # under ZeroGPU's emulation if not input_audio_path: raise gr.Error("Choose an audio file.") try: audio, input_rate = sf.read(input_audio_path, always_2d=True) except Exception as exc: raise gr.Error(f"Could not read the audio file: {exc}") from exc if len(audio) == 0: raise gr.Error("The audio file is empty.") duration = len(audio) / input_rate if duration > MAX_INPUT_SECONDS: raise gr.Error( f"The recording is {duration:.1f} s long. Trim it to {MAX_INPUT_SECONDS // 60} " "minutes or less." ) # Upstream divides by the input peak, so silence bypasses inference, and the GPU with it if not np.any(audio): output_rate = int(output_sample_rate) silence = torch.zeros(audio.shape[1], ceil(len(audio) * output_rate / input_rate)) return save(silence, output_rate, input_audio_path) return extend(input_audio_path, output_sample_rate) @spaces.GPU(duration=gpu_seconds) @torch.inference_mode() def extend(input_audio_path: str, output_sample_rate: str) -> str: audio, input_rate = torchaudio.load(input_audio_path) output_rate = int(output_sample_rate) output_length = ceil(audio.shape[-1] * output_rate / input_rate) native_length = ceil(audio.shape[-1] * SAMPLE_RATE / input_rate) minimum_length = ceil(2048 * input_rate / SAMPLE_RATE) if audio.shape[-1] < minimum_length: audio = F.pad(audio, (0, minimum_length - audio.shape[-1])) # Extending the bandwidth needs no other channel, so each is extended on its own and the # stereo image survives. The model normalizes what it is given, so each result is put back at # its channel's level relative to the loudest, and a silent channel, which it cannot # normalize, stays silent. peaks = audio.abs().amax(dim=-1) extended = { index: model.generate(audio[index:index + 1], input_rate, SAMPLE_RATE, timestep=1).cpu() * (peaks[index] / peaks.max()) for index in range(audio.shape[0]) if peaks[index] > 0 } length = next(iter(extended.values())).shape[-1] result = torch.cat([extended.get(index, torch.zeros(1, length)) for index in range(audio.shape[0])]) result = result[..., :native_length] if output_rate != SAMPLE_RATE: result = torchaudio.functional.resample(result, SAMPLE_RATE, output_rate) result = result[..., :output_length] peak = result.abs().max() if peak > 1: result *= 0.99 / peak return save(result, output_rate, input_audio_path) def save(result: torch.Tensor, output_rate: int, input_audio_path: str) -> str: # A WAV keeps the 32-bit float results, and FLAC takes 24-bit integers input_ext = Path(input_audio_path).suffix.lower() output_ext = input_ext if input_ext in LOSSLESS else ".wav" directory = Path(tempfile.mkdtemp(prefix="flowhigh-")) output_path = directory / f"{directory.name}{output_ext}" sf.write(output_path, result.T.numpy(), output_rate, subtype="FLOAT" if output_ext == ".wav" else "PCM_24") return str(output_path) if __name__ == "__main__": with gr.Blocks(title="FLowHigh", delete_cache=(600, 3600)) as demo: input_audio = ( gr.Audio(type="filepath", label="Input Audio") .harp_required(True) .set_info( "Low-bandwidth audio to extend, up to 2 minutes. Each channel is extended on its " "own." ) ) output_rate = gr.Dropdown( choices=["24000", "48000", "96000"], value="48000", label="Output Sample Rate", ).set_info("The model runs at 48 kHz, so 24 and 96 kHz outputs are resampled.") output_audio = gr.Audio(type="filepath", label="Enhanced Audio").set_info( "WAV at the selected sample rate, with the input's channels. Downloads are " "retained for one hour." ) build_endpoint( model_card=model_card, input_components=[input_audio, output_rate], output_components=[output_audio], process_fn=process_fn, ) demo.queue(default_concurrency_limit=1).launch(show_error=True, pwa=True, ssr_mode=False)