Spaces:
Running
on
Zero
Running
on
Zero
Create preprocess.py
Browse files- preprocess.py +87 -0
preprocess.py
ADDED
@@ -0,0 +1,87 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
from typing import Tuple, Dict
|
3 |
+
from joblib import Parallel, delayed
|
4 |
+
from resemble_enhance.enhancer.inference import denoise, enhance
|
5 |
+
|
6 |
+
|
7 |
+
def denoise_audio(audio: torch.Tensor, sr: int, device: str) -> Tuple[torch.Tensor, int]:
|
8 |
+
"""
|
9 |
+
Réduction de bruit sur un tenseur audio.
|
10 |
+
"""
|
11 |
+
try:
|
12 |
+
return denoise(audio, sr, device)
|
13 |
+
except Exception as e:
|
14 |
+
print(f"[denoise_audio] Erreur: {e}")
|
15 |
+
return audio, sr
|
16 |
+
|
17 |
+
|
18 |
+
def enhance_audio(
|
19 |
+
audio: torch.Tensor,
|
20 |
+
sr: int,
|
21 |
+
device: str,
|
22 |
+
solver: str = "midpoint",
|
23 |
+
nfe: int = 128,
|
24 |
+
lambd: float = 0.1,
|
25 |
+
tau: float = 0.01
|
26 |
+
) -> Tuple[torch.Tensor, int]:
|
27 |
+
"""
|
28 |
+
Amélioration audio (résolution d'ODE).
|
29 |
+
"""
|
30 |
+
try:
|
31 |
+
return enhance(
|
32 |
+
audio,
|
33 |
+
sr,
|
34 |
+
device,
|
35 |
+
solver=solver.lower(),
|
36 |
+
nfe=nfe,
|
37 |
+
lambd=lambd,
|
38 |
+
tau=tau
|
39 |
+
)
|
40 |
+
except Exception as e:
|
41 |
+
print(f"[enhance_audio] Erreur: {e}")
|
42 |
+
return audio, sr
|
43 |
+
|
44 |
+
|
45 |
+
def process_audio(
|
46 |
+
audio: torch.Tensor,
|
47 |
+
sr: int,
|
48 |
+
device: str = "cpu",
|
49 |
+
solver: str = "midpoint",
|
50 |
+
nfe: int = 128,
|
51 |
+
tau: float = 0.01,
|
52 |
+
denoise_before: bool = False,
|
53 |
+
) -> Dict[str, Tuple[torch.Tensor, int]]:
|
54 |
+
"""
|
55 |
+
Pipeline complet de traitement audio.
|
56 |
+
|
57 |
+
Args:
|
58 |
+
audio: Tensor d'entrée de forme (1, T).
|
59 |
+
sr: Fréquence d'échantillonnage.
|
60 |
+
device: 'cpu' ou 'cuda'.
|
61 |
+
solver: ODE solver.
|
62 |
+
nfe: nombre d'évaluations de la fonction.
|
63 |
+
tau: paramètre tau.
|
64 |
+
denoise_before: si True, lambd=0.9, sinon lambd=0.1.
|
65 |
+
parallel: si True, exécute denoise et enhance en parallèle (utile si indépendants).
|
66 |
+
|
67 |
+
Returns:
|
68 |
+
dict avec clés 'denoised' et 'enhanced', valeurs (tensor, sr).
|
69 |
+
"""
|
70 |
+
# Choix de lambda
|
71 |
+
lambd = 0.9 if denoise_before else 0.1
|
72 |
+
# Exécution séquentielle (débruitage puis amélioration)
|
73 |
+
denoised = denoise_audio(audio, sr, device)
|
74 |
+
# Applique enhance sur le résultat du débruitage
|
75 |
+
enhanced = enhance_audio(
|
76 |
+
denoised[0],
|
77 |
+
denoised[1],
|
78 |
+
device,
|
79 |
+
solver,
|
80 |
+
nfe,
|
81 |
+
lambd,
|
82 |
+
tau
|
83 |
+
)
|
84 |
+
return {
|
85 |
+
"denoised": denoised,
|
86 |
+
"enhanced": enhanced
|
87 |
+
}
|