import streamlit as st
from diffusers import StableDiffusionPipeline
import torch
import transformers

# Load Stable Diffusion model
model_id = "stabilityai/stable-diffusion-v1-4"
model = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")

# Streamlit UI
st.title("Generative AI Application")

# Input prompt
def main():
    st.sidebar.header("Input Settings")
    prompt = st.sidebar.text_input("Enter your prompt:", "A modern house floor plan with 3 bedrooms and a pool")

    if st.sidebar.button("Generate Image"):
        if not prompt:
            st.error("Please enter a prompt.")
        else:
            with st.spinner("Generating image..."):
                try:
                    images = model([prompt])
                    image = images[0]
                    st.image(image, caption="Generated Image", use_column_width=True)
                except Exception as e:
                    st.error(f"Error generating image: {str(e)}")

if __name__ == "__main__":
    main()