Spaces:

teamup-tech
/

DeepFilterNet

Running

Drop the comment on ssr_mode

aeec8e2 verified 3 days ago

6.59 kB

import spaces  # imported first, as ZeroGPU requires

import os
import tempfile
from pathlib import Path

import gradio as gr
import numpy as np
import soundfile as sf
import librosa
import torch
from huggingface_hub import snapshot_download

from pyharp import ModelCard, build_endpoint, SampleRate, Channels
from pyharp.tags import Subcategory
from df import enhance, init_df

# -----------------------------
# Model metadata for HARP
# -----------------------------
model_card = ModelCard(
    name="DeepFilterNet",
    description=(
        "Background noise suppression / speech enhancement using "
        "DeepFilterNet3. Slider controls strength through wet/dry blend. Denoising runs at "
        "48 kHz in mono. The result is resampled back to the input's rate and copied across its "
        "channels, so the output matches the format it came in with but carries "
        "no stereo image. Repository: https://github.com/Rikorose/DeepFilterNet "
        "Paper: https://arxiv.org/abs/2305.08227"
    ),
    author="Hendrik Schröter, Alberto N. Escalante-B., Tobias Rosenkranz, Andreas Maier",
    tags=[Subcategory.DENOISING, SampleRate(48000), Channels(1), 'domain:speech'],
)

# 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"}

# -----------------------------
# Device & Model Initialization
# -----------------------------
# Auto-detect GPU for Hugging Face deployment
DEVICE_STR = "cuda" if torch.cuda.is_available() else "cpu"
DEVICE = torch.device(DEVICE_STR)

# 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. deepfilternet would download it from the main branch
# of its GitHub repository, so it is read from a pinned mirror on the Hub instead, into a directory
# of its own, since loading writes beside the configuration.
if __name__ != "__main__":
    _weights = snapshot_download(
        "teamup-tech/DeepFilterNet-weights", revision="b2ec0e0c85db91c696d699a5cfc43ae7a3e3136a",
        local_dir=os.path.join(tempfile.gettempdir(), "deepfilternet-weights"),
    )
    MODEL, DF_STATE, _ = init_df(model_base_dir=os.path.join(_weights, "DeepFilterNet3"),
                                 log_file=None)
    MODEL = MODEL.to(DEVICE)


def load_audio_mono_48k(path: str):
    # Load audio from filepath, convert to mono float32, resample to 48kHz.
    try:
        audio, sr = sf.read(path, always_2d=False)
    except Exception:
        audio, sr = librosa.load(path, sr=None, mono=False)

    audio = np.asarray(audio)

    # stereo to mono
    if audio.ndim == 2:
        audio = audio.mean(axis=1)

    audio = audio.astype(np.float32)

    # resample to 48k
    if sr != 48000:
        audio = librosa.resample(audio, orig_sr=sr, target_sr=48000)
        sr = 48000

    return audio, sr


def apply_attenuation_db(noisy: np.ndarray, enhanced: np.ndarray, noise_atten_db: float, max_db: float = 30.0):
    """
    Map Noise Attenuation (dB) slider to a stable strength in [0,1]
    and crossfade between original and enhanced audio.
    """
    # convert dB slider into [0,1]
    s = float(noise_atten_db) / float(max_db)
    s = max(0.0, min(1.0, s))

    # make same length
    n = min(len(noisy), len(enhanced))
    noisy = noisy[:n]
    enhanced = enhanced[:n]

    out = (1.0 - s) * noisy + s * enhanced
    return out.astype(np.float32)


def process_fn(input_audio_path: str, noise_atten_db: float) -> str:
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if not input_audio_path:
        raise gr.Error("Provide the speech to denoise.")
    try:
        info = sf.info(input_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.")

    return denoise(input_audio_path, noise_atten_db)


@spaces.GPU(duration=120)
@torch.inference_mode()
def denoise(input_audio_path: str, noise_atten_db: float) -> str:
    # The session's format, read before the downmix and resample that
    # DeepFilterNet needs, and restored before saving
    info = sf.info(input_audio_path)
    original_sr = info.samplerate
    original_channels = info.channels

    # Load and normalize
    noisy, sr = load_audio_mono_48k(input_audio_path)

    # numpy to torch, add channel dim [T] to [1, T]
    noisy_t = torch.from_numpy(noisy).float().unsqueeze(0).to(DEVICE)

    # Denoise, then remove added channel and back to numpy
    enhanced_t = enhance(MODEL, DF_STATE, noisy_t)
    enhanced = enhanced_t.squeeze(0).detach().cpu().numpy()

    # Slider strength
    out = apply_attenuation_db(noisy, enhanced, noise_atten_db, max_db=30.0)

    if original_sr != sr:
        out = librosa.resample(out, orig_sr=sr, target_sr=original_sr)

    # Denoising is mono here, so the channels are copies rather than a
    # recovered image, which the model card states
    if original_channels > 1:
        out = np.tile(out[:, None], (1, original_channels))

    input_ext = Path(input_audio_path).suffix.lower()
    out_path = Path(tempfile.mkdtemp()) / f"denoised{input_ext if input_ext in LOSSLESS else '.wav'}"
    sf.write(str(out_path), out, original_sr)
    return str(out_path)


# -----------------------------
# Gradio endpoint
# -----------------------------
if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_components = [
            gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info((
                "A noisy speech recording. Converted to mono 48 kHz internally, so "
                "a stereo image is not preserved."
            )),
            gr.Slider(
                minimum=0,
                maximum=30,
                step=1,
                value=12,
                label="Noise Attenuation (dB)",
                info="0 = no change, 30 = strongest. Implemented as wet/dry strength.",
            ),
        ]

        output_components = [
            gr.Audio(type="filepath", label="Denoised Audio")
            .set_info("Denoised audio output."),
        ]

        app = 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 os
	import tempfile
	from pathlib import Path

	import gradio as gr
	import numpy as np
	import soundfile as sf
	import librosa
	import torch
	from huggingface_hub import snapshot_download

	from pyharp import ModelCard, build_endpoint, SampleRate, Channels
	from pyharp.tags import Subcategory
	from df import enhance, init_df

	# -----------------------------
	# Model metadata for HARP
	# -----------------------------
	model_card = ModelCard(
	name="DeepFilterNet",
	description=(
	"Background noise suppression / speech enhancement using "
	"DeepFilterNet3. Slider controls strength through wet/dry blend. Denoising runs at "
	"48 kHz in mono. The result is resampled back to the input's rate and copied across its "
	"channels, so the output matches the format it came in with but carries "
	"no stereo image. Repository: https://github.com/Rikorose/DeepFilterNet "
	"Paper: https://arxiv.org/abs/2305.08227"
	),
	author="Hendrik Schröter, Alberto N. Escalante-B., Tobias Rosenkranz, Andreas Maier",
	tags=[Subcategory.DENOISING, SampleRate(48000), Channels(1), 'domain:speech'],
	)

	# 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"}

	# -----------------------------
	# Device & Model Initialization
	# -----------------------------
	# Auto-detect GPU for Hugging Face deployment
	DEVICE_STR = "cuda" if torch.cuda.is_available() else "cpu"
	DEVICE = torch.device(DEVICE_STR)

	# 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. deepfilternet would download it from the main branch
	# of its GitHub repository, so it is read from a pinned mirror on the Hub instead, into a directory
	# of its own, since loading writes beside the configuration.
	if __name__ != "__main__":
	_weights = snapshot_download(
	"teamup-tech/DeepFilterNet-weights", revision="b2ec0e0c85db91c696d699a5cfc43ae7a3e3136a",
	local_dir=os.path.join(tempfile.gettempdir(), "deepfilternet-weights"),
	)
	MODEL, DF_STATE, _ = init_df(model_base_dir=os.path.join(_weights, "DeepFilterNet3"),
	log_file=None)
	MODEL = MODEL.to(DEVICE)


	def load_audio_mono_48k(path: str):
	# Load audio from filepath, convert to mono float32, resample to 48kHz.
	try:
	audio, sr = sf.read(path, always_2d=False)
	except Exception:
	audio, sr = librosa.load(path, sr=None, mono=False)

	audio = np.asarray(audio)

	# stereo to mono
	if audio.ndim == 2:
	audio = audio.mean(axis=1)

	audio = audio.astype(np.float32)

	# resample to 48k
	if sr != 48000:
	audio = librosa.resample(audio, orig_sr=sr, target_sr=48000)
	sr = 48000

	return audio, sr


	def apply_attenuation_db(noisy: np.ndarray, enhanced: np.ndarray, noise_atten_db: float, max_db: float = 30.0):
	"""
	Map Noise Attenuation (dB) slider to a stable strength in [0,1]
	and crossfade between original and enhanced audio.
	"""
	# convert dB slider into [0,1]
	s = float(noise_atten_db) / float(max_db)
	s = max(0.0, min(1.0, s))

	# make same length
	n = min(len(noisy), len(enhanced))
	noisy = noisy[:n]
	enhanced = enhanced[:n]

	out = (1.0 - s) * noisy + s * enhanced
	return out.astype(np.float32)


	def process_fn(input_audio_path: str, noise_atten_db: float) -> str:
	# Checked before the GPU is reserved, so a request the model cannot take spends none of the
	# caller's ZeroGPU allowance
	if not input_audio_path:
	raise gr.Error("Provide the speech to denoise.")
	try:
	info = sf.info(input_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.")

	return denoise(input_audio_path, noise_atten_db)


	@spaces.GPU(duration=120)
	@torch.inference_mode()
	def denoise(input_audio_path: str, noise_atten_db: float) -> str:
	# The session's format, read before the downmix and resample that
	# DeepFilterNet needs, and restored before saving
	info = sf.info(input_audio_path)
	original_sr = info.samplerate
	original_channels = info.channels

	# Load and normalize
	noisy, sr = load_audio_mono_48k(input_audio_path)

	# numpy to torch, add channel dim [T] to [1, T]
	noisy_t = torch.from_numpy(noisy).float().unsqueeze(0).to(DEVICE)

	# Denoise, then remove added channel and back to numpy
	enhanced_t = enhance(MODEL, DF_STATE, noisy_t)
	enhanced = enhanced_t.squeeze(0).detach().cpu().numpy()

	# Slider strength
	out = apply_attenuation_db(noisy, enhanced, noise_atten_db, max_db=30.0)

	if original_sr != sr:
	out = librosa.resample(out, orig_sr=sr, target_sr=original_sr)

	# Denoising is mono here, so the channels are copies rather than a
	# recovered image, which the model card states
	if original_channels > 1:
	out = np.tile(out[:, None], (1, original_channels))

	input_ext = Path(input_audio_path).suffix.lower()
	out_path = Path(tempfile.mkdtemp()) / f"denoised{input_ext if input_ext in LOSSLESS else '.wav'}"
	sf.write(str(out_path), out, original_sr)
	return str(out_path)


	# -----------------------------
	# Gradio endpoint
	# -----------------------------
	if __name__ == "__main__":
	with gr.Blocks() as demo:
	input_components = [
	gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info((
	"A noisy speech recording. Converted to mono 48 kHz internally, so "
	"a stereo image is not preserved."
	)),
	gr.Slider(
	minimum=0,
	maximum=30,
	step=1,
	value=12,
	label="Noise Attenuation (dB)",
	info="0 = no change, 30 = strongest. Implemented as wet/dry strength.",
	),
	]

	output_components = [
	gr.Audio(type="filepath", label="Denoised Audio")
	.set_info("Denoised audio output."),
	]

	app = 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)