Update app.py

app.py

@@ -13,7 +13,7 @@ generator = Generator()
 variables = generator.init(jax.random.PRNGKey(0), jnp.zeros([1, LATENT_DIM]), training=False)
 
 fs = HfFileSystem()
-with fs.open("PrakhAI/AIPlane2/g_checkpoint.msgpack", "rb") as f:
+with fs.open("PrakhAI/AIPlane3/g_checkpoint_200000.msgpack", "rb") as f:
   g_state = from_state_dict(variables, msgpack_restore(f.read()))
 
 def sample_latent(batch, key):
@@ -22,12 +22,12 @@ 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")
+st.write("The model and its details are at https://huggingface.co/PrakhAI/AIPlane3")
 if st.button('Generate Random'):
   st.session_state['generate'] = None
 
-ROWS = 4
-COLUMNS = 4
+ROWS = 2
+COLUMNS = 2
 
 def set_latent(latent):
   st.session_state['generate'] = latent
@@ -40,9 +40,9 @@ if 'generate' in st.session_state:
     if "similarity" not in st.session_state:
       st.session_state["similarity"] = 0.5
     similarity = st.number_input(label="Mutation (for \"Generate Similar\") - lower value generates more similar images", key="similarity", min_value=0.01, max_value=1.0)
-    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))
+    latents = np.repeat([previous], repeats=4, axis=0) + similarity * latents
+  g_out = generator.apply({'params': g_state['params']}, latents)
+  img = np.array(to_img(g_out))
   for row in range(ROWS):
     with st.container():
       for (col_idx, col) in enumerate(st.columns(COLUMNS)):