Update app.py

app.py

@@ -1,18 +1,23 @@
+#!/usr/bin/env python
 # ruff: noqa: E402
+
 import json
 import re
 import tempfile
-from importlib.resources import files
-from groq import Groq
 import os
+
 import click
 import gradio as gr
 import numpy as np
 import soundfile as sf
 import torchaudio
+
+from importlib.resources import files
+from groq import Groq
 from cached_path import cached_path
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+# Try to import spaces; if available, set USING_SPACES to True so we can decorate functions for GPU support.
 try:
     import spaces
 
@@ -22,12 +27,15 @@ except ImportError:
 
 
 def gpu_decorator(func):
+    """
+    Decorator that wraps a function with GPU acceleration if running in a Spaces environment.
+    """
     if USING_SPACES:
         return spaces.GPU(func)
-    else:
-        return func
+    return func
 
 
+# Local package imports
 from f5_tts.model import DiT, UNetT
 from f5_tts.infer.utils_infer import (
     load_vocoder,
@@ -38,58 +46,70 @@ from f5_tts.infer.utils_infer import (
     save_spectrogram,
 )
 
-
 DEFAULT_TTS_MODEL = "F5-TTS"
-tts_model_choice = DEFAULT_TTS_MODEL
-
 DEFAULT_TTS_MODEL_CFG = [
     "hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors",
     "hf://SWivid/F5-TTS/F5TTS_Base/vocab.txt",
     json.dumps(dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)),
 ]
 
-
-# Load models
+# Load vocoder and TTS model
 vocoder = load_vocoder()
 
-def load_f5tts(ckpt_path=str(cached_path("hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors"))):
+
+def load_f5tts(
+    ckpt_path: str = str(cached_path("hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors"))
+):
+    """
+    Load the F5-TTS model from the given checkpoint path.
+    """
     F5TTS_model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
     return load_model(DiT, F5TTS_model_cfg, ckpt_path)
 
-F5TTS_ema_model = load_f5tts()
-chat_model_state = None
-chat_tokenizer_state = None
 
+F5TTS_ema_model = load_f5tts()
 
 
+# Setup the Groq client for chat completions.
 groq_token = os.getenv("Groq_TOKEN", None)
-client = Groq(
-    api_key=groq_token,
-)
+client = Groq(api_key=groq_token)
+
 
 @gpu_decorator
 def generate_response(messages):
-    """Generate response using Groq"""
+    """
+    Generate a chat response using the Groq API.
+    If messages is a string, wrap it as a user message.
+    """
+    if isinstance(messages, str):
+        messages_payload = [{"role": "user", "content": messages}]
+    else:
+        messages_payload = messages
+
     chat_completion = client.chat.completions.create(
-        messages=[
-            {
-                "role": "user",
-                "content": messages,
-            }
-        ] if isinstance(messages, str) else messages,
+        messages=messages_payload,
         model="llama-3.3-70b-versatile",
         stream=False,
     )
-    return chat_completion.choices[0].message.content # this may need to be fixed
+    # Check that we got a valid response.
+    if chat_completion.choices and hasattr(chat_completion.choices[0].message, "content"):
+        return chat_completion.choices[0].message.content
+    return ""
 
 
 @gpu_decorator
 def process_audio_input(audio_path, text, history, conv_state):
+    """
+    Process audio and/or text input from the user:
+      - If an audio file is provided, its transcript is obtained.
+      - The conversation state and history are updated.
+    """
     if not audio_path and not text.strip():
         return history, conv_state, ""
 
     if audio_path:
-        text = preprocess_ref_audio_text(audio_path, text)[1]
+        # preprocess_ref_audio_text returns a tuple (audio, transcript).
+        _, text = preprocess_ref_audio_text(audio_path, text)
 
     if not text.strip():
         return history, conv_state, ""
@@ -102,19 +122,20 @@ def process_audio_input(audio_path, text, history, conv_state):
     return history, conv_state, ""
 
 
-
 @gpu_decorator
 def infer(
     ref_audio_orig,
     ref_text,
     gen_text,
-    model,
     remove_silence,
-    cross_fade_duration=0.15,
-    nfe_step=32,
-    speed=1,
-    show_info=gr.Info,
+    cross_fade_duration: float = 0.15,
+    nfe_step: int = 32,
+    speed: float = 1,
+    show_info=print,
 ):
+    """
+    Generate speech audio using the F5-TTS system based on a reference audio/text and generated text.
+    """
     if not ref_audio_orig:
         gr.Warning("Please provide reference audio.")
         return gr.update(), gr.update(), ref_text
@@ -123,8 +144,9 @@ def infer(
         gr.Warning("Please enter text to generate.")
         return gr.update(), gr.update(), ref_text
 
+    # Preprocess the reference audio and text.
     ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=show_info)
-    ema_model = F5TTS_ema_model  # Use F5-TTS by default
+    ema_model = F5TTS_ema_model  # Use the default F5-TTS model.
 
     final_wave, final_sample_rate, combined_spectrogram = infer_process(
         ref_audio,
@@ -140,12 +162,17 @@ def infer(
     )
 
     if remove_silence:
+        # Write the generated waveform to a temporary file.
         with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
-            sf.write(f.name, final_wave, final_sample_rate)
-            remove_silence_for_generated_wav(f.name)
-            final_wave, _ = torchaudio.load(f.name)
-        final_wave = final_wave.squeeze().cpu().numpy()
-
+            temp_audio_path = f.name
+            sf.write(temp_audio_path, final_wave, final_sample_rate)
+        # Process the file to remove silence.
+        remove_silence_for_generated_wav(temp_audio_path)
+        final_wave_tensor, _ = torchaudio.load(temp_audio_path)
+        final_wave = final_wave_tensor.squeeze().cpu().numpy()
+        os.unlink(temp_audio_path)  # Clean up the temporary file.
+
+    # Save the spectrogram as a temporary PNG file.
     with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
         spectrogram_path = tmp_spectrogram.name
         save_spectrogram(combined_spectrogram, spectrogram_path)
@@ -154,29 +181,27 @@ def infer(
 
 
 with gr.Blocks() as app_chat:
-    gr.Markdown("""
+    gr.Markdown(
+        """
 # Voice Chat
 Have a conversation with an AI using your reference voice! 
 1. Upload a reference audio clip and optionally its transcript.
 2. Load the chat model.
 3. Record your message through your microphone.
 4. The AI will respond using the reference voice.
-""")
+"""
+    )
 
-    if not USING_SPACES:
-        chat_interface_container = gr.Column(visible=False)
+    # The chat interface container is visible only if running in a Spaces environment.
+    chat_interface_container = gr.Column(visible=USING_SPACES)
 
-        
     with chat_interface_container:
         with gr.Row():
             with gr.Column():
                 ref_audio_chat = gr.Audio(label="Reference Audio", type="filepath")
             with gr.Column():
                 with gr.Accordion("Advanced Settings", open=False):
-                    remove_silence_chat = gr.Checkbox(
-                        label="Remove Silences",
-                        value=True,
-                    )
+                    remove_silence_chat = gr.Checkbox(label="Remove Silences", value=True)
                     ref_text_chat = gr.Textbox(
                         label="Reference Text",
                         info="Optional: Leave blank to auto-transcribe",
@@ -184,52 +209,52 @@ Have a conversation with an AI using your reference voice!
                     )
                     system_prompt_chat = gr.Textbox(
                         label="System Prompt",
-                        value="You are not an AI assistant, you are whoever the user says you are. You must stay in character. Keep your responses concise since they will be spoken out loud.",
+                        value=(
+                            "You are not an AI assistant, you are whoever the user says you are. "
+                            "You must stay in character. Keep your responses concise since they will be spoken out loud."
+                        ),
                         lines=2,
                     )
 
         chatbot_interface = gr.Chatbot(label="Conversation")
         with gr.Row():
             with gr.Column():
-                audio_input_chat = gr.Microphone(
-                    label="Speak your message",
-                    type="filepath",
-                )
+                audio_input_chat = gr.Microphone(label="Speak your message", type="filepath")
                 audio_output_chat = gr.Audio(autoplay=True)
             with gr.Column():
-                text_input_chat = gr.Textbox(
-                    label="Type your message",
-                    lines=1,
-                )
+                text_input_chat = gr.Textbox(label="Type your message", lines=1)
                 send_btn_chat = gr.Button("Send Message")
                 clear_btn_chat = gr.Button("Clear Conversation")
 
+        # Initialize the conversation state with the system prompt.
         conversation_state = gr.State(
             value=[
                 {
                     "role": "system",
-                    "content": "You are not an AI assistant, you are whoever the user says you are. You must stay in character. Keep your responses concise since they will be spoken out loud.",
+                    "content": (
+                        "You are not an AI assistant, you are whoever the user says you are. "
+                        "You must stay in character. Keep your responses concise since they will be spoken out loud."
+                    ),
                 }
             ]
         )
-        
-        
-
 
         @gpu_decorator
         def generate_audio_response(history, ref_audio, ref_text, remove_silence):
+            """
+            Generate an audio response from the last AI message in the conversation.
+            """
             if not history or not ref_audio:
-                return None
+                return None, ref_text
 
             last_user_message, last_ai_response = history[-1]
             if not last_ai_response:
-                return None
+                return None, ref_text
 
             audio_result, _, ref_text_out = infer(
                 ref_audio,
                 ref_text,
                 last_ai_response,
-                tts_model_choice,
                 remove_silence,
                 cross_fade_duration=0.15,
                 speed=1.0,
@@ -238,11 +263,28 @@ Have a conversation with an AI using your reference voice!
             return audio_result, ref_text_out
 
         def clear_conversation():
-            return [], [{"role": "system", "content": "You are not an AI assistant, you are whoever the user says you are. You must stay in character. Keep your responses concise since they will be spoken out loud."}]
+            """
+            Clear the chat conversation and reset the conversation state.
+            """
+            initial_state = [
+                {
+                    "role": "system",
+                    "content": (
+                        "You are not an AI assistant, you are whoever the user says you are. "
+                        "You must stay in character. Keep your responses concise since they will be spoken out loud."
+                    ),
+                }
+            ]
+            return [], initial_state
 
         def update_system_prompt(new_prompt):
-            return [], [{"role": "system", "content": new_prompt}]
+            """
+            Update the system prompt and reset the conversation.
+            """
+            initial_state = [{"role": "system", "content": new_prompt}]
+            return [], initial_state
 
+        # Set up callbacks so that when recording stops, or text is submitted, the chain of processing is run.
         audio_input_chat.stop_recording(
             process_audio_input,
             inputs=[audio_input_chat, text_input_chat, chatbot_interface, conversation_state],
@@ -274,7 +316,11 @@ Have a conversation with an AI using your reference voice!
         ).then(lambda: None, None, text_input_chat)
 
         clear_btn_chat.click(clear_conversation, outputs=[chatbot_interface, conversation_state])
-        system_prompt_chat.change(update_system_prompt, inputs=system_prompt_chat, outputs=[chatbot_interface, conversation_state])
+        system_prompt_chat.change(
+            update_system_prompt,
+            inputs=system_prompt_chat,
+            outputs=[chatbot_interface, conversation_state],
+        )
 
 
 app = app_chat
@@ -285,16 +331,19 @@ app = app_chat
 @click.option("--host", "-H", default=None, help="Host to run the app on")
 @click.option("--share", "-s", default=False, is_flag=True, help="Share the app via Gradio share link")
 @click.option("--api", "-a", default=True, is_flag=True, help="Allow API access")
-@click.option("--root_path", "-r", default=None, type=str, help='Root path for the application')
+@click.option("--root_path", "-r", default=None, type=str, help="Root path for the application")
 def main(port, host, share, api, root_path):
+    """
+    Launch the Gradio app.
+    """
     app.queue(api_open=api).launch(
         server_name=host,
         server_port=port,
         share=share,
         show_api=api,
-        root_path=root_path
+        root_path=root_path,
     )
 
 
 if __name__ == "__main__":
-    main()
+    main()