Update image_generation.py

image_generation.py

@@ -3,14 +3,13 @@ import mimetypes
 import json
 import streamlit as st
 import google.generativeai as genai
-from typing import List, Dict, Optional
+import google.api_core.exceptions
 from PIL import Image
 import io
 import traceback
 import time
 
-# Initialization
-client = None
+# --- INITIALIZATION ---
 try:
     api_key = st.secrets.get("GEMINI_API_KEY") or os.getenv("GEMINI_API_KEY")
     if api_key:
@@ -31,64 +30,54 @@ def save_binary_file(file_name: str, data: bytes):
         print(f"❌ Error saving file {file_name}: {e}")
 
 # --- IMAGE GENERATION FUNCTION ---
-# --- CORRECTED IMAGE GENERATION FUNCTION ---
 def generate_image_with_gemini(
     prompt: str,
     output_file_base: str,
     context_image: Optional[Image.Image] = None
 ) -> Optional[str]:
     """
-    Generates an image using the correct syntax for your library version.
+    Generates an image and now specifically handles ResourceExhausted errors.
     """
     print(f"--- 🎨 Generating image for prompt: '{prompt[:70]}...' ---")
     try:
-        # 1. Create the model instance directly
         model = genai.GenerativeModel(model_name="gemini-2.5-flash-image-preview")
 
-        # 2. Create the content list
         content_parts = []
         if context_image:
-            print(" -> Using previous image as context for consistent styling.")
             content_parts.extend([prompt, context_image])
         else:
             content_parts.append(prompt)
 
-        # 3. Make the call on the model object
         response = model.generate_content(
             contents=content_parts,
             stream=True
         )
 
         saved_file_path = None
-        text_responses = []
-
-        # 4. Process the response
         for chunk in response:
-            if chunk.parts:
-                for part in chunk.parts:
-                    if part.inline_data:
-                        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}")
-                        break 
-                    elif part.text:
-                        text_responses.append(part.text)
+            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
         
-        if not saved_file_path and text_responses:
-            print(f"⚠️ No image generated. API returned text: {''.join(text_responses)}")
-
-        return saved_file_path
+        return None # Return None if no image was generated
 
+    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}")
         traceback.print_exc()
         return None
 
+# --- MAIN FUNCTION ---
 def generate_all_images_from_file(
     json_path: str,
     output_dir: str = "generated_images",
@@ -101,67 +90,57 @@ def generate_all_images_from_file(
         print(f"❌ Error reading or parsing {json_path}: {e}")
         return []
 
-    if not os.path.exists(output_dir):
-        os.makedirs(output_dir)
+    # --- ADDED: Daily Limit Guard ---
+    if len(multimedia_data) > 100:
+        print(f"⚠️ WARNING: Your request for {len(multimedia_data)} images exceeds the daily limit of 100.")
+        print("Processing will stop to avoid hitting the daily rate limit.")
+        return multimedia_data
 
-    previous_image = None
-    successful_generations = 0
+    os.makedirs(output_dir, exist_ok=True)
 
+    previous_image = None
     for i, item in enumerate(multimedia_data):
         print(f"\n{'='*60}")
-        print(f"Processing item {i+1}/{len(multimedia_data)} - {'FIRST IMAGE' if i == 0 else 'CONTINUATION'}")
-        print(f"{'='*60}")
+        print(f"Processing item {i+1}/{len(multimedia_data)}")
         
         image_prompt = item.get("image_prompt")
         if not image_prompt or "Error:" in image_prompt:
-            print(f"⚠️ Skipping item {i}: No valid image prompt found")
             item["image_path"] = None
             continue
 
-        print(f"📝 Image prompt: {image_prompt[:100]}{'...' if len(image_prompt) > 100 else ''}")
         file_base_path = os.path.join(output_dir, f"image_{i:03d}")
 
-        # Your system prompts are now integrated into the main prompt for simplicity
-        if previous_image:
-             full_prompt = f"""You are a master storyboard artist creating a visual story sequence. IMPORTANT: You MUST generate an image. Maintain consistency with the provided image. Style: Cinematic, epic fantasy digital painting. Generate an image that illustrates the following scene: {image_prompt}"""
-        else:
-            full_prompt = f"""You are a master storyboard artist creating the opening scene. IMPORTANT: You MUST generate an image. Style: Stunning, cinematic, epic fantasy digital painting with rich detail and dramatic lighting. Generate an image that illustrates: {image_prompt}"""
-
-
-        saved_image_path = generate_image_with_gemini(
-            full_prompt,
-            file_base_path,
-            context_image=previous_image
-        )
-
+        # --- Retry Logic ---
+        max_retries = 3
+        for attempt in range(max_retries):
+            saved_image_path = generate_image_with_gemini(
+                image_prompt, file_base_path, context_image=previous_image
+            )
+
+            if saved_image_path == "RATE_LIMIT_EXCEEDED":
+                wait_time = 60 # Wait for 60 seconds
+                print(f"Retrying in {wait_time} seconds (Attempt {attempt + 1}/{max_retries})...")
+                time.sleep(wait_time)
+                continue
+            else:
+                break
+        
         item["image_path"] = saved_image_path
 
-        if saved_image_path and os.path.exists(saved_image_path):
+        if saved_image_path and saved_image_path != "RATE_LIMIT_EXCEEDED":
             try:
                 previous_image = Image.open(saved_image_path)
-                successful_generations += 1
-                print(f"✅ Loaded image {saved_image_path} as context for next generation.")
             except Exception as e:
-                print(f"⚠️ Could not load image {saved_image_path} for context. Error: {e}")
                 previous_image = None
         else:
             previous_image = None
-            print("❌ No image generated for this item!")
+            print("❌ Failed to generate an image for this item after retries.")
 
-        time.sleep(2)
-
-    try:
-        with open(output_json_path, 'w', encoding='utf-8') as f:
-            json.dump(multimedia_data, f, indent=2, ensure_ascii=False)
-        print(f"\n--- ✅ Image generation finished. Updated data saved to {output_json_path}. ---")
-        print(f"Successfully generated {successful_generations}/{len(multimedia_data)} images.")
-        
-        print(f"\n📊 GENERATION SUMMARY:")
-        for i, item in enumerate(multimedia_data):
-            status = "✅ SUCCESS" if item.get("image_path") else "❌ FAILED"
-            print(f" Item {i+1}: {status}")
-
-    except Exception as e:
-        print(f"❌ Error saving updated JSON: {e}")
+        # --- Increased delay for RPM limit ---
+        print("Waiting 7 seconds before next request...")
+        time.sleep(7)
 
-    return multimedia_data
+    with open(output_json_path, 'w', encoding='utf-8') as f:
+        json.dump(multimedia_data, f, indent=2, ensure_ascii=False)
+    
+    return multimedia_data