Blame for app.py

import spaces  # imported first, as ZeroGPU requires

import gradio as gr
import soundfile as sf
import torch
import torchaudio
import random
import tempfile
from pathlib import Path

from pyharp import ModelCard, build_endpoint, SampleRate, Channels
from pyharp.tags import Subcategory
from huggingface_hub import hf_hub_download

from remfx.models import RemFX, DemucsModel, DCUNetModel, FXClassifier
from remfx.classifier import Cnn14
from remfx.effects import Pedalboard_Effects

model_card = ModelCard(
    name="RemFX",
    description=(
        "Detects and removes audio effects from recordings. The model is "
        "fixed-rate and mono, so the result is resampled back to the input's "
        "rate and copied across its channels: the output has the format it came "
        "in with, but no stereo image. Repository: "
        "https://github.com/mhrice/RemFx Paper: "
        "https://arxiv.org/abs/2308.16177"
    ),
    author="Matthew Rice, Christian J. Steinmetz, George Fazekas, Joshua D. Reiss",
    tags=[Subcategory.EFFECT_REMOVAL, SampleRate(48000), Channels(1), 'domain:music'],
)

# 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"}

# ---- config ----

SAMPLE_RATE = 48000 # from remfx_detect.yaml
NUM_BINS = 1025       # n_fft // 2 + 1, where stft n_fft=2048
HF_REPO = "vanshc/remfx-checkpoints"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# order in which effects are removed; shuffled at runtime (matches training config)
EFFECTS_ORDER = [
    "RandomPedalboardDistortion",
    "RandomPedalboardCompressor",
    "RandomPedalboardReverb",
    "RandomPedalboardChorus",
    "RandomPedalboardDelay",
]


# ---- model loading ----

def download_and_load(model, ckpt_filename):
    ckpt_path = hf_hub_download(repo_id=HF_REPO, filename=ckpt_filename,
                                revision="c860ca07d4e4fa8d88700c1942c450a98ffd8943")
    # checkpoints contain the full pytorch-lightning state dict
    state_dict = torch.load(ckpt_path, map_location="cpu")["state_dict"]
    model.load_state_dict(state_dict)
    return model.to(DEVICE).eval()


def build_effect_model(network, ckpt_filename):
    # RemFX is the lightning wrapper around each effect-removal network
    model = RemFX(
        lr=1e-4, lr_beta1=0.95, lr_beta2=0.999, lr_eps=1e-6,
        lr_weight_decay=1e-3, sample_rate=SAMPLE_RATE, network=network,
    )
    return download_and_load(model, ckpt_filename)


# demucs handles distortion + compressor (time-domain)
# dcunet handles reverb, chorus, delay (frequency-domain)
# params come from cfg/exp/remfx_detect.yaml and cfg/model/demucs.yaml + dcunet.yaml

demucs_kwargs = dict(sources=["mixture"], audio_channels=1, nfft=4096, channels=48) # from demucs.yaml
dcunet_kwargs = dict(architecture="Large-DCUNet-20", stft_kernel_size=512, fix_length_mode="pad") # from dcunet.yaml

# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so every model is loaded there alone
if __name__ != "__main__":
    print("Loading effect models...")

    effect_models = {
        "RandomPedalboardDistortion": build_effect_model(
            DemucsModel(SAMPLE_RATE, **demucs_kwargs), "demucs_distortion_aug.ckpt"
        ),
        "RandomPedalboardCompressor": build_effect_model(
            DemucsModel(SAMPLE_RATE, **demucs_kwargs), "demucs_compressor_aug.ckpt"
        ),
        "RandomPedalboardReverb": build_effect_model(
            DCUNetModel(SAMPLE_RATE, NUM_BINS, **dcunet_kwargs), "dcunet_reverb_aug.ckpt"
        ),
        "RandomPedalboardChorus": build_effect_model(
            DCUNetModel(SAMPLE_RATE, NUM_BINS, **dcunet_kwargs), "dcunet_chorus_aug.ckpt"
        ),
        "RandomPedalboardDelay": build_effect_model(
            DCUNetModel(SAMPLE_RATE, NUM_BINS, **dcunet_kwargs), "dcunet_delay_aug.ckpt"
        ),
    }

    print("Loading classifier...")

    cnn14 = Cnn14( # taken from remfx_detect.yaml
        num_classes=5,
        n_fft=2048, hop_length=512, n_mels=128,
        sample_rate=SAMPLE_RATE, model_sample_rate=SAMPLE_RATE,
        specaugment=True,
    )
    classifier = FXClassifier( # taken from remfx_detect.yaml
        lr=3e-4, lr_weight_decay=1e-3, sample_rate=SAMPLE_RATE,
        mixup=False, network=cnn14,
    )
    classifier = download_and_load(classifier, "classifier.ckpt")

    print("All models loaded.")


# ---- audio helpers ----

def load_and_prep(audio_path):
    audio, sr = torchaudio.load(audio_path)
    if sr != SAMPLE_RATE:
        audio = torchaudio.transforms.Resample(sr, SAMPLE_RATE)(audio)
    audio = audio.mean(dim=0, keepdim=True)  # stereo to mono: (1, T)
    audio = audio.unsqueeze(0)               # add batch dim for model: (1, 1, T)
    return audio.to(DEVICE)


# ---- inference ----

def match_input_format(output_path: str, input_path: str) -> str:
    """
    Rewrite an output file at the sample rate and channel count it came in with.

    The models here are fixed-rate and mono, and hand the writing to vendored
    upstream code, so the format is corrected afterwards rather than threaded
    through inference. Duplicated channels are copies, not a recovered stereo
    image; the model card says so.

    Args:
        output_path (str): File the model wrote. Overwritten in place.
        input_path (str): File the request supplied.

    Returns:
        output_path (str): The same path, for use in a return statement.
    """

    import librosa
    import numpy as np
    import soundfile

    want = soundfile.info(input_path)
    have = soundfile.info(output_path)

    if have.samplerate == want.samplerate and have.channels == want.channels:
        return output_path

    data, rate = soundfile.read(output_path, always_2d=True)

    if rate != want.samplerate:
        data = librosa.resample(
            data.T, orig_sr=rate, target_sr=want.samplerate
        ).T

    if data.shape[1] == 1 and want.channels > 1:
        data = np.tile(data, (1, want.channels))
    elif data.shape[1] > 1 and want.channels == 1:
        data = data.mean(axis=1, keepdims=True)

    soundfile.write(output_path, data, want.samplerate)

    return output_path


def process_fn(audio_path, threshold):
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if not audio_path:
        raise gr.Error("Provide the audio to clean.")
    try:
        info = sf.info(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.")

    out_path = match_input_format(remove_effects(audio_path, threshold), audio_path)

    input_ext = Path(audio_path).suffix.lower()
    if input_ext in LOSSLESS and input_ext != ".wav":
        data, rate = sf.read(out_path, always_2d=True)
        out_path = str(Path(out_path).with_suffix(input_ext))
        sf.write(out_path, data, rate)
    return out_path


@spaces.GPU(duration=120)
@torch.inference_mode()
def remove_effects(audio_path, threshold):
    audio = load_and_prep(audio_path)

    outputs = classifier(audio)
    confidences = torch.hstack(outputs).squeeze(0)

    effects_to_remove = {
        Pedalboard_Effects[i].__name__
        for i, conf in enumerate(confidences)
        if conf.item() > threshold
    }

    # shuffle removal order (inference_effects_shuffle is set to True in training config)
    order = EFFECTS_ORDER.copy()
    random.shuffle(order)

    elem = audio
    for effect in order:
        if effect in effects_to_remove:
            # .model is the inner network (DemucsModel or DCUNetModel) inside RemFX
            elem = effect_models[effect].model.sample(elem)

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        out_path = f.name   

    torchaudio.save(out_path, elem.squeeze(0).cpu(), sample_rate=SAMPLE_RATE)

    return out_path


# ---- gradio ui ----

if __name__ == "__main__":
    with gr.Blocks() as demo:

        audio_input = gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info((
            "A recording with effects on it. The model detects which effects "
            "are present and removes them, so a dry recording has nothing to "
            "strip."
        ))

        threshold = gr.Slider(
            minimum=0.0, maximum=1.0, step=0.1, value=0.5,
            label="Detection Threshold",
            info="Effects with confidence above this value will be removed",
        )

        output_audio = gr.Audio(type="filepath", label="Effects Removed").set_info("Effect-removed audio.")

        # build_endpoint adds the process / cancel / view controls buttons for HARP
        build_endpoint(
            model_card=model_card,
            input_components=[audio_input, threshold],
            output_components=[output_audio],
            process_fn=process_fn,
        )

    demo.queue().launch(pwa=True, show_error=True, ssr_mode=False)