Update app.py

app.py

@@ -3,8 +3,8 @@ import os
 import mimetypes
 import tempfile
 import time
-from google import genai
-from google.genai import types
+import google.generativeai as genai
+from google.generativeai import types
 import gradio as gr
 
 def generate_image(api_key, prompt, file_name):
@@ -20,55 +20,49 @@ def generate_image(api_key, prompt, file_name):
     
     try:
         # Set up the client
-        client = genai.Client(api_key=api_key)
-        model = "gemini-2.0-flash-exp-image-generation"
+        genai.configure(api_key=api_key)
+        model = genai.GenerativeModel("gemini-2.0-flash-exp-image-generation")
         
-        # Create content
-        contents = [
-            types.Content(
-                role="user",
-                parts=[types.Part.from_text(text=prompt)],
-            ),
-        ]
-        
-        # Configure generation settings
-        generate_content_config = types.GenerateContentConfig(
-            response_modalities=["image", "text"],
-            safety_settings=[
-                types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="BLOCK_NONE"),
-                types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="BLOCK_NONE"),
-                types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="BLOCK_NONE"),
-                types.SafetySetting(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="BLOCK_NONE"),
-                types.SafetySetting(category="HARM_CATEGORY_CIVIC_INTEGRITY", threshold="BLOCK_NONE"),
-            ],
-            response_mime_type="text/plain",
-        )
+        # Create generation config
+        generation_config = {
+            "response_modalities": ["image", "text"],
+            "safety_settings": {
+                "HARM_CATEGORY_HARASSMENT": "BLOCK_NONE",
+                "HARM_CATEGORY_HATE_SPEECH": "BLOCK_NONE",
+                "HARM_CATEGORY_SEXUALLY_EXPLICIT": "BLOCK_NONE",
+                "HARM_CATEGORY_DANGEROUS_CONTENT": "BLOCK_NONE",
+                "HARM_CATEGORY_CIVIC_INTEGRITY": "BLOCK_NONE",
+            }
+        }
         
         # Create a temporary directory to store the image
         temp_dir = tempfile.mkdtemp()
         generated_image_path = None
         generation_text = ""
         
-        # Stream the content
-        for chunk in client.models.generate_content_stream(
-            model=model,
-            contents=contents,
-            config=generate_content_config,
-        ):
-            if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
-                continue
-                
-            if chunk.candidates[0].content.parts[0].inline_data:
-                inline_data = chunk.candidates[0].content.parts[0].inline_data
-                file_extension = mimetypes.guess_extension(inline_data.mime_type) or '.jpg'
-                generated_image_path = os.path.join(temp_dir, f"{file_name}{file_extension}")
-                
-                with open(generated_image_path, "wb") as f:
-                    f.write(inline_data.data)
-                
-                generation_text += f"Image of type {inline_data.mime_type} generated successfully."
-            else:
-                generation_text += chunk.text
+        # Generate the content
+        response = model.generate_content(
+            prompt,
+            generation_config=generation_config,
+            stream=True
+        )
+        
+        for chunk in response:
+            if hasattr(chunk, 'candidates') and chunk.candidates:
+                candidate = chunk.candidates[0]
+                if hasattr(candidate, 'content') and candidate.content:
+                    for part in candidate.content.parts:
+                        if hasattr(part, 'inline_data') and part.inline_data:
+                            inline_data = part.inline_data
+                            file_extension = mimetypes.guess_extension(inline_data.mime_type) or '.jpg'
+                            generated_image_path = os.path.join(temp_dir, f"{file_name}{file_extension}")
+                            
+                            with open(generated_image_path, "wb") as f:
+                                f.write(inline_data.data)
+                            
+                            generation_text += f"Image of type {inline_data.mime_type} generated successfully."
+                        elif hasattr(part, 'text') and part.text:
+                            generation_text += part.text
         
         if generated_image_path:
             return generated_image_path, generation_text