from inference import STTS2
import gradio as gr
from huggingface_hub import snapshot_download
import numpy as np

def synthesise(text, progress=gr.Progress()):
    if text.strip() == "":
        raise gr.Error("Please enter some text.")
        
    if len(text) > 300:
        raise gr.Error("Please enter text less than 300 characters.")

    audio = stts2.inference(text, ref_s, alpha=0.9, beta=0.9, diffusion_steps=16, embedding_scale=1)

    audio = (audio * 32767).astype(np.int16)
    
    return 24000, audio
    
if __name__ == '__main__':
    snapshot_download(repo_id="Scralius/StyleTTS2_SIWIS_French", local_dir="Models/SIWIS")
    stts2 = STTS2(config_path='Models/SIWIS/config.yml', model_folder='Models/SIWIS')
    ref_s = stts2.compute_style('Models/SIWIS/reference_audio.wav')
    
    demo = gr.Interface(
        fn=synthesise,
        inputs=[
            gr.Textbox(label='Enter Text:', lines=5, max_lines=10, placeholder="Type your text here..."),
        ],
        outputs=[
            gr.Audio(
                label="Generated Audio:",
                autoplay=False,
                streaming=False,
                type="numpy",
            ),
        ],
        allow_flagging='never',
        title="<div style='text-align: center;'>French StyleTTS2 demo</div>",
        description=(
            "This application uses a Text-to-Speech (TTS) model trained from scratch using StyleTTS2 on the SIWIS dataset. "
            "Enter some text in the input box below and let the model read it out loud with a natural and expressive voice. "
            "The model is capable of generating high-quality speech. "
            "this demo run on CPU, so it may take a while to generate the audio, if you want fast inference, you can run the model on a GPU."
        ),
        theme="origin",
        examples=[
            ["Voici une démonstration de synthèse vocale. Ce modèle est capable de lire du texte avec une voix naturelle. Essayez d'entrer votre propre texte pour voir ses capacités."],
            ["La technologie de synthèse vocale a beaucoup évolué ces dernières années. Elle est maintenant utilisée dans de nombreux domaines, y compris les assistants virtuels, les livres audio et les systèmes de navigation."],
            ["Le modèle que nous utilisons ici a été entraîné sur le dataset SIWIS, qui contient des enregistrements de voix en français."]
        ],
        cache_examples=False
    )

    demo.launch()