app.py

import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
import os
from huggingface_hub import login

# Login using your Hugging Face token
hf_token = os.getenv("HF_TOKEN")  # This should be set in your Hugging Face Space Secrets
login(token=hf_token)

# Access the NVIDIA API key from the environment
nvidia_api_key = os.getenv("NVIDIA_API_KEY")

# Load the model using the API key and move it to GPU
def load_model():
    model_id = "NVIDIA/sdxl-turbo"
    pipe = StableDiffusionPipeline.from_pretrained(
        model_id, 
        torch_dtype=torch.float16, 
        use_auth_token=nvidia_api_key  # Use the API key here
    )
    pipe = pipe.to("cuda")
    return pipe

pipe = load_model()

# Function to generate a Kindle cover
def generate_kindle_cover(prompt):
    image = pipe(prompt).images[0]
    return image

# 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
if __name__ == "__main__":
    iface.launch()