add postprocessing option

app.py

@@ -4,13 +4,17 @@ import torch
 
 from transformers import pipeline
 
+from resemble_enhance.enhancer.inference import denoise, enhance
+
 checkpoint_finetuned = "JackismyShephard/speecht5_tts-finetuned-nst-da"
 
+device = 0 if torch.cuda.is_available() else "cpu"
+
 pipe = pipeline(
     "text-to-speech",
     model=checkpoint_finetuned,
     use_fast=True,
-    device=0 if torch.cuda.is_available() else "cpu",
+    device=device,
 )
 
 
@@ -24,7 +28,7 @@ speaker_embeddings = {
 }
 
 
-def predict(text, speaker):
+def predict(text, speaker, post_process):
     if len(text.strip()) == 0:
         return (16000, np.zeros(0))
 
@@ -41,7 +45,10 @@ def predict(text, speaker):
     forward_params = {"speaker_embeddings": speaker_embedding}
     speech = pipe(text, forward_params=forward_params)
 
-    return (speech["sampling_rate"], speech["audio"])
+    if post_process:
+        return enhance_audio(speech["audio"], speech["sampling_rate"], device)
+    else:
+        return (speech["sampling_rate"], speech["audio"])
 
 
 def replace_danish_letters(text):
@@ -72,6 +79,17 @@ replacements = [
     ("ü", "y"),
 ]
 
+
+def enhance_audio(waveform, sr, device="cuda"):
+    tensor = torch.tensor(waveform).float()
+    denoised, new_sr = denoise(tensor, sr, device)
+    enhanced, new_sr = enhance(
+        denoised, new_sr, device, nfe=2, solver="midpoint", lambd=0.9, tau=0.95
+    )
+    enhanced_cpu = enhanced.cpu().numpy()
+    return new_sr, enhanced_cpu
+
+
 title = "Danish Speech Synthesis"
 
 description = (
@@ -84,6 +102,7 @@ examples = [
     [
         "I sin oprindelige før-kristne form blev alferne sandsynligvis opfattet som en personificering af det land og den natur, der omgav menneskene, dvs. den opdyrkede jord, gården og de naturressourcer, som hørte dertil. De var guddommelige eller delvis guddommelige væsener, der besad magiske kræfter, som de brugte både til fordel og ulempe for menneskene.",
         "F23 (Female, 23, Vestjylland)",
+        True,
     ],
 ]
 
@@ -103,6 +122,7 @@ demo = gr.Interface(
             ],
             value="F23 (Female, 23, Vestjylland)",
         ),
+        gr.Checkbox(label="Enhance audio (takes substantially longer)"),
     ],
     outputs=[
         gr.Audio(label="Generated Speech", type="numpy"),

requirements.txt

@@ -1,4 +1,5 @@
 --extra-index-url https://download.pytorch.org/whl/cu113
 torch
 transformers
-sentencepiece
+sentencepiece
+resemble-enhance