[Quantization] Allow loading TorchAO serialized Tensor objects with torch>=2.6

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

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()

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__)
 
 

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,

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.