From 0a943b792d0a9750c71cf18b622e86ce69283319 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Sat, 26 Sep 2026 11:40:01 +0000 Subject: [PATCH] style: apply pyink 23.10.0 to 4 files on main that fail CI formatting check CPUTests.yml runs 'pyink src/maxdiffusion --check' with pyink==23.10.0; on main aot_cache.py, tests/aot_cache_test.py, pipelines/wan/wan_pipeline.py and tests/wan/svg_attention_test.py fail it. Formatting only, no code changes. --- src/maxdiffusion/aot_cache.py | 9 +++------ src/maxdiffusion/pipelines/wan/wan_pipeline.py | 10 ++-------- src/maxdiffusion/tests/aot_cache_test.py | 4 +--- src/maxdiffusion/tests/wan/svg_attention_test.py | 12 +++--------- 4 files changed, 9 insertions(+), 26 deletions(-) diff --git a/src/maxdiffusion/aot_cache.py b/src/maxdiffusion/aot_cache.py index 0e7cc344e..1c27c2c92 100644 --- a/src/maxdiffusion/aot_cache.py +++ b/src/maxdiffusion/aot_cache.py @@ -131,9 +131,8 @@ def extract_svg_meta(config: Any, pipeline: Any = None) -> dict[str, Any]: t = getattr(pipeline, attr, None) if t is not None: t_cfg = getattr(t, "config", None) - attn_cfg = ( - getattr(t_cfg, "attention_config", None) - or (t_cfg.get("attention_config") if isinstance(t_cfg, dict) else None) + attn_cfg = getattr(t_cfg, "attention_config", None) or ( + t_cfg.get("attention_config") if isinstance(t_cfg, dict) else None ) if isinstance(attn_cfg, dict): meta[f"{attr}_svg_config"] = json.dumps( @@ -162,9 +161,7 @@ def _dynamic_signature(args: tuple, kwargs: dict) -> str: Array leaves contribute shape/dtype; non-array leaves (python scalars, None flags) contribute an address-stripped repr. """ - leaves_with_paths = jax.tree_util.tree_flatten_with_path( - (args, kwargs), is_leaf=_is_graphdef - )[0] + leaves_with_paths = jax.tree_util.tree_flatten_with_path((args, kwargs), is_leaf=_is_graphdef)[0] parts = [] for path, leaf in leaves_with_paths: if _is_graphdef(leaf): diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline.py b/src/maxdiffusion/pipelines/wan/wan_pipeline.py index fd673ca30..4006e2fd1 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline.py @@ -1422,10 +1422,7 @@ def _has_svg_enabled(obj: Any) -> bool: return True cfg = getattr(obj, "config", None) if cfg is not None: - attn_cfg = ( - getattr(cfg, "attention_config", None) - or (cfg.get("attention_config") if isinstance(cfg, dict) else None) - ) + attn_cfg = getattr(cfg, "attention_config", None) or (cfg.get("attention_config") if isinstance(cfg, dict) else None) if isinstance(attn_cfg, dict) and bool(attn_cfg.get("use_svg_attention", False)): return True if hasattr(obj, "attributes") and hasattr(obj, "nodes"): @@ -1434,10 +1431,7 @@ def _has_svg_enabled(obj: Any) -> bool: return True if k == "config": val = getattr(v, "value", None) - ac = ( - getattr(val, "attention_config", None) - or (val.get("attention_config") if isinstance(val, dict) else None) - ) + ac = getattr(val, "attention_config", None) or (val.get("attention_config") if isinstance(val, dict) else None) if isinstance(ac, dict) and bool(ac.get("use_svg_attention", False)): return True return False diff --git a/src/maxdiffusion/tests/aot_cache_test.py b/src/maxdiffusion/tests/aot_cache_test.py index 25ddd1ca4..a0a43170d 100644 --- a/src/maxdiffusion/tests/aot_cache_test.py +++ b/src/maxdiffusion/tests/aot_cache_test.py @@ -237,9 +237,7 @@ def test_extract_svg_meta_changes_fingerprint_on_svg_changes(self): def test_step_array_preserves_dynamic_signature_across_steps(self): sigs = { - aot_cache._dynamic_signature( - (self._a,), {"svg_step_index": jnp.asarray(step, dtype=jnp.int32)} - ) + aot_cache._dynamic_signature((self._a,), {"svg_step_index": jnp.asarray(step, dtype=jnp.int32)}) for step in range(40) } self.assertEqual(len(sigs), 1) diff --git a/src/maxdiffusion/tests/wan/svg_attention_test.py b/src/maxdiffusion/tests/wan/svg_attention_test.py index 638d03ac0..1be65b03e 100644 --- a/src/maxdiffusion/tests/wan/svg_attention_test.py +++ b/src/maxdiffusion/tests/wan/svg_attention_test.py @@ -598,21 +598,15 @@ def test_partially_specified_schedules(self): # 3. Fully specified schedule: steps [11, 40), layers [1, 40) self.assertIs( - svg_attention.is_svg_active( - step_index=11, layer_index=0, start_step=11, end_step=40, start_layer=1, end_layer=40 - ), + svg_attention.is_svg_active(step_index=11, layer_index=0, start_step=11, end_step=40, start_layer=1, end_layer=40), False, ) self.assertIs( - svg_attention.is_svg_active( - step_index=10, layer_index=1, start_step=11, end_step=40, start_layer=1, end_layer=40 - ), + svg_attention.is_svg_active(step_index=10, layer_index=1, start_step=11, end_step=40, start_layer=1, end_layer=40), False, ) self.assertIs( - svg_attention.is_svg_active( - step_index=11, layer_index=1, start_step=11, end_step=40, start_layer=1, end_layer=40 - ), + svg_attention.is_svg_active(step_index=11, layer_index=1, start_step=11, end_step=40, start_layer=1, end_layer=40), True, )