ruff and code_style

Author: ksikiric

src/maxdiffusion/generate_flux.py

@@ -135,19 +135,7 @@ def get_lin_function(x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: flo
 
 
 def run_inference(
-    states,
-    transformer,
-    vae,
-    config,
-    mesh,
-    latents,
-    latent_image_ids,
-    prompt_embeds,
-    txt_ids,
-    vec,
-    guidance_vec,
-    c_ts,
-    p_ts
+    states, transformer, vae, config, mesh, latents, latent_image_ids, prompt_embeds, txt_ids, vec, guidance_vec, c_ts, p_ts
 ):
 
   transformer_state = states["transformer"]
@@ -468,7 +456,7 @@ def validate_inputs(latents, latent_image_ids, prompt_embeds, text_ids, timestep
           vec=pooled_prompt_embeds,
           guidance_vec=guidance,
           c_ts=c_ts,
-          p_ts=p_ts
+          p_ts=p_ts,
       ),
       in_shardings=(state_shardings,),
       out_shardings=None,

src/maxdiffusion/models/attention_flax.py

@@ -55,19 +55,20 @@ class AttentionOp(nn.Module):
   def setup(self):
     if self.attention_kernel == "cudnn_flash_te":
       from transformer_engine.jax.flax.transformer import DotProductAttention  # pytype: disable=import-error
+
       self.dpa_layer = DotProductAttention(
-        head_dim=self.dim_head,
-        num_attention_heads=self.heads,
-        num_gqa_groups=self.heads,
-        attn_mask_type="no_mask",  # 'no_mask', 'padding', 'causal', or 'padding_causal'
-        attn_bias_type="NO_BIAS",  # 'no_bias', 'pre_scale_bias' or 'post_scale_bias'
-        # attention_dropout=self.dropout_rate,
-        dropout_rng_name="aqt",
-        dtype=self.dtype,
-        # float32_logits=self.float32_logits,
-        qkv_layout="BSHD_BSHD_BSHD",  # 'BS3HD', 'BSHD_BS2HD' or 'BSHD_BSHD_BSHD'
-        scale_factor=self.scale,
-        transpose_batch_sequence=False,
+          head_dim=self.dim_head,
+          num_attention_heads=self.heads,
+          num_gqa_groups=self.heads,
+          attn_mask_type="no_mask",  # 'no_mask', 'padding', 'causal', or 'padding_causal'
+          attn_bias_type="NO_BIAS",  # 'no_bias', 'pre_scale_bias' or 'post_scale_bias'
+          # attention_dropout=self.dropout_rate,
+          dropout_rng_name="aqt",
+          dtype=self.dtype,
+          # float32_logits=self.float32_logits,
+          qkv_layout="BSHD_BSHD_BSHD",  # 'BS3HD', 'BSHD_BS2HD' or 'BSHD_BSHD_BSHD'
+          scale_factor=self.scale,
+          transpose_batch_sequence=False,
       )
 
   def check_attention_inputs(self, query: Array, key: Array, value: Array) -> None:

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

@@ -513,9 +513,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, temb=temb, image_rotary_emb=image_rotary_emb
-      )
+      hidden_states = 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/modeling_flax_pytorch_utils.py

@@ -278,61 +278,6 @@ def convert_flux_lora_pytorch_state_dict_to_flax(config, pt_state_dict, params,
   return params, rank, network_alphas
 
 
-def convert_flux_lora_pytorch_state_dict_to_flax(config, pt_state_dict, params, adapter_name):
-  pt_state_dict = {k: v.float().numpy() for k, v in pt_state_dict.items()}
-  transformer_params = flatten_dict(unfreeze(params["transformer"]))
-  network_alphas = {}
-  rank = None
-  for pt_key, tensor in pt_state_dict.items():
-    renamed_pt_key = rename_key(pt_key)
-    renamed_pt_key = renamed_pt_key.replace("lora_unet_", "")
-    renamed_pt_key = renamed_pt_key.replace("lora_down", f"lora-{adapter_name}.down")
-    renamed_pt_key = renamed_pt_key.replace("lora_up", f"lora-{adapter_name}.up")
-
-    if "double_blocks" in renamed_pt_key:
-      renamed_pt_key = renamed_pt_key.replace("double_blocks.", "double_blocks_")
-      renamed_pt_key = renamed_pt_key.replace("processor.proj_lora1.down", f"attn.i_proj.lora-{adapter_name}.down")
-      renamed_pt_key = renamed_pt_key.replace("processor.proj_lora1.up", f"attn.i_proj.lora-{adapter_name}.up")
-      renamed_pt_key = renamed_pt_key.replace("processor.proj_lora2.down", f"attn.e_proj.lora-{adapter_name}.down")
-      renamed_pt_key = renamed_pt_key.replace("processor.proj_lora2.up", f"attn.e_proj.lora-{adapter_name}.up")
-      renamed_pt_key = renamed_pt_key.replace("processor.qkv_lora1.down", f"attn.i_qkv.lora-{adapter_name}.down")
-      renamed_pt_key = renamed_pt_key.replace("processor.qkv_lora1.up", f"attn.i_qkv.lora-{adapter_name}.up")
-      renamed_pt_key = renamed_pt_key.replace("processor.qkv_lora2.down", f"attn.e_qkv.lora-{adapter_name}.down")
-      renamed_pt_key = renamed_pt_key.replace("processor.qkv_lora2.up", f"attn.e_qkv.lora-{adapter_name}.up")
-
-      renamed_pt_key = renamed_pt_key.replace("_img_attn_proj", ".attn.i_proj")
-      renamed_pt_key = renamed_pt_key.replace("_img_attn_qkv", ".attn.i_qkv")
-      renamed_pt_key = renamed_pt_key.replace("_img_mlp_0", ".img_mlp.layers_0")
-      renamed_pt_key = renamed_pt_key.replace("_img_mlp_2", ".img_mlp.layers_2")
-      renamed_pt_key = renamed_pt_key.replace("_img_mod_lin", ".img_norm1.lin")
-      renamed_pt_key = renamed_pt_key.replace("_txt_attn_proj", ".attn.e_proj")
-      renamed_pt_key = renamed_pt_key.replace("_txt_attn_qkv", ".attn.e_qkv")
-      renamed_pt_key = renamed_pt_key.replace("_txt_mlp_0", ".txt_mlp.layers_0")
-      renamed_pt_key = renamed_pt_key.replace("_txt_mlp_2", ".txt_mlp.layers_2")
-      renamed_pt_key = renamed_pt_key.replace("_txt_mod_lin", ".txt_norm1.lin")
-    elif "single_blocks" in renamed_pt_key:
-      renamed_pt_key = renamed_pt_key.replace("_linear1", ".linear1")
-      renamed_pt_key = renamed_pt_key.replace("_linear2", ".linear2")
-      renamed_pt_key = renamed_pt_key.replace("_modulation_lin", ".norm.lin")
-
-    renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
-
-    pt_tuple_key = tuple(renamed_pt_key.split("."))
-    if "alpha" in pt_tuple_key:
-      pt_tuple_key = pt_tuple_key[:-1] + (f"lora-{adapter_name}", "down", "kernel")
-      network_alphas[tuple([*pt_tuple_key])] = tensor.item()  # noqa: C409
-      pt_tuple_key = pt_tuple_key[:-1] + (f"lora-{adapter_name}", "up", "kernel")
-      network_alphas[tuple([*pt_tuple_key])] = tensor.item()  # noqa: C409
-    else:
-      if pt_tuple_key[-2] == "up":
-        rank = tensor.shape[1]
-      transformer_params[tuple([*pt_tuple_key])] = jnp.asarray(tensor.T, dtype=config.weights_dtype)  # noqa: C409
-
-  params["transformer"] = unflatten_dict(transformer_params)
-
-  return params, rank, network_alphas
-
-
 def convert_lora_pytorch_state_dict_to_flax(pt_state_dict, params, network_alphas, adapter_name):
   # Step 1: Convert pytorch tensor to numpy
   # sometimes we load weights in bf16 and numpy doesn't support it