allow non list components_to_quantize.

Author: sayakpaul

src/diffusers/quantizers/pipe_quant_config.py

@@ -48,12 +48,15 @@ def __init__(
         self,
         quant_backend: str = None,
         quant_kwargs: Dict[str, Union[str, float, int, dict]] = None,
-        components_to_quantize: Optional[List[str]] = None,
+        components_to_quantize: Optional[Union[List[str], str]] = None,
         quant_mapping: Dict[str, Union[DiffQuantConfigMixin, "TransformersQuantConfigMixin"]] = None,
     ):
         self.quant_backend = quant_backend
         # Initialize kwargs to be {} to set to the defaults.
         self.quant_kwargs = quant_kwargs or {}
+        if components_to_quantize:
+            if isinstance(components_to_quantize, str):
+                components_to_quantize = [components_to_quantize]
         self.components_to_quantize = components_to_quantize
         self.quant_mapping = quant_mapping
         self.config_mapping = {}  # book-keeping Example: `{module_name: quant_config}`

tests/quantization/test_pipeline_level_quantization.py

@@ -298,3 +298,19 @@ def _parse_config_string(self, config_string: str) -> tuple[str, dict]:
         data = json.loads(json_part)
 
         return data
+
+    def test_single_component_to_quantize(self):
+        component_to_quantize = "transformer"
+        quant_config = PipelineQuantizationConfig(
+            quant_backend="bitsandbytes_8bit",
+            quant_kwargs={"load_in_8bit": True},
+            components_to_quantize=component_to_quantize,
+        )
+        pipe = DiffusionPipeline.from_pretrained(
+            self.model_name,
+            quantization_config=quant_config,
+            torch_dtype=torch.bfloat16,
+        )
+        for name, component in pipe.components.items():
+            if name == component_to_quantize:
+                self.assertTrue(hasattr(component.config, "quantization_config"))