Update app.py

app.py

@@ -1,4 +1,42 @@
+#import streamlit as st
+
+#x = st.slider('Select a value')
+#st.write(x)
+
+
+import torch
+from diffusers import StableDiffusionPipeline
 import streamlit as st
+from PIL import Image
+
+# Streamlit UI Setup
+st.title("Stable Diffusion Demo")
+
+# Text input for prompt
+prompt = st.text_input("Enter your prompt:", "A sunset over a mountain range")
+
+# Load the model once, so that it doesn't reload on every change
+@st.cache_resource
+def load_model():
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16
+    )
+    pipe = pipe.to("cuda")
+    return pipe
+
+pipe = load_model()
+
+# Generate image on button click
+if st.button("Generate Image"):
+    with st.spinner("Generating image..."):
+        image = pipe(prompt).images[0]
+
+        # Display the generated image in the app
+        st.image(image, caption=f"Generated image for prompt: {prompt}")
+
+# Optionally, save the image
+if st.button("Save Image"):
+    file_name = "generated_image.png"
+    image.save(file_name)
+    st.success(f"Image saved as {file_name}")
 
-x = st.slider('Select a value')
-st.write(x)