code: def ddpm_bit_scheduler_step(
      self,
      model_output: torch.FloatTensor,
      timestep: int,
      sample: torch.FloatTensor,
      prediction_type="epsilon",
      generator=None,
      return_dict: bool = True,
  ) -> Union[DDPMSchedulerOutput, Tuple]:
      """
      Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
      process from the learned model outputs (most often the predicted noise).
      Args:
          model_output (`torch.FloatTensor`): direct output from learned diffusion model.
          timestep (`int`): current discrete timestep in the diffusion chain.
          sample (`torch.FloatTensor`):
              current instance of sample being created by diffusion process.
          prediction_type (`str`, default `epsilon`):
              indicates whether the model predicts the noise (epsilon), or the samples (`sample`).
          generator: random number generator.
          return_dict (`bool`): option for returning tuple rather than DDPMSchedulerOutput class
      Returns:
          [`~schedulers.scheduling_utils.DDPMSchedulerOutput`] or `tuple`:
          [`~schedulers.scheduling_utils.DDPMSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. When
          returning a tuple, the first element is the sample tensor.
      """
      t = timestep
  
      if model_output.shape[1] == sample.shape[
          1
      ] * 2 and self.variance_type in ["learned", "learned_range"]:
          model_output, predicted_variance = torch.split(
              model_output, sample.shape[1], dim=1
          )
      else:
          predicted_variance = None
  
      # 1. compute alphas, betas
      alpha_prod_t = self.alphas_cumprod[t]
      alpha_prod_t_prev = (
          self.alphas_cumprod[t - 1] if t > 0 else self.one
      )
      beta_prod_t = 1 - alpha_prod_t
      beta_prod_t_prev = 1 - alpha_prod_t_prev
  
      # 2. compute predicted original sample from predicted noise also called
      # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
      if prediction_type == "epsilon":
          pred_original_sample = (
              sample - beta_prod_t ** (0.5) * model_output
          ) / alpha_prod_t ** (0.5)
      elif prediction_type == "sample":
          pred_original_sample = model_output
      else:
          raise ValueError(
              f"Unsupported prediction_type {prediction_type}."
          )
  
      # 3. Clip "predicted x_0"
      scale = self.bit_scale
      if self.config.clip_sample:
          pred_original_sample = torch.clamp(
              pred_original_sample, -scale, scale
          )
  
      # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
      # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
      pred_original_sample_coeff = (
          alpha_prod_t_prev ** (0.5) * self.betas[t]
      ) / beta_prod_t
      current_sample_coeff = (
          self.alphas[t] ** (0.5) * beta_prod_t_prev / beta_prod_t
      )
  
      # 5. Compute predicted previous sample µ_t
      # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
      pred_prev_sample = (
          pred_original_sample_coeff * pred_original_sample
          + current_sample_coeff * sample
      )
  
      # 6. Add noise
      variance = 0
      if t > 0:
          noise = torch.randn(
              model_output.size(),
              dtype=model_output.dtype,
              layout=model_output.layout,
              generator=generator,
          ).to(model_output.device)
          variance = (
              self._get_variance(
                  t, predicted_variance=predicted_variance
              )
              ** 0.5
          ) * noise
  
      pred_prev_sample = pred_prev_sample + variance
  
      if not return_dict:
          return (pred_prev_sample,)
  
      return DDPMSchedulerOutput(
          prev_sample=pred_prev_sample,
          pred_original_sample=pred_original_sample,
      )
  
file_path: llmga/llmga/diffusers/examples/community/bit_diffusion.py
imports: from typing import Optional, Tuple, Union
  import torch
  from einops import rearrange, reduce
  from diffusers import DDIMScheduler, DDPMScheduler, DiffusionPipeline, ImagePipelineOutput, UNet2DConditionModel
  from diffusers.schedulers.scheduling_ddim import DDIMSchedulerOutput
  from diffusers.schedulers.scheduling_ddpm import DDPMSchedulerOutput
label_desc: Use image node_id 4 for calling a global function with example usage: ddpm_bit_scheduler_step(model_output, timestep, sample, prediction_type, generator, return_dict) and returns: DDPMSchedulerOutput, pred_prev_sample