solle_img_generation / src /solle_img_generation.py
QuentinL52's picture
Create solle_img_generation.py
ca6e0ed verified
import os
import mimetypes
from google import genai
from google.genai import types
from .config import load_prompt_template
class imggenrator:
def __init__(self):
self.api_key = os.getenv("IMG_GEN_API_KEY")
self.model = os.getenv("model")
self.client = genai.Client(api_key=self.api_key)
# Notez que le paramètre `save_fn` n'est plus nécessaire.
def generate(self, image_path: str, user_input_text: str):
with open(image_path, "rb") as f:
image_data = f.read()
image_mime_type = mimetypes.guess_type(image_path)[0]
full_prompt = load_prompt_template(image_path, user_input_text)
contents = [
types.Content(
role="user",
parts=[
types.Part.from_bytes(data=image_data, mime_type=image_mime_type),
types.Part.from_text(text=full_prompt),
],
),
]
config = types.GenerateContentConfig(response_modalities=["IMAGE", "TEXT"])
generated_images_data = []
for chunk in self.client.models.generate_content_stream(
model=self.model,
contents=contents,
config=config,
):
if chunk.candidates and chunk.candidates[0].content and chunk.candidates[0].content.parts:
for part in chunk.candidates[0].content.parts:
if hasattr(part, "inline_data") and part.inline_data:
generated_images_data.append(part.inline_data.data)
elif hasattr(part, "text") and part.text:
print(part.text)
return generated_images_data