Update app.py

app.py

@@ -1,8 +1,5 @@
 import gradio as gr
 import google.generativeai as genai
-from PIL import Image
-import io
-import base64
 import requests
 
 # List of popular styles
@@ -20,7 +17,8 @@ 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(prompt, style):
+def enhance_prompt(google_api_key, prompt, style):
+    genai.configure(api_key=google_api_key)
     model = genai.GenerativeModel("gemini-2.0-flash-lite")
     enhanced_prompt_request = f"""
     Task: Enhance the following prompt for image generation.
@@ -39,10 +37,8 @@ def enhance_prompt(prompt, style):
     
     response = model.generate_content(enhanced_prompt_request)
     
-    # Extract only the enhanced prompt, removing any potential explanations or extra text
     enhanced_prompt = response.text.strip()
     
-    # If the response starts with "Enhanced prompt:" or similar, remove it
     prefixes_to_remove = ["Enhanced prompt:", "Here's the enhanced prompt:", "The enhanced prompt is:"]
     for prefix in prefixes_to_remove:
         if enhanced_prompt.lower().startswith(prefix.lower()):
@@ -50,7 +46,6 @@ def enhance_prompt(prompt, style):
     
     return enhanced_prompt
 
-
 def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt):
     url = "https://api.stability.ai/v2beta/stable-image/generate/sd3"
     
@@ -79,12 +74,13 @@ def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt):
     else:
         return f"Image generation failed: {response.text}"
 
-def process_and_generate(google_api_key, stability_api_key, prompt, style, negative_prompt):
-    genai.configure(api_key=google_api_key)
-    
-    enhanced_prompt = enhance_prompt(prompt, style)
+def process_prompt(google_api_key, prompt, style):
+    enhanced_prompt = enhance_prompt(google_api_key, prompt, style)
+    return enhanced_prompt
+
+def generate_from_enhanced(stability_api_key, enhanced_prompt, style, negative_prompt):
     image_url = generate_image(stability_api_key, enhanced_prompt, style, negative_prompt)
-    return image_url, enhanced_prompt
+    return image_url
 
 with gr.Blocks() as demo:
     gr.Markdown("# Stability AI Image Generator with Google Gemini Prompt Enhancement")
@@ -96,16 +92,25 @@ with gr.Blocks() as demo:
             prompt = gr.Textbox(label="Prompt")
             style = gr.Dropdown(label="Style", choices=STYLES)
             negative_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT)
+            enhance_btn = gr.Button("Enhance Prompt")
             generate_btn = gr.Button("Generate Image")
         
         with gr.Column(scale=1):
             image_output = gr.Image(label="Generated Image")
             enhanced_prompt_output = gr.Textbox(label="Enhanced Prompt")
     
+    enhanced_prompt_state = gr.State()
+    
+    enhance_btn.click(
+        process_prompt,
+        inputs=[google_api_key, prompt, style],
+        outputs=[enhanced_prompt_output, enhanced_prompt_state]
+    )
+    
     generate_btn.click(
-        process_and_generate,
-        inputs=[google_api_key, stability_api_key, prompt, style, negative_prompt],
-        outputs=[image_output, enhanced_prompt_output]
+        generate_from_enhanced,
+        inputs=[stability_api_key, enhanced_prompt_state, style, negative_prompt],
+        outputs=[image_output]
     )
 
 demo.launch()