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 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, files_list, model="gemini-2.0-flash-exp-image-generation"):
    """
    여러 이미지와 텍스트 프롬프트를 사용하여 Gemini API로 이미지를 생성하는 함수
    
    Args:
        text (str): 이미지 생성에 사용할 텍스트 프롬프트
        files_list (list): 이미지 파일 경로 리스트
        model (str): 사용할 Gemini 모델명
    
    Returns:
        str: 생성된 이미지의 파일 경로 또는 None (오류 발생 시)
    """
    logger.debug(f"generate 함수 시작 - 텍스트: '{text}', 파일 수: {len(files_list)}, 모델: '{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 클라이언트 초기화 완료.")

        # 여러 파일 업로드
        uploaded_files = []
        for file_path in files_list:
            if file_path and os.path.exists(file_path):
                uploaded_file = client.files.upload(file=file_path)
                uploaded_files.append(uploaded_file)
                logger.debug(f"파일 업로드 완료. URI: {uploaded_file.uri}, MIME 타입: {uploaded_file.mime_type}")
            else:
                logger.warning(f"파일이 존재하지 않거나 None입니다: {file_path}")
        
        if not uploaded_files:
            logger.error("업로드할 파일이 없습니다.")
            return None

        # 파트 리스트 생성 (모든 이미지 + 텍스트 프롬프트)
        parts = []
        for uploaded_file in uploaded_files:
            parts.append(
                types.Part.from_uri(
                    file_uri=uploaded_file.uri,
                    mime_type=uploaded_file.mime_type,
                )
            )
        
        # 텍스트 프롬프트 추가
        parts.append(types.Part.from_text(text=text))
        
        # 컨텐츠 객체 생성
        contents = [
            types.Content(
                role="user",
                parts=parts,
            ),
        ]
        logger.debug(f"컨텐츠 객체 생성 완료: {contents}")

        generate_content_config = types.GenerateContentConfig(
            temperature=1,
            top_p=0.95,
            top_k=40,
            max_output_tokens=8192,
            response_modalities=[
                "image",
                "text",
            ],
            response_mime_type="text/plain",
        )
        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_stream = client.models.generate_content_stream(
                model=model,
                contents=contents,
                config=generate_content_config,
            )

            logger.debug("응답 스트림 처리 시작...")
            for chunk in response_stream:
                if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
                    logger.warning("chunk에 후보, 컨텐츠, 또는 파트가 없습니다. 건너뜁니다.")
                    continue

                inline_data = chunk.candidates[0].content.parts[0].inline_data
                if inline_data:
                    save_binary_file(temp_path, inline_data.data)
                    logger.info(f"MIME 타입 {inline_data.mime_type}의 파일이 저장됨: {temp_path} (프롬프트: {text})")
                else:
                    logger.info(f"수신된 텍스트: {chunk.text}")
                    print(chunk.text)

                logger.debug(f"Raw chunk: {chunk}")

        # 업로드된 파일 정보 정리
        del uploaded_files
        logger.debug("업로드된 파일 정보 삭제 완료.")
        return temp_path

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


def save_image_temp(image):
    """이미지를 임시 파일로 저장하는 유틸리티 함수"""
    if image is None:
        return None
    
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        temp_path = tmp.name
        if isinstance(image, Image.Image):
            if image.mode == "RGBA":
                # 알파 채널이 있는 경우 처리
                image = image.convert("RGB")
            image.save(temp_path)
        else:
            logger.warning(f"지원되지 않는 이미지 타입: {type(image)}")
            return None
        
        logger.debug(f"이미지 저장 완료: {temp_path}")
        return temp_path


def process_with_background_and_style(main_image, background_image, style_image, prompt):
    """
    주 이미지, 배경 이미지, 스타일 이미지를 프롬프트와 함께 처리하는 함수
    
    Args:
        main_image (PIL.Image): 주 대상 이미지
        background_image (PIL.Image): 배경 이미지
        style_image (PIL.Image): 스타일 참조 이미지
        prompt (str): 이미지 편집 프롬프트
    
    Returns:
        list: 처리된 이미지 목록
    """
    logger.debug(f"process_with_background_and_style 함수 시작 - 프롬프트: '{prompt}'")
    try:
        # 각 이미지를 임시 파일로 저장
        files_list = []
        
        # 주 이미지 저장
        main_path = save_image_temp(main_image)
        if main_path:
            files_list.append(main_path)
        
        # 배경 이미지 저장 (있는 경우)
        if background_image is not None:
            bg_path = save_image_temp(background_image)
            if bg_path:
                files_list.append(bg_path)
        
        # 스타일 이미지 저장 (있는 경우)
        if style_image is not None:
            style_path = save_image_temp(style_image)
            if style_path:
                files_list.append(style_path)
        
        # 이미지가 하나도 없는 경우 처리
        if not files_list:
            logger.error("처리할 이미지가 없습니다.")
            return []
        
        # 프롬프트 보강 (배경 및 스타일 관련 정보 추가)
        enhanced_prompt = prompt
        if background_image is not None and style_image is not None:
            enhanced_prompt = f"{prompt}. 첫 번째 이미지는 주 대상 이미지이고, 두 번째 이미지는 배경으로 사용해주세요. 세 번째 이미지의 스타일을 적용해주세요."
        elif background_image is not None:
            enhanced_prompt = f"{prompt}. 첫 번째 이미지는 주 대상 이미지이고, 두 번째 이미지는 배경으로 사용해주세요."
        elif style_image is not None:
            enhanced_prompt = f"{prompt}. 첫 번째 이미지는 주 대상 이미지이고, 두 번째 이미지의 스타일을 적용해주세요."
        
        logger.debug(f"보강된 프롬프트: {enhanced_prompt}")
        model = "gemini-2.0-flash-exp-image-generation"

        # Gemini API 호출
        result_path = generate(text=enhanced_prompt, files_list=files_list, model=model)

        if result_path:
            logger.debug(f"이미지 생성 완료. 경로: {result_path}")
            result_img = Image.open(result_path)
            if result_img.mode == "RGBA":
                result_img = result_img.convert("RGB")
            return [result_img]
        else:
            logger.error("generate 함수에서 None 반환됨.")
            return []  # 오류 시 빈 리스트 반환

    except Exception as e:
        logger.exception("process_with_background_and_style 함수에서 오류 발생:")
        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>배경 이미지와 스타일 이미지를 추가로 적용할 수 있습니다.</p>
                <p>Gemini API 키는 환경변수(GEMINI_API_KEY)로 설정되어 있습니다.</p>
            </div>
        </div>
        """
    )
    gr.Markdown("이미지를 업로드하고, 편집할 내용을 입력하세요. 배경 이미지와 스타일 이미지는 선택사항입니다.")

    with gr.Row():
        with gr.Column():
            # 입력 컴포넌트
            main_image_input = gr.Image(type="pil", label="주 이미지 업로드", image_mode="RGB")
            background_image_input = gr.Image(type="pil", label="배경 이미지 업로드 (선택사항)", image_mode="RGB")
            style_image_input = gr.Image(type="pil", label="스타일 이미지 업로드 (선택사항)", image_mode="RGB")
            
            prompt_input = gr.Textbox(
                lines=3,
                placeholder="편집할 내용을 입력하세요...",
                label="편집 프롬프트"
            )
            
            # 프롬프트 예시
            gr.Markdown("### 프롬프트 예시")
            gr.Markdown("- '주 이미지를 배경 이미지 위에 자연스럽게 배치해주세요'")
            gr.Markdown("- '주 이미지를 스타일 이미지의 화풍으로 변환해주세요'")
            gr.Markdown("- '주 이미지의 배경을 두 번째 이미지로 대체하고, 세 번째 이미지의 색감을 적용해주세요'")
            
            submit_btn = gr.Button("이미지 편집 실행")
            
        with gr.Column():
            # 출력 컴포넌트
            output_gallery = gr.Gallery(label="편집 결과")
            output_info = gr.Textbox(label="처리 결과 정보", lines=2)

    # 처리 함수 연결
    submit_btn.click(
        fn=process_with_background_and_style,
        inputs=[main_image_input, background_image_input, style_image_input, prompt_input],
        outputs=[output_gallery],
    ).then(
        lambda: "편집이 완료되었습니다. 결과를 확인해주세요.",
        None,
        output_info
    )

    # 예제 이미지 추가
    gr.Examples(
        examples=[
            ["example_main.jpg", "example_background.jpg", "example_style.jpg", "주 이미지의 인물을 배경 이미지에 합성하고, 스타일 이미지의 화풍을 적용해주세요"],
            ["example_main.jpg", None, "example_style.jpg", "이미지를 스타일 이미지의 화풍으로 변환해주세요"],
            ["example_main.jpg", "example_background.jpg", None, "주 이미지의 배경을 제거하고 두 번째 이미지 위에 합성해주세요"],
        ],
        inputs=[main_image_input, background_image_input, style_image_input, prompt_input],
    )


# 애플리케이션 실행
if __name__ == "__main__":
    # 실제 운영에서는 테스트 코드 제거
    # 테스트용 더미 이미지 (실제 이미지로 대체 가능)
    # dummy_image = Image.new("RGB", (100, 100), color="red")
    # dummy_bg = Image.new("RGB", (100, 100), color="blue")
    # dummy_style = Image.new("RGB", (100, 100), color="green")
    # dummy_prompt = "이미지를 합성해줘"
    
    # logger.info("process_with_background_and_style 함수를 직접 호출합니다...")
    # result = process_with_background_and_style(dummy_image, dummy_bg, dummy_style, dummy_prompt)
    
    # if result:
    #     logger.info(f"직접 호출 성공. 결과: {result}")
    # else:
    #     logger.error("직접 호출 실패.")
    
    # Gradio 앱 실행
    demo.launch(share=True)