Update app.py

app.py

@@ -8,6 +8,7 @@ import requests
 import logging
 import threading
 import io
+import time
 
 # Set up logging
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -119,19 +120,22 @@ def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt):
         "width": 1024,
         "height": 1024,
         "num_images": 1,
-        "steps": 20,
+        "steps": 30,  # Increased from 20 to 30
         "cfg_scale": 7,
     }
     
     try:
-        response = requests.post(url, headers=headers, files={"none": ''}, data=data)
+        response = requests.post(url, headers=headers, files={"none": ''}, data=data, timeout=60)  # Added 60-second timeout
         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').startswith('image/'):
-            return response.content
+            image_data = response.content
+            if len(image_data) < 1000:  # Check if the image data is too small (likely incomplete)
+                raise Exception("Received incomplete image data")
+            return image_data
         else:
             error_message = response.text
             logging.error(f"Unexpected content type: {response.headers.get('content-type')}. Response: {error_message}")
@@ -147,10 +151,18 @@ def process_and_generate(google_api_key, stability_api_key, prompt, style, set_s
         enhanced_prompt = enhance_prompt(google_api_key, prompt, style)
         
         set_status("Generating image...")
-        image_bytes = generate_image(stability_api_key, enhanced_prompt, style, DEFAULT_NEGATIVE_PROMPT)
-        
-        set_status("Image generated successfully!")
-        return image_bytes, enhanced_prompt
+        max_attempts = 3
+        for attempt in range(max_attempts):
+            try:
+                image_bytes = generate_image(stability_api_key, enhanced_prompt, style, DEFAULT_NEGATIVE_PROMPT)
+                set_status("Image generated successfully!")
+                return image_bytes, enhanced_prompt
+            except Exception as e:
+                if attempt < max_attempts - 1:
+                    set_status(f"Attempt {attempt + 1} failed. Retrying...")
+                    time.sleep(2)  # Wait for 2 seconds before retrying
+                else:
+                    raise e
     except Exception as e:
         logging.error(f"Error in process_and_generate: {str(e)}")
         set_status(f"Error: {str(e)}")
@@ -187,12 +199,20 @@ def update_output(n_clicks, google_api_key, stability_api_key, prompt, style):
         else:
             return "", f"Error: {enhanced_prompt}", status["message"], True
     
-    # Run the process in a separate thread
-    thread = threading.Thread(target=run_process)
-    thread.start()
-    thread.join()  # Wait for the thread to complete
-    
-    return run_process()
+    try:
+        # Run the process in a separate thread
+        thread = threading.Thread(target=run_process)
+        thread.start()
+        thread.join(timeout=90)  # Wait for up to 90 seconds
+        
+        if thread.is_alive():
+            # If the thread is still running after 90 seconds, we consider it a timeout
+            return "", "Error: Image generation timed out", "Process timed out", True
+        
+        return run_process()
+    except Exception as e:
+        logging.error(f"Unexpected error in update_output: {str(e)}")
+        return "", f"Unexpected error: {str(e)}", "An unexpected error occurred", True
 
 @app.callback(
     Output("download-image", "data"),