def restore_finetune_checkpoint(
    checkpoint_path: str,
    train_state: state_mgr.VQGANTrainState,
    step: Optional[int] = None,
    not_restore_g_params_keys: Optional[List[str]] = None,
):
  """Restores a checkpoint for VQ finetuning.

  Args:
    checkpoint_path: Directory for saving the checkpoint.
    train_state: the target TrainState.
    step: Step number to load or None to load the latest. If specified,
      checkpoint_path must be a directory.
    not_restore_g_params_keys: List of g_params keys not to restore from the
      checkpoint. Default (None) loads every key.

  Returns:
    Training state and an int which is the current step.
  """
  restored_train_state = checkpoints.restore_checkpoint(
      ckpt_dir=checkpoint_path, target=None, step=step
  )

  g_params = train_state.g_params
  restored_g_params = restored_train_state['g_params']
  for k in restored_g_params.keys():
    if not_restore_g_params_keys and k not in not_restore_g_params_keys:
      g_params[k] = restored_g_params[k]

  train_state = train_state.replace(
      # global_step=restored_train_state['global_step'],  # step starts from 0
      g_model_state=restored_train_state.get('g_model_state'),
      d_model_state=restored_train_state.get('d_model_state'),
      rng=restored_train_state.get('rng'),
      metadata=restored_train_state.get('metadata'),
      g_params=restored_train_state['g_params'],
      d_params=restored_train_state['d_params'],
      ema_params=restored_train_state['ema_params'],
  )
  return train_state
Restores a checkpoint for VQ finetuning.

  Args:
    checkpoint_path: Directory for saving the checkpoint.
    train_state: the target TrainState.
    step: Step number to load or None to load the latest. If specified,
      checkpoint_path must be a directory.
    not_restore_g_params_keys: List of g_params keys not to restore from the
      checkpoint. Default (None) loads every key.

  Returns:
    Training state and an int which is the current step.
  
videogvt/train_lib/train_utils.py