Update app.py

app.py

@@ -1,53 +1,53 @@
-# prompt: i found this error please fix this "File "/home/user/app/app.py", line 14
-#       !pip install --quiet --upgrade diffusers transformers accelerate invisible_watermark mediapy
-#       ^
-# SyntaxError: invalid syntax"
-
-import mediapy as media
-import random
-import sys
-import torch
-from diffusers import DiffusionPipeline
 import streamlit as st
-# Remove this line:
-# %pip install --quiet --upgrade diffusers transformers accelerate invisible_watermark mediapy
-
-# Install necessary libraries if not already installed
-try:
-    import diffusers
-    import transformers
-    import accelerate
-    import invisible_watermark
-    import mediapy
-    import streamlit
-    import torch
-except ImportError:
-    !pip install --quiet --upgrade diffusers transformers accelerate invisible_watermark mediapy streamlit torch
-
-use_refiner = False
-
-
-pipe = DiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-    variant="fp16",
+from diffusers import StableDiffusionPipeline
+import torch
+from PIL import Image
+
+# Title and description of the app
+st.title("🖼️ Stable Diffusion Image Generator")
+st.write("Generate images from text using the Stable Diffusion v1.5 model!")
+
+# Sidebar for user inputs
+st.sidebar.title("Input Options")
+prompt = st.sidebar.text_input("Enter your prompt", "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k")
+generate_button = st.sidebar.button("Generate Image")
+
+# Load the pipeline when the app starts
+@st.cache_resource
+def load_pipeline():
+    # Use a smaller model or a more efficient pipeline
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5", 
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        revision="fp16" if torch.cuda.is_available() else None
     )
-
-if use_refiner:
-  refiner = DiffusionPipeline.from_pretrained(
-      "stabilityai/stable-diffusion-xl-refiner-1.0",
-      text_encoder_2=pipe.text_encoder_2,
-      vae=pipe.vae,
-      torch_dtype=torch.float16,
-      use_safetensors=True,
-      variant="fp16",
-  )
-
-  refiner = refiner.to("cuda")
-
-  pipe.enable_model_cpu_offload()
-else:
-  pipe = pipe.to("cuda")
-
-# ... (rest of your code)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pipe = pipe.to(device)
+    return pipe
+
+pipe = load_pipeline()
+
+# Generate image when button is clicked
+if generate_button:
+    st.write(f"### Prompt: {prompt}")
+    with st.spinner("Generating image... Please wait."):
+        try:
+            # Generate the image
+            with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
+                image = pipe(prompt).images[0]
+            
+            # Display the generated image
+            st.image(image, caption="Generated Image", use_column_width=True)
+            
+            # Option to download the image
+            img_path = "generated_image.png"
+            image.save(img_path)
+            with open(img_path, "rb") as img_file:
+                st.download_button(
+                    label="Download Image",
+                    data=img_file,
+                    file_name="generated_image.png",
+                    mime="image/png"
+                )
+        except Exception as e:
+            st.error(f"An error occurred: {e}")