docstr: Train step function for StableDiffusion.
  
  Parameters:
      text_embeddings (torch.Tensor): Embeddings tensor [B, 512].
      pred_rgb (torch.Tensor): Predicted RGB tensor [B, 3, 512, 512].
      guidance_scale (float): Guidance scaling factor.
      as_latent (bool): If True, considers pred_rgb as latent.
  
  Returns:
      float: Loss value.
docstr_tok: ["Train","step","function","for","StableDiffusion",".","Parameters",":","text_embeddings","(","torch.Tensor",")",":","Embeddings","tensor","[","B",",","512","]",".","pred_rgb","(","torch.Tensor",")",":","Predicted","RGB","tensor","[","B",",","3",",","512",",","512","]",".","guidance_scale","(","float",")",":","Guidance","scaling","factor",".","as_latent","(","bool",")",":","If","True",",","considers","pred_rgb","as","latent",".","Returns",":","float",":","Loss","value","."]
code: def train_step(
      self,
      text_embeddings: torch.Tensor,
      pred_rgb: torch.Tensor,
      guidance_scale: float = 100.0,
      as_latent: bool = False,
  ) -> float:
      """
      Train step function for StableDiffusion.
  
      Parameters:
          text_embeddings (torch.Tensor): Embeddings tensor [B, 512].
          pred_rgb (torch.Tensor): Predicted RGB tensor [B, 3, 512, 512].
          guidance_scale (float): Guidance scaling factor.
          as_latent (bool): If True, considers pred_rgb as latent.
  
      Returns:
          float: Loss value.
      """
      if as_latent:
          latents = (
              F.interpolate(
                  pred_rgb,
                  (64, 64),
                  mode="bilinear",
                  align_corners=False,
              )
              * 2
              - 1
          )
      else:
          pred_rgb_512 = F.interpolate(
              pred_rgb, (512, 512), mode="bilinear", align_corners=False
          )
          latents = self.encode_imgs(pred_rgb_512)
  
      t = torch.randint(
          self.min_step,
          self.max_step + 1,
          (latents.shape[0],),
          dtype=torch.long,
          device=self.device,
      )
  
      with torch.no_grad():
          # add noise
          noise = torch.randn_like(latents)
          latents_noisy = self.scheduler.add_noise(latents, noise, t)
          # pred noise
          latent_model_input = torch.cat([latents_noisy] * 2)
          td = torch.cat([t] * 2)
          noise_pred = self.unet(
              latent_model_input,
              td,
              encoder_hidden_states=text_embeddings,
          ).sample
  
          noise_pred_uncond, noise_pred_pos = noise_pred.chunk(2)
          noise_pred = noise_pred_uncond + guidance_scale * (
              noise_pred_pos - noise_pred_uncond
          )
  
      w = 1 - self.alphas[t]
      grad = w[:, None, None, None] * (noise_pred - noise)
      grad = torch.nan_to_num(grad)
  
      targets = (latents - grad).detach()
      loss = (
          0.5
          * F.mse_loss(latents.float(), targets, reduction="sum")
          / latents.shape[0]
      )
      return loss
  
file_path: NeMo/nemo/collections/multimodal/modules/nerf/guidance/stablediffusion_huggingface_pipeline.py
imports: from typing import List
  import torch
  import torch.nn.functional
  from diffusers import DDIMScheduler, StableDiffusionPipeline
  from nemo.collections.multimodal.modules.nerf.guidance.txt2img_guidance_base import Txt2ImgGuidanceBase
label_desc: Use image node_id 4 for calling the StableDiffusion obj's underlying member method code with example usage: obj.train_step(text_embeddings, pred_rgb, guidance_scale, as_latent) and returns: loss