Update app.py

app.py

@@ -37,7 +37,11 @@ if 'generate' in st.session_state:
   latents = sample_latent(ROWS * COLUMNS, jax.random.PRNGKey(unique_id))
   previous = st.session_state['generate']
   if previous is not None:
-    latents = np.repeat([previous], repeats=16, axis=0) + 0.25 * latents
+    if 'similarity' not in st.session_state:
+      st.session_state['similarity'] = 0.25
+    similarity = st.number_input(label="Similarity (for \"Generate Similar\"", min_value=0.01, max_value=1.0, value=st.session_state['similarity'])
+    st.session_state['similarity'] = similarity
+    latents = np.repeat([previous], repeats=16, axis=0) + similarity * 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):
@@ -46,4 +50,4 @@ if 'generate' in st.session_state:
         with col:
           idx = row*COLUMNS + col_idx
           st.image(Image.fromarray(img[idx]))
-          st.button(label="Generate similar", key="%d_%d" % (unique_id, idx), on_click=set_latent, args=(latents[idx],))
+          st.button(label="Generate Similar", key="%d_%d" % (unique_id, idx), on_click=set_latent, args=(latents[idx],))