app.py

import os
import logging
import re
import gradio as gr
import google.generativeai as genai
from dotenv import load_dotenv
from typing import Dict, Optional
from PIL import Image

load_dotenv()

# 로깅 설정
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Gemini API 설정
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
if not GEMINI_API_KEY:
    logger.error("Gemini API 키가 .env 파일에 설정되지 않았습니다.")
else:
    try:
        genai.configure(api_key=GEMINI_API_KEY)
        logger.info("Gemini API 키가 성공적으로 설정되었습니다.")
    except Exception as e:
        logger.error(f"Gemini API 키 설정 중 오류 발생: {e}")
        GEMINI_API_KEY = ""

# 번역 함수
def translate_with_gemini(text_to_translate: str) -> str:
    """Gemini API를 사용하여 한국어 텍스트를 영어로 번역합니다."""
    if not text_to_translate or not GEMINI_API_KEY:
        return text_to_translate

    try:
        model = genai.GenerativeModel('gemini-2.0-flash')
        prompt = f"""Translate the following Korean text into natural-sounding English.
Respond ONLY with the translated English text and nothing else.

Korean Text:
{text_to_translate}

English Translation:"""

        response = model.generate_content(
            prompt,
            generation_config=genai.types.GenerationConfig(
                temperature=0.2,
                max_output_tokens=len(text_to_translate) * 3
            )
        )

        translated_text = response.text.strip().strip('"')
        logger.info(f"번역 성공: '{text_to_translate}' -> '{translated_text}'")
        return translated_text

    except Exception as e:
        logger.error(f"번역 중 오류: {e}")
        return text_to_translate

# 프롬프트 생성 함수
def generate_prompt_with_gemini(
    translated_item_2: str, 
    translated_item_3: Optional[str] = None, 
    translated_custom: Optional[str] = None
) -> str:
    """Gemini API를 사용해 고급 프롬프트를 생성합니다."""
    if not GEMINI_API_KEY:
        return "오류: Gemini API 키가 없어 프롬프트 생성을 진행할 수 없습니다."

    try:
        # 아이템 설명 조합
        item_description = f"(From image #2)"
        if translated_item_3:
            item_description += f" and (From image #3)"

        # 프롬프트 생성 요청
        prompt_request = f"""Create a precise Midjourney prompt that:
- Exactly recreates (Reference Image #1) 
- Transforms ensemble with {item_description}
- Preserves original emotional essence

Key Requirements:
- Match original pose and body positioning
- Maintain specific lighting and mood
- Seamlessly integrate new styling
- Ensure hyperrealistic transformation

Contextual Notes: {translated_custom if translated_custom else 'Preserve intrinsic image character'}

Render with cinematic precision, 8k ultra-detail."""

        model = genai.GenerativeModel('gemini-2.0-flash')

        response = model.generate_content(
            prompt_request,
            generation_config=genai.types.GenerationConfig(
                temperature=0.7,
                top_p=0.95,
                max_output_tokens=8192
            )
        )

        # 프롬프트 추출 및 정제
        enhanced_prompt = response.text.strip()
        
        # 필수 파라미터 추가
        required_params = "--ar 9:16 --seed 123456 --q 3 --v 5.2 --style raw"
        if required_params not in enhanced_prompt:
            enhanced_prompt = f"/imagine Precise photorealistic recreation of (Reference Image #1) with meticulous transformation:\n- Exact pose and ethereal mood preserved\n- Seamless ensemble replacement from (image #2) and (image #3)\n- Original silhouette and body positioning maintained\n- Dramatic side-lighting matching reference image\n- Hyper-detailed fabric textures and subtle accessory integration\n- Soft, contemplative emotional landscape\n- Invisible stylistic transitions\n\nCinematic rendering focusing on form, texture, and emotional resonance, 8k ultra-detailed, professional studio lighting, artstation aesthetic. {required_params}"

        logger.info(f"생성된 프롬프트: {enhanced_prompt}")
        return enhanced_prompt

    except Exception as e:
        logger.exception("프롬프트 생성 중 오류:")
        return f"오류: 프롬프트 생성 실패 ({str(e)})"

# 최종 프롬프트 생성 함수
def generate_final_prompt(
    reference_image, 
    item_image_2, 
    item_image_3, 
    item_description_2_ko, 
    item_description_3_ko, 
    custom_prompt_ko
):
    # 필수 입력값 검증
    if not reference_image or not item_image_2:
        return "오류: 레퍼런스 이미지(#1)와 아이템 이미지(#2)를 업로드해주세요."
    
    if not item_description_2_ko:
        return "오류: 아이템(#2) 설명을 입력해주세요."

    if not GEMINI_API_KEY:
        return "오류: Gemini API 키가 설정되지 않았습니다."

    # 번역 수행
    logger.info("입력된 한국어 설명을 영어로 번역 시작...")
    translated_item_2 = translate_with_gemini(item_description_2_ko)
    
    # #3 아이템 선택적 처리
    translated_item_3 = None
    if item_image_3 and item_description_3_ko:
        translated_item_3 = translate_with_gemini(item_description_3_ko)

    # 커스텀 프롬프트 선택적 처리
    translated_custom = None
    if custom_prompt_ko:
        translated_custom = translate_with_gemini(custom_prompt_ko)

    # 프롬프트 생성
    try:
        generated_prompt = generate_prompt_with_gemini(
            translated_item_2,
            translated_item_3,
            translated_custom
        )

        return generated_prompt

    except Exception as e:
        logger.exception("최종 프롬프트 생성 과정에서 예외 발생:")
        return f"오류: 프롬프트 생성 실패 ({str(e)})"

# Gradio 인터페이스 구성
def create_app():
    with gr.Blocks(title="가상 피팅 스튜디오") as demo:
        gr.Markdown("# 가상 피팅 스튜디오")
        gr.Markdown("""
        레퍼런스 이미지를 기반으로 스타일을 완벽하게 재해석합니다.
        - **레퍼런스 이미지(#1):** 포즈, 분위기, 조명을 그대로 유지할 이미지
        - **아이템 이미지(#2):** 필수 스타일링 요소
        - **아이템 이미지(#3):** 선택적 액세서리 및 디테일
        - **설명 입력:** 각 이미지의 스타일링 특성을 한국어로 설명 (Gemini가 영어로 번역)
        - **주의:** Gemini API(gemini-2.0-flash)를 사용하여 고급 프롬프트 생성
        """)

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("## 입력 섹션")
                with gr.Row():
                    reference_image = gr.Image(label="레퍼런스 이미지 (#1)", type="pil", sources=["upload"])
                    item_image_2 = gr.Image(label="아이템 이미지 (#2)", type="pil", sources=["upload"])
                    item_image_3 = gr.Image(label="아이템 이미지 (#3, 선택)", type="pil", sources=["upload"])
                
                with gr.Row():
                    item_description_2_ko = gr.Textbox(
                        label="아이템설명(#2) (한국어)",
                        placeholder="예: 베이지 트위드 코트와 미니멀한 액세서리",
                        lines=1, interactive=True
                    )
                    item_description_3_ko = gr.Textbox(
                        label="아이템설명(#3) (한국어, 선택)",
                        placeholder="예: 실버 목걸이와 가죽 장갑",
                        lines=1, interactive=True
                    )
                
                custom_prompt_ko = gr.Textbox(
                    label="추가 스타일링 설명 (한국어, 선택 사항)",
                    placeholder="예: 파리지앵 세련됨, 차가운 겨울 아침 분위기",
                    lines=2, interactive=True
                )
                
                prompt_btn = gr.Button("프롬프트 생성 (번역 포함)", variant="primary")

            with gr.Column(scale=1):
                gr.Markdown("## 출력 섹션")
                prompt_output = gr.Textbox(
                    label="생성된 Midjourney 프롬프트 (영문)",
                    lines=15,
                    interactive=False
                )

        prompt_btn.click(
            fn=generate_final_prompt,
            inputs=[reference_image, item_image_2, item_image_3, 
                    item_description_2_ko, item_description_3_ko, 
                    custom_prompt_ko],
            outputs=[prompt_output]
        )

    return demo

if __name__ == "__main__":
    if not GEMINI_API_KEY:
        print("경고: Gemini API 키가 설정되지 않았습니다. 번역 및 프롬프트 생성이 제한될 수 있습니다.")
    
    print("사용할 Gemini 모델: gemini-2.0-flash (API 키 환경에서 사용 가능한지 확인 필요)")

    app = create_app()
    app.queue()
    app.launch()