app.py

import os
import gradio as gr
import torch
import numpy as np
import soundfile as sf
from outetts import Interface, ModelConfig, GenerationConfig, Backend, Models, GenerationType, SamplerConfig # SamplerConfigをインポート
import spaces # spacesライブラリをインポート
import uuid # 一時ファイル名生成用

# モデル設定
MODEL_PATH = "OuteAI/OuteTTS-1.0-0.6B"

# ディレクトリの作成
os.makedirs("samples", exist_ok=True)
os.makedirs("outputs", exist_ok=True)
os.makedirs("speakers", exist_ok=True)
os.makedirs("temp_speakers", exist_ok=True) # 一時スピーカープロファイル保存用

# GPUワーカープロセス内で使用されるグローバル変数
# これらの変数は、各GPUワーカープロセスごとに独立して存在します。
_gpu_interface = None

# インターフェースの初期化 (GPUワーカープロセス内で実行される)
@spaces.GPU # GPUリソースを必要とすることを明示的に指定
def _initialize_model_in_gpu_worker():
    global _gpu_interface

    if _gpu_interface is None:
        if torch.cuda.is_available():
            print("CUDA is available. OuteTTS will attempt to use GPU.")
        else:
            print("CUDA is not available. OuteTTS will run on CPU.")

        _gpu_interface = Interface(
            ModelConfig.auto_config(
                model=Models.VERSION_1_0_SIZE_0_6B,
                backend=Backend.HF,  # Hugging Face Transformersバックエンド
            )
        )
        print("OuteTTS model initialized successfully in GPU worker.")
    # この関数はモデルオブジェクトを返しません。
    # 代わりに、GPUワーカープロセス内のグローバル変数 _gpu_interface を設定します。

# 各Gradio関数から呼び出されるモデルロードヘルパー
def _get_gpu_interface():
    # _initialize_model_in_gpu_worker() を呼び出すことで、
    # モデルがGPUワーカープロセス内でロードされることを保証します。
    _initialize_model_in_gpu_worker()
    return _gpu_interface

# スピーカープロファイルの作成 (GPUワーカープロセス内で実行される)
@spaces.GPU
def create_speaker_profile_gpu(audio_file):
    if audio_file is None:
        return None, None, gr.update(value="エラー: 音声ファイルが選択されていません。", interactive=True), None # パスもNoneに
    
    audio_file_path = audio_file.name # GradioのFileコンポーネントはname属性でパスを提供
    
    interface = _get_gpu_interface()
    
    try:
        speaker = interface.create_speaker(audio_file_path)
        speaker_name = os.path.basename(audio_file_path)
        
        # 一時ファイルにスピーカープロファイルを保存
        temp_speaker_filename = f"temp_speaker_{uuid.uuid4()}.json"
        temp_speaker_path = os.path.join("temp_speakers", temp_speaker_filename)
        interface.save_speaker(speaker, temp_speaker_path)
        
        decoded_path = f"speakers/decoded_{speaker_name}.wav"
        interface.decode_and_save_speaker(speaker=speaker, path=decoded_path)
        
        status_message = f"スピーカープロファイル '{speaker_name}' を作成しました。"
        return temp_speaker_path, decoded_path, gr.update(value=status_message, interactive=True), temp_speaker_path
    except Exception as e:
        error_message = f"エラー: スピーカープロファイルの作成に失敗しました: {str(e)}"
        print(error_message)
        return None, None, gr.update(value=error_message, interactive=True), None


# デフォルトスピーカーのロード (GPUワーカープロセス内で実行される)
@spaces.GPU
def load_default_speaker_gpu(language, gender, style):
    interface = _get_gpu_interface()
    speaker_id = f"{language}-{gender}-1-{style}"
    try:
        # Python実装では "EN-FEMALE-1-NEUTRAL" のみがデフォルトで利用可能
        # UIで選択肢を限定しているため、このチェックは主に安全のため
        if speaker_id != "EN-FEMALE-1-NEUTRAL":
            error_message = f"エラー: スピーカー {speaker_id} はデフォルトでは利用できません。利用可能なのは 'EN-FEMALE-1-NEUTRAL' のみです。"
            print(error_message)
            return error_message, None

        speaker = interface.load_default_speaker(speaker_id)
        
        # ロードしたデフォルトスピーカーを一時ファイルに保存
        temp_speaker_filename = f"temp_default_speaker_{uuid.uuid4()}.json"
        temp_speaker_path = os.path.join("temp_speakers", temp_speaker_filename)
        interface.save_speaker(speaker, temp_speaker_path)

        status_message = f"デフォルトスピーカー {speaker_id} をロードしました"
        return status_message, temp_speaker_path
    except Exception as e:
        error_message = f"エラー: {str(e)}"
        print(error_message)
        return error_message, None

# 保存されたスピーカーのロード (GPUワーカープロセス内で実行される)
@spaces.GPU
def load_saved_speaker_gpu(speaker_path):
    if not speaker_path:
        return "スピーカーファイルを選択してください", None
    interface = _get_gpu_interface()
    try:
        speaker = interface.load_speaker(speaker_path)
        
        # ロードしたスピーカーを一時ファイルに保存 (念のため)
        temp_speaker_filename = f"temp_saved_speaker_{uuid.uuid4()}.json"
        temp_speaker_path = os.path.join("temp_speakers", temp_speaker_filename)
        interface.save_speaker(speaker, temp_speaker_path)

        status_message = f"スピーカー {os.path.basename(speaker_path)} をロードしました"
        return status_message, temp_speaker_path
    except Exception as e:
        error_message = f"エラー: {str(e)}"
        print(error_message)
        return error_message, None

# 音声生成 (GPUワーカープロセス内で実行される)
@spaces.GPU
def generate_speech_gpu(text, temperature, repetition_penalty, top_k, top_p, min_p, current_speaker_path):
    if not text:
        return "テキストを入力してください", None
    
    if current_speaker_path is None:
        return "スピーカープロファイルをロードしてください", None
    
    interface = _get_gpu_interface()
    
    try:
        # パスからスピーカーをロード
        speaker = interface.load_speaker(current_speaker_path)
        
        # SamplerConfigを生成
        sampler_config = SamplerConfig(
            temperature=temperature,
            repetition_penalty=repetition_penalty,
            top_k=top_k,
            top_p=top_p,
            min_p=min_p
            # mirostat_eta, mirostat_tau はUIにないので含めない
        )

        # GenerationConfigの引数をOuteTTSがサポートするものに限定
        # repetition_range は GenerationConfig の直接の引数ではないため削除
        output = interface.generate(
            GenerationConfig(
                text=text,
                speaker=speaker,
                sampler_config=sampler_config, # SamplerConfigオブジェクトを渡す
                # repetition_range=64, # ドキュメントに記載がない、またはSamplerConfig内にある可能性
            )
        )
        output_path = f"outputs/output_{hash(text)}.wav"
        output.save(output_path)
        return "音声生成が完了しました", output_path
    except Exception as e:
        error_message = f"エラー: {str(e)}"
        print(error_message)
        return error_message, None

# Gradio UI
def create_ui():
    with gr.Blocks(title="OuteTTS 音声生成デモ") as demo:
        gr.Markdown("# OuteTTS-1.0-0.6B 音声生成デモ")
        gr.Markdown("このデモでは、OuteTTS-1.0-0.6Bモデルを使用して、テキストから音声を生成します。音声クローニング機能も利用できます。")

        # 現在のスピーカープロファイルのパスを保持する隠しStateコンポーネント
        current_speaker_path_state = gr.State(value=None)

        with gr.Tab("デフォルトスピーカー"):
            with gr.Row():
                # Python実装では "EN-FEMALE-1-NEUTRAL" のみがデフォルトで利用可能
                language = gr.Dropdown(choices=["EN"], value="EN", label="言語")
                gender = gr.Dropdown(choices=["FEMALE"], value="FEMALE", label="性別")
                style = gr.Dropdown(choices=["NEUTRAL"], value="NEUTRAL", label="スタイル")

            load_default_btn = gr.Button("デフォルトスピーカーをロード")
            default_speaker_status = gr.Textbox(label="ステータス")

            # GPU対応関数を呼び出す
            load_default_btn.click(
                fn=load_default_speaker_gpu,
                inputs=[language, gender, style],
                outputs=[default_speaker_status, current_speaker_path_state] # 隠しStateも更新
            )

        with gr.Tab("音声クローニング"):
            gr.Markdown("### 音声ファイルからスピーカープロファイルを作成")
            audio_file = gr.File(label="音声ファイル（5〜10秒の明瞭な音声を推奨）")
            create_speaker_btn = gr.Button("スピーカープロファイルを作成")

            with gr.Row():
                speaker_path_output = gr.Textbox(label="保存されたスピーカープロファイルのパス")
                decoded_audio = gr.Audio(label="再構成された音声（品質確認用）")
                # スピーカー作成時のステータス表示用
                create_speaker_status = gr.Textbox(label="スピーカー作成ステータス") 

            # GPU対応関数を呼び出す
            create_speaker_btn.click(
                fn=create_speaker_profile_gpu,
                inputs=[audio_file],
                outputs=[speaker_path_output, decoded_audio, create_speaker_status, current_speaker_path_state] # 隠しStateも更新
            )

            gr.Markdown("### 保存されたスピーカープロファイルをロード")
            saved_speaker_path = gr.Textbox(label="スピーカープロファイルのパス")
            load_saved_btn = gr.Button("保存されたスピーカーをロード")
            saved_speaker_status = gr.Textbox(label="ステータス")

            # GPU対応関数を呼び出す
            load_saved_btn.click(
                fn=load_saved_speaker_gpu,
                inputs=[saved_speaker_path],
                outputs=[saved_speaker_status, current_speaker_path_state] # 隠しStateも更新
            )

        gr.Markdown("## テキストから音声を生成")

        with gr.Row():
            with gr.Column():
                text_input = gr.Textbox(label="テキスト", lines=5, placeholder="ここにテキストを入力してください...")

                with gr.Accordion("詳細設定", open=False):
                    # これらのスライダーはGradio UIに残しますが、
                    # GenerationConfigには渡しません。
                    # outettsのドキュメントでサポートされているか確認し、
                    # 必要であればGenerationConfigに含めてください。
                    temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Temperature")
                    repetition_penalty = gr.Slider(minimum=1.0, maximum=1.5, value=1.1, step=0.05, label="Repetition Penalty")
                    top_k = gr.Slider(minimum=1, maximum=100, value=40, step=1, label="Top-k")
                    top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p")
                    min_p = gr.Slider(minimum=0.01, maximum=0.2, value=0.05, step=0.01, label="Min-p")

                generate_btn = gr.Button("音声を生成")

            with gr.Column():
                status_output = gr.Textbox(label="ステータス")
                audio_output = gr.Audio(label="生成された音声")
                # GPU対応関数を呼び出す
                generate_btn.click(
            fn=generate_speech_gpu,
            inputs=[text_input, temperature, repetition_penalty, top_k, top_p, min_p, current_speaker_path_state], # 隠しStateをinputsに追加
            outputs=[status_output, audio_output]
        )

        gr.Markdown("## サンプル")

        with gr.Accordion("使用例", open=True):
            gr.Markdown("""
            ### 使用例:

            1. 「デフォルトスピーカー」タブで言語・性別・スタイルを選択し、「デフォルトスピーカーをロード」をクリックします。
            2. または「音声クローニング」タブで音声ファイルをアップロードし、「スピーカープロファイルを作成」をクリックします。
            3. テキスト入力欄にテキストを入力します。
            4. 「音声を生成」ボタンをクリックします。
            5. 生成された音声が表示されます。

            ### 注意事項:

            - 音声クローニングには、5〜10秒の明瞭な音声を使用することをお勧めします。
            - 日本語と英語の両方に対応していますが、最良の結果を得るには、使用する言語に合ったスピーカープロファイルを作成してください。
            - 生成される音声の品質は、スピーカープロファイルの品質に大きく依存します。
            """)
    return demo

# アプリケーションの起動
demo = create_ui()
demo.launch()