class_docstr: an UNet wrapper to implement additional operations in `forward()` method.
  
  ...
  
  Attributes
  ----------
  model : AcceleratedLightningModule
      the original nano-optimized model
docstr: Run inference.
code: def forward(
      self,
      sample: torch.FloatTensor,
      timestep: Union[torch.Tensor, float, int],
      encoder_hidden_states: torch.Tensor,
      class_labels: Optional[torch.Tensor] = None,
      timestep_cond: Optional[torch.Tensor] = None,
      attention_mask: Optional[torch.Tensor] = None,
      cross_attention_kwargs: Optional[Dict[str, Any]] = None,
      down_block_additional_residuals: Optional[
          Tuple[torch.Tensor]
      ] = None,
      mid_block_additional_residual: Optional[torch.Tensor] = None,
      return_dict: bool = True,
  ):
      """Run inference."""
      extra_args = []
      if cross_attention_kwargs is not None:
          attn_scale = cross_attention_kwargs.pop("scale", None)
          if attn_scale is not None:
              attn_scale = torch.Tensor([attn_scale]).float()
              extra_args.append(attn_scale)
              # print(f"Inference with lora scale: {attn_scale}")
          # TODO: warning
          # for key, value in cross_attention_kwargs.items():
          # print(f"Ignoring cross_attention_kwargs {key}, value: {value}")
  
      if isinstance(self.model, PytorchOpenVINOModel):
          timestep = timestep[None]
  
      temp = get_dummy_unet_additional_residuals()
      (
          down_block_additional_residuals,
          mid_block_additional_residual,
      ) = temp
      extra_args.extend(down_block_additional_residuals)
      extra_args.append(mid_block_additional_residual)
  
      noise_pred = self.model(
          sample, timestep, encoder_hidden_states, *extra_args
      )
      if hasattr(noise_pred, "sample"):
          noise_pred = noise_pred.sample
      if isinstance(noise_pred, tuple):
          noise_pred = noise_pred[0]
      elif isinstance(noise_pred, dict):
          noise_pred = noise_pred["sample"]
  
      return UNet2DConditionOutput(sample=noise_pred)
  
file_path: BigDL/python/nano/src/bigdl/nano/diffusion/diffusers/modules/unet.py
imports: import torch
  from typing import Any, Dict, List, Optional, Tuple, Union
  from diffusers.models.unet_2d_condition import UNet2DConditionOutput
  from bigdl.nano.pytorch.model import AcceleratedLightningModule
  from bigdl.nano.deps.openvino.pytorch.model import PytorchOpenVINOModel
  from bigdl.nano.diffusion.utils.pipeline import get_dummy_unet_additional_residuals
label_desc: Use image node_id 2 for calling the NanoUNet obj's underlying member method code with example usage: obj.forward(sample, timestep, encoder_hidden_states, class_labels, timestep_cond, attention_mask, cross_attention_kwargs, down_block_additional_residuals, mid_block_additional_residual, return_dict) and returns: UNet2DConditionOutput