From 3ee5f53a366c1a8321285184b0da14ec7b8b4e8c Mon Sep 17 00:00:00 2001 From: feng0w0 Date: Wed, 7 Jan 2026 11:31:22 +0800 Subject: [PATCH 1/2] [model][NPU]:Z-image model support NPU --- diffsynth/models/z_image_dit.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/diffsynth/models/z_image_dit.py b/diffsynth/models/z_image_dit.py index 7664fc5a3..4c5622c0a 100644 --- a/diffsynth/models/z_image_dit.py +++ b/diffsynth/models/z_image_dit.py @@ -8,6 +8,7 @@ from torch.nn import RMSNorm from ..core.attention import attention_forward +from ..core.device.npu_compatible_device import IS_NPU_AVAILABLE from ..core.gradient import gradient_checkpoint_forward @@ -274,7 +275,10 @@ def __call__(self, ids: torch.Tensor): result = [] for i in range(len(self.axes_dims)): index = ids[:, i] - result.append(self.freqs_cis[i][index]) + if IS_NPU_AVAILABLE: + result.append(self.freqs_cis[i][index]) + else: + result.append(torch.index_select(self.freqs_cis[i], 0, index)) return torch.cat(result, dim=-1) From c1c9a4853b2b862a3010637d16f5d9ddb241f87e Mon Sep 17 00:00:00 2001 From: feng0w0 Date: Wed, 7 Jan 2026 11:42:19 +0800 Subject: [PATCH 2/2] [model][NPU]:Z-image model support NPU --- diffsynth/models/z_image_dit.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/diffsynth/models/z_image_dit.py b/diffsynth/models/z_image_dit.py index 4c5622c0a..d20ec5193 100644 --- a/diffsynth/models/z_image_dit.py +++ b/diffsynth/models/z_image_dit.py @@ -276,9 +276,9 @@ def __call__(self, ids: torch.Tensor): for i in range(len(self.axes_dims)): index = ids[:, i] if IS_NPU_AVAILABLE: - result.append(self.freqs_cis[i][index]) - else: result.append(torch.index_select(self.freqs_cis[i], 0, index)) + else: + result.append(self.freqs_cis[i][index]) return torch.cat(result, dim=-1)