|
import streamlit as st |
|
from PIL import Image |
|
import jax |
|
import jax.numpy as jnp |
|
import numpy as np |
|
from flax import linen as nn |
|
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") |