code: def main(args):
      logging_dir = os.path.join(args.output_dir, args.logging_dir)
      accelerator_project_config = ProjectConfiguration(
          total_limit=args.checkpoints_total_limit,
          project_dir=args.output_dir,
          logging_dir=logging_dir,
      )
  
      accelerator = Accelerator(
          gradient_accumulation_steps=args.gradient_accumulation_steps,
          mixed_precision=args.mixed_precision,
          log_with=args.report_to,
          project_config=accelerator_project_config,
      )
  
      if args.logger == "tensorboard":
          if not is_tensorboard_available():
              raise ImportError(
                  "Make sure to install tensorboard if you want to use it for logging during training."
              )
  
      elif args.logger == "wandb":
          if not is_wandb_available():
              raise ImportError(
                  "Make sure to install wandb if you want to use it for logging during training."
              )
          import wandb
  
      # `accelerate` 0.16.0 will have better support for customized saving
      if version.parse(accelerate.__version__) >= version.parse(
          "0.16.0"
      ):
          # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
          def save_model_hook(models, weights, output_dir):
              if accelerator.is_main_process:
                  if args.use_ema:
                      ema_model.save_pretrained(
                          os.path.join(output_dir, "unet_ema")
                      )
  
                  for i, model in enumerate(models):
                      model.save_pretrained(
                          os.path.join(output_dir, "unet")
                      )
  
                      # make sure to pop weight so that corresponding model is not saved again
                      weights.pop()
  
          def load_model_hook(models, input_dir):
              if args.use_ema:
                  load_model = EMAModel.from_pretrained(
                      os.path.join(input_dir, "unet_ema"), UNet2DModel
                  )
                  ema_model.load_state_dict(load_model.state_dict())
                  ema_model.to(accelerator.device)
                  del load_model
  
              for i in range(len(models)):
                  # pop models so that they are not loaded again
                  model = models.pop()
  
                  # load diffusers style into model
                  load_model = UNet2DModel.from_pretrained(
                      input_dir, subfolder="unet"
                  )
                  model.register_to_config(**load_model.config)
  
                  model.load_state_dict(load_model.state_dict())
                  del load_model
  
          accelerator.register_save_state_pre_hook(save_model_hook)
          accelerator.register_load_state_pre_hook(load_model_hook)
  
      # Make one log on every process with the configuration for debugging.
      logging.basicConfig(
          format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
          datefmt="%m/%d/%Y %H:%M:%S",
          level=logging.INFO,
      )
      logger.info(accelerator.state, main_process_only=False)
      if accelerator.is_local_main_process:
          datasets.utils.logging.set_verbosity_warning()
          diffusers.utils.logging.set_verbosity_info()
      else:
          datasets.utils.logging.set_verbosity_error()
          diffusers.utils.logging.set_verbosity_error()
  
      # Handle the repository creation
      if accelerator.is_main_process:
          if args.push_to_hub:
              if args.hub_model_id is None:
                  repo_name = get_full_repo_name(
                      Path(args.output_dir).name, token=args.hub_token
                  )
              else:
                  repo_name = args.hub_model_id
              create_repo(
                  repo_name, exist_ok=True, token=args.hub_token
              )
              repo = Repository(
                  args.output_dir,
                  clone_from=repo_name,
                  token=args.hub_token,
              )
  
              with open(
                  os.path.join(args.output_dir, ".gitignore"), "w+"
              ) as gitignore:
                  if "step_*" not in gitignore:
                      gitignore.write("step_*\n")
                  if "epoch_*" not in gitignore:
                      gitignore.write("epoch_*\n")
          elif args.output_dir is not None:
              os.makedirs(args.output_dir, exist_ok=True)
  
      # Initialize the model
      if args.model_config_name_or_path is None:
          model = UNet2DModel(
              sample_size=args.resolution,
              in_channels=3,
              out_channels=3,
              layers_per_block=2,
              block_out_channels=(128, 128, 256, 256, 512, 512),
              down_block_types=(
                  "DownBlock2D",
                  "DownBlock2D",
                  "DownBlock2D",
                  "DownBlock2D",
                  "AttnDownBlock2D",
                  "DownBlock2D",
              ),
              up_block_types=(
                  "UpBlock2D",
                  "AttnUpBlock2D",
                  "UpBlock2D",
                  "UpBlock2D",
                  "UpBlock2D",
                  "UpBlock2D",
              ),
          )
      else:
          config = UNet2DModel.load_config(
              args.model_config_name_or_path
          )
          model = UNet2DModel.from_config(config)
  
      # Create EMA for the model.
      if args.use_ema:
          ema_model = EMAModel(
              model.parameters(),
              decay=args.ema_max_decay,
              use_ema_warmup=True,
              inv_gamma=args.ema_inv_gamma,
              power=args.ema_power,
              model_cls=UNet2DModel,
              model_config=model.config,
          )
  
      if args.enable_xformers_memory_efficient_attention:
          if is_xformers_available():
              import xformers
  
              xformers_version = version.parse(xformers.__version__)
              if xformers_version == version.parse("0.0.16"):
                  logger.warn(
                      "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
                  )
              model.enable_xformers_memory_efficient_attention()
          else:
              raise ValueError(
                  "xformers is not available. Make sure it is installed correctly"
              )
  
      # Initialize the scheduler
      accepts_prediction_type = "prediction_type" in set(
          inspect.signature(DDPMScheduler.__init__).parameters.keys()
      )
      if accepts_prediction_type:
          noise_scheduler = DDPMScheduler(
              num_train_timesteps=args.ddpm_num_steps,
              beta_schedule=args.ddpm_beta_schedule,
              prediction_type=args.prediction_type,
          )
      else:
          noise_scheduler = DDPMScheduler(
              num_train_timesteps=args.ddpm_num_steps,
              beta_schedule=args.ddpm_beta_schedule,
          )
  
      # Initialize the optimizer
      optimizer = torch.optim.AdamW(
          model.parameters(),
          lr=args.learning_rate,
          betas=(args.adam_beta1, args.adam_beta2),
          weight_decay=args.adam_weight_decay,
          eps=args.adam_epsilon,
      )
  
      optimizer = ORT_FP16_Optimizer(optimizer)
  
      # Get the datasets: you can either provide your own training and evaluation files (see below)
      # or specify a Dataset from the hub (the dataset will be downloaded automatically from the datasets Hub).
  
      # In distributed training, the load_dataset function guarantees that only one local process can concurrently
      # download the dataset.
      if args.dataset_name is not None:
          dataset = load_dataset(
              args.dataset_name,
              args.dataset_config_name,
              cache_dir=args.cache_dir,
              split="train",
          )
      else:
          dataset = load_dataset(
              "imagefolder",
              data_dir=args.train_data_dir,
              cache_dir=args.cache_dir,
              split="train",
          )
          # See more about loading custom images at
          # https://huggingface.co/docs/datasets/v2.4.0/en/image_load#imagefolder
  
      # Preprocessing the datasets and DataLoaders creation.
      augmentations = transforms.Compose(
          [
              transforms.Resize(
                  args.resolution,
                  interpolation=transforms.InterpolationMode.BILINEAR,
              ),
              transforms.CenterCrop(args.resolution)
              if args.center_crop
              else transforms.RandomCrop(args.resolution),
              transforms.RandomHorizontalFlip()
              if args.random_flip
              else transforms.Lambda(lambda x: x),
              transforms.ToTensor(),
              transforms.Normalize([0.5], [0.5]),
          ]
      )
  
      def transform_images(examples):
          images = [
              augmentations(image.convert("RGB"))
              for image in examples["image"]
          ]
          return {"input": images}
  
      logger.info(f"Dataset size: {len(dataset)}")
  
      dataset.set_transform(transform_images)
      train_dataloader = torch.utils.data.DataLoader(
          dataset,
          batch_size=args.train_batch_size,
          shuffle=True,
          num_workers=args.dataloader_num_workers,
      )
  
      # Initialize the learning rate scheduler
      lr_scheduler = get_scheduler(
          args.lr_scheduler,
          optimizer=optimizer,
          num_warmup_steps=args.lr_warmup_steps
          * args.gradient_accumulation_steps,
          num_training_steps=(len(train_dataloader) * args.num_epochs),
      )
  
      # Prepare everything with our `accelerator`.
      (
          model,
          optimizer,
          train_dataloader,
          lr_scheduler,
      ) = accelerator.prepare(
          model, optimizer, train_dataloader, lr_scheduler
      )
  
      if args.use_ema:
          ema_model.to(accelerator.device)
  
      # We need to initialize the trackers we use, and also store our configuration.
      # The trackers initializes automatically on the main process.
      if accelerator.is_main_process:
          run = os.path.split(__file__)[-1].split(".")[0]
          accelerator.init_trackers(run)
  
      model = ORTModule(model)
  
      total_batch_size = (
          args.train_batch_size
          * accelerator.num_processes
          * args.gradient_accumulation_steps
      )
      num_update_steps_per_epoch = math.ceil(
          len(train_dataloader) / args.gradient_accumulation_steps
      )
      max_train_steps = args.num_epochs * num_update_steps_per_epoch
  
      logger.info("***** Running training *****")
      logger.info(f"  Num examples = {len(dataset)}")
      logger.info(f"  Num Epochs = {args.num_epochs}")
      logger.info(
          f"  Instantaneous batch size per device = {args.train_batch_size}"
      )
      logger.info(
          f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}"
      )
      logger.info(
          f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}"
      )
      logger.info(f"  Total optimization steps = {max_train_steps}")
  
      global_step = 0
      first_epoch = 0
  
      # Potentially load in the weights and states from a previous save
      if args.resume_from_checkpoint:
          if args.resume_from_checkpoint != "latest":
              path = os.path.basename(args.resume_from_checkpoint)
          else:
              # Get the most recent checkpoint
              dirs = os.listdir(args.output_dir)
              dirs = [d for d in dirs if d.startswith("checkpoint")]
              dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
              path = dirs[-1] if len(dirs) > 0 else None
  
          if path is None:
              accelerator.print(
                  f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
              )
              args.resume_from_checkpoint = None
          else:
              accelerator.print(f"Resuming from checkpoint {path}")
              accelerator.load_state(
                  os.path.join(args.output_dir, path)
              )
              global_step = int(path.split("-")[1])
  
              resume_global_step = (
                  global_step * args.gradient_accumulation_steps
              )
              first_epoch = global_step // num_update_steps_per_epoch
              resume_step = resume_global_step % (
                  num_update_steps_per_epoch
                  * args.gradient_accumulation_steps
              )
  
      # Train!
      for epoch in range(first_epoch, args.num_epochs):
          model.train()
          progress_bar = tqdm(
              total=num_update_steps_per_epoch,
              disable=not accelerator.is_local_main_process,
          )
          progress_bar.set_description(f"Epoch {epoch}")
          for step, batch in enumerate(train_dataloader):
              # Skip steps until we reach the resumed step
              if (
                  args.resume_from_checkpoint
                  and epoch == first_epoch
                  and step < resume_step
              ):
                  if step % args.gradient_accumulation_steps == 0:
                      progress_bar.update(1)
                  continue
  
              clean_images = batch["input"]
              # Sample noise that we'll add to the images
              noise = torch.randn(
                  clean_images.shape,
                  dtype=(
                      torch.float32
                      if args.mixed_precision == "no"
                      else torch.float16
                  ),
              ).to(clean_images.device)
              bsz = clean_images.shape[0]
              # Sample a random timestep for each image
              timesteps = torch.randint(
                  0,
                  noise_scheduler.config.num_train_timesteps,
                  (bsz,),
                  device=clean_images.device,
              ).long()
  
              # Add noise to the clean images according to the noise magnitude at each timestep
              # (this is the forward diffusion process)
              noisy_images = noise_scheduler.add_noise(
                  clean_images, noise, timesteps
              )
  
              with accelerator.accumulate(model):
                  # Predict the noise residual
                  model_output = model(
                      noisy_images, timesteps, return_dict=False
                  )[0]
  
                  if args.prediction_type == "epsilon":
                      loss = F.mse_loss(
                          model_output, noise
                      )  # this could have different weights!
                  elif args.prediction_type == "sample":
                      alpha_t = _extract_into_tensor(
                          noise_scheduler.alphas_cumprod,
                          timesteps,
                          (clean_images.shape[0], 1, 1, 1),
                      )
                      snr_weights = alpha_t / (1 - alpha_t)
                      loss = snr_weights * F.mse_loss(
                          model_output, clean_images, reduction="none"
                      )  # use SNR weighting from distillation paper
                      loss = loss.mean()
                  else:
                      raise ValueError(
                          f"Unsupported prediction type: {args.prediction_type}"
                      )
  
                  accelerator.backward(loss)
  
                  if accelerator.sync_gradients:
                      accelerator.clip_grad_norm_(
                          model.parameters(), 1.0
                      )
                  optimizer.step()
                  lr_scheduler.step()
                  optimizer.zero_grad()
  
              # Checks if the accelerator has performed an optimization step behind the scenes
              if accelerator.sync_gradients:
                  if args.use_ema:
                      ema_model.step(model.parameters())
                  progress_bar.update(1)
                  global_step += 1
  
                  if global_step % args.checkpointing_steps == 0:
                      if accelerator.is_main_process:
                          save_path = os.path.join(
                              args.output_dir,
                              f"checkpoint-{global_step}",
                          )
                          accelerator.save_state(save_path)
                          logger.info(f"Saved state to {save_path}")
  
              logs = {
                  "loss": loss.detach().item(),
                  "lr": lr_scheduler.get_last_lr()[0],
                  "step": global_step,
              }
              if args.use_ema:
                  logs["ema_decay"] = ema_model.cur_decay_value
              progress_bar.set_postfix(**logs)
              accelerator.log(logs, step=global_step)
          progress_bar.close()
  
          accelerator.wait_for_everyone()
  
          # Generate sample images for visual inspection
          if accelerator.is_main_process:
              if (
                  epoch % args.save_images_epochs == 0
                  or epoch == args.num_epochs - 1
              ):
                  unet = accelerator.unwrap_model(model)
  
                  if args.use_ema:
                      ema_model.store(unet.parameters())
                      ema_model.copy_to(unet.parameters())
  
                  pipeline = DDPMPipeline(
                      unet=unet,
                      scheduler=noise_scheduler,
                  )
  
                  generator = torch.Generator(
                      device=pipeline.device
                  ).manual_seed(0)
                  # run pipeline in inference (sample random noise and denoise)
                  images = pipeline(
                      generator=generator,
                      batch_size=args.eval_batch_size,
                      num_inference_steps=args.ddpm_num_inference_steps,
                      output_type="numpy",
                  ).images
  
                  if args.use_ema:
                      ema_model.restore(unet.parameters())
  
                  # denormalize the images and save to tensorboard
                  images_processed = (
                      (images * 255).round().astype("uint8")
                  )
  
                  if args.logger == "tensorboard":
                      if is_accelerate_version(">=", "0.17.0.dev0"):
                          tracker = accelerator.get_tracker(
                              "tensorboard", unwrap=True
                          )
                      else:
                          tracker = accelerator.get_tracker(
                              "tensorboard"
                          )
                      tracker.add_images(
                          "test_samples",
                          images_processed.transpose(0, 3, 1, 2),
                          epoch,
                      )
                  elif args.logger == "wandb":
                      # Upcoming `log_images` helper coming in https://github.com/huggingface/accelerate/pull/962/files
                      accelerator.get_tracker("wandb").log(
                          {
                              "test_samples": [
                                  wandb.Image(img)
                                  for img in images_processed
                              ],
                              "epoch": epoch,
                          },
                          step=global_step,
                      )
  
              if (
                  epoch % args.save_model_epochs == 0
                  or epoch == args.num_epochs - 1
              ):
                  # save the model
                  unet = accelerator.unwrap_model(model)
  
                  if args.use_ema:
                      ema_model.store(unet.parameters())
                      ema_model.copy_to(unet.parameters())
  
                  pipeline = DDPMPipeline(
                      unet=unet,
                      scheduler=noise_scheduler,
                  )
  
                  pipeline.save_pretrained(args.output_dir)
  
                  if args.use_ema:
                      ema_model.restore(unet.parameters())
  
                  if args.push_to_hub:
                      repo.push_to_hub(
                          commit_message=f"Epoch {epoch}",
                          blocking=False,
                      )
  
      accelerator.end_training()
  
imports: import argparse
  import inspect
  import logging
  import math
  import os
  from pathlib import Path
  from typing import Optional
  import accelerate
  import datasets
  import torch
  import torch.nn.functional
  from accelerate import Accelerator
  from accelerate.logging import get_logger
  from accelerate.utils import ProjectConfiguration
  from datasets import load_dataset
  from huggingface_hub import HfFolder, Repository, create_repo, whoami
  from onnxruntime.training.optim.fp16_optimizer import FP16_Optimizer
  from onnxruntime.training.ortmodule import ORTModule
  from packaging import version
  from torchvision import transforms
  from tqdm.auto import tqdm
  import diffusers
  from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel
  from diffusers.optimization import get_scheduler
  from diffusers.training_utils import EMAModel
  from diffusers.utils import check_min_version, is_accelerate_version, is_tensorboard_available, is_wandb_available
  from diffusers.utils.import_utils import is_xformers_available
label_desc: Use image node_id 4 for calling a global function with example usage: main(args) and returns: dict