code: def upsample_2d(
      hidden_states: torch.FloatTensor,
      kernel: Optional[torch.FloatTensor] = None,
      factor: int = 2,
      gain: float = 1,
  ) -> torch.FloatTensor:
      r"""Upsample2D a batch of 2D images with the given filter.
      Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` and upsamples each image with the given
      filter. The filter is normalized so that if the input pixels are constant, they will be scaled by the specified
      `gain`. Pixels outside the image are assumed to be zero, and the filter is padded with zeros so that its shape is
      a: multiple of the upsampling factor.
  
      Args:
          hidden_states (`torch.FloatTensor`):
              Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
          kernel (`torch.FloatTensor`, *optional*):
              FIR filter of the shape `[firH, firW]` or `[firN]` (separable). The default is `[1] * factor`, which
              corresponds to nearest-neighbor upsampling.
          factor (`int`, *optional*, default to `2`):
              Integer upsampling factor.
          gain (`float`, *optional*, default to `1.0`):
              Scaling factor for signal magnitude (default: 1.0).
  
      Returns:
          output (`torch.FloatTensor`):
              Tensor of the shape `[N, C, H * factor, W * factor]`
      """
      assert isinstance(factor, int) and factor >= 1
      if kernel is None:
          kernel = [1] * factor
  
      kernel = torch.tensor(kernel, dtype=torch.float32)
      if kernel.ndim == 1:
          kernel = torch.outer(kernel, kernel)
      kernel /= torch.sum(kernel)
  
      kernel = kernel * (gain * (factor**2))
      pad_value = kernel.shape[0] - factor
      output = upfirdn2d_native(
          hidden_states,
          kernel.to(device=hidden_states.device),
          up=factor,
          pad=((pad_value + 1) // 2 + factor - 1, pad_value // 2),
      )
      return output
  
imports: from typing import Optional, Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  from ..utils import USE_PEFT_BACKEND
  from .lora import LoRACompatibleConv
  from .normalization import RMSNorm
label_desc: Use image node_id 2 for calling a global function with example usage: upsample_2d(hidden_states, kernel, factor, gain) and returns: output