Update app.py

app.py

@@ -1,8 +1,6 @@
 import gradio as gr
 import google.generativeai as genai
 import requests
-from PIL import Image
-import io
 import logging
 
 # Set up logging
@@ -58,34 +56,27 @@ def enhance_prompt(google_api_key, prompt, style):
         raise
 
 def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt):
-    url = "https://api.stability.ai/v2beta/stable-image/generate/sd3"
+    url = "https://api.stability.ai/v2beta/stable-image/generate/ultra"
     
     headers = {
-        "Accept": "image/jpeg",
+        "Accept": "image/*",
         "Authorization": f"Bearer {stability_api_key}"
     }
     
-    payload = {
+    data = {
         "prompt": f"{enhanced_prompt}, Style: {style}",
         "negative_prompt": negative_prompt,
-        "model": "sd3.5-large-turbo",
-        "width": 1024,
-        "height": 1024,
-        "num_images": 1,
-        "steps": 30,
-        "cfg_scale": 7.5,
-        "seed": 0,
-        "output_format": "jpeg"
+        "output_format": "webp"
     }
     
     try:
-        response = requests.post(url, headers=headers, json=payload)
+        response = requests.post(url, headers=headers, files={"none": ''}, data=data)
         response.raise_for_status()
         
         logging.debug(f"Response headers: {response.headers}")
         logging.debug(f"Response content type: {response.headers.get('content-type')}")
         
-        if response.headers.get('content-type') == 'image/jpeg':
+        if response.headers.get('content-type').startswith('image/'):
             return response.content
         else:
             error_message = response.text
@@ -102,17 +93,17 @@ def process_and_generate(google_api_key, stability_api_key, prompt, style, negat
         image_bytes = generate_image(stability_api_key, enhanced_prompt, style, negative_prompt)
         
         # Save image to a file
-        with open("generated_image.jpg", "wb") as f:
+        with open("generated_image.webp", "wb") as f:
             f.write(image_bytes)
         
         # Return the file path and enhanced prompt
-        return "generated_image.jpg", enhanced_prompt
+        return "generated_image.webp", enhanced_prompt
     except Exception as e:
         logging.error(f"Error in process_and_generate: {str(e)}")
         return str(e), enhanced_prompt if 'enhanced_prompt' in locals() else "Error occurred before prompt enhancement"
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Stability AI Image Generator with Google Gemini Prompt Enhancement")
+    gr.Markdown("# Stability AI Ultra Image Generator with Google Gemini Prompt Enhancement")
     
     with gr.Row():
         with gr.Column(scale=1):