def load_mm_vq_staged_decoder(
    config: ml_collections.ConfigDict, last_conv_name: str = 'Conv_4'
) -> Dict[str, TokenizerFn]:
  """Loads the two-staged VQ decoders and a standard encoder."""
  vq_model, vq_params = get_mm_vq_model_and_params(config)

  enc_params, _ = vq_params.pop('decoder')
  tokenizer_fn = functools.partial(
      vq_model.apply, {'params': enc_params}, method=vq_model.encode_to_indices
  )

  dec_params, _ = vq_params.pop('encoder')
  detokenizer_fn1 = functools.partial(
      vq_model.apply, {'params': dec_params}, method=vq_model.decode_stage1
  )

  # This is a hack for enc_dec_3dcnn or 3dcnn
  subset_params = {}
  for k, v in vq_params['decoder'].items():
    if k.startswith(last_conv_name):
      subset_params['Conv_0'] = v
      continue
  subset_params = {'decoder': subset_params}
  subset_params = flax.core.FrozenDict(subset_params)

  # parameter_overview.log_parameter_overview(new_params)
  detokenizer_fn2 = functools.partial(
      vq_model.apply, {'params': subset_params},
      method=vq_model.decode_stage2)
  return {'tokenizer': tokenizer_fn,
          'detokenizer1': detokenizer_fn1, 'detokenizer2': detokenizer_fn2}
Loads the two-staged VQ decoders and a standard encoder.
videogvt/interfaces/mm_vq.py