Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
-import google.generativeai as genai
+from google import genai
+from google.genai import types
 from PIL import Image
 from io import BytesIO
 import base64
@@ -19,40 +20,43 @@ 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(model, prompt, style):
+def enhance_prompt(client, prompt, style):
+    model = client.get_model('gemini-pro')
     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(model, enhanced_prompt, style, negative_prompt):
+def generate_image(client, enhanced_prompt, style, negative_prompt):
     full_prompt = f"{enhanced_prompt}\nStyle: {style}\nNegative prompt: {negative_prompt}"
     
-    response = model.generate_content([full_prompt, "generate an image"])
+    response = client.models.generate_images(
+        model='imagen-3.0-generate-002',
+        prompt=full_prompt,
+        config=types.GenerateImagesConfig(
+            number_of_images=1,
+        )
+    )
     
-    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
+    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()
+    
+    return f"data:image/png;base64,{img_str}"
 
 def process_and_generate(api_key, prompt, style, negative_prompt):
-    genai.configure(api_key=api_key)
-    model = genai.GenerativeModel('gemini-pro-vision')
-    enhanced_prompt = enhance_prompt(model, prompt, style)
-    image_url = generate_image(model, enhanced_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)
     return image_url, enhanced_prompt
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Google Gemini Image Generator")
+    gr.Markdown("# Google Imagen 3 Image Generator")
     
     with gr.Row():
         with gr.Column(scale=1):