def load_mm_vq_model(
    config: ml_collections.ConfigDict,
    rvq_quantize_steps: int = 0) -> Dict[str, TokenizerFn]:
  """Loads a pretrained multimodal vector quantization model.

  Args:
    config: Model config file.
    rvq_quantize_steps: Number of RVQ levels to perform if the RVQGAN is used.

  Returns:
    Dictionary of two functions: tokenizer and detokenizer.
  """
  vq_model, vq_params = get_mm_vq_model_and_params(config)

  if rvq_quantize_steps > 0:
    assert (
        'quantize_steps' in vq_model.config.vqvae
    ), 'quantize_steps NOT found in training.'
    vq_model.config.unlock()
    vq_model.config.vqvae.quantize_steps = rvq_quantize_steps
    vq_model.config.lock()

  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_fn = functools.partial(
      vq_model.apply,
      {'params': dec_params},
      method=vq_model.decode_from_indices,
  )
  # TODO(roadjiang): add audio tokenizers in this function
  return {'tokenizer': tokenizer_fn, 'detokenizer': detokenizer_fn}
Loads a pretrained multimodal vector quantization model.

  Args:
    config: Model config file.
    rvq_quantize_steps: Number of RVQ levels to perform if the RVQGAN is used.

  Returns:
    Dictionary of two functions: tokenizer and detokenizer.
  
videogvt/interfaces/mm_vq.py