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.baselines.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