From 626174d9575c92da20ea7405ff99a4d4eaf26818 Mon Sep 17 00:00:00 2001 From: li-lizhe <147392333@qq.com> Date: Sat, 19 Sep 2026 02:05:33 +0800 Subject: [PATCH] fix: support Ascend NPU stream in enable_group_offloading enable_group_offloading only checked cuda.is_available() then xpu. On Ascend NPU (torch_npu) with use_stream=True, this falls through to raise ValueError("requires CUDA or XPU"). Add elif branch for torch.npu.is_available() with torch.npu.Stream(). --- src/diffusers/hooks/group_offloading.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/diffusers/hooks/group_offloading.py b/src/diffusers/hooks/group_offloading.py index e6c8d8732113..1be261ca706b 100644 --- a/src/diffusers/hooks/group_offloading.py +++ b/src/diffusers/hooks/group_offloading.py @@ -675,8 +675,12 @@ def apply_group_offloading( stream = torch.cuda.Stream() elif hasattr(torch, "xpu") and torch.xpu.is_available(): stream = torch.Stream() + elif hasattr(torch, "npu") and torch.npu.is_available(): + stream = torch.npu.Stream() else: - raise ValueError("Using streams for data transfer requires a CUDA device, or an Intel XPU device.") + raise ValueError( + "Using streams for data transfer requires a CUDA device, an Intel XPU device, or an Ascend NPU device." + ) if not use_stream and record_stream: raise ValueError("`record_stream` cannot be True when `use_stream=False`.")