Create app.py

app.py

@@ -0,0 +1,44 @@
+import streamlit as st
+from PIL import Image
+import jax
+import numpy as np
+from flax import linen as nn  # Linen API
+from huggingface_hub import HfFileSystem
+from flax.serialization import msgpack_restore, from_state_dict
+import time
+
+LATENT_DIM = 100
+
+class Generator(nn.Module):
+  @nn.compact
+  def __call__(self, latent, training=True):
+    x = latent
+    x = nn.Dense(features=64)(x)
+    x = nn.BatchNorm(not training)(x)
+    x = nn.relu(x)
+    x = nn.Dense(features=2*2*512)(x)
+    x = nn.relu(x)
+    x = x.reshape((x.shape[0], 2, 2, -1))
+    x = nn.ConvTranspose(features=256, kernel_size=(2, 2), strides=(2, 2))(x)
+    x = nn.relu(x)
+    x = nn.ConvTranspose(features=128, kernel_size=(2, 2), strides=(2, 2))(x)
+    x = nn.relu(x)
+    x = nn.ConvTranspose(features=64, kernel_size=(2, 2), strides=(2, 2))(x)
+    x = nn.relu(x)
+    x = nn.ConvTranspose(features=1, kernel_size=(2, 2), strides=(2, 2))(x)
+    x = nn.tanh(x)
+
+generator = Generator()
+
+fs = HfFileSystem()
+with fs.open("PrakhAI/DigitGAN/g_checkpoint.msgpack", "rb") as f:
+  params = from_state_dict(params, msgpack_restore(f.read())["params"])
+
+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': 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)))