update

Author: sayakpaul

tests/lora/test_lora_layers_cogview4.py

@@ -13,10 +13,8 @@
 # limitations under the License.
 
 import sys
-import tempfile
 import unittest
 
-import numpy as np
 import torch
 from parameterized import parameterized
 from transformers import AutoTokenizer, GlmModel
@@ -27,7 +25,6 @@
     require_peft_backend,
     require_torch_accelerator,
     skip_mps,
-    torch_device,
 )
 
 
@@ -119,35 +116,6 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self):
     def test_simple_inference_with_text_denoiser_lora_unfused(self):
         super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3)
 
-    def test_simple_inference_save_pretrained(self):
-        """
-        Tests a simple usecase where users could use saving utilities for LoRA through save_pretrained
-        """
-        for scheduler_cls in self.scheduler_classes:
-            components, _, _ = self.get_dummy_components(scheduler_cls)
-            pipe = self.pipeline_class(**components)
-            pipe = pipe.to(torch_device)
-            pipe.set_progress_bar_config(disable=None)
-            _, _, inputs = self.get_dummy_inputs(with_generator=False)
-
-            output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
-            self.assertTrue(output_no_lora.shape == self.output_shape)
-
-            images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
-
-            with tempfile.TemporaryDirectory() as tmpdirname:
-                pipe.save_pretrained(tmpdirname)
-
-                pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname)
-                pipe_from_pretrained.to(torch_device)
-
-            images_lora_save_pretrained = pipe_from_pretrained(**inputs, generator=torch.manual_seed(0))[0]
-
-            self.assertTrue(
-                np.allclose(images_lora, images_lora_save_pretrained, atol=1e-3, rtol=1e-3),
-                "Loading from saved checkpoints should give same results.",
-            )
-
     @parameterized.expand([("block_level", True), ("leaf_level", False)])
     @require_torch_accelerator
     def test_group_offloading_inference_denoiser(self, offload_type, use_stream):

tests/lora/test_lora_layers_hunyuanvideo.py

@@ -156,11 +156,6 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self):
     def test_simple_inference_with_text_denoiser_lora_unfused(self):
         super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3)
 
-    # TODO(aryan): Fix the following test
-    @unittest.skip("This test fails with an error I haven't been able to debug yet.")
-    def test_simple_inference_save_pretrained(self):
-        pass
-
     @unittest.skip("Not supported in HunyuanVideo.")
     def test_simple_inference_with_text_denoiser_block_scale(self):
         pass

tests/lora/utils.py

@@ -329,20 +329,6 @@ def add_adapters_to_pipeline(self, pipe, text_lora_config=None, denoiser_lora_co
                 )
         return pipe, denoiser
 
-    def test_simple_inference(self):
-        """
-        Tests a simple inference and makes sure it works as expected
-        """
-        for scheduler_cls in self.scheduler_classes:
-            components, text_lora_config, _ = self.get_dummy_components(scheduler_cls)
-            pipe = self.pipeline_class(**components)
-            pipe = pipe.to(torch_device)
-            pipe.set_progress_bar_config(disable=None)
-
-            _, _, inputs = self.get_dummy_inputs()
-            output_no_lora = pipe(**inputs)[0]
-            self.assertTrue(output_no_lora.shape == self.output_shape)
-
     @require_peft_version_greater("0.13.1")
     def test_low_cpu_mem_usage_with_injection(self):
         """Tests if we can inject LoRA state dict with low_cpu_mem_usage."""