Update app.py

app.py

@@ -26,22 +26,27 @@ def generate_kindle_cover(prompt):
     
     if response.status_code == 200:
         response_body = response.json()
-        print("Response Body:", response_body)  # Debugging line
-        
-        # Check if image is in base64 format
-        image_data = response_body.get('image_base64')  # Replace 'image_base64' with the actual key if different
+        print("Response Body:", response_body)  # Debugging line to inspect the response
+
+        # Adjust based on actual response structure
+        image_data = response_body.get('image_base64')  # Replace 'image_base64' if necessary
         if image_data:
-            image_bytes = base64.b64decode(image_data)
-            image = Image.open(BytesIO(image_bytes))
-            return image
-        
-        # Check if image is returned as a URL
-        image_url = response_body.get('image_url')  # Replace 'image_url' with the actual key if different
+            try:
+                image_bytes = base64.b64decode(image_data)
+                image = Image.open(BytesIO(image_bytes))
+                return image
+            except Exception as e:
+                return f"Error decoding base64 image: {e}"
+
+        image_url = response_body.get('image_url')  # Replace 'image_url' if necessary
         if image_url:
-            image_response = requests.get(image_url)
-            image = Image.open(BytesIO(image_response.content))
-            return image
-        
+            try:
+                image_response = requests.get(image_url)
+                image = Image.open(BytesIO(image_response.content))
+                return image
+            except Exception as e:
+                return f"Error loading image from URL: {e}"
+
         return "No image data found in response."
     else:
         return f"Error: {response.status_code} - {response.text}"