Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
-from google import genai
-from google.genai import types
+import google.generativeai as genai
 from PIL import Image
 from io import BytesIO
 import base64
@@ -20,43 +19,41 @@ extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, si
 cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face
 """
 
-def enhance_prompt(client, prompt, style):
-    model = client.get_model('gemini-pro')
+def enhance_prompt(model, prompt, style):
     enhanced_prompt_request = f"Enhance the following prompt for image generation in the style of {style}. Make it more detailed and vivid, while keeping the original intent: '{prompt}'"
     
     response = model.generate_content(enhanced_prompt_request)
     
     return response.text
 
-def generate_image(client, enhanced_prompt, style, negative_prompt):
+def generate_image(model, enhanced_prompt, style, negative_prompt):
     full_prompt = f"{enhanced_prompt}\nStyle: {style}\nNegative prompt: {negative_prompt}"
     
-    response = client.models.generate_images(
-        model='imagen-3.0-generate-002',
-        prompt=full_prompt,
-        config=types.GenerateImagesConfig(
-            number_of_images=1,
-        )
-    )
-    
-    image_bytes = response.generated_images[0].image.image_bytes
-    image = Image.open(BytesIO(image_bytes))
-    
-    # Convert PIL Image to base64 string
-    buffered = BytesIO()
-    image.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
+    response = model.generate_content([full_prompt, "generate an image"])
     
-    return f"data:image/png;base64,{img_str}"
+    if response.parts[0].image:
+        image_bytes = response.parts[0].image.to_bytes()
+        image = Image.open(BytesIO(image_bytes))
+        
+        # Convert PIL Image to base64 string
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode()
+        
+        return f"data:image/png;base64,{img_str}"
+    else:
+        return None
 
 def process_and_generate(api_key, prompt, style, negative_prompt):
-    client = genai.Client(api_key=api_key)
-    enhanced_prompt = enhance_prompt(client, prompt, style)
-    image_url = generate_image(client, enhanced_prompt, style, negative_prompt)
+    genai.configure(api_key=api_key)
+    text_model = genai.GenerativeModel('gemini-pro')
+    image_model = genai.GenerativeModel('gemini-pro-vision')
+    enhanced_prompt = enhance_prompt(text_model, prompt, style)
+    image_url = generate_image(image_model, enhanced_prompt, style, negative_prompt)
     return image_url, enhanced_prompt
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Google Imagen 3 Image Generator")
+    gr.Markdown("# Google Gemini Image Generator")
     
     with gr.Row():
         with gr.Column(scale=1):