Move to torch 2.8.0, load in the worker from a Hub mirror, and keep a lossless container

Moves from torch 2.2.2, which ZeroGPU rejects, to 2.8.0, whose torchaudio still provides the module deepfilternet imports, loads the model in the worker alone from teamup-tech/DeepFilterNet-weights rather than the main branch of its GitHub repository, refuses an empty file, and returns a FLAC or AIFF input in its own container.

app.py

@@ -2,13 +2,14 @@ import spaces  # imported first, as ZeroGPU requires
 
 import os
 import tempfile
-from pathlib import Path  
+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
@@ -21,9 +22,8 @@ model_card = ModelCard(
     name="DeepFilterNet",
     description=(
         "Background noise suppression / speech enhancement using "
-        "DeepFilterNet3. Input is converted to mono 48kHz. 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 "
+        "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"
@@ -32,6 +32,10 @@ model_card = ModelCard(
     tags=[Subcategory.DENOISING, SampleRate(48000), Channels(1)],
 )
 
+# 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", ".aiff", ".aif"}
+
 # -----------------------------
 # Device & Model Initialization
 # -----------------------------
@@ -39,9 +43,18 @@ model_card = ModelCard(
 DEVICE_STR = "cuda" if torch.cuda.is_available() else "cpu"
 DEVICE = torch.device(DEVICE_STR)
 
-# Load model 1 time into global memory
-MODEL, DF_STATE, _ = init_df()
-MODEL = MODEL.to(DEVICE)
+# 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):
@@ -91,9 +104,11 @@ def process_fn(input_audio_path: str, noise_atten_db: float) -> str:
     if not input_audio_path:
         raise gr.Error("Provide the speech to denoise.")
     try:
-        sf.info(input_audio_path)
+        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)
 
@@ -120,11 +135,6 @@ def denoise(input_audio_path: str, noise_atten_db: float) -> str:
     # Slider strength
     out = apply_attenuation_db(noisy, enhanced, noise_atten_db, max_db=30.0)
 
-    # Save output WAV and return path
-    out_dir = Path(tempfile.gettempdir()) / "pyharp_dfnet_outputs"
-    out_dir.mkdir(parents=True, exist_ok=True)
-    out_path = out_dir / "denoised.wav"
-
     if original_sr != sr:
         out = librosa.resample(out, orig_sr=sr, target_sr=original_sr)
 
@@ -133,6 +143,8 @@ def denoise(input_audio_path: str, noise_atten_db: float) -> str:
     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)
 
@@ -169,4 +181,4 @@ if __name__ == "__main__":
             process_fn=process_fn,
         )
 
-    demo.queue().launch(show_error=True, pwa=True)
+    demo.queue().launch(show_error=True, pwa=True)

requirements.txt

@@ -1,6 +1,6 @@
 deepfilternet==0.5.6
-torch==2.2.2
-torchaudio==2.2.2
+torch==2.8.0
+torchaudio==2.8.0
 numpy==1.26.4
 soundfile==0.13.1
 librosa==0.11.0