Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import re
 import json
 import torch
@@ -8,24 +9,72 @@ import uroman as ur
 import numpy as np
 import torchaudio
 import gradio as gr
+import subprocess
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from outetts.wav_tokenizer.decoder import WavTokenizer
-from yarngpt.audiotokenizer import AudioTokenizerV2
 
-# Initialize paths and models
-tokenizer_path = "saheedniyi/YarnGPT2"
+# Check if yarngpt is installed, if not install it manually
+try:
+    from yarngpt.audiotokenizer import AudioTokenizerV2
+except ImportError:
+    print("YarnGPT not found, attempting to install...")
+    subprocess.run(["chmod", "+x", "install.sh"], check=True)
+    subprocess.run(["./install.sh"], check=True)
+    
+    # Add the yarngpt directory to the Python path
+    sys.path.append(os.path.join(os.getcwd(), "yarngpt"))
+    
+    # Try importing again
+    from yarngpt.audiotokenizer import AudioTokenizerV2
+
+# Check if model files exist
 wav_tokenizer_config_path = "wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml"
 wav_tokenizer_model_path = "wavtokenizer_large_speech_320_24k.ckpt"
 
+if not os.path.exists(wav_tokenizer_config_path) or not os.path.exists(wav_tokenizer_model_path):
+    print("Model files not found, downloading...")
+    if not os.path.exists(wav_tokenizer_config_path):
+        subprocess.run([
+            "wget", 
+            "https://huggingface.co/novateur/WavTokenizer-medium-speech-75token/resolve/main/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml"
+        ], check=True)
+    
+    if not os.path.exists(wav_tokenizer_model_path):
+        subprocess.run([
+            "wget", 
+            "https://huggingface.co/novateur/WavTokenizer-large-speech-75token/resolve/main/wavtokenizer_large_speech_320_24k.ckpt"
+        ], check=True)
+
+# Initialize paths and models
+tokenizer_path = "saheedniyi/YarnGPT2"
+
+# Add debug info
+print(f"Current directory: {os.getcwd()}")
+print(f"Files in directory: {os.listdir('.')}")
+print(f"Config exists: {os.path.exists(wav_tokenizer_config_path)}")
+print(f"Model exists: {os.path.exists(wav_tokenizer_model_path)}")
+
 # Initialize the audio tokenizer
-audio_tokenizer = AudioTokenizerV2(
-    tokenizer_path, wav_tokenizer_model_path, wav_tokenizer_config_path
-)
+try:
+    print("Initializing audio tokenizer...")
+    audio_tokenizer = AudioTokenizerV2(
+        tokenizer_path, wav_tokenizer_model_path, wav_tokenizer_config_path
+    )
+    print("Audio tokenizer initialized")
+except Exception as e:
+    print(f"Error initializing audio tokenizer: {str(e)}")
+    raise
 
 # Load the model
-model = AutoModelForCausalLM.from_pretrained(
-    tokenizer_path, torch_dtype="auto"
-).to(audio_tokenizer.device)
+try:
+    print("Loading model...")
+    model = AutoModelForCausalLM.from_pretrained(
+        tokenizer_path, torch_dtype="auto"
+    ).to(audio_tokenizer.device)
+    print("Model loaded")
+except Exception as e:
+    print(f"Error loading model: {str(e)}")
+    raise
 
 # Function to generate speech
 def generate_speech(text, language, speaker_name, temperature=0.1, repetition_penalty=1.1):
@@ -56,6 +105,7 @@ def generate_speech(text, language, speaker_name, temperature=0.1, repetition_pe
 # Create Gradio interface
 def tts_interface(text, language, speaker_name, temperature, repetition_penalty):
     try:
+        print(f"Generating speech for: {text[:30]}...")
         audio_path = generate_speech(
             text, 
             language, 
@@ -63,8 +113,10 @@ def tts_interface(text, language, speaker_name, temperature, repetition_penalty)
             temperature,
             repetition_penalty
         )
+        print("Speech generated successfully")
         return audio_path
     except Exception as e:
+        print(f"Error in tts_interface: {str(e)}")
         return f"Error: {str(e)}"
 
 # Define available languages and speakers
@@ -75,7 +127,7 @@ speakers = ["idera", "enitan", "abeo", "eniola", "kachi", "aisha", "amara", "bel
 demo = gr.Interface(
     fn=tts_interface,
     inputs=[
-        gr.Textbox(label="Text to convert to speech", lines=5),
+        gr.Textbox(label="Text to convert to speech", lines=5, value="Welcome to YarnGPT text-to-speech model for African languages."),
         gr.Dropdown(languages, label="Language", value="english"),
         gr.Dropdown(speakers, label="Speaker", value="idera"),
         gr.Slider(0.1, 1.0, value=0.1, label="Temperature"),
@@ -88,4 +140,5 @@ demo = gr.Interface(
 
 # Launch the app
 if __name__ == "__main__":
+    print("Starting Gradio interface...")
     demo.launch()