code: def get_max_memory(
      max_memory: Optional[
          Dict[Union[int, str], Union[int, str]]
      ] = None
  ):
      """
      Get the maximum memory available if nothing is passed, converts string to int otherwise.
      """
      import psutil
  
      if max_memory is None:
          if not (
              torch.cuda.is_available()
              or is_npu_available()
              or is_xpu_available()
          ):
              max_memory = {}
  
          else:
              # Make sure CUDA is initialized on each GPU to have the right memory info.
              if is_npu_available():
                  for i in range(torch.npu.device_count()):
                      _ = torch.tensor(0, device=torch.device("npu", i))
                  max_memory = {
                      i: torch.npu.mem_get_info(i)[0]
                      for i in range(torch.npu.device_count())
                  }
              elif is_xpu_available():
                  for i in range(torch.xpu.device_count()):
                      _ = torch.tensor(0, device=torch.device("xpu", i))
                  max_memory = {
                      i: torch.xpu.max_memory_allocated(i)
                      for i in range(torch.xpu.device_count())
                  }
              else:
                  for i in range(torch.cuda.device_count()):
                      _ = torch.tensor([0], device=i)
                  max_memory = {
                      i: torch.cuda.mem_get_info(i)[0]
                      for i in range(torch.cuda.device_count())
                  }
          # allocate everything in the mps device as the RAM is shared
          if is_mps_available():
              max_memory["mps"] = psutil.virtual_memory().available
          else:
              max_memory["cpu"] = psutil.virtual_memory().available
          return max_memory
  
      for key in max_memory:
          if isinstance(max_memory[key], str):
              max_memory[key] = convert_file_size_to_int(
                  max_memory[key]
              )
  
      # Need to sort the device by type to make sure that we allocate the gpu first.
      # As gpu/npu/xpu are represented by int, we need to sort them first.
      gpu_devices = [k for k in max_memory.keys() if isinstance(k, int)]
      gpu_devices.sort()
      # check if gpu/npu/xpu devices are available and if not, throw a warning
      if is_npu_available():
          num_devices = torch.npu.device_count()
      elif is_xpu_available():
          num_devices = torch.xpu.device_count()
      else:
          num_devices = torch.cuda.device_count()
      for device in gpu_devices:
          if device >= num_devices or device < 0:
              logger.warning(
                  f"Device {device} is not available, available devices are {list(range(num_devices))}"
              )
      # Add the other devices in the preset order if they are available
      all_devices = gpu_devices + [
          k for k in ["mps", "cpu", "disk"] if k in max_memory.keys()
      ]
      # Raise an error if a device is not recognized
      for k in max_memory.keys():
          if k not in all_devices:
              raise ValueError(
                  f"Device {k} is not recognized, available devices are integers(for GPU/XPU), 'mps', 'cpu' and 'disk'"
              )
      max_memory = {k: max_memory[k] for k in all_devices}
  
      return max_memory
  
imports: import contextlib
  import gc
  import inspect
  import json
  import logging
  import os
  import re
  import shutil
  import tempfile
  from collections import OrderedDict, defaultdict
  from typing import Dict, List, Optional, Tuple, Union
  import torch
  import torch.nn
  from ..state import AcceleratorState
  from .constants import SAFE_WEIGHTS_NAME, WEIGHTS_NAME
  from .dataclasses import AutocastKwargs, CustomDtype, DistributedType
  from .imports import is_mps_available, is_npu_available, is_xpu_available
  from .offload import load_offloaded_weight, offload_weight, save_offload_index
  from .tqdm import is_tqdm_available, tqdm
  from safetensors import safe_open
  from safetensors.torch import load_file
label_desc: Use image node_id 17 for calling a global function with example usage: get_max_memory(max_memory) and returns: max_memory, max_memory