import streamlit as st
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image

# Set up the app title and description
st.title("AI Art Generator")
st.write("This app generates beautiful art based on your text input using an AI model.")

# Load the Stable Diffusion model (make sure to use a proper model from Hugging Face)
@st.cache_resource
def load_model():
    model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v-1-4-original", torch_dtype=torch.float16)
    model.to("cuda")  # Use GPU if available
    return model

# Initialize the model
generator = load_model()

# Text input area for the user to describe the image they want
user_input = st.text_area("Describe your art here:", placeholder="E.g., A futuristic cityscape with flying cars and neon lights.", height=150)

# Generate the image based on text input
if st.button("Generate Art"):
    if user_input.strip():
        with st.spinner("Generating your artwork..."):
            # Generate image from text
            image = generator(user_input).images[0]
            # Display the generated image
            st.image(image, caption="Generated Art", use_column_width=True)
    else:
        st.warning("Please enter some text to generate art.")

# Footer
st.write("---")
st.write("This app uses Stable Diffusion, a Hugging Face model, to generate art from text input.")
st.write("Developed with ❤️ by [Your Name].")