Spaces:
Sleeping
Sleeping
File size: 1,424 Bytes
31b3285 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
import streamlit as st
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import io
import base64
# Model ID
model_id = "stabilityai/stable-diffusion-2-1"
# Load the model from Hugging Face Hub
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
# Streamlit app title and description
st.title("Stable Diffusion Image Generation")
st.write("Generate images using the Stability AI's Stable Diffusion model.")
# Input for the prompt
prompt = st.text_input("Enter a prompt for the image:")
# Generate button
if st.button("Generate Image"):
if prompt:
with st.spinner('Generating image...'):
# Generate image
image = pipe(prompt).images[0]
# Convert image to displayable format
buffered = io.BytesIO()
image.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode()
# Display the image
st.image(image, caption="Generated Image", use_column_width=True)
# Option to download the image
st.markdown(
f'<a href="data:image/png;base64,{img_str}" download="generated_image.png">Download Image</a>',
unsafe_allow_html=True,
)
else:
st.error("Please enter a prompt to generate an image.")
|