Art_Generator / app.py
Keiraj's picture
Update app.py
32e29ce verified
raw
history blame
780 Bytes
import streamlit as st
from diffusers import DiffusionPipeline
import torch
from PIL import Image
# Streamlit app title
st.title("Stable Diffusion Image Generator")
# Load the diffusion pipeline
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float32)
pipe.to("cpu") # Use CPU for inference
# Get user input (prompt)
prompt = st.text_input("What do you want to see?", "A beautiful landscape")
# Button to generate image
if st.button('Generate Image'):
with st.spinner('Generating...'):
# Generate image
image = pipe(prompt).images[0]
# Display image in Streamlit
st.image(image, caption="Generated Image", use_column_width=True)
st.success("Image generated successfully!")