Update app.py

app.py

@@ -3,6 +3,10 @@ import google.generativeai as genai
 import requests
 from PIL import Image
 import io
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # List of popular styles
 STYLES = [
@@ -37,22 +41,27 @@ def enhance_prompt(google_api_key, prompt, style):
     Enhanced prompt:
     """
     
-    response = model.generate_content(enhanced_prompt_request)
-    
-    enhanced_prompt = response.text.strip()
-    
-    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()):
-            enhanced_prompt = enhanced_prompt[len(prefix):].strip()
-    
-    return enhanced_prompt
+    try:
+        response = model.generate_content(enhanced_prompt_request)
+        
+        enhanced_prompt = response.text.strip()
+        
+        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()):
+                enhanced_prompt = enhanced_prompt[len(prefix):].strip()
+        
+        logging.info(f"Enhanced prompt: {enhanced_prompt}")
+        return enhanced_prompt
+    except Exception as e:
+        logging.error(f"Error in enhance_prompt: {str(e)}")
+        raise
 
 def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt):
     url = "https://api.stability.ai/v2beta/stable-image/generate/sd3"
     
     headers = {
-        "Accept": "image/*",
+        "Accept": "image/jpeg",
         "Authorization": f"Bearer {stability_api_key}"
     }
     
@@ -66,24 +75,41 @@ def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt):
         "steps": (None, "30"),
         "cfg_scale": (None, "7.5"),
         "seed": (None, "0"),
-        "output_format": (None, "png")
+        "output_format": (None, "jpeg")
     }
     
-    response = requests.post(url, headers=headers, files=files)
+    try:
+        response = requests.post(url, headers=headers, files=files)
+        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':
+            return response.content
+        else:
+            error_message = response.text
+            logging.error(f"Unexpected content type: {response.headers.get('content-type')}. Response: {error_message}")
+            raise Exception(f"Unexpected content type: {response.headers.get('content-type')}. Response: {error_message}")
     
-    if response.status_code == 200:
-        return response.content
-    else:
-        raise Exception(f"Image generation failed: {response.text}")
+    except requests.exceptions.RequestException as e:
+        logging.error(f"Request failed: {str(e)}")
+        raise Exception(f"Request failed: {str(e)}")
 
 def process_and_generate(google_api_key, stability_api_key, prompt, style, negative_prompt):
-    enhanced_prompt = enhance_prompt(google_api_key, prompt, style)
-    image_bytes = generate_image(stability_api_key, enhanced_prompt, style, negative_prompt)
-    
-    # Convert bytes to PIL Image
-    image = Image.open(io.BytesIO(image_bytes))
-    
-    return image, enhanced_prompt
+    try:
+        enhanced_prompt = enhance_prompt(google_api_key, prompt, style)
+        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:
+            f.write(image_bytes)
+        
+        # Return the file path and enhanced prompt
+        return "generated_image.jpg", 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")
@@ -98,7 +124,7 @@ with gr.Blocks() as demo:
             submit_btn = gr.Button("Generate Image")
         
         with gr.Column(scale=1):
-            image_output = gr.Image(label="Generated Image")
+            image_output = gr.Image(label="Generated Image", type="filepath")
             enhanced_prompt_output = gr.Textbox(label="Enhanced Prompt")
     
     submit_btn.click(