Update app.py

app.py

@@ -25,7 +25,8 @@ def to_img(normalized):
 ROWS = 4
 COLUMNS = 4
 def generate_images(previous=None):
-  latents = sample_latent(ROWS * COLUMNS, jax.random.PRNGKey(int(1_000_000 * time.time())))
+  unique_id = int(1_000_000 * time.time())
+  latents = sample_latent(ROWS * COLUMNS, jax.random.PRNGKey(unique_id))
   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)
@@ -36,7 +37,7 @@ def generate_images(previous=None):
         with col:
           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", key="%d_%d" % (unique_id, idx), 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'):