up

Author: sayakpaul

src/diffusers/quantizers/quantization_config.py

@@ -518,32 +518,20 @@ def post_init(self):
         TORCHAO_QUANT_TYPE_METHODS = self._get_torchao_quant_type_to_method()
         AO_VERSION = self._get_ao_version()
 
-        if isinstance(self.quant_type, str) and self.quant_type not in TORCHAO_QUANT_TYPE_METHODS.keys():
-            is_floating_quant_type = self.quant_type.startswith("float") or self.quant_type.startswith("fp")
-            if is_floating_quant_type and not self._is_xpu_or_cuda_capability_atleast_8_9():
+        if isinstance(self.quant_type, str):
+            if self.quant_type not in TORCHAO_QUANT_TYPE_METHODS.keys():
+                is_floating_quant_type = self.quant_type.startswith("float") or self.quant_type.startswith("fp")
+                if is_floating_quant_type and not self._is_xpu_or_cuda_capability_atleast_8_9():
+                    raise ValueError(
+                        f"Requested quantization type: {self.quant_type} is not supported on GPUs with CUDA capability <= 8.9. You "
+                        f"can check the CUDA capability of your GPU using `torch.cuda.get_device_capability()`."
+                    )
+
                 raise ValueError(
-                    f"Requested quantization type: {self.quant_type} is not supported on GPUs with CUDA capability <= 8.9. You "
-                    f"can check the CUDA capability of your GPU using `torch.cuda.get_device_capability()`."
+                    f"Requested quantization type: {self.quant_type} is not supported or is an incorrect `quant_type` name. If you think the "
+                    f"provided quantization type should be supported, please open an issue at https://github.com/huggingface/diffusers/issues."
                 )
 
-            raise ValueError(
-                f"Requested quantization type: {self.quant_type} is not supported or is an incorrect `quant_type` name. If you think the "
-                f"provided quantization type should be supported, please open an issue at https://github.com/huggingface/diffusers/issues."
-            )
-        elif AO_VERSION > version.parse("0.9.0"):
-            from torchao.quantization.quant_api import AOBaseConfig
-
-            if not isinstance(self.quant_type, AOBaseConfig):
-                raise TypeError(
-                    f"`quant_type` must be either a string or an `AOBaseConfig` instance, got {type(self.quant_type)}."
-                )
-        else:
-            raise ValueError(
-                f"In torchao <= 0.9.0, quant_type must be a string. Got {type(self.quant_type)}. "
-                f"Please upgrade to torchao > 0.9.0 to use `AOBaseConfig` instances."
-            )
-
-        if isinstance(self.quant_type, str):
             method = TORCHAO_QUANT_TYPE_METHODS[self.quant_type]
             signature = inspect.signature(method)
             all_kwargs = {
@@ -558,6 +546,18 @@ def post_init(self):
                     f'The quantization method "{self.quant_type}" does not support the following keyword arguments: '
                     f"{unsupported_kwargs}. The following keywords arguments are supported: {all_kwargs}."
                 )
+        elif AO_VERSION > version.parse("0.9.0"):
+            from torchao.quantization.quant_api import AOBaseConfig
+
+            if not isinstance(self.quant_type, AOBaseConfig):
+                raise TypeError(
+                    f"`quant_type` must be either a string or an `AOBaseConfig` instance, got {type(self.quant_type)}."
+                )
+        else:
+            raise ValueError(
+                f"In torchao <= 0.9.0, quant_type must be a string. Got {type(self.quant_type)}. "
+                f"Please upgrade to torchao > 0.9.0 to use `AOBaseConfig` instances."
+            )
 
     def to_dict(self):
         """Convert configuration to a dictionary."""

tests/quantization/torchao/test_torchao.py

@@ -590,7 +590,6 @@ def _test_original_model_expected_slice(self, quant_method, quant_method_kwargs,
         inputs = self.get_dummy_tensor_inputs(torch_device)
         output = quantized_model(**inputs)[0]
         output_slice = output.flatten()[-9:].detach().float().cpu().numpy()
-        print(f"{output_slice=}")
         weight = quantized_model.transformer_blocks[0].ff.net[2].weight
         self.assertTrue(isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor)))
         self.assertTrue(numpy_cosine_similarity_distance(output_slice, expected_slice) < 1e-3)