Formatting fixes

Author: coolkp

src/maxdiffusion/pyconfig.py

@@ -137,13 +137,12 @@ def wan_init(raw_keys):
       else:
         raise ValueError(f"{transformer_pretrained_model_name_or_path} transformer model is not supported for Wan 2.1")
     if "use_qwix_quantization" not in raw_keys:
-      raise ValueError(f"use_qwix_quantization is not set.")
+      raise ValueError("use_qwix_quantization is not set.")
     elif raw_keys["use_qwix_quantization"]:
       if "quantization" not in raw_keys:
-        raise ValueError(f"Quantization type is not set when use_qwix_quantization is enabled.")
+        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']}")
-      
 
   @staticmethod
   def calculate_global_batch_sizes(per_device_batch_size):

src/maxdiffusion/tests/wan_transformer_test.py

@@ -19,7 +19,7 @@
 import jax.numpy as jnp
 import pytest
 import unittest
-from unittest.mock import Mock, patch, MagicMock
+from unittest.mock import Mock, patch
 from absl.testing import absltest
 from flax import nnx
 from jax.sharding import Mesh
@@ -276,7 +276,7 @@ def test_wan_model(self):
           hidden_states=dummy_hidden_states, timestep=dummy_timestep, encoder_hidden_states=dummy_encoder_hidden_states
       )
     assert dummy_output.shape == hidden_states_shape
-  
+
   def test_get_qt_provider(self):
     """
     Tests the provider logic for all config branches.
@@ -301,7 +301,7 @@ def test_get_qt_provider(self):
     provider_fp8 = WanPipeline.get_qt_provider(config_fp8)
     self.assertIsNotNone(provider_fp8)
     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)
     config_fp8_full.use_qwix_quantization = True
@@ -329,11 +329,11 @@ def test_quantize_transformer_enabled(self, mock_get_dummy_inputs, mock_quantize
     mock_config.use_qwix_quantization = True
     mock_config.quantization = "fp8_full"
     mock_config.per_device_batch_size = 1
-    
+
     mock_model = Mock(spec=WanModel)
     mock_pipeline = Mock()
     mock_mesh = Mock()
-    
+
     # Mock the return values of dependencies
     mock_get_dummy_inputs.return_value = (Mock(), Mock(), Mock())
     mock_quantized_model_obj = Mock(spec=WanModel)
@@ -356,9 +356,9 @@ def test_quantize_transformer_disabled(self, mock_quantize_model):
     # Setup Mocks
     mock_config = Mock(spec=HyperParameters)
     mock_config.use_qwix_quantization = False # Main condition for this test
-    
+
     mock_model = Mock(spec=WanModel)
-    
+
     # Call the method under test
     result = WanPipeline.quantize_transformer(mock_config, mock_model, Mock(), Mock())