Update app.py

app.py

@@ -29,6 +29,9 @@ if st.button('Generate Random'):
 ROWS = 4
 COLUMNS = 4
 
+def set_latent(latent):
+  st.session_state['generate'] = latent
+
 if 'generate' in st.session_state:
   unique_id = int(1_000_000 * time.time())
   latents = sample_latent(ROWS * COLUMNS, jax.random.PRNGKey(unique_id))
@@ -46,5 +49,4 @@ if 'generate' in st.session_state:
         with col:
           idx = row*COLUMNS + col_idx
           st.image(Image.fromarray(img[idx]))
-          if st.button(label="Generate similar", key="%d_%d" % (unique_id, idx)):
-            st.session_state['generate'] = latents[idx]
+          st.button(label="Generate similar", key="%d_%d" % (unique_id, idx), on_click=set_latent, args=(latents[idx]))