Update app.py

app.py

@@ -2,6 +2,7 @@ import requests
 import gradio as gr
 from PIL import Image
 from io import BytesIO
+import base64
 import os
 
 # Replace with your NVIDIA API key
@@ -25,13 +26,23 @@ def generate_kindle_cover(prompt):
     
     if response.status_code == 200:
         response_body = response.json()
-        image_url = response_body.get('image_url')  # Adjust based on actual response structure
+        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
+        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
         if image_url:
             image_response = requests.get(image_url)
             image = Image.open(BytesIO(image_response.content))
             return image
-        else:
-            return "No image URL found in response."
+        
+        return "No image data found in response."
     else:
         return f"Error: {response.status_code} - {response.text}"