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" 대신 다음과 같이 수정
import google.generativeai as genai
from google.generativeai 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:
            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

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)