Update image_gen.py

image_gen.py

@@ -35,6 +35,7 @@ from PIL import ImageFont, ImageDraw, Image
 import seaborn as sns
 
 
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -207,4 +208,62 @@ def generate_image_with_retry(prompt_text, style, model="hf", max_retries=3):
             logger.error(f"Attempt {attempt+1} failed: {e}")
             if attempt == max_retries - 1:
                 raise
-    return None, None
+    return None, None
+
+# edit image function
+def edit_section_image(image_url: str, gemini_prompt: str):
+    """
+    Downloads the existing image from image_url, uses Google Gemini to edit it
+    according to gemini_prompt, and returns the new edited image (as a PIL.Image).
+    """
+    try:
+        # 1) Download the original image
+        resp = requests.get(image_url)
+        if resp.status_code != 200:
+            logger.error(f"Failed to download image from {image_url}")
+            return None
+
+        original_image = Image.open(io.BytesIO(resp.content))
+
+        # 2) Initialize Gemini client
+        g_api_key = os.getenv("GEMINI")
+        if not g_api_key:
+            logger.error("GEMINI_API_KEY not found in environment variables")
+            print("Google Gemini API key is missing. Please set the GEMINI_API_KEY environment variable.")
+            return None
+
+        genai.configure(api_key=g_api_key)
+        client = genai.Client(api_key=g_api_key)
+
+        # 3) Prepare the prompt_with_image: a list with [ prompt_text, PIL.Image ]
+        prompt_with_image = [
+            gemini_prompt,
+            original_image
+        ]
+
+        # 4) Call the Gemini model to edit the image
+        response = client.models.generate_content(
+            model="models/gemini-2.0-flash-exp",
+            contents=prompt_with_image,
+            config=types.GenerateContentConfig(
+                response_modalities=['Text', 'Image']
+            )
+        )
+
+        # 5) Extract the edited image from the response
+        #    Typically, the 'response' might have text + image. We want the image part.
+        edited_image = None
+        for part in response.candidates[0].content.parts:
+            if part.inline_data is not None:
+                edited_image = Image.open(io.BytesIO(part.inline_data.data))
+                break
+
+        if not edited_image:
+            logger.error("No edited image found in Gemini response")
+            return None
+
+        return edited_image
+
+    except Exception as e:
+        logger.error(f"Error editing section image: {e}")
+        return None