import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from datetime import datetime
from PIL import Image

# Load the model
model_id = "stabilityai/stable-diffusion-2"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cpu")

# Keep image history
image_history = []

# Style templates
STYLE_TEMPLATES = {
    "Photo": "highly detailed, realistic photo, natural lighting",
    "Anime": "anime style, vibrant colors, line art, cel shading",
    "Oil Painting": "oil painting style, brush strokes, classic fine art"
}

def enhance_prompt(prompt, style):
    style_suffix = STYLE_TEMPLATES.get(style, "")
    return f"{prompt}, {style_suffix}"

def generate_image(prompt, style, uploaded_image):
    if not prompt.strip():
        return None, image_history

    final_prompt = enhance_prompt(prompt, style)

    # Currently not using uploaded_image — could be used with img2img later
    image = pipe(final_prompt).images[0]

    # Add to history
    timestamp = datetime.now().strftime("%H:%M:%S")
    caption = f"{style} - {timestamp}"
    image_history.append((image, caption))
    
    return image, image_history[-5:]

def clear_history():
    global image_history
    image_history = []
    return None, []

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🎨 Realistic Text-to-Image Generator")
    gr.Markdown("Powered by **Stable Diffusion 2** | Now with style presets, prompt enhancer, image history, and optional image upload!")

    with gr.Row():
        prompt = gr.Textbox(label="Enter your prompt", placeholder="e.g., A golden retriever in the snow")
        style = gr.Dropdown(["Photo", "Anime", "Oil Painting"], label="Choose Style", value="Photo")
    
    uploaded_image = gr.Image(label="Upload Optional Image (future use)", type="pil", tool=None)

    with gr.Row():
        generate_btn = gr.Button("🎨 Generate Image")
        clear_btn = gr.Button("🧹 Clear History")

    output_image = gr.Image(label="Generated Image", type="pil")
    gallery = gr.Gallery(label="🖼️ Image History (last 5)", columns=3, object_fit="contain")

    generate_btn.click(generate_image, inputs=[prompt, style, uploaded_image], outputs=[output_image, gallery])
    clear_btn.click(clear_history, outputs=[output_image, gallery])

demo.launch()