Update app.py

app.py

@@ -19,23 +19,23 @@ with fs.open("PrakhAI/AIPlane2/g_checkpoint.msgpack", "rb") as f:
 def sample_latent(batch, key):
   return jax.random.normal(key, shape=(batch, LATENT_DIM))
 
-def gridify(images): # num x image_width x image_height x channels
-  # Every num can be padded to make a grid of size floor(sqrt(num)) x ceil(sqrt(num)) or ceil(sqrt(num)) x ceil(sqrt(num))
-  num = images.shape[0]
-  image_width = images.shape[1]
-  image_height = images.shape[2]
-  channels = images.shape[3]
-  width = math.floor(math.sqrt(num))
-  height = math.ceil(math.sqrt(num))
-  if width * height < num:
-    width += 1
-  padded = np.concatenate([images, np.zeros((width*height-num, image_width, image_height, channels))], axis=0)
-  return padded.reshape((width, height, image_width, image_height, -1)).transpose((0, 2, 1, 3, 4)).reshape((width * image_width, height * image_height, -1))
+def to_img(normalized):
+  return ((normalized+1)*255./2.).astype(np.uint8)
 
-st.write("The model and its details are at https://huggingface.co/PrakhAI/AIPlane2")
-num_images = st.number_input(label="Number of images to generate", min_value=1, max_value=256, value=16)
-if st.button('Generate Planes'):
-  latents = sample_latent(num_images, jax.random.PRNGKey(int(1_000_000 * time.time())))
+def generate_images(previous=None):
+  latents = sample_latent(16, jax.random.PRNGKey(int(1_000_000 * time.time())))
+  if previous:
+    latents = np.repeat([previous], repeats=16, axis=0) + 0.25 * latents
   (g_out128, _, _, _, _, _) = generator.apply({'params': g_state['params'], 'batch_stats': g_state['batch_stats']}, latents, training=False)
-  img = (np.array(gridify(g_out128)+1)*255./2.).astype(np.uint8)
-  st.image(Image.fromarray(img))
+  img = np.array(to_img(g_out128))
+  for row in range(4):
+    with st.container():
+      for (col_idx, col) in enumerate(st.columns(4)):
+        with col:
+          idx = row*grid_width + col_idx
+          st.image(Image.fromarray(img[idx]))
+          st.button(label="Generate similar", on_click=generate_images, args=latents[idx])
+
+st.write("The model and its details are at https://huggingface.co/PrakhAI/AIPlane2")
+if st.button('Generate Random'):
+  generate_images()