def get_mm_vq_model_and_params(
    config: ml_collections.ConfigDict) -> Tuple[nn.Module, PyTree]:
  """Loads a multimodal vector quantization model and pretrained parameters.

  Args:
    config: Model config file.

  Returns:
    A 2-tuple of a VQ model and the parameters for that model.
  """
  if config.eval_from.get('checkpoint_path') is not None:
    # uses the user-provided dir to find the model.
    ckpt_dir = config.eval_from.checkpoint_path
  assert gfile.exists(ckpt_dir)

  vq_train_state_dict = checkpoints.restore_checkpoint(ckpt_dir, None)
  vq_params = flax.core.freeze(vq_train_state_dict['ema_params'])

  vq_model = train_utils.get_vq_model(vq_config)['generator']
  parameter_overview.log_parameter_overview(vq_params)

  return vq_model, vq_params
Loads a multimodal vector quantization model and pretrained parameters.

  Args:
    config: Model config file.

  Returns:
    A 2-tuple of a VQ model and the parameters for that model.
  
videogvt/interfaces/mm_vq.py