up

Author: sayakpaul

src/diffusers/models/activations.py

@@ -100,11 +100,12 @@ def __init__(self, *args, **kwargs):
 
         activation = get_kernel("kernels-community/activation", revision="add_more_act")
         approximate = kwargs.get("approximate", "none")
+
+        super().__init__(*args, **kwargs)
         if approximate == "none":
-            self.act_fn = activation.gelu
+            self.act_fn = activation.layers.Gelu()
         elif approximate == "tanh":
-            self.act_fn = activation.gelu_tanh
-        super().__init__(*args, **kwargs)
+            self.act_fn = activation.layers.GeluTanh()
 
     def forward(self, hidden_states):
         hidden_states = self.proj(hidden_states)

src/diffusers/models/normalization.py

@@ -31,7 +31,7 @@
     from kernels import get_kernel
 
     activation = get_kernel("kernels-community/activation", revision="add_more_act")
-    silu_kernel = activation.silu
+    silu_kernel = activation.layers.Silu
 
 
 class AdaLayerNorm(nn.Module):
@@ -69,7 +69,7 @@ def __init__(
         if not DIFFUSERS_ENABLE_HUB_KERNELS:
             self.silu = nn.SiLU()
         else:
-            self.silu = silu_kernel
+            self.silu = silu_kernel()
         self.linear = nn.Linear(embedding_dim, output_dim)
         self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine)
 
@@ -159,7 +159,7 @@ def __init__(self, embedding_dim: int, num_embeddings: Optional[int] = None, nor
         if not DIFFUSERS_ENABLE_HUB_KERNELS:
             self.silu = nn.SiLU()
         else:
-            self.silu = silu_kernel
+            self.silu = silu_kernel()
         self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=bias)
         if norm_type == "layer_norm":
             self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
@@ -201,7 +201,7 @@ def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True):
         if not DIFFUSERS_ENABLE_HUB_KERNELS:
             self.silu = nn.SiLU()
         else:
-            self.silu = silu_kernel
+            self.silu = silu_kernel()
         self.linear = nn.Linear(embedding_dim, 3 * embedding_dim, bias=bias)
         if norm_type == "layer_norm":
             self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
@@ -356,7 +356,7 @@ def __init__(
         if not DIFFUSERS_ENABLE_HUB_KERNELS:
             self.silu = nn.SiLU()
         else:
-            self.silu = silu_kernel
+            self.silu = silu_kernel()
         self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
         if norm_type == "layer_norm":
             self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias)

src/diffusers/models/transformers/transformer_flux.py

@@ -45,7 +45,7 @@
     from kernels import get_kernel
 
     activation = get_kernel("kernels-community/activation", revision="add_more_act")
-    gelu_tanh_kernel = activation.gelu_tanh
+    gelu_tanh_kernel = activation.layers.GeluTanh
 
 
 def _get_projections(attn: "FluxAttention", hidden_states, encoder_hidden_states=None):
@@ -360,7 +360,7 @@ def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int,
         if not DIFFUSERS_ENABLE_HUB_KERNELS:
             self.act_mlp = nn.GELU(approximate="tanh")
         else:
-            self.act_mlp = gelu_tanh_kernel
+            self.act_mlp = gelu_tanh_kernel()
 
         self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)