Spaces:
Running
Running
| import streamlit as st | |
| from diffusers import StableDiffusionPipeline | |
| import torch | |
| from PIL import Image | |
| # Streamlit App Title | |
| st.title("🖼️ Stable Diffusion Image Generator") | |
| st.write("Generate images from text prompts using Stable Diffusion!") | |
| # Sidebar for prompt input and button | |
| st.sidebar.title("Image Generation Settings") | |
| prompt = st.sidebar.text_input("Enter your prompt", "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k") | |
| generate_button = st.sidebar.button("Generate Image") | |
| # Cache the Stable Diffusion model to load it only once | |
| def load_model(): | |
| model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16) | |
| model = model.to("cuda" if torch.cuda.is_available() else "cpu") | |
| return model | |
| # Load the model | |
| pipe = load_model() | |
| # Generate the image when the button is clicked | |
| if generate_button: | |
| if prompt.strip() == "": | |
| st.error("Prompt cannot be empty. Please enter a valid text prompt.") | |
| else: | |
| st.write(f"### Prompt: {prompt}") | |
| with st.spinner("Generating image... Please wait."): | |
| try: | |
| # Generate the image | |
| image = pipe(prompt).images[0] | |
| # Display the image | |
| st.image(image, caption="Generated Image", use_column_width=True) | |
| # Provide download option | |
| img_path = "generated_image.png" | |
| image.save(img_path) | |
| with open(img_path, "rb") as img_file: | |
| st.download_button( | |
| label="Download Image", | |
| data=img_file, | |
| file_name="generated_image.png", | |
| mime="image/png" | |
| ) | |
| except Exception as e: | |
| st.error(f"An error occurred: {e}") | |