Spaces:

teamup-tech
/

HTS-AT

Sleeping

HTS-AT / app.py

Drop the comment on ssr_mode

6d1773b verified 1 day ago

6.84 kB

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)

	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)