app.py

import json
import os
import time
import uuid
import tempfile
from PIL import Image
import gradio as gr
import base64
import mimetypes
import logging
from io import BytesIO

from google import genai
from google.genai import types

# .env 파일에 저장된 환경변수 로드 (python-dotenv 설치 필요: pip install python-dotenv)
from dotenv import load_dotenv
load_dotenv()

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


def save_binary_file(file_name, data):
    logger.debug(f"파일에 이진 데이터 저장 중: {file_name}")
    with open(file_name, "wb") as f:
        f.write(data)
    logger.debug(f"파일 저장 완료: {file_name}")


def generate(text, original_image_path, background_image_path=None, style_image_path=None, model="gemini-2.0-flash-exp-image-generation"):
    logger.debug(f"generate 함수 시작 - 텍스트: '{text}', 원본 파일명: '{original_image_path}', 모델: '{model}'")
    
    try:
        # API 키는 환경변수에서 불러옴
        effective_api_key = os.environ.get("GEMINI_API_KEY")
        if effective_api_key:
            logger.debug("환경변수에서 API 키 불러옴")
        else:
            logger.error("API 키가 환경변수에 설정되지 않았습니다.")
            raise ValueError("API 키가 필요합니다.")

        client = genai.Client(api_key=effective_api_key)
        logger.debug("Gemini 클라이언트 초기화 완료.")

        # PIL 이미지 객체로 변환
        original_img = Image.open(original_image_path)
        
        # 컨텐츠 리스트 생성 (공식 문서 방식대로)
        contents = []
        
        # 텍스트 프롬프트를 명확하게 작성
        prompt = text
        if background_image_path and "배경" not in text.lower():
            prompt += " 원본 이미지의 배경을 두 번째 업로드된 이미지로 완전히 교체해 주세요. 이미지를 업데이트하고 결과를 보여주세요."
        if style_image_path and "스타일" not in text.lower():
            prompt += " 세 번째 이미지의 스타일을 전체적으로 적용해 주세요."
            
        contents.append(prompt)
        contents.append(original_img)
        
        # 배경 이미지 추가 (있는 경우)
        if background_image_path:
            background_img = Image.open(background_image_path)
            contents.append(background_img)
            logger.debug("배경 이미지 추가됨")
            
        # 스타일 이미지 추가 (있는 경우)
        if style_image_path:
            style_img = Image.open(style_image_path)
            contents.append(style_img)
            logger.debug("스타일 이미지 추가됨")
            
        logger.debug(f"컨텐츠 객체 생성 완료: {len(contents)} 아이템")

        # 생성 설정
        generate_content_config = types.GenerateContentConfig(
            temperature=1,
            top_p=0.95,
            top_k=40,
            max_output_tokens=8192,
            response_modalities=["text", "image"],  # 문서에 맞게 대소문자 수정
        )
        logger.debug(f"생성 설정: {generate_content_config}")

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            temp_path = tmp.name
            logger.debug(f"임시 파일 생성됨: {temp_path}")

            # 스트리밍 대신 단일 요청으로 변경 (이미지 생성에 더 적합)
            response = client.models.generate_content(
                model=model,
                contents=contents,
                config=generate_content_config,
            )
            
            logger.debug("응답 처리 시작...")
            
            # 응답에서 이미지 추출
            image_saved = False
            for part in response.candidates[0].content.parts:
                if hasattr(part, 'text') and part.text:
                    logger.info(f"수신된 텍스트: {part.text}")
                    print(part.text)
                elif hasattr(part, 'inline_data') and part.inline_data:
                    save_binary_file(temp_path, part.inline_data.data)
                    logger.info(f"MIME 타입 {part.inline_data.mime_type}의 파일이 저장됨: {temp_path}")
                    image_saved = True

            if not image_saved:
                logger.warning("이미지가 생성되지 않았습니다.")
                return None

        logger.debug("이미지 생성 완료.")
        return temp_path

    except Exception as e:
        logger.exception("이미지 생성 중 오류 발생:")
        return None  # 오류 발생 시 None 반환


def process_image_and_prompt(original_pil, prompt, background_pil=None, style_pil=None):
    logger.debug(f"process_image_and_prompt 함수 시작 - 프롬프트: '{prompt}'")
    try:
        # 원본 이미지 저장
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            original_path = tmp.name
            original_pil.save(original_path)
            logger.debug(f"원본 이미지 저장 완료: {original_path}")

        # 배경 이미지 저장 (제공된 경우)
        background_path = None
        if background_pil is not None:
            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_bg:
                background_path = tmp_bg.name
                background_pil.save(background_path)
                logger.debug(f"배경 이미지 저장 완료: {background_path}")

        # 스타일 이미지 저장 (제공된 경우)
        style_path = None
        if style_pil is not None:
            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_style:
                style_path = tmp_style.name
                style_pil.save(style_path)
                logger.debug(f"스타일 이미지 저장 완료: {style_path}")

        # 프롬프트 보강 (영어와 한국어 두 언어로 제공)
        if prompt and not prompt.strip():
            if background_path and style_path:
                prompt = "원본 이미지의 인물을 유지하면서 배경을 두 번째 이미지로 교체하고 세 번째 이미지의 스타일을 적용해 주세요. Please replace the background with the second image while keeping the person, and apply the style of the third image."
            elif background_path:
                prompt = "원본 이미지의 인물을 유지하면서 배경을 두 번째 이미지로 교체해 주세요. Please replace the background with the second image while keeping the person."
            elif style_path:
                prompt = "원본 이미지에 두 번째 이미지의 스타일을 적용해 주세요. Please apply the style of the second image to the original image."

        model = "gemini-2.0-flash-exp-image-generation"

        gemma_edited_image_path = generate(
            text=prompt,
            original_image_path=original_path,
            background_image_path=background_path,
            style_image_path=style_path,
            model=model
        )

        if gemma_edited_image_path:
            logger.debug(f"이미지 생성 완료. 경로: {gemma_edited_image_path}")
            result_img = Image.open(gemma_edited_image_path)
            if result_img.mode == "RGBA":
                result_img = result_img.convert("RGB")
            
            # 임시 파일 정리
            try:
                os.unlink(original_path)
                if background_path:
                    os.unlink(background_path)
                if style_path:
                    os.unlink(style_path)
            except Exception as e:
                logger.warning(f"임시 파일 삭제 중 오류: {str(e)}")
            
            return [result_img]
        else:
            logger.error("generate 함수에서 None 반환됨.")
            return []  # 오류 시 빈 리스트 반환

    except Exception as e:
        logger.exception("process_image_and_prompt 함수에서 오류 발생:")
        return []  # 오류 시 빈 리스트 반환


# --- Gradio 인터페이스 구성 ---
with gr.Blocks() as demo:
    gr.HTML(
        """
        <div style='display: flex; align-items: center; justify-content: center; gap: 20px'>
            <div style="background-color: var(--block-background-fill); border-radius: 8px">
                <img src="https://www.gstatic.com/lamda/images/gemini_favicon_f069958c85030456e93de685481c559f160ea06b.png" style="width: 100px; height: 100px;">
            </div>
            <div>
                <h1>Gemini를 이용한 이미지 편집</h1>
                <p>Gemini API 키는 환경변수(GEMINI_API_KEY)로 설정되어 있습니다.</p>
            </div>
        </div>
        """
    )
    gr.Markdown("원본, 배경, 스타일 이미지를 업로드하고, 편집할 내용을 입력하세요.")

    with gr.Row():
        with gr.Column():
            original_input = gr.Image(type="pil", label="원본 이미지", image_mode="RGB")
            background_input = gr.Image(type="pil", label="배경 이미지", image_mode="RGB")
            style_input = gr.Image(type="pil", label="스타일 이미지", image_mode="RGB")
            prompt_input = gr.Textbox(
                lines=2,
                placeholder="편집할 내용을 입력하세요...",
                label="편집 프롬프트"
            )
            submit_btn = gr.Button("이미지 편집 실행")
        with gr.Column():
            output_gallery = gr.Gallery(label="편집 결과")
            output_text = gr.Textbox(label="API 응답 텍스트", visible=True)

    submit_btn.click(
        fn=process_image_and_prompt,
        inputs=[original_input, prompt_input, background_input, style_input],
        outputs=output_gallery,
    )

    gr.HTML("""
    <div style="margin-top: 20px; padding: 10px; background-color: #f8f9fa; border-radius: 8px;">
        <h3>사용 팁:</h3>
        <ul>
            <li><strong>프롬프트 작성 예시:</strong> "원본 인물을 유지하면서 배경을 두 번째 이미지로 교체해 주세요."</li>
            <li><strong>배경 교체 명확화:</strong> "배경만 완전히 교체하고 원본 인물은 그대로 유지해 주세요."</li>
            <li><strong>스타일 적용:</strong> "세 번째 이미지의 예술 스타일을 전체 이미지에 적용해 주세요."</li>
            <li><strong>영어 프롬프트:</strong> 더 나은 결과를 위해 영어와 한국어를 함께 사용해 보세요.</li>
        </ul>
    </div>
    """)

# --- 실행 ---
demo.launch(share=True)