app.py

import gradio as gr
import torch
import torchaudio
import random
import tempfile

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
import spaces


# ---- 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)
    # checkpoints contain the full pytorch-lightning state dict
    state_dict = torch.load(ckpt_path, map_location=DEVICE)["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

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 detect_effects(audio_path):
    """Run classifier and return per-effect confidence scores."""
    audio = load_and_prep(audio_path)
    with torch.no_grad():
        # classifier returns a list of 5 tensors (one per effect class), each [batch, 1]
        outputs = classifier(audio)
        confidences = torch.hstack(outputs).squeeze(0)  # -> [5]

    return {
        Pedalboard_Effects[i].__name__.replace("RandomPedalboard", ""): confidences[i].item()
        for i in range(len(Pedalboard_Effects))
    }


@torch.inference_mode()
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


@spaces.GPU(duration=120)
def process_fn(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 match_input_format(out_path, audio_path)


# ---- gradio ui ----

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)],
)

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",
        )

        # "Detect Effects" only runs the classifier so users can see scores before committing
        detect_btn = gr.Button("Detect Effects")
        detected = gr.Label(num_top_classes=5, label="Detected Effects")
        detect_btn.click(fn=detect_effects, inputs=audio_input, outputs=detected)

        output_audio = gr.Audio(type="filepath", label="Output Audio").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)