app.py

import json
import os
import time
import tempfile
from PIL import Image
import gradio as gr
import logging
import io

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_image_from_prompt(prompt, model="gemini-2.0-flash-exp-image-generation"):
    """텍스트 프롬프트만으로 이미지 생성"""
    logger.debug(f"generate_image_from_prompt 함수 시작 - 프롬프트: '{prompt}'")
    
    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 클라이언트 초기화 완료.")

        # 프롬프트 처리
        if not prompt or not prompt.strip():
            prompt = "A creative and visually appealing image that captures imagination. Use vibrant colors and interesting composition."
        
        # 한국어 프롬프트에 영어 설명 추가
        if not any(ord(c) < 128 for c in prompt):
            prompt += " Create a highly detailed, visually stunning image that captures the essence of the description."

        # 컨텐츠 리스트 생성
        parts = [
            types.Part.from_text(text=prompt)
        ]
        
        logger.debug(f"컨텐츠 객체 생성 완료: {len(parts)} 아이템")

        # 생성 설정
        generate_content_config = types.GenerateContentConfig(
            temperature=1,
            top_p=0.95,
            top_k=40,
            max_output_tokens=8192,
            response_modalities=["image", "text"],
        )
        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=[
                    types.Content(
                        role="user",
                        parts=parts,
                    ),
                ],
                config=generate_content_config,
            )
            
            logger.debug("응답 스트림 처리 시작...")
            
            # 응답에서 이미지와 텍스트 추출
            image_saved = False
            response_text = ""
            
            for chunk in response_stream:
                logger.debug(f"chunk 수신: {chunk}")
                
                # 응답 검증
                if not hasattr(chunk, 'candidates') or not chunk.candidates or len(chunk.candidates) == 0:
                    logger.warning("chunk에 candidates가 없습니다. 건너뜁니다.")
                    continue
                    
                if not hasattr(chunk.candidates[0], 'content') or chunk.candidates[0].content is None:
                    logger.warning("chunk.candidates[0]에 content가 없습니다. 건너뜁니다.")
                    continue
                    
                if not hasattr(chunk.candidates[0].content, 'parts') or not chunk.candidates[0].content.parts:
                    logger.warning("chunk.candidates[0].content에 parts가 없습니다. 건너뜁니다.")
                    continue
                
                for part in chunk.candidates[0].content.parts:
                    if hasattr(part, 'text') and part.text:
                        response_text += part.text
                        logger.info(f"수신된 텍스트: {part.text}")
                    elif hasattr(part, 'inline_data') and part.inline_data:
                        save_binary_file(temp_path, part.inline_data.data)
                        logger.info(f"MIME 타입 {part.inline_data.mime_type}의 파일이 저장됨: {temp_path}")
                        image_saved = True

            if not image_saved:
                logger.warning("이미지가 생성되지 않았습니다.")
                return None, response_text or "이미지가 생성되지 않았습니다. 다른 프롬프트로 시도해보세요."

        logger.debug("이미지 생성 완료.")
        return temp_path, response_text

    except Exception as e:
        logger.exception("이미지 생성 중 오류 발생:")
        return None, f"오류 발생: {str(e)}"  # 오류 발생 시 None과 오류 메시지 반환


def merge_images(person_img_path, product_img_path, background_img_path, prompt, model="gemini-2.0-flash-exp-image-generation"):
    """기존 이미지 합성 함수 (이전 코드와 동일)"""
    logger.debug(f"merge_images 함수 시작 - 프롬프트: '{prompt}'")
    
    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 클라이언트 초기화 완료.")

        # 이미지 파일 업로드
        person_file = client.files.upload(file=person_img_path)
        logger.debug(f"사람 이미지 업로드 완료: {person_file.uri}")
        
        product_file = client.files.upload(file=product_img_path)
        logger.debug(f"상품 이미지 업로드 완료: {product_file.uri}")
        
        # 컨텐츠 리스트 생성
        parts = []
        
        # 사람 이미지 추가
        parts.append(
            types.Part.from_uri(
                file_uri=person_file.uri,
                mime_type=person_file.mime_type
            )
        )
        
        # 상품 이미지 추가
        parts.append(
            types.Part.from_uri(
                file_uri=product_file.uri,
                mime_type=product_file.mime_type
            )
        )
        
        # 배경 이미지가 있으면 추가
        if background_img_path:
            background_file = client.files.upload(file=background_img_path)
            logger.debug(f"배경 이미지 업로드 완료: {background_file.uri}")
            parts.append(
                types.Part.from_uri(
                    file_uri=background_file.uri,
                    mime_type=background_file.mime_type
                )
            )
            logger.debug("배경 이미지 추가됨")
        
        # 마지막에 프롬프트 추가
        parts.append(
            types.Part.from_text(text=prompt)
        )
        
        logger.debug(f"컨텐츠 객체 생성 완료: {len(parts)} 아이템")

        # 생성 설정
        generate_content_config = types.GenerateContentConfig(
            temperature=1,
            top_p=0.95,
            top_k=40,
            max_output_tokens=8192,
            response_modalities=["image", "text"],
        )
        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=[
                    types.Content(
                        role="user",
                        parts=parts,
                    ),
                ],
                config=generate_content_config,
            )
            
            logger.debug("응답 스트림 처리 시작...")
            
            # 응답에서 이미지와 텍스트 추출
            image_saved = False
            response_text = ""
            
            for chunk in response_stream:
                logger.debug(f"chunk 수신: {chunk}")
                
                # 응답 검증
                if not hasattr(chunk, 'candidates') or not chunk.candidates or len(chunk.candidates) == 0:
                    logger.warning("chunk에 candidates가 없습니다. 건너뜁니다.")
                    continue
                    
                if not hasattr(chunk.candidates[0], 'content') or chunk.candidates[0].content is None:
                    logger.warning("chunk.candidates[0]에 content가 없습니다. 건너뜁니다.")
                    continue
                    
                if not hasattr(chunk.candidates[0].content, 'parts') or not chunk.candidates[0].content.parts:
                    logger.warning("chunk.candidates[0].content에 parts가 없습니다. 건너뜁니다.")
                    continue
                
                for part in chunk.candidates[0].content.parts:
                    if hasattr(part, 'text') and part.text:
                        response_text += part.text
                        logger.info(f"수신된 텍스트: {part.text}")
                    elif hasattr(part, 'inline_data') and part.inline_data:
                        save_binary_file(temp_path, part.inline_data.data)
                        logger.info(f"MIME 타입 {part.inline_data.mime_type}의 파일이 저장됨: {temp_path}")
                        image_saved = True

            if not image_saved:
                logger.warning("이미지가 생성되지 않았습니다.")
                return None, response_text or "이미지가 생성되지 않았습니다. 다른 프롬프트나 이미지로 시도해보세요."

        logger.debug("이미지 생성 완료.")
        return temp_path, response_text

    except Exception as e:
        logger.exception("이미지 생성 중 오류 발생:")
        return None, f"오류 발생: {str(e)}"  # 오류 발생 시 None과 오류 메시지 반환


def process_images_and_prompt(person_pil, product_pil, background_pil, prompt):
    """
    이미지 및 프롬프트를 처리하여 이미지를 생성하는 함수
    
    이미지 입력이 없는 경우 텍스트 기반 이미지 생성으로 전환
    이미지 입력이 있는 경우 기존의 이미지 합성 로직 유지
    """
    logger.debug(f"process_images_and_prompt 함수 시작 - 프롬프트: '{prompt}'")
    
    try:
        # 이미지 입력이 모두 없는 경우 텍스트 기반 이미지 생성
        if person_pil is None and product_pil is None and background_pil is None:
            logger.debug("이미지 없음. 텍스트 기반 이미지 생성으로 전환")
            result_path, response_text = generate_image_from_prompt(prompt)
            
            if result_path:
                try:
                    result_img = Image.open(result_path)
                    if result_img.mode == "RGBA":
                        result_img = result_img.convert("RGB")
                    
                    return [result_img], response_text or "이미지가 성공적으로 생성되었습니다."
                except Exception as e:
                    logger.exception(f"결과 이미지 로드 중 오류: {str(e)}")
                    return [], f"결과 이미지 처리 중 오류: {str(e)}"
            else:
                return [], response_text or "이미지 생성에 실패했습니다."
        
        # 필수 이미지 확인 (이미지 합성 모드)
        if person_pil is None:
            return [], "사람 이미지를 업로드해주세요."
        if product_pil is None:
            return [], "상품 이미지를 업로드해주세요."
            
        # 기본 프롬프트 설정 (비어있는 경우)
        if not prompt or not prompt.strip():
            if background_pil:
                prompt = "이 배경에 이 사람이 이 상품을 사용하는 모습을 자연스럽게 보여주세요. 상품을 잘 보이게 해주세요. Create a natural composite image showing this person using this product in this background setting. Make sure the product is clearly visible."
            else:
                prompt = "이 사람이 이 상품을 사용하는 모습을 자연스럽게 보여주세요. 상품을 잘 보이게 해주세요. Create a natural composite image showing this person using this product. Make sure the product is clearly visible."
        
        # 프롬프트에 영어가 없으면 영어 프롬프트 추가 (더 나은 결과를 위해)
        if not any(ord(c) < 128 for c in prompt):
            if background_pil:
                prompt += " Create a realistic composite image of this person with this product in this background."
            else:
                prompt += " Create a realistic composite image of this person with this product."
        
        # 이미지 저장
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_person:
            person_path = tmp_person.name
            # RGB로 변환하여 저장
            person_img = person_pil
            if person_img.mode != 'RGB':
                person_img = person_img.convert('RGB')
            person_img.save(person_path, 'JPEG')
            logger.debug(f"사람 이미지 저장 완료: {person_path}")
            
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_product:
            product_path = tmp_product.name
            # RGB로 변환하여 저장
            product_img = product_pil
            if product_img.mode != 'RGB':
                product_img = product_img.convert('RGB')
            product_img.save(product_path, 'JPEG')
            logger.debug(f"상품 이미지 저장 완료: {product_path}")
        
        # 배경 이미지 저장 (있는 경우)
        background_path = None
        if background_pil is not None:
            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_bg:
                background_path = tmp_bg.name
                # RGB로 변환하여 저장
                background_img = background_pil
                if background_img.mode != 'RGB':
                    background_img = background_img.convert('RGB')
                background_img.save(background_path, 'JPEG')
                logger.debug(f"배경 이미지 저장 완료: {background_path}")

        # 이미지 합성 실행
        result_path, response_text = merge_images(
            person_img_path=person_path,
            product_img_path=product_path,
            background_img_path=background_path,
            prompt=prompt
        )

        # 이미지 반환 및 임시 파일 정리
        if result_path:
            logger.debug(f"이미지 생성 완료. 경로: {result_path}")
            try:
                result_img = Image.open(result_path)
                if result_img.mode == "RGBA":
                    result_img = result_img.convert("RGB")
                
                # 임시 파일 정리
                try:
                    os.unlink(person_path)
                    os.unlink(product_path)
                    if background_path:
                        os.unlink(background_path)
                except Exception as e:
                    logger.warning(f"임시 파일 삭제 중 오류: {str(e)}")
                
                return [result_img], response_text or "이미지가 성공적으로 생성되었습니다."
            except Exception as e:
                logger.exception(f"결과 이미지 로드 중 오류: {str(e)}")
                return [], f"결과 이미지 처리 중 오류: {str(e)}"
        else:
            logger.error("merge_images 함수에서 None 반환됨.")
            return [], response_text or "이미지 생성에 실패했습니다. 다른 이미지나 프롬프트로 시도해보세요."

    except Exception as e:
        logger.exception("process_images_and_prompt 함수에서 오류 발생:")
        return [], f"오류 발생: {str(e)}"  # 오류 시 빈 리스트와 오류 메시지 반환