Update app.py

app.py

@@ -8,46 +8,15 @@ 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 = 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:
+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):
@@ -55,7 +24,7 @@ def sample_latent(key):
 
 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)
+  (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/AIPlane")
+  st.write("The model and its details are at https://huggingface.co/PrakhAI/AIPlane2")