Update app, add CLEVR images.

app.py

@@ -2,6 +2,7 @@ import functools
 import os
 
 from absl import flags
+from clu import checkpoint
 import gradio as gr
 import jax
 import jax.numpy as jnp
@@ -12,6 +13,7 @@ from invariant_slot_attention.lib import utils
 
 
 def load_model(config):
+  rng = jax.random.PRNGKey(42)
   rng, data_rng = jax.random.split(rng)
 
   # Initialize model
@@ -21,9 +23,7 @@ def load_model(config):
     rng, init_rng, model_rng, dropout_rng = jax.random.split(rng, num=4)
 
     init_conditioning = None
-    init_inputs = jnp.ones(
-        [1] + list(train_ds.element_spec["video"].shape)[2:],
-        jnp.float32)
+    init_inputs = jnp.ones([1, 1, 128, 128, 3], jnp.float32)
     initial_vars = model.init(
         {"params": model_rng, "state_init": init_rng, "dropout": dropout_rng},
         video=init_inputs, conditioning=init_conditioning,
@@ -40,27 +40,29 @@ def load_model(config):
 
   state_vars, initial_params = init_model(rng)
 
-  learning_rate_fn = lr_schedules.get_learning_rate_fn(config)
-  tx = optimizers.get_optimizer(
-      config.optimizer_configs, learning_rate_fn, params=initial_params)
-
-  opt_state = tx.init(initial_params)
-
+  opt_state = None
   state = utils.TrainState(
       step=1, opt_state=opt_state, params=initial_params, rng=rng,
       variables=state_vars)
 
-  loss_fn = functools.partial(
-      losses.compute_full_loss, loss_config=config.losses)
-
-  checkpoint_dir = os.path.join(workdir, "checkpoints")
+  checkpoint_dir = "clevr_isa_ts/checkpoints-0"
   ckpt = checkpoint.MultihostCheckpoint(checkpoint_dir)
   state = ckpt.restore_or_initialize(state)
 
+  init_inputs = jnp.ones([1, 1, 128, 128, 3], jnp.float32)
+  rng, init_rng = jax.random.split(rng, num=2)
+  out = model.apply(
+    {"params": state.params, **state.variables},
+    video=init_inputs,
+    rngs={"state_init": init_rng},
+    train=False)
+  print(out.keys())
+
 
 def greet(name):
   return "Hello " + name + "!"
 
 
+load_model(get_config())
 demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()