reformatted

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -1919,7 +1919,7 @@ def __call__(
 
     qkv_proj = self.qkv(hidden_states)
     qkv_proj = checkpoint_name(qkv_proj, "img_qkv_proj")
-    
+
     qkv_proj = qkv_proj.reshape(B, L, 3, H, D)
     query_proj, key_proj, value_proj = jnp.split(qkv_proj, 3, axis=2)
     query_proj = query_proj.squeeze(2)

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

@@ -30,7 +30,6 @@
 from ....common_types import BlockSizes
 from ....utils import BaseOutput
 from ...gradient_checkpoint import GradientCheckpointType, SKIP_GRADIENT_CHECKPOINT_KEY
-from jax import checkpoint_policies as cp
 from jax.ad_checkpoint import checkpoint_name
 
 AxisNames = common_types.AxisNames
@@ -83,9 +82,7 @@ def setup(self):
   def __call__(self, x, attn_output, gate, residual):
     mlp = self.lin_mlp(x)
     attn_mlp = jnp.concatenate([attn_output, self.mlp_act(mlp)], axis=2)
-    attn_mlp = nn.with_logical_constraint(
-        attn_mlp, ("activation_batch", None, "mlp")
-    )
+    attn_mlp = nn.with_logical_constraint(attn_mlp, ("activation_batch", None, "mlp"))
     hidden_states = self.linear2(attn_mlp)
     hidden_states = checkpoint_name(hidden_states, "lin2_hidden_states")
     hidden_states = gate * hidden_states
@@ -160,21 +157,19 @@ def setup(self):
 
   def __call__(self, hidden_states, temb, image_rotary_emb=None):
     residual = hidden_states
-    
+
     # FIX: Constrain inputs using valid config parameters (None skips sequence length axis parsing)
-    hidden_states = nn.with_logical_constraint(
-        hidden_states, ("activation_batch", None, "mlp")
-    )
-    
+    hidden_states = nn.with_logical_constraint(hidden_states, ("activation_batch", None, "mlp"))
+
     norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
-    
+
     qkv = self.lin_qkv(norm_hidden_states)
     qkv = checkpoint_name(qkv, "lin1_norm_hidden_states")
     qkv = nn.with_logical_constraint(qkv, ("activation_batch", None, "mlp"))
-    
+
     B, L = hidden_states.shape[:2]
     H, D, K = self.num_attention_heads, qkv.shape[-1] // (self.num_attention_heads * 3), 3
-    
+
     qkv_proj = qkv.reshape(B, L, K, H, D)
     q, k, v = jnp.split(qkv_proj, 3, axis=2)
     q = q.squeeze(2).swapaxes(1, 2)
@@ -196,7 +191,7 @@ def __call__(self, hidden_states, temb, image_rotary_emb=None):
     attn_output = checkpoint_name(attn_output, "attn_output")
 
     hidden_states = self.mlp_and_out(norm_hidden_states, attn_output, gate, residual)
-    
+
     if hidden_states.dtype == jnp.float16:
       hidden_states = jnp.clip(hidden_states, -65504, 65504)
 
@@ -308,12 +303,12 @@ def __call__(self, hidden_states, encoder_hidden_states, temb, image_rotary_emb=
     # --- IMAGE STREAM OPTIMIZATION (img_norm2) ---
     attn_output = gate_msa * attn_output
     hidden_states = hidden_states + attn_output
-    
+
     # Fully fused LayerNorm + scale_mlp + shift_mlp compilation block
     img_mean = jnp.mean(hidden_states, axis=-1, keepdims=True)
     img_var = jnp.mean(jnp.square(hidden_states - img_mean), axis=-1, keepdims=True)
     img_inv_std = jax.lax.rsqrt(img_var + self.eps)
-    
+
     norm_hidden_states = (hidden_states - img_mean) * img_inv_std * (1 + scale_mlp) + shift_mlp
     norm_hidden_states = nn.with_logical_constraint(norm_hidden_states, ("activation_batch", None, "mlp"))
 
@@ -334,48 +329,46 @@ def __call__(self, hidden_states, encoder_hidden_states, temb, image_rotary_emb=
 
     context_ff_output = self.txt_mlp(norm_encoder_hidden_states)
     encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output
-    
+
     # Safe numerical clipping limits for half precision math execution
     if encoder_hidden_states.dtype == jnp.float16 or encoder_hidden_states.dtype == jnp.bfloat16:
       encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
       hidden_states = hidden_states.clip(-65504, 65504)
 
     return hidden_states, encoder_hidden_states
 
+
 class ScannedDoubleBlockWrapper(nn.Module):
-    block_kwargs: dict
-    
-    @nn.compact
-    def __call__(self, carry, _):
-        hidden_states, encoder_hidden_states, temb, image_rotary_emb = carry
-        
-        # Instantiate the pure block (no remat here)
-        block = FluxTransformerBlock(**self.block_kwargs)
-        
-        h_out, e_out = block(
-            hidden_states=hidden_states,
-            encoder_hidden_states=encoder_hidden_states,
-            temb=temb,
-            image_rotary_emb=image_rotary_emb,
-        )
-        return (h_out, e_out, temb, image_rotary_emb), None
+  block_kwargs: dict
+
+  @nn.compact
+  def __call__(self, carry, _):
+    hidden_states, encoder_hidden_states, temb, image_rotary_emb = carry
+
+    # Instantiate the pure block (no remat here)
+    block = FluxTransformerBlock(**self.block_kwargs)
+
+    h_out, e_out = block(
+        hidden_states=hidden_states,
+        encoder_hidden_states=encoder_hidden_states,
+        temb=temb,
+        image_rotary_emb=image_rotary_emb,
+    )
+    return (h_out, e_out, temb, image_rotary_emb), None
 
 
 class ScannedSingleBlockWrapper(nn.Module):
-    block_kwargs: dict
-    
-    @nn.compact
-    def __call__(self, carry, _):
-        hidden_states, temb, image_rotary_emb = carry
-        
-        # Instantiate the pure block
-        block = FluxSingleTransformerBlock(**self.block_kwargs)
-        h_out = block(
-            hidden_states=hidden_states, 
-            temb=temb, 
-            image_rotary_emb=image_rotary_emb
-        )
-        return (h_out, temb, image_rotary_emb), None
+  block_kwargs: dict
+
+  @nn.compact
+  def __call__(self, carry, _):
+    hidden_states, temb, image_rotary_emb = carry
+
+    # Instantiate the pure block
+    block = FluxSingleTransformerBlock(**self.block_kwargs)
+    h_out = block(hidden_states=hidden_states, temb=temb, image_rotary_emb=image_rotary_emb)
+    return (h_out, temb, image_rotary_emb), None
+
 
 @flax_register_to_config
 class FluxTransformer2DModel(nn.Module, FlaxModelMixin, ConfigMixin):
@@ -446,25 +439,25 @@ def setup(self):
     self.gradient_checkpoint = GradientCheckpointType.from_str(self.remat_policy)
 
     # 2. Apply the policy to the Module classes
-    #RematDoubleBlock = self.gradient_checkpoint.apply_linen(FluxTransformerBlock)
-    #RematSingleBlock = self.gradient_checkpoint.apply_linen(FluxSingleTransformerBlock)
+    # RematDoubleBlock = self.gradient_checkpoint.apply_linen(FluxTransformerBlock)
+    # RematSingleBlock = self.gradient_checkpoint.apply_linen(FluxSingleTransformerBlock)
 
     # 1. Prepare the kwargs for the double blocks
     double_kwargs = {
-        '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,
-        'use_base2_exp': self.use_base2_exp,
-        'use_experimental_scheduler': self.use_experimental_scheduler,
+        "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,
+        "use_base2_exp": self.use_base2_exp,
+        "use_experimental_scheduler": self.use_experimental_scheduler,
     }
 
     double_policy = self.gradient_checkpoint.to_jax_policy(
@@ -480,27 +473,27 @@ def setup(self):
 
     self.scanned_double_blocks = nn.scan(
         RemattedDoubleWrapper,
-        variable_axes={'params': 0},
-        split_rngs={'params': True, 'dropout': True},
+        variable_axes={"params": 0},
+        split_rngs={"params": True, "dropout": True},
         length=self.num_layers,
-        metadata_params={'partition_name': None}
+        metadata_params={"partition_name": None},
     )(block_kwargs=double_kwargs)
 
     # 3. Define pure kwargs for single blocks
     single_kwargs = {
-        '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,
-        'use_base2_exp': self.use_base2_exp,
-        'use_experimental_scheduler': self.use_experimental_scheduler,
+        "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,
+        "use_base2_exp": self.use_base2_exp,
+        "use_experimental_scheduler": self.use_experimental_scheduler,
     }
 
     # 4. Force strict checkpointing on the Single Wrapper
@@ -517,10 +510,10 @@ def setup(self):
 
     self.scanned_single_blocks = nn.scan(
         RemattedSingleWrapper,
-        variable_axes={'params': 0},
-        split_rngs={'params': True, 'dropout': True},
+        variable_axes={"params": 0},
+        split_rngs={"params": True, "dropout": True},
         length=self.num_single_layers,
-        metadata_params={'partition_name': None}
+        metadata_params={"partition_name": None},
     )(block_kwargs=single_kwargs)
 
     self.norm_out = AdaLayerNormContinuous(
@@ -617,7 +610,7 @@ def __call__(
 
     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"))
-    
+
     # Execute the 38 Single Blocks
     carry = (hidden_states, temb, image_rotary_emb)
     carry, _ = self.scanned_single_blocks(carry, None)
@@ -640,23 +633,23 @@ def init_weights(self, rngs, max_sequence_length, eval_only=True):
     batch_size = 1 * num_devices
     batch_image_shape = (
         batch_size,
-        16, 
+        16,
         2 * resolution // scale_factor,
         2 * resolution // scale_factor,
     )
     text_shape = (
         batch_size,
         max_sequence_length,
-        4096, 
+        4096,
     )
     text_ids_shape = (
         batch_size,
         max_sequence_length,
-        3, 
+        3,
     )
     vec_shape = (
         batch_size,
-        768, 
+        768,
     )
     img = jnp.zeros(batch_image_shape, dtype=self.dtype)
     bs, _, h, w = img.shape
@@ -694,4 +687,4 @@ def init_weights(self, rngs, max_sequence_length, eval_only=True):
           pooled_projections=vec,
           timestep=t_vec,
           guidance=guidance_vec,
-      )["params"]
+      )["params"]

src/maxdiffusion/models/flux/util.py

@@ -204,8 +204,12 @@ def load_flow_model(name: str, eval_shapes: dict, device: str, hf_download: bool
             qkv_pt_key = renamed_pt_key.replace("linear1", "lin_qkv")
             mlp_pt_key = renamed_pt_key.replace("linear1", "mlp_and_out.lin_mlp")
 
-            flax_key_qkv, flax_tensor_qkv = rename_key_and_reshape_tensor(tuple(qkv_pt_key.split(".")), qkv_tensor, eval_shapes, scan_layers=True)
-            flax_key_mlp, flax_tensor_mlp = rename_key_and_reshape_tensor(tuple(mlp_pt_key.split(".")), mlp_tensor, eval_shapes, scan_layers=True)
+            flax_key_qkv, flax_tensor_qkv = rename_key_and_reshape_tensor(
+                tuple(qkv_pt_key.split(".")), qkv_tensor, eval_shapes, scan_layers=True
+            )
+            flax_key_mlp, flax_tensor_mlp = rename_key_and_reshape_tensor(
+                tuple(mlp_pt_key.split(".")), mlp_tensor, eval_shapes, scan_layers=True
+            )
 
             if flax_key_qkv not in single_blocks_tensors:
               single_blocks_tensors[flax_key_qkv] = {}

src/maxdiffusion/models/gradient_checkpoint.py

@@ -126,30 +126,25 @@ def apply(
     return nnx.remat(module, prevent_cse=prevent_cse, policy=policy, static_argnums=static_argnums)  # pylint: disable=invalid-name
 
   def apply_linen(
-        self,
-        module_class: type[linen_nn.Module],
-        names_which_can_be_saved: list = [],
-        names_which_can_be_offloaded: list = [],
-        static_argnums=(),
-        prevent_cse: bool = False,
-    ) -> type[linen_nn.Module]:
-      """
-      Applies a gradient checkpoint policy to a Linen module class.
-      If no policy is needed, it will return the module class as is.
-
-      Args:
-          module_class: the Linen Module class to apply the policy to
-
-      Returns:
-          The rematerialized Module class (or the original if no policy).
-      """
-      policy = self.to_jax_policy(names_which_can_be_saved, names_which_can_be_offloaded)
-      if policy == SKIP_GRADIENT_CHECKPOINT_KEY:
-        return module_class
-      
-      return linen_nn.remat(
-          module_class, 
-          prevent_cse=prevent_cse, 
-          policy=policy, 
-          static_argnums=static_argnums
-      )
+      self,
+      module_class: type[linen_nn.Module],
+      names_which_can_be_saved: list = [],
+      names_which_can_be_offloaded: list = [],
+      static_argnums=(),
+      prevent_cse: bool = False,
+  ) -> type[linen_nn.Module]:
+    """
+    Applies a gradient checkpoint policy to a Linen module class.
+    If no policy is needed, it will return the module class as is.
+
+    Args:
+        module_class: the Linen Module class to apply the policy to
+
+    Returns:
+        The rematerialized Module class (or the original if no policy).
+    """
+    policy = self.to_jax_policy(names_which_can_be_saved, names_which_can_be_offloaded)
+    if policy == SKIP_GRADIENT_CHECKPOINT_KEY:
+      return module_class
+
+    return linen_nn.remat(module_class, prevent_cse=prevent_cse, policy=policy, static_argnums=static_argnums)