app.py

import os, asyncio, json, tempfile, websockets, pdfplumber
import gradio as gr
import openai
from dotenv import load_dotenv
import numpy as np
import wave

# ─── 0. 초기화 ───────────────────────────────────────────────
load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")
if not openai.api_key:
    raise RuntimeError("OPENAI_API_KEY 가 .env 에 없습니다!")

LANG = ["Korean","English","Japanese","Chinese",
        "Thai","Russian","Vietnamese","Spanish","French"]
VOICE = {l: ("nova" if l in ["Korean","Japanese","Chinese"] else "alloy")
         for l in LANG}
FOUR = ["English","Chinese","Thai","Russian"]
WS_URL = "wss://api.openai.com/v1/realtime"  # 올바른 엔드포인트로 수정

# ─── 1. 공통 GPT 번역 / TTS ─────────────────────────────────
async def gpt_translate(text, src, tgt):
    rsp = await openai.AsyncClient().chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role":"system",
                   "content":f"Translate {src} → {tgt}. Return only the text."},
                  {"role":"user","content":text}],
        temperature=0.3,max_tokens=2048)
    return rsp.choices[0].message.content.strip()

async def gpt_tts(text, lang):
    rsp = await openai.AsyncClient().audio.speech.create(
        model="tts-1", voice=VOICE[lang], input=text[:4096])
    tmp = tempfile.NamedTemporaryFile(delete=False,suffix=".mp3")
    tmp.write(rsp.content); tmp.close(); return tmp.name

# ─── 2. PDF 번역 ────────────────────────────────────────────
def translate_pdf(file, src, tgt):
    if not file: return "⚠️ PDF 업로드 필요", ""
    with pdfplumber.open(file.name) as pdf:
        text = "\n".join(p.extract_text() or "" for p in pdf.pages[:5]).strip()
    if not text:
        return "⚠️ 텍스트 추출 실패", ""
    return text, asyncio.run(gpt_translate(text, src, tgt))

# ─── 2-1. 오디오 번역 (탭1용) ────────────────────────────────
async def translate_audio_async(file, src, tgt):
    if not file: return "⚠️ 오디오 업로드 필요", "", None
    
    # STT: Whisper API 사용
    with open(file, 'rb') as audio_file:
        transcript = await openai.AsyncClient().audio.transcriptions.create(
            model="whisper-1",
            file=audio_file,
            language=src[:2].lower()  # 언어 코드 간소화
        )
    
    orig_text = transcript.text
    trans_text = await gpt_translate(orig_text, src, tgt)
    audio_path = await gpt_tts(trans_text, tgt)
    
    return orig_text, trans_text, audio_path

def translate_audio(file, src, tgt):
    return asyncio.run(translate_audio_async(file, src, tgt))

# ─── 3. 실시간 STT (Whisper API 사용) ──────────────────────────
async def process_audio_chunk(audio_data, src_lang):
    """오디오 청크를 처리하여 텍스트로 변환"""
    if audio_data is None:
        return ""
    
    try:
        # Gradio는 (sample_rate, audio_array) 튜플을 반환
        if isinstance(audio_data, tuple):
            sample_rate, audio_array = audio_data
            # numpy array를 WAV 파일로 변환
            import numpy as np
            import wave
            
            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
                with wave.open(tmp.name, 'wb') as wav_file:
                    wav_file.setnchannels(1)  # mono
                    wav_file.setsampwidth(2)  # 16-bit
                    wav_file.setframerate(sample_rate)
                    
                    # numpy array를 16-bit PCM으로 변환
                    if audio_array.dtype == np.float32 or audio_array.dtype == np.float64:
                        audio_array = (audio_array * 32767).astype(np.int16)
                    wav_file.writeframes(audio_array.tobytes())
                tmp_path = tmp.name
        else:
            # bytes 데이터인 경우
            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
                tmp.write(audio_data)
                tmp_path = tmp.name
        
        # Whisper API로 변환
        with open(tmp_path, 'rb') as audio_file:
            transcript = await openai.AsyncClient().audio.transcriptions.create(
                model="whisper-1",
                file=audio_file,
                language=src_lang[:2].lower()
            )
        
        os.unlink(tmp_path)  # 임시 파일 삭제
        return transcript.text
    except Exception as e:
        print(f"STT 오류: {e}")
        return ""

# ─── 4. Gradio 스트림 핸들러 (동기 버전) ─────────────────────
def realtime_single_sync(audio, src, tgt, state):
    """동기 버전의 실시간 단일 언어 번역"""
    if state is None:
        state = {"orig": "", "trans": "", "audio_buffer": [], "sample_rate": None}
    
    if audio is None:
        # 스트림 종료 시 남은 버퍼 처리
        if state["audio_buffer"] and state["sample_rate"]:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                # 버퍼의 오디오 합치기
                combined_audio = np.concatenate(state["audio_buffer"])
                audio_data = (state["sample_rate"], combined_audio)
                
                text = loop.run_until_complete(process_audio_chunk(audio_data, src))
                if text:
                    state["orig"] = state["orig"] + " " + text if state["orig"] else text
                    trans = loop.run_until_complete(gpt_translate(text, src, tgt))
                    state["trans"] = state["trans"] + " " + trans if state["trans"] else trans
            finally:
                loop.close()
            state["audio_buffer"] = []
        
        return state["orig"], state["trans"], state
    
    # 오디오 데이터 버퍼링
    if isinstance(audio, tuple):
        sample_rate, audio_array = audio
        state["sample_rate"] = sample_rate
        state["audio_buffer"].append(audio_array)
        
        # 버퍼가 충분히 쌓였을 때만 처리 (약 1-2초 분량)
        buffer_duration = len(np.concatenate(state["audio_buffer"])) / sample_rate
        if buffer_duration >= 1.5:  # 1.5초마다 처리
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            
            try:
                # 버퍼의 오디오 합치기
                combined_audio = np.concatenate(state["audio_buffer"])
                audio_data = (sample_rate, combined_audio)
                
                # STT
                text = loop.run_until_complete(process_audio_chunk(audio_data, src))
                if text:
                    state["orig"] = state["orig"] + " " + text if state["orig"] else text
                    
                    # 번역
                    trans = loop.run_until_complete(gpt_translate(text, src, tgt))
                    state["trans"] = state["trans"] + " " + trans if state["trans"] else trans
                
                # 버퍼 초기화
                state["audio_buffer"] = []
            finally:
                loop.close()
    
    return state["orig"], state["trans"], state

def realtime_four_sync(audio, src, state):
    """동기 버전의 실시간 4언어 번역"""
    if state is None:
        state = {"orig": "", "English": "", "Chinese": "", "Thai": "", "Russian": "", 
                 "audio_buffer": [], "sample_rate": None}
    
    if audio is None:
        # 스트림 종료 시 남은 버퍼 처리
        if state["audio_buffer"] and state["sample_rate"]:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                combined_audio = np.concatenate(state["audio_buffer"])
                audio_data = (state["sample_rate"], combined_audio)
                
                text = loop.run_until_complete(process_audio_chunk(audio_data, src))
                if text:
                    state["orig"] = state["orig"] + " " + text if state["orig"] else text
                    
                    tasks = []
                    for lang in FOUR:
                        tasks.append(gpt_translate(text, src, lang))
                    
                    translations = loop.run_until_complete(asyncio.gather(*tasks))
                    
                    for lang, trans in zip(FOUR, translations):
                        state[lang] = state[lang] + " " + trans if state[lang] else trans
            finally:
                loop.close()
            state["audio_buffer"] = []
        
        return (state["orig"], state["English"], state["Chinese"], 
                state["Thai"], state["Russian"], state)
    
    # 오디오 데이터 버퍼링
    if isinstance(audio, tuple):
        sample_rate, audio_array = audio
        state["sample_rate"] = sample_rate
        state["audio_buffer"].append(audio_array)
        
        # 버퍼가 충분히 쌓였을 때만 처리
        buffer_duration = len(np.concatenate(state["audio_buffer"])) / sample_rate
        if buffer_duration >= 1.5:  # 1.5초마다 처리
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            
            try:
                combined_audio = np.concatenate(state["audio_buffer"])
                audio_data = (sample_rate, combined_audio)
                
                # STT
                text = loop.run_until_complete(process_audio_chunk(audio_data, src))
                if text:
                    state["orig"] = state["orig"] + " " + text if state["orig"] else text
                    
                    # 4개 언어로 번역
                    tasks = []
                    for lang in FOUR:
                        tasks.append(gpt_translate(text, src, lang))
                    
                    translations = loop.run_until_complete(asyncio.gather(*tasks))
                    
                    for lang, trans in zip(FOUR, translations):
                        state[lang] = state[lang] + " " + trans if state[lang] else trans
                
                state["audio_buffer"] = []
            finally:
                loop.close()
    
    return (state["orig"], state["English"], state["Chinese"], 
            state["Thai"], state["Russian"], state)

# ─── 5. UI ──────────────────────────────────────────────────
with gr.Blocks(title="SMARTok Demo") as demo:
    with gr.Tabs():
        # 탭 1 – 오디오 번역
        with gr.TabItem("🎙️ 오디오"):
            src1 = gr.Dropdown(LANG, value="Korean", label="입력 언어")
            tgt1 = gr.Dropdown(LANG, value="English", label="출력 언어")
            aud1 = gr.Audio(sources=["microphone", "upload"], type="filepath")
            btn1 = gr.Button("번역")
            o1 = gr.Textbox(label="원문")
            t1 = gr.Textbox(label="번역")
            a1 = gr.Audio(label="TTS", type="filepath", autoplay=True)
            
            btn1.click(translate_audio, [aud1, src1, tgt1], [o1, t1, a1])

        # 탭 2 – PDF 번역
        with gr.TabItem("📄 PDF"):
            src2 = gr.Dropdown(LANG, value="Korean", label="입력 언어")
            tgt2 = gr.Dropdown(LANG, value="English", label="출력 언어")
            pdf = gr.File(file_types=[".pdf"])
            btn2 = gr.Button("번역")
            o2 = gr.Textbox(label="추출 원문", lines=15)
            t2 = gr.Textbox(label="번역 결과", lines=15)
            
            btn2.click(translate_pdf, [pdf, src2, tgt2], [o2, t2])

        # 탭 3 – 실시간 1언어
        with gr.TabItem("⏱️ 실시간 1"):
            src3 = gr.Dropdown(LANG, value="Korean", label="입력 언어")
            tgt3 = gr.Dropdown(LANG, value="English", label="출력 언어")
            mic3 = gr.Audio(sources=["microphone"], streaming=True)
            o3 = gr.Textbox(label="원문(실시간)", lines=8)
            t3 = gr.Textbox(label="번역(실시간)", lines=8)
            st3 = gr.State()
            
            # stream 메서드 수정
            mic3.stream(
                realtime_single_sync,
                inputs=[mic3, src3, tgt3, st3],
                outputs=[o3, t3, st3]
            )

        # 탭 4 – 실시간 4언어
        with gr.TabItem("🌏 실시간 4"):
            src4 = gr.Dropdown(LANG, value="Korean", label="입력 언어")
            mic4 = gr.Audio(sources=["microphone"], streaming=True)
            o4 = gr.Textbox(label="원문", lines=8)
            e4 = gr.Textbox(label="English", lines=8)
            c4 = gr.Textbox(label="Chinese(简体)", lines=8)
            th4 = gr.Textbox(label="Thai", lines=8)
            r4 = gr.Textbox(label="Russian", lines=8)
            st4 = gr.State()
            
            # stream 메서드 수정
            mic4.stream(
                realtime_four_sync,
                inputs=[mic4, src4, st4],
                outputs=[o4, e4, c4, th4, r4, st4]
            )

demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)