refactor: fix lint errors

Author: YodaEmbedding

compressai/entropy_models/entropy_models.py

@@ -232,13 +232,18 @@ def _check_cdf_length(self):
         if len(self._cdf_length.size()) != 1:
             raise ValueError(f"Invalid offsets size {self._cdf_length.size()}")
 
-    def compress(self, inputs, indexes, means=None):
+    def compress(
+        self,
+        inputs: torch.Tensor,
+        indexes: torch.Tensor,
+        means: Optional[torch.Tensor] = None,
+    ):
         """
         Compress input tensors to char strings.
 
         Args:
             inputs (torch.Tensor): input tensors
-            indexes (torch.IntTensor): tensors CDF indexes
+            indexes (torch.Tensor): tensors CDF indexes
             means (torch.Tensor, optional): optional tensor means
         """
         symbols = self.quantize(inputs, "symbols", means)
@@ -269,17 +274,17 @@ def compress(self, inputs, indexes, means=None):
 
     def decompress(
         self,
-        strings: str,
-        indexes: torch.IntTensor,
+        strings: List[bytes],
+        indexes: torch.Tensor,
         dtype: torch.dtype = torch.float,
-        means: torch.Tensor = None,
+        means: Optional[torch.Tensor] = None,
     ):
         """
         Decompress char strings to tensors.
 
         Args:
-            strings (str): compressed tensors
-            indexes (torch.IntTensor): tensors CDF indexes
+            strings (list[bytes]): compressed tensors
+            indexes (torch.Tensor): tensors CDF indexes
             dtype (torch.dtype): type of dequantized output
             means (torch.Tensor, optional): optional tensor means
         """

compressai/entropy_models/entropy_models_vbr.py

@@ -29,7 +29,7 @@
 
 import warnings
 
-from typing import Any, Callable, Optional, Tuple
+from typing import Any, Callable, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -258,13 +258,19 @@ def _check_cdf_length(self):
         if len(self._cdf_length.size()) != 1:
             raise ValueError(f"Invalid offsets size {self._cdf_length.size()}")
 
-    def compress(self, inputs, indexes, means=None, qs=None):
+    def compress(
+        self,
+        inputs: torch.Tensor,
+        indexes: torch.Tensor,
+        means: Optional[torch.Tensor] = None,
+        qs: Optional[torch.Tensor] = None,
+    ):
         """
         Compress input tensors to char strings.
 
         Args:
             inputs (torch.Tensor): input tensors
-            indexes (torch.IntTensor): tensors CDF indexes
+            indexes (torch.Tensor): tensors CDF indexes
             means (torch.Tensor, optional): optional tensor means
             qs (torch.Tensor, optional): optional quantization step size
         """
@@ -299,18 +305,18 @@ def compress(self, inputs, indexes, means=None, qs=None):
 
     def decompress(
         self,
-        strings: str,
-        indexes: torch.IntTensor,
+        strings: List[bytes],
+        indexes: torch.Tensor,
         dtype: torch.dtype = torch.float,
-        means: torch.Tensor = None,
-        qs=None,
+        means: Optional[torch.Tensor] = None,
+        qs: Optional[torch.Tensor] = None,
     ):
         """
         Decompress char strings to tensors.
 
         Args:
-            strings (str): compressed tensors
-            indexes (torch.IntTensor): tensors CDF indexes
+            strings (list[bytes]): compressed tensors
+            indexes (torch.Tensor): tensors CDF indexes
             dtype (torch.dtype): type of dequantized output
             means (torch.Tensor, optional): optional tensor means
             qs (torch.Tensor, optional): optional quantization step size

compressai/latent_codecs/gaussian_conditional.py

@@ -85,7 +85,7 @@ def __init__(
         gaussian_conditional: Optional[GaussianConditional] = None,
         entropy_parameters: Optional[nn.Module] = None,
         quantizer: str = "noise",
-        chunks: Tuple[str] = ("scales", "means"),
+        chunks: Tuple[str, ...] = ("scales", "means"),
         **kwargs,
     ):
         super().__init__()