Update app.py

app.py

@@ -5,10 +5,15 @@ import torch
 import numpy as np
 import scipy.io.wavfile
 import gradio as gr
+import edge_tts
+import asyncio
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from groq import Groq # pip install groq
+from groq import Groq
+
+# API 키 및 환경 설정
+GENAI_KEY = os.getenv("GEMINI_KEY") # 기존 키 유지 시
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 
-# 1. 모델 싱글톤 관리 클래스 (메모리 효율화)
 class ModelManager:
     _llm_pipeline = None
     _music_pipeline = None
@@ -17,102 +22,84 @@ class ModelManager:
     @classmethod
     def get_qwen(cls):
         if cls._llm_pipeline is None:
+            # CPU 최적화 버전 Qwen 0.5B
             model_id = "Qwen/Qwen2.5-0.5B-Instruct"
             tokenizer = AutoTokenizer.from_pretrained(model_id)
-            # CPU 환경 최적화 로딩
-            model = AutoModelForCausalLM.from_pretrained(
-                model_id, 
-                device_map="cpu", 
-                torch_dtype=torch.float32,
-                low_cpu_mem_usage=True
-            )
+            model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cpu", torch_dtype=torch.float32)
             cls._llm_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
         return cls._llm_pipeline
 
     @classmethod
     def get_groq(cls):
-        api_key = os.getenv("GROQ_API_KEY")
-        if cls._groq_client is None and api_key:
-            cls._groq_client = Groq(api_key=api_key)
+        if cls._groq_client is None and GROQ_API_KEY:
+            cls._groq_client = Groq(api_key=GROQ_API_KEY)
         return cls._groq_client
 
     @classmethod
     def get_music(cls):
         if cls._music_pipeline is None:
-            # CPU에서 가장 가벼운 음악 모델 유지
             cls._music_pipeline = pipeline("text-to-audio", "facebook/musicgen-small", device="cpu")
         return cls._music_pipeline
 
-# 2. 핵심 비즈니스 로직
+# 8인 멤버 보이스 매핑
+MEMBERS_VOICE = {
+    "서윤 (Korea)": "ko-KR-SunHiNeural", "Chloe (USA)": "en-US-AriaNeural", 
+    "Naomi (Japan)": "ja-JP-NanamiNeural", "Beatrice (Brazil)": "pt-BR-FranciscaNeural",
+    "Elena (Spain)": "es-ES-ElviraNeural", "Amira (Egypt)": "ar-EG-SalmaNeural",
+    "Liwei (China)": "zh-CN-XiaoxiaoNeural", "Sophie (France)": "fr-FR-DeniseNeural"
+}
+
 async def band_consulting(user_input, member_name, lang_code):
     req_id = str(uuid.uuid4())[:8]
+    voice_path = f"/tmp/v_{req_id}.mp3"
     music_path = f"/tmp/m_{req_id}.wav"
     
-    # 시스템 프롬프트 설정
-    system_prompt = f"You are {member_name}, a rock star. Talk in {lang_code}. Format: 5 lines max + [TAB] details + [MUSIC: prompt]."
+    system_prompt = f"당신은 밴드 멤버 {member_name}입니다. {lang_code}로 답변하세요. 5줄 이내 요약, 상세 설명은 [TAB]에, 음악 프롬프트는 [MUSIC: 영어프롬프트]에 넣으세요."
     
     ai_text_raw = ""
-    
-    # 엔진 선택 (Groq 우선 -> 실패 시 로컬 Qwen)
+    # 1. Groq 시도
     groq_client = ModelManager.get_groq()
     if groq_client:
         try:
-            chat_completion = groq_client.chat.completions.create(
-                messages=[
-                    {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": user_input}
-                ],
-                model="llama-3.3-70b-versatile",
+            res = groq_client.chat.completions.create(
+                messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": user_input}],
+                model="llama-3.3-70b-versatile"
             )
-            ai_text_raw = chat_completion.choices[0].message.content
-        except Exception as e:
-            print(f"Groq Error, switching to Qwen: {e}")
+            ai_text_raw = res.choices[0].message.content
+        except: pass
 
-    # Groq 실패 시 혹은 미설정 시 로컬 Qwen 사용
+    # 2. 로컬 Qwen Fallback
     if not ai_text_raw:
-        try:
-            qwen = ModelManager.get_qwen()
-            messages = [
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": user_input}
-            ]
-            text = qwen.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-            output = qwen(text, max_new_tokens=512, do_sample=True, temperature=0.7)
-            ai_text_raw = output[0]['generated_text'].split("assistant\n")[-1]
-        except Exception as e:
-            ai_text_raw = "이봐 리더, 지금은 말보다 연주야! [TAB] 즉흥 연주 가동 [MUSIC: heavy rock guitar riff]"
+        qwen = ModelManager.get_qwen()
+        out = qwen(f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_input}<|im_end|>\nassistant\n", max_new_tokens=512)
+        ai_text_raw = out[0]['generated_text'].split("assistant\n")[-1]
 
-    # 3. 파싱 및 음악 생성
+    # 파싱
     tab_match = re.search(r'\[TAB\](.*?)(\[|$)', ai_text_raw, re.DOTALL | re.IGNORECASE)
     music_match = re.search(r'\[MUSIC:(.*?)\]', ai_text_raw, re.IGNORECASE)
-    
     tab_display = tab_match.group(1).strip() if tab_match else "상세 악보 준비 중..."
     clean_text = re.sub(r'\[TAB\].*?(\[|$)', '', ai_text_raw, flags=re.DOTALL | re.IGNORECASE)
     clean_text = re.sub(r'\[MUSIC:.*?\]', '', clean_text, flags=re.IGNORECASE).strip()
-    music_p = music_match.group(1).strip() if music_match else "rock music"
+    
+    # 3. TTS 음성 생성 (복구)
+    tts_text = "\n".join(clean_text.split('\n')[:5])
+    voice_name = MEMBERS_VOICE.get(member_name, "ko-KR-SunHiNeural")
+    communicate = edge_tts.Communicate(re.sub(r'[\*\#\-\_\~\|]', '', tts_text), voice_name)
+    await communicate.save(voice_path)
 
-    # 음악 생성 파이프라인
+    # 4. 음악 생성 (길이 연장: 512 토큰 = 약 12초)
     music_gen = ModelManager.get_music()
-    try:
-        # CPU 환경을 위해 토큰 수를 줄여 타임아웃 방지
-        music_output = music_gen(music_p, forward_params={"max_new_tokens": 256})
-        audio_data = np.squeeze(music_output["audio"])
-        audio_int16 = (audio_data * 32767).astype(np.int16)
-        scipy.io.wavfile.write(music_path, music_output["sampling_rate"], audio_int16)
-    except:
-        music_path = None
+    music_p = music_match.group(1).strip() if music_match else "energetic rock guitar solo"
+    music_output = music_gen(music_p, forward_params={"max_new_tokens": 512})
+    audio_data = np.squeeze(music_output["audio"])
+    audio_int16 = (audio_data * 32767).astype(np.int16)
+    scipy.io.wavfile.write(music_path, music_output["sampling_rate"], audio_int16)
 
-    # TTS용 텍스트 (기존 구조 유지)
-    tts_text = "\n".join(clean_text.split('\n')[:5])
-    
-    return tts_text, None, music_path, tab_display # voice_path는 edge_tts 등에서 처리
+    return tts_text, voice_path, music_path, tab_display
 
-# 3. Gradio 인터페이스 (Minimal)
 with gr.Blocks() as demo:
-    i1 = gr.Textbox(visible=False); i2 = gr.Textbox(visible=False); i3 = gr.Textbox(visible=False)
-    o1 = gr.Textbox(visible=False); o2 = gr.Audio(visible=False); o3 = gr.Audio(visible=False); o4 = gr.Textbox(visible=False)
-    
-    btn = gr.Button("GO", visible=False)
-    btn.click(band_consulting, [i1, i2, i3], [o1, o2, o3, o4], api_name="predict")
+    i1 = gr.Textbox(); i2 = gr.Textbox(); i3 = gr.Textbox()
+    o1 = gr.Textbox(); o2 = gr.Audio(); o3 = gr.Audio(); o4 = gr.Textbox()
+    btn = gr.Button("GO"); btn.click(band_consulting, [i1, i2, i3], [o1, o2, o3, o4], api_name="predict")
 
 demo.queue().launch()