fix lint issues

Author: csgoogle

src/maxdiffusion/models/wan/transformers/__init__.py

@@ -14,4 +14,4 @@
 limitations under the License.
 """
 
-from .transformer_wan_animate import NNXWanAnimateTransformer3DModel
+from .transformer_wan_animate import WanAnimateTransformer3DModel

src/maxdiffusion/models/wan/transformers/transformer_wan_animate.py

@@ -819,8 +819,8 @@ def __call__(
     return hidden_states
 
 
-class NNXWanAnimateTransformer3DModel(nnx.Module, FlaxModelMixin, ConfigMixin):
-  """NNX Wan Animate transformer with pose and face conditioning."""
+class WanAnimateTransformer3DModel(nnx.Module, FlaxModelMixin, ConfigMixin):
+  """Wan Animate transformer with pose and face conditioning."""
 
   @register_to_config
   def __init__(
@@ -1055,7 +1055,7 @@ def conditional_named_scope(self, name: str):
   def init_weights(self, rng: jax.Array, eval_only: bool = False) -> Dict[str, Any]:
     """NNX modules initialize parameters eagerly during construction."""
     del rng, eval_only
-    raise NotImplementedError("NNXWanAnimateTransformer3DModel initializes weights during construction.")
+    raise NotImplementedError("WanAnimateTransformer3DModel initializes weights during construction.")
 
   def _apply_face_adapter(self, hidden_states: jax.Array, motion_vec: Optional[jax.Array], block_idx) -> jax.Array:
     """Inject face-conditioning latents at the configured adapter blocks."""

src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py

@@ -46,7 +46,7 @@
 from maxdiffusion.max_utils import device_put_replicated, get_flash_block_sizes, get_precision
 from maxdiffusion.video_processor import VideoProcessor
 
-from ...models.wan.transformers.transformer_wan_animate import NNXWanAnimateTransformer3DModel
+from ...models.wan.transformers.transformer_wan_animate import WanAnimateTransformer3DModel
 from ...models.wan.wan_utils import load_wan_animate_transformer
 from ...pyconfig import HyperParameters
 from .wan_pipeline import WanPipeline, cast_with_exclusion
@@ -59,22 +59,22 @@ def create_sharded_animate_transformer(
     config: HyperParameters,
     restored_checkpoint=None,
     subfolder: str = "transformer",
-) -> NNXWanAnimateTransformer3DModel:
-  """Creates a sharded NNXWanAnimateTransformer3DModel on device.
+) -> WanAnimateTransformer3DModel:
+  """Creates a sharded WanAnimateTransformer3DModel on device.
 
   Follows the same pattern as create_sharded_logical_transformer in
-  wan_pipeline.py but uses NNXWanAnimateTransformer3DModel and the
+  wan_pipeline.py but uses WanAnimateTransformer3DModel and the
   animate-specific weight loader.
   """
 
   def _create_model(rngs: nnx.Rngs, wan_config: dict):
-    return NNXWanAnimateTransformer3DModel(**wan_config, rngs=rngs)
+    return WanAnimateTransformer3DModel(**wan_config, rngs=rngs)
 
   # 1. Load config.
   if restored_checkpoint:
     wan_config = restored_checkpoint["wan_config"]
   else:
-    wan_config = NNXWanAnimateTransformer3DModel.load_config(config.pretrained_model_name_or_path, subfolder=subfolder)
+    wan_config = WanAnimateTransformer3DModel.load_config(config.pretrained_model_name_or_path, subfolder=subfolder)
 
   wan_config["mesh"] = mesh
   wan_config["dtype"] = config.activations_dtype
@@ -215,15 +215,15 @@ class WanAnimatePipeline(WanPipeline):
 
   Args:
     config: HyperParameters configuration.
-    transformer: NNXWanAnimateTransformer3DModel instance (may be None for
+    transformer: WanAnimateTransformer3DModel instance (may be None for
       VAE-only mode).
     **kwargs: Passed to WanPipeline.__init__ (tokenizer, text_encoder, vae, etc.)
   """
 
   def __init__(
       self,
       config: HyperParameters,
-      transformer: Optional[NNXWanAnimateTransformer3DModel],
+      transformer: Optional[WanAnimateTransformer3DModel],
       **kwargs,
   ):
     super().__init__(config=config, **kwargs)
@@ -255,7 +255,7 @@ def load_animate_transformer(
       config: HyperParameters,
       restored_checkpoint=None,
       subfolder: str = "transformer",
-  ) -> NNXWanAnimateTransformer3DModel:
+  ) -> WanAnimateTransformer3DModel:
     with mesh:
       return create_sharded_animate_transformer(
           devices_array=devices_array,
@@ -273,7 +273,7 @@ def _load_and_init(
       restored_checkpoint=None,
       vae_only: bool = False,
       load_transformer: bool = True,
-  ) -> Tuple["WanAnimatePipeline", Optional[NNXWanAnimateTransformer3DModel]]:
+  ) -> Tuple["WanAnimatePipeline", Optional[WanAnimateTransformer3DModel]]:
     common_components = cls._create_common_components(config, vae_only)
     transformer = None
     if not vae_only and load_transformer:

src/maxdiffusion/tests/wan_animate_diffusers_parity_test.py

@@ -608,54 +608,54 @@ def test_mask_video_preprocessing_matches_diffusers(self):
 
   def test_check_inputs_matches_diffusers_validation(self):
     invalid_calls = [
-        dict(
-            prompt="prompt",
-            negative_prompt=None,
-            image=PIL.Image.new("RGB", (16, 16)),
-            pose_video=[PIL.Image.new("RGB", (16, 16))],
-            face_video=[PIL.Image.new("RGB", (16, 16))],
-            background_video=None,
-            mask_video=None,
-            height=16,
-            width=16,
-            prompt_embeds=jnp.zeros((1, 1, 1)),
-            negative_prompt_embeds=None,
-            image_embeds=None,
-            mode="animate",
-            prev_segment_conditioning_frames=1,
-        ),
-        dict(
-            prompt="prompt",
-            negative_prompt=None,
-            image=PIL.Image.new("RGB", (16, 16)),
-            pose_video=[PIL.Image.new("RGB", (16, 16))],
-            face_video=[PIL.Image.new("RGB", (16, 16))],
-            background_video=None,
-            mask_video=None,
-            height=18,
-            width=16,
-            prompt_embeds=None,
-            negative_prompt_embeds=None,
-            image_embeds=None,
-            mode="animate",
-            prev_segment_conditioning_frames=1,
-        ),
-        dict(
-            prompt="prompt",
-            negative_prompt=None,
-            image=PIL.Image.new("RGB", (16, 16)),
-            pose_video=[PIL.Image.new("RGB", (16, 16))],
-            face_video=[PIL.Image.new("RGB", (16, 16))],
-            background_video=None,
-            mask_video=None,
-            height=16,
-            width=16,
-            prompt_embeds=None,
-            negative_prompt_embeds=None,
-            image_embeds=None,
-            mode="replace",
-            prev_segment_conditioning_frames=3,
-        ),
+        {
+            "prompt": "prompt",
+            "negative_prompt": None,
+            "image": PIL.Image.new("RGB", (16, 16)),
+            "pose_video": [PIL.Image.new("RGB", (16, 16))],
+            "face_video": [PIL.Image.new("RGB", (16, 16))],
+            "background_video": None,
+            "mask_video": None,
+            "height": 16,
+            "width": 16,
+            "prompt_embeds": jnp.zeros((1, 1, 1)),
+            "negative_prompt_embeds": None,
+            "image_embeds": None,
+            "mode": "animate",
+            "prev_segment_conditioning_frames": 1,
+        },
+        {
+            "prompt": "prompt",
+            "negative_prompt": None,
+            "image": PIL.Image.new("RGB", (16, 16)),
+            "pose_video": [PIL.Image.new("RGB", (16, 16))],
+            "face_video": [PIL.Image.new("RGB", (16, 16))],
+            "background_video": None,
+            "mask_video": None,
+            "height": 18,
+            "width": 16,
+            "prompt_embeds": None,
+            "negative_prompt_embeds": None,
+            "image_embeds": None,
+            "mode": "animate",
+            "prev_segment_conditioning_frames": 1,
+        },
+        {
+            "prompt": "prompt",
+            "negative_prompt": None,
+            "image": PIL.Image.new("RGB", (16, 16)),
+            "pose_video": [PIL.Image.new("RGB", (16, 16))],
+            "face_video": [PIL.Image.new("RGB", (16, 16))],
+            "background_video": None,
+            "mask_video": None,
+            "height": 16,
+            "width": 16,
+            "prompt_embeds": None,
+            "negative_prompt_embeds": None,
+            "image_embeds": None,
+            "mode": "replace",
+            "prev_segment_conditioning_frames": 3,
+        },
     ]
 
     for kwargs in invalid_calls:
@@ -780,7 +780,7 @@ def _scalar(x):
     hf_negative = torch.tensor(to_numpy(max_negative))
     hf_image = torch.tensor(to_numpy(max_image))
 
-    scheduler_config = dict(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=5.0)
+    scheduler_config = {"prediction_type": "flow_prediction", "use_flow_sigmas": True, "flow_shift": 5.0}
     max_scheduler = FlaxUniPCMultistepScheduler(**scheduler_config)
     max_state = max_scheduler.create_state()
     max_state = max_scheduler.set_timesteps(max_state, num_inference_steps=timestep_count, shape=max_latents.shape)
@@ -852,7 +852,7 @@ def _scalar(x):
     np.testing.assert_allclose(to_numpy(max_next), hf_channel_first_to_last(hf_next), atol=1e-5, rtol=1e-5)
 
   def test_flax_unipc_flow_sigmas_match_diffusers(self):
-    scheduler_config = dict(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=5.0)
+    scheduler_config = {"prediction_type": "flow_prediction", "use_flow_sigmas": True, "flow_shift": 5.0}
 
     max_scheduler = FlaxUniPCMultistepScheduler(**scheduler_config)
     max_state = max_scheduler.create_state()

src/maxdiffusion/tests/wan_animate_module_parity_test.py

@@ -51,7 +51,7 @@
     MotionConv2d,
     MotionEncoderResBlock,
     MotionLinear,
-    NNXWanAnimateTransformer3DModel,
+    WanAnimateTransformer3DModel,
     WanAnimateFaceBlockCrossAttention,
     WanAnimateFaceEncoder,
     WanAnimateMotionEncoder,
@@ -320,7 +320,7 @@ def test_wan_animate_transformer_weight_mapping_covers_all_local_params(self):
     hf_model = HFWanAnimateTransformer3DModel(**cfg).eval()
 
     with self.mesh, nn_partitioning.axis_rules(self.logical_axis_rules):
-      max_model = NNXWanAnimateTransformer3DModel(rngs=self.rngs, scan_layers=False, mesh=self.mesh, **cfg)
+      max_model = WanAnimateTransformer3DModel(rngs=self.rngs, scan_layers=False, mesh=self.mesh, **cfg)
       missing_keys, flax_state_dict = map_hf_wan_animate_state_to_local(
           max_model, hf_model, num_layers=cfg["num_layers"], scan_layers=False
       )
@@ -359,7 +359,7 @@ def test_wan_animate_transformer_weight_mapping_covers_all_local_params_scanned(
     hf_model = HFWanAnimateTransformer3DModel(**cfg).eval()
 
     with self.mesh, nn_partitioning.axis_rules(self.logical_axis_rules):
-      max_model = NNXWanAnimateTransformer3DModel(rngs=self.rngs, scan_layers=True, mesh=self.mesh, **cfg)
+      max_model = WanAnimateTransformer3DModel(rngs=self.rngs, scan_layers=True, mesh=self.mesh, **cfg)
       missing_keys, flax_state_dict = map_hf_wan_animate_state_to_local(
           max_model, hf_model, num_layers=cfg["num_layers"], scan_layers=True
       )
@@ -432,7 +432,7 @@ def test_wan_animate_transformer_forward_parity(self):
     hf_model = HFWanAnimateTransformer3DModel(**cfg).eval()
 
     with self.mesh, nn_partitioning.axis_rules(self.logical_axis_rules):
-      max_model = NNXWanAnimateTransformer3DModel(rngs=self.rngs, scan_layers=False, mesh=self.mesh, **cfg)
+      max_model = WanAnimateTransformer3DModel(rngs=self.rngs, scan_layers=False, mesh=self.mesh, **cfg)
       missing_keys, _ = map_hf_wan_animate_state_to_local(
           max_model, hf_model, num_layers=cfg["num_layers"], scan_layers=False
       )
@@ -497,7 +497,7 @@ def test_wan_animate_transformer_forward_parity_scanned(self):
     hf_model = HFWanAnimateTransformer3DModel(**cfg).eval()
 
     with self.mesh, nn_partitioning.axis_rules(self.logical_axis_rules):
-      max_model = NNXWanAnimateTransformer3DModel(rngs=self.rngs, scan_layers=True, mesh=self.mesh, **cfg)
+      max_model = WanAnimateTransformer3DModel(rngs=self.rngs, scan_layers=True, mesh=self.mesh, **cfg)
       missing_keys, _ = map_hf_wan_animate_state_to_local(
           max_model, hf_model, num_layers=cfg["num_layers"], scan_layers=True
       )