code: def forward(self, x, time, text_embed, text_mask, **model_kwargs):
      bs = x.shape[0]
      assert (
          time.ndim <= 1
      ), "time should be in shape of either [bs] or scalar"
      sigma = time
      c_skip = self.sigma_data**2 / (
          sigma**2 + self.sigma_data**2
      )
      c_out = (
          sigma
          * self.sigma_data
          / (sigma**2 + self.sigma_data**2).sqrt()
      )
      c_in = 1 / (self.sigma_data**2 + sigma**2).sqrt()
      c_noise = sigma.log() / 4
  
      if c_noise.ndim < 1:
          c_noise = c_noise.repeat(
              bs,
          )
  
      if self.noise_cond_aug:
          # Applying noise conditioning augmentation
          assert (
              "x_low_res" in model_kwargs
          ), "x_low_res does not exist when attemping to apply noise augmentation"
          x_low_res = model_kwargs["x_low_res"]
          if self.inference:
              batch_size = x_low_res.shape[0]
              time_low_res = (
                  torch.ones(batch_size, device=x_low_res.device)
                  * 0.002
              )
              x_low_res_noisy, time_low_res = self.get_xt(
                  x0=x_low_res, t=time_low_res, epsilon=None
              )
          else:
              x_low_res_noisy, time_low_res = self.get_xt(
                  x0=x_low_res, t=None, epsilon=None
              )
          c_in_noise = (
              1 / (self.sigma_data**2 + time_low_res**2).sqrt()
          )
          c_noise_noise = time_low_res.log() / 4
          model_kwargs["x_low_res"] = batch_mul(
              c_in_noise, x_low_res_noisy
          )
          model_kwargs["time_low_res"] = c_noise_noise
  
      F_x = self.unet(
          batch_mul(c_in, x),
          c_noise,
          text_embed,
          text_mask,
          **model_kwargs,
      )
      D_x = batch_mul(c_skip, x) + batch_mul(c_out, F_x)
      return D_x
  
imports: import torch
  import torch.nn.functional
  from nemo.collections.multimodal.modules.imagen.sampler.batch_ops import batch_mul
  from nemo.collections.multimodal.modules.imagen.sampler.continuous_ddpm import GaussianDiffusionContinuousTimes
  from nemo.collections.multimodal.parts.utils import randn_like
label_desc: Use image node_id 2 for calling the EDMPrecond obj's underlying member method code with example usage: obj.forward(x, time, text_embed, text_mask) and returns: D_x