Update app.py

app.py

@@ -1,8 +1,12 @@
 import requests
 import gradio as gr
-import os
+from PIL import Image
+from io import BytesIO
+import base64
 import json
+import os
 
+# Get the API key from environment variable
 API_KEY = os.getenv("NVIDIA_API_KEY")
 invoke_url = "https://ai.api.nvidia.com/v1/genai/stabilityai/sdxl-turbo"
 headers = {
@@ -11,9 +15,6 @@ 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,
@@ -22,34 +23,36 @@ def generate_kindle_cover(prompt):
     }
     
     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}")
+        response.raise_for_status()
         
-        if response.status_code == 401:
-            return "Error: Unauthorized access (401). Please check your API key."
-        elif response.status_code != 200:
-            return f"Error: Received status code {response.status_code} from API. Response: {response.text}"
+        response_body = response.json()
         
-        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."
-        except json.JSONDecodeError:
-            return f"Error: Failed to parse JSON response. Raw response: {response.text}"
+        if 'results' in response_body and len(response_body['results']) > 0:
+            result = response_body['results'][0]
+            if 'image' in result:
+                image_data = result['image']
+                image_bytes = base64.b64decode(image_data)
+                image = Image.open(BytesIO(image_bytes))
+                return image
+            else:
+                return "No image data 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 json.JSONDecodeError:
+        return "Error parsing JSON response."
     except Exception as e:
-        print(f"An error occurred: {str(e)}")
-        return f"Error: {str(e)}"
+        return f"Unexpected error: {str(e)}"
 
 iface = gr.Interface(
     fn=generate_kindle_cover,
     inputs="text",
-    outputs="text",
+    outputs="image",
     title="Kindle Cover Generator",
     description="Generate high-quality covers for Amazon Kindle books using the NVIDIA SDXL-Turbo model."
 )
 
-iface.launch(debug=True)
+iface.launch()