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

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

@ -169,10 +169,16 @@ class AsymmetricAttention(nn.Module):
pass
else:
from .layers import RMSNorm
self.q_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.k_norm_y = RMSNorm(self.head_dim, device=device)
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.k_norm_x = RMSNorm(self.head_dim, device=device)
self.q_norm_y = RMSNorm(self.head_dim, device=device)
self.k_norm_y = RMSNorm(self.head_dim, device=device)
# Output layers. y features go back down from dim_x -> dim_y.
self.proj_x = nn.Linear(dim_x, dim_x, bias=out_bias, device=device)