-
Notifications
You must be signed in to change notification settings - Fork 5.3k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Fix cpu offload related slow tests #7618
Conversation
@@ -156,20 +156,6 @@ def __init__( | |||
) | |||
self.register_to_config(requires_safety_checker=requires_safety_checker) | |||
|
|||
def remove_all_hooks(self): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
we have a remove_all_hooks
method for all DiffusionPipeline
now
@@ -335,9 +321,6 @@ def run_safety_checker(self, image, device, dtype): | |||
nsfw_detected = None | |||
watermark_detected = None | |||
|
|||
if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None: | |||
self.unet_offload_hook.offload() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't think we need to manually offload anything anymore; I don't think unry_offload_hook
on pipeline is ever used
@@ -139,6 +138,7 @@ class IFImg2ImgPipeline(DiffusionPipeline, LoraLoaderMixin): | |||
|
|||
_optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor", "watermarker"] | |||
model_cpu_offload_seq = "text_encoder->unet" | |||
_exclude_from_cpu_offload = ["watermarker"] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
whatermarker
does not work with the sequential_cpu_offload hooks, since they are very small, I'm just excluding it
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nicely done 👍🏽
This PR cleans up the cpu_offloading code for kandinsky & IF
This should fix all cpu offload tests failure related to IF and Kandisnky, e.g https://github.com/huggingface/diffusers/actions/runs/8594679447/job/23548137528#step:9:17
https://github.com/huggingface/diffusers/actions/runs/8594679447/job/23548137612#step:9:60
https://github.com/huggingface/diffusers/actions/runs/8594679447/job/23548137335#step:9:17
I still have a few tests related to
from_pipe
I'm fixing here #7564 - but they are edge cases so don't have to worry about too much for now