def export_mm_vq_model_scenic(config: ml_collections.ConfigDict,
                              output_location: str,
                              batch_dim: Optional[int] = None):
  """Saves a multimodal vector quantization model as tf.SavedModel by scenic."""

  vq_model, vq_params = get_mm_vq_model_and_params(config)
  num_frames = config.get('dataset_configs.num_frames', 16)
  image_size = config.get('image_size', 256)
  latent_shape = config.get('transformer.latent_shape',
                            (num_frames // 4, image_size // 8, image_size // 8))

  enc_params, _ = vq_params.pop('decoder')
  enc_input_shape = (-1, num_frames, image_size, image_size, 3)
  def _tokenize(params, input_data):
    return vq_model.apply({'params': params},
                          input_data,
                          method=vq_model.encode_to_indices)
  export_model(
      _tokenize,
      enc_params,
      export_dir=f'{output_location}/tokenizer',
      with_gradient=False,
      batch_dim=batch_dim,
      input_shape=enc_input_shape)

  dec_params, _ = vq_params.pop('encoder')
  dec_input_shape = (-1, *latent_shape)

  def _detokenize(params, input_data):
    return vq_model.apply({'params': params},
                          input_data,
                          method=vq_model.decode_from_indices)
  export_model(
      _detokenize,
      dec_params,
      export_dir=f'{output_location}/detokenizer',
      with_gradient=False,
      input_shape=dec_input_shape,
      batch_dim=batch_dim,
      dtype=tf.int32)
Saves a multimodal vector quantization model as tf.SavedModel by scenic.
videogvt/interfaces/mm_vq.py