up

Author: yiyixuxu

tests/pipelines/wan/test_wan_22_image_to_video.py

@@ -195,10 +195,11 @@ def test_save_load_optional_components(self, expected_max_difference=1e-4):
             pipe_loaded.to(torch_device)
             pipe_loaded.set_progress_bar_config(disable=None)
 
-        self.assertTrue(
-            getattr(pipe_loaded, "transformer") is None,
-            "`transformer` did not stay set to None after loading.",
-        )
+        for component in optional_component:
+            self.assertTrue(
+                getattr(pipe_loaded, component) is None,
+                f"`{component}` did not stay set to None after loading.",
+            )
 
         inputs = self.get_dummy_inputs(generator_device)
         torch.manual_seed(0)