Update app.py

app.py

@@ -22,19 +22,21 @@ def sample_latent(batch, key):
 def to_img(normalized):
   return ((normalized+1)*255./2.).astype(np.uint8)
 
+ROWS = 4
+COLUMNS = 4
 def generate_images(previous=None):
-  latents = sample_latent(16, jax.random.PRNGKey(int(1_000_000 * time.time())))
+  latents = sample_latent(ROWS * COLUMNS, 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(to_img(g_out128))
-  for row in range(4):
+  for row in range(ROWS):
     with st.container():
-      for (col_idx, col) in enumerate(st.columns(4)):
+      for (col_idx, col) in enumerate(st.columns(COLUMNS)):
         with col:
-          idx = row*grid_width + col_idx
+          idx = row*COLUMNS + col_idx
           st.image(Image.fromarray(img[idx]))
-          st.button(label="Generate similar", on_click=generate_images, args=latents[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'):