Update image_generation.py

image_generation.py

@@ -4,6 +4,8 @@ import json
 import streamlit as st
 import google.generativeai as genai
 import google.api_core.exceptions
+from google import genai
+from google.genai import types
 from typing import List, Dict, Optional
 from PIL import Image
 import io
@@ -31,50 +33,107 @@ def save_binary_file(file_name: str, data: bytes):
         print(f"❌ Error saving file {file_name}: {e}")
 
 # --- IMAGE GENERATION FUNCTION ---
+def pil_image_to_part(image: Image.Image) -> types.Part:
+    img_byte_arr = io.BytesIO()
+    image.save(img_byte_arr, format="JPEG")
+    return types.Part.from_bytes(
+        data=img_byte_arr.getvalue(),
+        mime_type="image/jpeg",
+    )
+
 def generate_image_with_gemini(
     prompt: str,
     output_file_base: str,
     context_image: Optional[Image.Image] = None
 ) -> Optional[str]:
-    """
-    Generates an image and now specifically handles ResourceExhausted errors.
-    """
+    if not client:
+        return None
+
     print(f"--- 🎨 Generating image for prompt: '{prompt[:70]}...' ---")
+
     try:
-        model = genai.GenerativeModel(model_name="gemini-2.0-flash-preview-image-generation")
+        model = "gemini-2.0-flash-preview-image-generation"
 
-        content_parts = []
+        # Build contents
         if context_image:
-            content_parts.extend([prompt, context_image])
+            system_prompt = (
+                "You are a master storyboard artist creating a visual story sequence.\n"
+                "IMPORTANT: You MUST generate an image for every request.\n"
+                "Create a visually consistent image that follows the art style and character design of the provided reference image.\n"
+                "Maintain consistency in character appearance, art style, color palette, and lighting.\n"
+                "Style: Cinematic, epic fantasy digital painting with rich details and dramatic lighting.\n"
+                "Generate an image that illustrates the following scene:"
+            )
+            print(" -> Using previous image as context for consistent styling.")
         else:
-            content_parts.append(prompt)
+            system_prompt = (
+                "You are a master storyboard artist creating the opening scene of a visual story.\n"
+                "IMPORTANT: You MUST generate an image for this request.\n"
+                "Create a stunning, cinematic image in an epic fantasy digital painting style with rich, detailed artwork and dramatic lighting.\n"
+                "This is the first scene of the story. Generate an image that illustrates:"
+            )
+
+        content_parts = [
+            types.Part.from_text(text=system_prompt)
+        ]
+        if context_image:
+            content_parts.append(pil_image_to_part(context_image))
 
-        response = model.generate_content(
-            contents=content_parts,
-            stream=True
+        image_instruction = f"""CREATE AN IMAGE NOW:
+
+        {prompt}
+        
+        Remember: You must generate a visual image, not text. Create the artwork described above."""
+        content_parts.append(types.Part.from_text(text=image_instruction))
+
+        contents = [
+            types.Content(
+                role="user",
+                parts=content_parts,
+            )
+        ]
+
+        # CRITICAL: request both IMAGE and TEXT
+        generate_content_config = types.GenerateContentConfig(
+            response_modalities=["IMAGE", "TEXT"]
         )
 
         saved_file_path = None
-        for chunk in response:
-            if chunk.parts and chunk.parts[0].inline_data:
-                part = chunk.parts[0]
-                data = part.inline_data.data
-                mime_type = part.inline_data.mime_type
-                file_extension = mimetypes.guess_extension(mime_type) or ".jpg"
-                full_file_name = f"{output_file_base}{file_extension}"
-                save_binary_file(full_file_name, data)
-                saved_file_path = full_file_name
-                print(f"✅ Successfully generated and saved image: {full_file_name}")
-                return saved_file_path
-        
-        return None # Return None if no image was generated
+        text_responses = []
+
+        # Stream the response
+        for chunk in client.models.generate_content_stream(
+            model=model,
+            contents=contents,
+            config=generate_content_config,
+        ):
+            cand = getattr(chunk, "candidates", None)
+            if not cand or not cand.content or not cand.content.parts:
+                continue
+
+            for part in cand.content.parts:
+                # Image bytes
+                if getattr(part, "inline_data", None) and getattr(part.inline_data, "data", None):
+                    inline_data = part.inline_data
+                    data_buffer = inline_data.data
+                    file_extension = mimetypes.guess_extension(inline_data.mime_type) or ".jpg"
+                    full_file_name = f"{output_file_base}{file_extension}"
+                    save_binary_file(full_file_name, data_buffer)
+                    saved_file_path = full_file_name
+                    print(f"✅ Successfully generated and saved image: {full_file_name}")
+
+                # Text side-channel
+                if getattr(part, "text", None):
+                    text_responses.append(part.text)
+
+        if not saved_file_path and text_responses:
+            print(f"⚠️ No image generated. API returned text: {' '.join(text_responses)}")
+
+        return saved_file_path
 
-    except google.api_core.exceptions.ResourceExhausted as e:
-        # --- CATCH and signal the rate limit error ---
-        print(f"🔴 RATE LIMIT EXCEEDED. The script will wait and retry.")
-        return "RATE_LIMIT_EXCEEDED"
     except Exception as e:
         print(f"❌ An error occurred during the Gemini API call: {e}")
+        import traceback
         traceback.print_exc()
         return None