content: import math
  import torch
  import torch.nn as nn
  from ..attention import ImgToTriplaneTransformer
  import math
  from einops import rearrange
  
  
  class ImgToTriplaneModel(nn.Module):
      """
      The full UNet model with attention and timestep embedding.
      :param in_channels: channels in the input Tensor.
      :param model_channels: base channel count for the model.
      :param out_channels: channels in the output Tensor.
      :param num_res_blocks: number of residual blocks per downsample.
      :param attention_resolutions: a collection of downsample rates at which
          attention will take place. May be a set, list, or tuple.
          For example, if this contains 4, then at 4x downsampling, attention
          will be used.
      :param dropout: the dropout probability.
      :param channel_mult: channel multiplier for each level of the UNet.
      :param conv_resample: if True, use learned convolutions for upsampling and
          downsampling.
      :param dims: determines if the signal is 1D, 2D, or 3D.
      :param num_classes: if specified (as an int), then this model will be
          class-conditional with `num_classes` classes.
      :param use_checkpoint: use gradient checkpointing to reduce memory usage.
      :param num_heads: the number of attention heads in each attention layer.
      :param num_heads_channels: if specified, ignore num_heads and instead use
                                 a fixed channel width per attention head.
      :param num_heads_upsample: works with num_heads to set a different number
                                 of heads for upsampling. Deprecated.
      :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
      :param resblock_updown: use residual blocks for up/downsampling.
      :param use_new_attention_order: use a different attention pattern for potentially
                                      increased efficiency.
      """
  
      def __init__(
          self,
          pos_emb_size=32,
          pos_emb_dim=1024,
          cam_cond_dim=20,
          n_heads=16,
          d_head=64,
          depth=16,
          context_dim=768,
          triplane_dim=80,
          upsample_time=1,
          use_fp16=False,
          use_bf16=True,
      ):
          super().__init__()
  
          self.pos_emb_size = pos_emb_size
          self.pos_emb_dim  = pos_emb_dim
  
          # init embedding
          self.pos_emb = nn.Parameter(torch.zeros(1, 3 * pos_emb_size * pos_emb_size, pos_emb_dim))
          # TODO initialize pos_emb with a Gaussian random of zero-mean and std of 1/sqrt(1024).
  
          # build image to triplane decoder
          self.img_to_triplane_decoder = ImgToTriplaneTransformer(
              query_dim=pos_emb_dim, n_heads=n_heads,
              d_head=d_head, depth=depth, context_dim=context_dim,
              triplane_size=pos_emb_size, 
          )
  
          self.is_conv_upsampler = False
          # build upsampler
          self.triplane_dim = triplane_dim
          if self.is_conv_upsampler:
              upsamplers = []
              for i in range(upsample_time):
                  if i == 0:
                      upsampler = nn.ConvTranspose2d(in_channels=pos_emb_dim, out_channels=triplane_dim,
                                              kernel_size=2, stride=2,
                                              padding=0, output_padding=0)
                      upsamplers.append(upsampler)
                  else:
                      upsampler = nn.ConvTranspose2d(in_channels=triplane_dim, out_channels=triplane_dim,
                                              kernel_size=2, stride=2,
                                              padding=0, output_padding=0)
                      upsamplers.append(upsampler)
              if upsamplers:
                  self.upsampler = nn.Sequential(*upsamplers)
              else:
                  self.upsampler = nn.Conv2d(in_channels=pos_emb_dim, out_channels=triplane_dim,
                                              kernel_size=3, stride=1, padding=1)
          else:
              self.upsample_ratio = 4
              self.upsampler = nn.Linear(in_features=pos_emb_dim, out_features=triplane_dim*(self.upsample_ratio**2))
          
  
  
      def forward(self, x, cam_cond=None, **kwargs):
          """
          Apply the model to an input batch.
          :param x: an [N x C x ...] Tensor of inputs.
          :param timesteps: a 1-D batch of timesteps.
          :param context: conditioning plugged in via crossattn
          :param y: an [N] Tensor of labels, if class-conditional.
          :return: an [N x C x ...] Tensor of outputs.
          """
  
          B = x.shape[0]
          h = self.pos_emb.expand(B, -1, -1)
          context = x
  
          h = self.img_to_triplane_decoder(h, context=context)
  
          h = h.view(B * 3, self.pos_emb_size, self.pos_emb_size, self.pos_emb_dim)
          if self.is_conv_upsampler:
              h = rearrange(h, 'b h w c -> b c h w')
              h = self.upsampler(h)
              h = rearrange(h, '(b d) c h w-> b d c h w', d=3)
              h = h.type(x.dtype)
              return h 
          else:
              h = self.upsampler(h) #[b, h, w, triplane_dim*4]
              b, height, width, _ = h.shape
              h = h.view(b, height, width, self.triplane_dim, self.upsample_ratio, self.upsample_ratio) #[b, h, w, triplane_dim, 2, 2]
              h = h.permute(0,3,1,4,2,5).contiguous() #[b, triplane_dim, h, 2, w, 2]
              h = h.view(b, self.triplane_dim, height*self.upsample_ratio, width*self.upsample_ratio)
              h = rearrange(h, '(b d) c h w-> b d c h w', d=3)
              h = h.type(x.dtype)
              return h 
  