xinyun/ComfyUI-MochiWrapper
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

allow apex norm to work

Browse Source
This commit is contained in:
kijai 2024-11-05 19:40:22 +02:00
parent 7391389276
commit cdf9fdbc2d
1 changed files with 10 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
mochi_preview/dit/joint_model/asymm_models_joint.py
View File

@ -169,6 +169,12 @@ class AsymmetricAttention(nn.Module):
pass pass
else: else:
from .layers import RMSNorm from .layers import RMSNorm
if rms_norm_func == "apex":
self.q_norm_x = RMSNorm(self.head_dim)
self.k_norm_x = RMSNorm(self.head_dim)
self.q_norm_y = RMSNorm(self.head_dim)
self.k_norm_y = RMSNorm(self.head_dim)
else:
self.q_norm_x = RMSNorm(self.head_dim, device=device) self.q_norm_x = RMSNorm(self.head_dim, device=device)
self.k_norm_x = RMSNorm(self.head_dim, device=device) self.k_norm_x = RMSNorm(self.head_dim, device=device)
self.q_norm_y = RMSNorm(self.head_dim, device=device) self.q_norm_y = RMSNorm(self.head_dim, device=device)