import os
import gradio as gr
from diffusers import DiffusionPipeline
from PIL import Image
import io

pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
api_key = os.getenv("API_KEY")

def generate_image(prompt, style):
    input_prompt = prompt
    generated_image = pipeline(input_prompt)

    img = Image.fromarray(generated_image)
    return img

title = "Stable Diffusion XL"
description = "This app generates an image based on the provided prompt using the Stable Diffusion XL model."

styles = {
    "background": "linear-gradient(to bottom, #33ccff, #ff99cc)",
    "color": "black",
    "font-family": "Arial, sans-serif"
}

gr.Interface(
    fn=generate_image,
    inputs=["text", gr.Textbox(label="Style")],
    outputs="image",
    title=title,
    description=description,
    examples=[["Astronaut riding a horse", ""]],
    theme="compact",
    layout="vertical",
    allow_flagging=False,
    flagging_dir=None,
    flagging_host=None,
    capture_session=True,
    css={"body": styles}
).launch()