Update app.py

app.py

@@ -104,9 +104,9 @@ class Pipeline:
             return animations
 
         images = decode_images(samples)
-        # animations = decode_trajs(trajs)
+        animations = decode_trajs(trajs)
 
-        return images[0]
+        return images[0], animations[0]
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -155,8 +155,8 @@ if __name__ == "__main__":
             with gr.Column(scale=1):
                 btn = gr.Button("Generate")
                 output_sample = gr.Image(label="Images")
-            with gr.Column(scale=2):
-                 output_trajs = gr.Gallery(label="Trajs of Diffusion", columns=2, rows=2)
+            with gr.Column(scale=1):
+                 output_trajs = gr.Image(label="Trajs of Diffusion", columns=2, rows=2)
 
         btn.click(fn=pipeline,
                   inputs=[