Update app.py

app.py

@@ -47,6 +47,10 @@ def main(_):
         use_lock=False,
     )
 
+    checkerboard_r1 = np.concatenate([np.zeros((8, 8, 3)), np.ones((8, 8, 3)), np.zeros((8, 8, 3))], axis=1)
+    checkerboard_r2 = np.concatenate([np.ones((8, 8, 3)), np.zeros((8, 8, 3)), np.ones((8, 8, 3))], axis=1)
+    checkerboard = np.concatenate([checkerboard_r1, checkerboard_r2] * 16, axis=0).astype(np.float32)
+
     def generate_images(input_images, n_new_frames, n_candidates, temperature=1.0, top_p=0.9):
         assert len(input_images) > 0
         input_images = [
@@ -58,10 +62,14 @@ def main(_):
 
         generated_images = []
         for candidate in output_images:
-            combined_row = np.concatenate([input_images[0]] + list(candidate), axis=1)
+            concatenated_image = []
+            for i, img in enumerate(candidate):
+                concatenated_image.append(img)
+                if i < len(candidate) - 1:
+                    concatenated_image.append(checkerboard)
             generated_images.append(
                 Image.fromarray(
-                    (combined_row * 255).astype(np.uint8)
+                    (np.concatenate(concatenated_image, axis=1) * 255).astype(np.uint8)
                 )
             )
 
@@ -142,7 +150,7 @@ def main(_):
             )
             gr.Markdown('## Outputs (multi candidates)')
             output_gallery = gr.Gallery(
-                columns=1,
+                columns=4,
                 object_fit='scale-down',
                 label="Output image"
             )
@@ -238,4 +246,4 @@ def main(_):
     demo.launch()
 
 if __name__ == "__main__":
-    mlxu.run(main)
+    mlxu.run(main)