diff --git a/docs/source/en/quantization/torchao.md b/docs/source/en/quantization/torchao.md index c056876c2f09..19a8970fa9df 100644 --- a/docs/source/en/quantization/torchao.md +++ b/docs/source/en/quantization/torchao.md @@ -126,7 +126,7 @@ image = pipe(prompt, num_inference_steps=30, guidance_scale=7.0).images[0] image.save("output.png") ``` -Some quantization methods, such as `uint4wo`, cannot be loaded directly and may result in an `UnpicklingError` when trying to load the models, but work as expected when saving them. In order to work around this, one can load the state dict manually into the model. Note, however, that this requires using `weights_only=False` in `torch.load`, so it should be run only if the weights were obtained from a trustable source. +If you are using `torch<=2.6.0`, some quantization methods, such as `uint4wo`, cannot be loaded directly and may result in an `UnpicklingError` when trying to load the models, but work as expected when saving them. In order to work around this, one can load the state dict manually into the model. Note, however, that this requires using `weights_only=False` in `torch.load`, so it should be run only if the weights were obtained from a trustable source. ```python import torch diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index c482ed324179..6421ea871a75 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -2,20 +2,14 @@ from typing import TYPE_CHECKING -from diffusers.quantizers import quantization_config -from diffusers.utils import dummy_gguf_objects -from diffusers.utils.import_utils import ( - is_bitsandbytes_available, - is_gguf_available, - is_optimum_quanto_version, - is_torchao_available, -) - from .utils import ( DIFFUSERS_SLOW_IMPORT, OptionalDependencyNotAvailable, _LazyModule, + is_accelerate_available, + is_bitsandbytes_available, is_flax_available, + is_gguf_available, is_k_diffusion_available, is_librosa_available, is_note_seq_available, @@ -24,6 +18,7 @@ is_scipy_available, is_sentencepiece_available, is_torch_available, + is_torchao_available, is_torchsde_available, is_transformers_available, ) @@ -65,7 +60,7 @@ } try: - if not is_bitsandbytes_available(): + if not is_torch_available() and not is_accelerate_available() and not is_bitsandbytes_available(): raise OptionalDependencyNotAvailable() except OptionalDependencyNotAvailable: from .utils import dummy_bitsandbytes_objects @@ -77,7 +72,7 @@ _import_structure["quantizers.quantization_config"].append("BitsAndBytesConfig") try: - if not is_gguf_available(): + if not is_torch_available() and not is_accelerate_available() and not is_gguf_available(): raise OptionalDependencyNotAvailable() except OptionalDependencyNotAvailable: from .utils import dummy_gguf_objects @@ -89,7 +84,7 @@ _import_structure["quantizers.quantization_config"].append("GGUFQuantizationConfig") try: - if not is_torchao_available(): + if not is_torch_available() and not is_accelerate_available() and not is_torchao_available(): raise OptionalDependencyNotAvailable() except OptionalDependencyNotAvailable: from .utils import dummy_torchao_objects @@ -101,7 +96,7 @@ _import_structure["quantizers.quantization_config"].append("TorchAoConfig") try: - if not is_optimum_quanto_available(): + if not is_torch_available() and not is_accelerate_available() and not is_optimum_quanto_available(): raise OptionalDependencyNotAvailable() except OptionalDependencyNotAvailable: from .utils import dummy_optimum_quanto_objects @@ -112,7 +107,6 @@ else: _import_structure["quantizers.quantization_config"].append("QuantoConfig") - try: if not is_onnx_available(): raise OptionalDependencyNotAvailable() diff --git a/src/diffusers/quantizers/torchao/torchao_quantizer.py b/src/diffusers/quantizers/torchao/torchao_quantizer.py index 03cb29c6f037..f9fb217ed6bd 100644 --- a/src/diffusers/quantizers/torchao/torchao_quantizer.py +++ b/src/diffusers/quantizers/torchao/torchao_quantizer.py @@ -23,7 +23,14 @@ from packaging import version -from ...utils import get_module_from_name, is_torch_available, is_torch_version, is_torchao_available, logging +from ...utils import ( + get_module_from_name, + is_torch_available, + is_torch_version, + is_torchao_available, + is_torchao_version, + logging, +) from ..base import DiffusersQuantizer @@ -62,6 +69,43 @@ from torchao.quantization import quantize_ +def _update_torch_safe_globals(): + safe_globals = [ + (torch.uint1, "torch.uint1"), + (torch.uint2, "torch.uint2"), + (torch.uint3, "torch.uint3"), + (torch.uint4, "torch.uint4"), + (torch.uint5, "torch.uint5"), + (torch.uint6, "torch.uint6"), + (torch.uint7, "torch.uint7"), + ] + try: + from torchao.dtypes import NF4Tensor + from torchao.dtypes.floatx.float8_layout import Float8AQTTensorImpl + from torchao.dtypes.uintx.uint4_layout import UInt4Tensor + from torchao.dtypes.uintx.uintx_layout import UintxAQTTensorImpl, UintxTensor + + safe_globals.extend([UintxTensor, UInt4Tensor, UintxAQTTensorImpl, Float8AQTTensorImpl, NF4Tensor]) + + except (ImportError, ModuleNotFoundError) as e: + logger.warning( + "Unable to import `torchao` Tensor objects. This may affect loading checkpoints serialized with `torchao`" + ) + logger.debug(e) + + finally: + torch.serialization.add_safe_globals(safe_globals=safe_globals) + + +if ( + is_torch_available() + and is_torch_version(">=", "2.6.0") + and is_torchao_available() + and is_torchao_version(">=", "0.7.0") +): + _update_torch_safe_globals() + + logger = logging.get_logger(__name__) diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 1684c434f55e..50a470772772 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -94,6 +94,7 @@ is_torch_xla_available, is_torch_xla_version, is_torchao_available, + is_torchao_version, is_torchsde_available, is_torchvision_available, is_transformers_available, diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index b6aa8e96e619..5c3d27dd2e6c 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -868,6 +868,21 @@ def is_gguf_version(operation: str, version: str): return compare_versions(parse(_gguf_version), operation, version) +def is_torchao_version(operation: str, version: str): + """ + Compares the current torchao version to a given reference with an operation. + + Args: + operation (`str`): + A string representation of an operator, such as `">"` or `"<="` + version (`str`): + A version string + """ + if not _is_torchao_available: + return False + return compare_versions(parse(_torchao_version), operation, version) + + def is_k_diffusion_version(operation: str, version: str): """ Compares the current k-diffusion version to a given reference with an operation.