diff --git a/compressai/models/base.py b/compressai/models/base.py index 2fd63654..7be61187 100644 --- a/compressai/models/base.py +++ b/compressai/models/base.py @@ -205,7 +205,7 @@ def compress(self, x): def decompress(self, *args, **kwargs): y_out = self.latent_codec.decompress(*args, **kwargs) y_hat = y_out["y_hat"] - x_hat = self.g_s(y_hat).clamp_(0, 1) + x_hat = self.g_s(y_hat) return { "x_hat": x_hat, } diff --git a/compressai/models/cca.py b/compressai/models/cca.py index 106c519f..e3d00975 100644 --- a/compressai/models/cca.py +++ b/compressai/models/cca.py @@ -684,7 +684,7 @@ def decompress( shape: Dict[str, Tuple[int, ...]], ) -> Dict[str, Tensor]: y_out = self.latent_codec.decompress(strings, shape) - return {"x_hat": self.g_s(y_out["y_hat"]).clamp_(0, 1)} + return {"x_hat": self.g_s(y_out["y_hat"])} def update( self, scale_table: Optional[Tensor] = None, force: bool = False, **kwargs diff --git a/compressai/models/dcae.py b/compressai/models/dcae.py index 31f13056..05aac7a4 100644 --- a/compressai/models/dcae.py +++ b/compressai/models/dcae.py @@ -690,7 +690,7 @@ def decompress( self, strings: Sequence[Sequence[bytes]], shape: Sequence[int] ) -> Dict[str, Tensor]: out = self.latent_codec.decompress(strings, shape) - return {"x_hat": self.g_s(out["y_hat"]).clamp_(0, 1)} + return {"x_hat": self.g_s(out["y_hat"])} @classmethod def from_state_dict(cls, state_dict: Dict[str, Tensor]) -> "DCAE": diff --git a/compressai/models/google.py b/compressai/models/google.py index 712ce031..87702ec4 100644 --- a/compressai/models/google.py +++ b/compressai/models/google.py @@ -159,7 +159,7 @@ def compress(self, x): def decompress(self, strings, shape): assert isinstance(strings, list) and len(strings) == 1 y_hat = self.entropy_bottleneck.decompress(strings[0], shape) - x_hat = self.g_s(y_hat).clamp_(0, 1) + x_hat = self.g_s(y_hat) return {"x_hat": x_hat} @@ -329,7 +329,7 @@ def decompress(self, strings, shape): scales_hat = self.h_s(z_hat) indexes = self.gaussian_conditional.build_indexes(scales_hat) y_hat = self.gaussian_conditional.decompress(strings[0], indexes, z_hat.dtype) - x_hat = self.g_s(y_hat).clamp_(0, 1) + x_hat = self.g_s(y_hat) return {"x_hat": x_hat} @@ -426,7 +426,7 @@ def decompress(self, strings, shape): y_hat = self.gaussian_conditional.decompress( strings[0], indexes, means=means_hat ) - x_hat = self.g_s(y_hat).clamp_(0, 1) + x_hat = self.g_s(y_hat) return {"x_hat": x_hat} @@ -633,7 +633,7 @@ def _compress_ar(self, y_hat, params, height, width, kernel_size, padding): y_crop = y_crop[:, :, padding, padding] y_q = self.gaussian_conditional.quantize(y_crop, "symbols", means_hat) - y_hat[:, :, h + padding, w + padding] = y_q + means_hat + y_hat[:, :, h + padding, w + padding] = self.gaussian_conditional.dequantize(y_q, means_hat) symbols_list.extend(y_q.squeeze().tolist()) indexes_list.extend(indexes.squeeze().tolist()) @@ -688,7 +688,7 @@ def decompress(self, strings, shape): ) y_hat = F.pad(y_hat, (-padding, -padding, -padding, -padding)) - x_hat = self.g_s(y_hat).clamp_(0, 1) + x_hat = self.g_s(y_hat) return {"x_hat": x_hat} def _decompress_ar( @@ -724,7 +724,7 @@ def _decompress_ar( rv = decoder.decode_stream( indexes.squeeze().tolist(), cdf, cdf_lengths, offsets ) - rv = torch.Tensor(rv).reshape(1, -1, 1, 1) + rv = torch.tensor(rv, dtype=means_hat.dtype, device=means_hat.device).reshape(1, -1, 1, 1) rv = self.gaussian_conditional.dequantize(rv, means_hat) hp = h + padding diff --git a/compressai/models/mlic.py b/compressai/models/mlic.py index 0299685f..2e1fee5e 100644 --- a/compressai/models/mlic.py +++ b/compressai/models/mlic.py @@ -340,7 +340,7 @@ def decompress( shape: Dict[str, Union[List[Tuple[int, ...]], Tuple[int, ...]]], ) -> Dict[str, Tensor]: y_out = self.latent_codec.decompress(strings, shape) - return {"x_hat": self.g_s(y_out["y_hat"]).clamp_(0, 1)} + return {"x_hat": self.g_s(y_out["y_hat"])} @classmethod def from_state_dict(cls, state_dict: Dict[str, Tensor]) -> "_BaseMLIC": diff --git a/compressai/models/saaf.py b/compressai/models/saaf.py index 88387ed5..37b43bbe 100644 --- a/compressai/models/saaf.py +++ b/compressai/models/saaf.py @@ -878,7 +878,7 @@ def decompress( self, strings: Sequence[Sequence[bytes]], shape: Sequence[int] ) -> Dict[str, Tensor]: out = self.latent_codec.decompress(strings, shape) - return {"x_hat": self._decode(out["y_hat"]).clamp_(0, 1)} + return {"x_hat": self._decode(out["y_hat"])} @classmethod def from_state_dict(cls, state_dict: Dict[str, Tensor]) -> "SAAF": diff --git a/compressai/models/stf.py b/compressai/models/stf.py index e01ac586..b7b0ce36 100644 --- a/compressai/models/stf.py +++ b/compressai/models/stf.py @@ -741,7 +741,7 @@ def decompress( y_out = self.latent_codec.decompress(strings, shape) y_hat = y_out["y_hat"] height, width = y_hat.shape[2:] - return {"x_hat": self._synthesis_transform(y_hat, height, width).clamp_(0, 1)} + return {"x_hat": self._synthesis_transform(y_hat, height, width)} @classmethod def from_state_dict(cls, state_dict: Dict[str, Tensor]) -> "SymmetricalTransFormer": diff --git a/compressai/models/vbr.py b/compressai/models/vbr.py index 35769333..4fe94c5f 100644 --- a/compressai/models/vbr.py +++ b/compressai/models/vbr.py @@ -297,7 +297,7 @@ def decompress(self, strings, shape, stage: int = 2, s: int = 1, inputscale=0): y_hat = signs * (q_abs + q_offsets) y_ch_means = 0 y_hat = y_hat * rescale + y_ch_means - x_hat = self.g_s(y_hat).clamp_(0, 1) + x_hat = self.g_s(y_hat) return {"x_hat": x_hat} @@ -499,7 +499,7 @@ def decompress(self, strings, shape, stage: int = 2, s: int = 1, inputscale=0): y_hat = signs * (q_abs + q_offsets) y_hat = y_hat * rescale + means_hat - x_hat = self.g_s(y_hat).clamp_(0, 1) + x_hat = self.g_s(y_hat) return {"x_hat": x_hat} @@ -866,7 +866,7 @@ def decompress(self, strings, shape, stage: int = 2, s: int = 1, inputscale=0): ) y_hat = F.pad(y_hat, (-padding, -padding, -padding, -padding)) - x_hat = self.g_s(y_hat).clamp_(0, 1) + x_hat = self.g_s(y_hat) return {"x_hat": x_hat} def _decompress_ar( # noqa: C901 @@ -921,9 +921,8 @@ def _decompress_ar( # noqa: C901 rv = decoder.decode_stream( indexes.squeeze().tolist(), cdf, cdf_lengths, offsets ) - rv = ( - torch.Tensor(rv).reshape(1, -1, 1, 1).to(scales_hat.device) - ) # TODO: move rv to gpu ? + # TODO: move rv to gpu ? + rv = torch.tensor(rv, dtype=means_hat.dtype, device=means_hat.device).reshape(1, -1, 1, 1) if stage == 1: rv = self.gaussian_conditional.dequantize(rv, means_hat)