fixed small things, also gif and live progress

app.py

@@ -3,6 +3,7 @@ from PIL import Image
 from inference import inference
 import torch
 import io
+from diffusion import DiffusionImageAPI
 
 def main():
 
@@ -22,7 +23,7 @@ def main():
     'Thriller': 13
 }
 
-    st.title("Image Display App")
+    st.title("Movie Diffusion")
     cond = torch.tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
     
     # Add a sidebar for genre selection
@@ -31,6 +32,9 @@ def main():
 
     selected_genres = st.sidebar.multiselect('Select Genres', list(genres_dict.keys()))
 
+    progress_placeholder = st.empty()
+    image_placeholder = st.empty()
+
     
 
     # Button to trigger image generation
@@ -38,20 +42,18 @@ def main():
         for genre in selected_genres:
             code = genres_dict[genre]
             cond[code-1] = code
-        # Display loading sign while generating image
-        with st.spinner('Generating Image...'):
-            # Call the function from inference.py with selected genre
-            image = inference(cond)
-            #image = inference(genre)
-
-        # Convert Pillow image to bytes for display in Streamlit
-        img_buffer = io.BytesIO()
-        #"""0,0,0,0,0,0,0,1, 2, 7, 4, 0, 0, 0"""
-        image.save(img_buffer, format="PNG")
-        img_buffer.seek(0)
-
-        # Display the generated image
-        st.image(img_buffer, caption='Generated Image', use_column_width=True)
+
+        def callback(image, progress):
+            image = DiffusionImageAPI(None).tensor_to_image(image.squeeze(0))
+            img_buffer = io.BytesIO()
+            image.save(img_buffer, format="PNG")
+            img_buffer.seek(0)
+
+            # Update the content of the placeholders
+            progress_placeholder.write(f"Generating Image...\nProgress: {(progress * 100):.2f}%")
+            image_placeholder.image(img_buffer, caption='Generated Image', width=300)
+
+        inference(cond, callback=callback)
 
 if __name__ == "__main__":
     main()

diffusion.py

@@ -160,7 +160,7 @@ class GaussianDiffusion:
 
     return x_t_minus_1
   
-  def sample(self, num_samples, show_progress=True, cond=None, x0=None):
+  def sample(self, num_samples, show_progress=True, cond=None, x0=None, cb=None):
     """
     Sample from the model
     """
@@ -193,6 +193,9 @@ class GaussianDiffusion:
       if show_progress:
         it = tqdm(it)
       for t in it:
+        temp_image = self.denormalize_image(torch.clip(x, -1, 1)).clone().squeeze(0)
+        if cb is not None:
+          cb(temp_image, 1-t/self.noise_steps)
         image_versions.append(self.denormalize_image(torch.clip(x, -1, 1)).clone().squeeze(0))
         
         if x0 is not None and t > 80:

inference.py

@@ -19,7 +19,7 @@ def inference1():
   image = requests.get("https://picsum.photos/120/80").content
   return Image.open(io.BytesIO(image))
 
-def inference(cond, x0=None, gif=False):
+def inference(cond, x0=None, gif=False, callback=None):
   model = Unet(
     image_channels=3,
     dropout=0.1,
@@ -48,7 +48,7 @@ def inference(cond, x0=None, gif=False):
 
   imageAPI = DiffusionImageAPI(diffusion)
 
-  new_images, versions = diffusion.sample(1,cond=cond,x0=x0)
+  new_images, versions = diffusion.sample(1,cond=cond,x0=x0, cb=callback)
   if gif:
     images = []
     for image in versions: