Spaces:
Sleeping
Sleeping
| 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") | |
| # The model will be loaded in the future | |
| 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) | |