Create solle_img_generation.py

src/solle_img_generation.py

@@ -0,0 +1,43 @@
+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