Skip to content

Commit aa1f00f

Browse files
yiyixuxuyiyixuxu
yiyixuxu
and
yiyixuxu
authored
Fix cpu offload related slow tests (huggingface#7618)
* fix * up --------- Co-authored-by: yiyixuxu <yixu310@gmail,com>
1 parent d95b993 commit aa1f00f

10 files changed

+12
-145
lines changed

‎src/diffusers/pipelines/deepfloyd_if/pipeline_if.py‎

Lines changed: 1 addition & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
from ...schedulers import DDPMScheduler
1313
from ...utils import (
1414
BACKENDS_MAPPING,
15-
is_accelerate_available,
1615
is_bs4_available,
1716
is_ftfy_available,
1817
logging,
@@ -115,6 +114,7 @@ class IFPipeline(DiffusionPipeline, LoraLoaderMixin):
115114

116115
_optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor", "watermarker"]
117116
model_cpu_offload_seq = "text_encoder->unet"
117+
_exclude_from_cpu_offload = ["watermarker"]
118118

119119
def __init__(
120120
self,
@@ -156,20 +156,6 @@ def __init__(
156156
)
157157
self.register_to_config(requires_safety_checker=requires_safety_checker)
158158

159-
def remove_all_hooks(self):
160-
if is_accelerate_available():
161-
from accelerate.hooks import remove_hook_from_module
162-
else:
163-
raise ImportError("Please install accelerate via `pip install accelerate`")
164-
165-
for model in [self.text_encoder, self.unet, self.safety_checker]:
166-
if model is not None:
167-
remove_hook_from_module(model, recurse=True)
168-
169-
self.unet_offload_hook = None
170-
self.text_encoder_offload_hook = None
171-
self.final_offload_hook = None
172-
173159
@torch.no_grad()
174160
def encode_prompt(
175161
self,
@@ -335,9 +321,6 @@ def run_safety_checker(self, image, device, dtype):
335321
nsfw_detected = None
336322
watermark_detected = None
337323

338-
if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None:
339-
self.unet_offload_hook.offload()
340-
341324
return image, nsfw_detected, watermark_detected
342325

343326
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs

‎src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py‎

Lines changed: 1 addition & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
from ...utils import (
1616
BACKENDS_MAPPING,
1717
PIL_INTERPOLATION,
18-
is_accelerate_available,
1918
is_bs4_available,
2019
is_ftfy_available,
2120
logging,
@@ -139,6 +138,7 @@ class IFImg2ImgPipeline(DiffusionPipeline, LoraLoaderMixin):
139138

140139
_optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor", "watermarker"]
141140
model_cpu_offload_seq = "text_encoder->unet"
141+
_exclude_from_cpu_offload = ["watermarker"]
142142

143143
def __init__(
144144
self,
@@ -180,21 +180,6 @@ def __init__(
180180
)
181181
self.register_to_config(requires_safety_checker=requires_safety_checker)
182182

183-
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.remove_all_hooks
184-
def remove_all_hooks(self):
185-
if is_accelerate_available():
186-
from accelerate.hooks import remove_hook_from_module
187-
else:
188-
raise ImportError("Please install accelerate via `pip install accelerate`")
189-
190-
for model in [self.text_encoder, self.unet, self.safety_checker]:
191-
if model is not None:
192-
remove_hook_from_module(model, recurse=True)
193-
194-
self.unet_offload_hook = None
195-
self.text_encoder_offload_hook = None
196-
self.final_offload_hook = None
197-
198183
@torch.no_grad()
199184
def encode_prompt(
200185
self,
@@ -361,9 +346,6 @@ def run_safety_checker(self, image, device, dtype):
361346
nsfw_detected = None
362347
watermark_detected = None
363348

364-
if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None:
365-
self.unet_offload_hook.offload()
366-
367349
return image, nsfw_detected, watermark_detected
368350

369351
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.prepare_extra_step_kwargs

‎src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py‎

Lines changed: 1 addition & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
from ...utils import (
1717
BACKENDS_MAPPING,
1818
PIL_INTERPOLATION,
19-
is_accelerate_available,
2019
is_bs4_available,
2120
is_ftfy_available,
2221
logging,
@@ -143,6 +142,7 @@ class IFImg2ImgSuperResolutionPipeline(DiffusionPipeline, LoraLoaderMixin):
143142

144143
_optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor"]
145144
model_cpu_offload_seq = "text_encoder->unet"
145+
_exclude_from_cpu_offload = ["watermarker"]
146146

147147
def __init__(
148148
self,
@@ -191,21 +191,6 @@ def __init__(
191191
)
192192
self.register_to_config(requires_safety_checker=requires_safety_checker)
193193

194-
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.remove_all_hooks
195-
def remove_all_hooks(self):
196-
if is_accelerate_available():
197-
from accelerate.hooks import remove_hook_from_module
198-
else:
199-
raise ImportError("Please install accelerate via `pip install accelerate`")
200-
201-
for model in [self.text_encoder, self.unet, self.safety_checker]:
202-
if model is not None:
203-
remove_hook_from_module(model, recurse=True)
204-
205-
self.unet_offload_hook = None
206-
self.text_encoder_offload_hook = None
207-
self.final_offload_hook = None
208-
209194
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing
210195
def _text_preprocessing(self, text, clean_caption=False):
211196
if clean_caption and not is_bs4_available():
@@ -513,9 +498,6 @@ def run_safety_checker(self, image, device, dtype):
513498
nsfw_detected = None
514499
watermark_detected = None
515500

516-
if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None:
517-
self.unet_offload_hook.offload()
518-
519501
return image, nsfw_detected, watermark_detected
520502

521503
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.prepare_extra_step_kwargs
@@ -1012,8 +994,6 @@ def __call__(
1012994
nsfw_detected = None
1013995
watermark_detected = None
1014996

1015-
if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None:
1016-
self.unet_offload_hook.offload()
1017997
else:
1018998
# 10. Post-processing
1019999
image = (image / 2 + 0.5).clamp(0, 1)

‎src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py‎

Lines changed: 1 addition & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
from ...utils import (
1616
BACKENDS_MAPPING,
1717
PIL_INTERPOLATION,
18-
is_accelerate_available,
1918
is_bs4_available,
2019
is_ftfy_available,
2120
logging,
@@ -142,6 +141,7 @@ class IFInpaintingPipeline(DiffusionPipeline, LoraLoaderMixin):
142141

143142
_optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor", "watermarker"]
144143
model_cpu_offload_seq = "text_encoder->unet"
144+
_exclude_from_cpu_offload = ["watermarker"]
145145

146146
def __init__(
147147
self,
@@ -183,21 +183,6 @@ def __init__(
183183
)
184184
self.register_to_config(requires_safety_checker=requires_safety_checker)
185185

186-
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.remove_all_hooks
187-
def remove_all_hooks(self):
188-
if is_accelerate_available():
189-
from accelerate.hooks import remove_hook_from_module
190-
else:
191-
raise ImportError("Please install accelerate via `pip install accelerate`")
192-
193-
for model in [self.text_encoder, self.unet, self.safety_checker]:
194-
if model is not None:
195-
remove_hook_from_module(model, recurse=True)
196-
197-
self.unet_offload_hook = None
198-
self.text_encoder_offload_hook = None
199-
self.final_offload_hook = None
200-
201186
@torch.no_grad()
202187
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt
203188
def encode_prompt(
@@ -365,9 +350,6 @@ def run_safety_checker(self, image, device, dtype):
365350
nsfw_detected = None
366351
watermark_detected = None
367352

368-
if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None:
369-
self.unet_offload_hook.offload()
370-
371353
return image, nsfw_detected, watermark_detected
372354

373355
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.prepare_extra_step_kwargs

‎src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py‎

Lines changed: 1 addition & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
from ...utils import (
1717
BACKENDS_MAPPING,
1818
PIL_INTERPOLATION,
19-
is_accelerate_available,
2019
is_bs4_available,
2120
is_ftfy_available,
2221
logging,
@@ -145,6 +144,7 @@ class IFInpaintingSuperResolutionPipeline(DiffusionPipeline, LoraLoaderMixin):
145144

146145
model_cpu_offload_seq = "text_encoder->unet"
147146
_optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor", "watermarker"]
147+
_exclude_from_cpu_offload = ["watermarker"]
148148

149149
def __init__(
150150
self,
@@ -193,21 +193,6 @@ def __init__(
193193
)
194194
self.register_to_config(requires_safety_checker=requires_safety_checker)
195195

196-
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.remove_all_hooks
197-
def remove_all_hooks(self):
198-
if is_accelerate_available():
199-
from accelerate.hooks import remove_hook_from_module
200-
else:
201-
raise ImportError("Please install accelerate via `pip install accelerate`")
202-
203-
for model in [self.text_encoder, self.unet, self.safety_checker]:
204-
if model is not None:
205-
remove_hook_from_module(model, recurse=True)
206-
207-
self.unet_offload_hook = None
208-
self.text_encoder_offload_hook = None
209-
self.final_offload_hook = None
210-
211196
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing
212197
def _text_preprocessing(self, text, clean_caption=False):
213198
if clean_caption and not is_bs4_available():
@@ -515,9 +500,6 @@ def run_safety_checker(self, image, device, dtype):
515500
nsfw_detected = None
516501
watermark_detected = None
517502

518-
if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None:
519-
self.unet_offload_hook.offload()
520-
521503
return image, nsfw_detected, watermark_detected
522504

523505
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.prepare_extra_step_kwargs

‎src/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py‎

Lines changed: 1 addition & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
from ...schedulers import DDPMScheduler
1616
from ...utils import (
1717
BACKENDS_MAPPING,
18-
is_accelerate_available,
1918
is_bs4_available,
2019
is_ftfy_available,
2120
logging,
@@ -101,6 +100,7 @@ class IFSuperResolutionPipeline(DiffusionPipeline, LoraLoaderMixin):
101100

102101
_optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor", "watermarker"]
103102
model_cpu_offload_seq = "text_encoder->unet"
103+
_exclude_from_cpu_offload = ["watermarker"]
104104

105105
def __init__(
106106
self,
@@ -149,21 +149,6 @@ def __init__(
149149
)
150150
self.register_to_config(requires_safety_checker=requires_safety_checker)
151151

152-
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.remove_all_hooks
153-
def remove_all_hooks(self):
154-
if is_accelerate_available():
155-
from accelerate.hooks import remove_hook_from_module
156-
else:
157-
raise ImportError("Please install accelerate via `pip install accelerate`")
158-
159-
for model in [self.text_encoder, self.unet, self.safety_checker]:
160-
if model is not None:
161-
remove_hook_from_module(model, recurse=True)
162-
163-
self.unet_offload_hook = None
164-
self.text_encoder_offload_hook = None
165-
self.final_offload_hook = None
166-
167152
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing
168153
def _text_preprocessing(self, text, clean_caption=False):
169154
if clean_caption and not is_bs4_available():
@@ -471,9 +456,6 @@ def run_safety_checker(self, image, device, dtype):
471456
nsfw_detected = None
472457
watermark_detected = None
473458

474-
if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None:
475-
self.unet_offload_hook.offload()
476-
477459
return image, nsfw_detected, watermark_detected
478460

479461
# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.prepare_extra_step_kwargs

‎src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,7 @@ class KandinskyCombinedPipeline(DiffusionPipeline):
143143

144144
_load_connected_pipes = True
145145
model_cpu_offload_seq = "text_encoder->unet->movq->prior_prior->prior_image_encoder->prior_text_encoder"
146+
_exclude_from_cpu_offload = ["prior_prior"]
146147

147148
def __init__(
148149
self,
@@ -360,6 +361,7 @@ class KandinskyImg2ImgCombinedPipeline(DiffusionPipeline):
360361

361362
_load_connected_pipes = True
362363
model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->prior_prior->" "text_encoder->unet->movq"
364+
_exclude_from_cpu_offload = ["prior_prior"]
363365

364366
def __init__(
365367
self,
@@ -600,6 +602,7 @@ class KandinskyInpaintCombinedPipeline(DiffusionPipeline):
600602

601603
_load_connected_pipes = True
602604
model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->prior_prior->text_encoder->unet->movq"
605+
_exclude_from_cpu_offload = ["prior_prior"]
603606

604607
def __init__(
605608
self,

‎src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,7 @@ class KandinskyV22CombinedPipeline(DiffusionPipeline):
135135

136136
model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->unet->movq"
137137
_load_connected_pipes = True
138+
_exclude_from_cpu_offload = ["prior_prior"]
138139

139140
def __init__(
140141
self,
@@ -362,6 +363,7 @@ class KandinskyV22Img2ImgCombinedPipeline(DiffusionPipeline):
362363

363364
model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->unet->movq"
364365
_load_connected_pipes = True
366+
_exclude_from_cpu_offload = ["prior_prior"]
365367

366368
def __init__(
367369
self,
@@ -610,6 +612,7 @@ class KandinskyV22InpaintCombinedPipeline(DiffusionPipeline):
610612

611613
model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->unet->movq"
612614
_load_connected_pipes = True
615+
_exclude_from_cpu_offload = ["prior_prior"]
613616

614617
def __init__(
615618
self,

‎src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py‎

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
from ...schedulers import DDPMScheduler
99
from ...utils import (
1010
deprecate,
11-
is_accelerate_available,
1211
logging,
1312
replace_example_docstring,
1413
)
@@ -72,20 +71,6 @@ def __init__(
7271
tokenizer=tokenizer, text_encoder=text_encoder, unet=unet, scheduler=scheduler, movq=movq
7372
)
7473

75-
def remove_all_hooks(self):
76-
if is_accelerate_available():
77-
from accelerate.hooks import remove_hook_from_module
78-
else:
79-
raise ImportError("Please install accelerate via `pip install accelerate`")
80-
81-
for model in [self.text_encoder, self.unet, self.movq]:
82-
if model is not None:
83-
remove_hook_from_module(model, recurse=True)
84-
85-
self.unet_offload_hook = None
86-
self.text_encoder_offload_hook = None
87-
self.final_offload_hook = None
88-
8974
def process_embeds(self, embeddings, attention_mask, cut_context):
9075
if cut_context:
9176
embeddings[attention_mask == 0] = torch.zeros_like(embeddings[attention_mask == 0])

0 commit comments

Comments
 (0)