def init_from_pretrained_checkpoint(checkpoint_path: str,
                                    train_state: TrainState,
                                    config: ml_collections.ConfigDict):
  """Initialize the train state with a pretrained checkpoint.

  First restores the checkpoint, then inflates the parameters if requested,
  finally build a new train state with updated parameters.

  Args:
    checkpoint_path: Directory for saving the checkpoint.
    train_state: An instance of TrainState that holds the state of training.
    config: Model config.

  Returns:
    Training state.
  """
  assert train_state is not None, 'train_state cannot be None.'
  restored_train_state = checkpoints.restore_checkpoint(checkpoint_path, None)
  if restored_train_state is None:
    raise ValueError('No checkpoint for the pretrained model is found in: '
                     f'{checkpoint_path}')
  restored_train_state_type = train_state_manager.check_train_state_type(
      restored_train_state)
  restored_train_state = flax.core.freeze(restored_train_state)
  if restored_train_state_type == TrainStateDeprecated:
    # deprecated train_state format
    g_restored_params = restored_train_state['g_optimizer']['target']
    g_restored_model_state = restored_train_state.get('generator_state')
    d_restored_params = restored_train_state['d_optimizer']['target']
    d_restored_model_state = restored_train_state.get('discriminator_state')
  elif restored_train_state_type == TrainState:
    # model trained using new optax optimizers.
    g_restored_params = restored_train_state['g_params']
    g_restored_model_state = restored_train_state.get('g_model_state')
    d_restored_params = restored_train_state['d_params']
    d_restored_model_state = restored_train_state.get('d_model_state')
  else:
    raise ValueError('restored_train_state is not supported.')
  ema_restored_params = restored_train_state['ema_params']

  inflation = config.init_from.get('inflation')
  if inflation is not None:
    if '/' in inflation:
      direction, mode = inflation.split('/')
    else:
      direction, mode = inflation, 'central'
    if direction == '2d->3d':
      inflate_fn = functools.partial(inflate_2d_to_3d, mode=mode)
    else:
      raise NotImplementedError(f'Inflation {inflation}')
    inflate_fn = jax.jit(inflate_fn, backend='cpu')
    g_restored_params = inflate_fn(g_restored_params,
                                   train_state.g_params)
    g_restored_model_state = inflate_fn(
        g_restored_model_state or flax.core.FrozenDict(),
        train_state.g_model_state)
    d_restored_params = inflate_fn(d_restored_params,
                                   train_state.d_params)
    d_restored_model_state = inflate_fn(
        d_restored_model_state or flax.core.FrozenDict(),
        train_state.d_model_state)
    ema_restored_params = inflate_fn(ema_restored_params,
                                     train_state.ema_params)
  if config.init_from.get('embedding_path') is not None:
    embedding = _load_embedding(
        config.init_from.embedding_path,
        config.init_from.get('embedding_start'),
        config.init_from.get('embedding_end'),
    )
    g_restored_params = _update_embedding(g_restored_params, embedding)
    ema_restored_params = _update_embedding(ema_restored_params, embedding)
  # pytype: disable=attribute-error
  new_train_state = train_state.replace(
      g_params=g_restored_params,
      d_params=d_restored_params,
      g_model_state=g_restored_model_state,
      d_model_state=d_restored_model_state,
      ema_params=ema_restored_params)

  # pytype: enable=attribute-error
  return new_train_state
Initialize the train state with a pretrained checkpoint.

  First restores the checkpoint, then inflates the parameters if requested,
  finally build a new train state with updated parameters.

  Args:
    checkpoint_path: Directory for saving the checkpoint.
    train_state: An instance of TrainState that holds the state of training.
    config: Model config.

  Returns:
    Training state.
  
videogvt/train_lib/vqgan_train_lib.py