Update app.py

app.py

@@ -2,23 +2,30 @@ import streamlit as st
 from diffusers import DDPMPipeline
 import PIL
 import time
-
-def generate_image():
-    pipeline = DDPMPipeline.from_pretrained('dvoils/sd-class-butterflies-32')
-    image = pipeline().images[0]
-    return image
+import threading
 
 st.title('DDPMPipeline Image Generator')
 
+class ImageGenerator:
+    def __init__(self):
+        self.done = False
+        self.image = None
+
+    def long_running_task(self): 
+        pipeline = DDPMPipeline.from_pretrained('dvoils/sd-class-butterflies-32')
+        self.image = pipeline().images[0]
+        self.done = True
+
 if st.button('Generate Image'):
+    generator = ImageGenerator()
+    thread = threading.Thread(target=generator.long_running_task)
+    thread.start()
+
     with st.spinner('Generating Image...'):
-        # Create a progress bar
         progress_bar = st.progress(0)
-
-        for i in range(100):
-            # Update the progress bar with each iteration.
-            time.sleep(0.01)
-            progress_bar.progress(i + 1)
-
-        image = generate_image()
-        st.image(image, caption='Generated Image', use_column_width=True)
+        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)