update

Author: dsikka

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -180,7 +180,7 @@ def from_pretrained_model(
         # assume multiple compression formats means mixed-precision
         # as we currently only support one compressor per precision type and scheme
         if len(quantization_format) > 1:
-            quantization_format = CompressionFormat.mixed_precision
+            quantization_format = CompressionFormat.mixed_precision.value
         else:
             quantization_format = quantization_format[0]
 
@@ -258,15 +258,15 @@ def parse_quantization_config(
 
         return quantization_config
 
-    def _fetch_unique_quantization_formats(self):
+    def _fetch_unique_quantization_formats(self) -> List[str]:
         """
-        Get all unique compression formats used in
-        model
+        Get all unique compression formats present in a model
+        :return: list of quantization formats
         """
         quantization_formats = []
         for _, scheme in self.quantization_config.config_groups.items():
             if scheme.format not in quantization_formats:
-                quantization_formats.append(scheme)
+                quantization_formats.append(scheme.format)
         return quantization_formats
 
     def __init__(
@@ -287,7 +287,6 @@ def __init__(
             )
 
         quantization_formats = self._fetch_unique_quantization_formats()
-
         if quantization_config is not None:
             self.quantization_compressor = {}
             for format in quantization_formats:
@@ -694,6 +693,7 @@ def update_config(self, save_directory: str):
         config_data[QUANTIZATION_CONFIG_NAME][
             COMPRESSION_VERSION_NAME
         ] = compressed_tensors.__version__
+
         if self.quantization_config is not None:
             self.quantization_config.quant_method = DEFAULT_QUANTIZATION_METHOD
         else:

src/compressed_tensors/quantization/quant_scheme.py

@@ -49,7 +49,7 @@ class QuantizationScheme(BaseModel):
     weights: Optional[QuantizationArgs] = None
     input_activations: Optional[QuantizationArgs] = None
     output_activations: Optional[QuantizationArgs] = None
-    format: Optional[CompressionFormat] = None
+    format: Optional[str] = None
 
     @model_validator(mode="after")
     def validate_model_after(model: "QuantizationScheme") -> "QuantizationScheme":