class_docstr: The DDPOTrainer uses Deep Diffusion Policy Optimization to optimise diffusion models.
  Note, this trainer is heavily inspired by the work here: https://github.com/kvablack/ddpo-pytorch
  As of now only Stable Diffusion based pipelines are supported
  
  Attributes:
      **config** (`DDPOConfig`) -- Configuration object for DDPOTrainer. Check the documentation of `PPOConfig` for more
       details.
      **reward_function** (Callable[[torch.Tensor, Tuple[str], Tuple[Any]], torch.Tensor]) -- Reward function to be used
      **prompt_function** (Callable[[], Tuple[str, Any]]) -- Function to generate prompts to guide model
      **sd_pipeline** (`DDPOStableDiffusionPipeline`) -- Stable Diffusion pipeline to be used for training.
      **image_samples_hook** (Optional[Callable[[Any, Any, Any], Any]]) -- Hook to be called to log images
code: def __init__(
      self,
      config: DDPOConfig,
      reward_function: Callable[
          [torch.Tensor, Tuple[str], Tuple[Any]], torch.Tensor
      ],
      prompt_function: Callable[[], Tuple[str, Any]],
      sd_pipeline: DDPOStableDiffusionPipeline,
      image_samples_hook: Optional[
          Callable[[Any, Any, Any], Any]
      ] = None,
  ):
      if image_samples_hook is None:
          warn(
              "No image_samples_hook provided; no images will be logged"
          )
  
      self.prompt_fn = prompt_function
      self.reward_fn = reward_function
      self.config = config
      self.image_samples_callback = image_samples_hook
  
      accelerator_project_config = ProjectConfiguration(
          **self.config.project_kwargs
      )
  
      if self.config.resume_from:
          self.config.resume_from = os.path.normpath(
              os.path.expanduser(self.config.resume_from)
          )
          if "checkpoint_" not in os.path.basename(
              self.config.resume_from
          ):
              # get the most recent checkpoint in this directory
              checkpoints = list(
                  filter(
                      lambda x: "checkpoint_" in x,
                      os.listdir(self.config.resume_from),
                  )
              )
              if len(checkpoints) == 0:
                  raise ValueError(
                      f"No checkpoints found in {self.config.resume_from}"
                  )
              checkpoint_numbers = sorted(
                  [int(x.split("_")[-1]) for x in checkpoints]
              )
              self.config.resume_from = os.path.join(
                  self.config.resume_from,
                  f"checkpoint_{checkpoint_numbers[-1]}",
              )
  
              accelerator_project_config.iteration = (
                  checkpoint_numbers[-1] + 1
              )
  
      # number of timesteps within each trajectory to train on
      self.num_train_timesteps = int(
          self.config.sample_num_steps
          * self.config.train_timestep_fraction
      )
  
      self.accelerator = Accelerator(
          log_with=self.config.log_with,
          mixed_precision=self.config.mixed_precision,
          project_config=accelerator_project_config,
          # we always accumulate gradients across timesteps; we want config.train.gradient_accumulation_steps to be the
          # number of *samples* we accumulate across, so we need to multiply by the number of training timesteps to get
          # the total number of optimizer steps to accumulate across.
          gradient_accumulation_steps=self.config.train_gradient_accumulation_steps
          * self.num_train_timesteps,
          **self.config.accelerator_kwargs,
      )
  
      is_okay, message = self._config_check()
      if not is_okay:
          raise ValueError(message)
  
      is_using_tensorboard = (
          config.log_with is not None
          and config.log_with == "tensorboard"
      )
  
      if self.accelerator.is_main_process:
          self.accelerator.init_trackers(
              self.config.tracker_project_name,
              config=dict(ddpo_trainer_config=config.to_dict())
              if not is_using_tensorboard
              else config.to_dict(),
              init_kwargs=self.config.tracker_kwargs,
          )
  
      logger.info(f"\n{config}")
  
      set_seed(self.config.seed, device_specific=True)
  
      self.sd_pipeline = sd_pipeline
  
      self.sd_pipeline.set_progress_bar_config(
          position=1,
          disable=not self.accelerator.is_local_main_process,
          leave=False,
          desc="Timestep",
          dynamic_ncols=True,
      )
  
      # For mixed precision training we cast all non-trainable weights (vae, non-lora text_encoder and non-lora unet) to half-precision
      # as these weights are only used for inference, keeping weights in full precision is not required.
      if self.accelerator.mixed_precision == "fp16":
          inference_dtype = torch.float16
      elif self.accelerator.mixed_precision == "bf16":
          inference_dtype = torch.bfloat16
      else:
          inference_dtype = torch.float32
  
      self.sd_pipeline.vae.to(
          self.accelerator.device, dtype=inference_dtype
      )
      self.sd_pipeline.text_encoder.to(
          self.accelerator.device, dtype=inference_dtype
      )
      self.sd_pipeline.unet.to(
          self.accelerator.device, dtype=inference_dtype
      )
  
      trainable_layers = self.sd_pipeline.get_trainable_layers()
  
      self.accelerator.register_save_state_pre_hook(
          self._save_model_hook
      )
      self.accelerator.register_load_state_pre_hook(
          self._load_model_hook
      )
  
      # Enable TF32 for faster training on Ampere GPUs,
      # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
      if self.config.allow_tf32:
          torch.backends.cuda.matmul.allow_tf32 = True
  
      self.optimizer = self._setup_optimizer(
          trainable_layers.parameters()
          if not isinstance(trainable_layers, list)
          else trainable_layers
      )
  
      self.neg_prompt_embed = self.sd_pipeline.text_encoder(
          self.sd_pipeline.tokenizer(
              [""]
              if self.config.negative_prompts is None
              else self.config.negative_prompts,
              return_tensors="pt",
              padding="max_length",
              truncation=True,
              max_length=self.sd_pipeline.tokenizer.model_max_length,
          ).input_ids.to(self.accelerator.device)
      )[0]
  
      if config.per_prompt_stat_tracking:
          self.stat_tracker = PerPromptStatTracker(
              config.per_prompt_stat_tracking_buffer_size,
              config.per_prompt_stat_tracking_min_count,
          )
  
      # NOTE: for some reason, autocast is necessary for non-lora training but for lora training it isn't necessary and it uses
      # more memory
      self.autocast = (
          self.sd_pipeline.autocast or self.accelerator.autocast
      )
  
      if (
          hasattr(self.sd_pipeline, "use_lora")
          and self.sd_pipeline.use_lora
      ):
          unet, self.optimizer = self.accelerator.prepare(
              trainable_layers, self.optimizer
          )
          self.trainable_layers = list(
              filter(lambda p: p.requires_grad, unet.parameters())
          )
      else:
          (
              self.trainable_layers,
              self.optimizer,
          ) = self.accelerator.prepare(trainable_layers, self.optimizer)
  
      if self.config.async_reward_computation:
          self.executor = futures.ThreadPoolExecutor(
              max_workers=config.max_workers
          )
  
      if config.resume_from:
          logger.info(f"Resuming from {config.resume_from}")
          self.accelerator.load_state(config.resume_from)
          self.first_epoch = int(config.resume_from.split("_")[-1]) + 1
      else:
          self.first_epoch = 0
  
imports: import os
  import warnings
  from collections import defaultdict
  from concurrent import futures
  from typing import Any, Callable, Optional, Tuple
  from warnings import warn
  import torch
  from accelerate import Accelerator
  from accelerate.logging import get_logger
  from accelerate.utils import ProjectConfiguration, set_seed
  from huggingface_hub import whoami
  from ..models import DDPOStableDiffusionPipeline
  from .None import BaseTrainer, DDPOConfig
  from .utils import PerPromptStatTracker
label_desc: Use image node_id 1 to create a new DDPOTrainer object from inherited base classes: BaseTrainer with example: obj = DDPOTrainer(config, reward_function, prompt_function, sd_pipeline, image_samples_hook)