From d4638e76d977898eea5a68931303bb71d67ace62 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Tue, 13 Oct 2026 08:03:45 +0000 Subject: [PATCH] style: fix pyink formatting on main (2 files) --- src/maxdiffusion/configuration_utils.py | 4 +--- src/maxdiffusion/trainers/wan_trainer_2_2.py | 1 + 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/maxdiffusion/configuration_utils.py b/src/maxdiffusion/configuration_utils.py index dacd21f2b..7de137632 100644 --- a/src/maxdiffusion/configuration_utils.py +++ b/src/maxdiffusion/configuration_utils.py @@ -496,9 +496,7 @@ def extract_init_dict(cls, config_dict, **kwargs): diffusers_library = importlib.import_module(__name__.split(".")[0]) if cls.has_compatibles: - compatible_classes = [ - c for c in cls._get_compatibles() if not isinstance(c, DummyObject) # type: ignore - ] + compatible_classes = [c for c in cls._get_compatibles() if not isinstance(c, DummyObject)] # type: ignore else: compatible_classes = [] diff --git a/src/maxdiffusion/trainers/wan_trainer_2_2.py b/src/maxdiffusion/trainers/wan_trainer_2_2.py index 6ea105114..b2c2533c5 100644 --- a/src/maxdiffusion/trainers/wan_trainer_2_2.py +++ b/src/maxdiffusion/trainers/wan_trainer_2_2.py @@ -554,6 +554,7 @@ def train_step_2_2(state_low, state_high, data, rng, scheduler_state, is_high_no is_high_noise = jax.random.uniform(cond_rng) >= u_boundary else: _, new_rng, timestep_rng, dropout_rng, noise_rng = jax.random.split(rng, num=5) + def compute_loss_high(high_params, s_high): u = jax.random.uniform(timestep_rng, (bsz,), minval=u_boundary, maxval=1.0) t_shifted = (u * shift) / (1.0 + (shift - 1.0) * u)