pyink test failures

Author: amepas

src/maxdiffusion/loaders/ltx2_lora_nnx_loader.py

@@ -76,8 +76,6 @@ def translate_fn(nnx_path_str):
     # the merge_fn warns about unmatched keys in each dict, so we only warn about any leftovers
     unmatched_keys = set(h_state_dict) - set(transformer_state_dict) - set(connector_state_dict)
     if unmatched_keys:
-      max_logging.log(
-          f"{len(unmatched_keys)} key(s) in LoRA dictionary routed to no merge target: {unmatched_keys}"
-      )
+      max_logging.log(f"{len(unmatched_keys)} key(s) in LoRA dictionary routed to no merge target: {unmatched_keys}")
 
     return pipeline

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

@@ -36,6 +36,7 @@
 HEAD = common_types.HEAD
 D_KV = common_types.D_KV
 
+
 @flax.struct.dataclass
 class Transformer2DModelOutput(BaseOutput):
   """
@@ -111,8 +112,8 @@ class FluxSingleTransformerBlock(nn.Module):
   dtype: jnp.dtype = jnp.float32
   weights_dtype: jnp.dtype = jnp.float32
   precision: jax.lax.Precision = None
-  use_global_modulation: bool = False # Added flag!
-  use_swiglu: bool = False # Added flag!
+  use_global_modulation: bool = False  # Added flag!
+  use_swiglu: bool = False  # Added flag!
 
   def setup(self):
     self.mlp_hidden_dim = int(self.dim * self.mlp_ratio)
@@ -170,7 +171,7 @@ def __call__(self, hidden_states, temb=None, image_rotary_emb=None, temb_mod=Non
       shift_msa = jnp.expand_dims(shift_msa, axis=1)
       scale_msa = jnp.expand_dims(scale_msa, axis=1)
       gate = jnp.expand_dims(gate, axis=1)
-      
+
       norm_hidden_states = self.norm(hidden_states)
       norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa
     else:
@@ -244,16 +245,24 @@ class FluxTransformerBlock(nn.Module):
   mlp_ratio: float = 4.0
   qkv_bias: bool = False
   attention_kernel: str = "dot_product"
-  use_global_modulation: bool = False # Added flag!
-  use_swiglu: bool = False # Added flag!
+  use_global_modulation: bool = False  # Added flag!
+  use_swiglu: bool = False  # Added flag!
 
   def setup(self):
     if self.use_global_modulation:
-      self.img_norm1 = nn.LayerNorm(use_bias=False, use_scale=False, epsilon=self.eps, dtype=self.dtype, param_dtype=self.weights_dtype)
-      self.txt_norm1 = nn.LayerNorm(use_bias=False, use_scale=False, epsilon=self.eps, dtype=self.dtype, param_dtype=self.weights_dtype)
+      self.img_norm1 = nn.LayerNorm(
+          use_bias=False, use_scale=False, epsilon=self.eps, dtype=self.dtype, param_dtype=self.weights_dtype
+      )
+      self.txt_norm1 = nn.LayerNorm(
+          use_bias=False, use_scale=False, epsilon=self.eps, dtype=self.dtype, param_dtype=self.weights_dtype
+      )
     else:
-      self.img_norm1 = AdaLayerNormZero(self.dim, dtype=self.dtype, weights_dtype=self.weights_dtype, precision=self.precision)
-      self.txt_norm1 = AdaLayerNormZero(self.dim, dtype=self.dtype, weights_dtype=self.weights_dtype, precision=self.precision)
+      self.img_norm1 = AdaLayerNormZero(
+          self.dim, dtype=self.dtype, weights_dtype=self.weights_dtype, precision=self.precision
+      )
+      self.txt_norm1 = AdaLayerNormZero(
+          self.dim, dtype=self.dtype, weights_dtype=self.weights_dtype, precision=self.precision
+      )
 
     self.attn = FlaxFluxAttention(
         query_dim=self.dim,
@@ -285,27 +294,30 @@ def setup(self):
           name="img_mlp",
       )
     else:
-      self.img_mlp = nn.Sequential([
-          nn.Dense(
-              int(self.dim * self.mlp_ratio),
-              use_bias=True,
-              kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
-              bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
-              dtype=self.dtype,
-              param_dtype=self.weights_dtype,
-              precision=self.precision,
-          ),
-          nn.gelu,
-          nn.Dense(
-              self.dim,
-              use_bias=True,
-              kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")),
-              bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
-              dtype=self.dtype,
-              param_dtype=self.weights_dtype,
-              precision=self.precision,
-          ),
-      ], name="img_mlp")
+      self.img_mlp = nn.Sequential(
+          [
+              nn.Dense(
+                  int(self.dim * self.mlp_ratio),
+                  use_bias=True,
+                  kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
+                  bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
+                  dtype=self.dtype,
+                  param_dtype=self.weights_dtype,
+                  precision=self.precision,
+              ),
+              nn.gelu,
+              nn.Dense(
+                  self.dim,
+                  use_bias=True,
+                  kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")),
+                  bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
+                  dtype=self.dtype,
+                  param_dtype=self.weights_dtype,
+                  precision=self.precision,
+              ),
+          ],
+          name="img_mlp",
+      )
 
     self.txt_norm2 = nn.LayerNorm(
         use_bias=False,
@@ -325,56 +337,60 @@ def setup(self):
           name="txt_mlp",
       )
     else:
-      self.txt_mlp = nn.Sequential([
-          nn.Dense(
-              int(self.dim * self.mlp_ratio),
-              use_bias=True,
-              kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
-              bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
-              dtype=self.dtype,
-              param_dtype=self.weights_dtype,
-              precision=self.precision,
-          ),
-          nn.gelu,
-          nn.Dense(
-              self.dim,
-              use_bias=True,
-              kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")),
-              bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
-              dtype=self.dtype,
-              param_dtype=self.weights_dtype,
-              precision=self.precision,
-          ),
-      ], name="txt_mlp")
+      self.txt_mlp = nn.Sequential(
+          [
+              nn.Dense(
+                  int(self.dim * self.mlp_ratio),
+                  use_bias=True,
+                  kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
+                  bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
+                  dtype=self.dtype,
+                  param_dtype=self.weights_dtype,
+                  precision=self.precision,
+              ),
+              nn.gelu,
+              nn.Dense(
+                  self.dim,
+                  use_bias=True,
+                  kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")),
+                  bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
+                  dtype=self.dtype,
+                  param_dtype=self.weights_dtype,
+                  precision=self.precision,
+              ),
+          ],
+          name="txt_mlp",
+      )
 
     # let chunk size default to None
     self._chunk_size = None
     self._chunk_dim = 0
 
-  def __call__(self, hidden_states, encoder_hidden_states, temb=None, image_rotary_emb=None,
-               temb_mod_img=None, temb_mod_txt=None):
+  def __call__(
+      self, hidden_states, encoder_hidden_states, temb=None, image_rotary_emb=None, temb_mod_img=None, temb_mod_txt=None
+  ):
     if self.use_global_modulation:
       (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp) = jnp.split(temb_mod_img, 6, axis=-1)
       (c_shift_msa, c_scale_msa, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp) = jnp.split(temb_mod_txt, 6, axis=-1)
-      
+
       # Unsqueeze sequence dimension for broadcasting when batch_size > 1
       shift_msa = jnp.expand_dims(shift_msa, axis=1)
       scale_msa = jnp.expand_dims(scale_msa, axis=1)
       gate_msa = jnp.expand_dims(gate_msa, axis=1)
       shift_mlp = jnp.expand_dims(shift_mlp, axis=1)
       scale_mlp = jnp.expand_dims(scale_mlp, axis=1)
       gate_mlp = jnp.expand_dims(gate_mlp, axis=1)
-      
+
       c_shift_msa = jnp.expand_dims(c_shift_msa, axis=1)
       c_scale_msa = jnp.expand_dims(c_scale_msa, axis=1)
       c_gate_msa = jnp.expand_dims(c_gate_msa, axis=1)
       c_shift_mlp = jnp.expand_dims(c_shift_mlp, axis=1)
       c_scale_mlp = jnp.expand_dims(c_scale_mlp, axis=1)
       c_gate_mlp = jnp.expand_dims(c_gate_mlp, axis=1)
-      
+
       norm_hidden_states = self.img_norm1(hidden_states)
       norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa
-      
+
       norm_encoder_hidden_states = self.txt_norm1(encoder_hidden_states)
       norm_encoder_hidden_states = (1 + c_scale_msa) * norm_encoder_hidden_states + c_shift_msa
     else:
@@ -465,8 +481,8 @@ class FluxTransformer2DModel(nn.Module, FlaxModelMixin, ConfigMixin):
   joint_attention_bias: bool = True
   x_embedder_bias: bool = True
   proj_out_bias: bool = True
-  use_global_modulation: bool = False # Added config flag!
-  use_swiglu: bool = False # Added config flag!
+  use_global_modulation: bool = False  # Added config flag!
+  use_swiglu: bool = False  # Added config flag!
 
   def setup(self):
     self.out_channels = self.in_channels
@@ -699,7 +715,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"))
-    
+
     for single_block in self.single_blocks:
       hidden_states = single_block(
           hidden_states=hidden_states,