code: def train_fn(config):
      # Load pre-trained models.
      (
          text_encoder,
          noise_scheduler,
          vae,
          unet,
          unet_trainable_parameters,
          text_trainable_parameters,
      ) = load_models(config)
  
      text_encoder = train.torch.prepare_model(text_encoder)
      unet = train.torch.prepare_model(unet)
      # manually move to device as `prepare_model` can't be used on
      # non-training models.
      vae = vae.to(train.torch.get_device())
  
      # Use the regular AdamW optimizer to work with bfloat16 weights.
      optimizer = torch.optim.AdamW(
          itertools.chain(
              unet_trainable_parameters, text_trainable_parameters
          ),
          lr=config["lr"],
      )
  
      train_dataset = train.get_dataset_shard("train")
  
      # Train!
      num_train_epochs = config["num_epochs"]
  
      print(f"Running {num_train_epochs} epochs.")
  
      global_step = 0
      for _ in range(num_train_epochs):
          if global_step >= config["max_train_steps"]:
              print(
                  f"Stopping training after reaching {global_step} steps..."
              )
              break
  
          for _, batch in enumerate(
              train_dataset.iter_torch_batches(
                  batch_size=config["train_batch_size"],
                  device=train.torch.get_device(),
              )
          ):
              batch = collate(batch, torch.bfloat16)
  
              optimizer.zero_grad()
  
              # Convert images to latent space
              latents = (
                  vae.encode(batch["images"]).latent_dist.sample()
                  * 0.18215
              )
  
              # Sample noise that we'll add to the latents
              noise = torch.randn_like(latents)
              bsz = latents.shape[0]
              # Sample a random timestep for each image
              timesteps = torch.randint(
                  0,
                  noise_scheduler.config.num_train_timesteps,
                  (bsz,),
                  device=latents.device,
              )
              timesteps = timesteps.long()
  
              # Add noise to the latents according to the noise magnitude at each timestep
              # (this is the forward diffusion process)
              noisy_latents = noise_scheduler.add_noise(
                  latents, noise, timesteps
              )
  
              # Get the text embedding for conditioning
              encoder_hidden_states = text_encoder(batch["prompt_ids"])[
                  0
              ]
  
              # Predict the noise residual.
              model_pred = unet(
                  noisy_latents.to(train.torch.get_device()),
                  timesteps.to(train.torch.get_device()),
                  encoder_hidden_states.to(train.torch.get_device()),
              ).sample
              target = get_target(
                  noise_scheduler, noise, latents, timesteps
              )
  
              loss = prior_preserving_loss(
                  model_pred, target, config["prior_loss_weight"]
              )
              loss.backward()
  
              # Gradient clipping before optimizer stepping.
              clip_grad_norm_(
                  itertools.chain(
                      unet_trainable_parameters,
                      text_trainable_parameters,
                  ),
                  config["max_grad_norm"],
              )
  
              optimizer.step()  # Step all optimizers.
  
              global_step += 1
              results = {
                  "step": global_step,
                  "loss": loss.detach().item(),
              }
              train.report(results)
  
              if global_step >= config["max_train_steps"]:
                  break
      # END: Training loop
  
      # Create pipeline using the trained modules and save it.
      if train.get_context().get_world_rank() == 0:
          if not config["use_lora"]:
              pipeline = DiffusionPipeline.from_pretrained(
                  config["model_dir"],
                  text_encoder=text_encoder.module,
                  unet=unet.module,
              )
              pipeline.save_pretrained(config["output_dir"])
          else:
              save_lora_weights(
                  unet.module, text_encoder.module, config["output_dir"]
              )
  
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 5 for calling a global function with example usage: train_fn(config) without return types