Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import google.generativeai as genai
 from PIL import Image
-from io import BytesIO
+import io
 import base64
 
 # List of popular styles
@@ -19,41 +19,42 @@ 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(prompt, style):
+    model = genai.GenerativeModel("gemini-2.0-flash-lite")
     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(enhanced_prompt, style, negative_prompt):
+    model = genai.GenerativeModel("imagen-3.0-generate-002")
     full_prompt = f"{enhanced_prompt}\nStyle: {style}\nNegative prompt: {negative_prompt}"
     
-    response = model.generate_content([full_prompt, "generate an image"])
+    response = model.generate_images(
+        prompt=full_prompt,
+        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.images[0].image_bytes
+    image = Image.open(io.BytesIO(image_bytes))
+    
+    # Convert PIL Image to base64 string
+    buffered = io.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)
-    text_model = genai.GenerativeModel('gemini-2.0-flash-lite')
-    image_model = genai.GenerativeModel('imagen-3.0-generate-002')
-    enhanced_prompt = enhance_prompt(text_model, prompt, style)
-    image_url = generate_image(image_model, enhanced_prompt, style, negative_prompt)
+    
+    enhanced_prompt = enhance_prompt(prompt, style)
+    image_url = generate_image(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):