code: def get_model_cls(model_name):
      """ "Selects MTV model type."""
      if model_name == "mtv_multiclass_classification":
          return MTVClassificationModel
      elif model_name == "mtv_multihead_classification":
          return MTVMultiheadClassificationModel
      else:
          raise ValueError("Unrecognized model: {}".format(model_name))
  
imports: import functools
  from typing import Any, List, Sequence, Tuple, Union
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.mtv import model_utils
  from scenic.projects.vivit import model
  from scenic.train_lib import train_utils
label_desc: Use image node_id 1 for calling a global function with example usage: get_model_cls(model_name) and returns: MTVClassificationModel, MTVMultiheadClassificationModel