theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-24 01:48:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1176 from Feng0w0/z-image-rope

Browse Source 
[model][NPU]: Z-image model support NPU
This commit is contained in:
Zhongjie Duan
2026-01-12 10:11:22 +08:00
committed by GitHub
parent 00f2d1aa5d c1c9a4853b
commit 381067515c
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
diffsynth/models/z_image_dit.py
View File

@@ -8,6 +8,7 @@ from torch.nn.utils.rnn import pad_sequence
from torch.nn import RMSNorm from torch.nn import RMSNorm
from ..core.attention import attention_forward from ..core.attention import attention_forward
from ..core.device.npu_compatible_device import IS_NPU_AVAILABLE
from ..core.gradient import gradient_checkpoint_forward from ..core.gradient import gradient_checkpoint_forward
@@ -315,7 +316,10 @@ class RopeEmbedder:
result = [] result = []
for i in range(len(self.axes_dims)): for i in range(len(self.axes_dims)):
index = ids[:, i] index = ids[:, i]
result.append(self.freqs_cis[i][index]) if IS_NPU_AVAILABLE:
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) return torch.cat(result, dim=-1)