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]