File size: 2,608 Bytes
c84c172 |
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 49 50 51 52 53 54 55 56 57 58 59 60 61 |
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
from local_response_norm import LocalResponseNorm
LATENT_DIM = 100
class Generator(nn.Module):
@nn.compact
def __call__(self, latent, training=True):
x = nn.Dense(features=32)(latent)
# x = nn.BatchNorm(not training)(x)
x = nn.relu(x)
x = nn.Dense(features=2*2*256)(x)
x = nn.BatchNorm(not training)(x)
x = nn.relu(x)
x = nn.Dropout(0.5, deterministic=not training)(x)
x = x.reshape((x.shape[0], 2, 2, -1))
x4o = nn.ConvTranspose(features=3, kernel_size=(2, 2), strides=(2, 2))(x)
x4 = nn.ConvTranspose(features=128, kernel_size=(2, 2), strides=(2, 2))(x)
x4 = LocalResponseNorm()(x4)
# x4 = nn.BatchNorm(not training)(x4)
x8 = nn.relu(x4)
# x8 = nn.Dropout(0.5, deterministic=not training)(x8)
x8o = nn.ConvTranspose(features=3, kernel_size=(2, 2), strides=(2, 2))(x8)
x8 = nn.ConvTranspose(features=64, kernel_size=(2, 2), strides=(2, 2))(x8)
x8 = LocalResponseNorm()(x8)
# x8 = nn.BatchNorm(not training)(x8)
x16 = nn.relu(x8)
# x16 = nn.Dropout(0.5, deterministic=not training)(x16)
x16o = nn.ConvTranspose(features=3, kernel_size=(2, 2), strides=(2, 2))(x16)
x16 = nn.ConvTranspose(features=32, kernel_size=(2, 2), strides=(2, 2))(x16)
x16 = LocalResponseNorm()(x16)
# x16 = nn.BatchNorm(not training)(x16)
x32 = nn.relu(x16)
# x32 = nn.Dropout(0.5, deterministic=not training)(x32)
x32o = nn.ConvTranspose(features=3, kernel_size=(2, 2), strides=(2, 2))(x32)
return (nn.tanh(x32o), nn.tanh(x16o), nn.tanh(x8o), nn.tanh(x4o))
generator = Generator()
variables = generator.init(jax.random.PRNGKey(0), jnp.zeros([1, LATENT_DIM]), training=False)
fs = HfFileSystem()
with fs.open("PrakhAI/AIPlane/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 Plane'):
latents = sample_latent(jax.random.PRNGKey(int(1_000_000 * time.time())))
(g_out32, g_out16, g_out8, g_out4) = 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/AIPlane") |