pyink checks

Author: prishajain1

src/maxdiffusion/__init__.py

@@ -202,27 +202,29 @@ def setup(self):
         dtype=self.dtype,
         param_dtype=self.weights_dtype,
     )
-    self.img_mlp = nn.Sequential([
-        nn.Dense(
-            int(self.dim * self.mlp_ratio),
-            use_bias=True,
-            kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
-            bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
-            dtype=self.dtype,
-            param_dtype=self.weights_dtype,
-            precision=self.precision,
-        ),
-        nn.gelu,
-        nn.Dense(
-            self.dim,
-            use_bias=True,
-            kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")),
-            bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
-            dtype=self.dtype,
-            param_dtype=self.weights_dtype,
-            precision=self.precision,
-        ),
-    ])
+    self.img_mlp = nn.Sequential(
+        [
+            nn.Dense(
+                int(self.dim * self.mlp_ratio),
+                use_bias=True,
+                kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
+                bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
+                dtype=self.dtype,
+                param_dtype=self.weights_dtype,
+                precision=self.precision,
+            ),
+            nn.gelu,
+            nn.Dense(
+                self.dim,
+                use_bias=True,
+                kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")),
+                bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
+                dtype=self.dtype,
+                param_dtype=self.weights_dtype,
+                precision=self.precision,
+            ),
+        ]
+    )
 
     self.txt_norm2 = nn.LayerNorm(
         use_bias=False,
@@ -231,27 +233,29 @@ def setup(self):
         dtype=self.dtype,
         param_dtype=self.weights_dtype,
     )
-    self.txt_mlp = nn.Sequential([
-        nn.Dense(
-            int(self.dim * self.mlp_ratio),
-            use_bias=True,
-            kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
-            bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
-            dtype=self.dtype,
-            param_dtype=self.weights_dtype,
-            precision=self.precision,
-        ),
-        nn.gelu,
-        nn.Dense(
-            self.dim,
-            use_bias=True,
-            kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")),
-            bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
-            dtype=self.dtype,
-            param_dtype=self.weights_dtype,
-            precision=self.precision,
-        ),
-    ])
+    self.txt_mlp = nn.Sequential(
+        [
+            nn.Dense(
+                int(self.dim * self.mlp_ratio),
+                use_bias=True,
+                kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
+                bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
+                dtype=self.dtype,
+                param_dtype=self.weights_dtype,
+                precision=self.precision,
+            ),
+            nn.gelu,
+            nn.Dense(
+                self.dim,
+                use_bias=True,
+                kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")),
+                bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)),
+                dtype=self.dtype,
+                param_dtype=self.weights_dtype,
+                precision=self.precision,
+            ),
+        ]
+    )
 
     # let chunk size default to None
     self._chunk_size = None

src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import inspect
 from importlib import import_module
 from typing import Any, Dict, Optional, Tuple

src/maxdiffusion/models/ltx_video/transformers_pytorch/attention.py

@@ -460,11 +460,13 @@ def __call__(
 
     control_hidden_states = self.vace_patch_embedding(control_hidden_states)
     control_hidden_states = jax.lax.collapse(control_hidden_states, 1, -1)
-    control_hidden_states_padding = jnp.zeros((
-        batch_size,
-        control_hidden_states.shape[1],
-        hidden_states.shape[2] - control_hidden_states.shape[2],
-    ))
+    control_hidden_states_padding = jnp.zeros(
+        (
+            batch_size,
+            control_hidden_states.shape[1],
+            hidden_states.shape[2] - control_hidden_states.shape[2],
+        )
+    )
 
     control_hidden_states = jnp.concatenate([control_hidden_states, control_hidden_states_padding], axis=2)
 

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

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import os
 import time
 

src/maxdiffusion/pedagogical_examples/attention_comparison.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import os
 import argparse
 import tensorflow as tf

src/maxdiffusion/pedagogical_examples/dataset_tf_cache_to_tfrecord.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 from typing import Sequence
 from absl import app
 import jax

src/maxdiffusion/pedagogical_examples/parameter_count.py

@@ -54,12 +54,14 @@
 dl_manager = tfds.download.DownloadManager(download_dir="/tmp")
 tmp_dataset = "dataset"
 
-TRANSFORMS = transforms.Compose([
-    transforms.ToTensor(),
-    transforms.Resize(size=512, interpolation=transforms.InterpolationMode.BICUBIC),
-    transforms.CenterCrop(size=512),
-    transforms.Normalize([0.5], [0.5]),
-])
+TRANSFORMS = transforms.Compose(
+    [
+        transforms.ToTensor(),
+        transforms.Resize(size=512, interpolation=transforms.InterpolationMode.BICUBIC),
+        transforms.CenterCrop(size=512),
+        transforms.Normalize([0.5], [0.5]),
+    ]
+)
 
 
 def delete_files(path):

src/maxdiffusion/pedagogical_examples/to_tfrecords.py

@@ -51,14 +51,16 @@
 
   _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_and_onnx_objects))
 else:
-  _import_structure["stable_diffusion"].extend([
-      "OnnxStableDiffusionImg2ImgPipeline",
-      "OnnxStableDiffusionInpaintPipeline",
-      "OnnxStableDiffusionInpaintPipelineLegacy",
-      "OnnxStableDiffusionPipeline",
-      "OnnxStableDiffusionUpscalePipeline",
-      "StableDiffusionOnnxPipeline",
-  ])
+  _import_structure["stable_diffusion"].extend(
+      [
+          "OnnxStableDiffusionImg2ImgPipeline",
+          "OnnxStableDiffusionInpaintPipeline",
+          "OnnxStableDiffusionInpaintPipelineLegacy",
+          "OnnxStableDiffusionPipeline",
+          "OnnxStableDiffusionUpscalePipeline",
+          "StableDiffusionOnnxPipeline",
+      ]
+  )
 
 try:
   if not is_flax_available():
@@ -80,14 +82,18 @@
   _import_structure["controlnet"].extend(
       ["FlaxStableDiffusionControlNetPipeline", "FlaxStableDiffusionXLControlNetPipeline"]
   )
-  _import_structure["stable_diffusion"].extend([
-      "FlaxStableDiffusionImg2ImgPipeline",
-      "FlaxStableDiffusionInpaintPipeline",
-      "FlaxStableDiffusionPipeline",
-  ])
-  _import_structure["stable_diffusion_xl"].extend([
-      "FlaxStableDiffusionXLPipeline",
-  ])
+  _import_structure["stable_diffusion"].extend(
+      [
+          "FlaxStableDiffusionImg2ImgPipeline",
+          "FlaxStableDiffusionInpaintPipeline",
+          "FlaxStableDiffusionPipeline",
+      ]
+  )
+  _import_structure["stable_diffusion_xl"].extend(
+      [
+          "FlaxStableDiffusionXLPipeline",
+      ]
+  )
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
   try:
     if not is_onnx_available():

src/maxdiffusion/pipelines/__init__.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 from typing import TYPE_CHECKING
 
 from ...utils import (