Update app.py

app.py

@@ -1,6 +1,9 @@
 import gradio as gr
-from openai import OpenAI
-import httpx
+from google import genai
+from google.genai import types
+from PIL import Image
+from io import BytesIO
+import base64
 
 # List of popular styles
 STYLES = [
@@ -18,43 +21,46 @@ cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, di
 """
 
 def enhance_prompt(client, prompt, style):
-    enhanced_prompt_request = f"Enhance the following prompt for DALL-E 3 image generation in the style of {style}. Make it more detailed and vivid, while keeping the original intent: '{prompt}'"
+    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 = client.chat.completions.create(
-        model="gpt-4",
-        messages=[
-            {"role": "system", "content": "You are an expert at creating detailed, vivid prompts for image generation."},
-            {"role": "user", "content": enhanced_prompt_request}
-        ]
-    )
+    response = model.generate_content(enhanced_prompt_request)
     
-    return response.choices[0].message.content.strip()
+    return response.text
 
 def generate_image(client, enhanced_prompt, style, negative_prompt):
     full_prompt = f"{enhanced_prompt}\nStyle: {style}\nNegative prompt: {negative_prompt}"
     
-    response = client.images.generate(
-        model="dall-e-3",
+    response = client.models.generate_images(
+        model='imagen-3.0-generate-002',
         prompt=full_prompt,
-        size="1024x1024",
-        quality="standard",
-        n=1
+        config=types.GenerateImagesConfig(
+            number_of_images=1,
+        )
     )
     
-    return response.data[0].url
+    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):
-    client = OpenAI(api_key=api_key, http_client=httpx.Client())
+    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("# DALL-E 3 Image Generator")
+    gr.Markdown("# Google Imagen Image Generator")
     
     with gr.Row():
         with gr.Column(scale=1):
-            api_key = gr.Textbox(label="OpenAI API Key", type="password")
+            api_key = gr.Textbox(label="Google AI API Key", type="password")
             prompt = gr.Textbox(label="Prompt")
             style = gr.Dropdown(label="Style", choices=STYLES)
             negative_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT)