Update app.py

app.py

@@ -32,10 +32,12 @@ Use this space to generate long-form speech up to around ~2 minutes in length. T
 """.strip()
 
 # Create pipelines, downloading required files as necessary
+speech_enhancement = Task.get("speech-enhancement", model="deep-filter-net-v3", available_only=False)
+speech_enhancement.download_required_files(text_callback=print)
 hybrid_task = Task.get("speech-synthesis", model="zonos-hybrid", available_only=False)
 hybrid_task.download_required_files(text_callback=print)
 hybrid_pipe = hybrid_task()
-hybrid_pipe.load()
+hybrid_pipe.load(allow_optional=True)
 
 transformer_task = Task.get(
     "speech-synthesis", model="zonos-transformer", available_only=False
@@ -45,7 +47,7 @@ transformer_pipe = transformer_task()
 
 if is_hf_spaces:
     # Must load all models on GPU when using ZERO
-    transformer_pipe.load()
+    transformer_pipe.load(allow_optional=True)
 
 # Global state
 pipelines = {
@@ -126,6 +128,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 +150,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 +190,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 +206,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 +228,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 +285,9 @@ if __name__ == "__main__":
                     label="Optional Prefix Audio (continue from this audio)",
                     type="filepath",
                 )
+                prefix_equalize_checkbox = gr.Checkbox(label="Equalize Prefix Audio", value=True)
+                prefix_enhance_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 +302,31 @@ 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="Speaker Noised",
+                    value=False,
+                    interactive=False,
+                    info="'Speaker Noised' is a conditioning value that the model understands, not a processing step. Check this box if your input audio is noisy."
+                )
+                speaker_equalize_checkbox = gr.Checkbox(label="Equalize Speaker Audio", value=True)
+                speaker_enhance_checkbox = gr.Checkbox(label="Enhance Speaker Audio with DeepFilterNet", value=True)
+
+                def on_enhanced_change(use_enhance: bool) -> Dict[str, Any]:
+                    update_dict = {"interactive": 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 +449,14 @@ if __name__ == "__main__":
                 speaking_rate_uncond,
                 emotion_uncond,
                 speaker_uncond,
+                speaker_pitch_shift,
+                speaker_equalize_checkbox,
+                speaker_enhance_checkbox,
+                prefix_equalize_checkbox,
+                prefix_enhance_checkbox,
+                enhanced_checkbox,
             ],
             outputs=[output_audio, seed_number],
         )
-    
+
         demo.launch()