import gradio as gr
import os
import traceback
import base64  # Add this
from Zonos_main.tts import ZonosTTS

# Initialize TTS with error handling
try:
    tts = ZonosTTS()
    print("✅ TTS model loaded successfully")
except Exception as e:
    print(f"❌ Model loading failed: {str(e)}")
    raise

def generate_audio(text):
    try:
        if not text or len(text.strip()) == 0:
            raise ValueError("Input text cannot be empty")
            
        output_path = "/tmp/output.wav"
        tts.generate(text, output_path)
        
        if not os.path.exists(output_path):
            raise FileNotFoundError("Audio file was not generated")
        
        # Return audio as base64 string
        with open(output_path, "rb") as audio_file:
            audio_base64 = base64.b64encode(audio_file.read()).decode('utf-8')
            return f"data:audio/wav;base64,{audio_base64}"
            
    except Exception as e:
        traceback.print_exc()
        raise gr.Error(f"Audio generation failed: {str(e)}")

# Gradio interface
demo = gr.Interface(
    fn=generate_audio,
    inputs=gr.Textbox(label="Input Text", placeholder="Enter text here..."),
    outputs=gr.Audio(label="Generated Speech"),
    title="Zonos TTS",
    examples=[["Hello world"], ["This is a test"]]
)

demo.launch(show_error=True)