import gradio as gr
from diffusers import AutoPipelineForText2Image
import torch
from datetime import datetime

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

prompt_history = []

def generate_image(prompt, seed):
    if not prompt:
        return None, "Please enter a prompt!"
    if seed is None:
        seed = torch.seed()
    generator = torch.Generator().manual_seed(int(seed))
    image = pipe(
        prompt,
        generator=generator,
        num_inference_steps=4, 
        guidance_scale=0.0
    ).images[0]
    prompt_history.append(f"{datetime.now().strftime('%H:%M:%S')} — {prompt}")
    return image, "\n".join(prompt_history[-5:])

def clear_history():
    prompt_history.clear()
    return "History cleared."

# Example prompt suggestions
example_prompts = [
    "a futuristic city skyline at sunset, cinematic lighting",
    "a cat astronaut painting on the moon, watercolor style",
    "a robot reading a book under a tree, soft lighting",
    "a realistic portrait of a medieval queen, detailed textures",
    "a cozy cabin in the snowy mountains, warm lights"
]

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft(), title="AI Image Generator") as demo:
    gr.Markdown("""
    <h1 style="text-align:center;">Text-to-Image Generator</h1>
    <p style="text-align:center;">Type your prompt or pick a suggestion below to see generative AI in action.</p>
    """)

    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(
                label="Enter your text prompt",
                placeholder="e.g., a futuristic city with flying cars",
                lines=2,
            )
            seed = gr.Number(label="Seed (optional)", value=None)
            generate_btn = gr.Button("Generate Image", variant="primary")
            clear_btn = gr.Button("Clear History")

            gr.Markdown("Try a sample prompt:")
            with gr.Row():
                for p in example_prompts:
                    gr.Button(p).click(lambda x=p: x, None, prompt, queue=False)

        with gr.Column(scale=2):
            output_image = gr.Image(label="Generated Image", type="pil", show_download_button=True)

    gr.Markdown("Recent Prompts")
    history_box = gr.Textbox(
        label="Prompt History (last 5)",
        interactive=False,
        lines=5,
        placeholder="History will appear here...",
    )

    # Button actions
    generate_btn.click(fn=generate_image, inputs=[prompt, seed], outputs=[output_image, history_box])
    clear_btn.click(fn=clear_history, inputs=None, outputs=history_box)

    gr.Markdown("<p style='text-align:center; color:gray;'>Created using Gradio + Stable Diffusion Turbo</p>")

demo.launch()