content: import warnings
  warnings.filterwarnings("ignore")  # ignore all warnings
  import diffusers.utils.logging as diffusion_logging
  diffusion_logging.set_verbosity_error()  # ignore diffusers warnings
  
  from partcrafter_src.utils.typing_utils import *
  
  import os
  import argparse
  import logging
  import time
  import math
  import gc
  from packaging import version
  
  import trimesh
  from PIL import Image
  import numpy as np
  import wandb
  from tqdm import tqdm
  
  import torch
  import torch.nn.functional as tF
  import accelerate
  from accelerate import Accelerator
  from accelerate.logging import get_logger as get_accelerate_logger
  from accelerate import DataLoaderConfiguration, DeepSpeedPlugin
  from diffusers.training_utils import (
      compute_density_for_timestep_sampling,
      compute_loss_weighting_for_sd3
  )
  
  from transformers import (
      BitImageProcessor,
      Dinov2Model,
  )
  from partcrafter_src.schedulers import RectifiedFlowScheduler
  from partcrafter_src.models.autoencoders import TripoSGVAEModel
  from partcrafter_src.models.transformers import PartCrafterDiTModel
  from partcrafter_src.pipelines.pipeline_partcrafter import PartCrafterPipeline
  
  from partcrafter_src.datasets import (
      ObjaversePartDataset, 
      BatchedObjaversePartDataset, 
      MultiEpochsDataLoader, 
      yield_forever
  )
  from partcrafter_src.utils.data_utils import get_colored_mesh_composition
  from partcrafter_src.utils.train_utils import (
      MyEMAModel, 
      get_configs,
      get_optimizer,
      get_lr_scheduler,
      save_experiment_params,
      save_model_architecture,
  )
  from partcrafter_src.utils.render_utils import (
      render_views_around_mesh, 
      render_normal_views_around_mesh, 
      make_grid_for_images_or_videos,
      export_renderings
  )
  from partcrafter_src.utils.metric_utils import compute_cd_and_f_score_in_training
  
  def main():
      PROJECT_NAME = "PartCrafter"
  
      parser = argparse.ArgumentParser(
          description="Train a diffusion model for 3D object generation",
      )
  
      parser.add_argument(
          "--config",
          type=str,
          required=True,
          help="Path to the config file"
      )
      parser.add_argument(
          "--tag",
          type=str,
          default=None,
          help="Tag that refers to the current experiment"
      )
      parser.add_argument(
          "--output_dir",
          type=str,
          default="output",
          help="Path to the output directory"
      )
      parser.add_argument(
          "--resume_from_iter",
          type=int,
          default=None,
          help="The iteration to load the checkpoint from"
      )
      parser.add_argument(
          "--seed",
          type=int,
          default=0,
          help="Seed for the PRNG"
      )
      parser.add_argument(
          "--offline_wandb",
          action="store_true",
          help="Use offline WandB for experiment tracking"
      )
  
      parser.add_argument(
          "--max_train_steps",
          type=int,
          default=None,
          help="The max iteration step for training"
      )
      parser.add_argument(
          "--max_val_steps",
          type=int,
          default=2,
          help="The max iteration step for validation"
      )
      parser.add_argument(
          "--num_workers",
          type=int,
          default=32,
          help="The number of processed spawned by the batch provider"
      )
      parser.add_argument(
          "--pin_memory",
          action="store_true",
          help="Pin memory for the data loader"
      )
  
      parser.add_argument(
          "--use_ema",
          action="store_true",
          help="Use EMA model for training"
      )
      parser.add_argument(
          "--scale_lr",
          action="store_true",
          help="Scale lr with total batch size (base batch size: 256)"
      )
      parser.add_argument(
          "--max_grad_norm",
          type=float,
          default=1.,
          help="Max gradient norm for gradient clipping"
      )
      parser.add_argument(
          "--gradient_accumulation_steps",
          type=int,
          default=1,
          help="Number of updates steps to accumulate before performing a backward/update pass"
      )
      parser.add_argument(
          "--mixed_precision",
          type=str,
          default="fp16",
          choices=["no", "fp16", "bf16"],
          help="Type of mixed precision training"
      )
      parser.add_argument(
          "--allow_tf32",
          action="store_true",
          help="Enable TF32 for faster training on Ampere GPUs"
      )
  
      parser.add_argument(
          "--val_guidance_scales",
          type=list,
          nargs="+",
          default=[7.0],
          help="CFG scale used for validation"
      )
  
      parser.add_argument(
          "--use_deepspeed",
          action="store_true",
          help="Use DeepSpeed for training"
      )
      parser.add_argument(
          "--zero_stage",
          type=int,
          default=1,
          choices=[1, 2, 3],  # https://huggingface.co/docs/accelerate/usage_guides/deepspeed
          help="ZeRO stage type for DeepSpeed"
      )
  
      parser.add_argument(
          "--from_scratch",
          action="store_true",
          help="Train from scratch"
      )
      parser.add_argument(
          "--load_pretrained_model",
          type=str,
          default=None,
          help="Tag of a pretrained PartCrafterDiTModel in this project"
      )
      parser.add_argument(
          "--load_pretrained_model_ckpt",
          type=int,
          default=-1,
          help="Iteration of the pretrained PartCrafterDiTModel checkpoint"
      )
  
      # Parse the arguments
      args, extras = parser.parse_known_args()
      # Parse the config file
      configs = get_configs(args.config, extras)  # change yaml configs by `extras`
  
      args.val_guidance_scales = [float(x[0]) if isinstance(x, list) else float(x) for x in args.val_guidance_scales]
      if args.max_val_steps > 0: 
          # If enable validation, the max_val_steps must be a multiple of nrow
          # Always keep validation batchsize 1
          divider = configs["val"]["nrow"]
          args.max_val_steps = max(args.max_val_steps, divider)
          if args.max_val_steps % divider != 0:
              args.max_val_steps = (args.max_val_steps // divider + 1) * divider
  
      # Create an experiment directory using the `tag`
      if args.tag is None:
          args.tag = time.strftime("%Y%m%d_%H_%M_%S")
      exp_dir = os.path.join(args.output_dir, args.tag)
      ckpt_dir = os.path.join(exp_dir, "checkpoints")
      eval_dir = os.path.join(exp_dir, "evaluations")
      os.makedirs(ckpt_dir, exist_ok=True)
      os.makedirs(eval_dir, exist_ok=True)
  
      # Initialize the logger
      logging.basicConfig(
          format="%(asctime)s - %(message)s",
          datefmt="%Y/%m/%d %H:%M:%S",
          level=logging.INFO
      )
      logger = get_accelerate_logger(__name__, log_level="INFO")
      file_handler = logging.FileHandler(os.path.join(exp_dir, "log.txt"))  # output to file
      file_handler.setFormatter(logging.Formatter(
          fmt="%(asctime)s - %(message)s",
          datefmt="%Y/%m/%d %H:%M:%S"
      ))
      logger.logger.addHandler(file_handler)
      logger.logger.propagate = True  # propagate to the root logger (console)
  
      # Set DeepSpeed config
      if args.use_deepspeed:
          deepspeed_plugin = DeepSpeedPlugin(
              gradient_accumulation_steps=args.gradient_accumulation_steps,
              gradient_clipping=args.max_grad_norm,
              zero_stage=int(args.zero_stage),
              offload_optimizer_device="cpu",  # hard-coded here, TODO: make it configurable
          )
      else:
          deepspeed_plugin = None
  
      # Initialize the accelerator
      accelerator = Accelerator(
          project_dir=exp_dir,
          gradient_accumulation_steps=args.gradient_accumulation_steps,
          mixed_precision=args.mixed_precision,
          split_batches=False,  # batch size per GPU
          dataloader_config=DataLoaderConfiguration(non_blocking=args.pin_memory),
          deepspeed_plugin=deepspeed_plugin,
      )
      logger.info(f"Accelerator state:\n{accelerator.state}\n")
  
      # Set the random seed
      if args.seed >= 0:
          accelerate.utils.set_seed(args.seed)
          logger.info(f"You have chosen to seed([{args.seed}]) the experiment [{args.tag}]\n")
  
      # Enable TF32 for faster training on Ampere GPUs
      if args.allow_tf32:
          torch.backends.cuda.matmul.allow_tf32 = True
  
      train_dataset = BatchedObjaversePartDataset(
          configs=configs,
          batch_size=configs["train"]["batch_size_per_gpu"],
          is_main_process=accelerator.is_main_process,
          shuffle=True,
          training=True,
      )
      val_dataset = ObjaversePartDataset(
          configs=configs,
          training=False,
      )
      train_loader = MultiEpochsDataLoader(
          train_dataset,
          batch_size=configs["train"]["batch_size_per_gpu"],
          num_workers=args.num_workers,
          drop_last=True,
          pin_memory=args.pin_memory,
          collate_fn=train_dataset.collate_fn,
      )
      val_loader = MultiEpochsDataLoader(
          val_dataset,
          batch_size=configs["val"]["batch_size_per_gpu"],
          num_workers=args.num_workers,
          drop_last=True,
          pin_memory=args.pin_memory,
      )
      random_val_loader = MultiEpochsDataLoader(
          val_dataset,
          batch_size=configs["val"]["batch_size_per_gpu"],
          shuffle=True,
          num_workers=args.num_workers,
          drop_last=True,
          pin_memory=args.pin_memory,
      )
  
      logger.info(f"Loaded [{len(train_dataset)}] training samples and [{len(val_dataset)}] validation samples\n")
  
      # Compute the effective batch size and scale learning rate
      total_batch_size = configs["train"]["batch_size_per_gpu"] * \
          accelerator.num_processes * args.gradient_accumulation_steps
      configs["train"]["total_batch_size"] = total_batch_size
      if args.scale_lr:
          configs["optimizer"]["lr"] *= (total_batch_size / 256)
          configs["lr_scheduler"]["max_lr"] = configs["optimizer"]["lr"]
      
      # Initialize the model
      logger.info("Initializing the model...")
      vae = TripoSGVAEModel.from_pretrained(
          configs["model"]["pretrained_model_name_or_path"],
          subfolder="vae"
      )
      feature_extractor_dinov2 = BitImageProcessor.from_pretrained(
          configs["model"]["pretrained_model_name_or_path"],
          subfolder="feature_extractor_dinov2"
      )
      image_encoder_dinov2 = Dinov2Model.from_pretrained(
          configs["model"]["pretrained_model_name_or_path"],
          subfolder="image_encoder_dinov2"
      )
  
      enable_part_embedding = configs["model"]["transformer"].get("enable_part_embedding", True)
      enable_local_cross_attn = configs["model"]["transformer"].get("enable_local_cross_attn", True)
      enable_global_cross_attn = configs["model"]["transformer"].get("enable_global_cross_attn", True)
      global_attn_block_ids = configs["model"]["transformer"].get("global_attn_block_ids", None)
      if global_attn_block_ids is not None:
          global_attn_block_ids = list(global_attn_block_ids)
      global_attn_block_id_range = configs["model"]["transformer"].get("global_attn_block_id_range", None)
      if global_attn_block_id_range is not None:
          global_attn_block_id_range = list(global_attn_block_id_range)
      if args.from_scratch:
          logger.info(f"Initialize PartCrafterDiTModel from scratch\n")
          transformer = PartCrafterDiTModel.from_config(
              os.path.join(
                  configs["model"]["pretrained_model_name_or_path"],
                  "transformer"
              ), 
              enable_part_embedding=enable_part_embedding,
              enable_local_cross_attn=enable_local_cross_attn,
              enable_global_cross_attn=enable_global_cross_attn,
              global_attn_block_ids=global_attn_block_ids,
              global_attn_block_id_range=global_attn_block_id_range,
          )
      elif args.load_pretrained_model is None:
          logger.info(f"Load pretrained TripoSGDiTModel to initialize PartCrafterDiTModel from [{configs['model']['pretrained_model_name_or_path']}]\n")
          transformer, loading_info = PartCrafterDiTModel.from_pretrained(
              configs["model"]["pretrained_model_name_or_path"],
              subfolder="transformer",
              low_cpu_mem_usage=False, 
              output_loading_info=True, 
              enable_part_embedding=enable_part_embedding,
              enable_local_cross_attn=enable_local_cross_attn,
              enable_global_cross_attn=enable_global_cross_attn,
              global_attn_block_ids=global_attn_block_ids,
              global_attn_block_id_range=global_attn_block_id_range,
          )
      else:
          logger.info(f"Load PartCrafterDiTModel EMA checkpoint from [{args.load_pretrained_model}] iteration [{args.load_pretrained_model_ckpt:06d}]\n")
          path = os.path.join(
              args.output_dir,
              args.load_pretrained_model, 
              "checkpoints", 
              f"{args.load_pretrained_model_ckpt:06d}"
          )
          transformer, loading_info = PartCrafterDiTModel.from_pretrained(
              path, 
              subfolder="transformer_ema",
              low_cpu_mem_usage=False, 
              output_loading_info=True, 
              enable_part_embedding=enable_part_embedding,
              enable_local_cross_attn=enable_local_cross_attn,
              enable_global_cross_attn=enable_global_cross_attn,
              global_attn_block_ids=global_attn_block_ids,
              global_attn_block_id_range=global_attn_block_id_range,
          )
      if not args.from_scratch:
          for v in loading_info.values():
              if v and len(v) > 0:
                  logger.info(f"Loading info of PartCrafterDiTModel: {loading_info}\n")
                  break
  
      noise_scheduler = RectifiedFlowScheduler.from_pretrained(
          configs["model"]["pretrained_model_name_or_path"],
          subfolder="scheduler"
      )
  
      if args.use_ema:
          ema_transformer = MyEMAModel(
              transformer.parameters(),
              model_cls=PartCrafterDiTModel,
              model_config=transformer.config,
              **configs["train"]["ema_kwargs"]
          )
  
      # Freeze VAE and image encoder
      vae.requires_grad_(False)
      image_encoder_dinov2.requires_grad_(False)
      vae.eval()
      image_encoder_dinov2.eval()
  
      trainable_modules = configs["train"].get("trainable_modules", None)
      if trainable_modules is None:
          transformer.requires_grad_(True)
      else:
          trainable_module_names = []
          transformer.requires_grad_(False)
          for name, module in transformer.named_modules():
              for module_name in tuple(trainable_modules.split(",")):
                  if module_name in name:
                      for params in module.parameters():
                          params.requires_grad = True
                      trainable_module_names.append(name)
          logger.info(f"Trainable parameter names: {trainable_module_names}\n")
  
      # transformer.enable_xformers_memory_efficient_attention()  # use `tF.scaled_dot_product_attention` instead
  
      # `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_transformer.save_pretrained(os.path.join(output_dir, "transformer_ema"))
  
                  for i, model in enumerate(models):
                      model.save_pretrained(os.path.join(output_dir, "transformer"))
  
                      # Make sure to pop weight so that corresponding model is not saved again
                      if weights:
                          weights.pop()
  
          def load_model_hook(models, input_dir):
              if args.use_ema:
                  load_model = MyEMAModel.from_pretrained(os.path.join(input_dir, "transformer_ema"), PartCrafterDiTModel)
                  ema_transformer.load_state_dict(load_model.state_dict())
                  ema_transformer.to(accelerator.device)
                  del load_model
  
              for _ in range(len(models)):
                  # Pop models so that they are not loaded again
                  model = models.pop()
  
                  # Load diffusers style into model
                  load_model = PartCrafterDiTModel.from_pretrained(input_dir, subfolder="transformer")
                  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)
  
      if configs["train"]["grad_checkpoint"]:
          transformer.enable_gradient_checkpointing()
  
      # Initialize the optimizer and learning rate scheduler
      logger.info("Initializing the optimizer and learning rate scheduler...\n")
      name_lr_mult = configs["train"].get("name_lr_mult", None)
      lr_mult = configs["train"].get("lr_mult", 1.0)
      params, params_lr_mult, names_lr_mult = [], [], []
      for name, param in transformer.named_parameters():
          if name_lr_mult is not None:
              for k in name_lr_mult.split(","):
                  if k in name:
                      params_lr_mult.append(param)
                      names_lr_mult.append(name)
              if name not in names_lr_mult:
                  params.append(param)
          else:
              params.append(param)
      optimizer = get_optimizer(
          params=[
              {"params": params, "lr": configs["optimizer"]["lr"]},
              {"params": params_lr_mult, "lr": configs["optimizer"]["lr"] * lr_mult}
          ],
          **configs["optimizer"]
      )
      if name_lr_mult is not None:
          logger.info(f"Learning rate x [{lr_mult}] parameter names: {names_lr_mult}\n")
  
      configs["lr_scheduler"]["total_steps"] = configs["train"]["epochs"] * math.ceil(
          len(train_loader) // accelerator.num_processes / args.gradient_accumulation_steps)  # only account updated steps
      configs["lr_scheduler"]["total_steps"] *= accelerator.num_processes  # for lr scheduler setting
      if "num_warmup_steps" in configs["lr_scheduler"]:
          configs["lr_scheduler"]["num_warmup_steps"] *= accelerator.num_processes  # for lr scheduler setting
      lr_scheduler = get_lr_scheduler(optimizer=optimizer, **configs["lr_scheduler"])
      configs["lr_scheduler"]["total_steps"] //= accelerator.num_processes  # reset for multi-gpu
      if "num_warmup_steps" in configs["lr_scheduler"]:
          configs["lr_scheduler"]["num_warmup_steps"] //= accelerator.num_processes  # reset for multi-gpu
  
      # Prepare everything with `accelerator`
      transformer, optimizer, lr_scheduler, train_loader, val_loader, random_val_loader = accelerator.prepare(
          transformer, optimizer, lr_scheduler, train_loader, val_loader, random_val_loader
      )
      # Set classes explicitly for everything
      transformer: DistributedDataParallel
      optimizer: AcceleratedOptimizer
      lr_scheduler: AcceleratedScheduler
      train_loader: DataLoaderShard
      val_loader: DataLoaderShard
      random_val_loader: DataLoaderShard
  
      if args.use_ema:
          ema_transformer.to(accelerator.device)
  
      # For mixed precision training we cast all non-trainable weigths to half-precision
      # as these weights are only used for inference, keeping weights in full precision is not required.
      weight_dtype = torch.float32
      if accelerator.mixed_precision == "fp16":
          weight_dtype = torch.float16
      elif accelerator.mixed_precision == "bf16":
          weight_dtype = torch.bfloat16
  
      # Move `vae` and `image_encoder_dinov2` to gpu and cast to `weight_dtype`
      vae.to(accelerator.device, dtype=weight_dtype)
      image_encoder_dinov2.to(accelerator.device, dtype=weight_dtype)
  
      # Training configs after distribution and accumulation setup
      updated_steps_per_epoch = math.ceil(len(train_loader) / args.gradient_accumulation_steps)
      total_updated_steps = configs["lr_scheduler"]["total_steps"]
      if args.max_train_steps is None:
          args.max_train_steps = total_updated_steps
      assert configs["train"]["epochs"] * updated_steps_per_epoch == total_updated_steps
      if accelerator.num_processes > 1 and accelerator.is_main_process:
          print()
      accelerator.wait_for_everyone()
      logger.info(f"Total batch size: [{total_batch_size}]")
      logger.info(f"Learning rate: [{configs['optimizer']['lr']}]")
      logger.info(f"Gradient Accumulation steps: [{args.gradient_accumulation_steps}]")
      logger.info(f"Total epochs: [{configs['train']['epochs']}]")
      logger.info(f"Total steps: [{total_updated_steps}]")
      logger.info(f"Steps for updating per epoch: [{updated_steps_per_epoch}]")
      logger.info(f"Steps for validation: [{len(val_loader)}]\n")
  
      # (Optional) Load checkpoint
      global_update_step = 0
      if args.resume_from_iter is not None:
          if args.resume_from_iter < 0:
              args.resume_from_iter = int(sorted(os.listdir(ckpt_dir))[-1])
          logger.info(f"Load checkpoint from iteration [{args.resume_from_iter}]\n")
          # Load everything
          if version.parse(torch.__version__) >= version.parse("2.4.0"):
              torch.serialization.add_safe_globals([
                  int, list, dict, 
                  defaultdict,
                  Any,
                  DictConfig, ListConfig, Metadata, ContainerMetadata, AnyNode
              ]) # avoid deserialization error when loading optimizer state
          accelerator.load_state(os.path.join(ckpt_dir, f"{args.resume_from_iter:06d}"))  # torch < 2.4.0 here for `weights_only=False`
          global_update_step = int(args.resume_from_iter)
  
      # Save all experimental parameters and model architecture of this run to a file (args and configs)
      if accelerator.is_main_process:
          exp_params = save_experiment_params(args, configs, exp_dir)
          save_model_architecture(accelerator.unwrap_model(transformer), exp_dir)
  
      # WandB logger
      if accelerator.is_main_process:
          if args.offline_wandb:
              os.environ["WANDB_MODE"] = "offline"
          wandb.init(
              project=PROJECT_NAME, name=args.tag,
              config=exp_params, dir=exp_dir,
              resume=True
          )
          # Wandb artifact for logging experiment information
          arti_exp_info = wandb.Artifact(args.tag, type="exp_info")
          arti_exp_info.add_file(os.path.join(exp_dir, "params.yaml"))
          arti_exp_info.add_file(os.path.join(exp_dir, "model.txt"))
          arti_exp_info.add_file(os.path.join(exp_dir, "log.txt"))  # only save the log before training
          wandb.log_artifact(arti_exp_info)
  
      def get_sigmas(timesteps: Tensor, n_dim: int, dtype=torch.float32):
          sigmas = noise_scheduler.sigmas.to(dtype=dtype, device=accelerator.device)
          schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device)
          timesteps = timesteps.to(accelerator.device)
  
          step_indices = [(schedule_timesteps == t).nonzero()[0].item() for t in timesteps]
  
          sigma = sigmas[step_indices].flatten()
          while len(sigma.shape) < n_dim:
              sigma = sigma.unsqueeze(-1)
          return sigma
  
      # Start training
      if accelerator.is_main_process:
          print()
      logger.info(f"Start training into {exp_dir}\n")
      logger.logger.propagate = False  # not propagate to the root logger (console)
      progress_bar = tqdm(
          range(total_updated_steps),
          initial=global_update_step,
          desc="Training",
          ncols=125,
          disable=not accelerator.is_main_process
      )
      for batch in yield_forever(train_loader):
  
          if global_update_step == args.max_train_steps:
              progress_bar.close()
              logger.logger.propagate = True  # propagate to the root logger (console)
              if accelerator.is_main_process:
                  wandb.finish()
              logger.info("Training finished!\n")
              return
  
          transformer.train()
  
          with accelerator.accumulate(transformer):
              
              images = batch["images"] # [N, H, W, 3]
              with torch.no_grad():
                  images = feature_extractor_dinov2(images=images, return_tensors="pt").pixel_values
              images = images.to(device=accelerator.device, dtype=weight_dtype)
              with torch.no_grad():
                  image_embeds = image_encoder_dinov2(images).last_hidden_state
              negative_image_embeds = torch.zeros_like(image_embeds)
  
              part_surfaces = batch["part_surfaces"] # [N, P, 6]
              part_surfaces = part_surfaces.to(device=accelerator.device, dtype=weight_dtype)
  
              num_parts = batch["num_parts"] # [M, ] The shape of num_parts is not fixed
              num_objects = num_parts.shape[0] # M
  
              with torch.no_grad():
                  latents = vae.encode(
                      part_surfaces, 
                      **configs["model"]["vae"]
                  ).latent_dist.sample()
  
              noise = torch.randn_like(latents)
              # For weighting schemes where we sample timesteps non-uniformly
              u = compute_density_for_timestep_sampling(
                  weighting_scheme=configs["train"]["weighting_scheme"],
                  batch_size=num_objects,
                  logit_mean=configs["train"]["logit_mean"],
                  logit_std=configs["train"]["logit_std"],
                  mode_scale=configs["train"]["mode_scale"],
              )
              indices = (u * noise_scheduler.config.num_train_timesteps).long()
              timesteps = noise_scheduler.timesteps[indices].to(accelerator.device) # [M, ]
              # Repeat the timesteps for each part
              timesteps = timesteps.repeat_interleave(num_parts) # [N, ]
  
              sigmas = get_sigmas(timesteps, len(latents.shape), weight_dtype)
              latent_model_input = noisy_latents = (1. - sigmas) * latents + sigmas * noise
  
              if configs["train"]["cfg_dropout_prob"] > 0:
                  # We use the same dropout mask for the same part
                  dropout_mask = torch.rand(num_objects, device=accelerator.device) < configs["train"]["cfg_dropout_prob"] # [M, ]
                  dropout_mask = dropout_mask.repeat_interleave(num_parts) # [N, ]
                  if dropout_mask.any():
                      image_embeds[dropout_mask] = negative_image_embeds[dropout_mask]
  
              model_pred = transformer(
                  hidden_states=latent_model_input,
                  timestep=timesteps,
                  encoder_hidden_states=image_embeds, 
                  attention_kwargs={"num_parts": num_parts}
              ).sample
  
              if configs["train"]["training_objective"] == "x0":  # Section 5 of https://arxiv.org/abs/2206.00364
                  model_pred = model_pred * (-sigmas) + noisy_latents  # predicted x_0
                  target = latents
              elif configs["train"]["training_objective"] == 'v':  # flow matching
                  target = noise - latents
              elif configs["train"]["training_objective"] == '-v':  # reverse flow matching
                  # The training objective for TripoSG is the reverse of the flow matching objective. 
                  # It uses "different directions", i.e., the negative velocity. 
                  # This is probably a mistake in engineering, not very harmful. 
                  # In TripoSG's rectified flow scheduler, prev_sample = sample + (sigma - sigma_next) * model_output
                  # See TripoSG's scheduler https://github.com/VAST-AI-Research/TripoSG/blob/main/triposg/schedulers/scheduling_rectified_flow.py#L296
                  # While in diffusers's flow matching scheduler, prev_sample = sample + (sigma_next - sigma) * model_output
                  # See https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py#L454
                  target = latents - noise
              else:
                  raise ValueError(f"Unknown training objective [{configs['train']['training_objective']}]")
  
              # For these weighting schemes use a uniform timestep sampling, so post-weight the loss
              weighting = compute_loss_weighting_for_sd3(
                  configs["train"]["weighting_scheme"],
                  sigmas
              )
  
              loss = weighting * tF.mse_loss(model_pred.float(), target.float(), reduction="none")
              loss = loss.mean(dim=list(range(1, len(loss.shape))))
  
              # Backpropagate
              accelerator.backward(loss.mean())
              if accelerator.sync_gradients:
                  accelerator.clip_grad_norm_(transformer.parameters(), args.max_grad_norm)
  
              optimizer.step()
              lr_scheduler.step()
              optimizer.zero_grad()
  
          # Checks if the accelerator has performed an optimization step behind the scenes
          if accelerator.sync_gradients:
              # Gather the losses across all processes for logging (if we use distributed training)
              loss = accelerator.gather(loss.detach()).mean()
  
              logs = {
                  "loss": loss.item(),
                  "lr": lr_scheduler.get_last_lr()[0]
              }
              if args.use_ema:
                  ema_transformer.step(transformer.parameters())
                  logs.update({"ema": ema_transformer.cur_decay_value})
  
              progress_bar.set_postfix(**logs)
              progress_bar.update(1)
              global_update_step += 1
  
              logger.info(
                  f"[{global_update_step:06d} / {total_updated_steps:06d}] " +
                  f"loss: {logs['loss']:.4f}, lr: {logs['lr']:.2e}" +
                  f", ema: {logs['ema']:.4f}" if args.use_ema else ""
              )
  
              # Log the training progress
              if (
                  global_update_step % configs["train"]["log_freq"] == 0 
                  or global_update_step == 1
                  or global_update_step % updated_steps_per_epoch == 0 # last step of an epoch
              ):  
                  if accelerator.is_main_process:
                      wandb.log({
                          "training/loss": logs["loss"],
                          "training/lr": logs["lr"],
                      }, step=global_update_step)
                      if args.use_ema:
                          wandb.log({
                              "training/ema": logs["ema"]
                          }, step=global_update_step)
  
              # Save checkpoint
              if (
                  global_update_step % configs["train"]["save_freq"] == 0  # 1. every `save_freq` steps
                  or global_update_step % (configs["train"]["save_freq_epoch"] * updated_steps_per_epoch) == 0  # 2. every `save_freq_epoch` epochs
                  or global_update_step == total_updated_steps # 3. last step of an epoch
                  # or global_update_step == 1 # 4. first step
              ): 
  
                  gc.collect()
                  if accelerator.distributed_type == accelerate.utils.DistributedType.DEEPSPEED:
                      # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues
                      accelerator.save_state(os.path.join(ckpt_dir, f"{global_update_step:06d}"))
                  elif accelerator.is_main_process:
                      accelerator.save_state(os.path.join(ckpt_dir, f"{global_update_step:06d}"))
                  accelerator.wait_for_everyone()  # ensure all processes have finished saving
                  gc.collect()
  
              # Evaluate on the validation set
              if args.max_val_steps > 0 and (
                  (global_update_step % configs["train"]["early_eval_freq"] == 0 and global_update_step < configs["train"]["early_eval"])  # 1. more frequently at the beginning
                  or global_update_step % configs["train"]["eval_freq"] == 0  # 2. every `eval_freq` steps
                  or global_update_step % (configs["train"]["eval_freq_epoch"] * updated_steps_per_epoch) == 0  # 3. every `eval_freq_epoch` epochs
                  or global_update_step == total_updated_steps # 4. last step of an epoch
                  or global_update_step == 1 # 5. first step
              ):  
  
                  # Use EMA parameters for evaluation
                  if args.use_ema:
                      # Store the Transformer parameters temporarily and load the EMA parameters to perform inference
                      ema_transformer.store(transformer.parameters())
                      ema_transformer.copy_to(transformer.parameters())
  
                  transformer.eval()
  
                  log_validation(
                      val_loader, random_val_loader,
                      feature_extractor_dinov2, image_encoder_dinov2,
                      vae, transformer,
                      global_update_step, eval_dir,
                      accelerator, logger,
                      args, configs
                  )
  
                  if args.use_ema:
  