class_docstr: Standard solver for AudioCraft.
  
  The standard solver implements a base training loop with the following stages:
  train, valid, evaluate and generate that are expected to be all defined for
  solvers in AudioCraft. It also provides a nice default management of Dora history replay,
  checkpoint management across epoch, and logging configuration.
  
  AudioCraft solvers must inherit from the StandardSolver and define the methods
  associated to each stage as well as the show, build_model and build_dataloaders methods.
class_docstr_tok: ["Standard","solver","for","AudioCraft",".","The","standard","solver","implements","a","base","training","loop","with","the","following","stages",":","train",",","valid",",","evaluate","and","generate","that","are","expected","to","be","all","defined","for","solvers","in","AudioCraft",".","It","also","provides","a","nice","default","management","of","Dora","history","replay",",","checkpoint","management","across","epoch",",","and","logging","configuration",".","AudioCraft","solvers","must","inherit","from","the","StandardSolver","and","define","the","methods","associated","to","each","stage","as","well","as","the","show",",","build_model","and","build_dataloaders","methods","."]
docstr: Load last checkpoint or the one specified in continue_from.
  
  Args:
      load_best (bool): Whether to load from best state dict or not.
          Best state dict is always used when not loading the current xp.
      ignore_state_keys (list of str): List of sources to ignore when loading the state, e.g. `optimizer`.
  Returns:
      state (dict, optional): The loaded state dictionary.
docstr_tok: ["Load","last","checkpoint","or","the","one","specified","in","continue_from",".","Args",":","load_best","(","bool",")",":","Whether","to","load","from","best","state","dict","or","not",".","Best","state","dict","is","always","used","when","not","loading","the","current","xp",".","ignore_state_keys","(","list","of","str",")",":","List","of","sources","to","ignore","when","loading","the","state",",","e.g",".","`","optimizer","`",".","Returns",":","state","(","dict",",","optional",")",":","The","loaded","state","dictionary","."]
code: def load_checkpoints(
      self,
      load_best: bool = False,
      ignore_state_keys: tp.List[str] = [],
  ) -> tp.Optional[dict]:
      """Load last checkpoint or the one specified in continue_from.
  
      Args:
          load_best (bool): Whether to load from best state dict or not.
              Best state dict is always used when not loading the current xp.
          ignore_state_keys (list of str): List of sources to ignore when loading the state, e.g. `optimizer`.
      Returns:
          state (dict, optional): The loaded state dictionary.
      """
      # load checkpoints from xp folder or cfg.continue_from
      is_sharded = self.cfg.fsdp.use
      load_from_path: tp.Optional[Path] = None
      checkpoint_source: tp.Optional[checkpoint.CheckpointSource] = None
  
      if load_best:
          self.logger.info("Trying to load state_dict from best state.")
  
      state: tp.Optional[dict] = None
      rank0_checkpoint_path = self.checkpoint_path(use_fsdp=False)
      current_checkpoint_path = self.checkpoint_path()
      _pretrained_prefix = "//pretrained/"
      continue_pretrained = (self.cfg.continue_from or "").startswith(
          _pretrained_prefix
      )
      if rank0_checkpoint_path.exists():
          self.logger.info(
              f"Loading existing checkpoint: {current_checkpoint_path}"
          )
          load_from_path = current_checkpoint_path
          checkpoint.check_sharded_checkpoint(
              current_checkpoint_path, rank0_checkpoint_path
          )
          checkpoint_source = checkpoint.CheckpointSource.CURRENT_XP
      elif self.cfg.continue_from and not continue_pretrained:
          self.logger.info(
              f"Continuing from provided checkpoint: {self.cfg.continue_from}"
          )
          # we're always continuing from consolidated checkpoints: self.cfg.use_fsdp and not continue_best
          load_from_path = checkpoint.resolve_checkpoint_path(
              self.cfg.continue_from, use_fsdp=False
          )
          if load_from_path is None:
              self.logger.error(
                  "Could not resolve the continue_from checkpoint %s",
                  self.cfg.continue_from,
              )
              raise RuntimeError(
                  f"Could not resolve continue_from checkpoint {self.cfg.continue_from}"
              )
          checkpoint_source = checkpoint.CheckpointSource.OTHER
  
      if load_from_path is not None:
          state = checkpoint.load_checkpoint(load_from_path, is_sharded)
      elif continue_pretrained:
          self.logger.info(
              "Loading a pretrained model. Ignoring 'load_best' and 'ignore_state_keys' params."
          )
          state = self.load_from_pretrained(
              self.cfg.continue_from[len(_pretrained_prefix) :]
          )
          checkpoint_source = checkpoint.CheckpointSource.PRETRAINED
          load_best = True
  
      # checkpoints are not from the current xp, we only retrieve the best state
      if (
          checkpoint_source is not None
          and checkpoint_source
          != checkpoint.CheckpointSource.CURRENT_XP
      ):
          assert state is not None
          self.logger.info(
              "Checkpoint source is not the current xp: Load state_dict from best state."
          )
          load_best = True
          state = {
              key: state[key]
              for key in self._continue_best_source_keys
              if key in state
          }
          # loaded checkpoints are FSDP checkpoints: we're reading the best state
          # from FSDP and we drop the regular best_state
          if "fsdp_best_state" in state and state["fsdp_best_state"]:
              state.pop("best_state", None)
              self.logger.info(
                  "... Loaded checkpoint has FSDP best state"
              )
          # FSDP is enabled in the solver, if the loaded checkpoints do not have FSDP support
          # then we're initializing FSDP best state with the regular best state
          elif self.cfg.fsdp.use:
              if (
                  "fsdp_best_state" not in state
                  or not state["fsdp_best_state"]
              ):
                  # we swap non-FSDP checkpoints best_state to FSDP-compatible best state
                  state["fsdp_best_state"] = state.pop("best_state")
                  self.logger.info(
                      "... Loaded checkpoint does not have FSDP best state. Use regular best state"
                  )
  
      if state is not None:
          if load_best:
              self.logger.info(
                  "Ignoring keys when loading best %r",
                  ignore_state_keys,
              )
              for key in set(ignore_state_keys):
                  if key in state:
                      state.pop(key)
              has_best_state = (
                  "best_state" in state or "fsdp_best_state" in state
              )
              assert has_best_state, (
                  "Trying to load best state but neither 'best_state'",
                  " or 'fsdp_best_state' found in checkpoints.",
              )
          self.load_state_dict(state)
  
      # for FSDP, let's make extra sure nothing bad happened with out of sync
      # checkpoints across workers.
      epoch = float(self.epoch)
      avg_epoch = flashy.distrib.average_metrics({"epoch": epoch})[
          "epoch"
      ]
      if avg_epoch != epoch:
          raise RuntimeError(
              f"Inconsistent loading of checkpoints happened, our epoch is {epoch} "
              f"but average of epochs is {avg_epoch}, at least one gpu must have a "
              "different epoch number."
          )
  
      # on load_best, properly reinitialize state_dict, best states and ema
      # otherwise we load from the current xp and don't alter anything
      if load_best:
          self.logger.info("Loading state_dict from best state.")
          if not self.cfg.fsdp.use and self.fsdp_best_state:
              # loading from an FSDP checkpoint but with FSDP deactivated
              self.logger.info("... Loading from FSDP best state dict.")
              self.best_state.load_state_dict(self.fsdp_best_state)
  
          # if load_best, we permanently override the regular state_dict with the best state
          if self.cfg.fsdp.use:
              self.logger.info(
                  "FSDP is used, loading from FSDP best state."
              )
              with fsdp.switch_to_full_state_dict(self._fsdp_modules):
                  # this might be really fragile but okay for now.
                  self.load_state_dict(self.fsdp_best_state)
          else:
              # we permanently swap the stateful objects to their best state
              self._load_new_state_dict(self.best_state.state_dict())
  
          # the EMA modules should also be instantiated with best state.
          # the easiest way to do so is to reinitialize a new EMA with best state loaded.
          if self.ema is not None:
              self.logger.info("Re-initializing EMA from best state")
              self.initialize_ema()
  
          if self.cfg.fsdp.use:
              self.logger.info(
                  "Re-initializing best state after using FSDP best state."
              )
              for name in self.best_state.states.keys():
                  state_source = self._get_state_source(name)
                  self.best_state.update(name, state_source)
  
      return state
  
imports: from abc import ABC, abstractmethod
  from contextlib import contextmanager
  from pathlib import Path
  import typing
  import flashy
  import omegaconf
  import torch
  from torch import nn
  from ..None import optim
  from ..optim import fsdp
  from ..utils import checkpoint
  from ..utils.autocast import TorchAutocast
  from ..utils.best_state import BestStateDictManager
  from ..utils.deadlock import DeadlockDetect
  from ..utils.profiler import Profiler
  from ..utils.utils import copy_state, dict_from_config, model_hash, with_rank_rng
label_desc: Use image node_id 24 for calling the StandardSolver obj's underlying member method code with example usage: obj.load_checkpoints(load_best, ignore_state_keys) and returns: state