cra5/models/vaeformer/vaeformer.py

# Copyright (c) 2021-2022, InterDigital Communications, Inc
# All rights reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted (subject to the limitations in the disclaimer
# below) provided that the following conditions are met:

# * Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
# * Neither the name of InterDigital Communications, Inc nor the names of its
#   contributors may be used to endorse or promote products derived from this
#   software without specific prior written permission.

# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT
# NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import warnings
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import os.path as osp
from typing import Optional
from compressai.ans import BufferedRansEncoder, RansDecoder
from cra5.models.compressai.entropy_models import EntropyBottleneck, GaussianConditional
from cra5.models.compressai.layers import GDN, MaskedConv2d
from cra5.models.compressai.registry import register_model

from cra5.models.compressai.models.base import (
    SCALES_LEVELS,
    SCALES_MAX,
    SCALES_MIN,
    CompressionModel,
    get_scale_table,
)
from cra5.models.compressai.models.utils import conv, deconv
from cra5.models.vaeformer.vit_nlc import Encoder, Decoder,HyperPriorEncoder, HyperPriorDecoder
from cra5.models.vaeformer.modules.distributions import DiagonalGaussianDistribution
from cra5.models.compressai.models.google import ScaleHyperprior
from collections import OrderedDict
        
__all__ = [
    "CompressionModel",
    "FactorizedPrior",
    "FactorizedPriorReLU",
    "ScaleHyperprior",
    "MeanScaleHyperprior",
    "JointAutoregressiveHierarchicalPriors",
    "get_scale_table",
    "SCALES_MIN",
    "SCALES_MAX",
    "SCALES_LEVELS",
    "VaritionInVaration_CNN_Prior"
]


class VAEformer(CompressionModel):
    """
    Args:
    N (int): Number of channels
    M (int): Number of channels in the expansion layers (last layer of the
    encoder and last layer of the hyperprior decoder)
    """

    def __init__(self, 
                 model_version,
                 embed_dim=None, 
                 z_channels=None,
                 y_channels=None,
                 sample_posterior=None, 
                 pretrained_vae=None, 
                 frozen_encoder=None, 
                 ddconfig=None, 
                 priorconfig=None,
                 rate_distortion_loss=None, 
                 kl_loss=None, 
                 ignore_keys:list=[], 
                 lower_dim= False,
                 **kwargs):
        if model_version == 268:
            embed_dim=256
            z_channels=256
            y_channels=1024
            lower_dim=True
            sample_posterior =False
            pretrained_vae = None #'./exp/comp/era5_autoencoder_ps10_159v/iter_150000.pth',
            frozen_encoder=False
            ddconfig=dict(
                arch = 'vit_large',
                pretrained_model = '',
                patch_size=(11,10),
                patch_stride=(10,10),
                in_chans=268,
                out_chans=268,
                kwargs=dict(
                    z_dim =  None,
                    learnable_pos= True,
                    window= True,
                    window_size = [(24, 24), (12, 48), (48, 12)],
                    interval = 4,
                    drop_path_rate= 0.,
                    round_padding= True,
                    pad_attn_mask= True , # to_do: ablation
                    test_pos_mode= 'learnable_simple_interpolate', # to_do: ablation
                    lms_checkpoint_train= True,
                    img_size= (721, 1440)
                ),
            )
            priorconfig = dict(
                pretrained_model = '', # '../PretrainedModels/maevit/mae_pretrain_vit_large.pth',
                patch_size=(4,4),
                in_chans=256,
                out_chans=256,
                kwargs=dict(
                    z_dim = 256,
                    embed_dim=360,
                    depth=8,
                    num_heads=5,
                    interval=1,
                    learnable_pos= True,
                    window= False,
                    drop_path_rate= 0.,
                    round_padding= True,
                    pad_attn_mask= True , # to_do: ablation
                    test_pos_mode= 'learnable_simple_interpolate', # to_do: ablation
                    lms_checkpoint_train= False,
                    img_size= (72,144)
                ),
            )
        super().__init__(**kwargs)
        self.sample_posterior = sample_posterior
        self.lower_dim = lower_dim
        self.frozen_encoder = frozen_encoder
        
        self.entropy_bottleneck = EntropyBottleneck(z_channels)
        
        self.g_a = Encoder(**ddconfig)
        self.g_s = Decoder(**ddconfig)
        
        if  self.lower_dim:
            self.quant_conv = torch.nn.Conv2d(2*y_channels, 2*embed_dim, 1)
            self.post_quant_conv = torch.nn.Conv2d(embed_dim, y_channels, 1)
            
        self.h_a = HyperPriorEncoder(**priorconfig)
        self.h_s = HyperPriorDecoder(**priorconfig)
        self.gaussian_conditional = GaussianConditional(None)
        
        if rate_distortion_loss is not None:
            self.criterion = MODELS.build(rate_distortion_loss)
        if kl_loss is not None:
            self.kl_loss = MODELS.build(kl_loss)
        if pretrained_vae is not None:
            self.init_from_ckpt(pretrained_vae, ignore_keys=ignore_keys)

    @classmethod
    def from_state_dict(cls, state_dict):
        """Return a new model instance from `state_dict`."""

        variable_num = state_dict["backbone.g_a.patch_embed.proj.weight"].size(1)
        
        new_state_dict = OrderedDict()
        for k, v in state_dict.items():
            if 'kl_loss.logvar' not in k:
                new_state_dict[k.replace("backbone.", "")] = v
            

        net = cls(variable_num)
        # net.update(force=True)
        
        net.load_state_dict(new_state_dict)
        
        return net

    def init_from_ckpt(self, ckpt, ignore_keys=list()):
        last_saved: Optional[str]
        if isinstance(ckpt, str):
            if ckpt.endswith('.pth'):
                last_saved = ckpt
            else:
                save_file = osp.join(ckpt, 'last_checkpoint')

                if osp.exists(save_file):
                    with open(save_file) as f:
                        last_saved = f.read().strip()
                else:
                    raise ValueError(f"You do not have a saved checkpoint to restore, "
                                    f"please set the load path: {ckpt} as None in config file")

            sd = torch.load(last_saved, map_location="cpu")["state_dict"]
        else:
            sd=ckpt
            
        ga_state_dict = OrderedDict()
        gs_state_dict = OrderedDict()
        quant_conv_state = OrderedDict()
        post_quant_conv_state = OrderedDict()
        
        loss_state_dict = OrderedDict()

        for k, v in sd.items(): #if k in model_dict.keys()
            skip = [ True  for ik in ignore_keys if k.startswith(ik)]
            if len(skip) > 0:
                print("Deleting key {} from state_dict.".format(k))
                continue
            if 'encoder' in k:
                ga_state_dict[k.replace("backbone.encoder.", "")] = v
            if 'decoder' in k:
                gs_state_dict[k.replace("backbone.decoder.", "")] = v
            if 'logvar' in k:
                loss_state_dict[k.replace("backbone.loss.", "")] = v
            if 'quant_conv' in k and 'post_quant_conv' not in k:
                quant_conv_state[k.replace("backbone.quant_conv.", "")] = v
            if 'post_quant_conv' in k:
                post_quant_conv_state[k.replace("backbone.post_quant_conv.", "")] = v

        self.g_a.load_state_dict(ga_state_dict, strict=True)
        self.g_s.load_state_dict(gs_state_dict, strict=True)
        self.quant_conv.load_state_dict(quant_conv_state,strict=True)
        self.post_quant_conv.load_state_dict(post_quant_conv_state, strict=True)
        self.kl_loss.load_state_dict(loss_state_dict, strict=True)
        if self.frozen_encoder:
            for param in self.g_a.parameters():
                param.requires_grad = False


        print(f"Restored from {last_saved}, and make the frozen_encoder as {self.frozen_encoder}" )
    
    @property
    def downsampling_factor(self) -> int:
        return 2 ** (4 + 2)

    def training_step(self, inputs, batch_idx, optimizer_idx):
        out_net = self(inputs)
        out_criterion = self.criterion(out_net, inputs)
        discloss = self.kl_loss(inputs, out_net['x_hat'], out_net['posterior'],
                                optimizer_idx, 0,
                                last_layer=self.get_last_layer(), split="train")


        return  {**discloss, **out_criterion, "aux_loss":self.aux_loss()}
    def prediction(self, inputs):

        t1 = time.time()
        out = self.compress(inputs)
        t2 =  time.time()
        x_hat = self.decompress(out['strings'], out['shape'])
        t3 =  time.time()


        return {
            **x_hat,
            "strings":out['strings'],
            "z_shape":out['shape'],
            'x_shape':inputs.shape,
            'encoding_time':(t2-t1)/inputs.size(0),
            'decoding_time':(t3-t2)/inputs.size(0)}


    def encode_latent(self, x, type='quantized'):
        moments = self.g_a(x)
        posterior = None
        if self.lower_dim:
            moments = self.quant_conv(moments)

        posterior = DiagonalGaussianDistribution(moments)
        if self.sample_posterior:
            y = posterior.sample()
        else:
            y = posterior.mode()
            
        if type == "quantized":
            z = self.h_a(y.detach())
            z_hat, z_likelihoods = self.entropy_bottleneck(z)
            gaussian_params = self.h_s(z_hat)
            scales_hat, means_hat = gaussian_params.chunk(2, 1)
            y_hat, y_likelihoods = self.gaussian_conditional(y, scales_hat, means=means_hat)
            return y, y_hat, y_likelihoods
        else:
            return y, None, None
        
    def decode_latent(self, y, type='quantized'):
        if self.lower_dim:
            y_hat = self.post_quant_conv(y)

        x_hat = self.g_s(y_hat)
        
        return x_hat
         
    def forward(self, x):
        moments = self.g_a(x)
        posterior = None
        if self.lower_dim:
            moments = self.quant_conv(moments)

        posterior = DiagonalGaussianDistribution(moments)
        if self.sample_posterior:
            y = posterior.sample()
        else:
            y = posterior.mode()
        
        z = self.h_a(y.detach())

        z_hat, z_likelihoods = self.entropy_bottleneck(z)
        gaussian_params = self.h_s(z_hat)
        scales_hat, means_hat = gaussian_params.chunk(2, 1)
        y_hat, y_likelihoods = self.gaussian_conditional(y, scales_hat, means=means_hat)

        
        
        if self.lower_dim:
            y_hat = self.post_quant_conv(y_hat)

        x_hat = self.g_s(y_hat)

        return {

            "x_hat": x_hat,
            "likelihoods": {"y": y_likelihoods, "z": z_likelihoods},
            "posterior": posterior
        }
    def compress_from_latent(self, 
                 y):
        
        z = self.h_a(y)

        z_strings = self.entropy_bottleneck.compress(z)
        z_hat = self.entropy_bottleneck.decompress(z_strings, z.size()[-2:])

        gaussian_params = self.h_s(z_hat)
        scales_hat, means_hat = gaussian_params.chunk(2, 1)
        indexes = self.gaussian_conditional.build_indexes(scales_hat)

        y_strings = self.gaussian_conditional.compress(y, indexes, means=means_hat)

        return {"strings": [y_strings, z_strings], "z_shape": z.size()[-2:]}

    def compress(self, 
                 x,
             ):
        moments = self.g_a(x)
        if self.lower_dim:
            moments = self.quant_conv(moments)
        
        posterior = DiagonalGaussianDistribution(moments)
        if self.sample_posterior:
            y = posterior.sample()
        else:
            y = posterior.mode()            

        z = self.h_a(y)

        z_strings = self.entropy_bottleneck.compress(z)
        z_hat = self.entropy_bottleneck.decompress(z_strings, z.size()[-2:])

        gaussian_params = self.h_s(z_hat)
        scales_hat, means_hat = gaussian_params.chunk(2, 1)
        indexes = self.gaussian_conditional.build_indexes(scales_hat)

        y_strings = self.gaussian_conditional.compress(y, indexes, means=means_hat)

        # bpp = sum(len(s[0]) for s in [y_strings, z_strings]) * 8.0 #/ num_pixels

        return {"strings": [y_strings, z_strings], "z_shape": z.size()[-2:]}

    def decompress(self, 
                   strings, 
                   shape,
                   return_format: str='reconstructed'):
        assert isinstance(strings, list) and len(strings) == 2
        z_hat = self.entropy_bottleneck.decompress(strings[1], shape)
        gaussian_params = self.h_s(z_hat)
        scales_hat, means_hat = gaussian_params.chunk(2, 1)
        indexes = self.gaussian_conditional.build_indexes(scales_hat)
        y_hat = self.gaussian_conditional.decompress(
            strings[0], indexes, means=means_hat
        )
        if return_format=='latent':
            return y_hat
        
        if self.lower_dim:
            y_hat = self.post_quant_conv(y_hat)
            

        
        x_hat = self.g_s(y_hat) #.clamp_(0, 1)

        return {"x_hat": x_hat}

    def get_last_layer(self):
        return self.g_s.final.weight