Spaces:
Running
Running
| import spaces # imported first, as ZeroGPU requires | |
| import sys | |
| import numpy as np | |
| sys.stdout.reconfigure(line_buffering=True) # flush print() calls immediately so HF Space logs are live | |
| import os | |
| import tempfile | |
| import traceback | |
| from pathlib import Path | |
| import gradio as gr | |
| import librosa | |
| import torch | |
| import soundfile as sf | |
| from huggingface_hub import hf_hub_download | |
| from pyharp import ModelCard, build_endpoint, SampleRate, Channels | |
| from pyharp.tags import Subcategory | |
| from models.resunet import ResUNet30 | |
| from models.clap_encoder import CLAP_Encoder | |
| model_card = ModelCard( | |
| name="AudioSep", | |
| description=( | |
| "Separate any sound from a mixture using a natural language text " | |
| "description. Separation runs at 32 kHz in mono. The result is " | |
| "resampled back to the input's rate and copied across its channels, so " | |
| "the output has the format it came in with but not its stereo image. " | |
| "Repository: https://github.com/Audio-AGI/AudioSep Paper: " | |
| "https://arxiv.org/abs/2308.05037" | |
| ), | |
| author="Xubo Liu, Qiuqiang Kong, Yan Zhao, Haohe Liu, Yi Yuan, Yuzhuo Liu, Rui Xia, Yuxuan Wang, Mark D. Plumbley, Wenwu Wang", | |
| tags=[Subcategory.TARGET_SOURCE_EXTRACTION, SampleRate(32000), Channels(1), 'text-queried', 'domain:general-audio'], | |
| ) | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| # 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"} | |
| ss_model = None | |
| query_encoder = None | |
| model_error = None | |
| def load_model(): | |
| global ss_model, query_encoder, model_error | |
| try: | |
| print("Downloading checkpoints...") | |
| main_ckpt_path = hf_hub_download( | |
| repo_id="Audio-AGI/AudioSep", | |
| repo_type="space", | |
| filename="checkpoint/audiosep_base_4M_steps.ckpt", | |
| revision="0807a713283035d8b0d0628f1b8e3d6580459f7b", | |
| ) | |
| clap_ckpt_path = hf_hub_download( | |
| repo_id="Audio-AGI/AudioSep", | |
| repo_type="space", | |
| filename="checkpoint/music_speech_audioset_epoch_15_esc_89.98.pt", | |
| revision="0807a713283035d8b0d0628f1b8e3d6580459f7b", | |
| ) | |
| query_encoder = CLAP_Encoder(pretrained_path=clap_ckpt_path).eval() | |
| _ss_model = ResUNet30(input_channels=1, output_channels=1, condition_size=512) | |
| # Load weights from the Lightning checkpoint — keys are prefixed with "ss_model." | |
| state_dict = torch.load(main_ckpt_path, map_location="cpu", weights_only=False)["state_dict"] | |
| weights = { | |
| k.removeprefix("ss_model."): v | |
| for k, v in state_dict.items() | |
| if k.startswith("ss_model.") | |
| } | |
| _ss_model.load_state_dict(weights) | |
| _ss_model.eval().to(DEVICE) | |
| ss_model = _ss_model | |
| print("Model loaded successfully.") | |
| except Exception as e: | |
| model_error = str(e) | |
| print(f"Error loading model: {traceback.format_exc()}") | |
| # pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the | |
| # server, so the model is loaded there alone, before process_fn can reserve a GPU. A failure is | |
| # kept rather than raised, so that each request can report it. | |
| if __name__ != "__main__": | |
| load_model() | |
| def process_fn(audio_path: str, text_query: str): | |
| # Checked before the GPU is reserved, so a request that cannot run spends none of the caller's | |
| # ZeroGPU allowance | |
| if ss_model is None: | |
| raise gr.Error(f"Model failed to load: {model_error}") | |
| if not audio_path: | |
| raise gr.Error("Provide the mixture to separate.") | |
| 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.") | |
| if not (text_query or "").strip(): | |
| raise gr.Error("Describe the sound to separate, such as \"acoustic guitar\".") | |
| return separate(audio_path, text_query.strip()) | |
| def separate(audio_path: str, text_query: str): | |
| print(f"Separating [{audio_path}] with query [{text_query}]") | |
| # The model expects 32 kHz mono. Both are restored after inference, so the | |
| # result drops into the session it came from rather than at the model's rate | |
| info = sf.info(audio_path) | |
| original_sr = info.samplerate | |
| original_channels = info.channels | |
| mixture, _ = librosa.load(audio_path, sr=32000, mono=True) | |
| conditions = query_encoder.get_query_embed( | |
| modality="text", | |
| text=[text_query], | |
| device=DEVICE, | |
| ) | |
| input_dict = { | |
| "mixture": torch.tensor(mixture, dtype=torch.float32)[None, None, :].to(DEVICE), | |
| "condition": conditions, | |
| } | |
| # Note: using ss_model.forward() directly. | |
| # chunk_inference() is not used — it has a latent self.sampling_rate AttributeError | |
| # in the original source and is only needed for very long audio to reduce peak memory. | |
| sep_segment = ss_model(input_dict)["waveform"] | |
| sep_np = sep_segment.squeeze(0).squeeze(0).cpu().numpy() | |
| if original_sr != 32000: | |
| sep_np = librosa.resample(sep_np, orig_sr=32000, target_sr=original_sr) | |
| # Separation here is mono, so the channels are rebuilt by copying rather | |
| # than recovered, as the model card says | |
| if original_channels > 1: | |
| sep_np = np.tile(sep_np[:, None], (1, original_channels)) | |
| input_ext = Path(audio_path).suffix.lower() | |
| with tempfile.NamedTemporaryFile(suffix=input_ext if input_ext in LOSSLESS else ".wav", | |
| delete=False) as f: | |
| out_path = f.name | |
| sf.write(out_path, sep_np, original_sr) | |
| return out_path | |
| if __name__ == "__main__": | |
| with gr.Blocks() as demo: | |
| input_components = [ | |
| gr.Audio(type="filepath", label="Input Audio (Mixture)").harp_required(True).set_info(( | |
| "The mixture containing the sound to isolate." | |
| )), | |
| gr.Textbox(label="Text Query", placeholder="e.g. a dog barking", | |
| info=( | |
| "Describe the sound to isolate, such as \"acoustic guitar\" or \"a " | |
| "dog barking\". Plain words work better than studio terms." | |
| )), | |
| ] | |
| output_components = [ | |
| gr.Audio(type="filepath", label="Separated Audio").set_info( | |
| "The sound matching the description, at the sample rate and channel count of " | |
| "the input." | |
| ), | |
| ] | |
| build_endpoint( | |
| model_card=model_card, | |
| input_components=input_components, | |
| output_components=output_components, | |
| process_fn=process_fn, | |
| ) | |
| demo.queue().launch(pwa=True, show_error=True, ssr_mode=False) | |