RemFx / app.py
mattricesound's picture
Update app.py
ccbf851
import gradio as gr
import torch
import torchaudio
import hydra
from hydra import compose, initialize
import random
import os
from remfx import effects
cfg = None
classifier = None
models = {}
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
ALL_EFFECTS = effects.Pedalboard_Effects
def init_hydra():
global cfg
initialize(config_path="cfg", job_name="remfx", version_base="2.0")
cfg = compose(config_name="config", overrides=["+exp=remfx_detect"])
def load_models():
global classifier
print("Loading models")
classifier = hydra.utils.instantiate(cfg.classifier, _convert_="partial")
ckpt_path = cfg.classifier_ckpt
state_dict = torch.load(ckpt_path, map_location=device)["state_dict"]
classifier.load_state_dict(state_dict)
classifier.to(device)
for effect in cfg.ckpts:
model = hydra.utils.instantiate(cfg.ckpts[effect].model, _convert_="partial")
ckpt_path = cfg.ckpts[effect].ckpt_path
state_dict = torch.load(ckpt_path, map_location=device)["state_dict"]
model.load_state_dict(state_dict)
model.to(device)
models[effect] = model
def audio_classification(audio_file):
audio, sr = torchaudio.load(audio_file)
audio = torchaudio.transforms.Resample(sr, cfg.sample_rate)(audio)
# Convert to mono
audio = audio.mean(0, keepdim=True)
# Add dimension for batch
audio = audio.unsqueeze(0)
audio = audio.to(device)
with torch.no_grad():
# Classify
print("Detecting effects")
labels = torch.tensor(classifier(audio))
labels_dict = {
ALL_EFFECTS[i].__name__.replace("RandomPedalboard", ""): labels[i].item()
for i in range(len(ALL_EFFECTS))
}
return labels_dict
def audio_removal(audio_file, labels, threshold):
audio, sr = torchaudio.load(audio_file)
audio = torchaudio.transforms.Resample(sr, cfg.sample_rate)(audio)
# Convert to mono
audio = audio.mean(0, keepdim=True)
# Add dimension for batch
audio = audio.unsqueeze(0)
audio = audio.to(device)
label_names = [f"RandomPedalboard{lab['label']}" for lab in labels["confidences"]]
logits = torch.tensor([lab["confidence"] for lab in labels["confidences"]])
rem_fx_labels = torch.where(logits > threshold, 1.0, 0.0)
effects_present = [
name for name, effect in zip(label_names, rem_fx_labels) if effect == 1.0
]
print("Removing effects:", effects_present)
# Remove effects
# Shuffle effects order
effects_order = cfg.inference_effects_ordering
random.shuffle(effects_order)
# Get the correct effect by search for names in effects_order
effects = [effect for effect in effects_order if effect in effects_present]
elem = audio
with torch.no_grad():
for effect in effects:
# Sample the model
elem = models[effect].model.sample(elem)
output = elem.squeeze(0)
waveform = gr.make_waveform((cfg.sample_rate, output[0].numpy()))
return waveform
def ui():
css = """
#classifier {
padding-top: 40px;
}
#classifier .output-class {
display: none;
}
"""
with gr.Blocks(css=css) as interface:
gr.HTML(
"""
<div style="text-align: center; max-width: 700px; margin: 0 auto;">
<div
style="
display: inline-flex; align-items: center; gap: 0.8rem; font-size: 1.75rem;
"
>
<h1 style="font-weight: 900; margin-bottom: 7px; line-height: normal;">
RemFx: General Purpose Audio Effect Removal
</h1>
</div> <p style="margin-bottom: 10px; font-size: 94%">
<a href="https://arxiv.org/abs/2308.16177">[Paper]</a> <a href="https://csteinmetz1.github.io/RemFX/">[Project
page]</a>
</p>
</div>
"""
)
gr.HTML(
"""
<div style="text-align: left;"> This is our demo for the paper General Purpose Audio Effect Removal. It uses the RemFX Detect system described in the paper to detect the audio effects that are present and remove them. <br>
To use the demo, use one of our curated examples or upload your own audio file and click submit. The system will then detect the effects present in the audio remove them if they meet the threshold. </div>
"""
)
with gr.Row():
with gr.Column():
effected_audio = gr.Audio(
source="upload",
type="filepath",
label="File",
interactive=True,
elem_id="melody-input",
)
submit = gr.Button("Submit")
threshold = gr.Slider(
minimum=0.0,
maximum=1.0,
step=0.1,
value=0.5,
label="Detection Threshold",
)
with gr.Column():
classifier = gr.Label(
num_top_classes=5, label="Effects Present", elem_id="classifier"
)
audio_output = gr.Video(label="Output")
gr.Examples(
fn=audio_removal,
examples=[
["./input_examples/guitar.wav"],
["./input_examples/vocal.wav"],
["./input_examples/bass.wav"],
["./input_examples/drums.wav"],
["./input_examples/crazy_guitar.wav"],
],
inputs=effected_audio,
)
submit.click(
audio_classification,
inputs=[effected_audio],
outputs=[classifier],
queue=False,
show_progress=False,
).then(
audio_removal,
inputs=[effected_audio, classifier, threshold],
outputs=[audio_output],
)
interface.queue().launch()
if __name__ == "__main__":
init_hydra()
load_models()
ui()