Skip to content

Commit ec1aded

Browse files
authored
Optimize test files by fixing CPU-offloading usage (huggingface#8409)
* 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
1 parent 151a56b commit ec1aded

16 files changed

+1
-21
lines changed

docs/source/en/api/pipelines/animatediff.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,7 @@ from PIL import Image
165165
adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2", torch_dtype=torch.float16)
166166
# load SD 1.5 based finetuned model
167167
model_id = "SG161222/Realistic_Vision_V5.1_noVAE"
168-
pipe = AnimateDiffVideoToVideoPipeline.from_pretrained(model_id, motion_adapter=adapter, torch_dtype=torch.float16).to("cuda")
168+
pipe = AnimateDiffVideoToVideoPipeline.from_pretrained(model_id, motion_adapter=adapter, torch_dtype=torch.float16)
169169
scheduler = DDIMScheduler.from_pretrained(
170170
model_id,
171171
subfolder="scheduler",

tests/pipelines/stable_diffusion/test_stable_diffusion.py

-1
Original file line numberDiff line numberDiff line change
@@ -1135,7 +1135,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
11351135
torch.cuda.reset_peak_memory_stats()
11361136

11371137
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
1138-
pipe = pipe.to(torch_device)
11391138
pipe.set_progress_bar_config(disable=None)
11401139
pipe.enable_attention_slicing(1)
11411140
pipe.enable_sequential_cpu_offload()

tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py

-1
Original file line numberDiff line numberDiff line change
@@ -513,7 +513,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
513513
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
514514
"CompVis/stable-diffusion-v1-4", safety_checker=None, torch_dtype=torch.float16
515515
)
516-
pipe = pipe.to(torch_device)
517516
pipe.set_progress_bar_config(disable=None)
518517
pipe.enable_attention_slicing(1)
519518
pipe.enable_sequential_cpu_offload()

tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py

-2
Original file line numberDiff line numberDiff line change
@@ -705,7 +705,6 @@ def test_stable_diffusion_inpaint_with_sequential_cpu_offloading(self):
705705
pipe = StableDiffusionInpaintPipeline.from_pretrained(
706706
"runwayml/stable-diffusion-inpainting", safety_checker=None, torch_dtype=torch.float16
707707
)
708-
pipe = pipe.to(torch_device)
709708
pipe.set_progress_bar_config(disable=None)
710709
pipe.enable_attention_slicing(1)
711710
pipe.enable_sequential_cpu_offload()
@@ -911,7 +910,6 @@ def test_stable_diffusion_inpaint_with_sequential_cpu_offloading(self):
911910
"runwayml/stable-diffusion-inpainting", safety_checker=None, torch_dtype=torch.float16
912911
)
913912
pipe.vae = vae
914-
pipe = pipe.to(torch_device)
915913
pipe.set_progress_bar_config(disable=None)
916914
pipe.enable_attention_slicing(1)
917915
pipe.enable_sequential_cpu_offload()

tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py

-1
Original file line numberDiff line numberDiff line change
@@ -394,7 +394,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
394394
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
395395
"timbrooks/instruct-pix2pix", safety_checker=None, torch_dtype=torch.float16
396396
)
397-
pipe = pipe.to(torch_device)
398397
pipe.set_progress_bar_config(disable=None)
399398
pipe.enable_attention_slicing(1)
400399
pipe.enable_sequential_cpu_offload()

tests/pipelines/stable_diffusion_2/test_stable_diffusion.py

-1
Original file line numberDiff line numberDiff line change
@@ -462,7 +462,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
462462
pipe = StableDiffusionPipeline.from_pretrained(
463463
"stabilityai/stable-diffusion-2-base", torch_dtype=torch.float16
464464
)
465-
pipe = pipe.to(torch_device)
466465
pipe.set_progress_bar_config(disable=None)
467466
pipe.enable_attention_slicing(1)
468467
pipe.enable_sequential_cpu_offload()

tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py

-1
Original file line numberDiff line numberDiff line change
@@ -506,7 +506,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
506506
pipe = StableDiffusionDepth2ImgPipeline.from_pretrained(
507507
"stabilityai/stable-diffusion-2-depth", safety_checker=None, torch_dtype=torch.float16
508508
)
509-
pipe = pipe.to(torch_device)
510509
pipe.set_progress_bar_config(disable=None)
511510
pipe.enable_attention_slicing(1)
512511
pipe.enable_sequential_cpu_offload()

tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py

-1
Original file line numberDiff line numberDiff line change
@@ -261,7 +261,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
261261
scheduler=pndm,
262262
torch_dtype=torch.float16,
263263
)
264-
pipe.to(torch_device)
265264
pipe.set_progress_bar_config(disable=None)
266265
pipe.enable_attention_slicing(1)
267266
pipe.enable_sequential_cpu_offload()

tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py

-1
Original file line numberDiff line numberDiff line change
@@ -472,7 +472,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
472472
model_id,
473473
torch_dtype=torch.float16,
474474
)
475-
pipe.to(torch_device)
476475
pipe.set_progress_bar_config(disable=None)
477476
pipe.enable_attention_slicing(1)
478477
pipe.enable_sequential_cpu_offload()

tests/pipelines/stable_diffusion_2/test_stable_diffusion_v_pred.py

-1
Original file line numberDiff line numberDiff line change
@@ -537,7 +537,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading_v_pred(self):
537537
prompt = "Andromeda galaxy in a bottle"
538538

539539
pipeline = StableDiffusionPipeline.from_pretrained(pipeline_id, torch_dtype=torch.float16)
540-
pipeline = pipeline.to(torch_device)
541540
pipeline.enable_attention_slicing(1)
542541
pipeline.enable_sequential_cpu_offload()
543542

tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py

-2
Original file line numberDiff line numberDiff line change
@@ -809,7 +809,6 @@ def test_stable_diffusion_adapter_zoedepth_sd_v15(self):
809809
adapter = T2IAdapter.from_pretrained(adapter_model, torch_dtype=torch.float16)
810810

811811
pipe = StableDiffusionAdapterPipeline.from_pretrained(sd_model, adapter=adapter, safety_checker=None)
812-
pipe.to(torch_device)
813812
pipe.set_progress_bar_config(disable=None)
814813
pipe.enable_model_cpu_offload()
815814
generator = torch.Generator(device="cpu").manual_seed(0)
@@ -942,7 +941,6 @@ def test_stable_diffusion_adapter_pipeline_with_sequential_cpu_offloading(self):
942941
pipe = StableDiffusionAdapterPipeline.from_pretrained(
943942
"CompVis/stable-diffusion-v1-4", adapter=adapter, safety_checker=None
944943
)
945-
pipe = pipe.to(torch_device)
946944
pipe.set_progress_bar_config(disable=None)
947945
pipe.enable_attention_slicing(1)
948946
pipe.enable_sequential_cpu_offload()

tests/pipelines/stable_diffusion_image_variation/test_stable_diffusion_image_variation.py

-1
Original file line numberDiff line numberDiff line change
@@ -263,7 +263,6 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
263263
pipe = StableDiffusionImageVariationPipeline.from_pretrained(
264264
"lambdalabs/sd-image-variations-diffusers", safety_checker=None, torch_dtype=torch.float16
265265
)
266-
pipe = pipe.to(torch_device)
267266
pipe.set_progress_bar_config(disable=None)
268267
pipe.enable_attention_slicing(1)
269268
pipe.enable_sequential_cpu_offload()

tests/pipelines/stable_diffusion_panorama/test_stable_diffusion_panorama.py

-1
Original file line numberDiff line numberDiff line change
@@ -415,7 +415,6 @@ def test_stable_diffusion_panorama_pipeline_with_sequential_cpu_offloading(self)
415415
model_ckpt = "stabilityai/stable-diffusion-2-base"
416416
scheduler = DDIMScheduler.from_pretrained(model_ckpt, subfolder="scheduler")
417417
pipe = StableDiffusionPanoramaPipeline.from_pretrained(model_ckpt, scheduler=scheduler, safety_checker=None)
418-
pipe = pipe.to(torch_device)
419418
pipe.set_progress_bar_config(disable=None)
420419
pipe.enable_attention_slicing(1)
421420
pipe.enable_sequential_cpu_offload()

tests/pipelines/stable_unclip/test_stable_unclip.py

-2
Original file line numberDiff line numberDiff line change
@@ -206,7 +206,6 @@ def test_stable_unclip(self):
206206
)
207207

208208
pipe = StableUnCLIPPipeline.from_pretrained("fusing/stable-unclip-2-1-l", torch_dtype=torch.float16)
209-
pipe.to(torch_device)
210209
pipe.set_progress_bar_config(disable=None)
211210
# stable unclip will oom when integration tests are run on a V100,
212211
# so turn on memory savings
@@ -228,7 +227,6 @@ def test_stable_unclip_pipeline_with_sequential_cpu_offloading(self):
228227
torch.cuda.reset_peak_memory_stats()
229228

230229
pipe = StableUnCLIPPipeline.from_pretrained("fusing/stable-unclip-2-1-l", torch_dtype=torch.float16)
231-
pipe = pipe.to(torch_device)
232230
pipe.set_progress_bar_config(disable=None)
233231
pipe.enable_attention_slicing()
234232
pipe.enable_sequential_cpu_offload()

tests/pipelines/stable_unclip/test_stable_unclip_img2img.py

-3
Original file line numberDiff line numberDiff line change
@@ -233,7 +233,6 @@ def test_stable_unclip_l_img2img(self):
233233
pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
234234
"fusing/stable-unclip-2-1-l-img2img", torch_dtype=torch.float16
235235
)
236-
pipe.to(torch_device)
237236
pipe.set_progress_bar_config(disable=None)
238237
# stable unclip will oom when integration tests are run on a V100,
239238
# so turn on memory savings
@@ -261,7 +260,6 @@ def test_stable_unclip_h_img2img(self):
261260
pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
262261
"fusing/stable-unclip-2-1-h-img2img", torch_dtype=torch.float16
263262
)
264-
pipe.to(torch_device)
265263
pipe.set_progress_bar_config(disable=None)
266264
# stable unclip will oom when integration tests are run on a V100,
267265
# so turn on memory savings
@@ -289,7 +287,6 @@ def test_stable_unclip_img2img_pipeline_with_sequential_cpu_offloading(self):
289287
pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
290288
"fusing/stable-unclip-2-1-h-img2img", torch_dtype=torch.float16
291289
)
292-
pipe = pipe.to(torch_device)
293290
pipe.set_progress_bar_config(disable=None)
294291
pipe.enable_attention_slicing()
295292
pipe.enable_sequential_cpu_offload()

tests/pipelines/unclip/test_unclip.py

-1
Original file line numberDiff line numberDiff line change
@@ -500,7 +500,6 @@ def test_unclip_pipeline_with_sequential_cpu_offloading(self):
500500
torch.cuda.reset_peak_memory_stats()
501501

502502
pipe = UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha", torch_dtype=torch.float16)
503-
pipe = pipe.to(torch_device)
504503
pipe.set_progress_bar_config(disable=None)
505504
pipe.enable_attention_slicing()
506505
pipe.enable_sequential_cpu_offload()

0 commit comments

Comments
 (0)