File size: 1,782 Bytes
de78f83 1436bf3 de78f83 9c9ee7c de78f83 31266d7 de78f83 31266d7 de78f83 39de372 de78f83 63d67a7 |
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 42 43 44 45 46 47 48 |
import streamlit as st
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
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)
return x
generator = Generator()
variables = generator.init(jax.random.PRNGKey(0), jnp.zeros([1, LATENT_DIM]), training=False)
fs = HfFileSystem()
with fs.open("PrakhAI/DigitGAN/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))
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)))
st.write("The model's details are at https://huggingface.co/PrakhAI/DigitGAN/blob/main/README.md") |