Update app.py

app.py

@@ -13,76 +13,73 @@ import subprocess
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from outetts.wav_tokenizer.decoder import WavTokenizer
 
-# 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
+# Clone the YarnGPT repository if it doesn't exist
+if not os.path.exists("yarngpt"):
+    print("Cloning YarnGPT repository...")
+    subprocess.run(["git", "clone", "https://github.com/saheedniyi02/yarngpt.git"], check=True)
+
+# Add the yarngpt directory to the Python path
+yarngpt_path = os.path.abspath("yarngpt")
+if yarngpt_path not in sys.path:
+    sys.path.append(yarngpt_path)
+    print(f"Added {yarngpt_path} to Python path")
+
+# Now try importing from yarngpt
+from yarngpt.audiotokenizer import AudioTokenizerV2
 
-# Check if model files exist
+# Download model files if they don't 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_config_path):
+    print(f"Downloading {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)
+if not os.path.exists(wav_tokenizer_model_path):
+    print(f"Downloading {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 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
-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
+print("Initializing audio tokenizer...")
+audio_tokenizer = AudioTokenizerV2(
+    tokenizer_path, wav_tokenizer_model_path, wav_tokenizer_config_path
+)
+print("Audio tokenizer initialized")
 
 # Load the model
-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
+print("Loading model...")
+model = AutoModelForCausalLM.from_pretrained(
+    tokenizer_path, torch_dtype="auto"
+).to(audio_tokenizer.device)
+print("Model loaded successfully")
 
 # Function to generate speech
 def generate_speech(text, language, speaker_name, temperature=0.1, repetition_penalty=1.1):
+    print(f"Generating speech for: '{text[:50]}...'")
+    print(f"Parameters: language={language}, speaker={speaker_name}, temp={temperature}, rep_penalty={repetition_penalty}")
+    
     # Create prompt
     prompt = audio_tokenizer.create_prompt(text, lang=language, speaker_name=speaker_name)
+    print("Prompt created")
     
     # Tokenize prompt
     input_ids = audio_tokenizer.tokenize_prompt(prompt)
+    print("Prompt tokenized")
     
     # Generate output
     output = model.generate(
@@ -91,21 +88,25 @@ def generate_speech(text, language, speaker_name, temperature=0.1, repetition_pe
         repetition_penalty=repetition_penalty,
         max_length=4000,
     )
+    print("Model generation complete")
     
     # Get audio codes and convert to audio
     codes = audio_tokenizer.get_codes(output)
+    print("Audio codes extracted")
+    
     audio = audio_tokenizer.get_audio(codes)
+    print("Audio generated")
     
     # Save audio to file
     output_path = "output.wav"
     torchaudio.save(output_path, audio, sample_rate=24000)
+    print(f"Audio saved to {output_path}")
     
     return output_path
 
 # 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, 
@@ -113,10 +114,11 @@ 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)}")
+        import traceback
+        error_details = traceback.format_exc()
+        print(f"Error in tts_interface: {str(e)}\n{error_details}")
         return f"Error: {str(e)}"
 
 # Define available languages and speakers
@@ -135,7 +137,12 @@ demo = gr.Interface(
     ],
     outputs=gr.Audio(type="filepath"),
     title="YarnGPT Text-to-Speech",
-    description="Convert text to speech using YarnGPT model for various African languages",
+    description="Convert text to speech using YarnGPT model for various African languages.",
+    examples=[
+        ["The election was won by businessman and politician, Moshood Abiola, but Babangida annulled the results, citing concerns over national security.", "english", "idera", 0.1, 1.1],
+        ["Hello, how are you today?", "english", "enitan", 0.1, 1.1],
+        ["Bawo ni?", "yoruba", "eniola", 0.2, 1.2],
+    ]
 )
 
 # Launch the app