import spaces # 修正1: これを絶対に「一番最初」に書く必要があります！
import gradio as gr
import os, torch, io
import tempfile 
import nltk

# 必要なデータをダウンロード
os.system('python -m unidic download')
nltk.download('averaged_perceptron_tagger_eng')

from melo.api import TTS

speed = 1.0

# ZeroGPU環境では、グローバルでのdevice判定が難しい場合がありますが、
# spacesを最初にインポートしておけばうまくハンドリングしてくれます。
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# モデルのロード
models = {
    'EN': TTS(language='EN', device=device),
    'ES': TTS(language='ES', device=device),
    'FR': TTS(language='FR', device=device),
    'ZH': TTS(language='ZH', device=device),
    'JP': TTS(language='JP', device=device),
    'KR': TTS(language='KR', device=device),
}

all_speaker_ids = []
for m in models.values():
    all_speaker_ids.extend(list(m.hps.data.spk2id.keys()))
# 重複を消してリスト化
all_speaker_ids = list(set(all_speaker_ids))

default_text_dict = {
    'EN': 'The field of text-to-speech has seen rapid development recently.',
    'JP': 'テキスト読み上げの分野は最近急速な発展を遂げています',
}

# 修正2: 関数の一行上にデコレーターを追加して、この処理中だけGPUを割り当てます
@spaces.GPU
def synthesize(text, speaker, speed, language, progress=gr.Progress()):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
        temp_path = tmp_file.name

    # 生成処理
    models[language].tts_to_file(
        text, 
        models[language].hps.data.spk2id[speaker], 
        temp_path,
        speed=speed, 
        pbar=progress.tqdm, 
        format='wav'
    )
    return temp_path

def load_speakers(language, text):
    if text in list(default_text_dict.values()):
        newtext = default_text_dict.get(language, text)
    else:
        newtext = text
    return gr.update(value=list(models[language].hps.data.spk2id.keys())[0], choices=list(models[language].hps.data.spk2id.keys())), newtext

with gr.Blocks() as demo:
    gr.Markdown('# MeloTTS Demo')
    with gr.Group():
        speaker = gr.Dropdown(choices=all_speaker_ids, interactive=True, value='EN-US', label='Speaker')
        language = gr.Radio(['EN', 'ES', 'FR', 'ZH', 'JP', 'KR'], label='Language', value='EN')
        speed = gr.Slider(label='Speed', minimum=0.1, maximum=10.0, value=1.0, interactive=True, step=0.1)
        text = gr.Textbox(label="Text to speak", value=default_text_dict['EN'])
        language.input(load_speakers, inputs=[language, text], outputs=[speaker, text])
    
    btn = gr.Button('Synthesize', variant='primary')
    aud = gr.Audio(interactive=False)
    
    btn.click(synthesize, inputs=[text, speaker, speed, language], outputs=[aud])

if __name__ == "__main__":
    try:
        # SSR無効化などのオプションをつけて起動
        demo.queue(api_open=True).launch(server_name="0.0.0.0", ssr_mode=False)
    except Exception:
        demo.queue(api_open=True).launch(server_name="0.0.0.0")