Update app.py

app.py

@@ -2,6 +2,9 @@ import requests
 import gradio as gr
 import os
 import json
+import base64
+from PIL import Image
+import io
 
 API_KEY = os.getenv("NVIDIA_API_KEY")
 invoke_url = "https://ai.api.nvidia.com/v1/genai/stabilityai/sdxl-turbo"
@@ -13,41 +16,48 @@ headers = {
 def generate_kindle_cover(prompt):
     print(f"Generating cover for prompt: {prompt}")
     print(f"API Key (first 5 characters): {API_KEY[:5]}..." if API_KEY else "API Key is not set")
-    
+
     payload = {
         "text_prompts": [{"text": prompt}],
         "seed": 0,
         "sampler": "K_EULER_ANCESTRAL",
         "steps": 2
     }
-    
+
     try:
         print("Sending request to NVIDIA API...")
         response = requests.post(invoke_url, headers=headers, json=payload)
         print(f"Received response with status code: {response.status_code}")
         print(f"Response headers: {json.dumps(dict(response.headers), indent=2)}")
-        print(f"Response content: {response.text}")
-        
+
         if response.status_code == 401:
-            return "Error: Unauthorized access (401). Please check your API key."
+            return "Error: Unauthorized access (401). Please check your API key.", None
         elif response.status_code != 200:
-            return f"Error: Received status code {response.status_code} from API. Response: {response.text}"
-        
+            return f"Error: Received status code {response.status_code} from API. Response: {response.text}", None
+
         try:
             response_body = response.json()
             print("Response Body:", json.dumps(response_body, indent=2))
-            return f"Success: API responded with status code 200. Check console for full response details."
+            
+            # Extract and decode the image data
+            image_data = base64.b64decode(response_body['results'][0]['image'])
+            image = Image.open(io.BytesIO(image_data))
+            
+            return "Success: Image generated successfully.", image
+
         except json.JSONDecodeError:
-            return f"Error: Failed to parse JSON response. Raw response: {response.text}"
-    
+            return f"Error: Failed to parse JSON response. Raw response: {response.text}", None
+        except KeyError:
+            return "Error: Expected image data not found in the response.", None
+
     except Exception as e:
         print(f"An error occurred: {str(e)}")
-        return f"Error: {str(e)}"
+        return f"Error: {str(e)}", None
 
 iface = gr.Interface(
     fn=generate_kindle_cover,
     inputs="text",
-    outputs="text",
+    outputs=["text", gr.Image()],
     title="Kindle Cover Generator",
     description="Generate high-quality covers for Amazon Kindle books using the NVIDIA SDXL-Turbo model."
 )