File size: 6,585 Bytes
17c5062
 
bc9667c
 
4c67860
bc9667c
 
 
 
 
 
4c67860
bc9667c
84b2c2b
61911d1
bc9667c
 
 
 
 
 
234c9a9
bc9667c
f64bb96
4c67860
9ac3fae
f64bb96
 
 
bc9667c
7c5a78e
268cf1b
bc9667c
 
4c67860
 
9ac3fae
4c67860
bc9667c
 
 
 
 
 
 
4c67860
 
 
 
 
 
 
 
 
 
 
 
bc9667c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
17c5062
 
bc9667c
17c5062
 
4c67860
17c5062
 
4c67860
 
bc9667c
17c5062
 
 
 
 
 
19cd1f0
 
 
 
 
 
bc9667c
 
 
 
 
 
 
 
 
 
 
 
 
19cd1f0
 
 
 
 
 
 
 
4c67860
 
19cd1f0
bc9667c
 
 
 
 
 
61388cf
 
 
608ec8d
 
 
 
61388cf
 
 
 
 
 
 
 
 
 
 
e47cfcc
61388cf
 
 
 
 
 
 
 
 
 
68afc31
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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
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)