|
from pathlib import Path |
|
from threading import Thread |
|
|
|
import gdown |
|
import gradio as gr |
|
import librosa |
|
import numpy as np |
|
import torch |
|
|
|
from pipeline import build_audiosep |
|
|
|
CHECKPOINTS_DIR = Path("checkpoint") |
|
|
|
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
|
|
MODEL_NAME = CHECKPOINTS_DIR / "audiosep_base_4M_steps.ckpt" |
|
MODEL = None |
|
|
|
|
|
description = """ |
|
# AudioSep: Separate Anything You Describe |
|
[[Project Page]](https://audio-agi.github.io/Separate-Anything-You-Describe) [[Paper]](https://audio-agi.github.io/Separate-Anything-You-Describe/AudioSep_arXiv.pdf) [[Code]](https://github.com/Audio-AGI/AudioSep) |
|
|
|
AudioSep is a foundation model for open-domain sound separation with natural language queries. |
|
AudioSep demonstrates strong separation performance and impressivezero-shot generalization ability on |
|
numerous tasks such as audio event separation, musical instrument separation, and speech enhancement. |
|
""" |
|
|
|
|
|
def get_model(): |
|
model = build_audiosep( |
|
config_yaml="config/audiosep_base.yaml", |
|
checkpoint_path=MODEL_NAME, |
|
device=DEVICE, |
|
) |
|
return model |
|
|
|
|
|
def inference(audio_file_path: str, text: str): |
|
print(f"Separate audio from [{audio_file_path}] with textual query [{text}]") |
|
mixture, _ = librosa.load(audio_file_path, sr=32000, mono=True) |
|
|
|
with torch.no_grad(): |
|
text = [text] |
|
|
|
conditions = MODEL.query_encoder.get_query_embed( |
|
modality="text", text=text, device=DEVICE |
|
) |
|
|
|
input_dict = { |
|
"mixture": torch.Tensor(mixture)[None, None, :].to(DEVICE), |
|
"condition": conditions, |
|
} |
|
|
|
sep_segment = MODEL.ss_model(input_dict)["waveform"] |
|
|
|
sep_segment = sep_segment.squeeze(0).squeeze(0).data.cpu().numpy() |
|
|
|
return 32000, np.round(sep_segment * 32767).astype(np.int16) |
|
|
|
|
|
def download_models(): |
|
CHECKPOINTS_DIR.mkdir(exist_ok=True) |
|
success_file = CHECKPOINTS_DIR / "_SUCCESS" |
|
|
|
models = ( |
|
( |
|
"https://drive.google.com/file/d/1wQuXThdATXrkmkPM2sRGaNapJ4mTqmlY/view?usp=sharing", |
|
MODEL_NAME, |
|
), |
|
( |
|
"https://drive.google.com/file/d/11oj8_tPG6SXgw5fIEsZ5HiWZnJOrvdhw/view?usp=sharing", |
|
CHECKPOINTS_DIR / "music_speech_audioset_epoch_15_esc_89.98.pt", |
|
), |
|
) |
|
|
|
def download(models): |
|
for model_url, model_path in models: |
|
gdown.download(model_url, str(model_path), quiet=False, fuzzy=True) |
|
|
|
success_file.touch() |
|
|
|
global MODEL |
|
MODEL = get_model() |
|
button.update(value="Separate", interactive=True) |
|
|
|
if not success_file.exists(): |
|
thread = Thread(target=download, args=[models]) |
|
thread.start() |
|
|
|
|
|
with gr.Blocks(title="AudioSep") as demo: |
|
gr.Markdown(description) |
|
with gr.Row(): |
|
with gr.Column(): |
|
input_audio = gr.Audio() |
|
text = gr.Textbox() |
|
with gr.Column(): |
|
with gr.Column(): |
|
output_audio = gr.Audio(scale=10) |
|
button = gr.Button( |
|
"Downloading the models...", |
|
variant="primary", |
|
scale=2, |
|
size="lg", |
|
interactive=False, |
|
) |
|
button.click( |
|
fn=inference, inputs=[input_audio, text], outputs=[output_audio] |
|
) |
|
|
|
download_models() |
|
|
|
demo.queue().launch(share=True) |
|
|
