lint.

Author: entrpn

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -233,11 +233,11 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
   def get_basic_config(cls, dtype):
     rules = [
         qwix.QtRule(
-          module_path='.*',  # Apply to all modules
-          weight_qtype=dtype,
-          act_qtype=dtype,
+            module_path=".*",  # Apply to all modules
+            weight_qtype=dtype,
+            act_qtype=dtype,
         )
-      ]
+    ]
     return rules
 
   @classmethod
@@ -249,17 +249,17 @@ def get_fp8_config(cls, quantization_calibration_method: str):
     """
     rules = [
         qwix.QtRule(
-          module_path='.*',  # Apply to all modules
-          weight_qtype=jnp.float8_e4m3fn,
-          act_qtype=jnp.float8_e4m3fn,
-          bwd_qtype=jnp.float8_e5m2,
-          bwd_use_original_residuals=True,
-          disable_channelwise_axes=True, # per_tensor calibration
-          weight_calibration_method = quantization_calibration_method,
-          act_calibration_method = quantization_calibration_method,
-          bwd_calibration_method = quantization_calibration_method,
+            module_path=".*",  # Apply to all modules
+            weight_qtype=jnp.float8_e4m3fn,
+            act_qtype=jnp.float8_e4m3fn,
+            bwd_qtype=jnp.float8_e5m2,
+            bwd_use_original_residuals=True,
+            disable_channelwise_axes=True,  # per_tensor calibration
+            weight_calibration_method=quantization_calibration_method,
+            act_calibration_method=quantization_calibration_method,
+            bwd_calibration_method=quantization_calibration_method,
         )
-      ]
+    ]
     return rules
 
   @classmethod
@@ -288,7 +288,7 @@ def quantize_transformer(cls, config: HyperParameters, model: WanModel, pipeline
 
     batch_size = int(config.per_device_batch_size * jax.local_device_count())
     latents, prompt_embeds, timesteps = get_dummy_wan_inputs(config, pipeline, batch_size)
-    model_inputs= (latents, timesteps, prompt_embeds)
+    model_inputs = (latents, timesteps, prompt_embeds)
     with mesh:
       quantized_model = qwix.quantize_model(model, q_rules, *model_inputs)
     max_logging.log("Qwix Quantization complete.")

src/maxdiffusion/pyconfig.py

@@ -142,7 +142,9 @@ def wan_init(raw_keys):
       if "quantization" not in raw_keys:
         raise ValueError("Quantization type is not set when use_qwix_quantization is enabled.")
       elif raw_keys["quantization"] not in ["int8", "fp8", "fp8_full"]:
-        raise ValueError(f"Quantization type is not supported when use_qwix_quantization is enabled: {raw_keys['quantization']}")
+        raise ValueError(
+            f"Quantization type is not supported when use_qwix_quantization is enabled: {raw_keys['quantization']}"
+        )
 
   @staticmethod
   def calculate_global_batch_sizes(per_device_batch_size):

src/maxdiffusion/tests/wan_transformer_test.py

@@ -13,6 +13,7 @@
  See the License for the specific language governing permissions and
  limitations under the License.
  """
+
 from qwix import QtProvider
 import os
 import jax
@@ -290,7 +291,7 @@ def test_get_qt_provider(self):
     config_int8 = Mock(spec=HyperParameters)
     config_int8.use_qwix_quantization = True
     config_int8.quantization = "int8"
-    provider_int8:QtProvider = WanPipeline.get_qt_provider(config_int8)
+    provider_int8: QtProvider = WanPipeline.get_qt_provider(config_int8)
     self.assertIsNotNone(provider_int8)
     self.assertEqual(provider_int8._rules[0].weight_qtype, jnp.int8)
 
@@ -300,7 +301,7 @@ def test_get_qt_provider(self):
     config_fp8.quantization = "fp8"
     provider_fp8 = WanPipeline.get_qt_provider(config_fp8)
     self.assertIsNotNone(provider_fp8)
-    self.assertEqual(provider_fp8.rules[0].kwargs['weight_qtype'], jnp.float8_e4m3fn)
+    self.assertEqual(provider_fp8.rules[0].kwargs["weight_qtype"], jnp.float8_e4m3fn)
 
     # Case 4: Quantization enabled, type 'fp8_full'
     config_fp8_full = Mock(spec=HyperParameters)
@@ -309,7 +310,7 @@ def test_get_qt_provider(self):
     config_fp8_full.quantization_calibration_method = "absmax"
     provider_fp8_full = WanPipeline.get_qt_provider(config_fp8_full)
     self.assertIsNotNone(provider_fp8_full)
-    self.assertEqual(provider_fp8_full.rules[0].kwargs['bwd_qtype'], jnp.float8_e5m2)
+    self.assertEqual(provider_fp8_full.rules[0].kwargs["bwd_qtype"], jnp.float8_e5m2)
 
     # Case 5: Invalid quantization type
     config_invalid = Mock(spec=HyperParameters)
@@ -318,8 +319,8 @@ def test_get_qt_provider(self):
     self.assertIsNone(WanPipeline.get_qt_provider(config_invalid))
 
   # To test quantize_transformer, we patch its external dependencies
-  @patch('maxdiffusion.pipelines.wan.wan_pipeline.qwix.quantize_model')
-  @patch('maxdiffusion.pipelines.wan.wan_pipeline.get_dummy_wan_inputs')
+  @patch("maxdiffusion.pipelines.wan.wan_pipeline.qwix.quantize_model")
+  @patch("maxdiffusion.pipelines.wan.wan_pipeline.get_dummy_wan_inputs")
   def test_quantize_transformer_enabled(self, mock_get_dummy_inputs, mock_quantize_model):
     """
     Tests that quantize_transformer calls qwix when quantization is enabled.
@@ -348,14 +349,14 @@ def test_quantize_transformer_enabled(self, mock_get_dummy_inputs, mock_quantize
     # Check that the model returned is the new quantized model
     self.assertIs(result, mock_quantized_model_obj)
 
-  @patch('maxdiffusion.pipelines.wan.wan_pipeline.qwix.quantize_model')
+  @patch("maxdiffusion.pipelines.wan.wan_pipeline.qwix.quantize_model")
   def test_quantize_transformer_disabled(self, mock_quantize_model):
     """
     Tests that quantize_transformer is skipped when quantization is disabled.
     """
     # Setup Mocks
     mock_config = Mock(spec=HyperParameters)
-    mock_config.use_qwix_quantization = False # Main condition for this test
+    mock_config.use_qwix_quantization = False  # Main condition for this test
 
     mock_model = Mock(spec=WanModel)
 

src/maxdiffusion/trainers/wan_trainer.py

@@ -55,7 +55,9 @@ def generate_sample(config, pipeline, filename_prefix):
   Generates a video to validate training did not corrupt the model
   """
   if not hasattr(pipeline, "vae"):
-    wan_vae, vae_cache = WanPipeline.load_vae(pipeline.mesh.devices, pipeline.mesh, nnx.Rngs(jax.random.key(config.seed)), config)
+    wan_vae, vae_cache = WanPipeline.load_vae(
+        pipeline.mesh.devices, pipeline.mesh, nnx.Rngs(jax.random.key(config.seed)), config
+    )
     pipeline.vae = wan_vae
     pipeline.vae_cache = vae_cache
   return generate_wan(config, pipeline, filename_prefix)
@@ -147,7 +149,7 @@ def start_training(self):
     pipeline = self.load_checkpoint()
     # Generate a sample before training to compare against generated sample after training.
     pretrained_video_path = generate_sample(self.config, pipeline, filename_prefix="pre-training-")
-    
+
     # save some memory.
     del pipeline.vae
     del pipeline.vae_cache