code: def _load_model(
      model_ckpt: str,
      model_cfg: str,
      eval_mode: bool = True,
      trainer: Trainer = None,
  ):
      assert model_ckpt is not None, "model ckpt cannot be None"
      if model_ckpt.endswith(".nemo"):
          model_cfg = MegatronImagen.restore_from(
              restore_path=model_ckpt,
              trainer=trainer,
              return_config=True,
          )
          model_cfg.unet.flash_attention = False
          model_cfg.micro_batch_size = 1
          model_cfg.global_batch_size = 1
          model = MegatronImagen.restore_from(
              restore_path=model_ckpt,
              override_config_path=model_cfg,
              trainer=trainer,
          )
      elif model_ckpt.endswith(".ckpt"):
          model_cfg = OmegaConf.load(model_cfg)
          model_cfg.model.unet.flash_attention = False
          model_cfg.model.micro_batch_size = 1
          model_cfg.model.global_batch_size = 1
          model = MegatronImagen(cfg=model_cfg.model, trainer=trainer)
          checkpoint = torch.load(
              model_ckpt, map_location=lambda storage, loc: storage
          )
  
          # Change weight keys if training using TorchInductor
          state_dict = checkpoint["state_dict"]
          del_keys = []
          for k, v in state_dict.items():
              if "._orig_mod" in k:
                  del_keys.append(k)
          if len(del_keys) != 0:
              print(
                  "ckpt was saved with TorchInductor. Renaming weights.."
              )
          for k in del_keys:
              new_k = k.replace("._orig_mod", "")
              state_dict[new_k] = state_dict[k]
              del state_dict[k]
          model.load_state_dict(state_dict, strict=True)
      else:
          raise Exception(
              "Invalid ckpt type. Should be either .nemo or .ckpt with cfg"
          )
  
      model = (
          model.model
      )  # We do not need Megatron Instance for inference
      model.model.set_inference_mode(
          True
      )  # Used for adding the least noise for EDM inference for SR model.
      if eval_mode:
          model.unet.cuda().eval()
      return model
  
imports: import os
  import time
  from dataclasses import dataclass, field
  from typing import Callable, List, Optional, Union
  import torch
  from omegaconf.omegaconf import OmegaConf
  from pytorch_lightning import Trainer
  from torch.cuda.amp import autocast
  from nemo.collections.multimodal.models.text_to_image.imagen.imagen import Imagen, MegatronImagen
  from nemo.collections.multimodal.parts.utils import numpy_to_pil, setup_trainer_and_models_for_inference
label_desc: Use image node_id 2 for calling the ImagenPipeline obj's underlying member method code with example usage: obj._load_model(model_ckpt, model_cfg, eval_mode, trainer) and returns: model