andreas/ComfyUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Switch mochi and wan modes to use pytorch RMSNorm. (#7925)

Browse Source 
* Switch genmo model to native RMSNorm.

* Switch WAN to native RMSNorm.
This commit is contained in:
comfyanonymous
2025-05-03 16:07:55 -07:00
committed by GitHub
parent 7689917113
commit 3041e5c354
3 changed files with 7 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
comfy/ldm/genmo/joint_model/asymm_models_joint.py
View File

@@ -13,7 +13,6 @@ from comfy.ldm.modules.attention import optimized_attention
from .layers import (
FeedForward,
PatchEmbed,
RMSNorm,
TimestepEmbedder,
)
@@ -90,10 +89,10 @@ class AsymmetricAttention(nn.Module):
# Query and key normalization for stability.
assert qk_norm
self.q_norm_x = RMSNorm(self.head_dim, device=device, dtype=dtype)
self.k_norm_x = RMSNorm(self.head_dim, device=device, dtype=dtype)
self.q_norm_y = RMSNorm(self.head_dim, device=device, dtype=dtype)
self.k_norm_y = RMSNorm(self.head_dim, device=device, dtype=dtype)
self.q_norm_x = operations.RMSNorm(self.head_dim, eps=1e-5, device=device, dtype=dtype)
self.k_norm_x = operations.RMSNorm(self.head_dim, eps=1e-5, device=device, dtype=dtype)
self.q_norm_y = operations.RMSNorm(self.head_dim, eps=1e-5, device=device, dtype=dtype)
self.k_norm_y = operations.RMSNorm(self.head_dim, eps=1e-5, device=device, dtype=dtype)
# Output layers. y features go back down from dim_x -> dim_y.
self.proj_x = operations.Linear(dim_x, dim_x, bias=out_bias, device=device, dtype=dtype)