Update app.py

app.py

@@ -126,6 +126,12 @@ def generate_audio(
     skip_speaking_rate: bool,
     skip_emotion: bool,
     skip_speaker: bool,
+    speaker_pitch_shift: float,
+    speaker_equalize: bool,
+    speaker_enhance: bool,
+    prefix_equalize: bool,
+    prefix_enhance: bool,
+    enhance: bool,
     progress=gr.Progress(),
 ) -> Tuple[Tuple[int, np.ndarray[Any, Any]], int]:
     """
@@ -142,9 +148,15 @@ def generate_audio(
     try:
         wav_out = selected_pipeline(
             text=text,
+            enhance=enhance,
             language=language,
             reference_audio=speaker_audio,
+            reference_audio_pitch_shift=speaker_pitch_shift,
+            equalize_reference_audio=speaker_equalize,
+            enhance_reference_audio=speaker_enhance,
             prefix_audio=prefix_audio,
+            equalize_prefix_audio=prefix_equalize,
+            enhance_prefix_audio=prefix_enhance,
             seed=seed,
             max_chunk_length=max_chunk_length,
             cross_fade_duration=cross_fade_duration,
@@ -176,7 +188,13 @@ def generate_audio(
             output_format="float",
         )
 
-        return (44100, wav_out.squeeze().numpy()), seed
+        return (
+            (
+                48000 if enhance else 44100,
+                wav_out.squeeze().numpy()
+            ),
+            seed
+        )
     finally:
         selected_pipeline.off_progress()
 
@@ -186,6 +204,7 @@ if __name__ == "__main__":
         with gr.Row():
             with gr.Column(scale=3):
                 gr.Markdown(header_markdown)
+                
             gr.Image(
                 value="https://raw.githubusercontent.com/Zyphra/Zonos/refs/heads/main/assets/ZonosHeader.png",
                 container=False,
@@ -207,6 +226,10 @@ if __name__ == "__main__":
                 value="en-us",
                 label="Language",
             )
+            enhanced_checkbox = gr.Checkbox(
+                value=True,
+                label="Enhance Output with DeepFilterNet"
+            )
     
         with gr.Row():
             if not is_hf_spaces:
@@ -260,6 +283,8 @@ if __name__ == "__main__":
                     label="Optional Prefix Audio (continue from this audio)",
                     type="filepath",
                 )
+                prefix_equalized_checkbox = gr.Checkbox(label="Equalize Prefix Audio", value=True)
+                prefix_enhanced_checkbox = gr.Checkbox(label="Enhance Prefix Audio with DeepFilterNet", value=True)
             with gr.Column(scale=3):
                 cfg_scale_slider = gr.Slider(1.0, 5.0, 2.0, 0.1, label="CFG Scale")
                 min_p_slider = gr.Slider(0.0, 1.0, 0.15, 0.01, label="Min P")
@@ -274,8 +299,29 @@ if __name__ == "__main__":
         with gr.Row(variant="panel", equal_height=True) as speaker_row:
             with gr.Column():
                 speaker_uncond = gr.Checkbox(label="Skip Speaker")
-                speaker_noised_checkbox = gr.Checkbox(label="Denoise Speaker", value=False)
-    
+                speaker_noised_checkbox = gr.Checkbox(
+                    label="Denoise Speaker",
+                    value=False,
+                    info="Note; this is not a pre-processing step, it is a conditioning value that the model understands. Check this box if your input audio is noisy."
+                )
+                speaker_equalized_checkbox = gr.Checkbox(label="Equalize Speaker Audio", value=True)
+                speaker_enhanced_checkbox = gr.Checkbox(label="Enhance Speaker Audio with DeepFilterNet", value=True)
+
+                def on_enhanced_change(use_enhance: bool) -> Dict[str, Any]:
+                    update_dict = {"enabled": not use_enhance}
+                    if use_enhance:
+                        update_dict["value"] = False
+                    return gr.update(**update_dict)
+
+                speaker_enhance_checkbox.change(
+                    fn=on_enhanced_change,
+                    inputs=[speaker_enhance_checkbox],
+                    outputs=[speaker_noised_checkbox]
+                )
+                speaker_pitch_shift = gr.Slider(
+                    -1200, 1200, -44.99, 0.01, label="Speaker Pitch Shift (Cents)",
+                    info="A pitch shift to apply to speaker audio before extracting embeddings. A slight down-shift of ~45 cents tends to produce a more accurate voice cloning."
+                )
             speaker_audio = gr.Audio(
                 label="Optional Speaker Audio (for cloning)",
                 type="filepath",
@@ -398,8 +444,14 @@ if __name__ == "__main__":
                 speaking_rate_uncond,
                 emotion_uncond,
                 speaker_uncond,
+                speaker_pitch_shift,
+                speaker_equalized_checkbox,
+                speaker_enhanced_checkbox,
+                prefix_equalized_checkbox,
+                prefix_enhanced_checkbox,
+                enhanced_checkbox,
             ],
             outputs=[output_audio, seed_number],
         )
-    
+
         demo.launch()