app.py

import os
import tempfile
from PIL import Image
import gradio as gr
import logging

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 process_images_with_prompt(image1, image2, image3, prompt):
    """
    3개의 이미지와 프롬프트를 처리하는 함수
    """
    try:
        # API 키 확인
        api_key = os.environ.get("GEMINI_API_KEY")
        if not api_key:
            return None, "API 키가 설정되지 않았습니다. 환경변수를 확인해주세요."
            
        # Gemini 클라이언트 초기화
        client = genai.Client(api_key=api_key)
        
        # 프롬프트 기본값 설정
        if not prompt or not prompt.strip():
            prompt = "이 이미지들을 활용하여 새로운 이미지를 생성해주세요."
        
        # 컨텐츠 리스트 생성 (이미지와 프롬프트 결합)
        parts = []
        
        # 텍스트 프롬프트 추가
        parts.append(types.Part.from_text(text=prompt))
        
        # 사용 가능한 이미지 추가
        for img in [image1, image2, image3]:
            if img is not None:
                # PIL 이미지를 바이트로 변환
                with tempfile.NamedTemporaryFile(suffix=".png") as tmp:
                    img_path = tmp.name
                    img.save(img_path, format="PNG")
                    with open(img_path, "rb") as f:
                        image_bytes = f.read()
                
                # 이미지를 파트로 추가
                parts.append(types.Part.from_data(data=image_bytes, mime_type="image/png"))
        
        # 생성 설정
        generate_content_config = types.GenerateContentConfig(
            temperature=1,
            response_modalities=["image"],
        )
        
        # 임시 파일 생성
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            temp_path = tmp.name
            
            # Gemini 모델로 요청 전송
            response = client.models.generate_content(
                model="gemini-2.0-flash-exp-image-generation",
                contents=[types.Content(role="user", parts=parts)],
                config=generate_content_config,
            )
            
            # 응답에서 이미지 추출
            for part in response.candidates[0].content.parts:
                if hasattr(part, 'inline_data') and part.inline_data:
                    save_binary_file(temp_path, part.inline_data.data)
                    
            # 결과 이미지 반환
            result_img = Image.open(temp_path)
            if result_img.mode == "RGBA":
                result_img = result_img.convert("RGB")
                
            return result_img, "이미지가 성공적으로 생성되었습니다."
            
    except Exception as e:
        logger.exception("이미지 생성 중 오류 발생:")
        return None, f"오류 발생: {str(e)}"

# 간소화된 Gradio 인터페이스
with gr.Blocks() as demo:
    gr.HTML("<h1>간단한 이미지 생성기</h1><p>이미지 3개와 프롬프트를 입력하세요</p>")
    
    with gr.Row():
        with gr.Column():
            # 3개의 이미지 입력
            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="프롬프트"
            )
            
            # 생성 버튼
            submit_btn = gr.Button("이미지 생성")
        
        with gr.Column():
            # 결과 출력
            output_image = gr.Image(label="생성된 이미지")
            output_text = gr.Textbox(label="상태 메시지")
    
    # 버튼 클릭 이벤트
    submit_btn.click(
        fn=process_images_with_prompt,
        inputs=[image1_input, image2_input, image3_input, prompt_input],
        outputs=[output_image, output_text],
    )

# 애플리케이션 실행
if __name__ == "__main__":
    demo.launch(share=True)