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
from generator import Generator, LATENT_DIM

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(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_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/AIPlane2")