import os
from flask import jsonify, send_file, request
from main import *
import torch
import torchaudio
import uuid

def text_to_speech_func(text):
    if tts_model is None:
        return {"error": "TTS model not initialized."}
    input_tokens = tts_model.tokenizer(text, return_tensors="pt", padding=True).to(device)
    with torch.no_grad():
        audio_output = tts_model(input_tokens['input_ids'])
    temp_audio_path = f"temp_audio_{uuid.uuid4()}.wav"
    torchaudio.save(temp_audio_path, audio_output.cpu(), 16000)
    return temp_audio_path

def tts_api():
    data = request.get_json()
    text = data.get('text')
    if not text:
        return jsonify({"error": "Text is required"}), 400
    output_file = text_to_speech_func(text)
    if "error" in output:
        return jsonify({"error": output["error"]}), 500
    return send_file(output_file, mimetype="audio/wav", as_attachment=True, download_name="output.wav")