Spaces:
Running on Zero
Running on Zero
| 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) | |
| 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) | |