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)