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)