From 66d5f342e883ad2350e7d623ae9317c98b850142 Mon Sep 17 00:00:00 2001
From: James Huang <syhuang1201@gmail.com>
Date: Tue, 14 Jul 2026 23:45:00 +0000
Subject: [PATCH] Fix pyink 23.10.0 formatting for files introduced in #443

---
 src/maxdiffusion/generate_wan.py              | 20 +++++------
 .../pipelines/wan/wan_pipeline_2_1.py         | 16 ++++-----
 .../pipelines/wan/wan_pipeline_i2v_2p2.py     | 24 ++++++-------
 src/maxdiffusion/tests/aot_cache_test.py      | 36 +++++++++----------
 4 files changed, 44 insertions(+), 52 deletions(-)

diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py
index 49f35f490..b3b1f5416 100644
--- a/src/maxdiffusion/generate_wan.py
+++ b/src/maxdiffusion/generate_wan.py
@@ -390,17 +390,15 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
     vae_decode_total = trace.get("vae_decode", 0.0)
     vae_decode_tpu = trace.get("vae_decode_tpu", 0.0)
     vae_decode_post = vae_decode_total - vae_decode_tpu
-    summary.extend(
-        [
-            f"  {'─' * 40}",
-            f"  Conditioning:        {trace.get('conditioning', 0.0):>7.1f}s",
-            f"    - VAE Encode:      {trace.get('vae_encode', 0.0):>7.1f}s",
-            f"  Denoise Total:       {trace.get('denoise_total', 0.0):>7.1f}s",
-            f"  VAE Decode:          {vae_decode_total:>7.1f}s",
-            f"    - TPU Compute:     {vae_decode_tpu:>7.1f}s",
-            f"    - Host Formatting: {vae_decode_post:>7.1f}s",
-        ]
-    )
+    summary.extend([
+        f"  {'─' * 40}",
+        f"  Conditioning:        {trace.get('conditioning', 0.0):>7.1f}s",
+        f"    - VAE Encode:      {trace.get('vae_encode', 0.0):>7.1f}s",
+        f"  Denoise Total:       {trace.get('denoise_total', 0.0):>7.1f}s",
+        f"  VAE Decode:          {vae_decode_total:>7.1f}s",
+        f"    - TPU Compute:     {vae_decode_tpu:>7.1f}s",
+        f"    - Host Formatting: {vae_decode_post:>7.1f}s",
+    ])
   summary.append(f"{'=' * 50}")
   max_logging.log("\n".join(summary))
 
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py
index 54629d181..f5851efeb 100644
--- a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py
+++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py
@@ -290,15 +290,13 @@ def run_inference_2_1(
   transformer_obj = nnx.merge(graphdef, sharded_state, rest_of_state)
 
   # Compute RoPE once as it only depends on shape
-  dummy_hidden_states = jnp.zeros(
-      (
-          latents.shape[0],
-          latents.shape[2],
-          latents.shape[3],
-          latents.shape[4],
-          latents.shape[1],
-      )
-  )
+  dummy_hidden_states = jnp.zeros((
+      latents.shape[0],
+      latents.shape[2],
+      latents.shape[3],
+      latents.shape[4],
+      latents.shape[1],
+  ))
   rotary_emb = transformer_obj.rope(dummy_hidden_states)
 
   kv_cache = None
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py
index f156fedd7..b1c09d146 100644
--- a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py
+++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py
@@ -971,19 +971,17 @@ def scan_body(carry, t):
     # tracing both 14B branches per step and keeps the AOT cache usable.
     use_high_noise = bool(np.asarray(scheduler_state.timesteps)[step] >= np.asarray(boundary))
     branch = high_noise_branch if use_high_noise else low_noise_branch
-    noise_pred, _ = branch(
-        (
-            latent_model_input,
-            timestep,
-            prompt_embeds_combined,
-            image_embeds_combined,
-            kv_cache_high,
-            kv_cache_low,
-            rotary_emb,
-            encoder_attention_mask_high,
-            encoder_attention_mask_low,
-        )
-    )
+    noise_pred, _ = branch((
+        latent_model_input,
+        timestep,
+        prompt_embeds_combined,
+        image_embeds_combined,
+        kv_cache_high,
+        kv_cache_low,
+        rotary_emb,
+        encoder_attention_mask_high,
+        encoder_attention_mask_low,
+    ))
     noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1))
     latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
 
diff --git a/src/maxdiffusion/tests/aot_cache_test.py b/src/maxdiffusion/tests/aot_cache_test.py
index bdfa0a589..2154b26ac 100644
--- a/src/maxdiffusion/tests/aot_cache_test.py
+++ b/src/maxdiffusion/tests/aot_cache_test.py
@@ -137,25 +137,23 @@ def test_signature_deterministic_across_processes(self):
     import subprocess
     import sys
 
-    snippet = "\n".join(
-        (
-            "import os",
-            "os.environ['JAX_PLATFORMS'] = 'cpu'",
-            "import jax",
-            "import jax.numpy as jnp",
-            "from flax import nnx",
-            "from maxdiffusion import aot_cache",
-            "",
-            "class T(nnx.Module):",
-            "  def __init__(self, rngs):",
-            "    self.lin = nnx.Linear(4, 4, rngs=rngs)",
-            "",
-            "graphdef, state = nnx.split(T(nnx.Rngs(0)))",
-            "sig = aot_cache._dynamic_signature(",
-            "    (graphdef, state.to_pure_dict(), jnp.ones((2, 4))), {})",
-            "print(sig)",
-        )
-    )
+    snippet = "\n".join((
+        "import os",
+        "os.environ['JAX_PLATFORMS'] = 'cpu'",
+        "import jax",
+        "import jax.numpy as jnp",
+        "from flax import nnx",
+        "from maxdiffusion import aot_cache",
+        "",
+        "class T(nnx.Module):",
+        "  def __init__(self, rngs):",
+        "    self.lin = nnx.Linear(4, 4, rngs=rngs)",
+        "",
+        "graphdef, state = nnx.split(T(nnx.Rngs(0)))",
+        "sig = aot_cache._dynamic_signature(",
+        "    (graphdef, state.to_pure_dict(), jnp.ones((2, 4))), {})",
+        "print(sig)",
+    ))
     outs = [
         subprocess.run(
             [sys.executable, "-c", snippet],