Update app.py

app.py

@@ -34,14 +34,13 @@ params = generator.init(jax.random.PRNGKey(0), jnp.zeros([1, LATENT_DIM]), train
 
 fs = HfFileSystem()
 with fs.open("PrakhAI/DigitGAN/g_checkpoint.msgpack", "rb") as f:
-  params = from_state_dict(params, msgpack_restore(f.read())["params"])
-  batch_stats = from_state_dict(params, msgpack_restore(f.read())["batch_stats"])
+  g_state = from_state_dict(params, msgpack_restore(f.read()))
 
 def sample_latent(key):
   return jax.random.normal(key, shape=(1, LATENT_DIM))
 
 if st.button('Generate Digit'):
   latents = sample_latent(jax.random.PRNGKey(int(1_000_000 * time.time())))
-  g_out = Generator().apply({'params': params, 'batch_stats': batch_stats}, latents, training=False)
+  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]
   st.image(Image.fromarray(np.repeat(img, repeats=3, axis=2)))