app.py

import requests
import gradio as gr
from PIL import Image
from io import BytesIO
import base64
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",
}

def generate_kindle_cover(prompt):
    payload = {
        "text_prompts": [{"text": prompt}],
        "seed": 0,
        "sampler": "K_EULER_ANCESTRAL",
        "steps": 2
    }
    
    response = requests.post(invoke_url, headers=headers, json=payload)
    
    if response.status_code == 200:
        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:
                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}"

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_kindle_cover,
    inputs="text",
    outputs="image",
    title="Kindle Cover Generator",
    description="Generate high-quality covers for Amazon Kindle books using the NVIDIA SDXL-Turbo model."
)

# Launch the Gradio interface
iface.launch()