Update app.py

app.py

@@ -55,7 +55,7 @@ def sample_latent(key):
 
 if st.button('Generate Plane'):
   latents = sample_latent(jax.random.PRNGKey(int(1_000_000 * time.time())))
-  g_out = generator.apply({'params': g_state['params'], 'batch_stats': g_state['batch_stats']}, latents, training=False)
-  img = ((np.array(g_out)+1)*255./2.).astype(np.uint8)[0]
+  (g_out32, g_out16, g_out8, g_out4) = generator.apply({'params': g_state['params'], 'batch_stats': g_state['batch_stats']}, latents, training=False)
+  img = ((np.array(g_out32[0])+1)*255./2.).astype(np.uint8)
   st.image(Image.fromarray(np.repeat(img, repeats=3, axis=2)))
   st.write("The model's details are at https://huggingface.co/PrakhAI/AIPlane/blob/main/README.md")