import os
import gradio as gr
import torchaudio as ta
from chatterbox.tts import ChatterboxTTS

# Explicitly force CPU computation to remain within the 16GB free instance limit
device = "cpu"
print("Loading NAMAA Egyptian TTS model on CPU...")
model = ChatterboxTTS.from_pretrained(repo_id="NAMAA-Space/NAMAA-Egyptian-TTS", device=device)
print("Model loaded successfully.")

def generate_speech(text, reference_audio):
    if not text.strip():
        return None, "Error: Text input cannot be empty."
    
    try:
        output_path = "output.wav"
        
        # Check if reference voice clip is uploaded
        if reference_audio is not None:
            wav = model.generate(text, audio_prompt_path=reference_audio)
        else:
            wav = model.generate(text)
            
        ta.save(output_path, wav, model.sr)
        return output_path, "Success"
        
    except Exception as e:
        return None, f"Error during generation: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("# NAMAA Egyptian Arabic TTS")
    
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                label="Input Text (Egyptian Arabic & English mixed)", 
                placeholder="اكتب هنا مثلاً: أنا مروح البيت دلوقتي حالا"
            )
            audio_input = gr.Audio(
                label="Optional Reference Audio for Voice Cloning (WAV/MP3)", 
                type="filepath"
            )
            submit_btn = gr.Button("Generate Audio")
            
        with gr.Column():
            audio_output = gr.Audio(label="Generated Speech Output", type="filepath")
            status_output = gr.Textbox(label="Status Logging", interactive=False)

    submit_btn.click(
        fn=generate_speech, 
        inputs=[text_input, audio_input], 
        outputs=[audio_output, status_output]
    )

if __name__ == "__main__":
    demo.launch()