Update app.py

app.py

@@ -1,90 +1,118 @@
 import os
 import re
-import requests
 import uuid
-import scipy.io.wavfile
-import numpy as np
 import torch
-from transformers import pipeline
+import numpy as np
+import scipy.io.wavfile
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from groq import Groq # pip install groq
 
-# Hugging Face 토큰 및 모델 설정
-HF_TOKEN = os.getenv("HF_TOKEN")
-MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
+# 1. 모델 싱글톤 관리 클래스 (메모리 효율화)
+class ModelManager:
+    _llm_pipeline = None
+    _music_pipeline = None
+    _groq_client = None
 
-# 음악 모델 로드 (기존 유지)
-try:
-    music_synthesiser = pipeline("text-to-audio", "facebook/musicgen-small")
-except Exception as e:
-    print(f"Music Model Load Error: {e}")
-    music_synthesiser = None
+    @classmethod
+    def get_qwen(cls):
+        if cls._llm_pipeline is None:
+            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
+            )
+            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)
+        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. 핵심 비즈니스 로직
 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"
     
-    # 1. Qwen용 시스템 프롬프트 구성 (Gemini보다 더 강력하게 지시)
-    system_instruction = f"""
-    You are '{member_name}', a professional rock band member. 
-    [Language Rules] Respond ONLY in {lang_code}.
-    [Output Format] 
-    1. Short advice (under 5 lines).
-    2. Detailed sheet music/advice inside [TAB] tags.
-    3. Music generation prompt inside [MUSIC: English Prompt] tags.
-    """
-
-    try:
-        # 2. Hugging Face Inference API 호출
-        api_url = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
-        headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-        
-        prompt = f"<|im_start|>system\n{system_instruction}<|im_end|>\n<|im_start|>user\n{user_input}<|im_end|>\n<|im_start|>assistant\n"
-        
-        payload = {
-            "inputs": prompt,
-            "parameters": {"max_new_tokens": 1024, "temperature": 0.7, "top_p": 0.9}
-        }
+    # 시스템 프롬프트 설정
+    system_prompt = f"You are {member_name}, a rock star. Talk in {lang_code}. Format: 5 lines max + [TAB] details + [MUSIC: prompt]."
+    
+    ai_text_raw = ""
+    
+    # 엔진 선택 (Groq 우선 -> 실패 시 로컬 Qwen)
+    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",
+            )
+            ai_text_raw = chat_completion.choices[0].message.content
+        except Exception as e:
+            print(f"Groq Error, switching to Qwen: {e}")
 
-        response = requests.post(api_url, headers=headers, json=payload, timeout=20)
-        res_json = response.json()
+    # Groq 실패 시 혹은 미설정 시 로컬 Qwen 사용
+    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]"
 
-        # 응답이 비어있거나 에러일 경우 'Improv Mode' 발동
-        if not res_json or 'error' in res_json:
-            ai_text_raw = "이봐 리더, 지금은 말보다 연주가 필요한 순간이야! 내 즉흥 연주를 들어봐. [TAB] 즉흥 연주 모드 가동 [MUSIC: aggressive hard rock guitar solo]"
-        else:
-            # Qwen은 입력 프롬프트를 포함해서 반환할 수 있으므로 assistant 응답만 추출
-            full_text = res_json[0]['generated_text']
-            ai_text_raw = full_text.split("assistant\n")[-1].strip()
+    # 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. 파싱 로직 (Regex 강화)
-        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()
-        
-        # 텍스트가 비어있으면 Fallback
-        if len(clean_text) < 5:
-            clean_text = "말이 필요 없는 끝내주는 라인이 생각났어. 바로 들어보자고!"
+    # 음악 생성 파이프라인
+    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
 
-        # 4. 음악 생성 (MusicGen 연동)
-        music_p = music_match.group(1).strip() if music_match else "energetic rock riff"
-        
-        if music_synthesiser:
-            # 안전 필터 단어 우회 (Vocal 등 차단 방지)
-            safe_p = re.sub(r'(laugh|vocal|voice|human)', 'rock guitar', music_p.lower())
-            music_output = music_synthesiser(safe_p, forward_params={"max_new_tokens": 512})
-            
-            # 오디오 처리 (Float32 -> Int16)
-            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)
-        else:
-            music_path = None
+    # TTS용 텍스트 (기존 구조 유지)
+    tts_text = "\n".join(clean_text.split('\n')[:5])
+    
+    return tts_text, None, music_path, tab_display # voice_path는 edge_tts 등에서 처리
 
-        return clean_text, None, music_path, tab_display # voice_path는 필요시 edge_tts 추가
+# 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")
 
-    except Exception as e:
-        print(f"Critical System Error: {e}")
-        return "시스템 과부하인가? 하지만 락은 멈추지 않아! [MUSIC: basic rock drum beat]", None, None, "Retry Mode"
+demo.queue().launch()