ilsa15's picture
Create app.py
80b2ebe verified
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")