app.py

"""
Gradio UI that calls the remote animation server.
ANIM_API_URL 환경변수를 설정해 서버 주소를 주입하세요.
   ANIM_API_URL=http://211.233.58.201:7788/
"""

import os, logging
from datetime import datetime

import gradio as gr
import httpx
from gradio_client import Client, handle_file

# ──────────────────────── 로깅 ────────────────────────── #
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s"
)
log = logging.getLogger(__name__)

# ──────────────── 서버 URL 결정 ─────────────────────── #
DEFAULT_URL = "http://127.0.0.1:7862/"
REMOTE_URL  = "http://211.233.58.201:7862/"
API_URL     = os.getenv("ANIM_API_URL", REMOTE_URL)

if "127.0.0.1" in API_URL and os.getenv("HF_SPACE") == "true":
    raise RuntimeError(
        "HF Space 컨테이너 안에서는 공인 IP나 도메인을 ANIM_API_URL 로 지정해야 합니다."
    )

# ──────────────── 헬스체크 함수 ─────────────────────── #
TIMEOUT = httpx.Timeout(connect=30.0, read=120.0, write=120.0, pool=30.0)

def test_api_connection():
    now = datetime.now().strftime("%H:%M:%S")
    try:
        resp  = httpx.get(f"{API_URL.rstrip('/')}/healthz", timeout=TIMEOUT)
        ready = resp.json().get("ready", False)
        msg   = f"[{now}] 서버 연결 성공 ✅ (ready={ready})"
        log.info(msg)
        return True, msg
    except Exception as e:
        msg = f"[{now}] 서버 연결 실패 ❌ : {e}"
        log.error(msg)
        return False, msg

# ──────────────── 애니메이션 생성 ───────────────────── #
def generate_animation(image, audio, guidance_scale, steps, progress=gr.Progress()):
    start = datetime.now().strftime("%H:%M:%S")
    logs  = [f"[{start}] 요청 시작"]

    try:
        if image is None or audio is None:
            raise ValueError("이미지와 오디오를 모두 업로드하세요.")

        progress(0.05, desc="파일 준비")
        client = Client(API_URL)

        progress(0.15, desc="서버 호출 중… (수 분 소요 가능)")
        result = client.predict(
            image_path=handle_file(image),
            audio_path=handle_file(audio),
            guidance_scale=guidance_scale,
            steps=steps,
            api_name="/generate_animation"
        )

        progress(0.95, desc="결과 정리")
        
        # 결과 처리 - dict 형태 처리 추가
        def extract_video_path(obj):
            """비디오 객체에서 경로 추출"""
            if isinstance(obj, str):
                return obj
            elif isinstance(obj, dict):
                # Gradio의 FileData dict 처리
                if 'video' in obj:
                    return obj['video']  # {'video': '경로', 'subtitles': None} 형태 처리
                elif 'path' in obj:
                    return obj['path']
                elif 'url' in obj:
                    return obj['url']
                elif 'name' in obj:
                    return obj['name']
                else:
                    log.warning(f"Unexpected dict structure: {obj.keys()}")
                    return None
            else:
                log.warning(f"Unexpected type: {type(obj)}")
                return None
        
        if isinstance(result, (list, tuple)) and len(result) >= 2:
            anim_path = extract_video_path(result[0])
            comp_path = extract_video_path(result[1])
            
            if anim_path and comp_path:
                logs.append(f"[{datetime.now().strftime('%H:%M:%S')}] 성공")
                return anim_path, comp_path, "\n".join(logs)
            else:
                raise RuntimeError(f"비디오 경로 추출 실패: {result}")
        else:
            raise RuntimeError(f"예상치 못한 반환 형식: {type(result)}")

    except Exception as e:
        logs.append(f"[{datetime.now().strftime('%H:%M:%S')}] 오류: {e}")
        log.error(f"Animation generation error: {e}", exc_info=True)
        return None, None, "\n".join(logs)

# ─────────────────────── UI 정의 ────────────────────── #
with gr.Blocks(title="Animation Generator Client") as demo:
    gr.Markdown("# 🎬 Animation Generator – Client UI")

    # 서버 상태 체크
    status_box = gr.Textbox(label="API 상태", interactive=False)
    test_btn   = gr.Button("서버 연결 테스트")
    test_btn.click(test_api_connection, outputs=[status_box, status_box])

    gr.Markdown("---")

    with gr.Row():
        with gr.Column():
            img_in  = gr.Image(type="filepath", label="Portrait Image")
            aud_in  = gr.Audio(type="filepath", label="Driving Audio")
            scale   = gr.Slider(1, 10, value=3.0, step=0.1, label="Guidance Scale")
            steps   = gr.Slider(5, 30, value=10, step=1, label="Inference Steps")
            gen_btn = gr.Button("🚀 Generate")
        with gr.Column():
            anim_out = gr.Video(label="Animation Result")
            comp_out = gr.Video(label="Side-by-Side")

    with gr.Accordion("실행 로그", open=False):
        log_out = gr.Textbox(label="Logs", lines=12, max_lines=20, interactive=False)

    gen_btn.click(
        generate_animation,
        inputs=[img_in, aud_in, scale, steps],
        outputs=[anim_out, comp_out, log_out]
    )

# ─────────────────────── 실행 ───────────────────────── #
if __name__ == "__main__":
    demo.queue(max_size=4).launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_api=False
    )