remove rmsnorm assert

Author: a-r-r-o-w

src/diffusers/models/transformers/transformer_chroma.py

@@ -235,7 +235,6 @@ def __init__(
             out_dim=dim,
             bias=True,
             processor=processor,
-            qk_norm="rms_norm",
             eps=1e-6,
             pre_only=True,
         )
@@ -296,7 +295,6 @@ def __init__(
             context_pre_only=False,
             bias=True,
             processor=FluxAttnProcessor(),
-            qk_norm=qk_norm,
             eps=eps,
         )
 

src/diffusers/models/transformers/transformer_flux.py

@@ -277,7 +277,6 @@ def __init__(
         dim_head: int = 64,
         dropout: float = 0.0,
         bias: bool = False,
-        qk_norm: Optional[str] = None,
         added_kv_proj_dim: Optional[int] = None,
         added_proj_bias: Optional[bool] = True,
         out_bias: bool = True,
@@ -289,7 +288,6 @@ def __init__(
         processor=None,
     ):
         super().__init__()
-        assert qk_norm == "rms_norm", "Flux uses RMSNorm"
 
         self.head_dim = dim_head
         self.inner_dim = out_dim if out_dim is not None else dim_head * heads
@@ -375,7 +373,6 @@ def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int,
             out_dim=dim,
             bias=True,
             processor=processor,
-            qk_norm="rms_norm",
             eps=1e-6,
             pre_only=True,
         )
@@ -431,7 +428,6 @@ def __init__(
             context_pre_only=False,
             bias=True,
             processor=FluxAttnProcessor(),
-            qk_norm=qk_norm,
             eps=eps,
         )