docstr: Loads ModelPT from checkpoint, with some maintenance of restoration.
  For documentation, please refer to LightningModule.load_from_checkpoin() documentation.
docstr_tok: ["Loads","ModelPT","from","checkpoint",",","with","some","maintenance","of","restoration",".","For","documentation",",","please","refer","to","LightningModule.load_from_checkpoin","(",")","documentation","."]
code: def load_from_checkpoint(
      cls,
      checkpoint_path: str,
      map_location: Any = None,
      hparams_file: Optional[str] = None,
      strict: bool = True,
      **kwargs,
  ):
      """
      Loads ModelPT from checkpoint, with some maintenance of restoration.
      For documentation, please refer to LightningModule.load_from_checkpoin() documentation.
      """
      checkpoint = None
      try:
          cls._set_model_restore_state(is_being_restored=True)
          # TODO: replace with proper PTL API
          with pl_legacy_patch():
              if map_location is not None:
                  checkpoint = pl_load(
                      checkpoint_path, map_location=map_location
                  )
              else:
                  checkpoint = pl_load(
                      checkpoint_path,
                      map_location=lambda storage, loc: storage,
                  )
  
          if hparams_file is not None:
              extension = hparams_file.split(".")[-1]
              if extension.lower() == "csv":
                  hparams = load_hparams_from_tags_csv(hparams_file)
              elif extension.lower() in ("yml", "yaml"):
                  hparams = load_hparams_from_yaml(hparams_file)
              else:
                  raise ValueError(
                      ".csv, .yml or .yaml is required for `hparams_file`"
                  )
  
              hparams["on_gpu"] = False
  
              # overwrite hparams by the given file
              checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY] = hparams
  
          # for past checkpoint need to add the new key
          if cls.CHECKPOINT_HYPER_PARAMS_KEY not in checkpoint:
              checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY] = {}
          # override the hparams with values that were passed in
          cfg = checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY].get(
              "cfg", checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY]
          )
          # TODO: can we do this without overriding?
          config_kwargs = kwargs.copy()
          if "trainer" in config_kwargs:
              config_kwargs.pop("trainer")
          cfg.update(config_kwargs)
  
          # Disable individual unet/vae weights loading otherwise the model will look for these partial ckpts and raise error
          if cfg:
              if cfg.get("unet_config") and cfg.get("unet_config").get(
                  "from_pretrained"
              ):
                  cfg.unet_config.from_pretrained = None
              if cfg.get("first_stage_config") and cfg.get(
                  "first_stage_config"
              ).get("from_pretrained"):
                  cfg.first_stage_config.from_pretrained = None
              ## Now when we covert ckpt to nemo, let's always get rid of those _orig_mod
              if cfg.get("inductor"):
                  cfg.inductor = False
              ## Append some dummy configs that DB didn't support
              if not cfg.get("channels_last"):
                  cfg.channels_last = True
              if not cfg.get("capture_cudagraph_iters"):
                  cfg.capture_cudagraph_iters = -1
  
          # compatibility for stable diffusion old checkpoint tweaks
          first_key = list(checkpoint["state_dict"].keys())[0]
          if first_key == "betas":
              # insert "model." into for megatron wrapper
              new_state_dict = {}
              for key in checkpoint["state_dict"].keys():
                  new_key = "model." + key
                  new_state_dict[new_key] = checkpoint["state_dict"][
                      key
                  ]
              checkpoint["state_dict"] = new_state_dict
          elif (
              first_key
              == "model.text_encoder.transformer.text_model.embeddings.position_ids"
              or first_key
              == "model.text_encoder.model.language_model.embedding.position_embeddings"
          ):
              # remap state keys from dreambooth when using HF clip
              new_state_dict = {}
              for key in checkpoint["state_dict"].keys():
                  new_key = key.replace("._orig_mod", "")
                  new_key = new_key.replace(
                      "unet", "model.diffusion_model"
                  )
                  new_key = new_key.replace("vae", "first_stage_model")
                  new_key = new_key.replace(
                      "text_encoder", "cond_stage_model"
                  )
                  new_key = new_key.replace(".noise_scheduler", "")
                  new_state_dict[new_key] = checkpoint["state_dict"][
                      key
                  ]
              checkpoint["state_dict"] = new_state_dict
  
          # compatibility for inductor in inference
          if not cfg.get("inductor", False):
              new_state_dict = {}
              for key in checkpoint["state_dict"].keys():
                  new_key = key.replace("._orig_mod", "", 1)
                  new_state_dict[new_key] = checkpoint["state_dict"][
                      key
                  ]
              checkpoint["state_dict"] = new_state_dict
  
          if cfg.get("megatron_amp_O2", False):
              new_state_dict = {}
              for key in checkpoint["state_dict"].keys():
                  new_key = key.replace("model.", "model.module.", 1)
                  new_state_dict[new_key] = checkpoint["state_dict"][
                      key
                  ]
              checkpoint["state_dict"] = new_state_dict
  
          if "cfg" in kwargs:
              model = ptl_load_state(
                  cls, checkpoint, strict=strict, **kwargs
              )
          else:
              model = ptl_load_state(
                  cls, checkpoint, strict=strict, cfg=cfg, **kwargs
              )
              # cfg = checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY].cfg
  
          checkpoint = model
  
      finally:
          cls._set_model_restore_state(is_being_restored=False)
      return checkpoint
  
imports: from functools import partial
  from typing import Any, Optional
  import torch
  from omegaconf import DictConfig
  from pytorch_lightning import Trainer
  from torch._inductor import config
  from nemo.collections.multimodal.data.dreambooth.dreambooth_dataset import DreamBoothDataset
  from nemo.collections.multimodal.modules.stable_diffusion.distributions.distributions import DiagonalGaussianDistribution
  from nemo.collections.multimodal.parts.utils import randn_like
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.core.classes.common import Serialization
  from nemo.utils import logging
label_desc: Use image node_id 21 for calling the MegatronDreamBooth obj's underlying member method code with example usage: obj.load_from_checkpoint(cls, checkpoint_path, map_location, hparams_file, strict) and returns: checkpoint