code: def _load_customized_model(
      cfg: ImagenPipelineConfig,
      trainer=None,
      megatron_loading=False,
      megatron_cfg=None,
  ):
      if megatron_loading:
          assert megatron_cfg
  
          def model_cfg_modifier(model_cfg):
              model_cfg.inductor = False
              model_cfg.unet.flash_attention = False
              model_cfg.micro_batch_size = (
                  megatron_cfg.fid.ncaptions_per_batch
              )
              model_cfg.global_batch_size = (
                  model_cfg.micro_batch_size
                  * megatron_cfg.fid.ntasks_per_node
              )
  
          (
              trainer,
              megatron_models,
          ) = setup_trainer_and_models_for_inference(
              MegatronImagen,
              cfg=megatron_cfg,
              model_cfg_modifier=model_cfg_modifier,
          )
          models = [mm.model for mm in megatron_models]
          for model in models:
              model.cuda().eval()
              model.model.set_inference_mode(True)
          return models
      customized_models = cfg.customized_model
      models = []
      print("Load base model.")
      model = ImagenPipeline._load_model(
          model_ckpt=customized_models.base_ckpt,
          model_cfg=customized_models.base_cfg,
          trainer=trainer,
      )
      models.append(model)
  
      if cfg.target_resolution >= 256:
          print("Load SR256 model.")
          model = ImagenPipeline._load_model(
              model_ckpt=customized_models.sr256_ckpt,
              model_cfg=customized_models.sr256_cfg,
              trainer=trainer,
          )
          models.append(model)
  
      if cfg.target_resolution >= 1024:
          print("Load SR1024 model.")
          model = ImagenPipeline._load_model(
              model_ckpt=customized_models.sr1024_ckpt,
              model_cfg=customized_models.sr1024_cfg,
              trainer=trainer,
          )
          models.append(model)
      return models
  
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 3 for calling the ImagenPipeline obj's underlying member method code with example usage: obj._load_customized_model(cfg, trainer, megatron_loading, megatron_cfg) and returns: models, models