import gradio as gr
import os
import traceback
from Zonos_main.tts import ZonosTTS  # Adjust import path as needed

# 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 output_path
    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)