49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
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 |