app.py

import os, re, uuid, torch, scipy.io.wavfile, edge_tts, asyncio
import numpy as np
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from groq import Groq

# 모델 관리 싱글톤
class ModelManager:
    _llm_pipeline = None
    _music_pipeline = None
    _groq_client = 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)
            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):
        key = os.getenv("GROQ_API_KEY")
        if cls._groq_client is None and key:
            cls._groq_client = Groq(api_key=key)
        return cls._groq_client

    @classmethod
    def get_music(cls):
        if cls._music_pipeline is None:
            cls._music_pipeline = pipeline("text-to-audio", "facebook/musicgen-small", device="cpu")
        return cls._music_pipeline

# 멤버별 고유 성격/말투 페르소나
PERSONA_MAP = {
    "서윤 (Korea)": "서정적이면서도 강렬한 K-Rock 리더. 진지하고 철학적이며 따뜻하게 조언함.",
    "Chloe (USA)": "자유분방한 캘리포니아 펑크 락스타. 거침없고 쿨하며 에너제틱한 말투.",
    "Naomi (Japan)": "완벽주의 J-Rock 기타리스트. 섬세하고 예의 바르며 음악 이론에 날카로움.",
    "Beatrice (Brazil)": "열정적인 삼바 락커. 긍정적이고 리듬감이 넘치며 인생의 즐거움을 강조함.",
    "Elena (Spain)": "드라마틱한 감성의 소유자. 플라멩코의 정열과 락의 파괴력을 동시에 가짐.",
    "Amira (Egypt)": "신비로운 오리엔탈 락커. 깊은 지혜와 고전적인 무게감을 담아 조언함.",
    "Liwei (China)": "전통과 현대의 조화를 중시하는 퓨전 락커. 절제되고 힘 있는 말투.",
    "Sophie (France)": "예술적 자존심이 강한 아방가르드 락커. 시적이고 세련된 표현을 즐겨 씀."
}

# 언어별 멤버/보이스 매핑
LANG_MEMBER_MAP = {
    "Korean": {"name": "서윤 (Korea)", "voice": "ko-KR-SunHiNeural"},
    "English": {"name": "Chloe (USA)", "voice": "en-US-AriaNeural"},
    "Japanese": {"name": "Naomi (Japan)", "voice": "ja-JP-NanamiNeural"},
    "Portuguese": {"name": "Beatrice (Brazil)", "voice": "pt-BR-FranciscaNeural"},
    "Spanish": {"name": "Elena (Spain)", "voice": "es-ES-ElviraNeural"},
    "Arabic": {"name": "Amira (Egypt)", "voice": "ar-EG-SalmaNeural"},
    "Chinese": {"name": "Liwei (China)", "voice": "zh-CN-XiaoxiaoNeural"},
    "French": {"name": "Sophie (France)", "voice": "fr-FR-DeniseNeural"}
}

async def band_consulting(user_input, selected_lang, g_inst, b_inst, d_inst, chords):
    try:
        req_id = str(uuid.uuid4())[:8]
        voice_path = f"/tmp/v_{req_id}.mp3"
        music_path = f"/tmp/m_{req_id}.wav"
        
        m_info = LANG_MEMBER_MAP.get(selected_lang, LANG_MEMBER_MAP["Korean"])
        persona = PERSONA_MAP.get(m_info['name'], "열정적인 락스타")
        
        # 본문 요약 및 상세 탭 분리 지시 프롬프트
        system_prompt = f"""You are the rock star '{m_info['name']}'. Respond ONLY in {selected_lang}.
        Persona: {persona}.

        CRITICAL RULES:
        1. MAIN ADVICE: Provide exactly 3 to 5 deep, soulful sentences for the main chat.
        2. [TAB] SECTION: Put ALL long explanations, music theory, Drum patterns (e.g., H|x-x-x-x| S|--o---o-|), and detailed Guitar Solo tabs here. This must be very technical and professional.
        3. [TRANSLATION]: English translation of your 3-5 sentence MAIN ADVICE only.
        4. [MUSIC]: English prompt for MusicGen reflecting: Guitar:{g_inst}, Bass:{b_inst}, Drums:{d_inst}, Chords:{chords}
        """
        
        ai_text_raw = ""
        groq_client = ModelManager.get_groq()
        if groq_client:
            try:
                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 = res.choices[0].message.content
            except Exception as e:
                print(f"Groq Error: {e}")

        if not ai_text_raw:
            qwen = ModelManager.get_qwen()
            input_t = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_input}<|im_end|>\nassistant\n"
            out = qwen(input_t, max_new_tokens=1024)
            ai_text_raw = out[0]['generated_text'].split("assistant\n")[-1]

        # 파싱 로직
        tab_match = re.search(r'\[TAB\](.*?)(\[|$)', ai_text_raw, re.DOTALL | re.IGNORECASE)
        music_match = re.search(r'\[MUSIC:(.*?)\]', ai_text_raw, re.IGNORECASE)
        trans_match = re.search(r'\[TRANSLATION\](.*?)(\[|$)', ai_text_raw, re.DOTALL | re.IGNORECASE)
        
        tab_display = tab_match.group(1).strip() if tab_match else "No Detailed Data"
        eng_translation = trans_match.group(1).strip() if trans_match else ""

        # 음성에서 번역문/탭/음악프롬프트 완벽 제거
        speech_text = re.sub(r'\[TAB\].*?(\[|$)', '', ai_text_raw, flags=re.DOTALL | re.IGNORECASE)
        speech_text = re.sub(r'\[MUSIC:.*?\]', '', speech_text, flags=re.IGNORECASE)
        speech_text = re.sub(r'\[TRANSLATION\].*?(\[|$)', '', speech_text, flags=re.DOTALL | re.IGNORECASE).strip()
        
        # TTS 생성
        tts_text = re.sub(r'[\*\#\-\_\~\|]', '', speech_text)
        communicate = edge_tts.Communicate(tts_text, m_info["voice"])
        await communicate.save(voice_path)

        # MusicGen 음악 생성
        music_gen = ModelManager.get_music()
        music_p = music_match.group(1).strip() if music_match else "rock music"
        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)

        return speech_text, voice_path, music_path, tab_display, eng_translation

    except Exception as e:
        print(f"Final Error: {e}")
        return f"밴드 시스템 에러: {str(e)}", None, None, "No Data", "Error occurred"

with gr.Blocks() as demo:
    inputs = [gr.Textbox(visible=False) for _ in range(6)]
    outputs = [
        gr.Textbox(visible=False), 
        gr.Audio(visible=False), 
        gr.Audio(visible=False), 
        gr.Textbox(visible=False), 
        gr.Textbox(visible=False)
    ]
    btn = gr.Button("API", visible=False)
    btn.click(band_consulting, inputs, outputs, api_name="predict")

demo.queue().launch()