from diffusers import AutoPipelineForText2Image
import torch
import gradio as gr

# Load a CPU-friendly model (works without GPU)
pipe = AutoPipelineForText2Image.from_pretrained(
    "stabilityai/sd-turbo",
    torch_dtype=torch.float32
).to("cpu")

# Image generation function
def generate_image(prompt):
    image = pipe(prompt, num_inference_steps=6).images[0]
    return image

# Gradio UI
gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter a prompt"),
    outputs=gr.Image(label="Generated Image"),
    title="CPU Image Generator (SD-Turbo)"
).launch()