import gradio as gr
from diffusers import AutoPipelineForText2Image
import torch

# Function to generate the image
def generate_image(prompt):
    # Load the pipeline with default settings which should default to a CPU compatible setting if `fp32` is available
    pipe = AutoPipelineForText2Image.from_pretrained(
        "stabilityai/sdxl-turbo",
        torch_dtype=torch.float32  # Using float32 for CPU compatibility
    )
    pipe = pipe.to("cpu")  # Ensure the pipeline is using the CPU

    # Generate the image based on the prompt
    image = pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
    return image

# Define the Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter a description for the image"),
    outputs=gr.Image(type="pil", label="Generated Image"),
    title="Image Generator",
    description="This interface generates images based on your descriptions using the Stability AI SDXL-Turbo model."
)

# Prepare to run in Hugging Face Spaces
if __name__ == "__main__":
    interface.launch()