Update app.py

app.py

@@ -1,31 +1,17 @@
 import streamlit as st
 from diffusers import DDPMPipeline
 import PIL
-import time
-import threading
 
-st.title('DDPMPipeline Image Generator')
-
-class ImageGenerator:
-    def __init__(self):
-        self.done = False
-        self.image = None
+# Use the cache decorator to only re-run when the inputs change
+@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+def generate_image():
+    pipeline = DDPMPipeline.from_pretrained('dvoils/sd-class-butterflies-32')
+    image = pipeline().images[0]
+    return image
 
-    def long_running_task(self): 
-        pipeline = DDPMPipeline.from_pretrained('dvoils/sd-class-butterflies-32')
-        self.image = pipeline().images[0]
-        self.done = True
+st.title('DDPMPipeline Image Generator')
 
 if st.button('Generate Image'):
-    generator = ImageGenerator()
-    thread = threading.Thread(target=generator.long_running_task)
-    thread.start()
-
     with st.spinner('Generating Image...'):
-        progress_bar = st.progress(0)
-        while not generator.done:
-            for i in range(100):
-                time.sleep(0.01)
-                progress_bar.progress(i + 1)
-            progress_bar.progress(0)
-        st.image(generator.image, caption='Generated Image', use_column_width=True)
+        image = generate_image()
+    st.image(image, caption='Generated Image', use_column_width=True)