Duplicate from PrakhAI/AIPlane

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: AIPlane
+emoji: 🌖
+colorFrom: green
+colorTo: indigo
+sdk: streamlit
+sdk_version: 1.25.0
+app_file: app.py
+pinned: false
+duplicated_from: PrakhAI/AIPlane
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,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")

local_response_norm.py

@@ -0,0 +1,11 @@
+from flax import linen as nn
+import jax
+import jax.numpy as jnp
+
+class LocalResponseNorm(nn.Module):
+  @nn.compact
+  def __call__(
+      self,
+      value: jax.Array
+  ) -> jax.Array:
+    return value / jnp.repeat(jnp.expand_dims((1e-8 + (value**2).mean(axis=-1))**0.5, axis=-1), repeats=value.shape[-1], axis=-1)

requirements.txt

@@ -0,0 +1 @@
+flax