def expand_pos_embedding(tree_old, tree_new, *, latent_seq_len):
  """Add condition token positions before latent sequence."""
  cur = tree_new['BertEmbed_0']['position_embeddings']['embedding']
  loaded = tree_old['BertEmbed_0']['position_embeddings']['embedding']
  assert cur.shape[0] >= loaded.shape[0]
  # Prefix positions: class or uncond token
  prefix_length = loaded.shape[0] - latent_seq_len
  cur = cur.at[:prefix_length].set(loaded[:prefix_length])
  # Latent sequence
  suffix_length = latent_seq_len
  cur = cur.at[-suffix_length:].set(loaded[-suffix_length:])
  tree_old['BertEmbed_0']['position_embeddings']['embedding'] = cur
  return tree_old
Add condition token positions before latent sequence.
videogvt/train_lib/maskgvt_train_lib.py