code: def init_from_mtv_checkpoint(
      config: ml_collections.ConfigDict,
      model: model_lib.MTVClassificationModel,
      train_state: train_utils.TrainState,
  ) -> train_utils.TrainState:
      """Initialize train state from a MTV checkpoint."""
      if config.init_from.get("model_cfg") is None:
          logging.info("model_cfg is empty. Using current model_cfg.")
          restored_model_cfg = copy.deepcopy(config)
      else:
          restored_model_cfg = config.init_from.model_cfg
      init_checkpoint_path = config.init_from.get("checkpoint_path")
      restored_train_state = (
          pretrain_utils.restore_pretrained_checkpoint(
              init_checkpoint_path, train_state, assert_exist=True
          )
      )
      return model.init_from_train_state(
          train_state,
          restored_train_state,
          restored_model_cfg,
          restore_output_proj=config.init_from.get(
              "restore_output_projection", False
          ),
      )
  
imports: import copy
  import functools
  import os
  from typing import Any, Dict, Tuple
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from clu import platform
  from flax import jax_utils
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.mtv import model
  from scenic.projects.mtv import train_utils
  from scenic.projects.vivit import evaluation_lib
  from scenic.projects.vivit import train_utils
  from scenic.train_lib import lr_schedules
  from scenic.train_lib import optimizers
  from scenic.train_lib import pretrain_utils
  from scenic.train_lib import train_utils
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: init_from_mtv_checkpoint(config, model, train_state) and returns: model