Optimize test files by fixing CPU-offloading usage (huggingface#8409)

tolgacangoz · web-flow · commit ec1aded12e8f · 2024-06-06T09:51:26.000-10:00
* Refactor code to remove unnecessary calls to `to(torch_device)`

* Refactor code to remove unnecessary calls to `to("cuda")`

* Update pipeline_stable_diffusion_diffedit.py
diff --git a/docs/source/en/api/pipelines/animatediff.md b/docs/source/en/api/pipelines/animatediff.md
@@ -165,7 +165,7 @@ from PIL import Image
 adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2", torch_dtype=torch.float16)
 # load SD 1.5 based finetuned model
 model_id = "SG161222/Realistic_Vision_V5.1_noVAE"
-pipe = AnimateDiffVideoToVideoPipeline.from_pretrained(model_id, motion_adapter=adapter, torch_dtype=torch.float16).to("cuda")
+pipe = AnimateDiffVideoToVideoPipeline.from_pretrained(model_id, motion_adapter=adapter, torch_dtype=torch.float16)
 scheduler = DDIMScheduler.from_pretrained(
     model_id,
     subfolder="scheduler",
diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion.py b/tests/pipelines/stable_diffusion/test_stable_diffusion.py
@@ -1135,7 +1135,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
         torch.cuda.reset_peak_memory_stats()
 
         pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
-        pipe = pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing(1)
         pipe.enable_sequential_cpu_offload()
diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py
@@ -513,7 +513,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
         pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4", safety_checker=None, torch_dtype=torch.float16
         )
-        pipe = pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing(1)
         pipe.enable_sequential_cpu_offload()
diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py
@@ -705,7 +705,6 @@ def test_stable_diffusion_inpaint_with_sequential_cpu_offloading(self):
         pipe = StableDiffusionInpaintPipeline.from_pretrained(
             "runwayml/stable-diffusion-inpainting", safety_checker=None, torch_dtype=torch.float16
         )
-        pipe = pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing(1)
         pipe.enable_sequential_cpu_offload()
@@ -911,7 +910,6 @@ def test_stable_diffusion_inpaint_with_sequential_cpu_offloading(self):
             "runwayml/stable-diffusion-inpainting", safety_checker=None, torch_dtype=torch.float16
         )
         pipe.vae = vae
-        pipe = pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing(1)
         pipe.enable_sequential_cpu_offload()
diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py
@@ -394,7 +394,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
         pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
             "timbrooks/instruct-pix2pix", safety_checker=None, torch_dtype=torch.float16
         )
-        pipe = pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing(1)
         pipe.enable_sequential_cpu_offload()
diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion.py
@@ -462,7 +462,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
         pipe = StableDiffusionPipeline.from_pretrained(
             "stabilityai/stable-diffusion-2-base", torch_dtype=torch.float16
         )
-        pipe = pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing(1)
         pipe.enable_sequential_cpu_offload()
diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py
@@ -506,7 +506,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
         pipe = StableDiffusionDepth2ImgPipeline.from_pretrained(
             "stabilityai/stable-diffusion-2-depth", safety_checker=None, torch_dtype=torch.float16
         )
-        pipe = pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing(1)
         pipe.enable_sequential_cpu_offload()
diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py
@@ -261,7 +261,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
             scheduler=pndm,
             torch_dtype=torch.float16,
         )
-        pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing(1)
         pipe.enable_sequential_cpu_offload()
diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py
@@ -472,7 +472,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
             model_id,
             torch_dtype=torch.float16,
         )
-        pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing(1)
         pipe.enable_sequential_cpu_offload()
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
@@ -537,7 +537,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading_v_pred(self):
         prompt = "Andromeda galaxy in a bottle"
 
         pipeline = StableDiffusionPipeline.from_pretrained(pipeline_id, torch_dtype=torch.float16)
-        pipeline = pipeline.to(torch_device)
         pipeline.enable_attention_slicing(1)
         pipeline.enable_sequential_cpu_offload()
 
diff --git a/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py b/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py
@@ -809,7 +809,6 @@ def test_stable_diffusion_adapter_zoedepth_sd_v15(self):
         adapter = T2IAdapter.from_pretrained(adapter_model, torch_dtype=torch.float16)
 
         pipe = StableDiffusionAdapterPipeline.from_pretrained(sd_model, adapter=adapter, safety_checker=None)
-        pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_model_cpu_offload()
         generator = torch.Generator(device="cpu").manual_seed(0)
@@ -942,7 +941,6 @@ def test_stable_diffusion_adapter_pipeline_with_sequential_cpu_offloading(self):
         pipe = StableDiffusionAdapterPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4", adapter=adapter, safety_checker=None
         )
-        pipe = pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing(1)
         pipe.enable_sequential_cpu_offload()
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
@@ -263,7 +263,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
         pipe = StableDiffusionImageVariationPipeline.from_pretrained(
             "lambdalabs/sd-image-variations-diffusers", safety_checker=None, torch_dtype=torch.float16
         )
-        pipe = pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing(1)
         pipe.enable_sequential_cpu_offload()
diff --git a/tests/pipelines/stable_diffusion_panorama/test_stable_diffusion_panorama.py b/tests/pipelines/stable_diffusion_panorama/test_stable_diffusion_panorama.py
@@ -415,7 +415,6 @@ def test_stable_diffusion_panorama_pipeline_with_sequential_cpu_offloading(self)
         model_ckpt = "stabilityai/stable-diffusion-2-base"
         scheduler = DDIMScheduler.from_pretrained(model_ckpt, subfolder="scheduler")
         pipe = StableDiffusionPanoramaPipeline.from_pretrained(model_ckpt, scheduler=scheduler, safety_checker=None)
-        pipe = pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing(1)
         pipe.enable_sequential_cpu_offload()
diff --git a/tests/pipelines/stable_unclip/test_stable_unclip.py b/tests/pipelines/stable_unclip/test_stable_unclip.py
@@ -206,7 +206,6 @@ def test_stable_unclip(self):
         )
 
         pipe = StableUnCLIPPipeline.from_pretrained("fusing/stable-unclip-2-1-l", torch_dtype=torch.float16)
-        pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         # stable unclip will oom when integration tests are run on a V100,
         # so turn on memory savings
@@ -228,7 +227,6 @@ def test_stable_unclip_pipeline_with_sequential_cpu_offloading(self):
         torch.cuda.reset_peak_memory_stats()
 
         pipe = StableUnCLIPPipeline.from_pretrained("fusing/stable-unclip-2-1-l", torch_dtype=torch.float16)
-        pipe = pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing()
         pipe.enable_sequential_cpu_offload()
diff --git a/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py b/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py
@@ -233,7 +233,6 @@ def test_stable_unclip_l_img2img(self):
         pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
             "fusing/stable-unclip-2-1-l-img2img", torch_dtype=torch.float16
         )
-        pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         # stable unclip will oom when integration tests are run on a V100,
         # so turn on memory savings
@@ -261,7 +260,6 @@ def test_stable_unclip_h_img2img(self):
         pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
             "fusing/stable-unclip-2-1-h-img2img", torch_dtype=torch.float16
         )
-        pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         # stable unclip will oom when integration tests are run on a V100,
         # so turn on memory savings
@@ -289,7 +287,6 @@ def test_stable_unclip_img2img_pipeline_with_sequential_cpu_offloading(self):
         pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
             "fusing/stable-unclip-2-1-h-img2img", torch_dtype=torch.float16
         )
-        pipe = pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing()
         pipe.enable_sequential_cpu_offload()
diff --git a/tests/pipelines/unclip/test_unclip.py b/tests/pipelines/unclip/test_unclip.py
@@ -500,7 +500,6 @@ def test_unclip_pipeline_with_sequential_cpu_offloading(self):
         torch.cuda.reset_peak_memory_stats()
 
         pipe = UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha", torch_dtype=torch.float16)
-        pipe = pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
         pipe.enable_attention_slicing()
         pipe.enable_sequential_cpu_offload()

Original file line number	Diff line number	Diff line change
`@@ -513,7 +513,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):`
`513`	`513`	`pipe = StableDiffusionImg2ImgPipeline.from_pretrained(`
`514`	`514`	`"CompVis/stable-diffusion-v1-4", safety_checker=None, torch_dtype=torch.float16`
`515`	`515`	`)`
`516`		`- pipe = pipe.to(torch_device)`
`517`	`516`	`pipe.set_progress_bar_config(disable=None)`
`518`	`517`	`pipe.enable_attention_slicing(1)`
`519`	`518`	`pipe.enable_sequential_cpu_offload()`
Original file line number	Diff line number	Diff line change
`@@ -394,7 +394,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):`
`394`	`394`	`pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(`
`395`	`395`	`"timbrooks/instruct-pix2pix", safety_checker=None, torch_dtype=torch.float16`
`396`	`396`	`)`
`397`		`- pipe = pipe.to(torch_device)`
`398`	`397`	`pipe.set_progress_bar_config(disable=None)`
`399`	`398`	`pipe.enable_attention_slicing(1)`
`400`	`399`	`pipe.enable_sequential_cpu_offload()`
Original file line number	Diff line number	Diff line change
`@@ -462,7 +462,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):`
`462`	`462`	`pipe = StableDiffusionPipeline.from_pretrained(`
`463`	`463`	`"stabilityai/stable-diffusion-2-base", torch_dtype=torch.float16`
`464`	`464`	`)`
`465`		`- pipe = pipe.to(torch_device)`
`466`	`465`	`pipe.set_progress_bar_config(disable=None)`
`467`	`466`	`pipe.enable_attention_slicing(1)`
`468`	`467`	`pipe.enable_sequential_cpu_offload()`
Original file line number	Diff line number	Diff line change
`@@ -506,7 +506,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):`
`506`	`506`	`pipe = StableDiffusionDepth2ImgPipeline.from_pretrained(`
`507`	`507`	`"stabilityai/stable-diffusion-2-depth", safety_checker=None, torch_dtype=torch.float16`
`508`	`508`	`)`
`509`		`- pipe = pipe.to(torch_device)`
`510`	`509`	`pipe.set_progress_bar_config(disable=None)`
`511`	`510`	`pipe.enable_attention_slicing(1)`
`512`	`511`	`pipe.enable_sequential_cpu_offload()`
Original file line number	Diff line number	Diff line change
`@@ -261,7 +261,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):`
`261`	`261`	`scheduler=pndm,`
`262`	`262`	`torch_dtype=torch.float16,`
`263`	`263`	`)`
`264`		`- pipe.to(torch_device)`
`265`	`264`	`pipe.set_progress_bar_config(disable=None)`
`266`	`265`	`pipe.enable_attention_slicing(1)`
`267`	`266`	`pipe.enable_sequential_cpu_offload()`
Original file line number	Diff line number	Diff line change
`@@ -472,7 +472,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):`
`472`	`472`	`model_id,`
`473`	`473`	`torch_dtype=torch.float16,`
`474`	`474`	`)`
`475`		`- pipe.to(torch_device)`
`476`	`475`	`pipe.set_progress_bar_config(disable=None)`
`477`	`476`	`pipe.enable_attention_slicing(1)`
`478`	`477`	`pipe.enable_sequential_cpu_offload()`