File size: 1,891 Bytes
c84c172 d9e2ba0 827c2c4 c84c172 0053124 c84c172 c30102f c84c172 c30102f 0053124 38e7538 c30102f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import streamlit as st
from PIL import Image
import jax
import jax.numpy as jnp # JAX NumPy
import numpy as np
from huggingface_hub import HfFileSystem
from flax.serialization import msgpack_restore, from_state_dict
import time
from generator import Generator, LATENT_DIM
import math
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:
g_state = from_state_dict(variables, msgpack_restore(f.read()))
def sample_latent(batch, key):
return jax.random.normal(key, shape=(batch, LATENT_DIM))
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(gridify(g_out128)+1)*255./2.).astype(np.uint8)
st.image(Image.fromarray(img)) |