Update image_generation.py

image_generation.py

@@ -1,3 +1,5 @@
+# image_generation.py
+
 import os
 import mimetypes
 import json
@@ -7,11 +9,14 @@ import time
 import traceback
 from PIL import Image
 from typing import List, Dict, Optional
+
+# CORRECT IMPORTS FOR THE 'google-genai' SDK
 from google import genai
 from google.generativeai import types
 from google.api_core import exceptions
 
 # --- Client Initialization ---
+# This section initializes the client using the secrets from your HF Space.
 client = None
 try:
     api_key = st.secrets.get("GEMINI_API_KEY")
@@ -24,10 +29,11 @@ try:
         st.stop()
 except Exception as e:
     print(f"❌ Error initializing Google AI client: {e}")
-    st.error(f"Error initializing Google AI client: {e}")
+    st.error(f"An unexpected error occurred during client initialization: {e}")
     st.stop()
 
 # --- Helper Functions ---
+
 def save_binary_file(file_name: str, data: bytes):
     """Saves binary data to a file."""
     try:
@@ -42,20 +48,14 @@ def pil_image_to_part(image: Image.Image) -> types.Part:
     img_byte_arr = io.BytesIO()
     image.save(img_byte_arr, format='JPEG')
     img_bytes = img_byte_arr.getvalue()
-    # This call was correct and remains the same.
-    return types.Part.from_bytes(
-        data=img_bytes,
-        mime_type='image/jpeg'
-    )
+    return types.Part(inline_data=types.Blob(mime_type="image/jpeg", data=img_bytes))
 
 def generate_image_with_gemini(
     prompt: str,
     output_file_base: str,
     context_image: Optional[Image.Image] = None
 ) -> Optional[str]:
-    """
-    Generates an image using the Gemini API, now with all corrections.
-    """
+    """Generates an image using the Gemini API with the corrected SDK calls."""
     if not client:
         print("❌ Gemini client not initialized.")
         return None
@@ -66,7 +66,6 @@ def generate_image_with_gemini(
         model_name = "gemini-2.0-flash-preview-image-generation"
         content_parts = []
 
-        # Define system prompts
         if context_image:
             system_prompt = """You are a master storyboard artist creating a visual story sequence.
             IMPORTANT: You MUST generate an image for every request. Create a visually consistent image that follows the art style and character design of the provided reference image. Maintain consistency in:
@@ -83,14 +82,14 @@ def generate_image_with_gemini(
             - Dramatic lighting and atmosphere
             - High-quality digital painting aesthetic
             This is the first scene of the story. Generate an image that illustrates:"""
-
-        content_parts.append(types.Part.from_text(system_prompt))
+        
+        content_parts.append(types.Part(text=system_prompt))
         
         if context_image:
             content_parts.append(pil_image_to_part(context_image))
         
-        image_instruction = f"CREATE AN IMAGE NOW:\n{prompt}\nRemember: You must generate a visual image, not text."
-        content_parts.append(types.Part.from_text(image_instruction))
+        image_instruction = f"CREATE AN IMAGE NOW:\n{prompt}\nRemember: You must generate a visual image."
+        content_parts.append(types.Part(text=image_instruction))
 
         contents = [types.Content(role="user", parts=content_parts)]
 
@@ -105,7 +104,6 @@ def generate_image_with_gemini(
         )
 
         saved_file_path = None
-        text_responses = []
         for chunk in stream:
             if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
                 continue
@@ -117,15 +115,11 @@ def generate_image_with_gemini(
                     full_file_name = f"{output_file_base}{file_extension}"
                     save_binary_file(full_file_name, inline_data.data)
                     saved_file_path = full_file_name
-                elif part.text:
-                    text_responses.append(part.text)
-
+        
         if saved_file_path:
             print(f"✅ Successfully generated and saved image: {saved_file_path}")
-        elif text_responses:
-            print(f"⚠️ No image generated. API returned text only: {' '.join(text_responses)}")
         else:
-            print("⚠️ No image or text was returned from the API.")
+            print("⚠️ No image was returned from the API.")
 
         return saved_file_path
 
@@ -133,27 +127,19 @@ def generate_image_with_gemini(
         print(f"❌ API Invalid Argument Error: {e}")
         traceback.print_exc()
         return None
-    except TypeError as e:
-        print(f"❌ TypeError during API call: {e}")
-        print("  -> This often indicates an incorrect way of creating API objects like 'Part'. Please double-check SDK documentation.")
-        traceback.print_exc()
-        return None
     except Exception as e:
-        print(f"❌ An unexpected error occurred: {e}")
+        print(f"❌ An unexpected error occurred during the Gemini API call: {e}")
         traceback.print_exc()
         return None
 
-def generate_all_images_from_file(
-    json_path: str,
-    output_dir: str = "generated_images",
-    output_json_path: str = "multimedia_data_with_images.json"
-) -> List[Dict]:
+def generate_all_images_from_file(json_path: str, output_dir: str, output_json_path: str):
+    """Main loop to process a JSON file and generate images."""
     try:
         with open(json_path, 'r', encoding='utf-8') as f:
             multimedia_data = json.load(f)
     except (FileNotFoundError, json.JSONDecodeError) as e:
         print(f"❌ Error reading or parsing {json_path}: {e}")
-        return []
+        return
 
     if not os.path.exists(output_dir):
         os.makedirs(output_dir)
@@ -164,42 +150,29 @@ def generate_all_images_from_file(
         print(f"\n{'='*60}\nProcessing item {i+1}/{len(multimedia_data)}\n{'='*60}")
         image_prompt = item.get("image_prompt")
 
-        if not image_prompt or "Error:" in image_prompt:
-            print(f"⚠️ Skipping item {i+1}: No valid image prompt found.")
+        if not image_prompt:
             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}")
-
         saved_image_path = generate_image_with_gemini(
-            image_prompt,
-            file_base_path,
-            context_image=previous_image
+            image_prompt, file_base_path, context_image=previous_image
         )
 
         item["image_path"] = saved_image_path
 
-        if saved_image_path and os.path.exists(saved_image_path):
+        if saved_image_path:
             try:
                 previous_image = Image.open(saved_image_path)
                 successful_generations += 1
-                print(f"✅ Loaded image {saved_image_path} as context for the next generation.")
             except Exception as e:
-                print(f"⚠️ Could not load image {saved_image_path} for context. Error: {e}")
                 previous_image = None
         else:
-            print("❌ No image was generated for this item. Context will be reset.")
             previous_image = None
-
+        
         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.")
-    except Exception as e:
-        print(f"❌ Error saving updated JSON: {e}")
-
-    return multimedia_data
+    with open(output_json_path, 'w', encoding='utf-8') as f:
+        json.dump(multimedia_data, f, indent=2, ensure_ascii=False)
+    
+    print(f"\n--- ✅ Finished. Generated {successful_generations}/{len(multimedia_data)} images. ---")