Update app.py

app.py

@@ -1,50 +1,36 @@
-import gradio as gr
-import torch
-import torchaudio
 import os
 import re
-import subprocess
-from transformers import AutoModelForCausalLM
-from yarngpt_utils import AudioTokenizer
+import json
+import torch
+import inflect
+import random
+import uroman as ur
+import numpy as np
+import torchaudio
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from outetts.wav_tokenizer.decoder import WavTokenizer
+from yarngpt.audiotokenizer import AudioTokenizerV2
 
-# Download model files if they don't exist
-def download_if_not_exists(url, filename):
-    if not os.path.exists(filename):
-        print(f"Downloading {filename}...")
-        subprocess.run(["wget", url, "-O", filename])
-        print(f"Downloaded {filename}")
+# Initialize paths and models
+tokenizer_path = "saheedniyi/YarnGPT2"
+wav_tokenizer_config_path = "wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml"
+wav_tokenizer_model_path = "wavtokenizer_large_speech_320_24k.ckpt"
 
-# Download necessary files
-download_if_not_exists(
-    "https://huggingface.co/novateur/WavTokenizer-medium-speech-75token/resolve/main/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml",
-    "wavtokenizer_config.yaml"
-)
-download_if_not_exists(
-    "https://huggingface.co/novateur/WavTokenizer-large-speech-75token/blob/main/wavtokenizer_large_speech_320_v2.ckpt",
-    "wavtokenizer_model.ckpt"
+# Initialize the audio tokenizer
+audio_tokenizer = AudioTokenizerV2(
+    tokenizer_path, wav_tokenizer_model_path, wav_tokenizer_config_path
 )
 
-# Initialize the model (this runs when the app starts)
-def initialize_model():
-    # Set paths
-    hf_path = "saheedniyi/YarnGPT"
-    wav_tokenizer_config_path = "wavtokenizer_config.yaml"
-    wav_tokenizer_model_path = "wavtokenizer_model.ckpt"
-    
-    # Create AudioTokenizer
-    audio_tokenizer = AudioTokenizer(
-        hf_path, wav_tokenizer_model_path, wav_tokenizer_config_path
-    )
-    
-    # Load model
-    model = AutoModelForCausalLM.from_pretrained(hf_path, torch_dtype="auto").to(audio_tokenizer.device)
-    
-    return model, audio_tokenizer
+# Load the model
+model = AutoModelForCausalLM.from_pretrained(
+    tokenizer_path, torch_dtype="auto"
+).to(audio_tokenizer.device)
 
-# Generate audio from text
-def generate_speech(text, speaker_name):
+# Function to generate speech
+def generate_speech(text, language, speaker_name, temperature=0.1, repetition_penalty=1.1):
     # Create prompt
-    prompt = audio_tokenizer.create_prompt(text, speaker_name)
+    prompt = audio_tokenizer.create_prompt(text, lang=language, speaker_name=speaker_name)
     
     # Tokenize prompt
     input_ids = audio_tokenizer.tokenize_prompt(prompt)
@@ -52,45 +38,54 @@ def generate_speech(text, speaker_name):
     # Generate output
     output = model.generate(
         input_ids=input_ids,
-        temperature=0.1,
-        repetition_penalty=1.1,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
         max_length=4000,
     )
     
-    # Convert to audio codes
+    # Get audio codes and convert to audio
     codes = audio_tokenizer.get_codes(output)
-    
-    # Convert codes to audio
     audio = audio_tokenizer.get_audio(codes)
     
-    # Save audio temporarily
-    temp_path = "output.wav"
-    torchaudio.save(temp_path, audio, sample_rate=24000)
+    # Save audio to file
+    output_path = "output.wav"
+    torchaudio.save(output_path, audio, sample_rate=24000)
     
-    return temp_path
-
-# Load model globally
-print("Loading model...")
-model, audio_tokenizer = initialize_model()
-print("Model loaded!")
-
-# Add this before initializing the model
-from inspect import signature
-from outetts.wav_tokenizer.decoder import WavTokenizer
-print("WavTokenizer parameters:", signature(WavTokenizer.__init__))
+    return output_path
 
 # Create Gradio interface
-speakers = ["idera", "emma", "jude", "osagie", "tayo", "zainab", "joke", "regina", "remi", "umar", "chinenye"]
+def tts_interface(text, language, speaker_name, temperature, repetition_penalty):
+    try:
+        audio_path = generate_speech(
+            text, 
+            language, 
+            speaker_name,
+            temperature,
+            repetition_penalty
+        )
+        return audio_path
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Define available languages and speakers
+languages = ["english", "igbo", "yoruba", "hausa", "pidgin"]
+speakers = ["idera", "enitan", "abeo", "eniola", "kachi", "aisha", "amara", "bello", "chidi"]
 
+# Create the Gradio interface
 demo = gr.Interface(
-    fn=generate_speech,
+    fn=tts_interface,
     inputs=[
-        gr.Textbox(lines=5, placeholder="Enter text here..."),
-        gr.Dropdown(choices=speakers, label="Speaker", value="idera")
+        gr.Textbox(label="Text to convert to speech", lines=5),
+        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"),
+        gr.Slider(1.0, 2.0, value=1.1, label="Repetition Penalty"),
     ],
     outputs=gr.Audio(type="filepath"),
-    title="YarnGPT: Nigerian Accented Text-to-Speech",
-    description="Generate natural-sounding Nigerian accented speech from text."
+    title="YarnGPT Text-to-Speech",
+    description="Convert text to speech using YarnGPT model for various African languages",
 )
 
-demo.launch()
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()