Spaces:
Running on Zero
Running on Zero
| 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) | |
| 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) | |
| 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 | |
| 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 | |
| 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 | |
| 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] | |