Update app.py

app.py

@@ -10,15 +10,16 @@ from api import BaseSpeakerTTS, ToneColorConverter
 import langid
 import traceback
 from dotenv import load_dotenv
-from fastapi import FastAPI, UploadFile, Form
-from fastapi.responses import JSONResponse
-from gradio.routes import mount_gradio_app
 
 # Load environment variables
 load_dotenv()
 
-# Initialize FastAPI app
-app = FastAPI()
+# 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():
@@ -38,53 +39,52 @@ def download_and_extract_checkpoints():
         os.remove(zip_path)
         print("Checkpoints are ready.")
 
-# Call the function to ensure checkpoints are available
-download_and_extract_checkpoints()
+# Initialize models and resources
+def initialize_resources():
+    global en_base_speaker_tts, zh_base_speaker_tts, tone_color_converter, target_se
+    print("Initializing resources...")
 
-# 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.")
+    # 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'
-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)}")
+    # 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')
+
+    # 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')
+
+    # 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)
+
+    # 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()
 
 # Supported languages
 supported_languages = ['zh', 'en']
 
-# Predict function (shared between FastAPI and Gradio)
+# Predict function
 def predict(audio_file_pth, agree):
     text_hint = ''
     synthesized_audio_path = None
@@ -124,25 +124,19 @@ def predict(audio_file_pth, agree):
     print(f"Detected language: {language_predicted}")
 
     if language_predicted not in supported_languages:
-        text_hint += f"[ERROR] The detected language '{language_predicted}' is not supported.\n"
+        text_hint += f"[ERROR] Unsupported language: {language_predicted}\n"
         return (text_hint, None)
 
-    # 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'
+    # 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'
 
     # 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 friendly character."},
+                {"role": "system", "content": "You are Mickey Mouse, a cheerful character who responds to children's queries."},
                 {"role": "user", "content": input_text}
             ]
         )
@@ -155,44 +149,25 @@ def predict(audio_file_pth, agree):
     # Synthesize reply text to audio
     try:
         src_path = os.path.join(output_dir, 'tmp_reply.wav')
-        tts_model.tts(reply_text, src_path, speaker=speaker_style, language=language)
+        tts_model.tts(reply_text, src_path, speaker='default', language=language)
 
         save_path = os.path.join(output_dir, 'output_reply.wav')
         tone_color_converter.convert(
-            audio_src_path=src_path, 
-            src_se=en_source_default_se if language == 'English' else zh_source_se,
+            audio_src_path=src_path,
+            src_se=target_se,
             tgt_se=target_se,
             output_path=save_path
         )
 
-        text_hint += "Response generated successfully."
+        text_hint += "Response generated successfully.\n"
         synthesized_audio_path = save_path
-
     except Exception as e:
         text_hint += f"[ERROR] Synthesis failed: {str(e)}\n"
+        traceback.print_exc()
         return (text_hint, None)
 
     return (text_hint, synthesized_audio_path)
 
-
-# FastAPI endpoint for prediction
-@app.post("/predict")
-async def predict_endpoint(file: UploadFile, agree: bool = Form(...)):
-    # Save uploaded file
-    temp_file_path = f"temp_{file.filename}"
-    with open(temp_file_path, "wb") as temp_file:
-        temp_file.write(await file.read())
-
-    # Call predict
-    info, audio_path = predict(temp_file_path, agree)
-    os.remove(temp_file_path)
-
-    if audio_path:
-        return JSONResponse({"info": info, "audio": audio_path})
-    else:
-        return JSONResponse({"info": info}, status_code=400)
-
-
 # Gradio UI
 with gr.Blocks(analytics_enabled=False) as demo:
     gr.Markdown("# Mickey Mouse Voice Assistant")
@@ -209,5 +184,11 @@ with gr.Blocks(analytics_enabled=False) as demo:
 
     submit_button.click(predict, inputs=[audio_input, tos_checkbox], outputs=[info_output, audio_output])
 
-# Mount Gradio app to FastAPI
-mount_gradio_app(app, demo, path="/")
+demo.queue()
+demo.launch(
+    server_name="0.0.0.0",
+    server_port=int(os.environ.get("PORT", 7860)),
+    debug=True,
+    show_api=True,
+    share=False
+)