Spaces:
Sleeping
Sleeping
| import os | |
| import streamlit as st | |
| import matplotlib.pyplot as plt | |
| from PIL import Image | |
| from diffusers import StableDiffusionPipeline | |
| import torch | |
| # Function to generate and display image | |
| def generate_and_display_image(prompt_text): | |
| # Set device to CPU if GPU is not available | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # Create the diffusion pipeline | |
| pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32) # Changed torch.float16 to torch.float32 | |
| pipe = pipe.to(device) | |
| try: | |
| # Generate the image | |
| image = pipe(prompt_text).images[0] | |
| # Save the generated image with the name of the prompt | |
| image_path = f"{prompt_text}.png" | |
| image.save(image_path) | |
| # Show success message | |
| st.write("Image generated successfully!") | |
| # Display the generated image | |
| display_image(image_path) | |
| except Exception as e: | |
| # Show error message if an error occurs | |
| st.error(f"An error occurred: {str(e)}") | |
| # Function to display image | |
| def display_image(image_path): | |
| # Open the generated image | |
| generated_image = Image.open(image_path) | |
| st.image(generated_image, use_column_width=True) | |
| # Model parameters | |
| model_id = "CompVis/stable-diffusion-v1-4" | |
| # Streamlit text input for prompt | |
| prompt_text = st.text_input("Enter the prompt:") | |
| # Generate and display image if prompt is provided | |
| if prompt_text: | |
| generate_and_display_image(prompt_text) | |