Spaces:
Sleeping
Sleeping
| 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.") | |