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)