import spaces # imported first, as ZeroGPU requires import sys sys.stdout.reconfigure(line_buffering=True) from types import SimpleNamespace import gradio as gr import librosa import numpy as np import pandas as pd import soundfile as sf import torch from huggingface_hub import hf_hub_download from pyharp import ModelCard, AudioLabel, LabelList, build_endpoint from pyharp.tags import Channels, SampleRate, Subcategory from model.htsat import HTSAT_Swin_Transformer model_card = ModelCard( name="HTS-AT", description=( "Tags the sounds present in an audio clip across 527 AudioSet classes " "(speech, music, instruments, animals, environmental sounds, etc). " "Clips longer than 10 seconds are not currently supported, and the input is mixed to " "mono and resampled to 32 kHz internally. Repository: " "https://github.com/RetroCirce/HTS-Audio-Transformer Paper: " "https://arxiv.org/abs/2202.00874" ), author="Ke Chen, Xingjian Du, Bilei Zhu, Zejun Ma, Taylor Berg-Kirkpatrick, Shlomo Dubnov", tags=[Subcategory.GENERAL_AUDIO_ANALYSIS, SampleRate(32000), Channels(1), 'audioset', 'domain:general-audio'], ) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # AudioSet configuration values from HTS-AT-repo/config.py CONFIG = SimpleNamespace( sample_rate=32000, window_size=1024, hop_size=320, mel_bins=64, fmin=50, fmax=14000, enable_tscam=True, htsat_attn_heatmap=False, loss_type="clip_bce", enable_repeat_mode=False, ) # pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the # server, so the model is built and placed there alone, on its main thread, where ZeroGPU's # emulation records the placement. The checkpoint lives in a model repository rather than in this # Space, which keeps the Space small enough to clone quickly, pinned to a revision so the Space # cannot change behavior under us. if __name__ != "__main__": model = HTSAT_Swin_Transformer(config=CONFIG) checkpoint = torch.load(hf_hub_download( repo_id="teamup-tech/HTS-AT-weights", filename="HTSAT_AudioSet_Saved_1.ckpt", revision="2ecf7c32a8b6666c704da84b3677a1883cae69ff", ), map_location="cpu") state_dict = {key.replace("sed_model.", ""): value for key, value in checkpoint["state_dict"].items()} model.load_state_dict(state_dict) model.eval().to(DEVICE) # The model always produces this many framewise steps, regardless of clip length, # by stretching shorter clips onto this fixed grid (so frame i maps linearly onto # real time as i / MODEL_FRAMES * clip_duration). Longer clips instead use a # sliding-window path that averages overlapping crops together, which no longer # maps onto real time at all - framewise localization is only meaningful below # this length until that path is implemented separately. MODEL_FRAMES = model.spec_size * model.freq_ratio labels_df = pd.read_csv("class_label_indice.csv") idx_to_label = dict(zip(labels_df["index"], labels_df["display_name"])) def find_active_segments(frame_scores: np.ndarray, threshold: float) -> list[tuple[int, int]]: segments = [] start = None for i, score in enumerate(frame_scores): if score > threshold and start is None: start = i elif score <= threshold and start is not None: segments.append((start, i)) start = None if start is not None: segments.append((start, len(frame_scores))) return segments def process_fn(input_audio_path: str, detection_threshold: float): # Checked before the GPU is reserved, so a request that cannot run spends none of the caller's # ZeroGPU allowance if not input_audio_path: raise gr.Error("Provide the clip to tag.") try: duration = sf.info(input_audio_path).duration except Exception as exc: raise gr.Error(f"Could not read the audio file: {exc}") from exc if duration == 0: raise gr.Error("The audio file is empty.") max_seconds = MODEL_FRAMES * CONFIG.hop_size / CONFIG.sample_rate if duration > max_seconds: raise gr.Error( f"Clip is too long ({duration:.1f}s). This model currently supports " f"clips up to about {max_seconds:.1f}s." ) return tag(input_audio_path, detection_threshold) @spaces.GPU @torch.inference_mode() def tag(input_audio_path: str, detection_threshold: float): waveform, _ = librosa.load(input_audio_path, sr=CONFIG.sample_rate) duration = len(waveform) / CONFIG.sample_rate audio_tensor = torch.from_numpy(waveform).float().to(DEVICE)[None, :] mel_frames = model.logmel_extractor(model.spectrogram_extractor(audio_tensor)).shape[2] if mel_frames > MODEL_FRAMES: raise gr.Error( f"Clip is too long ({duration:.1f}s). This model currently supports " f"clips up to about {MODEL_FRAMES * CONFIG.hop_size / CONFIG.sample_rate:.1f}s." ) framewise_output = model(audio_tensor, None, False)["framewise_output"][0].cpu().numpy() label_list = LabelList() for class_idx in range(framewise_output.shape[1]): frame_scores = framewise_output[:, class_idx] for start, end in find_active_segments(frame_scores, detection_threshold): peak_confidence = frame_scores[start:end].max() label_list.append(AudioLabel( t=start / MODEL_FRAMES * duration, duration=(end - start) / MODEL_FRAMES * duration, label=idx_to_label[class_idx], description=f"confidence {peak_confidence:.0%}", amplitude=float(peak_confidence), )) return input_audio_path, label_list if __name__ == "__main__": with gr.Blocks() as demo: input_components = [ gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info(( "A clip of up to about ten seconds to tag. A longer one is refused, " "since the model takes one window of that length." )), gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5, label="Detection Threshold", info="Minimum confidence for a sound to be tagged"), ] output_components = [ gr.Audio(type="filepath", label="Audio with Sound Tags").set_info("Input audio, unchanged."), gr.JSON(label="Detected Sounds").set_info("Top predicted sound tags with confidence scores."), ] build_endpoint( model_card=model_card, input_components=input_components, output_components=output_components, process_fn=process_fn, ) demo.queue().launch(show_error=True, pwa=True, ssr_mode=False)