Update app.py

app.py

@@ -47,10 +47,6 @@ 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 = [
@@ -62,14 +58,14 @@ def main(_):
 
         generated_images = []
         for candidate in output_images:
-            concatenated_image = []
-            for i, img in enumerate(candidate):
-                concatenated_image.append(img)
-                if i < len(candidate) - 1:
-                    concatenated_image.append(checkerboard)
+            concatenated_rows = []
+            for i in range(len(candidate)):
+                concatenated_row = np.concatenate([input_images[i]] + candidate, axis=1)
+                concatenated_rows.append(concatenated_row)
+            combined_image = np.concatenate(concatenated_rows, axis=0)
             generated_images.append(
                 Image.fromarray(
-                    (np.concatenate(concatenated_image, axis=1) * 255).astype(np.uint8)
+                    (combined_image * 255).astype(np.uint8)
                 )
             )
 
@@ -247,4 +243,3 @@ def main(_):
 
 if __name__ == "__main__":
     mlxu.run(main)
-