Incompatibility with torchao >= 0.17.0 due to hardcoded legacy import paths in torchao_quantizer.py

[lisi31415926]

Describe the bug

In src/diffusers/quantizers/torchao/torchao_quantizer.py (lines 99-102), the imports inside _update_torch_safe_globals() are hardcoded to legacy torchao paths: Pythonfrom torchao.dtypes import NF4Tensor
from torchao.dtypes.uintx.uintx_layout import UintxAQTTensorImpl, UintxTensor
Starting from torchao >= 0.17.0 (and 0.18.0), these modules were refactored/relocated:UintxTensor / UintxAQTTensorImpl moved to torchao.prototype.dtypesNF4Tensor moved to torchao.quantizationImporting Diffusers with torchao >= 0.17.0 installed leads to ModuleNotFoundError / ImportError.

Reproduction

Proposed Fix
Wrap the imports in try-except fallback blocks inside _update_torch_safe_globals():
Python
safe_globals = []

1. NF4Tensor fallback

try:
from torchao.dtypes import NF4Tensor
safe_globals.append(NF4Tensor)
except ImportError:
try:
from torchao.quantization import NF4Tensor
safe_globals.append(NF4Tensor)
except ImportError:
pass

2. UintxTensor fallback

try:
from torchao.dtypes.uintx.uintx_layout import UintxAQTTensorImpl, UintxTensor
safe_globals.extend([UintxTensor, UintxAQTTensorImpl])
except ImportError:
try:
from torchao.prototype.dtypes import UintxAQTTensorImpl, UintxTensor
safe_globals.extend([UintxTensor, UintxAQTTensorImpl])
except ImportError:
pass

if safe_globals:
torch.serialization.add_safe_globals(safe_globals=safe_globals)

Logs

System Info

py 3.13 win11

Who can help?

No response