feat: draft pipeline-level quant config.

Author: sayakpaul

src/diffusers/pipelines/pipeline_loading_utils.py

@@ -667,9 +667,12 @@ def load_sub_model(
     use_safetensors: bool,
     dduf_entries: Optional[Dict[str, DDUFEntry]],
     provider_options: Any,
+    quantization_config: Optional[Any] = None,
 ):
     """Helper method to load the module `name` from `library_name` and `class_name`"""
 
+    from ..quantizers import PipelineQuantizationConfig
+
     # retrieve class candidates
 
     class_obj, class_candidates = get_class_obj_and_candidates(
@@ -761,6 +764,12 @@ def load_sub_model(
         else:
             loading_kwargs["low_cpu_mem_usage"] = False
 
+    if quantization_config is not None and isinstance(quantization_config, PipelineQuantizationConfig):
+        exclude_modules = quantization_config.exclude_modules or []
+        if name not in exclude_modules:
+            model_quant_config = _resolve_quant_config(quantization_config, is_diffusers=is_diffusers_model)
+            loading_kwargs["quantization_config"] = model_quant_config
+
     # check if the module is in a subdirectory
     if dduf_entries:
         loading_kwargs["dduf_entries"] = dduf_entries
@@ -1070,3 +1079,22 @@ def _maybe_raise_error_for_incorrect_transformers(config_dict):
                 break
     if has_transformers_component and not is_transformers_version(">", "4.47.1"):
         raise ValueError("Please upgrade your `transformers` installation to the latest version to use DDUF.")
+
+
+def _resolve_quant_config(quant_config, is_diffusers=True):
+    if is_diffusers:
+        from ..quantizers.auto import AUTO_QUANTIZATION_CONFIG_MAPPING
+    else:
+        from transformers.quantizers.auto import AUTO_QUANTIZATION_CONFIG_MAPPING
+
+    quant_backend = quant_config.quant_backend
+    if quant_backend not in AUTO_QUANTIZATION_CONFIG_MAPPING:
+        raise ValueError(
+            f"Provided {quant_backend=} was not found in the support quantizers. Available ones are: {AUTO_QUANTIZATION_CONFIG_MAPPING.keys()}."
+        )
+
+    quant_config_cls = AUTO_QUANTIZATION_CONFIG_MAPPING[quant_backend]
+
+    quant_kwargs = quant_config.quant_kwargs
+    quant_config = quant_config_cls(**quant_kwargs)
+    return quant_config

src/diffusers/pipelines/pipeline_utils.py

@@ -702,6 +702,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
         use_safetensors = kwargs.pop("use_safetensors", None)
         use_onnx = kwargs.pop("use_onnx", None)
         load_connected_pipeline = kwargs.pop("load_connected_pipeline", False)
+        quantization_config = kwargs.pop("quantization_config", None)
 
         if not isinstance(torch_dtype, torch.dtype):
             torch_dtype = torch.float32
@@ -973,6 +974,7 @@ def load_module(name, value):
                     use_safetensors=use_safetensors,
                     dduf_entries=dduf_entries,
                     provider_options=provider_options,
+                    quantization_config=quantization_config,
                 )
                 logger.info(
                     f"Loaded {name} as {class_name} from `{name}` subfolder of {pretrained_model_name_or_path}."

src/diffusers/quantizers/__init__.py

@@ -14,3 +14,10 @@
 
 from .auto import DiffusersAutoQuantizer
 from .base import DiffusersQuantizer
+
+
+class PipelineQuantizationConfig:
+    def __init__(self, quant_backend, quant_kwargs, exclude_modules):
+        self.quant_backend = quant_backend
+        self.quant_kwargs = quant_kwargs
+        self.exclude_modules = exclude_modules