| |
| |
|
|
| import torch |
| from torch import nn |
|
|
| from ..layers.cuda_inference import combine_for_reading_2x, \ |
| restore_y_2x, restore_y_2x_with_cat_after, add_and_multiply, \ |
| replicate_pad, restore_y_4x, clamp_reciprocal_with_quant |
| from .entropy_models import BitEstimator, GaussianEncoder, EntropyCoder |
|
|
|
|
| class CompressionModel(nn.Module): |
| def __init__(self, z_channel, extra_qp=0): |
| super().__init__() |
|
|
| self.z_channel = z_channel |
| self.entropy_coder = None |
| self.bit_estimator_z = BitEstimator(64 + extra_qp, z_channel) |
| self.gaussian_encoder = GaussianEncoder() |
|
|
| self.masks = {} |
| self.cuda_streams = {} |
|
|
| def get_cuda_stream(self, device, idx=0, priority=0): |
| key = f"{device}_{priority}_{idx}" |
| if key not in self.cuda_streams: |
| self.cuda_streams[key] = torch.cuda.Stream(device, priority=priority) |
| return self.cuda_streams[key] |
|
|
| @staticmethod |
| def get_qp_num(): |
| return 64 |
|
|
| @staticmethod |
| def get_padding_size(height, width, p=64): |
| new_h = (height + p - 1) // p * p |
| new_w = (width + p - 1) // p * p |
| padding_right = new_w - width |
| padding_bottom = new_h - height |
| return padding_right, padding_bottom |
|
|
| @staticmethod |
| def get_downsampled_shape(height, width, p): |
| new_h = (height + p - 1) // p * p |
| new_w = (width + p - 1) // p * p |
| return int(new_h / p + 0.5), int(new_w / p + 0.5) |
|
|
| def update(self, force_zero_thres=None): |
| self.entropy_coder = EntropyCoder() |
| self.gaussian_encoder.update(self.entropy_coder, force_zero_thres=force_zero_thres) |
| self.bit_estimator_z.update(self.entropy_coder) |
|
|
| def set_use_two_entropy_coders(self, use_two_entropy_coders): |
| self.entropy_coder.set_use_two_entropy_coders(use_two_entropy_coders) |
|
|
| def pad_for_y(self, y): |
| _, _, H, W = y.size() |
| padding_r, padding_b = self.get_padding_size(H, W, 4) |
| y_pad = replicate_pad(y, padding_b, padding_r) |
| return y_pad |
|
|
| def separate_prior(self, params, is_video=False): |
| if is_video: |
| quant_step, scales, means = params.chunk(3, 1) |
| quant_step = torch.clamp_min(quant_step, 0.5) |
| q_enc = 1. / quant_step |
| q_dec = quant_step |
| else: |
| q = params[:, :2, :, :] |
| q_enc, q_dec = (torch.sigmoid(q) * 1.5 + 0.5).chunk(2, 1) |
| scales, means = params[:, 2:, :, :].chunk(2, 1) |
| return q_enc, q_dec, scales, means |
|
|
| @staticmethod |
| def separate_prior_for_video_encoding(params, y): |
| q_dec, scales, means = params.chunk(3, 1) |
| q_dec, y = clamp_reciprocal_with_quant(q_dec, y, 0.5) |
| return y, q_dec, scales, means |
|
|
| @staticmethod |
| def separate_prior_for_video_decoding(params): |
| quant_step, scales, means = params.chunk(3, 1) |
| quant_step = torch.clamp_min(quant_step, 0.5) |
| return quant_step, scales, means |
|
|
| def process_with_mask(self, y, scales, means, mask): |
| return self.gaussian_encoder.process_with_mask(y, scales, means, mask) |
|
|
| @staticmethod |
| def get_one_mask(micro_mask, height, width, dtype, device): |
| mask = torch.tensor(micro_mask, dtype=dtype, device=device) |
| mask = mask.repeat((height + 1) // 2, (width + 1) // 2) |
| mask = mask[:height, :width] |
| mask = torch.unsqueeze(mask, 0) |
| mask = torch.unsqueeze(mask, 0) |
| return mask |
|
|
| def get_mask_4x(self, batch, channel, height, width, dtype, device): |
| curr_mask_str = f"{batch}_{channel}_{width}_{height}_4x" |
| with torch.no_grad(): |
| if curr_mask_str not in self.masks: |
| assert channel % 4 == 0 |
| m = torch.ones((batch, channel // 4, height, width), dtype=dtype, device=device) |
| m0 = self.get_one_mask(((1, 0), (0, 0)), height, width, dtype, device) |
| m1 = self.get_one_mask(((0, 1), (0, 0)), height, width, dtype, device) |
| m2 = self.get_one_mask(((0, 0), (1, 0)), height, width, dtype, device) |
| m3 = self.get_one_mask(((0, 0), (0, 1)), height, width, dtype, device) |
|
|
| mask_0 = torch.cat((m * m0, m * m1, m * m2, m * m3), dim=1) |
| mask_1 = torch.cat((m * m3, m * m2, m * m1, m * m0), dim=1) |
| mask_2 = torch.cat((m * m2, m * m3, m * m0, m * m1), dim=1) |
| mask_3 = torch.cat((m * m1, m * m0, m * m3, m * m2), dim=1) |
|
|
| self.masks[curr_mask_str] = [mask_0, mask_1, mask_2, mask_3] |
| return self.masks[curr_mask_str] |
|
|
| def get_mask_2x(self, batch, channel, height, width, dtype, device): |
| curr_mask_str = f"{batch}_{channel}_{width}_{height}_2x" |
| with torch.no_grad(): |
| if curr_mask_str not in self.masks: |
| assert channel % 2 == 0 |
| m = torch.ones((batch, channel // 2, height, width), dtype=dtype, device=device) |
| m0 = self.get_one_mask(((1, 0), (0, 1)), height, width, dtype, device) |
| m1 = self.get_one_mask(((0, 1), (1, 0)), height, width, dtype, device) |
|
|
| mask_0 = torch.cat((m * m0, m * m1), dim=1) |
| mask_1 = torch.cat((m * m1, m * m0), dim=1) |
|
|
| self.masks[curr_mask_str] = [mask_0, mask_1] |
| return self.masks[curr_mask_str] |
|
|
| @staticmethod |
| def single_part_for_writing_4x(x): |
| x0, x1, x2, x3 = x.chunk(4, 1) |
| return (x0 + x1) + (x2 + x3) |
|
|
| @staticmethod |
| def single_part_for_writing_2x(x): |
| x0, x1 = x.chunk(2, 1) |
| return x0 + x1 |
|
|
| def compress_prior_2x(self, y, common_params, y_spatial_prior): |
| y, q_dec, scales, means = self.separate_prior_for_video_encoding(common_params, y) |
| dtype = y.dtype |
| device = y.device |
| B, C, H, W = y.size() |
| mask_0, mask_1 = self.get_mask_2x(B, C, H, W, dtype, device) |
|
|
| _, y_q_0, y_hat_0, s_hat_0 = self.process_with_mask(y, scales, means, mask_0) |
| cat_params = torch.cat((y_hat_0, common_params), dim=1) |
| scales, means = y_spatial_prior(cat_params).chunk(2, 1) |
| _, y_q_1, y_hat_1, s_hat_1 = self.process_with_mask(y, scales, means, mask_1) |
|
|
| y_hat = add_and_multiply(y_hat_0, y_hat_1, q_dec) |
|
|
| y_q_w_0 = self.single_part_for_writing_2x(y_q_0) |
| y_q_w_1 = self.single_part_for_writing_2x(y_q_1) |
| s_w_0 = self.single_part_for_writing_2x(s_hat_0) |
| s_w_1 = self.single_part_for_writing_2x(s_hat_1) |
| return y_q_w_0, y_q_w_1, s_w_0, s_w_1, y_hat |
|
|
| def decompress_prior_2x(self, common_params, y_spatial_prior): |
| infos = self.decompress_prior_2x_part1(common_params) |
| y_hat = self.decompress_prior_2x_part2(common_params, y_spatial_prior, infos) |
| return y_hat |
|
|
| def decompress_prior_2x_part1(self, common_params): |
| q_dec, scales, means = self.separate_prior_for_video_decoding(common_params) |
| dtype = means.dtype |
| device = means.device |
| B, C, H, W = means.size() |
| mask_0, mask_1 = self.get_mask_2x(B, C, H, W, dtype, device) |
|
|
| scales_r = combine_for_reading_2x(scales, mask_0, inplace=False) |
| indexes, skip_cond = self.gaussian_encoder.build_indexes_decoder(scales_r) |
| self.gaussian_encoder.decode_y(indexes) |
| infos = { |
| "q_dec": q_dec, |
| "mask_0": mask_0, |
| "mask_1": mask_1, |
| "means": means, |
| "scales_r": scales_r, |
| "skip_cond": skip_cond, |
| "indexes": indexes, |
| } |
| return infos |
|
|
| def decompress_prior_2x_part2(self, common_params, y_spatial_prior, infos): |
| dtype = common_params.dtype |
| device = common_params.device |
| y_q_r = self.gaussian_encoder.get_y(infos["scales_r"].shape, |
| infos["scales_r"].numel(), |
| dtype, device, |
| infos["skip_cond"], infos["indexes"]) |
| y_hat_0, cat_params = restore_y_2x_with_cat_after(y_q_r, infos["means"], infos["mask_0"], |
| common_params) |
| scales, means = y_spatial_prior(cat_params).chunk(2, 1) |
| scales_r = combine_for_reading_2x(scales, infos["mask_1"], inplace=True) |
| y_q_r = self.gaussian_encoder.decode_and_get_y(scales_r, dtype, device) |
| y_hat_1 = restore_y_2x(y_q_r, means, infos["mask_1"]) |
|
|
| y_hat = add_and_multiply(y_hat_0, y_hat_1, infos["q_dec"]) |
| return y_hat |
|
|
| def compress_prior_4x(self, y, common_params, y_spatial_prior_reduction, |
| y_spatial_prior_adaptor_1, y_spatial_prior_adaptor_2, |
| y_spatial_prior_adaptor_3, y_spatial_prior): |
| ''' |
| y_0 means split in channel, the 0/4 quater |
| y_1 means split in channel, the 1/4 quater |
| y_2 means split in channel, the 2/4 quater |
| y_3 means split in channel, the 3/4 quater |
| y_?_0, means multiply with mask_0 |
| y_?_1, means multiply with mask_1 |
| y_?_2, means multiply with mask_2 |
| y_?_3, means multiply with mask_3 |
| ''' |
| q_enc, q_dec, scales, means = self.separate_prior(common_params, False) |
| common_params = y_spatial_prior_reduction(common_params) |
| dtype = y.dtype |
| device = y.device |
| B, C, H, W = y.size() |
| mask_0, mask_1, mask_2, mask_3 = self.get_mask_4x(B, C, H, W, dtype, device) |
|
|
| y = y * q_enc |
|
|
| _, y_q_0, y_hat_0, s_hat_0 = self.process_with_mask(y, scales, means, mask_0) |
|
|
| y_hat_so_far = y_hat_0 |
| params = torch.cat((y_hat_so_far, common_params), dim=1) |
| scales, means = y_spatial_prior(y_spatial_prior_adaptor_1(params)).chunk(2, 1) |
| _, y_q_1, y_hat_1, s_hat_1 = self.process_with_mask(y, scales, means, mask_1) |
|
|
| y_hat_so_far = y_hat_so_far + y_hat_1 |
| params = torch.cat((y_hat_so_far, common_params), dim=1) |
| scales, means = y_spatial_prior(y_spatial_prior_adaptor_2(params)).chunk(2, 1) |
| _, y_q_2, y_hat_2, s_hat_2 = self.process_with_mask(y, scales, means, mask_2) |
|
|
| y_hat_so_far = y_hat_so_far + y_hat_2 |
| params = torch.cat((y_hat_so_far, common_params), dim=1) |
| scales, means = y_spatial_prior(y_spatial_prior_adaptor_3(params)).chunk(2, 1) |
| _, y_q_3, y_hat_3, s_hat_3 = self.process_with_mask(y, scales, means, mask_3) |
|
|
| y_hat = y_hat_so_far + y_hat_3 |
| y_hat = y_hat * q_dec |
|
|
| y_q_w_0 = self.single_part_for_writing_4x(y_q_0) |
| y_q_w_1 = self.single_part_for_writing_4x(y_q_1) |
| y_q_w_2 = self.single_part_for_writing_4x(y_q_2) |
| y_q_w_3 = self.single_part_for_writing_4x(y_q_3) |
| s_w_0 = self.single_part_for_writing_4x(s_hat_0) |
| s_w_1 = self.single_part_for_writing_4x(s_hat_1) |
| s_w_2 = self.single_part_for_writing_4x(s_hat_2) |
| s_w_3 = self.single_part_for_writing_4x(s_hat_3) |
| return y_q_w_0, y_q_w_1, y_q_w_2, y_q_w_3, s_w_0, s_w_1, s_w_2, s_w_3, y_hat |
|
|
| def decompress_prior_4x(self, common_params, y_spatial_prior_reduction, |
| y_spatial_prior_adaptor_1, y_spatial_prior_adaptor_2, |
| y_spatial_prior_adaptor_3, y_spatial_prior): |
| _, quant_step, scales, means = self.separate_prior(common_params, False) |
| common_params = y_spatial_prior_reduction(common_params) |
| dtype = means.dtype |
| device = means.device |
| B, C, H, W = means.size() |
| mask_0, mask_1, mask_2, mask_3 = self.get_mask_4x(B, C, H, W, dtype, device) |
|
|
| scales_r = self.single_part_for_writing_4x(scales * mask_0) |
| y_q_r = self.gaussian_encoder.decode_and_get_y(scales_r, dtype, device) |
| y_hat_curr_step = restore_y_4x(y_q_r, means, mask_0) |
| y_hat_so_far = y_hat_curr_step |
|
|
| params = torch.cat((y_hat_so_far, common_params), dim=1) |
| scales, means = y_spatial_prior(y_spatial_prior_adaptor_1(params)).chunk(2, 1) |
| scales_r = self.single_part_for_writing_4x(scales * mask_1) |
| y_q_r = self.gaussian_encoder.decode_and_get_y(scales_r, dtype, device) |
| y_hat_curr_step = restore_y_4x(y_q_r, means, mask_1) |
| y_hat_so_far = y_hat_so_far + y_hat_curr_step |
|
|
| params = torch.cat((y_hat_so_far, common_params), dim=1) |
| scales, means = y_spatial_prior(y_spatial_prior_adaptor_2(params)).chunk(2, 1) |
| scales_r = self.single_part_for_writing_4x(scales * mask_2) |
| y_q_r = self.gaussian_encoder.decode_and_get_y(scales_r, dtype, device) |
| y_hat_curr_step = restore_y_4x(y_q_r, means, mask_2) |
| y_hat_so_far = y_hat_so_far + y_hat_curr_step |
|
|
| params = torch.cat((y_hat_so_far, common_params), dim=1) |
| scales, means = y_spatial_prior(y_spatial_prior_adaptor_3(params)).chunk(2, 1) |
| scales_r = self.single_part_for_writing_4x(scales * mask_3) |
| y_q_r = self.gaussian_encoder.decode_and_get_y(scales_r, dtype, device) |
| y_hat_curr_step = restore_y_4x(y_q_r, means, mask_3) |
| y_hat_so_far = y_hat_so_far + y_hat_curr_step |
|
|
| y_hat = y_hat_so_far * quant_step |
|
|
| return y_hat |
|
|
