'''
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