Update app.py

app.py

@@ -8,7 +8,6 @@ import os
 # Replace with your NVIDIA API key
 API_KEY = os.getenv("NVIDIA_API_KEY")  # Ensure the key is set in Space secrets
 invoke_url = "https://ai.api.nvidia.com/v1/genai/stabilityai/sdxl-turbo"
-
 headers = {
     "Authorization": f"Bearer {API_KEY}",
     "Accept": "application/json",
@@ -22,34 +21,34 @@ def generate_kindle_cover(prompt):
         "steps": 2
     }
     
-    response = requests.post(invoke_url, headers=headers, json=payload)
-    
-    if response.status_code == 200:
+    try:
+        response = requests.post(invoke_url, headers=headers, json=payload)
+        response.raise_for_status()  # Raise an exception for non-200 status codes
+        
         response_body = response.json()
         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:
-            try:
+        
+        # Check for 'results' key in response
+        if 'results' in response_body and len(response_body['results']) > 0:
+            result = response_body['results'][0]
+            
+            # Check for 'image' key in result
+            if 'image' in result:
+                image_data = result['image']
                 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:
-            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}"
+            else:
+                return "No 'image' key found in the result."
+        else:
+            return "No 'results' found in the response."
+    
+    except requests.exceptions.RequestException as e:
+        return f"Error making request: {str(e)}"
+    except ValueError as e:
+        return f"Error parsing JSON response: {str(e)}"
+    except Exception as e:
+        return f"Unexpected error: {str(e)}"
 
 # Create the Gradio interface
 iface = gr.Interface(
@@ -61,4 +60,4 @@ iface = gr.Interface(
 )
 
 # Launch the Gradio interface
-iface.launch()
+iface.launch()