From 4d39b7483d405474e913754c3a22903cca0d7fbf Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Fri, 29 Mar 2024 14:23:28 +0530 Subject: [PATCH] Memory clean up on all Slow Tests (#7514) * update * update --------- Co-authored-by: Sayak Paul --- tests/lora/test_lora_layers_sd.py | 10 ++++++++++ tests/lora/test_lora_layers_sdxl.py | 10 ++++++++++ tests/models/autoencoders/test_models_vae.py | 6 ++++++ .../pipelines/animatediff/test_animatediff.py | 6 ++++++ tests/pipelines/audioldm/test_audioldm.py | 10 ++++++++++ tests/pipelines/audioldm2/test_audioldm2.py | 5 +++++ .../test_consistency_models.py | 5 +++++ tests/pipelines/controlnet/test_controlnet.py | 10 ++++++++++ .../controlnet/test_controlnet_img2img.py | 5 +++++ .../controlnet/test_controlnet_inpaint.py | 5 +++++ .../controlnet/test_controlnet_sdxl.py | 5 +++++ .../dance_diffusion/test_dance_diffusion.py | 6 ++++++ tests/pipelines/dit/test_dit.py | 5 +++++ tests/pipelines/i2vgen_xl/test_i2vgenxl.py | 6 ++++++ .../kandinsky/test_kandinsky_img2img.py | 6 ++++++ .../latent_diffusion/test_latent_diffusion.py | 10 ++++++++++ .../test_ledits_pp_stable_diffusion.py | 5 +++++ tests/pipelines/musicldm/test_musicldm.py | 5 +++++ .../paint_by_example/test_paint_by_example.py | 6 ++++++ tests/pipelines/pixart_alpha/test_pixart.py | 5 +++++ .../test_semantic_diffusion.py | 12 ++++++++++++ tests/pipelines/shap_e/test_shap_e.py | 6 ++++++ tests/pipelines/shap_e/test_shap_e_img2img.py | 6 ++++++ .../test_stable_cascade_decoder.py | 6 ++++++ .../test_stable_cascade_prior.py | 6 ++++++ .../stable_diffusion/test_stable_diffusion.py | 10 ++++++++++ .../test_stable_diffusion_img2img.py | 10 ++++++++++ .../test_stable_diffusion_inpaint.py | 5 +++++ ...est_stable_diffusion_instruction_pix2pix.py | 5 +++++ .../test_stable_diffusion_attend_and_excite.py | 5 +++++ .../test_stable_diffusion_depth.py | 10 ++++++++++ .../test_stable_diffusion_diffedit.py | 10 ++++++++++ .../test_stable_diffusion_inpaint.py | 6 ++++++ .../test_stable_diffusion_latent_upscale.py | 5 +++++ .../test_stable_diffusion_upscale.py | 12 ++++++++++++ .../test_stable_diffusion_v_pred.py | 12 ++++++++++++ .../test_stable_diffusion_adapter.py | 5 +++++ .../test_stable_diffusion_image_variation.py | 10 ++++++++++ .../test_stable_diffusion_k_diffusion.py | 6 ++++++ .../test_stable_diffusion_ldm3d.py | 10 ++++++++++ .../test_stable_diffusion_panorama.py | 5 +++++ .../test_safe_diffusion.py | 12 ++++++++++++ .../test_stable_diffusion_sag.py | 6 ++++++ .../test_stable_diffusion_xl.py | 5 +++++ .../test_stable_diffusion_xl_adapter.py | 5 +++++ .../test_stable_diffusion_xl_k_diffusion.py | 6 ++++++ .../stable_unclip/test_stable_unclip.py | 6 ++++++ .../test_stable_unclip_img2img.py | 6 ++++++ .../test_stable_video_diffusion.py | 6 ++++++ tests/pipelines/test_pipelines.py | 18 ++++++++++++++++++ tests/pipelines/unclip/test_unclip.py | 12 ++++++++++++ .../unclip/test_unclip_image_variation.py | 6 ++++++ .../pipelines/unidiffuser/test_unidiffuser.py | 10 ++++++++++ 53 files changed, 391 insertions(+) diff --git a/tests/lora/test_lora_layers_sd.py b/tests/lora/test_lora_layers_sd.py index 1ead6c2f7001..ebf46e396284 100644 --- a/tests/lora/test_lora_layers_sd.py +++ b/tests/lora/test_lora_layers_sd.py @@ -81,6 +81,11 @@ class StableDiffusionLoRATests(PeftLoraLoaderMixinTests, unittest.TestCase): "latent_channels": 4, } + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() @@ -150,6 +155,11 @@ def test_integration_move_lora_cpu(self): @require_torch_gpu @require_peft_backend class LoraIntegrationTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/lora/test_lora_layers_sdxl.py b/tests/lora/test_lora_layers_sdxl.py index d603994b66dc..b46b887d10fb 100644 --- a/tests/lora/test_lora_layers_sdxl.py +++ b/tests/lora/test_lora_layers_sdxl.py @@ -90,6 +90,11 @@ class StableDiffusionXLLoRATests(PeftLoraLoaderMixinTests, unittest.TestCase): "sample_size": 128, } + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() @@ -100,6 +105,11 @@ def tearDown(self): @require_torch_gpu @require_peft_backend class LoraSDXLIntegrationTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/models/autoencoders/test_models_vae.py b/tests/models/autoencoders/test_models_vae.py index e70874f23b63..b0c24b8d4315 100644 --- a/tests/models/autoencoders/test_models_vae.py +++ b/tests/models/autoencoders/test_models_vae.py @@ -1017,6 +1017,12 @@ def test_stable_diffusion_encode_sample(self, seed, expected_slice): @slow class ConsistencyDecoderVAEIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/animatediff/test_animatediff.py b/tests/pipelines/animatediff/test_animatediff.py index ca6045ac397f..802da19ba654 100644 --- a/tests/pipelines/animatediff/test_animatediff.py +++ b/tests/pipelines/animatediff/test_animatediff.py @@ -342,6 +342,12 @@ def test_vae_slicing(self): @slow @require_torch_gpu class AnimateDiffPipelineSlowTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/audioldm/test_audioldm.py b/tests/pipelines/audioldm/test_audioldm.py index 84b578846242..f83dc8158e83 100644 --- a/tests/pipelines/audioldm/test_audioldm.py +++ b/tests/pipelines/audioldm/test_audioldm.py @@ -371,6 +371,11 @@ def test_xformers_attention_forwardGenerator_pass(self): @nightly class AudioLDMPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() @@ -411,6 +416,11 @@ def test_audioldm(self): @nightly class AudioLDMPipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/audioldm2/test_audioldm2.py b/tests/pipelines/audioldm2/test_audioldm2.py index 58b1aef5e967..ca8652a0b555 100644 --- a/tests/pipelines/audioldm2/test_audioldm2.py +++ b/tests/pipelines/audioldm2/test_audioldm2.py @@ -493,6 +493,11 @@ def test_sequential_cpu_offload_forward_pass(self): @nightly class AudioLDM2PipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/consistency_models/test_consistency_models.py b/tests/pipelines/consistency_models/test_consistency_models.py index 2cf7c0adb451..e255cb510c42 100644 --- a/tests/pipelines/consistency_models/test_consistency_models.py +++ b/tests/pipelines/consistency_models/test_consistency_models.py @@ -170,6 +170,11 @@ def test_consistency_model_pipeline_onestep_class_cond(self): @nightly @require_torch_gpu class ConsistencyModelPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/controlnet/test_controlnet.py b/tests/pipelines/controlnet/test_controlnet.py index 9d2c3f2edd0d..8429d7b2d5cf 100644 --- a/tests/pipelines/controlnet/test_controlnet.py +++ b/tests/pipelines/controlnet/test_controlnet.py @@ -702,6 +702,11 @@ def test_save_pretrained_raise_not_implemented_exception(self): @slow @require_torch_gpu class ControlNetPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() @@ -1153,6 +1158,11 @@ def test_single_file_component_configs(self): @slow @require_torch_gpu class StableDiffusionMultiControlNetPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/controlnet/test_controlnet_img2img.py b/tests/pipelines/controlnet/test_controlnet_img2img.py index 85312656e825..fe6a0126118d 100644 --- a/tests/pipelines/controlnet/test_controlnet_img2img.py +++ b/tests/pipelines/controlnet/test_controlnet_img2img.py @@ -394,6 +394,11 @@ def test_save_pretrained_raise_not_implemented_exception(self): @slow @require_torch_gpu class ControlNetImg2ImgPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/controlnet/test_controlnet_inpaint.py b/tests/pipelines/controlnet/test_controlnet_inpaint.py index 6a4ec79d1eb2..69379b0fa3c1 100644 --- a/tests/pipelines/controlnet/test_controlnet_inpaint.py +++ b/tests/pipelines/controlnet/test_controlnet_inpaint.py @@ -445,6 +445,11 @@ def test_save_pretrained_raise_not_implemented_exception(self): @slow @require_torch_gpu class ControlNetInpaintPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/controlnet/test_controlnet_sdxl.py b/tests/pipelines/controlnet/test_controlnet_sdxl.py index 3bbb54b0a433..41fb9a1b62fc 100644 --- a/tests/pipelines/controlnet/test_controlnet_sdxl.py +++ b/tests/pipelines/controlnet/test_controlnet_sdxl.py @@ -893,6 +893,11 @@ def test_negative_conditions(self): @slow @require_torch_gpu class ControlNetSDXLPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/dance_diffusion/test_dance_diffusion.py b/tests/pipelines/dance_diffusion/test_dance_diffusion.py index e40d81301650..1f60c0b421f3 100644 --- a/tests/pipelines/dance_diffusion/test_dance_diffusion.py +++ b/tests/pipelines/dance_diffusion/test_dance_diffusion.py @@ -118,6 +118,12 @@ def test_inference_batch_single_identical(self): @nightly @require_torch_gpu class PipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/dit/test_dit.py b/tests/pipelines/dit/test_dit.py index 13e72f2ea287..937265ab05e0 100644 --- a/tests/pipelines/dit/test_dit.py +++ b/tests/pipelines/dit/test_dit.py @@ -109,6 +109,11 @@ def test_xformers_attention_forwardGenerator_pass(self): @nightly @require_torch_gpu class DiTPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/i2vgen_xl/test_i2vgenxl.py b/tests/pipelines/i2vgen_xl/test_i2vgenxl.py index aeda67174ad5..0273e972a620 100644 --- a/tests/pipelines/i2vgen_xl/test_i2vgenxl.py +++ b/tests/pipelines/i2vgen_xl/test_i2vgenxl.py @@ -229,6 +229,12 @@ def test_num_videos_per_prompt(self): @slow @require_torch_gpu class I2VGenXLPipelineSlowTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/kandinsky/test_kandinsky_img2img.py b/tests/pipelines/kandinsky/test_kandinsky_img2img.py index 35ced793fccb..ea289c5ccd71 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_img2img.py +++ b/tests/pipelines/kandinsky/test_kandinsky_img2img.py @@ -365,6 +365,12 @@ def test_kandinsky_img2img(self): @nightly @require_torch_gpu class KandinskyImg2ImgPipelineNightlyTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/latent_diffusion/test_latent_diffusion.py b/tests/pipelines/latent_diffusion/test_latent_diffusion.py index b1ff68400769..e751240e43b0 100644 --- a/tests/pipelines/latent_diffusion/test_latent_diffusion.py +++ b/tests/pipelines/latent_diffusion/test_latent_diffusion.py @@ -138,6 +138,11 @@ def test_inference_text2img(self): @nightly @require_torch_gpu class LDMTextToImagePipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() @@ -174,6 +179,11 @@ def test_ldm_default_ddim(self): @nightly @require_torch_gpu class LDMTextToImagePipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion.py b/tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion.py index 9ff75a918084..26417768843c 100644 --- a/tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion.py +++ b/tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion.py @@ -204,6 +204,11 @@ def test_ledits_pp_warmup_steps(self): @slow @require_torch_gpu class LEditsPPPipelineStableDiffusionSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/musicldm/test_musicldm.py b/tests/pipelines/musicldm/test_musicldm.py index 779b0cbf31ed..e51f5103933a 100644 --- a/tests/pipelines/musicldm/test_musicldm.py +++ b/tests/pipelines/musicldm/test_musicldm.py @@ -408,6 +408,11 @@ def test_to_dtype(self): @nightly @require_torch_gpu class MusicLDMPipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/paint_by_example/test_paint_by_example.py b/tests/pipelines/paint_by_example/test_paint_by_example.py index cb76570d6b9a..c71e2d4761c2 100644 --- a/tests/pipelines/paint_by_example/test_paint_by_example.py +++ b/tests/pipelines/paint_by_example/test_paint_by_example.py @@ -174,6 +174,12 @@ def test_inference_batch_single_identical(self): @nightly @require_torch_gpu class PaintByExamplePipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/pixart_alpha/test_pixart.py b/tests/pipelines/pixart_alpha/test_pixart.py index 3d6db5c32ece..dd358af08395 100644 --- a/tests/pipelines/pixart_alpha/test_pixart.py +++ b/tests/pipelines/pixart_alpha/test_pixart.py @@ -332,6 +332,11 @@ class PixArtAlphaPipelineIntegrationTests(unittest.TestCase): ckpt_id_512 = "PixArt-alpha/PixArt-XL-2-512x512" prompt = "A small cactus with a happy face in the Sahara desert." + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/semantic_stable_diffusion/test_semantic_diffusion.py b/tests/pipelines/semantic_stable_diffusion/test_semantic_diffusion.py index 0c9bb630c562..1cc3111c2631 100644 --- a/tests/pipelines/semantic_stable_diffusion/test_semantic_diffusion.py +++ b/tests/pipelines/semantic_stable_diffusion/test_semantic_diffusion.py @@ -37,6 +37,12 @@ class SafeDiffusionPipelineFastTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() @@ -267,6 +273,12 @@ def test_semantic_diffusion_fp16(self): @nightly @require_torch_gpu class SemanticDiffusionPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/shap_e/test_shap_e.py b/tests/pipelines/shap_e/test_shap_e.py index 11595fe3328b..39d1ead17ef1 100644 --- a/tests/pipelines/shap_e/test_shap_e.py +++ b/tests/pipelines/shap_e/test_shap_e.py @@ -224,6 +224,12 @@ def test_sequential_cpu_offload_forward_pass(self): @nightly @require_torch_gpu class ShapEPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/shap_e/test_shap_e_img2img.py b/tests/pipelines/shap_e/test_shap_e_img2img.py index c666b01025dd..f3661355e9dd 100644 --- a/tests/pipelines/shap_e/test_shap_e_img2img.py +++ b/tests/pipelines/shap_e/test_shap_e_img2img.py @@ -250,6 +250,12 @@ def test_sequential_cpu_offload_forward_pass(self): @nightly @require_torch_gpu class ShapEImg2ImgPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/stable_cascade/test_stable_cascade_decoder.py b/tests/pipelines/stable_cascade/test_stable_cascade_decoder.py index 4a8cab77079c..07e4244e3c68 100644 --- a/tests/pipelines/stable_cascade/test_stable_cascade_decoder.py +++ b/tests/pipelines/stable_cascade/test_stable_cascade_decoder.py @@ -311,6 +311,12 @@ def test_stable_cascade_decoder_single_prompt_multiple_image_embeddings_with_gui @slow @require_torch_gpu class StableCascadeDecoderPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/stable_cascade/test_stable_cascade_prior.py b/tests/pipelines/stable_cascade/test_stable_cascade_prior.py index 56e58aeec295..75dc4883119d 100644 --- a/tests/pipelines/stable_cascade/test_stable_cascade_prior.py +++ b/tests/pipelines/stable_cascade/test_stable_cascade_prior.py @@ -313,6 +313,12 @@ def test_stable_cascade_decoder_prompt_embeds(self): @slow @require_torch_gpu class StableCascadePriorPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion.py b/tests/pipelines/stable_diffusion/test_stable_diffusion.py index b2821f71c058..9a71cc462b10 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion.py @@ -1244,6 +1244,11 @@ def test_stable_diffusion_lcm(self): @slow @require_torch_gpu class StableDiffusionPipelineCkptTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() @@ -1338,6 +1343,11 @@ def test_single_file_component_configs(self): @nightly @require_torch_gpu class StableDiffusionPipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py index 48e520e9030b..63534118343a 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py @@ -395,6 +395,11 @@ def callback_on_step_end(pipe, i, t, callback_kwargs): @slow @require_torch_gpu class StableDiffusionImg2ImgPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() @@ -654,6 +659,11 @@ def test_img2img_compile(self): @nightly @require_torch_gpu class StableDiffusionImg2ImgPipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py index 8c6d0caed5c9..955e84e465a8 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py @@ -1089,6 +1089,11 @@ def test_download_ckpt_diff_format_is_same(self): @nightly @require_torch_gpu class StableDiffusionInpaintPipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py index fc6bd2f4e043..a85ea9c2605d 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py @@ -271,6 +271,11 @@ def callback_no_cfg(pipe, i, t, callback_kwargs): @slow @require_torch_gpu class StableDiffusionInstructPix2PixPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py index 0957fd587e1f..5e296f945ded 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py @@ -205,6 +205,11 @@ def tearDownClass(cls): super().tearDownClass() torch.use_deterministic_algorithms(True) + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py index d247efe581d7..97dc88cd3d26 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py @@ -378,6 +378,11 @@ def test_inference_batch_single_identical(self): @slow @require_torch_gpu class StableDiffusionDepth2ImgPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() @@ -517,6 +522,11 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self): @nightly @require_torch_gpu class StableDiffusionImg2ImgPipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_diffedit.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_diffedit.py index 8e7b9b56e0ce..e9a9f79aa989 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_diffedit.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_diffedit.py @@ -293,6 +293,11 @@ def test_inversion_dpm(self): @require_torch_gpu @nightly class StableDiffusionDiffEditPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() @@ -365,6 +370,11 @@ def test_stable_diffusion_diffedit_full(self): @nightly @require_torch_gpu class StableDiffusionDiffEditPipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py index 563d518680d9..e1dc363878c4 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py @@ -156,6 +156,12 @@ def test_inference_batch_single_identical(self): @slow @require_torch_gpu class StableDiffusionInpaintPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_latent_upscale.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_latent_upscale.py index 8434a6245d32..70a6e444bf13 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_latent_upscale.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_latent_upscale.py @@ -243,6 +243,11 @@ def test_float16_inference(self): @require_torch_gpu @slow class StableDiffusionLatentUpscalePipelineIntegrationTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py index 4dd612102166..9118c60181fd 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py @@ -40,6 +40,12 @@ class StableDiffusionUpscalePipelineFastTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() @@ -377,6 +383,12 @@ def test_stable_diffusion_upscale_from_save_pretrained(self): @slow @require_torch_gpu class StableDiffusionUpscalePipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_v_pred.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_v_pred.py index 8ef8e8ab5a34..97b69043e11e 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_v_pred.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_v_pred.py @@ -45,6 +45,12 @@ class StableDiffusion2VPredictionPipelineFastTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() @@ -254,6 +260,12 @@ def test_stable_diffusion_v_pred_fp16(self): @slow @require_torch_gpu class StableDiffusion2VPredictionPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py b/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py index bd721e6c5064..cb951d5dd833 100644 --- a/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py +++ b/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py @@ -593,6 +593,11 @@ def test_inference_batch_single_identical( @slow @require_torch_gpu class StableDiffusionAdapterPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion_image_variation/test_stable_diffusion_image_variation.py b/tests/pipelines/stable_diffusion_image_variation/test_stable_diffusion_image_variation.py index 78d414496562..99cd8b2e7d1b 100644 --- a/tests/pipelines/stable_diffusion_image_variation/test_stable_diffusion_image_variation.py +++ b/tests/pipelines/stable_diffusion_image_variation/test_stable_diffusion_image_variation.py @@ -164,6 +164,11 @@ def test_inference_batch_single_identical(self): @slow @require_torch_gpu class StableDiffusionImageVariationPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() @@ -274,6 +279,11 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self): @nightly @require_torch_gpu class StableDiffusionImageVariationPipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion_k_diffusion/test_stable_diffusion_k_diffusion.py b/tests/pipelines/stable_diffusion_k_diffusion/test_stable_diffusion_k_diffusion.py index 65b4f2317c60..fe78f0ec3c1a 100644 --- a/tests/pipelines/stable_diffusion_k_diffusion/test_stable_diffusion_k_diffusion.py +++ b/tests/pipelines/stable_diffusion_k_diffusion/test_stable_diffusion_k_diffusion.py @@ -29,6 +29,12 @@ @nightly @require_torch_gpu class StableDiffusionPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/stable_diffusion_ldm3d/test_stable_diffusion_ldm3d.py b/tests/pipelines/stable_diffusion_ldm3d/test_stable_diffusion_ldm3d.py index a5de5eff5cab..8f07d02aad5e 100644 --- a/tests/pipelines/stable_diffusion_ldm3d/test_stable_diffusion_ldm3d.py +++ b/tests/pipelines/stable_diffusion_ldm3d/test_stable_diffusion_ldm3d.py @@ -207,6 +207,11 @@ def test_stable_diffusion_negative_prompt(self): @nightly @require_torch_gpu class StableDiffusionLDM3DPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() @@ -253,6 +258,11 @@ def test_ldm3d_stable_diffusion(self): @nightly @require_torch_gpu class StableDiffusionPipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion_panorama/test_stable_diffusion_panorama.py b/tests/pipelines/stable_diffusion_panorama/test_stable_diffusion_panorama.py index ede4552e3499..f275c59c7ca5 100644 --- a/tests/pipelines/stable_diffusion_panorama/test_stable_diffusion_panorama.py +++ b/tests/pipelines/stable_diffusion_panorama/test_stable_diffusion_panorama.py @@ -253,6 +253,11 @@ def test_stable_diffusion_panorama_pndm(self): @nightly @require_torch_gpu class StableDiffusionPanoramaNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion_safe/test_safe_diffusion.py b/tests/pipelines/stable_diffusion_safe/test_safe_diffusion.py index 478e465efe79..14100ea03dc1 100644 --- a/tests/pipelines/stable_diffusion_safe/test_safe_diffusion.py +++ b/tests/pipelines/stable_diffusion_safe/test_safe_diffusion.py @@ -28,6 +28,12 @@ class SafeDiffusionPipelineFastTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() @@ -258,6 +264,12 @@ def test_stable_diffusion_fp16(self): @nightly @require_torch_gpu class SafeDiffusionPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/stable_diffusion_sag/test_stable_diffusion_sag.py b/tests/pipelines/stable_diffusion_sag/test_stable_diffusion_sag.py index 94a5616471ab..8123df3da9c0 100644 --- a/tests/pipelines/stable_diffusion_sag/test_stable_diffusion_sag.py +++ b/tests/pipelines/stable_diffusion_sag/test_stable_diffusion_sag.py @@ -146,6 +146,12 @@ def test_pipeline_different_schedulers(self): @nightly @require_torch_gpu class StableDiffusionPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py index 9670f069a21f..95d6dd69ca9e 100644 --- a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py +++ b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py @@ -1011,6 +1011,11 @@ def callback_on_step_end(pipe, i, t, callback_kwargs): @slow class StableDiffusionXLPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_adapter.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_adapter.py index eb687f8035d4..08bd2a0444cd 100644 --- a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_adapter.py +++ b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_adapter.py @@ -683,6 +683,11 @@ def test_adapter_sdxl_lcm_custom_timesteps(self): @slow @require_torch_gpu class AdapterSDXLPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() diff --git a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_k_diffusion.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_k_diffusion.py index d5d2abf49861..60207d9a0e76 100644 --- a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_k_diffusion.py +++ b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_k_diffusion.py @@ -31,6 +31,12 @@ class StableDiffusionXLKPipelineIntegrationTests(unittest.TestCase): dtype = torch.float16 + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/stable_unclip/test_stable_unclip.py b/tests/pipelines/stable_unclip/test_stable_unclip.py index 080fa5bb3267..eaf4573c123a 100644 --- a/tests/pipelines/stable_unclip/test_stable_unclip.py +++ b/tests/pipelines/stable_unclip/test_stable_unclip.py @@ -188,6 +188,12 @@ def test_inference_batch_single_identical(self): @nightly @require_torch_gpu class StableUnCLIPPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py b/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py index 12f6a9101709..cd2787b2f5e3 100644 --- a/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py +++ b/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py @@ -209,6 +209,12 @@ def test_xformers_attention_forwardGenerator_pass(self): @nightly @require_torch_gpu class StableUnCLIPImg2ImgPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/stable_video_diffusion/test_stable_video_diffusion.py b/tests/pipelines/stable_video_diffusion/test_stable_video_diffusion.py index 33cf4c72863b..199ed57bc27b 100644 --- a/tests/pipelines/stable_video_diffusion/test_stable_video_diffusion.py +++ b/tests/pipelines/stable_video_diffusion/test_stable_video_diffusion.py @@ -516,6 +516,12 @@ def test_disable_cfg(self): @slow @require_torch_gpu class StableVideoDiffusionPipelineSlowTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/test_pipelines.py b/tests/pipelines/test_pipelines.py index 930d78aad031..69d06d980af1 100644 --- a/tests/pipelines/test_pipelines.py +++ b/tests/pipelines/test_pipelines.py @@ -1056,6 +1056,12 @@ def test_save_pipeline_change_config(self): class PipelineFastTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() @@ -1673,6 +1679,12 @@ def test_pipe_same_device_id_offload(self): @slow @require_torch_gpu class PipelineSlowTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() @@ -1898,6 +1910,12 @@ def test_weighted_prompts_compel(self): @nightly @require_torch_gpu class PipelineNightlyTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/unclip/test_unclip.py b/tests/pipelines/unclip/test_unclip.py index e3b2222b07f0..d3d026a448ca 100644 --- a/tests/pipelines/unclip/test_unclip.py +++ b/tests/pipelines/unclip/test_unclip.py @@ -421,6 +421,12 @@ def test_float16_inference(self): @nightly class UnCLIPPipelineCPUIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() @@ -453,6 +459,12 @@ def test_unclip_karlo_cpu_fp32(self): @nightly @require_torch_gpu class UnCLIPPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/unclip/test_unclip_image_variation.py b/tests/pipelines/unclip/test_unclip_image_variation.py index ab3aea59631a..dfc3acc0c0f2 100644 --- a/tests/pipelines/unclip/test_unclip_image_variation.py +++ b/tests/pipelines/unclip/test_unclip_image_variation.py @@ -496,6 +496,12 @@ def test_float16_inference(self): @nightly @require_torch_gpu class UnCLIPImageVariationPipelineIntegrationTests(unittest.TestCase): + def setUp(self): + # clean up the VRAM before each test + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): # clean up the VRAM after each test super().tearDown() diff --git a/tests/pipelines/unidiffuser/test_unidiffuser.py b/tests/pipelines/unidiffuser/test_unidiffuser.py index 11a02a656d3f..561b82aafbda 100644 --- a/tests/pipelines/unidiffuser/test_unidiffuser.py +++ b/tests/pipelines/unidiffuser/test_unidiffuser.py @@ -574,6 +574,11 @@ def test_unidiffuser_default_img2text_v1_cuda_fp16(self): @nightly @require_torch_gpu class UniDiffuserPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect() @@ -690,6 +695,11 @@ def test_unidiffuser_compile(self, seed=0): @nightly @require_torch_gpu class UniDiffuserPipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() + gc.collect() + torch.cuda.empty_cache() + def tearDown(self): super().tearDown() gc.collect()