Spaces:

teamup-tech
/

RemFX

Running

RemFX / app.py

Drop the comment on ssr_mode

7d7ef0e verified 3 days ago

9.17 kB

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)

	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)