Blame for app.py

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)