 :128, 128:].reshape(-1), x_wt[:, :, :128, 128:].reshape(-1), reduction='sum')
        BCE_wt /= x_wt_hat.numel()
        
        logvar = torch.log(var)
        KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) * kl_weight
        KLD /= x.shape[0] * 3 * 512 * 512

        return BCE + BCE_wt + KLD, BCE + BCE_wt, KLD

    def set_device(self, device):
        self.device = device
        if 'cuda' in self.device:
            self.cuda = True
    
    def set_filters(self, filters):
        self.iwt = IWT(iwt=iwt, num_iwt=self.num_iwt)
        self.iwt.set_filters(filters)


class IWTVAE_512_Mask_1(nn.Module):
    def __init__(self, image_channels=3, z_dim=500, num_iwt=2):
        super(IWTVAE_512_Mask_1, self).__init__()
        # Resolution of images (512 x 512)
        self.res = 512
        self.device = None
        self.cuda = False
        
        self.z_dim = z_dim
        self.num_iwt = num_iwt
        self.leakyrelu = nn.LeakyReLU(0.2)
        self.sigmoid = nn.Sigmoid()

        # Z Encoder - Decoder                                                                [b, 3, 512, 512]
        self.e1 = nn.Conv2d(3, 64, 4, stride=2, padding=1, bias=True, padding_mode='zeros') #[b, 64, 256, 256]
        weights_init(self.e1)
        self.instance_norm_e1 = nn.BatchNorm2d(num_features=64)

        self.e2 = nn.Conv2d(64, 128, 4, stride=2, padding=1, bias=True, padding_mode='zeros') #[b, 128, 128, 128]
        weights_init(self.e2)
        self.instance_norm_e2 = nn.BatchNorm2d(num_features=128)

        self.e3 = nn.Conv2d(128, 256, 4, stride=2, padding=1, bias=True, padding_mode='zeros') #[b, 256, 64, 64]
        weights_init(self.e3)
        self.instance_norm_e3 = nn.BatchNorm2d(num_features=256)

        self.e4 = nn.Conv2d(256, 512, 4, stride=2, padding=1, bias=True, padding_mode='zeros') #[b, 512, 32, 32]
        weights_init(self.e4)
        self.instance_norm_e4 = nn.BatchNorm2d(num_features=512)

        self.e5 = nn.Conv2d(512, 1024, 4, stride=2, padding=1, bias=True, padding_mode='zeros') #[b, 512, 16, 16]
        weights_init(self.e5)
        self.instance_norm_e5 = nn.BatchNorm2d(num_features=1024)

        self.e6 = nn.Conv2d(1024, 2048, 4, stride=2, padding=1, bias=True, padding_mode='zeros') #[b, 2048, 8, 8]
        weights_init(self.e6)
        self.instance_norm_e6 = nn.BatchNorm2d(num_features=2048)
        
        self.fc_enc = nn.Linear(2048 * 8 * 8, 1024)
        weights_init(self.fc_enc)
        
        self.fc_mean = nn.Linear(1024, z_dim)
        weights_init(self.fc_mean)
        
        self.fc_var = nn.Linear(1024, z_dim)
        weights_init(self.fc_var)
        
        self.fc_dec = nn.Linear(z_dim, 2048 * 8 * 8)
        weights_init(self.fc_dec)

        self.d1 = nn.ConvTranspose2d(2048, 1024, 4, stride=2, padding=1, bias=True) #[b, 1024, 16, 16]
        weights_init(self.d1)
        self.instance_norm_d1 = nn.BatchNorm2d(num_features=1024)

        self.d2 = nn.ConvTranspose2d(1024, 512, 4, stride=2, padding=1, bias=True) #[b, 512, 32, 32]
        weights_init(self.d2)
        self.instance_norm_d2 = nn.BatchNorm2d(num_features=512)

        self.d3 = nn.ConvTranspose2d(512, 256, 4, stride=2, padding=1, bias=True) #[b, 256, 64, 64]
        weights_init(self.d3)
        self.instance_norm_d3 = nn.BatchNorm2d(num_features=256)

        self.d4 = nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1, bias=True) #[b, 128, 128, 128]
        weights_init(self.d4)
        self.instance_norm_d4 = nn.BatchNorm2d(num_features=128)

        self.d5 = nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1, bias=True) #[b, 32, 256, 256]
        weights_init(self.d5)
        self.instance_norm_d5 = nn.BatchNorm2d(num_features=64)

        self.d6 = nn.ConvTranspose2d(64, 3, 4, stride=2, padding=1, bias=True) #[b, 3, 512, 512]
        weights_init(self.d6)
        self.instance_norm_d6 = nn.BatchNorm2d(num_features=3)
        
        self.iwt = None
    
      
    def encode(self, x):
        h = self.leakyrelu(self.instance_norm_e1(self.e1(x)))                       #[b, 64, 256, 256]
        h = self.leakyrelu(self.instance_norm_e2(self.e2(h)))                       #[b, 128, 128, 128]
        h = self.leakyrelu(self.instance_norm_e3(self.e3(h)))                       #[b, 256, 64, 64]
        h = self.leakyrelu(self.instance_norm_e4(self.e4(h)))                       #[b, 512, 32, 32]
        h = self.leakyrelu(self.instance_norm_e5(self.e5(h)))                       #[b, 1024, 16, 16]
        h = self.leakyrelu(self.instance_norm_e6(self.e6(h)))                       #[b, 1024, 8, 8]

        h = self.leakyrelu(self.fc_enc(h.reshape(-1,2048*8*8)))                     #[b, z_dim]

        return self.fc_mean(h), F.softplus(self.fc_var(h))                          #[b, z_dim]
    
    def reparameterize(self, mu, var):
        std = torch.sqrt(var)
        if self.cuda:
            eps = torch.FloatTensor(std.size()).normal_().to(self.device)
        else:
            eps = torch.FloatTensor(std.size()).normal_()
        eps = Variable(eps)

        return eps.mul(std).add_(mu) 
    
    def decode(self, y, z):
        h = self.leakyrelu(self.fc_dec(z))                                              #[b, 1024*8*8]
        h = self.leakyrelu(self.instance_norm_d1(self.d1(h.reshape(-1, 2048, 8, 8))))   #[b, 1024, 16, 16]
        h = self.leakyrelu(self.instance_norm_d2