Update app.py

app.py

@@ -31,20 +31,22 @@ 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_images(
-        prompt=full_prompt,
-        number_of_images=1,
-    )
-    
-    image_bytes = response.images[0].image_bytes
-    image = Image.open(io.BytesIO(image_bytes))
+    response = model.generate_content(full_prompt)
     
-    # Convert PIL Image to base64 string
-    buffered = io.BytesIO()
-    image.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
+    if response.parts:
+        for part in response.parts:
+            if part.image:
+                image_bytes = part.image.to_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}"
     
-    return f"data:image/png;base64,{img_str}"
+    return "Image generation failed"
 
 def process_and_generate(api_key, prompt, style, negative_prompt):
     genai.configure(api_key=api_key)