import gradio as gr
from transformers import VitsModel, AutoTokenizer
import torch
import scipy.io.wavfile

# Charger le modèle Vits
model = VitsModel.from_pretrained("joefox/tts_vits_ru_hf")
tokenizer = AutoTokenizer.from_pretrained("joefox/tts_vits_ru_hf")

def generate_audio(text):
    # Prétraiter le texte
    text = text.lower()
    inputs = tokenizer(text, return_tensors="pt")
    inputs['speaker_id'] = 3
    
    # Générer l'audio
    with torch.no_grad():
        output = model(**inputs).waveform
        
    # Enregistrer l'audio dans un fichier
    scipy.io.wavfile.write("output.wav", rate=model.config.sampling_rate, data=output[0].cpu().numpy())
    
    return "output.wav"

# Créer une interface Gradio
iface = gr.Interface(fn=generate_audio, inputs="text", outputs="audio", title="Text to Speech (Vits russe)", description="Saisissez votre texte russe pour convertir en discours audio russe.")
iface.launch()