Update app.py

app.py

@@ -3,11 +3,9 @@ from PIL import Image
 import jax
 import jax.numpy as jnp  # JAX NumPy
 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
-from local_response_norm import LocalResponseNorm
 from generator import Generator, LATENT_DIM
 
 generator = Generator()
@@ -17,12 +15,26 @@ fs = HfFileSystem()
 with fs.open("PrakhAI/AIPlane2/g_checkpoint.msgpack", "rb") as f:
   g_state = from_state_dict(variables, msgpack_restore(f.read()))
 
-def sample_latent(key):
-  return jax.random.normal(key, shape=(1, LATENT_DIM))
+def sample_latent(batch, key):
+  return jax.random.normal(key, shape=(batch, LATENT_DIM))
 
-if st.button('Generate Plane'):
-  latents = sample_latent(jax.random.PRNGKey(int(1_000_000 * time.time())))
+def gridify(images): # num x image_width x image_height x channels
+  # Every num can be padded to make a grid of size floor(sqrt(num)) x ceil(sqrt(num)) or ceil(sqrt(num)) x ceil(sqrt(num))
+  num = images.shape[0]
+  image_width = images.shape[1]
+  image_height = images.shape[2]
+  channels = images.shape[3]
+  width = math.floor(math.sqrt(num))
+  height = math.ceil(math.sqrt(num))
+  if width * height < num:
+    width += 1
+  padded = np.concatenate([images, np.zeros((width*height-num, image_width, image_height, channels))], axis=0)
+  return padded.reshape((width, height, image_width, image_height, -1)).transpose((0, 2, 1, 3, 4)).reshape((width * image_width, height * image_height, -1))
+
+st.write("The model and its details are at https://huggingface.co/PrakhAI/AIPlane2")
+num_images = st.number_input(label="Number of images to generate", min_value=1, max_value=256, value=16)
+if st.button('Generate Planes'):
+  latents = sample_latent(num_images, jax.random.PRNGKey(int(1_000_000 * time.time())))
   (g_out128, _, _, _, _, _) = 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(img))
-  st.write("The model and its details are at https://huggingface.co/PrakhAI/AIPlane2")
+  img = ((np.array(gridify(g_out128)+1)*255./2.).astype(np.uint8)
+  st.image(Image.fromarray(img))