import spaces
import gradio as gr
import torch
from TTS.api import TTS
from TTS.tts.utils.text.tokenizer import TTSTokenizer
from TTS.tts.utils.text.phonemizer import Phonemizer
import os
import json
import scipy.io.wavfile as wavfile
import numpy as np

os.environ["COQUI_TOS_AGREED"] = "1"

device = "cuda"

tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)

# Inicjalizacja tokenizera i fonemizera
tokenizer = TTSTokenizer(use_phonemes=False)
phonemizer = Phonemizer()

@spaces.GPU(enable_queue=True)
def clone(text, audio):
    # Generowanie mowy
    wav = tts.tts(text=text, speaker_wav=audio, language="pl")
    
    # Konwersja do numpy array i zapisanie jako plik WAV
    wav_np = np.array(wav)
    wavfile.write("./output.wav", 24000, (wav_np * 32767).astype(np.int16))
    
    # Przetwarzanie tekstu na fonemy
    tokens = tokenizer.text_to_ids(text)
    phonemes = phonemizer.phonemize(tokens, language="pl")
    
    # Przygotowanie informacji o fonemach
    phonemes_data = []
    for i, phoneme in enumerate(phonemes):
        phonemes_data.append({
            "phoneme": phoneme,
            "index": i
        })
    
    # Zapisywanie informacji o fonemach do pliku JSON
    with open("./phonemes_info.json", "w", encoding="utf-8") as f:
        json.dump(phonemes_data, f, ensure_ascii=False, indent=2)
    
    return "./output.wav", "./phonemes_info.json"

# Interfejs Gradio
iface = gr.Interface(
    fn=clone, 
    inputs=[
        gr.Textbox(label='Tekst do syntezy'),
        gr.Audio(type='filepath', label='Plik audio z głosem referencyjnym')
    ], 
    outputs=[
        gr.Audio(type='filepath', label='Zsyntezowana mowa'),
        gr.File(label='Informacje o fonemach (JSON)')
    ],
    title='Klonowanie Głosu z Informacjami o Fonemach',
    theme=gr.themes.Base(primary_hue="teal", secondary_hue="teal", neutral_hue="slate")
)

iface.launch()