Blame for app.py

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)