add validation to the input quantization_config.

Author: sayakpaul

src/diffusers/pipelines/pipeline_utils.py

@@ -47,6 +47,7 @@
 from ..models import AutoencoderKL
 from ..models.attention_processor import FusedAttnProcessor2_0
 from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, ModelMixin
+from ..quantizers import PipelineQuantizationConfig
 from ..quantizers.bitsandbytes.utils import _check_bnb_status
 from ..schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME
 from ..utils import (
@@ -742,6 +743,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                 " install accelerate\n```\n."
             )
 
+        if quantization_config is not None and not isinstance(quantization_config, PipelineQuantizationConfig):
+            raise ValueError("`quantization_config` must be an instance of `PipelineQuantizationConfig`.")
+
         if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"):
             raise NotImplementedError(
                 "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set"

tests/quantization/test_pipeline_level_quantization.py

@@ -74,7 +74,7 @@ def test_quant_config_set_correctly_through_kwargs(self):
 
         _ = pipe(self.prompt, num_inference_steps=self.num_inference_steps)
 
-    def test_quant_config_set_correctly_granular(self):
+    def test_quant_config_set_correctly_through_granular(self):
         quant_config = PipelineQuantizationConfig(
             quant_mapping={
                 "transformer": QuantoConfig(weights_dtype="int8"),
@@ -128,6 +128,21 @@ def test_validation_for_kwargs(self):
             "The signatures of the __init__ methods of the quantization config classes" in str(err_context.exception)
         )
 
+    def test_raises_error_for_wrong_config_class(self):
+        quant_config = {
+            "transformer": QuantoConfig(weights_dtype="int8"),
+            "text_encoder_2": TranBitsAndBytesConfig(load_in_4bit=True, compute_dtype=torch.bfloat16),
+        }
+        with self.assertRaises(ValueError) as err_context:
+            _ = DiffusionPipeline.from_pretrained(
+                self.model_name,
+                quantization_config=quant_config,
+                torch_dtype=torch.bfloat16,
+            )
+        self.assertTrue(
+            str(err_context.exception) == "`quantization_config` must be an instance of `PipelineQuantizationConfig`."
+        )
+
     def test_validation_for_mapping(self):
         with self.assertRaises(ValueError) as err_context:
             _ = PipelineQuantizationConfig(