import torch
import gradio as gr
from diffusers import StableDiffusionPipeline

model_id = "segmind/tiny-sd"

pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float32
)

pipe = pipe.to("cpu")

pipe.enable_attention_slicing()

def generate(prompt):
    image = pipe(
        prompt,
        num_inference_steps=10,
        height=512,
        width=512
    ).images[0]

    return image

demo = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(
        label="Prompt",
        placeholder="A futuristic cyberpunk city..."
    ),
    outputs=gr.Image(label="Generated Image"),
    title="Tiny SD CPU Demo",
    description="Tiny Stable Diffusion running on Hugging Face CPU Space"
)

demo.launch()