app.py

import os
import tempfile
from PIL import Image
import gradio as gr
import logging
import re
from io import BytesIO

from google import genai
from google.genai import types

# 환경변수 로드
from dotenv import load_dotenv
load_dotenv()

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

def save_binary_file(file_name, data):
    with open(file_name, "wb") as f:
        f.write(data)

def translate_prompt_to_english(prompt):
    """
    입력된 프롬프트에 한글이 포함되어 있으면 Gemini‑2.0‑flash 모델을 사용하여 영어로 번역합니다.
    한글이 없으면 원본 프롬프트를 그대로 반환합니다.
    """
    if not re.search("[가-힣]", prompt):
        return prompt
    try:
        api_key = os.environ.get("GEMINI_API_KEY")
        if not api_key:
            logger.error("Gemini API 키가 설정되지 않았습니다.")
            return prompt
        client = genai.Client(api_key=api_key)
        translation_prompt = f"Translate the following Korean text to English:\n\n{prompt}"
        logger.info(f"Translation prompt: {translation_prompt}")
        response = client.models.generate_content(
            model="gemini-2.0-flash",
            contents=[translation_prompt],
            config=types.GenerateContentConfig(
                response_modalities=['Text'],
                temperature=0.2,
                top_p=0.95,
                top_k=40,
                max_output_tokens=512
            )
        )
        translated_text = ""
        for part in response.candidates[0].content.parts:
            if hasattr(part, 'text') and part.text:
                translated_text += part.text
        if translated_text.strip():
            logger.info(f"Translated text: {translated_text.strip()}")
            return translated_text.strip()
        else:
            logger.warning("번역 결과가 없습니다. 원본 프롬프트 사용")
            return prompt
    except Exception as e:
        logger.exception("번역 중 오류 발생:")
        return prompt

def preprocess_prompt(prompt, image1, image2, image3, bg_options=None):
    """
    프롬프트를 처리하고 기능 명령을 해석
    """
    has_img1 = image1 is not None
    has_img2 = image2 is not None
    has_img3 = image3 is not None

    if "#1" in prompt and not has_img1:
        prompt = prompt.replace("#1", "첫 번째 이미지(없음)")
    else:
        prompt = prompt.replace("#1", "첫 번째 이미지")

    if "#2" in prompt and not has_img2:
        prompt = prompt.replace("#2", "두 번째 이미지(없음)")
    else:
        prompt = prompt.replace("#2", "두 번째 이미지")

    if "#3" in prompt and not has_img3:
        prompt = prompt.replace("#3", "세 번째 이미지(없음)")
    else:
        prompt = prompt.replace("#3", "세 번째 이미지")

    if "1. 이미지 변경" in prompt:
        desc_match = re.search(r'#1을 "(.*?)"으로 바꿔라', prompt)
        if desc_match:
            description = desc_match.group(1)
            prompt = f"첫 번째 이미지를 {description}으로 변경해주세요. 원본 이미지의 주요 내용은 유지하되 새로운 스타일과 분위기로 재해석해주세요."
        else:
            prompt = "첫 번째 이미지를 창의적으로 변형해주세요. 더 생생하고 예술적인 버전으로 만들어주세요."

    elif "2. 글자지우기" in prompt:
        text_match = re.search(r'#1에서 "(.*?)"를 지워라', prompt)
        if text_match:
            text_to_remove = text_match.group(1)
            prompt = f"첫 번째 이미지에서 '{text_to_remove}' 텍스트를 찾아 자연스럽게 제거해주세요. 텍스트가 있던 부분을 배경과 조화롭게 채워주세요."
        else:
            prompt = "첫 번째 이미지에서 모든 텍스트를 찾아 자연스럽게 제거해주세요. 깔끔한 이미지로 만들어주세요."

    elif "4. 옷바꾸기" in prompt:
        prompt = "첫 번째 이미지의 인물 의상을 두 번째 이미지의 의상으로 변경해주세요. 의상의 스타일과 색상은 두 번째 이미지를 따르되, 신체 비율과 포즈는 첫 번째 이미지를 유지해주세요."

    elif "5. 배경바꾸기" in prompt or "배경 변경" in prompt:
        # 배경 옵션이 있으면 그것을 기반으로 프롬프트 생성
        if bg_options and all([bg_options.get("product_name"), bg_options.get("background_type"), bg_options.get("background_style")]):
            product_name = bg_options.get("product_name", "")
            product_category = bg_options.get("product_category", "일반 상품")
            background_type = bg_options.get("background_type", "자연 배경")
            background_style = bg_options.get("background_style", "숲")
            lighting = bg_options.get("lighting", "자연광")
            texture = bg_options.get("texture", "")
            color_tone = bg_options.get("color_tone", "")
            mood = bg_options.get("mood", "")
            
            # 상세 프롬프트 생성
            prompt = f"{product_name} 상품 이미지의 배경을 {background_type} 중 {background_style}으로 자연스럽게 변경해주세요. "
            
            # 추가 세부 사항 포함
            if lighting:
                prompt += f"{lighting} 조명을 사용하고, "
            if texture:
                prompt += f"{texture} 질감을 표현하며, "
            if color_tone:
                prompt += f"{color_tone} : 전체 색감으로 처리하고, "
            if mood:
                prompt += f"{mood}한 분위기로 "
            
            prompt += f"연출해주세요. {product_category} 상품의 색상, 질감, 형태는 완벽하게 보존하고 고품질 상품 사진 스타일로 제작해주세요."
        else:
            prompt = "첫 번째 이미지의 배경을 두 번째 이미지의 배경으로 변경해주세요. 첫 번째 이미지의 주요 피사체는 유지하고, 두 번째 이미지의 배경과 조화롭게 합성해주세요."

    elif "상품합성" in prompt:
        prompt = "#1 상품 이미지에 #2 이미지 또는 #3 이미지를 [모든 이미지의 주요 요소를 포함하고, 특히 상품이 돋보이도록 조화롭게 통합하여]자연스럽게 합성하라."

    prompt += " 이미지를 생성해주세요."
    return prompt

def generate_with_images(prompt, images):
    """
    API 호출을 통해 이미지를 생성하고 결과 이미지를 반환합니다.
    """
    try:
        api_key = os.environ.get("GEMINI_API_KEY")
        if not api_key:
            return None, "API 키가 설정되지 않았습니다. 환경변수를 확인해주세요."

        client = genai.Client(api_key=api_key)
        logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt}")

        contents = [prompt]
        for idx, img in enumerate(images, 1):
            if img is not None:
                contents.append(img)
                logger.info(f"이미지 #{idx} 추가됨")

        response = client.models.generate_content(
            model="gemini-2.0-flash-exp-image-generation",
            contents=contents,
            config=types.GenerateContentConfig(
                response_modalities=['Text', 'Image'],
                temperature=1,
                top_p=0.95,
                top_k=40,
                max_output_tokens=8192
            )
        )

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            temp_path = tmp.name
            result_text = ""
            image_found = False
            for part in response.candidates[0].content.parts:
                if hasattr(part, 'text') and part.text:
                    result_text += part.text
                    logger.info(f"응답 텍스트: {part.text}")
                elif hasattr(part, 'inline_data') and part.inline_data:
                    save_binary_file(temp_path, part.inline_data.data)
                    image_found = True
                    logger.info("응답에서 이미지 추출 성공")
            if not image_found:
                return None, f"API에서 이미지를 생성하지 못했습니다. 응답 텍스트: {result_text}"
            result_img = Image.open(temp_path)
            if result_img.mode == "RGBA":
                result_img = result_img.convert("RGB")
            return result_img, f"이미지가 성공적으로 생성되었습니다. {result_text}"
    except Exception as e:
        logger.exception("이미지 생성 중 오류 발생:")
        return None, f"오류 발생: {str(e)}"

def process_images_with_prompt(image1, image2, image3, prompt, bg_options=None):
    """
    3개의 이미지와 프롬프트를 처리하여 최종 영어 프롬프트(final_prompt)를 생성한 후,
    API를 호출하여 결과 이미지를 반환합니다.
    """
    try:
        images = [image1, image2, image3]
        valid_images = [img for img in images if img is not None]
        if not valid_images:
            return None, "적어도 하나의 이미지를 업로드해주세요.", ""

        if prompt and prompt.strip():
            processed_prompt = preprocess_prompt(prompt, image1, image2, image3, bg_options)
            if re.search("[가-힣]", processed_prompt):
                final_prompt = translate_prompt_to_english(processed_prompt)
            else:
                final_prompt = processed_prompt
        else:
            if len(valid_images) == 1:
                final_prompt = "Please creatively transform this image into a more vivid and artistic version."
                logger.info("Default prompt generated for single image")
            elif len(valid_images) == 2:
                final_prompt = "Please seamlessly composite these two images, integrating their key elements harmoniously into a single image."
                logger.info("Default prompt generated for two images")
            else:
                final_prompt = "Please creatively composite these three images, combining their main elements into a cohesive and natural scene."
                logger.info("Default prompt generated for three images")

        result_img, status = generate_with_images(final_prompt, valid_images)
        return result_img, status, final_prompt
    except Exception as e:
        logger.exception("이미지 처리 중 오류 발생:")
        return None, f"오류 발생: {str(e)}", prompt

def process_and_show_prompt(image1, image2, image3, prompt, product_name, product_category, 
                            background_type, background_style, lighting, texture, color_tone, mood):
    try:
        # 배경 옵션을 딕셔너리로 구성
        bg_options = {
            "product_name": product_name,
            "product_category": product_category,
            "background_type": background_type,
            "background_style": background_style,
            "lighting": lighting,
            "texture": texture,
            "color_tone": color_tone,
            "mood": mood
        }
            
        result_img, status, final_prompt = process_images_with_prompt(image1, image2, image3, prompt, bg_options)
        
        # 선택된 옵션을 문자열로 구성
        options_summary = f"상품명: {product_name}\n"
        options_summary += f"제품군: {product_category}\n"
        options_summary += f"배경 유형: {background_type}\n"
        options_summary += f"세부 스타일: {background_style}\n"
        
        if lighting:
            options_summary += f"조명: {lighting}\n"
        if texture:
            options_summary += f"배경 질감: {texture}\n"
        if color_tone:
            options_summary += f"색감: {color_tone}\n"
        if mood:
            options_summary += f"분위기: {mood}\n"
        
        # 한국어와 영어 프롬프트 비교를 위해 원본 프롬프트 확인
        kr_prompt = preprocess_prompt(prompt, image1, image2, image3, bg_options)
        
        return result_img, status, final_prompt, options_summary, kr_prompt
    except Exception as e:
        logger.exception("처리 중 오류 발생:")
        return None, f"오류 발생: {str(e)}", prompt, "", ""

# 선택 옵션 표시/숨김 처리 함수
def toggle_background_options(selection):
    if selection == "배경바꾸기":
        return (gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), 
                gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), 
                gr.update(visible=True), gr.update(visible=True), gr.update(visible=True),
                gr.update(visible=True))
    else:
        return (gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), 
                gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), 
                gr.update(visible=False), gr.update(visible=False), gr.update(visible=False),
                gr.update(visible=False))

# 배경 유형에 따른 세부 스타일 옵션 업데이트
def update_background_style_options(background_type):
    if background_type == "자연 배경":
        return gr.update(choices=["숲", "해변", "산", "들판", "계곡", "열대 정원", "사막", "눈 풍경"], value="숲")
    elif background_type == "실내 배경":
        return gr.update(choices=["거실", "주방", "서재", "화장대", "호텔 로비", "카페", "미니멀 공간"], value="거실")
    elif background_type == "특수 배경":
        return gr.update(choices=["도시 풍경", "루프탑", "해안가", "모던 스튜디오", "빈티지 공간", "미니어처 공간"], value="도시 풍경")
    else:
        return gr.update(choices=["선택하세요"], value="선택하세요")

# 선택 내용 미리보기 업데이트 함수
def update_preview(product_name, product_category, background_type, background_style, lighting, texture, color_tone, mood):
    preview = f"상품명: {product_name}\n"
    preview += f"제품군: {product_category}\n"
    preview += f"배경 유형: {background_type}\n"
    preview += f"세부 스타일: {background_style}\n"
    
    if lighting:
        preview += f"조명: {lighting}\n"
    if texture:
        preview += f"배경 질감: {texture}\n"
    if color_tone:
        preview += f"색감: {color_tone}\n"
    if mood:
        preview += f"분위기: {mood}\n"
    
    return preview

# Gradio 인터페이스
with gr.Blocks() as demo:
    gr.HTML(
        """
        <div style="text-align: center; margin-bottom: 1rem;">
            <h1>이커머스 상품 이미지 배경 생성기</h1>
            <p>상품 이미지를 업로드하고 원하는 배경 옵션을 선택하여 새로운 배경 이미지를 생성하세요.</p>
        </div>
        """
    )
    
    current_selection = gr.State("")
    
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Row():
                image1_input = gr.Image(type="pil", label="#1 상품 이미지", image_mode="RGB")
                image2_input = gr.Image(type="pil", label="#2 (선택사항)", image_mode="RGB")
                image3_input = gr.Image(type="pil", label="#3 (선택사항)", image_mode="RGB")
            
            prompt_input = gr.Textbox(
                lines=3,
                placeholder="프롬프트를 입력하거나 비워두면 자동 합성됩니다.",
                label="프롬프트 (선택 사항)"
            )
            
            with gr.Row():
                image_change_btn = gr.Button("이미지 변경")
                text_remove_btn = gr.Button("글자지우기")
                clothes_change_btn = gr.Button("옷바꾸기")
                background_change_btn = gr.Button("배경바꾸기", variant="primary")
                composite_product_btn = gr.Button("상품합성")
            
            # 배경 변경 관련 섹션 (기본적으로 숨김 상태)
            with gr.Row(visible=False) as background_header:
                gr.HTML("<h3 style='margin-top: 20px; margin-bottom: 10px; width: 100%; text-align: center;'>배경 변경 옵션</h3>")
            
            # 상품명 입력
            product_name = gr.Textbox(
                label="상품명 입력",
                placeholder="상품명을 입력해주세요",
                visible=False
            )
            
            # 제품군 카테고리 드롭다운
            product_category = gr.Dropdown(
                choices=["화장품", "가방", "액세서리", "보석류", "홈 데코", "신발", "의류", "주방용품", "전자기기", "기타"],
                label="제품군 선택",
                value="기타",
                visible=False
            )
            
            # 배경 유형 드롭다운
            background_type = gr.Dropdown(
                choices=["자연 배경", "실내 배경", "특수 배경"],
                label="배경 유형 선택",
                value="자연 배경",
                visible=False
            )
            
            # 배경 세부 스타일 드롭다운 (배경 유형에 따라 동적 변경)
            background_style = gr.Dropdown(
                choices=["숲", "해변", "산", "들판", "계곡", "열대 정원", "사막", "눈 풍경"],
                label="배경 세부 스타일 선택",
                value="숲",
                visible=False
            )
            
            # 추가 선택 옵션 헤더
            with gr.Row(visible=False) as additional_header:
                gr.HTML("<h4 style='margin-top: 15px; margin-bottom: 10px; width: 100%; text-align: center;'>추가 선택 옵션</h4>")
            
            # 조명 드롭다운
            lighting = gr.Dropdown(
                choices=["자연광", "따뜻한 조명", "그림자 있는 조명", "미니멀 조명"],
                label="조명 선택",
                value="자연광",
                visible=False
            )
            
            # 배경 질감 드롭다운
            texture = gr.Dropdown(
                choices=["흰 벽", "대리석", "나무", "벽돌", "흰 천", "시멘트"],
                label="배경 질감 선택",
                value="",
                visible=False
            )
            
            # 색감 드롭다운
            color_tone = gr.Dropdown(
                choices=["화이트 톤", "블루 톤", "따뜻한 색상", "모노톤", "파스텔"],
                label="색감 선택",
                value="",
                visible=False
            )
            
            # 분위기 드롭다운
            mood = gr.Dropdown(
                choices=["포근한", "심플한", "모던한", "내추럴", "미니멀"],
                label="분위기 선택",
                value="",
                visible=False
            )
            
            # 선택 내용 미리보기
            selection_preview = gr.Textbox(
                label="선택 내용 미리보기",
                lines=8,
                readonly=True,
                visible=False
            )
            
            submit_btn = gr.Button("이미지 생성", variant="primary")
        
        with gr.Column(scale=1):
            output_image = gr.Image(label="생성된 이미지")
            output_text = gr.Textbox(label="상태 메시지")
            prompt_display = gr.Textbox(label="사용된 프롬프트 (영어)", visible=True)
            kr_prompt_display = gr.Textbox(label="사용된 프롬프트 (한국어)", visible=True)
            options_summary = gr.Textbox(label="선택된 옵션 요약", lines=8, visible=True)
    
    # 선택 옵션 버튼 클릭 시 프롬프트 입력란 업데이트 (한국어 문구)
    image_change_btn.click(
        fn=lambda: "#1 이미지의 [다른 모습]으로 변경하라.",
        inputs=[],
        outputs=prompt_input
    ).then(
        fn=lambda: "이미지 변경",
        inputs=[],
        outputs=current_selection
    ).then(
        fn=toggle_background_options,
        inputs=[current_selection],
        outputs=[background_header, product_name, product_category, background_type, 
                 background_style, additional_header, lighting, texture, color_tone, 
                 mood, selection_preview]
    )
    
    text_remove_btn.click(
        fn=lambda: "#1 이미지의 [중국어를 모두]를 제거하라.",
        inputs=[],
        outputs=prompt_input
    ).then(
        fn=lambda: "글자지우기",
        inputs=[],
        outputs=current_selection
    ).then(
        fn=toggle_background_options,
        inputs=[current_selection],
        outputs=[background_header, product_name, product_category, background_type, 
                 background_style, additional_header, lighting, texture, color_tone, 
                 mood, selection_preview]
    )
    
    clothes_change_btn.click(
        fn=lambda: "#1이미지에서 [신체 비율과 포즈는 유지한 체] 의상[스타일과 색상을]#2, #3으로 변경해주세요.",
        inputs=[],
        outputs=prompt_input
    ).then(
        fn=lambda: "옷바꾸기",
        inputs=[],
        outputs=current_selection
    ).then(
        fn=toggle_background_options,
        inputs=[current_selection],
        outputs=[background_header, product_name, product_category, background_type, 
                 background_style, additional_header, lighting, texture, color_tone, 
                 mood, selection_preview]
    )
    
    background_change_btn.click(
        fn=lambda: "배경 변경",
        inputs=[],
        outputs=prompt_input
    ).then(
        fn=lambda: "배경바꾸기",
        inputs=[],
        outputs=current_selection
    ).then(
        fn=toggle_background_options,
        inputs=[current_selection],
        outputs=[background_header, product_name, product_category, background_type, 
                 background_style, additional_header, lighting, texture, color_tone, 
                 mood, selection_preview]
    )
    
    composite_product_btn.click(
        fn=lambda: "이미지 #1의 모델이 필요에 따라 자세나 표정을 유동적으로 변형한 후, 이미지 #2(또는 #3)의 상품을 최대한 원래 형태와 디테일을 유지하면서 자연스럽게 착용한 모습을 합성하라.",
        inputs=[],
        outputs=prompt_input
    ).then(
        fn=lambda: "상품합성",
        inputs=[],
        outputs=current_selection
    ).then(
        fn=toggle_background_options,
        inputs=[current_selection],
        outputs=[background_header, product_name, product_category, background_type, 
                 background_style, additional_header, lighting, texture, color_tone, 
                 mood, selection_preview]
    )
    
    # 배경 유형에 따라 세부 스타일 옵션 업데이트
    background_type.change(
        fn=update_background_style_options,
        inputs=[background_type],
        outputs=[background_style]
    )
    
    # 선택 내용 미리보기 업데이트 (모든 선택 컴포넌트의 변경 이벤트에 연결)
    product_name.change(
        fn=update_preview,
        inputs=[product_name, product_category, background_type, background_style, lighting, texture, color_tone, mood],
        outputs=[selection_preview]
    )
    product_category.change(
        fn=update_preview,
        inputs=[product_name, product_category, background_type, background_style, lighting, texture, color_tone, mood],
        outputs=[selection_preview]
    )
    background_type.change(
        fn=update_preview,
        inputs=[product_name, product_category, background_type, background_style, lighting, texture, color_tone, mood],
        outputs=[selection_preview]
    )
    background_style.change(
        fn=update_preview,
        inputs=[product_name, product_category, background_type, background_style, lighting, texture, color_tone, mood],
        outputs=[selection_preview]
    )
    lighting.change(
        fn=update_preview,
        inputs=[product_name, product_category, background_type, background_style, lighting, texture, color_tone, mood],
        outputs=[selection_preview]
    )
    texture.change(
        fn=update_preview,
        inputs=[product_name, product_category, background_type, background_style, lighting, texture, color_tone, mood],
        outputs=[selection_preview]
    )
    color_tone.change(
        fn=update_preview,
        inputs=[product_name, product_category, background_type, background_style, lighting, texture, color_tone, mood],
        outputs=[selection_preview]
    )
    mood.change(
        fn=update_preview,
        inputs=[product_name, product_category, background_type, background_style, lighting, texture, color_tone, mood],
        outputs=[selection_preview]
    )
    
    # 이미지 생성 버튼 클릭 이벤트
    submit_btn.click(
        fn=process_and_show_prompt,
        inputs=[image1_input, image2_input, image3_input, prompt_input, 
                product_name, product_category, background_type, background_style, 
                lighting, texture, color_tone, mood],
        outputs=[output_image, output_text, prompt_display, options_summary, kr_prompt_display],
    )
    
    
    gr.Markdown(
        """
        ### 사용 방법:
        
        1. **자동 합성**: 이미지를 업로드하고 프롬프트를 비워두면 자동으로 합성됩니다.
        2. **이미지 참조**: #1, #2, #3으로 각 이미지를 참조할 수 있습니다.
        3. **선택 옵션**: 위의 버튼을 클릭하면 프롬프트 입력란에 한국어 문구로 입력하시면 됩니다.
        4. **배경 변경**: 배경바꾸기 버튼을 클릭하면 추가 옵션이 나타납니다.
        
        > **팁**: 프롬프트를 직접 수정할 수도 있습니다.
        """
    )

if __name__ == "__main__":
    demo.launch(share=True)