import gradio as gr
from transformers import VitsModel, AutoTokenizer
import torch

# Load model and tokenizer
print("Loading model...")
model = VitsModel.from_pretrained("facebook/mms-tts-kin")
tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-kin")
print("Model loaded successfully!")

def text_to_speech(text):
    """
    Generate speech from Kinyarwanda text
    
    Args:
        text: Input text in Kinyarwanda
    
    Returns:
        Tuple of (sample_rate, audio_array)
    """
    if not text or text.strip() == "":
        return None
    
    try:
        # Tokenize input text
        inputs = tokenizer(text, return_tensors="pt")
        
        # Generate speech
        with torch.no_grad():
            output = model(**inputs).waveform
        
        # Convert to numpy array
        audio_data = output.squeeze().cpu().numpy()
        
        # Return sample rate and audio data
        return (model.config.sampling_rate, audio_data)
    
    except Exception as e:
        print(f"Error: {e}")
        return None

# Create Gradio interface
demo = gr.Interface(
    fn=text_to_speech,
    inputs=gr.Textbox(
        label="Kinyarwanda Text",
        placeholder="Enter text in Kinyarwanda (e.g., 'Muraho Rwanda')",
        lines=3
    ),
    outputs=gr.Audio(label="Generated Speech"),
    title="🗣️ Kinyarwanda Text-to-Speech",
    description="""
    Generate natural speech from Kinyarwanda text using Meta's MMS-TTS model.
    
    Simply enter text in Kinyarwanda and click Submit to hear it spoken.
    
    Model: [facebook/mms-tts-kin](https://huggingface.co/facebook/mms-tts-kin)
    """,
    examples=[
        ["Muraho"],
        ["Mwaramutse"],
        ["Amakuru"],
        ["Murakoze cyane"],
    ],
    theme=gr.themes.Soft()
)

# Launch the app
if __name__ == "__main__":
    demo.launch(share=True)