Duplicate from PrakhAI/AIPlane2

.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: AIPlane2
+emoji: 🌖
+colorFrom: green
+colorTo: indigo
+sdk: streamlit
+sdk_version: 1.25.0
+app_file: app.py
+pinned: false
+duplicated_from: PrakhAI/AIPlane2
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,52 @@
+import streamlit as st
+from PIL import Image
+import jax
+import jax.numpy as jnp  # JAX NumPy
+import numpy as np
+from huggingface_hub import HfFileSystem
+from flax.serialization import msgpack_restore, from_state_dict
+import time
+from generator import Generator, LATENT_DIM
+import math
+
+generator = Generator()
+variables = generator.init(jax.random.PRNGKey(0), jnp.zeros([1, LATENT_DIM]), training=False)
+
+fs = HfFileSystem()
+with fs.open("PrakhAI/AIPlane2/g_checkpoint.msgpack", "rb") as f:
+  g_state = from_state_dict(variables, msgpack_restore(f.read()))
+
+def sample_latent(batch, key):
+  return jax.random.normal(key, shape=(batch, LATENT_DIM))
+
+def to_img(normalized):
+  return ((normalized+1)*255./2.).astype(np.uint8)
+
+st.write("The model and its details are at https://huggingface.co/PrakhAI/AIPlane2")
+if st.button('Generate Random'):
+  st.session_state['generate'] = None
+
+ROWS = 4
+COLUMNS = 4
+
+def set_latent(latent):
+  st.session_state['generate'] = latent
+
+if 'generate' in st.session_state:
+  unique_id = int(1_000_000 * time.time())
+  latents = sample_latent(ROWS * COLUMNS, jax.random.PRNGKey(unique_id))
+  previous = st.session_state['generate']
+  if previous is not None:
+    if "similarity" not in st.session_state:
+      st.session_state["similarity"] = 0.5
+    similarity = st.number_input(label="Mutation (for \"Generate Similar\") - lower value generates more similar images", key="similarity", min_value=0.01, max_value=1.0)
+    latents = np.repeat([previous], repeats=16, axis=0) + similarity * latents
+  (g_out128, _, _, _, _, _) = generator.apply({'params': g_state['params'], 'batch_stats': g_state['batch_stats']}, latents, training=False)
+  img = np.array(to_img(g_out128))
+  for row in range(ROWS):
+    with st.container():
+      for (col_idx, col) in enumerate(st.columns(COLUMNS)):
+        with col:
+          idx = row*COLUMNS + col_idx
+          st.image(Image.fromarray(img[idx]))
+          st.button(label="Generate Similar", key="%d_%d" % (unique_id, idx), on_click=set_latent, args=(latents[idx],))

generator.py

@@ -0,0 +1,59 @@
+from flax import linen as nn
+import jax
+import jax.numpy as jnp
+from local_response_norm import LocalResponseNorm
+
+LATENT_DIM = 500
+EPSILON = 1e-8
+
+class Generator(nn.Module):
+  @nn.compact
+  def __call__(self, latent, training=True):
+    x = nn.Dense(features=64)(latent)
+    # x = nn.BatchNorm(not training)(x)
+    x = nn.relu(x)
+    x = nn.Dense(features=2*2*1024)(x)
+    x = nn.BatchNorm(not training)(x)
+    x = nn.relu(x)
+    x = nn.Dropout(0.25, deterministic=not training)(x)
+    x = x.reshape((x.shape[0], 2, 2, -1))
+    x4 = nn.ConvTranspose(features=512, kernel_size=(3, 3), strides=(2, 2))(x)
+    x4 = LocalResponseNorm()(x4)
+    x4 = nn.relu(x4)
+    x4o = nn.ConvTranspose(features=3, kernel_size=(3, 3))(x4)
+    x4 = nn.ConvTranspose(features=512, kernel_size=(3, 3))(x4)
+    x4 = LocalResponseNorm()(x4)
+    x4 = nn.relu(x4)
+    x8 = nn.ConvTranspose(features=256, kernel_size=(3, 3), strides=(2, 2))(x4)
+    x8 = LocalResponseNorm()(x8)
+    x8 = nn.relu(x8)
+    x8o = nn.ConvTranspose(features=3, kernel_size=(3, 3))(x8)
+    x8 = nn.ConvTranspose(features=256, kernel_size=(3, 3))(x8)
+    x8 = LocalResponseNorm()(x8)
+    x8 = nn.relu(x8)
+    x16 = nn.ConvTranspose(features=128, kernel_size=(3, 3), strides=(2, 2))(x8)
+    x16 = LocalResponseNorm()(x16)
+    x16 = nn.relu(x16)
+    x16o = nn.ConvTranspose(features=3, kernel_size=(3, 3))(x16)
+    x16 = nn.ConvTranspose(features=128, kernel_size=(3, 3))(x16)
+    x16 = LocalResponseNorm()(x16)
+    x16 = nn.relu(x16)
+    x32 = nn.ConvTranspose(features=64, kernel_size=(3, 3), strides=(2, 2))(x16)
+    x32 = LocalResponseNorm()(x32)
+    x32 = nn.relu(x32)
+    x32o = nn.ConvTranspose(features=3, kernel_size=(3, 3))(x32)
+    x32 = nn.ConvTranspose(features=64, kernel_size=(3, 3))(x32)
+    x32 = LocalResponseNorm()(x32)
+    x32 = nn.relu(x32)
+    x64 = nn.ConvTranspose(features=32, kernel_size=(3, 3), strides=(2, 2))(x32)
+    x64 = LocalResponseNorm()(x64)
+    x64 = nn.relu(x64)
+    x64o = nn.ConvTranspose(features=3, kernel_size=(3, 3))(x64)
+    x64 = nn.ConvTranspose(features=32, kernel_size=(3, 3))(x64)
+    x64 = LocalResponseNorm()(x64)
+    x64 = nn.relu(x64)
+    x128 = nn.ConvTranspose(features=64, kernel_size=(3, 3), strides=(2, 2))(x64)
+    x128 = LocalResponseNorm()(x128)
+    x128 = nn.relu(x128)
+    x128o = nn.ConvTranspose(features=3, kernel_size=(3, 3))(x128)
+    return (nn.tanh(x128o), nn.tanh(x64o), nn.tanh(x32o), nn.tanh(x16o), nn.tanh(x8o), nn.tanh(x4o))

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