fix repeated double and single blocks.

Author: jfacevedo-google

src/maxdiffusion/configs/base_flux_dev.yml

@@ -55,16 +55,7 @@ precision: "DEFAULT"
 from_pt: True
 split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash
-flash_block_sizes: {
-  "block_q" : 128,
-  "block_kv" : 128,
-  "block_kv_compute" : 128,
-  "block_q_dkv" : 128,
-  "block_kv_dkv" : 128,
-  "block_kv_dkv_compute" : 128,
-  "block_q_dq" : 128,
-  "block_kv_dq" : 128
-}
+flash_block_sizes: {}
 # GroupNorm groups
 norm_num_groups: 32
 

src/maxdiffusion/configs/base_flux_schnell.yml

@@ -54,16 +54,7 @@ precision: "DEFAULT"
 from_pt: True
 split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash
-flash_block_sizes: {
-  "block_q" : 128,
-  "block_kv" : 128,
-  "block_kv_compute" : 128,
-  "block_q_dkv" : 128,
-  "block_kv_dkv" : 128,
-  "block_kv_dkv_compute" : 128,
-  "block_q_dq" : 128,
-  "block_kv_dq" : 128
-}
+flash_block_sizes: {}
 # GroupNorm groups
 norm_num_groups: 32
 

src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py

@@ -383,48 +383,43 @@ def setup(self):
         precision=self.precision,
     )
 
-    self.double_blocks = nn.Sequential(
-        [
-            *[
-                FluxTransformerBlock(
-                    dim=self.inner_dim,
-                    num_attention_heads=self.num_attention_heads,
-                    attention_head_dim=self.attention_head_dim,
-                    attention_kernel=self.attention_kernel,
-                    flash_min_seq_length=self.flash_min_seq_length,
-                    flash_block_sizes=self.flash_block_sizes,
-                    mesh=self.mesh,
-                    dtype=self.dtype,
-                    weights_dtype=self.weights_dtype,
-                    precision=self.precision,
-                    mlp_ratio=self.mlp_ratio,
-                    qkv_bias=self.qkv_bias,
-                )
-                for _ in range(self.num_layers)
-            ]
-        ]
-    )
-
-    self.single_blocks = nn.Sequential(
-        [
-            *[
-                FluxSingleTransformerBlock(
-                    dim=self.inner_dim,
-                    num_attention_heads=self.num_attention_heads,
-                    attention_head_dim=self.attention_head_dim,
-                    attention_kernel=self.attention_kernel,
-                    flash_min_seq_length=self.flash_min_seq_length,
-                    flash_block_sizes=self.flash_block_sizes,
-                    mesh=self.mesh,
-                    dtype=self.dtype,
-                    weights_dtype=self.weights_dtype,
-                    precision=self.precision,
-                    mlp_ratio=self.mlp_ratio,
-                )
-                for _ in range(self.num_single_layers)
-            ]
-        ]
-    )
+    double_blocks = []
+    for _ in range(self.num_layers):
+      double_block = FluxTransformerBlock(
+        dim=self.inner_dim,
+        num_attention_heads=self.num_attention_heads,
+        attention_head_dim=self.attention_head_dim,
+        attention_kernel=self.attention_kernel,
+        flash_min_seq_length=self.flash_min_seq_length,
+        flash_block_sizes=self.flash_block_sizes,
+        mesh=self.mesh,
+        dtype=self.dtype,
+        weights_dtype=self.weights_dtype,
+        precision=self.precision,
+        mlp_ratio=self.mlp_ratio,
+        qkv_bias=self.qkv_bias,
+      )
+      double_blocks.append(double_block)
+    self.double_blocks = double_blocks
+
+    single_blocks = []
+    for _ in range(self.num_single_layers):
+      single_block = FluxSingleTransformerBlock(
+        dim=self.inner_dim,
+        num_attention_heads=self.num_attention_heads,
+        attention_head_dim=self.attention_head_dim,
+        attention_kernel=self.attention_kernel,
+        flash_min_seq_length=self.flash_min_seq_length,
+        flash_block_sizes=self.flash_block_sizes,
+        mesh=self.mesh,
+        dtype=self.dtype,
+        weights_dtype=self.weights_dtype,
+        precision=self.precision,
+        mlp_ratio=self.mlp_ratio,
+      )
+      single_blocks.append(single_block)
+    
+    self.single_blocks = single_blocks
 
     self.norm_out = AdaLayerNormContinuous(
         self.inner_dim,
@@ -509,18 +504,19 @@ def __call__(
     image_rotary_emb = self.pe_embedder(ids)
     image_rotary_emb = nn.with_logical_constraint(image_rotary_emb, ("activation_batch", "activation_embed"))
 
-    hidden_states, encoder_hidden_states, temb, image_rotary_emb = self.double_blocks(
-        hidden_states=hidden_states,
-        encoder_hidden_states=encoder_hidden_states,
-        temb=temb,
-        image_rotary_emb=image_rotary_emb,
-    )
+    for double_block in self.double_blocks:
+      hidden_states, encoder_hidden_states, temb, image_rotary_emb = double_block(
+          hidden_states=hidden_states,
+          encoder_hidden_states=encoder_hidden_states,
+          temb=temb,
+          image_rotary_emb=image_rotary_emb,
+      )
     hidden_states = jnp.concatenate([encoder_hidden_states, hidden_states], axis=1)
     hidden_states = nn.with_logical_constraint(hidden_states, ("activation_batch", "activation_length", "activation_embed"))
-
-    hidden_states, temb, image_rotary_emb = self.single_blocks(
-        hidden_states=hidden_states, temb=temb, image_rotary_emb=image_rotary_emb
-    )
+    for single_block in self.single_blocks:
+      hidden_states, temb, image_rotary_emb = single_block(
+          hidden_states=hidden_states, temb=temb, image_rotary_emb=image_rotary_emb
+      )
     hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
 
     hidden_states = self.norm_out(hidden_states, temb)

src/maxdiffusion/models/flux/util.py

@@ -159,7 +159,6 @@ def load_flow_model(name: str, eval_shapes: dict, device: str, hf_download: bool
             for pt_key, tensor in tensors.items():
               renamed_pt_key = rename_key(pt_key)
               if "double_blocks" in renamed_pt_key:
-                renamed_pt_key = renamed_pt_key.replace("double_blocks_", "double_blocks.layers_")
                 renamed_pt_key = renamed_pt_key.replace("img_mlp_", "img_mlp.layers_")
                 renamed_pt_key = renamed_pt_key.replace("txt_mlp_", "txt_mlp.layers_")
                 renamed_pt_key = renamed_pt_key.replace("img_mod", "img_norm1")
@@ -176,7 +175,6 @@ def load_flow_model(name: str, eval_shapes: dict, device: str, hf_download: bool
                 renamed_pt_key = renamed_pt_key.replace("in_layer", "linear_1")
                 renamed_pt_key = renamed_pt_key.replace("out_layer", "linear_2")
               elif "single_blocks" in renamed_pt_key:
-                renamed_pt_key = renamed_pt_key.replace("single_blocks_", "single_blocks.layers_")
                 renamed_pt_key = renamed_pt_key.replace("modulation", "norm")
                 renamed_pt_key = renamed_pt_key.replace("norm.key_norm", "attn.key_norm")
                 renamed_pt_key = renamed_pt_key.replace("norm.query_norm", "attn.query_norm")
@@ -188,7 +186,6 @@ def load_flow_model(name: str, eval_shapes: dict, device: str, hf_download: bool
               elif "final_layer" in renamed_pt_key:
                 renamed_pt_key = renamed_pt_key.replace("final_layer.linear", "proj_out")
                 renamed_pt_key = renamed_pt_key.replace("final_layer.adaLN_modulation_1", "norm_out.Dense_0")
-
               pt_tuple_key = tuple(renamed_pt_key.split("."))
               flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, eval_shapes)
               flax_state_dict[flax_key] = jax.device_put(jnp.asarray(flax_tensor), device=cpu)