Update app.py

app.py

@@ -14,13 +14,6 @@ from dotenv import load_dotenv
 # Load environment variables
 load_dotenv()
 
-# Global variables for preloaded resources
-en_base_speaker_tts = None
-zh_base_speaker_tts = None
-tone_color_converter = None
-target_se = None
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
 # Function to download and extract checkpoints
 def download_and_extract_checkpoints():
     zip_url = "https://huggingface.co/camenduru/OpenVoice/resolve/main/checkpoints_1226.zip"
@@ -39,52 +32,56 @@ def download_and_extract_checkpoints():
         os.remove(zip_path)
         print("Checkpoints are ready.")
 
-# Initialize models and resources
-def initialize_resources():
-    global en_base_speaker_tts, zh_base_speaker_tts, tone_color_converter, target_se
-    print("Initializing resources...")
-
-    # Download and extract checkpoints
-    download_and_extract_checkpoints()
-
-    # Define paths to checkpoints
-    en_ckpt_base = 'checkpoints/base_speakers/EN'
-    zh_ckpt_base = 'checkpoints/base_speakers/ZH'
-    ckpt_converter = 'checkpoints/converter'
-
-    # Load TTS models
-    en_base_speaker_tts = BaseSpeakerTTS(f'{en_ckpt_base}/config.json', device=device)
-    en_base_speaker_tts.load_ckpt(f'{en_ckpt_base}/checkpoint.pth')
-    zh_base_speaker_tts = BaseSpeakerTTS(f'{zh_ckpt_base}/config.json', device=device)
-    zh_base_speaker_tts.load_ckpt(f'{zh_ckpt_base}/checkpoint.pth')
+# Call the function to ensure checkpoints are available
+download_and_extract_checkpoints()
 
-    # Load tone color converter
-    tone_color_converter = ToneColorConverter(f'{ckpt_converter}/config.json', device=device)
-    tone_color_converter.load_ckpt(f'{ckpt_converter}/checkpoint.pth')
+# Initialize OpenAI API key
+openai.api_key = os.getenv("OPENAI_API_KEY")
+if not openai.api_key:
+    raise ValueError("Please set the OPENAI_API_KEY environment variable.")
 
-    # Load speaker embeddings
-    en_source_default_se = torch.load(f'{en_ckpt_base}/en_default_se.pth').to(device)
-    zh_source_se = torch.load(f'{zh_ckpt_base}/zh_default_se.pth').to(device)
+parser = argparse.ArgumentParser()
+parser.add_argument("--share", action='store_true', default=False, help="make link public")
+args = parser.parse_args()
 
-    # Extract speaker embedding from the default Mickey Mouse audio
-    default_speaker_audio = "resources/output.wav"
-    try:
-        target_se, _ = se_extractor.get_se(
-            default_speaker_audio,
-            tone_color_converter,
-            target_dir='processed',
-            vad=True
-        )
-        print("Speaker embedding extracted successfully.")
-    except Exception as e:
-        raise RuntimeError(f"Failed to extract speaker embedding from {default_speaker_audio}: {str(e)}")
-
-initialize_resources()
+# Define paths to checkpoints
+en_ckpt_base = 'checkpoints/base_speakers/EN'
+zh_ckpt_base = 'checkpoints/base_speakers/ZH'
+ckpt_converter = 'checkpoints/converter'
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+output_dir = 'outputs'
+os.makedirs(output_dir, exist_ok=True)
+
+# Load TTS models
+en_base_speaker_tts = BaseSpeakerTTS(f'{en_ckpt_base}/config.json', device=device)
+en_base_speaker_tts.load_ckpt(f'{en_ckpt_base}/checkpoint.pth')
+zh_base_speaker_tts = BaseSpeakerTTS(f'{zh_ckpt_base}/config.json', device=device)
+zh_base_speaker_tts.load_ckpt(f'{zh_ckpt_base}/checkpoint.pth')
+
+tone_color_converter = ToneColorConverter(f'{ckpt_converter}/config.json', device=device)
+tone_color_converter.load_ckpt(f'{ckpt_converter}/checkpoint.pth')
+
+# Load speaker embeddings
+en_source_default_se = torch.load(f'{en_ckpt_base}/en_default_se.pth').to(device)
+en_source_style_se = torch.load(f'{en_ckpt_base}/en_style_se.pth').to(device)
+zh_source_se = torch.load(f'{zh_ckpt_base}/zh_default_se.pth').to(device)
+
+# Extract speaker embedding from the default Mickey Mouse audio
+default_speaker_audio = "resources/output.wav"
+try:
+    target_se, _ = se_extractor.get_se(
+        default_speaker_audio,
+        tone_color_converter,
+        target_dir='processed',
+        vad=True
+    )
+    print("Speaker embedding extracted successfully.")
+except Exception as e:
+    raise RuntimeError(f"Failed to extract speaker embedding from {default_speaker_audio}: {str(e)}")
 
 # Supported languages
 supported_languages = ['zh', 'en']
 
-# Predict function
 def predict(audio_file_pth, agree):
     text_hint = ''
     synthesized_audio_path = None
@@ -98,7 +95,7 @@ def predict(audio_file_pth, agree):
     if audio_file_pth is not None:
         speaker_wav = audio_file_pth
     else:
-        text_hint += "[ERROR] Please provide an audio file.\n"
+        text_hint += "[ERROR] Please record your voice using the Microphone.\n"
         return (text_hint, None)
 
     # Transcribe audio to text using OpenAI Whisper
@@ -124,66 +121,105 @@ def predict(audio_file_pth, agree):
     print(f"Detected language: {language_predicted}")
 
     if language_predicted not in supported_languages:
-        text_hint += f"[ERROR] Unsupported language: {language_predicted}\n"
+        text_hint += f"[ERROR] The detected language '{language_predicted}' is not supported. Supported languages are: {supported_languages}\n"
         return (text_hint, None)
 
-    # Select TTS model
-    tts_model = zh_base_speaker_tts if language_predicted == "zh" else en_base_speaker_tts
-    language = 'Chinese' if language_predicted == "zh" else 'English'
+    # Select TTS model based on language
+    if language_predicted == "zh":
+        tts_model = zh_base_speaker_tts
+        language = 'Chinese'
+        speaker_style = 'default'
+    else:
+        tts_model = en_base_speaker_tts
+        language = 'English'
+        speaker_style = 'default'
 
+    # Generate response using OpenAI GPT-4
     # Generate response using OpenAI GPT-4
     try:
         response = openai.chat.completions.create(
             model="gpt-4o-mini",
             messages=[
-                {"role": "system", "content": "You are Mickey Mouse, a cheerful character who responds to children's queries."},
+                {"role": "system", "content": "You are Mickey Mouse, a friendly and cheerful character who responds to children's queries in a simple and engaging manner. Please keep your response up to 200 characters."},
                 {"role": "user", "content": input_text}
-            ]
+            ],
+            max_tokens=200,
+            n=1,
+            stop=None,
+            temperature=0.7,
         )
-        reply_text = response['choices'][0]['message']['content'].strip()
+        # Correctly access the response content
+        reply_text = response.choices[0].message.content.strip()
         print(f"GPT-4 Reply: {reply_text}")
     except Exception as e:
-        text_hint += f"[ERROR] GPT-4 response failed: {str(e)}\n"
+        text_hint += f"[ERROR] Failed to get response from OpenAI GPT-4: {str(e)}\n"
         return (text_hint, None)
 
     # Synthesize reply text to audio
     try:
         src_path = os.path.join(output_dir, 'tmp_reply.wav')
-        tts_model.tts(reply_text, src_path, speaker='default', language=language)
+
+        tts_model.tts(reply_text, src_path, speaker=speaker_style, language=language)
+        print(f"Audio synthesized and saved to {src_path}")
 
         save_path = os.path.join(output_dir, 'output_reply.wav')
+
         tone_color_converter.convert(
-            audio_src_path=src_path,
-            src_se=target_se,
+            audio_src_path=src_path, 
+            src_se=en_source_default_se if language == 'English' else zh_source_se,
             tgt_se=target_se,
-            output_path=save_path
+            output_path=save_path,
+            message="@MickeyMouse"
         )
+        print(f"Tone color conversion completed and saved to {save_path}")
 
         text_hint += "Response generated successfully.\n"
         synthesized_audio_path = save_path
+
     except Exception as e:
-        text_hint += f"[ERROR] Synthesis failed: {str(e)}\n"
+        text_hint += f"[ERROR] Failed to synthesize audio: {str(e)}\n"
         traceback.print_exc()
         return (text_hint, None)
 
     return (text_hint, synthesized_audio_path)
 
-# Gradio UI
 with gr.Blocks(analytics_enabled=False) as demo:
     gr.Markdown("# Mickey Mouse Voice Assistant")
 
     with gr.Row():
         with gr.Column():
-            audio_input = gr.Audio(source="microphone", type="filepath", label="Record Your Voice")
-            tos_checkbox = gr.Checkbox(label="Agree to Terms & Conditions", value=False)
+            audio_input = gr.Audio(
+                source="microphone",
+                type="filepath",
+                label="Record Your Voice",
+                info="Click the microphone button to record your voice."
+            )
+            tos_checkbox = gr.Checkbox(
+                label="Agree to Terms & Conditions",
+                value=False,
+                info="I agree to the terms of service."
+            )
             submit_button = gr.Button("Send")
 
         with gr.Column():
-            info_output = gr.Textbox(label="Info", interactive=False, lines=4)
-            audio_output = gr.Audio(label="Mickey's Response", interactive=False, autoplay=True)
+            info_output = gr.Textbox(
+                label="Info",
+                interactive=False,
+                lines=4,
+            )
+            audio_output = gr.Audio(
+                label="Mickey's Response",
+                interactive=False,
+                autoplay=True,
+            )
 
-    submit_button.click(predict, inputs=[audio_input, tos_checkbox], outputs=[info_output, audio_output])
+    submit_button.click(
+        predict, 
+        inputs=[audio_input, tos_checkbox], 
+        outputs=[info_output, audio_output]
+    )
 
+# Launch the Gradio app
 demo.queue()
 demo.launch(
     server_name="0.0.0.0",