support mapping.

Author: sayakpaul

src/diffusers/pipelines/pipeline_loading_utils.py

@@ -769,9 +769,10 @@ def load_sub_model(
         and isinstance(quantization_config, PipelineQuantizationConfig)
         and issubclass(class_obj, torch.nn.Module)
     ):
-        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)
+        model_quant_config = quantization_config._resolve_quant_config(
+            is_diffusers=is_diffusers_model, module_name=name
+        )
+        if model_quant_config is not None:
             loading_kwargs["quantization_config"] = model_quant_config
 
     # check if the module is in a subdirectory
@@ -1085,20 +1086,33 @@ def _maybe_raise_error_for_incorrect_transformers(config_dict):
         raise ValueError("Please upgrade your `transformers` installation to the latest version to use DDUF.")
 
 
-def _resolve_quant_config(quant_config, is_diffusers=True):
+def _resolve_quant_config(quant_config, is_diffusers=True, module_name=None):
     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]
+    # Granular case.
+    if getattr(quant_config, "is_granular", False):
+        config = quant_config.mapping.get(module_name)
+        quant_backend = config.get("quant_backend")
+        if quant_backend not in AUTO_QUANTIZATION_CONFIG_MAPPING:
+            raise ValueError(
+                f"Module '{module_name}': Provided quant_backend={quant_backend} was not found. "
+                f"Available ones are: {list(AUTO_QUANTIZATION_CONFIG_MAPPING.keys())}."
+            )
+        quant_config_cls = AUTO_QUANTIZATION_CONFIG_MAPPING[quant_backend]
+        quant_kwargs = config.get("quant_kwargs")
 
-    quant_kwargs = quant_config.quant_kwargs
-    quant_config = quant_config_cls(**quant_kwargs)
-    return quant_config
+        return quant_config_cls(**quant_kwargs)
+    else:
+        # Global config case.
+        quant_backend = quant_config.quant_backend
+        if quant_backend not in AUTO_QUANTIZATION_CONFIG_MAPPING:
+            raise ValueError(
+                f"Provided quant_backend={quant_backend} was not found. "
+                f"Available ones are: {list(AUTO_QUANTIZATION_CONFIG_MAPPING.keys())}."
+            )
+        quant_config_cls = AUTO_QUANTIZATION_CONFIG_MAPPING[quant_backend]
+        quant_kwargs = quant_config.quant_kwargs
+        return quant_config_cls(**quant_kwargs)

src/diffusers/pipelines/pipeline_utils.py

@@ -875,6 +875,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
         }
         init_kwargs = {**init_kwargs, **passed_pipe_kwargs}
 
+        # TODO: add checking for quantization_config `mapping` i.e., if the modules specified there actually exist.
+        #########################
+
         # remove `null` components
         def load_module(name, value):
             if value[0] is None:

src/diffusers/quantizers/__init__.py

@@ -12,12 +12,63 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from ..utils import is_transformers_available
 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
+    def __init__(self, quant_backend: str, quant_kwargs: dict, exclude_modules: list, mapping: dict = None):
+        if mapping is not None:
+            self.mapping = mapping
+            self.is_granular = True
+        else:
+            self.quant_backend = quant_backend
+            self.quant_kwargs = quant_kwargs or {}
+            self.exclude_modules = exclude_modules or []
+            self.is_granular = False
+
+        self.post_init()
+
+    def post_init(self):
+        if self.is_granular and self.mapping is None:
+            raise ValueError(
+                "In the granular case, a `mapping` defining the quantization configs"
+                " for the desired modules have to be defined."
+            )
+
+    def _resolve_quant_config(self, is_diffusers=True, module_name=None):
+        if is_diffusers:
+            from ..quantizers.auto import AUTO_QUANTIZATION_CONFIG_MAPPING
+        else:
+            if is_transformers_available():
+                from transformers.quantizers.auto import AUTO_QUANTIZATION_CONFIG_MAPPING
+
+        # Granular case.
+        if self.is_granular:
+            config = self.mapping.get(module_name)
+            quant_backend = config.get("quant_backend")
+            if quant_backend not in AUTO_QUANTIZATION_CONFIG_MAPPING:
+                raise ValueError(
+                    f"Module '{module_name}': Provided quant_backend={quant_backend} was not found. "
+                    f"Available ones are: {list(AUTO_QUANTIZATION_CONFIG_MAPPING.keys())}."
+                )
+            quant_config_cls = AUTO_QUANTIZATION_CONFIG_MAPPING[quant_backend]
+            quant_kwargs = config.get("quant_kwargs")
+
+            return quant_config_cls(**quant_kwargs)
+
+        # Global config case.
+        elif module_name not in self.exclude_modules:
+            quant_backend = self.quant_backend
+            if quant_backend not in AUTO_QUANTIZATION_CONFIG_MAPPING:
+                raise ValueError(
+                    f"Provided quant_backend={quant_backend} was not found. "
+                    f"Available ones are: {list(AUTO_QUANTIZATION_CONFIG_MAPPING.keys())}."
+                )
+            quant_config_cls = AUTO_QUANTIZATION_CONFIG_MAPPING[quant_backend]
+            quant_kwargs = self.quant_kwargs
+            return quant_config_cls(**quant_kwargs)
+
+        else:
+            return None