Create preprocess.py

preprocess.py

@@ -0,0 +1,87 @@
+import torch
+from typing import Tuple, Dict
+from joblib import Parallel, delayed
+from resemble_enhance.enhancer.inference import denoise, enhance
+
+
+def denoise_audio(audio: torch.Tensor, sr: int, device: str) -> Tuple[torch.Tensor, int]:
+    """
+    Réduction de bruit sur un tenseur audio.
+    """
+    try:
+        return denoise(audio, sr, device)
+    except Exception as e:
+        print(f"[denoise_audio] Erreur: {e}")
+        return audio, sr
+
+
+def enhance_audio(
+    audio: torch.Tensor,
+    sr: int,
+    device: str,
+    solver: str = "midpoint",
+    nfe: int = 128,
+    lambd: float = 0.1,
+    tau: float = 0.01
+) -> Tuple[torch.Tensor, int]:
+    """
+    Amélioration audio (résolution d'ODE).
+    """
+    try:
+        return enhance(
+            audio,
+            sr,
+            device,
+            solver=solver.lower(),
+            nfe=nfe,
+            lambd=lambd,
+            tau=tau
+        )
+    except Exception as e:
+        print(f"[enhance_audio] Erreur: {e}")
+        return audio, sr
+
+
+def process_audio(
+    audio: torch.Tensor,
+    sr: int,
+    device: str = "cpu",
+    solver: str = "midpoint",
+    nfe: int = 128,
+    tau: float = 0.01,
+    denoise_before: bool = False,
+) -> Dict[str, Tuple[torch.Tensor, int]]:
+    """
+    Pipeline complet de traitement audio.
+
+    Args:
+        audio: Tensor d'entrée de forme (1, T).
+        sr: Fréquence d'échantillonnage.
+        device: 'cpu' ou 'cuda'.
+        solver: ODE solver.
+        nfe: nombre d'évaluations de la fonction.
+        tau: paramètre tau.
+        denoise_before: si True, lambd=0.9, sinon lambd=0.1.
+        parallel: si True, exécute denoise et enhance en parallèle (utile si indépendants).
+
+    Returns:
+        dict avec clés 'denoised' et 'enhanced', valeurs (tensor, sr).
+    """
+    # Choix de lambda
+    lambd = 0.9 if denoise_before else 0.1
+    # Exécution séquentielle (débruitage puis amélioration)
+    denoised = denoise_audio(audio, sr, device)
+    # Applique enhance sur le résultat du débruitage
+    enhanced = enhance_audio(
+        denoised[0],
+        denoised[1],
+        device,
+        solver,
+        nfe,
+        lambd,
+        tau
+    )
+    return {
+        "denoised": denoised,
+        "enhanced": enhanced
+    }