45 lines
1.3 KiB
Python
45 lines
1.3 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
|
|
|
|
|
|
def gemini_image(prompt, folder=None):
|
|
from custom_logger.main_log import main_logger as LOG
|
|
|
|
image_path = ''
|
|
client = genai.Client(api_key="AIzaSyB7tu67y9gOkJkpQtvI5OAYSzUzwv9qwnE")
|
|
|
|
for i in range(3):
|
|
response = client.models.generate_content(
|
|
model="gemini-2.0-flash-preview-image-generation",
|
|
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 |