Spaces:
Running
Running
File size: 4,191 Bytes
69c07f8 550f5fe 69c07f8 550f5fe 69c07f8 550f5fe 69c07f8 550f5fe 69c07f8 550f5fe 69c07f8 550f5fe 69c07f8 550f5fe 69c07f8 550f5fe 69c07f8 550f5fe 69c07f8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | import sys
sys.stdout.reconfigure(line_buffering=True) # flush print() calls immediately so HF Space logs are live
import os
import tempfile
import threading
import traceback
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
from models.resunet import ResUNet30
from models.clap_encoder import CLAP_Encoder
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# Model loading state — populated by the background thread.
ss_model = None
query_encoder = None
model_loading = True
model_error = None
def load_model():
global ss_model, query_encoder, model_loading, 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",
)
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",
)
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()}")
finally:
model_loading = False
threading.Thread(target=load_model, daemon=True).start()
model_card = ModelCard(
name="AudioSep",
description="Separate any sound from a mixture using a natural language text description.",
author="Xubo Liu, Qiuqiang Kong, Yan Zhao, Haohe Liu, Yi Yuan, Yuzhuo Liu, Rui Xia, Yuxuan Wang, Mark D. Plumbley, Wenwu Wang",
tags=["audio separation", "text-queried", "source separation"],
)
@torch.inference_mode()
def process_fn(audio_path: str, text_query: str):
if model_loading:
raise gr.Error("Model is still loading, please wait a moment and try again.")
if ss_model is None:
raise gr.Error(f"Model failed to load: {model_error}")
print(f"Separating [{audio_path}] with query [{text_query}]")
mixture, _ = librosa.load(audio_path, sr=32000, mono=True) # model expects 32k audio
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()
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
out_path = f.name
sf.write(out_path, sep_np, 32000)
return out_path
with gr.Blocks() as demo:
input_components = [
gr.Audio(type="filepath", label="Input Audio (Mixture)").harp_required(True),
gr.Textbox(label="Text Query", placeholder="e.g. a dog barking"),
]
output_components = [
gr.Audio(type="filepath", label="Separated Audio").set_info(
"Separated audio at 32 kHz matching the text description."
),
]
build_endpoint(
model_card=model_card,
input_components=input_components,
output_components=output_components,
process_fn=process_fn,
)
demo.queue().launch()
|