code: def forward(
      self,
      unet,
      noise_map,
      text_encoding,
      text_mask,
      x_low_res=None,
      cond_scale=1.0,
      sampling_steps=None,
      thresholding_method="dynamic",
  ):
      if self.unet_type == "base":
          assert x_low_res is None
      elif self.unet_type == "sr":
          assert x_low_res is not None
      low_res_cond = (
          {"x_low_res": x_low_res} if x_low_res is not None else {}
      )
      thresholding_method = "dynamic"
      sigma_min = self.sigma_min
      sigma_max = self.sigma_max
      print(
          f"Sampling with sigma in [{sigma_min}, {sigma_max}], cfg={cond_scale}"
      )
      # Time step discretization
      num_steps = sampling_steps if sampling_steps else self.num_steps
      step_indices = torch.arange(num_steps, device=noise_map.device)
      # Table 1: Sampling - Time steps
      t_steps = (
          sigma_max ** (1 / self.rho)
          + step_indices
          / (num_steps - 1)
          * (sigma_min ** (1 / self.rho) - sigma_max ** (1 / self.rho))
      ) ** self.rho
      t_steps = torch.cat(
          [t_steps, torch.zeros_like(t_steps[:1])]
      )  # t_N = 0
  
      # Main sampling loop.
      x_next = noise_map * t_steps[0]
      for i, (t_cur, t_next) in tqdm(
          enumerate(zip(t_steps[:-1], t_steps[1:])),
          total=len(t_steps[:-1]),
      ):  # 0, ..., N-1
          x_cur = x_next
  
          # Increase noise temporarily.
          gamma = (
              min(self.S_churn / num_steps, np.sqrt(2) - 1)
              if self.S_min <= t_cur <= self.S_max
              else 0
          )
          t_hat = (t_cur + gamma * t_cur).to(x_cur.device)
          x_hat = x_cur + (
              t_hat**2 - t_cur**2
          ).sqrt() * self.S_noise * torch.randn_like(x_cur)
  
          # Euler step.
          denoised = unet.forward_with_cond_scale(
              x=x_hat.to(torch.float32),
              time=t_hat.to(torch.float32),
              text_embed=text_encoding,
              text_mask=text_mask,
              cond_scale=cond_scale,
              **low_res_cond,
          )
          d_cur = (x_hat - denoised) / t_hat
          d_cur = thresholding_derivative(
              x_hat,
              t_hat,
              d_cur,
              thresholding_method=thresholding_method,
          )
          x_next = x_hat + (t_next - t_hat) * d_cur
  
          # Apply 2nd order correction.
          if i < num_steps - 1:
              denoised = unet.forward_with_cond_scale(
                  x=x_next.to(torch.float32),
                  time=t_next.to(torch.float32),
                  text_embed=text_encoding,
                  text_mask=text_mask,
                  cond_scale=cond_scale,
                  **low_res_cond,
              )
              d_prime = (x_next - denoised) / t_next
              d_prime = thresholding_derivative(
                  x_next,
                  t_next,
                  d_prime,
                  thresholding_method=thresholding_method,
              )
              x_next = x_hat + (t_next - t_hat) * (
                  0.5 * d_cur + 0.5 * d_prime
              )
      return x_next
  
imports: import numpy
  import torch
  from einops import rearrange
  from tqdm import tqdm
  from nemo.collections.multimodal.modules.imagen.sampler.batch_ops import batch_div, batch_mul
  from nemo.collections.multimodal.modules.imagen.sampler.continuous_ddpm import GaussianDiffusionContinuousTimes
label_desc: Use image node_id 2 for calling the EDMSampler obj's underlying member method code with example usage: obj.forward(unet, noise_map, text_encoding, text_mask, x_low_res, cond_scale, sampling_steps, thresholding_method) and returns: x_next