Update app.py

app.py

@@ -22,24 +22,26 @@ def sample_latent(batch, key):
 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")
+if st.button('Generate Random'):
+  st.session_state['generate'] = None
+
 ROWS = 4
 COLUMNS = 4
-def generate_images(previous=None):
+
+if 'generate' in st.session_state:
   unique_id = int(1_000_000 * time.time())
   latents = sample_latent(ROWS * COLUMNS, jax.random.PRNGKey(unique_id))
+  previous = st.session_state['generate']
   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(ROWS):
-    with st.container():
-      for (col_idx, col) in enumerate(st.columns(COLUMNS)):
-        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)):
-            generate_images(latents[idx])
-
-st.write("The model and its details are at https://huggingface.co/PrakhAI/AIPlane2")
-if st.button('Generate Random'):
-  generate_images()
+  with st.container():
+    for (col_idx, col) in enumerate(st.columns(COLUMNS)):
+      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]