content: '''
  Codes are from:
  https://github.com/jaxony/unet-pytorch/blob/master/model.py
  '''
  
  import torch
  import torch.nn as nn
  from diffusers import UNet2DModel
  import einops
  class UNetPP(nn.Module):
      '''
          Wrapper for UNet in diffusers
      '''
      def __init__(self, in_channels):
          super(UNetPP, self).__init__()
          self.in_channels = in_channels
          self.unet = UNet2DModel(
                  sample_size=[256, 256*3],
                  in_channels=in_channels,
                  out_channels=32,
                  layers_per_block=2,
                  block_out_channels=(64, 128, 128, 128*2, 128*2, 128*4, 128*4),
                  down_block_types=(
                      "DownBlock2D",
                      "DownBlock2D",
                      "DownBlock2D",
                      "AttnDownBlock2D",
                      "AttnDownBlock2D",
                      "AttnDownBlock2D",
                      "DownBlock2D",
                  ),
                  up_block_types=(
                      "UpBlock2D",
                      "AttnUpBlock2D",
                      "AttnUpBlock2D",
                      "AttnUpBlock2D",
                      "UpBlock2D",
                      "UpBlock2D",
                      "UpBlock2D",
                  ),
              )
               
          self.unet.enable_xformers_memory_efficient_attention()    
          if in_channels > 12:
              self.learned_plane = torch.nn.parameter.Parameter(torch.zeros([1,in_channels-12,256,256*3]))
  
      def forward(self, x, t=256):
          learned_plane = self.learned_plane
          if x.shape[1] < self.in_channels:
              learned_plane = einops.repeat(learned_plane, '1 C H W -> B C H W', B=x.shape[0]).to(x.device)
              x = torch.cat([x, learned_plane], dim = 1)
          return self.unet(x, t).sample
  
  