import io
import torch
import soundfile as sf
from flask import Flask, request, send_file
from transformers import AutoModelForTextToSpeech, AutoProcessor

app = Flask(__name__)

print("⏳ SPRINGLab/F5-Hindi-24KHz मॉडल लोड हो रहा है...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# मॉडल और प्रोसेसर को हगिंग फेस से लोड करें
model = AutoModelForTextToSpeech.from_pretrained("SPRINGLab/F5-Hindi-24KHz").to(device)
processor = AutoProcessor.from_pretrained("SPRINGLab/F5-Hindi-24KHz")

print("✅ मॉडल तैयार है!")

@app.route('/tts', methods=['POST'])
def tts_api():
    text = request.form.get('text', '')
    if not text:
        return "Error: 'text' parameter is missing", 400

    # टेक्स्ट को इनपुट फॉर्मेट में बदलें
    inputs = processor(text=text, return_tensors="pt").to(device)

    with torch.no_grad():
        output = model(**inputs)
        waveform = output.waveform.cpu().numpy().squeeze()
        sample_rate = output.sample_rate

    # WAV फ़ाइल के रूप में सेव करें
    audio_buffer = io.BytesIO()
    sf.write(audio_buffer, waveform, sample_rate, format='WAV')
    audio_buffer.seek(0)

    return send_file(
        audio_buffer,
        mimetype='audio/wav',
        as_attachment=True,
        download_name='output.wav'
    )

@app.route('/')
def home():
    return "🎙️ SPRINGLab/F5-Hindi-24KHz TTS API चल रही है!"