app.py

import sys
import os
import tempfile
import soundfile as sf
import torch
from audiosr import build_model, super_resolution_long_audio
from pyharp import ModelCard, build_endpoint
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

model = None

def get_model():
    global model
    if model is None:
        print("Loading AudioSR model...", flush=True)
        model = build_model(model_name="basic", device="auto")
        print("Model loaded.", flush=True)
    return model

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

model_card = ModelCard(
    name="AudioSR",
    description="Upsample any audio to 48kHz using audio super-resolution.",
    author="Haohe Liu, Ke Chen, Qiao Tian, Wenwu Wang, Mark D. Plumbley",
    tags=["audio", "super-resolution", "upsampling", "48kHz"],
)

@torch.inference_mode()
def process_fn(input_audio_path: str, ddim_steps: int, guidance_scale: float, seed: str, multiband_ensemble: bool, input_cutoff: int) -> str:
    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 = sf.info(input_audio_path).samplerate
    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

    waveform = super_resolution_long_audio(
        get_model(),
        input_audio_path,
        seed=seed_val,
        guidance_scale=float(guidance_scale),
        ddim_steps=int(ddim_steps),
    )

    output = waveform.cpu().numpy()

    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

    output_path = tempfile.mktemp(suffix=".wav")
    sf.write(output_path, output, samplerate=48000)
    return output_path

with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio").harp_required(True),
        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="Output Audio (48kHz)").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)