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, file_name, background_file=None, style_file=None, model="gemini-2.0-flash-exp-image-generation"):
    logger.debug(f"generate 함수 시작 - 텍스트: '{text}', 파일명: '{file_name}', 모델: '{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 = []
        uploaded_files.append(client.files.upload(file=file_name))
        logger.debug(f"원본 파일 업로드 완료. URI: {uploaded_files[0].uri}, MIME 타입: {uploaded_files[0].mime_type}")

        if background_file is not None:
            bg_file = client.files.upload(file=background_file)
            uploaded_files.append(bg_file)
            logger.debug(f"배경 파일 업로드 완료. URI: {bg_file.uri}, MIME 타입: {bg_file.mime_type}")
        if style_file is not None:
            style_uploaded = client.files.upload(file=style_file)
            uploaded_files.append(style_uploaded)
            logger.debug(f"스타일 파일 업로드 완료. URI: {style_uploaded.uri}, MIME 타입: {style_uploaded.mime_type}")

        # 컨텐츠 객체 생성: 업로드한 이미지들을 순서대로(원본, 배경, 스타일) 추가하고 마지막에 텍스트 프롬프트 추가
        parts = []
        # 원본 이미지 파트
        parts.append(
            types.Part.from_uri(
                file_uri=uploaded_files[0].uri,
                mime_type=uploaded_files[0].mime_type,
            )
        )
        # 배경 이미지 파트 (존재 시)
        if background_file is not None:
            parts.append(
                types.Part.from_uri(
                    file_uri=uploaded_files[1].uri,
                    mime_type=uploaded_files[1].mime_type,
                )
            )
        # 스타일 이미지 파트 (존재 시)
        if style_file is not None:
            # 배경이 없는 경우 uploaded_files[1]가 스타일 이미지가 됨에 주의
            style_index = 2 if background_file is not None else 1
            parts.append(
                types.Part.from_uri(
                    file_uri=uploaded_files[style_index].uri,
                    mime_type=uploaded_files[style_index].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 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}")

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

        gemma_edited_image_path = generate(text=input_text,
                                           file_name=original_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():
            original_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=[original_input, prompt_input, background_input, style_input],
        outputs=output_gallery,
    )

# --- 테스트 코드 ---
# 테스트용 더미 이미지 (실제 이미지로 대체 가능)
dummy_original = 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_original, dummy_prompt, dummy_background, dummy_style)

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

demo.launch(share=True)