import streamlit as st
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import io

st.set_page_config(page_title="AI Book Cover Designer", layout="centered")
st.title("📚 AI Book Cover Designer")

@st.cache_resource
def load_model():
    model_id = "stabilityai/stable-diffusion-2"
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
    pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
    return pipe

pipe = load_model()

title = st.text_input("Book Title", "The Last Moonwalker")
desc = st.text_area("Short Description", "An astronaut trapped on the moon battles loneliness and discovers a strange portal.")
genre = st.selectbox("Genre", ["Fantasy", "Romance", "Sci-Fi", "Horror", "Mystery", "Adventure", "Historical Fiction"])

if st.button("🎨 Generate Cover"):
    with st.spinner("Generating your beautiful book cover..."):
        prompt = f"A book cover for a {genre} novel titled '{title}': {desc}"
        image = pipe(prompt, num_inference_steps=25).images[0]
        st.image(image, caption="Your AI-generated book cover", use_column_width=True)

        # Save button
        img_bytes = io.BytesIO()
        image.save(img_bytes, format='PNG')
        st.download_button(label="📥 Download Cover", data=img_bytes.getvalue(), file_name="book_cover.png", mime="image/png")