AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-24 20:59:27 +03:00
Code Activity

Enable BitsAndBytes quantization in MPS (#13915)

Browse Source 
Fix BitsAndBytes quantization in MPS

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
This commit is contained in:
Lucas Ste
2026-06-23 02:20:14 -03:00
committed by GitHub
parent aed5c64154
commit e97a1ad2a2
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py
View File

@@ -61,7 +61,7 @@ class BnB4BitDiffusersQuantizer(DiffusersQuantizer):
self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules
def validate_environment(self, *args, **kwargs):
if not (torch.cuda.is_available() or torch.xpu.is_available()):
if not (torch.cuda.is_available() or torch.xpu.is_available() or torch.mps.is_available()):
raise RuntimeError("No GPU found. A GPU is needed for quantization.")
if not is_accelerate_available() or is_accelerate_version("<", "0.26.0"):
raise ImportError(
@@ -240,6 +240,8 @@ class BnB4BitDiffusersQuantizer(DiffusersQuantizer):
if device_map is None:
if torch.xpu.is_available():
current_device = f"xpu:{torch.xpu.current_device()}"
elif torch.mps.is_available():
current_device = "mps"
else:
current_device = f"cuda:{torch.cuda.current_device()}"
device_map = {"": current_device}
@@ -411,6 +413,8 @@ class BnB8BitDiffusersQuantizer(DiffusersQuantizer):
if device_map is None:
if torch.xpu.is_available():
current_device = f"xpu:{torch.xpu.current_device()}"
elif torch.mps.is_available():
current_device = "mps"
else:
current_device = f"cuda:{torch.cuda.current_device()}"
device_map = {"": current_device}