Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import google.generativeai as genai
 import requests
+import base64
 
 # List of popular styles
 STYLES = [
@@ -63,24 +64,23 @@ def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt):
         "height": 1024,
         "num_images": 1,
         "steps": 30,
-        "cfg_scale": 7.5
+        "cfg_scale": 7.5,
+        "seed": 0,
+        "output_format": "png"
     }
     
     response = requests.post(url, json=payload, headers=headers)
     
     if response.status_code == 200:
         image_data = response.json()['images'][0]['image_base64']
-        return f"data:image/png;base64,{image_data}"
+        return base64.b64decode(image_data)
     else:
-        return f"Image generation failed: {response.text}"
+        raise Exception(f"Image generation failed: {response.text}")
 
-def process_prompt(google_api_key, prompt, style):
+def process_and_generate(google_api_key, stability_api_key, prompt, style, negative_prompt):
     enhanced_prompt = enhance_prompt(google_api_key, prompt, style)
-    return enhanced_prompt, enhanced_prompt  # Return the enhanced prompt twice
-
-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
+    image_data = generate_image(stability_api_key, enhanced_prompt, style, negative_prompt)
+    return image_data, enhanced_prompt
 
 with gr.Blocks() as demo:
     gr.Markdown("# Stability AI Image Generator with Google Gemini Prompt Enhancement")
@@ -92,25 +92,16 @@ 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")
+            submit_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(
-        generate_from_enhanced,
-        inputs=[stability_api_key, enhanced_prompt_state, style, negative_prompt],
-        outputs=[image_output]
+    submit_btn.click(
+        process_and_generate,
+        inputs=[google_api_key, stability_api_key, prompt, style, negative_prompt],
+        outputs=[image_output, enhanced_prompt_output]
     )
 
 demo.launch()