Apply pyink 23.10.0 formatting to 5 FLUX.2-klein source files

Author: amepas

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

@@ -722,7 +722,6 @@ def init_weights(self, rngs, max_sequence_length, eval_only=True):
       )["params"]
 
 
-
 class FlaxSwiGluFeedForward(nn.Module):
   dim: int
   hidden_dim: int
@@ -1717,29 +1716,33 @@ def __init__(
         rngs=rngs,
     )
 
-    self.double_blocks = nnx.List([
-        NNXFluxDoubleTransformerBlock(
-            rngs=rngs,
-            dim=self.inner_dim,
-            num_attention_heads=num_attention_heads,
-            attention_head_dim=attention_head_dim,
-            dtype=dtype,
-            weights_dtype=weights_dtype,
-        )
-        for _ in range(num_layers)
-    ])
-
-    self.single_blocks = nnx.List([
-        NNXFluxSingleTransformerBlock(
-            rngs=rngs,
-            dim=self.inner_dim,
-            num_attention_heads=num_attention_heads,
-            attention_head_dim=attention_head_dim,
-            dtype=dtype,
-            weights_dtype=weights_dtype,
-        )
-        for _ in range(num_single_layers)
-    ])
+    self.double_blocks = nnx.List(
+        [
+            NNXFluxDoubleTransformerBlock(
+                rngs=rngs,
+                dim=self.inner_dim,
+                num_attention_heads=num_attention_heads,
+                attention_head_dim=attention_head_dim,
+                dtype=dtype,
+                weights_dtype=weights_dtype,
+            )
+            for _ in range(num_layers)
+        ]
+    )
+
+    self.single_blocks = nnx.List(
+        [
+            NNXFluxSingleTransformerBlock(
+                rngs=rngs,
+                dim=self.inner_dim,
+                num_attention_heads=num_attention_heads,
+                attention_head_dim=attention_head_dim,
+                dtype=dtype,
+                weights_dtype=weights_dtype,
+            )
+            for _ in range(num_single_layers)
+        ]
+    )
 
     self.norm_out = NNXAdaLayerNormContinuous(
         rngs=rngs,

src/maxdiffusion/models/normalization_flax.py

@@ -197,7 +197,12 @@ def __init__(
         num_features=embedding_dim, epsilon=eps, use_bias=False, use_scale=False, dtype=dtype, rngs=rngs
     )
     self.linear = nnx.Linear(
-        in_features=embedding_dim, out_features=embedding_dim * 2, use_bias=True, dtype=dtype, param_dtype=weights_dtype, rngs=rngs
+        in_features=embedding_dim,
+        out_features=embedding_dim * 2,
+        use_bias=True,
+        dtype=dtype,
+        param_dtype=weights_dtype,
+        rngs=rngs,
     )
 
   def __call__(self, x: jax.Array, conditioning_embedding: jax.Array) -> jax.Array:
@@ -209,7 +214,9 @@ def __call__(self, x: jax.Array, conditioning_embedding: jax.Array) -> jax.Array
 
 class NNXAdaLayerNormZero(nnx.Module):
 
-  def __init__(self, embedding_dim: int, eps: float = 1e-6, dtype: jnp.dtype = jnp.float32, weights_dtype: jnp.dtype = jnp.float32):
+  def __init__(
+      self, embedding_dim: int, eps: float = 1e-6, dtype: jnp.dtype = jnp.float32, weights_dtype: jnp.dtype = jnp.float32
+  ):
     self.embedding_dim = embedding_dim
     self.eps = eps
     self.dtype = dtype
@@ -227,7 +234,9 @@ def __call__(self, x: jax.Array, emb: jax.Array):
 
 class NNXAdaLayerNormZeroSingle(nnx.Module):
 
-  def __init__(self, embedding_dim: int, eps: float = 1e-6, dtype: jnp.dtype = jnp.float32, weights_dtype: jnp.dtype = jnp.float32):
+  def __init__(
+      self, embedding_dim: int, eps: float = 1e-6, dtype: jnp.dtype = jnp.float32, weights_dtype: jnp.dtype = jnp.float32
+  ):
     self.embedding_dim = embedding_dim
     self.eps = eps
     self.dtype = dtype

src/maxdiffusion/models/qwen3_flax.py

@@ -587,15 +587,11 @@ def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config):
     self.embed_tokens = nnx.Embed(
         num_embeddings=config.vocab_size,
         features=config.hidden_size,
-        embedding_init=nnx.with_partitioning(
-            nnx.initializers.normal(stddev=config.hidden_size**-0.5), ("vocab", "embed")
-        ),
+        embedding_init=nnx.with_partitioning(nnx.initializers.normal(stddev=config.hidden_size**-0.5), ("vocab", "embed")),
         dtype=config.dtype,
         rngs=rngs,
     )
-    self.layers = nnx.List(
-        [NNXFlaxQwen3DecoderLayer(rngs=rngs, config=config) for _ in range(config.num_hidden_layers)]
-    )
+    self.layers = nnx.List([NNXFlaxQwen3DecoderLayer(rngs=rngs, config=config) for _ in range(config.num_hidden_layers)])
     self.norm = NNXFlaxQwen3RMSNorm(rngs=rngs, dim=config.hidden_size, eps=config.rms_norm_eps, dtype=config.dtype)
 
   def __call__(

src/maxdiffusion/models/vae_flax.py

@@ -1006,7 +1006,9 @@ def __init__(
         rngs=rngs,
     )
 
-  def decode(self, latents: jax.Array, decoder_params: Optional[dict] = None, deterministic: bool = True, return_dict: bool = True):
+  def decode(
+      self, latents: jax.Array, decoder_params: Optional[dict] = None, deterministic: bool = True, return_dict: bool = True
+  ):
     if latents.shape[-1] != self.latent_channels:
       latents = jnp.transpose(latents, (0, 2, 3, 1))
 

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -215,10 +215,7 @@ def __call__(
         tokenizer = Qwen2TokenizerFast.from_pretrained(tokenizer_path, subfolder="tokenizer", local_files_only=True)
 
       # Tokenize using deterministic explicit template string (version-agnostic across transformers versions)
-      templated_texts = [
-          f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
-          for p in prompts
-      ]
+      templated_texts = [f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n" for p in prompts]
       max_logging.log(f"DEBUG: templated_texts[0] = {repr(templated_texts[0])}")
       inputs = tokenizer(templated_texts, return_tensors="np", padding="max_length", truncation=True, max_length=seq_len_txt)
       prompt_ids = jnp.array(inputs["input_ids"])
@@ -246,7 +243,9 @@ def __call__(
       # ---------------------------------------------------------------------
       # PHASE B: Denoising Loop (Flux Transformer)
       # ---------------------------------------------------------------------
-      max_logging.log(f"[PHASE B] Running {num_inference_steps}-step E2E Denoising Loop on a batch of {batch_size} images...")
+      max_logging.log(
+          f"[PHASE B] Running {num_inference_steps}-step E2E Denoising Loop on a batch of {batch_size} images..."
+      )
       t0 = time.perf_counter()
 
       guidance_vec_val = None
@@ -271,7 +270,9 @@ def __call__(
 
         # Print progress
         sigma_val = scheduler_state.sigmas[step_idx]
-        max_logging.log(f" -> Step {step_idx}: Timestep = {scheduler_state.timesteps[step_idx]:.4f}, Sigma = {sigma_val:.4f}")
+        max_logging.log(
+            f" -> Step {step_idx}: Timestep = {scheduler_state.timesteps[step_idx]:.4f}, Sigma = {sigma_val:.4f}"
+        )
 
       latents_jax.block_until_ready()