Spaces:
Sleeping
Sleeping
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 |