shenglin.net/ComfyUI-MochiWrapper-2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update asymm_models_joint.py

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

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

@ -169,16 +169,14 @@ class AsymmetricAttention(nn.Module):
pass
else:
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.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)
norm_kwargs = {}
if rms_norm_func != "apex":
norm_kwargs['device'] = device
self.q_norm_x = RMSNorm(self.head_dim, **norm_kwargs)
self.k_norm_x = RMSNorm(self.head_dim, **norm_kwargs)
self.q_norm_y = RMSNorm(self.head_dim, **norm_kwargs)
self.k_norm_y = RMSNorm(self.head_dim, **norm_kwargs)
# 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)