src/model/encoder/vae/vae.py

import math
import struct

import torch
import torch.nn as nn
import torch.nn.functional as F

from . import common
from . import coding
from .entropy_coding import DiscretizedGaussian

from .config import Config


one_div_sqrt2 = 1 / math.sqrt(2)
half_log2pi = 0.5 * math.log(2 * math.pi)


@torch.autocast('cuda', enabled=False)
def gaussian_log_prob_mass(mean: torch.Tensor, scale: torch.Tensor, x: torch.Tensor, prob_clamp: float = 1e-6):
    mean = mean.float()
    scale = scale.float()
    x = x.float()
    scale_sqrt2_r = scale.reciprocal() * one_div_sqrt2
    mean_sub_x = mean - x
    prob_mass = 0.5 * (
        (mean_sub_x - 0.5).mul_(scale_sqrt2_r).erfc_() -
        (mean_sub_x + 0.5).mul_(scale_sqrt2_r).erfc_())
    mask = prob_mass > prob_clamp
    log_prob_mass = prob_mass.clamp_min(1e-8).log_()
    normal_log_prob = -(mean_sub_x * scale_sqrt2_r).square_() - scale.log() - half_log2pi
    log_prob = torch.where(mask, log_prob_mass, normal_log_prob)
    return log_prob


class VRLVBlockBase(nn.Module):
    def __init__(self, width, z_dim, enc_key, enc_width, embed_dim=256, kernel_size=7, mlp_ratio: float = 2.0):
        super().__init__()
        self.in_channels = width
        self.out_channels = width
        self.enc_key = enc_key

        block = common.ConvNeXtBlockAdaLN
        self.resnet_front = block(width, embed_dim, kernel_size=kernel_size, mlp_ratio=mlp_ratio)
        self.resnet_end = block(width, embed_dim, kernel_size=kernel_size, mlp_ratio=mlp_ratio)
        self.posterior0 = block(enc_width, embed_dim, kernel_size=kernel_size)
        self.posterior1 = block(width, embed_dim, kernel_size=kernel_size)
        self.posterior2 = block(width, embed_dim, kernel_size=kernel_size)
        self.post_merge = common.conv_k1s1(width + enc_width, width)
        self.posterior = common.conv_k3s1(width, z_dim)
        self.z_proj = common.conv_k1s1(z_dim, width)
        self.prior = common.conv_k1s1(width, z_dim * 2)

        self.discrete_gaussian = DiscretizedGaussian()
        self.is_latent_block = True

    def transform_prior(self, feature, lmb_embedding):
        feature = self.resnet_front(feature, lmb_embedding)
        pm, plogv = self.prior(feature).chunk(2, dim=1)
        with torch.autocast(device_type=plogv.device.type, enabled=False):
            # Keep probability scale construction in FP32 under mixed precision.
            pv = F.softplus(plogv.float() + 2.3).sub(2.3).clip(max=4.6).exp()
        return feature, pm, pv

    def transform_posterior(self, feature, enc_feature, lmb_embedding):
        enc_feature = self.posterior0(enc_feature, lmb_embedding)
        feature = self.posterior1(feature, lmb_embedding)
        merged = torch.cat([feature, enc_feature], dim=1)
        merged = self.post_merge(merged)
        merged = self.posterior2(merged, lmb_embedding)
        qm = self.posterior(merged)
        return qm

    def fuse_feature_and_z(self, feature, z):
        # add the new information carried by z to the feature
        return feature + self.z_proj(z)

    def forward(self, feature, lmb_embedding, enc_feature=None, mode='trainval',
                get_latent=False, strings=None):
        feature, pm, pv = self.transform_prior(feature, lmb_embedding)

        additional = dict()
        if mode == 'trainval':
            qm = self.transform_posterior(feature, enc_feature, lmb_embedding)
            if self.training:
                z = qm + (torch.rand_like(qm) - 0.5)
                kl = -gaussian_log_prob_mass(pm, pv, x=z)
            else:
                z, probs = self.discrete_gaussian(qm, scales=pv, means=pm)
                kl = probs.log().neg_()
            additional['kl'] = kl
        elif mode == 'compress':
            qm = self.transform_posterior(feature, enc_feature, lmb_embedding)
            indexes = self.discrete_gaussian.build_indexes(pv)
            strings = self.discrete_gaussian.compress(qm, indexes, means=pm)[0]
            z = self.discrete_gaussian.quantize(qm, mode='dequantize', means=pm)
            additional['strings'] = strings
        elif mode == 'decompress':
            assert strings is not None
            indexes = self.discrete_gaussian.build_indexes(pv)
            z = self.discrete_gaussian.decompress([strings], indexes, means=pm)
        else:
            raise ValueError(f'Unknown mode={mode}')

        feature = self.fuse_feature_and_z(feature, z)
        feature = self.resnet_end(feature, lmb_embedding)
        if get_latent:
            additional['z'] = z.detach()
        return feature, additional

    def update(self):
        self.discrete_gaussian.update()


class VariableRateLossyVAE(nn.Module):
    log2_e = math.log2(math.e)

    def __init__(self, in_ch, cfg: Config):
        super().__init__()
        self.cfg = cfg

        enc_dims = cfg.enc_channels
        res_block = common.ConvNeXtBlockAdaLN
        embed_dim = cfg.lmb_embed_dim[1]
        enc_blocks = [
            # 64x64
            common.patch_downsample(in_ch, enc_dims[0], rate=4),
            # 16x16
            *[res_block(enc_dims[0], embed_dim, kernel_size=7) for _ in range(3)],  # original: 6
            res_block(enc_dims[0]),
            common.patch_downsample(enc_dims[0], enc_dims[1]),
            # 8x8
            *[res_block(enc_dims[1], embed_dim, kernel_size=7) for _ in range(3)],  # original: 6
            common.SetKey('enc_s8'),
            res_block(enc_dims[1], embed_dim),
            common.patch_downsample(enc_dims[1], enc_dims[2]),
            # 4x4
            *[res_block(enc_dims[2], embed_dim, kernel_size=5) for _ in range(3)],  # original: 6
            common.SetKey('enc_s16'),
            res_block(enc_dims[2], embed_dim),
            common.patch_downsample(enc_dims[2], enc_dims[3]),
            # 2x2
            *[res_block(enc_dims[3], embed_dim, kernel_size=3) for _ in range(2)],  # original: 4
            common.SetKey('enc_s32'),
            res_block(enc_dims[3], embed_dim),
            common.patch_downsample(enc_dims[3], enc_dims[4]),
            # 1x1
            *[res_block(enc_dims[4], embed_dim, kernel_size=1) for _ in range(2)],  # original: 4
            common.SetKey('enc_s64'),
        ]

        dec_dims = cfg.dec_channels
        z_dims = cfg.z_dec_channels
        dec_blocks = [
            # 1x1
            *[VRLVBlockBase(dec_dims[0], z_dims[0], enc_key='enc_s64', enc_width=enc_dims[-1], kernel_size=1,
                            mlp_ratio=4.0) for _ in range(1)],
            res_block(dec_dims[0], embed_dim, kernel_size=1, mlp_ratio=4.0),
            common.patch_upsample(dec_dims[0], dec_dims[1], rate=2),
            # 2x2
            res_block(dec_dims[1], embed_dim, kernel_size=3, mlp_ratio=3.0),
            *[VRLVBlockBase(dec_dims[1], z_dims[1], enc_key='enc_s32', enc_width=enc_dims[-2], kernel_size=3,
                            mlp_ratio=3.0) for _ in range(1)],  # original: 2
            res_block(dec_dims[1], embed_dim, kernel_size=3, mlp_ratio=3.0),
            common.patch_upsample(dec_dims[1], dec_dims[2], rate=2),
            # 4x4
            res_block(dec_dims[2], embed_dim, kernel_size=5, mlp_ratio=2.0),
            *[VRLVBlockBase(dec_dims[2], z_dims[2], enc_key='enc_s16', enc_width=enc_dims[-3], kernel_size=5,
                            mlp_ratio=2.0) for _ in range(2)],  # original: 3
            res_block(dec_dims[2], embed_dim, kernel_size=5, mlp_ratio=2.0),
            common.patch_upsample(dec_dims[2], dec_dims[3], rate=2),
            # 8x8
            res_block(dec_dims[3], kernel_size=7, mlp_ratio=1.75),
            *[VRLVBlockBase(dec_dims[3], z_dims[3], enc_key='enc_s8', enc_width=enc_dims[-4], kernel_size=7,
                            mlp_ratio=1.75) for _ in range(2)],  # original: 3
            common.CompressionStopFlag(),  # no need to execute remaining blocks when compressing
            res_block(dec_dims[3], kernel_size=7, mlp_ratio=1.75),
            common.patch_upsample(dec_dims[3], dec_dims[4], rate=2),
            # 16x16
            *[res_block(dec_dims[4], kernel_size=7, mlp_ratio=1.5) for _ in range(4)],  # original: 8
            common.patch_upsample(dec_dims[4], in_ch, rate=4)
        ]
        self.encoder = common.FeatureExtractorWithEmbedding(enc_blocks)
        self.dec_blocks = nn.ModuleList(dec_blocks)
        self.bias = nn.Parameter(torch.zeros(1, self.dec_blocks[0].in_channels, 1, 1))
        self.lmb_embedding = nn.Sequential(
            nn.Linear(cfg.lmb_embed_dim[0], cfg.lmb_embed_dim[1]),
            nn.GELU(),
            nn.Linear(cfg.lmb_embed_dim[1], cfg.lmb_embed_dim[1]),
            nn.GELU()
        )

        self.register_buffer(
            'sinusoidal_freqs',
            torch.pow(self.cfg.sin_period, 1 - torch.linspace(0, 1, steps=(cfg.lmb_embed_dim[0] // 2))).reshape(1, -1),
            persistent=False)
        self.sinusoidal_freqs: torch.Tensor
        self.num_latent_blocks = sum(
            1 for block in self.dec_blocks if getattr(block, "is_latent_block", False)
        )

        self.compress_mode(True)

    @torch.no_grad()
    def sample_lmb(self, n, device):
        low, high = self.cfg.lmb_range
        if low == high:
            return torch.full((n,), float(low), device=device)
        p = 3.0
        low, high = math.pow(low, 1/p), math.pow(high, 1/p)  # transformed space
        lmb = low + (high-low) * torch.rand(n, device=device)
        lmb.pow_(p)
        return lmb

    def expand_to_tensor(self, input_, n, device):
        if isinstance(input_, (float, int)):
            input_ = torch.full(size=(n,), fill_value=float(input_), device=device)
        assert input_.shape == (n,), f'{input_=}, {input_.shape=}'
        return input_

    def get_lmb_embedding(self, lmb, n, device):
        lmb = self.expand_to_tensor(lmb, n, device)
        lmb_limit = self.cfg.lmb_limit
        scaled = torch.log(lmb) / math.log(lmb_limit)
        args = scaled.reshape(-1, 1) * self.sinusoidal_freqs
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        embedding = self.lmb_embedding(embedding)
        return embedding

    def get_bias(self, bhw_repeat=(1, 1, 1)):
        nB, nH, nW = bhw_repeat
        feature = self.bias.expand(nB, -1, nH, nW)
        return feature

    def forward_end2end(self, features: torch.Tensor, lmb, mode='trainval', get_latent=False):
        x = features
        # ================ get lambda embedding ================
        emb = self.get_lmb_embedding(lmb, features.shape[0], x.device)
        # ================ Forward pass ================
        enc_features = self.encoder(x, emb)[1]
        batch_size, _, height, width = x.shape
        feature = self.get_bias(bhw_repeat=(batch_size, height//self.cfg.max_stride, width//self.cfg.max_stride))
        lv_block_results = []  # all latent variable block results
        for i, block in enumerate(self.dec_blocks):
            if getattr(block, 'is_latent_block', False):
                f_enc = enc_features[block.enc_key]
                feature, stats = block(feature, emb, enc_feature=f_enc, mode=mode, get_latent=get_latent)
                lv_block_results.append(stats)
            elif getattr(block, 'requires_embedding', False):
                feature = block(feature, emb)
            elif isinstance(block, common.CompressionStopFlag) and (mode == 'compress'):
                # no need to execute remaining blocks when compressing
                return lv_block_results
            else:
                feature = block(feature)
        return feature, lv_block_results

    def forward(self, features, lmb=None):
        if self.training:
            return self.train_forward(features, lmb=lmb)
        return self.test_forward(features)

    def train_forward(self, features, lmb=None):
        height, width = features.shape[2:]

        if lmb is None:
            lmb = self.sample_lmb(features.shape[0], features.device)
        else:
            lmb = self.expand_to_tensor(lmb, features.shape[0], features.device)
        x_hat, stats_all = self.forward_end2end(features, lmb)

        kl_loss = sum([stat['kl'].sum(dim=(1, 2, 3)) for stat in stats_all]) * (
            1 / (height * width)
        )

        return kl_loss * self.log2_e, x_hat, lmb

    @torch.compiler.disable
    @torch.no_grad()
    def test_forward(self, features):
        bitstreams = []
        reconstructions = []
        num_bytes = []

        for i in range(features.shape[0]):
            feature_i = features[i : i + 1]
            strings = self.compress(feature_i)
            feature_hat = self.decompress(strings)
            bitstreams.append(strings)
            reconstructions.append(feature_hat)
            num_bytes.append(len(strings))

        features_hat = torch.cat(reconstructions, dim=0)
        num_bytes_tensor = torch.tensor(
            num_bytes, device=features.device, dtype=torch.float32
        )
        return bitstreams, features_hat, num_bytes_tensor

    def train(self, mode: bool = True):
        if mode is False:
            self.compress_mode(True)
        return super(VariableRateLossyVAE, self).train(mode)

    def compress_mode(self, mode=True):
        if mode:
            for block in self.dec_blocks:
                if hasattr(block, 'update'):
                    block.update()

    def compress_once(self, features, padded_features, lmb):
        lv_block_results = self.forward_end2end(
            padded_features, lmb, mode='compress', get_latent=True
        )
        all_lv_strings = [res['strings'] for res in lv_block_results]
        strings = struct.pack('f', torch.tensor(lmb, dtype=torch.float32).item()) + \
            struct.pack('2H', features.shape[2], features.shape[3]) + \
            coding.concat_bytes_list(all_lv_strings)
        return strings, lv_block_results

    @torch.no_grad()
    def compress(self, features):
        assert (
            features.shape[0] == 1
        ), "Feature codec bitstream export currently supports batch size 1."
        padded_features = coding.pad_tensor_divisible_by(features, div=self.cfg.max_stride)
        lmb = self.cfg.lmb
        strings, _ = self.compress_once(features, padded_features, lmb)
        return strings

    @torch.no_grad()
    def decompress(self, strings):
        (lmb,), strings = struct.unpack("f", strings[:4]), strings[4:]
        (height, width), strings = struct.unpack("2H", strings[:4]), strings[4:]
        downsampled_height = math.ceil(height / self.cfg.max_stride)
        downsampled_width = math.ceil(width / self.cfg.max_stride)
        all_lv_strings = coding.split_bytes_list(strings, self.num_latent_blocks)

        device = next(self.parameters()).device
        lmb_embedding = self.get_lmb_embedding(lmb, 1, device)
        feature = self.get_bias(
            bhw_repeat=(1, downsampled_height, downsampled_width)
        )

        str_i = 0
        for block in self.dec_blocks:
            if getattr(block, "is_latent_block", False):
                feature, _ = block(
                    feature,
                    lmb_embedding,
                    mode="decompress",
                    strings=all_lv_strings[str_i],
                )
                str_i += 1
            elif getattr(block, "requires_embedding", False):
                feature = block(feature, lmb_embedding)
            else:
                feature = block(feature)

        assert str_i == len(all_lv_strings), f"str_i={str_i}, len={len(all_lv_strings)}"
        return feature[:, :, :height, :width]