Upload app.py

app.py

@@ -1,58 +1,90 @@
 import streamlit as st
-import subprocess
-import sys
-from pathlib import Path
-import base64
-import io
-import soundfile as sf
+import torch
+from diffusers import StableBeluga2Pipeline
 import numpy as np
+import soundfile as sf
+import io
+import base64
 from huggingface_hub import hf_hub_download
+import asyncio
+import time
+from concurrent.futures import ThreadPoolExecutor
 
-# AIvtuber.pyの機能をインポート
-sys.path.append(str(Path.cwd()))
-from AIvtuber import chat_session, tts_to_wav
+# モデルのダウンロードと初期化
+@st.cache_resource
+def load_model():
+    model_id = "stabilityai/stable-beluga-2"
+    pipe = StableBeluga2Pipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+        variant="fp16"
+    )
+    if torch.cuda.is_available():
+        pipe = pipe.to("cuda")
+    return pipe
 
-# TTSモデルの設定
-model_file = "https://huggingface.co/buchi-stdesign/3DAItuber-model/resolve/main/Anneli_e116_s32000.safetensors"
-config_file = "https://huggingface.co/buchi-stdesign/3DAItuber-model/resolve/main/config.json"
-style_file = "https://huggingface.co/buchi-stdesign/3DAItuber-model/resolve/main/style_vectors.npy"
+# 音声生成を非同期で実行する関数
+def generate_audio_async(pipe, text, progress_bar):
+    try:
+        # 音声生成
+        audio = pipe(
+            text,
+            num_inference_steps=50,
+            guidance_scale=7.5
+        ).audio[0]
+        
+        # プログレスバーを更新
+        progress_bar.progress(1.0)
+        
+        return audio
+    except Exception as e:
+        st.error(f"音声生成中にエラーが発生しました: {str(e)}")
+        return None
 
-def generate_audio(text):
-    tts = TTSModel(
-        model_path=model_file,
-        config_path=config_file,
-        style_vec_path=style_file,
-        device="cpu"  # クラウド環境ではCPUを使用
-    )
-    sr, wav = tts.infer(text=text, length=0.85)
+# メインアプリケーション
+def main():
+    st.title("AI VTuber チャット")
     
-    # 音声データをBase64エンコード
-    buffer = io.BytesIO()
-    sf.write(buffer, wav, sr, format='WAV')
-    audio_base64 = base64.b64encode(buffer.getvalue()).decode()
+    # モデルの読み込み
+    pipe = load_model()
     
-    return audio_base64
-
-st.title("AI VTuber Chat")
-
-# Vroid Hubをiframeで埋め込み
-st.components.v1.iframe(
-    "https://hub.vroid.com/",
-    height=600,
-    scrolling=True
-)
-
-# チャットインターフェース
-user_input = st.text_input("メッセージを入力してください：")
-
-if user_input:
-    # AIvtuber.pyの機能を使用
-    resp = chat_session.send_message(user_input)
-    st.write("AI:", resp.text)
+    # チャット履歴の初期化
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    
+    # チャット入力
+    user_input = st.chat_input("メッセージを入力してください")
     
-    # 音声生成と再生
-    sr, wav = tts_to_wav(resp.text)
-    buffer = io.BytesIO()
-    sf.write(buffer, wav, sr, format='WAV')
-    audio_base64 = base64.b64encode(buffer.getvalue()).decode()
-    st.audio(f"data:audio/wav;base64,{audio_base64}") 
+    if user_input:
+        # ユーザーメッセージを表示
+        st.session_state.messages.append({"role": "user", "content": user_input})
+        
+        # 音声生成の進捗バー
+        progress_bar = st.progress(0.0)
+        status_text = st.empty()
+        status_text.text("音声を生成中...")
+        
+        # 音声生成を非同期で実行
+        with ThreadPoolExecutor() as executor:
+            future = executor.submit(generate_audio_async, pipe, user_input, progress_bar)
+            audio = future.result()
+        
+        if audio is not None:
+            # 音声データを保存
+            audio_path = "generated_audio.wav"
+            sf.write(audio_path, audio.cpu().numpy(), 44100)
+            
+            # 音声プレーヤーを表示
+            st.audio(audio_path)
+            
+            # ステータステキストを更新
+            status_text.text("音声生成完了！")
+            
+            # チャット履歴を表示
+            for message in st.session_state.messages:
+                with st.chat_message(message["role"]):
+                    st.write(message["content"])
+
+if __name__ == "__main__":
+    main()