Update app.py

app.py

@@ -1,50 +1,72 @@
+# prompt: write this notebook code for me  i want to deploy this model on hugging space using streamlit
+
 import streamlit as st
-from diffusers import StableDiffusionPipeline
+import mediapy as media
+import random
+import sys
 import torch
-from PIL import Image
-
-# Streamlit App Title
-st.title("🖼️ Stable Diffusion Image Generator")
-st.write("Generate images from text prompts using Stable Diffusion!")
-
-# Sidebar for prompt input and button
-st.sidebar.title("Image Generation Settings")
-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")
-
-# Cache the Stable Diffusion model to load it only once
-@st.cache_resource
-def load_model():
-    model = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float16)
-    model = model.to("cuda" if torch.cuda.is_available() else "cpu")
-    return model
-
-# Load the model
-pipe = load_model()
-
-# Generate the image when the button is clicked
-if generate_button:
-    if prompt.strip() == "":
-        st.error("Prompt cannot be empty. Please enter a valid text prompt.")
+from diffusers import DiffusionPipeline
+
+# Install necessary libraries if not already installed
+try:
+    import diffusers
+except ImportError:
+    !pip install --quiet --upgrade diffusers transformers accelerate invisible_watermark mediapy
+
+# Set use_refiner to False for initial deployment. You can add a checkbox later.
+use_refiner = False
+
+@st.cache_resource  # Cache the pipeline to avoid reloading every time
+def load_pipeline():
+    pipe = DiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+        variant="fp16",
+    )
+
+    if use_refiner:  # Keep the refiner logic but disabled initially
+        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:
-        st.write(f"### Prompt: {prompt}")
-        with st.spinner("Generating image... Please wait. This may take a few moments depending on system performance."):
-            try:
-                # Generate the image
-                image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
-
-                # Display the image
-                st.image(image, caption="Generated Image", use_column_width=True)
-
-                # Provide download option
-                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}")
+        pipe = pipe.to("cuda")
+    return pipe
+
+# Load the pipeline
+pipe = load_pipeline()
+
+st.title("Stable Diffusion XL Image Generator")
+
+# Input text field for the prompt
+prompt = st.text_input("Enter your prompt:", "a photo of Pikachu fine dining with a view to the Eiffel Tower")
+
+# Seed input (optional)
+seed = st.number_input("Enter a seed (optional):", min_value=0, value=random.randint(0, sys.maxsize))
+
+if st.button("Generate Image"):
+    if not prompt:
+        st.warning("Please enter a prompt.")
+    else:
+        with st.spinner("Generating image..."):
+            images = pipe(
+                prompt=prompt,
+                output_type="latent" if use_refiner else "pil",
+                generator=torch.Generator("cuda").manual_seed(seed),
+            ).images
+            if use_refiner:  # Condition still here, but the refiner is not used
+                images = refiner(prompt=prompt, image=images).images
+            st.write(f"Prompt:\t{prompt}\nSeed:\t{seed}")
+            st.image(images[0]) #Display in Streamlit
+            # You might want to add a download button here
+
+# Add a checkbox for the refiner
+#use_refiner = st.checkbox("Use Refiner (Experimental)", value=False)
+