[refactor] make set_attention_slice recursive #1532
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
The documentation is not available anymore as the PR was closed or merged. |
patil-suraj
changed the title
| set_requires_grad(self.text_encoder, False) | ||
| set_requires_grad(self.clip_model, False) | ||
|
|
||
| def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"): |
Contributor
There was a problem hiding this comment.
Let's not change community pipelines. This will remove functionality at the moment.
src/diffusers/models/attention.py
Outdated
| ) | ||
|
|
||
| def _set_attention_slice(self, slice_size): | ||
| def set_attention_slice(self, slice_size): |
Contributor
There was a problem hiding this comment.
Nice! works for me
patrickvonplaten
approved these changes
Dec 5, 2022
Contributor
patrickvonplaten
left a comment
patrickvonplaten
left a comment
There was a problem hiding this comment.
Thanks for doing this PR! Let's just remove the changes from the community pipelines.
pcuenca
approved these changes
Dec 5, 2022
Member
pcuenca
left a comment
pcuenca
left a comment
There was a problem hiding this comment.
Very nice! Same comment as Patrick, this would currently break community pipelines as they are loaded from main.
patil-suraj
commented
Dec 5, 2022
Contributor
Author
patil-suraj
left a comment
patil-suraj
left a comment
There was a problem hiding this comment.
Looks good to me! Just left one comment
Comment on lines
+260
to
+267
| if slice_size == "auto": | ||
| # half the attention head size is usually a good trade-off between | ||
| # speed and memory | ||
| slice_size = [dim // 2 for dim in sliceable_head_dims] | ||
| elif slice_size == "max": | ||
| # make smallest slice possible | ||
| slice_size = num_slicable_layers * [1] | ||
|
|
Contributor
Author
There was a problem hiding this comment.
As discussed offline, this logic could be moved to the attention block itself. That way if would be easy to support this functionality in new models which use the attention blocks.
Contributor
There was a problem hiding this comment.
That makes sense. As discussed offline, I'll leave as is now though so that we can test that number of passed slices equals number of slicable layers. Otherwise we cannot really test it.
Comment on lines
+338
to
+346
| if slice_size == "auto": | ||
| # half the attention head size is usually a good trade-off between | ||
| # speed and memory | ||
| slice_size = [dim // 2 for dim in sliceable_head_dims] | ||
| elif slice_size == "max": | ||
| # make smallest slice possible | ||
| slice_size = num_slicable_layers * [1] | ||
|
|
||
| slice_size = num_slicable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size |
Contributor
Author
There was a problem hiding this comment.
same comment as above.
Comment on lines
+337
to
+374
| def test_model_attention_slicing(self): | ||
| init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() | ||
|
|
||
| init_dict["attention_head_dim"] = (8, 16) | ||
|
|
||
| model = self.model_class(**init_dict) | ||
| model.to(torch_device) | ||
| model.eval() | ||
|
|
||
| model.set_attention_slice("auto") | ||
| with torch.no_grad(): | ||
| output = model(**inputs_dict) | ||
| assert output is not None | ||
|
|
||
| model.set_attention_slice("max") | ||
| with torch.no_grad(): | ||
| output = model(**inputs_dict) | ||
| assert output is not None | ||
|
|
||
| model.set_attention_slice(2) | ||
| with torch.no_grad(): | ||
| output = model(**inputs_dict) | ||
| assert output is not None | ||
|
|
||
| def test_model_slicable_head_dim(self): | ||
| init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() | ||
|
|
||
| init_dict["attention_head_dim"] = (8, 16) | ||
|
|
||
| model = self.model_class(**init_dict) | ||
|
|
||
| def check_slicable_dim_attr(module: torch.nn.Module): | ||
| if hasattr(module, "set_attention_slice"): | ||
| assert isinstance(module.sliceable_head_dim, int) | ||
|
|
||
| for child in module.children(): | ||
| check_slicable_dim_attr(child) | ||
|
|
Contributor
Author
There was a problem hiding this comment.
Thanks a lot for adding the tests!
tcapelle
pushed a commit
to tcapelle/diffusers
that referenced
this pull request
Dec 12, 2022
* make attn slice recursive * remove set_attention_slice from blocks * fix copies * make enable_attention_slicing base class method of DiffusionPipeline * fix set_attention_slice * fix set_attention_slice * fix copies * add tests * up * up * up * update * up * uP Co-authored-by: Patrick von Platen <[email protected]>
sliard
pushed a commit
to sliard/diffusers
that referenced
this pull request
Dec 21, 2022
* make attn slice recursive * remove set_attention_slice from blocks * fix copies * make enable_attention_slicing base class method of DiffusionPipeline * fix set_attention_slice * fix set_attention_slice * fix copies * add tests * up * up * up * update * up * uP Co-authored-by: Patrick von Platen <[email protected]>
yoonseokjin
pushed a commit
to yoonseokjin/diffusers
that referenced
this pull request
Dec 25, 2023
* make attn slice recursive * remove set_attention_slice from blocks * fix copies * make enable_attention_slicing base class method of DiffusionPipeline * fix set_attention_slice * fix set_attention_slice * fix copies * add tests * up * up * up * update * up * uP Co-authored-by: Patrick von Platen <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Inspired by #1493, this PR makes
set_attention_slicemethod recursive to be able to easily apply it to various attention blocks without much boilerplate.set_attention_slicea method ofDiffusionPipeline, so it can be applied to all pipelines, since all of them have attention.autoslice logic insideUNet2DConditionModeland recurse there to apply the slice to all blocks.The logic differs a bit from #1493, in that, the we keep the
set_attention_slicemethod on main model classes likeUNet2DConditionModeland move the recursive application logic there. This is so that we can control this functionality using individual model classes.