GLASSES_AI_SERVER/custom_apps/gemini/main.py

import os
import time

from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO

from main_rest.app.utils.date_utils import D
from const import TEMP_FOLDER, API_KEY_MANAGER

def gemini_image(prompt, folder=None):
    from custom_logger.main_log import main_logger as LOG
    
    image_path = ''
    api_key = API_KEY_MANAGER.get_api_key()
    
    if api_key is None:
        raise Exception("API 키 세팅 필요! - 서버를 다시 구동하거나, API키 파일을 확인")
    
    client = genai.Client(api_key=API_KEY_MANAGER.get_api_key()) # a2tec key

    for i in range(3):
        response = client.models.generate_content(
            model="gemini-2.5-flash-image",
            contents=prompt,
            config=types.GenerateContentConfig(
            response_modalities=['TEXT', 'IMAGE']
            )
        )
        if folder == None:
            folder = TEMP_FOLDER
        
        if not os.path.exists(folder):
            os.makedirs(folder)
        
        for part in response.candidates[0].content.parts:
            if part.inline_data is not None:
                image = Image.open(BytesIO((part.inline_data.data)))
                image_path = os.path.join(folder,f"gemini_{D.date_file_name()}_query.png")
                image.save(image_path)
        
        LOG.info(f"image generate : {image_path}")
        time.sleep(2)
        
        if os.path.exists(image_path):
            return image_path
    
    return image_path