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 파일에 저장된 환경변수 로드 (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, composite_file, background_file=None, style_file=None, model="gemini-2.0-flash-exp-image-generation"):
    logger.debug(f"generate 함수 시작 - 텍스트: '{text}', composite_file: '{composite_file}', "
                 f"background_file: '{background_file}', style_file: '{style_file}', 모델: '{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 클라이언트 초기화 완료.")

        # 업로드할 파일들 리스트 생성
        files = []
        # 원본 이미지 업로드 (필수)
        composite_upload = client.files.upload(file=composite_file)
        files.append(composite_upload)
        logger.debug(f"원본 이미지 업로드 완료. URI: {composite_upload.uri}, MIME 타입: {composite_upload.mime_type}")

        # 배경 이미지 업로드 (선택)
        if background_file is not None:
            background_upload = client.files.upload(file=background_file)
            files.append(background_upload)
            logger.debug(f"배경 이미지 업로드 완료. URI: {background_upload.uri}, MIME 타입: {background_upload.mime_type}")
        # 스타일 이미지 업로드 (선택)
        if style_file is not None:
            style_upload = client.files.upload(file=style_file)
            files.append(style_upload)
            logger.debug(f"스타일 이미지 업로드 완료. URI: {style_upload.uri}, MIME 타입: {style_upload.mime_type}")

        # 파일 업로드 결과를 바탕으로 콘텐츠 파트 생성
        parts = []
        for file_obj in files:
            parts.append(types.Part.from_uri(file_uri=file_obj.uri, mime_type=file_obj.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 files
        logger.debug("업로드된 파일 정보 삭제 완료.")
        return temp_path

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


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

        background_path = None
        style_path = None

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

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

        model = "gemini-2.0-flash-exp-image-generation"
        gemma_edited_image_path = generate(text=prompt,
                                           composite_file=composite_path,
                                           background_file=background_path,
                                           style_file=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")
            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():
            composite_input = gr.Image(type="pil", label="원본 이미지", image_mode="RGBA")
            background_input = gr.Image(type="pil", label="배경 이미지", image_mode="RGBA")
            style_input = gr.Image(type="pil", label="스타일 이미지", image_mode="RGBA")
            prompt_input = gr.Textbox(
                lines=2,
                placeholder="편집할 내용을 입력하세요...",
                label="편집 프롬프트"
            )
            submit_btn = gr.Button("이미지 편집 실행")
        with gr.Column():
            output_gallery = gr.Gallery(label="편집 결과")

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

# --- 테스트 코드 ---
dummy_composite = Image.new("RGBA", (100, 100), color="red")
dummy_background = Image.new("RGBA", (100, 100), color="green")
dummy_style = Image.new("RGBA", (100, 100), color="blue")
dummy_prompt = "이미지에 새로운 스타일을 적용해줘"

logger.info("process_image_and_prompt 함수를 직접 호출합니다...")
result = process_image_and_prompt(dummy_composite, dummy_background, dummy_style, dummy_prompt)

if result:
    logger.info(f"직접 호출 성공. 결과: {result}")
else:
    logger.error("직접 호출 실패.")

demo.launch(share=True)