fix: fix wrong dtype

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,
         }

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

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":

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

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":

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":

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":

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)