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

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Example text in Bambara
example = "An filɛ ni ye yɔrɔ minna ni an ye an sigi ka a layɛ yala an bɛ ka baara min kɛ ɛsike a kɛlen don ka Ɲɛ wa ?"

# Load model and tokenizer
try:
    logger.info("Loading Bambara TTS model and tokenizer...")
    model = VitsModel.from_pretrained("sudoping01/bambara-tts")
    tokenizer = AutoTokenizer.from_pretrained("sudoping01/bambara-tts")
    logger.info("Successfully loaded model and tokenizer")
except Exception as e:
    logger.error(f"Failed to load model: {str(e)}")
    raise Exception(f"Model loading failed: {str(e)}")

def generate_audio(text):
    if not text.strip():
        return None, "Please enter some text to synthesize."
    
    try:
        inputs = tokenizer(text, return_tensors="pt")
        
        with torch.no_grad():
            output = model(**inputs).waveform
        
        waveform = output.squeeze().cpu().numpy()
        sample_rate = model.config.sampling_rate
        
        return (sample_rate, waveform), None
    except Exception as e:
        logger.error(f"Error during inference: {str(e)}")
        return None, f"Error generating audio: {str(e)}"

def load_example():
    return example

# Create Gradio interface
with gr.Blocks(title="Bambara TTS") as demo:
    gr.Markdown(
        """
        # Bambara TTS: Text-to-Speech for Bambara Language 🇲🇱
        
        High-quality text-to-speech for Bambara (Bamanankan), spoken by over 14 million people.
        
        - ✅ Real-time TTS with fast response
        - ✅ Runs on CPU
        
        ## How to Use
        1. Enter your Bambara text or load the example
        2. Click **"Generate Audio"** to listen
        """
    )
    
    with gr.Row():
        with gr.Column():
            text = gr.Textbox(label="Bambara Text", lines=5, placeholder="Type your text in Bambara here...")
            example_btn = gr.Button("Load Example")
    
    generate_btn = gr.Button("Generate Audio", variant="primary")
    audio_output = gr.Audio(label="Generated Audio", type="numpy")
    error_msg = gr.Textbox(label="Status", visible=False)
    
    # Footer
    gr.Markdown(
        """
        By [sudoping01](https://huggingface.co/sudoping01). Model: [sudoping01/bambara-tts](https://huggingface.co/sudoping01/bambara-tts). Fine-tuned on Meta's MMS TTS. License: CC BY-NC 4.0 (non-commercial use).
        """
    )
    
    # Connect buttons to functions
    generate_btn.click(
        fn=generate_audio,
        inputs=[text],
        outputs=[audio_output, error_msg]
    )
    example_btn.click(
        fn=load_example,
        inputs=[],
        outputs=text
    )

# Launch the demo
demo.launch()