app.py

import os
import tempfile
from PIL import Image
import gradio as gr
import logging
import re
import io
from io import BytesIO

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 preprocess_prompt(prompt, image1, image2, image3):
    """
    입력된 프롬프트가 영어로만 되어 있는지 확인하고,
    단순히 "Please generate the image."를 덧붙여 반환합니다.
    """
    # 프롬프트에 한글이 포함되어 있으면 에러 발생
    if re.search(r'[ㄱ-ㅎ가-힣]', prompt):
        raise ValueError("Error: Prompt must be in English only.")

    # 불필요한 기능 관련 처리는 제거하고 단순 명령어 추가
    prompt = prompt.strip() + " Please generate the image."
    return prompt

def generate_with_images(prompt, images):
    """
    공식 문서에 기반한 올바른 API 호출 방식 구현
    """
    try:
        # API 키 확인
        api_key = os.environ.get("GEMINI_API_KEY")
        if not api_key:
            return None, "API key is not set. Please check your environment variables."
            
        # Gemini 클라이언트 초기화
        client = genai.Client(api_key=api_key)
        
        logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt}")
        
        # 컨텐츠 준비
        contents = []
        
        # 텍스트 프롬프트 추가
        contents.append(prompt)
        
        # 이미지 추가 (이미지가 없으면 텍스트만 전송)
        for idx, img in enumerate(images, 1):
            if img is not None:
                contents.append(img)
                logger.info(f"Image #{idx} added.")
        
        # 생성 설정 - 공식 문서에 따라 responseModalities 설정
        response = client.models.generate_content(
            model="gemini-2.0-flash-exp-image-generation",
            contents=contents,
            config=types.GenerateContentConfig(
                response_modalities=['Text', 'Image'],
                temperature=1,
                top_p=0.95,
                top_k=40,
                max_output_tokens=8192
            )
        )
        
        # 임시 파일 생성
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            temp_path = tmp.name
            
            result_text = ""
            image_found = False
            
            # 응답 처리
            for part in response.candidates[0].content.parts:
                if hasattr(part, 'text') and part.text:
                    result_text += part.text
                    logger.info(f"Response text: {part.text}")
                elif hasattr(part, 'inline_data') and part.inline_data:
                    save_binary_file(temp_path, part.inline_data.data)
                    image_found = True
                    logger.info("Image extracted successfully from response.")
            
            if not image_found:
                return None, f"API did not generate an image. Response text: {result_text}"
                    
            # 결과 이미지 반환
            result_img = Image.open(temp_path)
            if result_img.mode == "RGBA":
                result_img = result_img.convert("RGB")
                
            return result_img, f"Image generated successfully. {result_text}"
            
    except Exception as e:
        logger.exception("Error during image generation:")
        return None, f"Error occurred: {str(e)}"

def process_images_with_prompt(image1, image2, image3, prompt):
    """
    3개의 이미지와 프롬프트를 처리하는 함수.
    이미지 입력이 없어도 프롬프트만으로 API 호출을 진행합니다.
    """
    try:
        # 이미지 개수 확인 (이미지 없이도 진행)
        images = [image1, image2, image3]
        valid_images = [img for img in images if img is not None]
        
        # 프롬프트가 없으면 업로드된 이미지 수에 따라 자동 합성 프롬프트 생성 (영어)
        if not prompt or not prompt.strip():
            if len(valid_images) == 0:
                prompt = "Please generate an image based on the description."
                logger.info("Auto prompt generated for no image input.")
            elif len(valid_images) == 1:
                prompt = "Please creatively transform this image into a more vivid and artistic version."
                logger.info("Auto prompt generated for a single image.")
            elif len(valid_images) == 2:
                prompt = "Please seamlessly blend these two images, integrating their elements harmoniously into a single image."
                logger.info("Auto prompt generated for two images.")
            else:
                prompt = "Please creatively composite these three images, incorporating their key elements into a natural and coherent scene."
                logger.info("Auto prompt generated for three images.")
        else:
            prompt = preprocess_prompt(prompt, image1, image2, image3)
        
        # 새로운 API 호출 방식 사용
        return generate_with_images(prompt, valid_images)
            
    except Exception as e:
        logger.exception("Error during image processing:")
        return None, f"Error occurred: {str(e)}"

# Gradio 인터페이스 (기능 선택, 기능 적용, 커스텀 텍스트 관련 요소 제거)
with gr.Blocks() as demo:
    gr.HTML(
        """
        <div style="text-align: center; margin-bottom: 1rem;">
            <h1>Simple Image Generator</h1>
            <p>Upload an image (or leave empty) and click generate to create an image based on the English prompt.</p>
        </div>
        """
    )
    
    with gr.Row():
        with gr.Column():
            # 3개의 이미지 입력 (이미지 없이도 실행 가능)
            with gr.Row():
                image1_input = gr.Image(type="pil", label="Image 1", image_mode="RGB")
                image2_input = gr.Image(type="pil", label="Image 2", image_mode="RGB")
                image3_input = gr.Image(type="pil", label="Image 3", image_mode="RGB")
            
            # 프롬프트 입력 (영어로만 입력)
            prompt_input = gr.Textbox(
                lines=3,
                placeholder="Enter the prompt in English.",
                label="Prompt (Required: English only)"
            )
            
            # 생성 버튼
            submit_btn = gr.Button("Generate Image", variant="primary")
        
        with gr.Column():
            # 결과 출력
            output_image = gr.Image(label="Generated Image")
            output_text = gr.Textbox(label="Status Message")
            prompt_display = gr.Textbox(label="Used Prompt", visible=True)
    
    # 이미지 생성 버튼 클릭 이벤트
    def process_and_show_prompt(image1, image2, image3, prompt):
        # 이미지 개수 확인
        images = [image1, image2, image3]
        valid_images = [img for img in images if img is not None]
        
        try:
            # 만약 사용자가 프롬프트를 입력했다면 영어만 포함되었는지 확인
            if prompt and re.search(r'[ㄱ-ㅎ가-힣]', prompt):
                return None, "Error: Prompt must be in English only.", prompt
            
            # 프롬프트가 없으면 자동 생성 (영어)
            auto_prompt = prompt
            if not prompt or not prompt.strip():
                if len(valid_images) == 0:
                    auto_prompt = "Please generate an image based on the description."
                elif len(valid_images) == 1:
                    auto_prompt = "Please creatively transform this image into a more vivid and artistic version."
                elif len(valid_images) == 2:
                    auto_prompt = "Please seamlessly blend these two images, integrating their elements harmoniously into a single image."
                else:
                    auto_prompt = "Please creatively composite these three images, incorporating their key elements into a natural and coherent scene."
            else:
                auto_prompt = preprocess_prompt(prompt, image1, image2, image3)
            
            result_img, status = process_images_with_prompt(image1, image2, image3, prompt)
            
            return result_img, status, auto_prompt
        except Exception as e:
            logger.exception("Error during processing:")
            return None, f"Error occurred: {str(e)}", prompt
    
    submit_btn.click(
        fn=process_and_show_prompt,
        inputs=[image1_input, image2_input, image3_input, prompt_input],
        outputs=[output_image, output_text, prompt_display],
    )

    gr.Markdown(
        """
        ### Instructions:
        
        1. **Auto Generation**: You can leave the image upload empty and the system will generate an image based solely on the prompt.
        2. **Prompt Requirement**: Enter the prompt in English only.
        3. **Image Reference**: The app supports up to three image inputs.
        """
    )

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