app.py

import requests
import gradio as gr
from PIL import Image
from io import BytesIO
import base64
import os
import time

API_KEY = os.getenv("NVIDIA_API_KEY")
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):
    print(f"Generating cover for prompt: {prompt}")
    
    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}")
        
        if response.status_code != 200:
            return f"Error: Received status code {response.status_code} from API"
        
        response_body = response.json()
        print("Response Body:", response_body)
        
        if 'results' not in response_body or not response_body['results']:
            return "Error: No results found in the response"
        
        result = response_body['results'][0]
        if 'image' not in result:
            return "Error: No image data found in the result"
        
        image_data = result['image']
        print("Image data length:", len(image_data))
        
        image_bytes = base64.b64decode(image_data)
        print("Decoded image bytes length:", len(image_bytes))
        
        image = Image.open(BytesIO(image_bytes))
        print("Image successfully created")
        
        # Save the image
        timestamp = int(time.time())
        filename = f"kindle_cover_{timestamp}.png"
        image.save(filename)
        print(f"Image saved as {filename}")
        
        return filename
    
    except Exception as e:
        print(f"An error occurred: {str(e)}")
        return f"Error: {str(e)}"

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

iface.launch(debug=True)