code: def add_lora_layers(unet, text_encoder):
      """Add LoRA layers for unet and text encoder.
  
      `unet` and `text_encoder` will be modified in place.
  
      Returns:
          The LoRA parameters for unet and text encoder correspondingly.
      """
      unet_lora_attn_procs = {}
      unet_lora_parameters = []
      for name, attn_processor in unet.attn_processors.items():
          cross_attention_dim = (
              None
              if name.endswith("attn1.processor")
              else unet.config.cross_attention_dim
          )
          if name.startswith("mid_block"):
              hidden_size = unet.config.block_out_channels[-1]
          elif name.startswith("up_blocks"):
              block_id = int(name[len("up_blocks.")])
              hidden_size = list(
                  reversed(unet.config.block_out_channels)
              )[block_id]
          elif name.startswith("down_blocks"):
              block_id = int(name[len("down_blocks.")])
              hidden_size = unet.config.block_out_channels[block_id]
  
          if isinstance(
              attn_processor,
              (
                  AttnAddedKVProcessor,
                  SlicedAttnAddedKVProcessor,
                  AttnAddedKVProcessor2_0,
              ),
          ):
              lora_attn_processor_class = LoRAAttnAddedKVProcessor
          else:
              lora_attn_processor_class = (
                  LoRAAttnProcessor2_0
                  if hasattr(F, "scaled_dot_product_attention")
                  else LoRAAttnProcessor
              )
  
          module = lora_attn_processor_class(
              hidden_size=hidden_size,
              cross_attention_dim=cross_attention_dim,
              rank=LORA_RANK,
          )
          unet_lora_attn_procs[name] = module
          unet_lora_parameters.extend(module.parameters())
  
      unet.set_attn_processor(unet_lora_attn_procs)
  
      text_lora_parameters = LoraLoaderMixin._modify_text_encoder(
          text_encoder, dtype=torch.float32, rank=LORA_RANK
      )
  
      return unet_lora_parameters, text_lora_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 3 for calling a global function with example usage: add_lora_layers(unet, text_encoder) and returns: unet_lora_parameters,  text_lora_parameters