adjustments to doublestreamblock

keras_hub/src/models/flux/flux_layers.py

@@ -181,56 +181,36 @@ def call(self, x):
         )
 
 
+
 class DoubleStreamBlock(keras.Model):
-    def __init__(
-        self,
-        hidden_size: int,
-        num_heads: int,
-        mlp_ratio: float,
-        qkv_bias: bool = False,
-    ):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
         super().__init__()
 
         mlp_hidden_dim = int(hidden_size * mlp_ratio)
         self.num_heads = num_heads
         self.hidden_size = hidden_size
+
         self.img_mod = Modulation(hidden_size, double=True)
-        self.img_norm1 = keras.layers.LayerNormalization(
-            elementwise_affine=False, eps=1e-6
-        )
-        self.img_attn = SelfAttention(
-            dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias
-        )
+        self.img_norm1 = keras.layers.LayerNormalization(epsilon=1e-6)
+        self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
 
-        self.img_norm2 = keras.layers.LayerNormalization(
-            elementwise_affine=False, eps=1e-6
-        )
-        self.img_mlp = keras.Sequential(
-            [
-                keras.layers.Dense(mlp_hidden_dim, use_bias=True),
-                keras.layers.Activation("tanh"),
-                keras.layers.Dense(hidden_size, use_bias=True),
-            ]
-        )
+        self.img_norm2 = keras.layers.LayerNormalization(epsilon=1e-6)
+        self.img_mlp = keras.Sequential([
+            keras.layers.Dense(mlp_hidden_dim, use_bias=True),
+            keras.layers.Activation("tanh"),
+            keras.layers.Dense(hidden_size, use_bias=True)
+        ])
 
         self.txt_mod = Modulation(hidden_size, double=True)
-        self.txt_norm1 = keras.layers.LayerNormalization(
-            elementwise_affine=False, eps=1e-6
-        )
-        self.txt_attn = SelfAttention(
-            dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias
-        )
+        self.txt_norm1 = keras.layers.LayerNormalization(epsilon=1e-6)
+        self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
 
-        self.txt_norm2 = keras.layers.LayerNormalization(
-            elementwise_affine=False, eps=1e-6
-        )
-        self.txt_mlp = keras.Sequential(
-            [
-                keras.layers.Dense(mlp_hidden_dim, use_bias=True),
-                keras.layers.Activation("tanh"),
-                keras.layers.Dense(hidden_size, use_bias=True),
-            ]
-        )
+        self.txt_norm2 = keras.layers.LayerNormalization(epsilon=1e-6)
+        self.txt_mlp = keras.Sequential([
+            keras.layers.Dense(mlp_hidden_dim, use_bias=True),
+            keras.layers.Activation("tanh"),
+            keras.layers.Dense(hidden_size, use_bias=True)
+        ])
 
     def call(self, img, txt, vec, pe):
         img_mod1, img_mod2 = self.img_mod(vec)
@@ -240,19 +220,15 @@ def call(self, img, txt, vec, pe):
         img_modulated = self.img_norm1(img)
         img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
         img_qkv = self.img_attn.qkv(img_modulated)
-        img_q, img_k, img_v = rearrange(
-            img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
-        )
-        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
+        img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        img_q, img_k = self.img_attn.norm(img_q, img_k)
 
         # prepare txt for attention
         txt_modulated = self.txt_norm1(txt)
         txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
         txt_qkv = self.txt_attn.qkv(txt_modulated)
-        txt_q, txt_k, txt_v = rearrange(
-            txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
-        )
-        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
+        txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k)
 
         # run actual attention
         q = keras.ops.concatenate((txt_q, img_q), axis=2)
@@ -264,13 +240,11 @@ def call(self, img, txt, vec, pe):
 
         # calculate the img bloks
         img = img + img_mod1.gate * self.img_attn.proj(img_attn)
-        img = img + img_mod2.gate * self.img_mlp(
-            (1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift
-        )
+        img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
 
         # calculate the txt bloks
         txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
-        txt = txt + txt_mod2.gate * self.txt_mlp(
-            (1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift
-        )
+        txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
         return img, txt
+
+