code: def load_models(config):
      """Load pre-trained Stable Diffusion models."""
      # Load all models in bfloat16 to save GRAM.
      # For models that are only used for inferencing,
      # full precision is also not required.
      dtype = torch.bfloat16
  
      text_encoder = CLIPTextModel.from_pretrained(
          args.model_dir,
          subfolder="text_encoder",
          torch_dtype=dtype,
      )
  
      noise_scheduler = DDPMScheduler.from_pretrained(
          config["model_dir"],
          subfolder="scheduler",
          torch_dtype=dtype,
      )
  
      # VAE is only used for inference, keeping weights in full precision is not required.
      vae = AutoencoderKL.from_pretrained(
          config["model_dir"],
          subfolder="vae",
          torch_dtype=dtype,
      )
      # We are not training VAE part of the model.
      vae.requires_grad_(False)
  
      # Convert unet to bf16 to save GRAM.
      unet = UNet2DConditionModel.from_pretrained(
          config["model_dir"],
          subfolder="unet",
          torch_dtype=dtype,
      )
  
      if is_xformers_available():
          unet.enable_xformers_memory_efficient_attention()
  
      if not config["use_lora"]:
          unet_trainable_parameters = unet.parameters()
          text_trainable_parameters = text_encoder.parameters()
      else:
          text_encoder.requires_grad_(False)
          unet.requires_grad_(False)
          (
              unet_trainable_parameters,
              text_trainable_parameters,
          ) = add_lora_layers(unet, text_encoder)
  
      text_encoder.train()
      unet.train()
  
      torch.cuda.empty_cache()
  
      return (
          text_encoder,
          noise_scheduler,
          vae,
          unet,
          unet_trainable_parameters,
          text_trainable_parameters,
      )
  
imports: from typing import Dict
  import itertools
  from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, UNet2DConditionModel
  from diffusers.loaders import LoraLoaderMixin, text_encoder_lora_state_dict
  from diffusers.models.attention_processor import AttnAddedKVProcessor, AttnAddedKVProcessor2_0, LoRAAttnAddedKVProcessor, LoRAAttnProcessor, LoRAAttnProcessor2_0, SlicedAttnAddedKVProcessor
  from diffusers.utils.import_utils import is_xformers_available
  from ray.train import ScalingConfig
  from ray import train
  from ray.train.torch import TorchTrainer
  import torch
  import torch.nn.functional
  from torch.nn.utils import clip_grad_norm_
  from transformers import CLIPTextModel
  from dataset import collate, get_train_dataset
  from flags import train_arguments
label_desc: Use image node_id 4 for calling a global function with example usage: load_models(config) and returns: text_encoder,  noise_scheduler,  vae,  unet,  unet_trainable_parameters,  text_trainable_parameters