New model.

app.py

@@ -1,19 +1,14 @@
-import functools
 import os
 
-from absl import flags
 from clu import checkpoint
 import gradio as gr
 import jax
 import jax.numpy as jnp
 import numpy as np
 from PIL import Image
-import tensorflow as tf
 from huggingface_hub import snapshot_download
 
 from invariant_slot_attention.configs.clevr_with_masks.equiv_transl_scale import get_config
-from invariant_slot_attention.lib import input_pipeline
-from invariant_slot_attention.lib import preprocessing
 from invariant_slot_attention.lib import utils
 
 
@@ -33,7 +28,7 @@ def load_model(config, checkpoint_dir):
         {"params": model_rng, "state_init": init_rng, "dropout": dropout_rng},
         video=init_inputs, conditioning=init_conditioning,
         padding_mask=jnp.ones(init_inputs.shape[:-1], jnp.int32))
-
+ 
     # Split into state variables (e.g. for batchnorm stats) and model params.
     # Note that `pop()` on a FrozenDict performs a deep copy.
     state_vars, initial_params = initial_vars.pop("params")  # pytype: disable=attribute-error
@@ -47,11 +42,11 @@ def load_model(config, checkpoint_dir):
 
   opt_state = None
   state = utils.TrainState(
-      step=42, opt_state=opt_state, params=initial_params, rng=rng,
+      step=1, opt_state=opt_state, params=initial_params, rng=rng,
       variables=state_vars)
 
-  ckpt = checkpoint.MultihostCheckpoint(checkpoint_dir) 
-  state = ckpt.restore(state)
+  ckpt = checkpoint.Checkpoint(checkpoint_dir)
+  state = ckpt.restore(state, checkpoint=checkpoint_dir + "/ckpt-0")
 
   return model, state, rng
 
@@ -67,7 +62,7 @@ def load_image(name):
 
 
 download_path = snapshot_download(repo_id="ondrejbiza/isa")
-checkpoint_dir = os.path.join(download_path, "clevr_isa_ts", "checkpoints")
+checkpoint_dir = os.path.join(download_path, "clevr_isa_ts")
 
 model, state, rng = load_model(get_config(), checkpoint_dir)