code: def load_model_vars(
      model_name: str,
      checkpoint_path: Optional[str] = None,
      download_dir: str = download.DEFAULT_DOWNLOAD_DIR,
  ) -> PyTree:
      """Load model variables from a checkpoint, downloading if necessary."""
      checkpoint_path = checkpoint_path or CHECKPOINTS.get(model_name)
      if checkpoint_path is None:
          checkpoint_path = os.path.join(
              download_dir, model_name + ".npy"
          )
  
          if not gfile.exists(checkpoint_path):
              # Download PyTorch checkpoint
              url = CHECKPOINTS_TORCH.get(model_name)
              logging.info(
                  "Downloading checkpoint from %s to %s",
                  url,
                  download_dir,
              )
              checkpoint_path_torch = download.download(
                  url, download_dir, expected_sha256=url.split("/")[-2]
              )
  
              # Load and convert checkpoint to numpy
              logging.info(
                  "Converting checkpoint %s to numpy",
                  checkpoint_path_torch,
              )
              try:
                  import torch
              except ImportError as e:
                  logging.error(
                      "Could not import torch for CLIP checkpoint conversion"
                  )
              params = torch.jit.load(
                  checkpoint_path_torch, map_location="cpu"
              ).state_dict()
              params = jax.tree_util.tree_map(
                  lambda p: p.cpu().numpy(), params
              )
  
              # Save converted checkpoint
              with gfile.GFile(checkpoint_path, "wb") as f:
                  np.save(f, params)
              del params
              gfile.remove(checkpoint_path_torch)
  
      with gfile.GFile(checkpoint_path, "rb") as f:
          np_params = np.load(f, allow_pickle=True).tolist()
      return _convert_vars(np_params)
  
imports: import os
  from typing import Any, Mapping, Optional
  from absl import logging
  import flax
  import jax
  import jax.numpy
  import numpy
  from scenic.projects.owl_vit.clip import layers
  from scenic.projects.baselines.clip import download
  from tensorflow.io import gfile
label_desc: Use image node_id 1 for calling a global function with example usage: load_model_vars(model_name, checkpoint_path, download_dir) and returns: _convert_vars