# 🤖✨ CPU Image Generator (Stable Diffusion)

import torch
from diffusers import StableDiffusionPipeline
import gradio as gr

# Load model once at startup (≈3.4 GB; fits HF Space free tier CPU)
PIPELINE_ID = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(
    PIPELINE_ID,
    torch_dtype=torch.float32,
    safety_checker=None,            # skip NSFW checker
)
pipe = pipe.to("cpu")

def generate_image(prompt: str, steps: int):
    if not prompt:
        return None
    image = pipe(prompt, num_inference_steps=steps).images[0]
    return image

with gr.Blocks(title="🤖✨ AI Image Generator (CPU)") as demo:
    gr.Markdown(
        "# 🤖✨ AI Image Generator\n"
        "Enter a creative prompt, adjust inference steps, and generate a unique image—**100% CPU**."
    )
    with gr.Row():
        prompt_in = gr.Textbox(
            label="Prompt",
            placeholder="e.g. A photorealistic portrait of a cyberpunk fox"
        )
        steps_in = gr.Slider(
            minimum=1, maximum=50, value=25, step=1, label="Inference Steps"
        )
    run_btn = gr.Button("Generate 🖼️", variant="primary")
    img_out = gr.Image(label="Generated Image")

    run_btn.click(generate_image, [prompt_in, steps_in], img_out)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0")