Skip to content

[PR 1/5] A few pyink lint fixes from previous PR #494

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
jitendra-jalwaniya wants to merge 1 commit into from
Closed

[PR 1/5] A few pyink lint fixes from previous PR #494

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 3 additions & 6 deletions src/maxdiffusion/aot_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Comment thread
jitendra-jalwaniya marked this conversation as resolved.
if isinstance(attn_cfg, dict):
meta[f"{attr}_svg_config"] = json.dumps(
Expand Down Expand Up @@ -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):
Expand Down
10 changes: 2 additions & 8 deletions src/maxdiffusion/pipelines/wan/wan_pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment thread
jitendra-jalwaniya marked this conversation as resolved.
if isinstance(attn_cfg, dict) and bool(attn_cfg.get("use_svg_attention", False)):
return True
if hasattr(obj, "attributes") and hasattr(obj, "nodes"):
Expand All @@ -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)
Comment thread
jitendra-jalwaniya marked this conversation as resolved.
if isinstance(ac, dict) and bool(ac.get("use_svg_attention", False)):
return True
return False
Expand Down
4 changes: 1 addition & 3 deletions src/maxdiffusion/tests/aot_cache_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
12 changes: 3 additions & 9 deletions src/maxdiffusion/tests/wan/svg_attention_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand Down
Loading