def restore_legacy_checkpoint(
    checkpoint_path: str,
    train_state: state_mgr.TrainState,
    step: Optional[int] = None,
):
  """Restores a legacy checkpoint.

  Args:
    checkpoint_path: Directory for saving the checkpoint.
    train_state: An instance of TrainState that holds the state of training.
    step: Step number to load or None to load latest. If specified,
      checkpoint_path must be a directory.

  Returns:
    Training state and an int which is the current step.
  """
  restored_train_state = checkpoints.restore_checkpoint(checkpoint_path, None,
                                                        step)
  if restored_train_state is None:
    raise ValueError('No checkpoint for the pretrained model is found in: '
                     f'{checkpoint_path}')
  restored_train_state_type = state_mgr.check_train_state_type(
      restored_train_state)
  if isinstance(train_state, restored_train_state_type):
    if (restored_train_state_type == state_mgr.ScenicTrainState) and (
        'metadata' not in restored_train_state):
      restored_train_state.pop('accum_train_time', None)
    if restored_train_state.get('metadata') is None:
      restored_train_state['metadata'] = {}
    restored_train_state['metadata'].setdefault('lecam_ema_real', None)
    restored_train_state['metadata'].setdefault('lecam_ema_fake', None)
    train_state = flax.serialization.from_state_dict(train_state,
                                                     restored_train_state)
  elif restored_train_state_type == state_mgr.VQGANTrainStateDeprecated:
    assert isinstance(train_state, state_mgr.VQGANTrainState)
    logging.warn('Optimizer state cannot be restored for %s',
                 restored_train_state_type)
    train_state = train_state.replace(  # pytype: disable=attribute-error
        global_step=restored_train_state['global_step'],
        g_model_state=restored_train_state.get('generator_state'),
        d_model_state=restored_train_state.get('discriminator_state'),
        rng=restored_train_state.get('rng'),
        metadata=restored_train_state.get('metadata'),
        g_params=restored_train_state['g_optimizer']['target'],
        d_params=restored_train_state['d_optimizer']['target'],
        ema_params=restored_train_state['ema_params'],
    )
  else:
    raise ValueError(f'Train_state {type(train_state)} is not supported.')
  return train_state
Restores a legacy checkpoint.

  Args:
    checkpoint_path: Directory for saving the checkpoint.
    train_state: An instance of TrainState that holds the state of training.
    step: Step number to load or None to load latest. If specified,
      checkpoint_path must be a directory.

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