revert back

Author: faaany

tests/pipelines/pag/test_pag_sd_img2img.py

@@ -269,7 +269,7 @@ def test_pag_uncond(self):
         pipeline.enable_model_cpu_offload(device=torch_device)
         pipeline.set_progress_bar_config(disable=None)
 
-        inputs = self.get_inputs(torch_device)
+        inputs = self.get_inputs(torch_device, guidance_scale=0.0)
         image = pipeline(**inputs).images
 
         image_slice = image[0, -3:, -3:, -1].flatten()

tests/pipelines/stable_diffusion/test_stable_diffusion.py

@@ -46,6 +46,7 @@
     CaptureLogger,
     backend_empty_cache,
     backend_max_memory_allocated,
+    backend_reset_max_memory_allocated,
     backend_reset_peak_memory_stats,
     enable_full_determinism,
     is_torch_compile,
@@ -1142,6 +1143,7 @@ def test_stable_diffusion_low_cpu_mem_usage(self):
 
     def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
         backend_empty_cache(torch_device)
+        backend_reset_max_memory_allocated(torch_device)
         backend_reset_peak_memory_stats(torch_device)
 
         pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)

tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py

@@ -37,6 +37,7 @@
 from diffusers.utils.testing_utils import (
     backend_empty_cache,
     backend_max_memory_allocated,
+    backend_reset_max_memory_allocated,
     backend_reset_peak_memory_stats,
     enable_full_determinism,
     floats_tensor,
@@ -511,6 +512,7 @@ def callback_fn(step: int, timestep: int, latents: torch.Tensor) -> None:
     def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
         backend_empty_cache(torch_device)
         backend_reset_peak_memory_stats(torch_device)
+        backend_reset_max_memory_allocated(torch_device)
 
         pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4", safety_checker=None, torch_dtype=torch.float16
@@ -529,6 +531,7 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
     def test_stable_diffusion_pipeline_with_model_offloading(self):
         backend_empty_cache(torch_device)
         backend_reset_peak_memory_stats(torch_device)
+        backend_reset_max_memory_allocated(torch_device)
 
         inputs = self.get_inputs(torch_device, dtype=torch.float16)
 

tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py

@@ -39,6 +39,7 @@
 from diffusers.utils.testing_utils import (
     backend_empty_cache,
     backend_max_memory_allocated,
+    backend_reset_max_memory_allocated,
     backend_reset_peak_memory_stats,
     enable_full_determinism,
     floats_tensor,
@@ -702,6 +703,7 @@ def test_stable_diffusion_inpaint_k_lms(self):
     def test_stable_diffusion_inpaint_with_sequential_cpu_offloading(self):
         backend_empty_cache(torch_device)
         backend_reset_peak_memory_stats(torch_device)
+        backend_reset_max_memory_allocated(torch_device)
 
         pipe = StableDiffusionInpaintPipeline.from_pretrained(
             "botp/stable-diffusion-v1-5-inpainting", safety_checker=None, torch_dtype=torch.float16
@@ -904,6 +906,7 @@ def test_stable_diffusion_inpaint_k_lms(self):
     def test_stable_diffusion_inpaint_with_sequential_cpu_offloading(self):
         backend_empty_cache(torch_device)
         backend_reset_peak_memory_stats(torch_device)
+        backend_reset_max_memory_allocated(torch_device)
 
         vae = AsymmetricAutoencoderKL.from_pretrained(
             "cross-attention/asymmetric-autoencoder-kl-x-1-5", torch_dtype=torch.float16

tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py

@@ -35,6 +35,7 @@
 from diffusers.utils.testing_utils import (
     backend_empty_cache,
     backend_max_memory_allocated,
+    backend_reset_max_memory_allocated,
     backend_reset_peak_memory_stats,
     enable_full_determinism,
     floats_tensor,
@@ -389,6 +390,7 @@ def callback_fn(step: int, timestep: int, latents: torch.Tensor) -> None:
     def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
         backend_empty_cache(torch_device)
         backend_reset_peak_memory_stats(torch_device)
+        backend_reset_max_memory_allocated(torch_device)
 
         pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
             "timbrooks/instruct-pix2pix", safety_checker=None, torch_dtype=torch.float16

tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py

@@ -25,6 +25,7 @@
 from diffusers import AutoencoderKL, PNDMScheduler, StableDiffusionInpaintPipeline, UNet2DConditionModel
 from diffusers.utils.testing_utils import (
     backend_empty_cache,
+    backend_reset_max_memory_allocated,
     backend_reset_peak_memory_stats,
     enable_full_determinism,
     floats_tensor,
@@ -245,6 +246,7 @@ def test_stable_diffusion_inpaint_pipeline_fp16(self):
     def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
         backend_empty_cache(torch_device)
         backend_reset_peak_memory_stats(torch_device)
+        backend_reset_max_memory_allocated(torch_device)
 
         init_image = load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"

tests/pipelines/stable_diffusion_2/test_stable_diffusion_v_pred.py

@@ -33,6 +33,7 @@
 from diffusers.utils.testing_utils import (
     backend_empty_cache,
     backend_max_memory_allocated,
+    backend_reset_max_memory_allocated,
     backend_reset_peak_memory_stats,
     enable_full_determinism,
     load_numpy,
@@ -535,6 +536,7 @@ def test_stable_diffusion_low_cpu_mem_usage_v_pred(self):
     def test_stable_diffusion_pipeline_with_sequential_cpu_offloading_v_pred(self):
         backend_empty_cache(torch_device)
         backend_reset_peak_memory_stats(torch_device)
+        backend_reset_max_memory_allocated(torch_device)
 
         pipeline_id = "stabilityai/stable-diffusion-2"
         prompt = "Andromeda galaxy in a bottle"

tests/pipelines/stable_diffusion_image_variation/test_stable_diffusion_image_variation.py

@@ -32,6 +32,7 @@
 from diffusers.utils.testing_utils import (
     backend_empty_cache,
     backend_max_memory_allocated,
+    backend_reset_max_memory_allocated,
     backend_reset_peak_memory_stats,
     enable_full_determinism,
     floats_tensor,
@@ -263,6 +264,7 @@ def callback_fn(step: int, timestep: int, latents: torch.Tensor) -> None:
     def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
         backend_empty_cache(torch_device)
         backend_reset_peak_memory_stats(torch_device)
+        backend_reset_max_memory_allocated(torch_device)
 
         pipe = StableDiffusionImageVariationPipeline.from_pretrained(
             "lambdalabs/sd-image-variations-diffusers", safety_checker=None, torch_dtype=torch.float16