Skip to content

Helper functions to return skip-layer compatible layers #12048

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.

Sign up for GitHub

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

Jump to bottom
Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Helper functions to return skip-layer compatible layers #12048

wants to merge 1 commit into from
+44 −0

Conversation

a-r-r-o-w
Copy link
Member 

from diffusers import AutoModel
from diffusers.hooks.utils import _get_identifiable_transformer_blocks_in_module, _get_identifiable_attention_layers_in_module, _get_identifiable_feedforward_layers_in_module

model = AutoModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet")

transformer_blocks = _get_identifiable_transformer_blocks_in_module(model)
attention_layers = _get_identifiable_attention_layers_in_module(model)
feedforward_layers = _get_identifiable_feedforward_layers_in_module(model)

print("transformer blocks")
for name, module in transformer_blocks:
    print(name)
print()
print("attention layers")
for name, module in attention_layers:
    print(name)
print()
print("feedforward layers")
for name, module in feedforward_layers:
    print(name)
output 
transformer blocks
down_blocks
down_blocks.1.attentions.0.transformer_blocks
down_blocks.1.attentions.1.transformer_blocks
down_blocks.2.attentions.0.transformer_blocks
down_blocks.2.attentions.1.transformer_blocks
up_blocks
up_blocks.0.attentions.0.transformer_blocks
up_blocks.0.attentions.1.transformer_blocks
up_blocks.0.attentions.2.transformer_blocks
up_blocks.1.attentions.0.transformer_blocks
up_blocks.1.attentions.1.transformer_blocks
up_blocks.1.attentions.2.transformer_blocks
mid_block.attentions.0.transformer_blocks

attention layers
down_blocks.1.attentions.0.transformer_blocks.0.attn1
down_blocks.1.attentions.0.transformer_blocks.0.attn2
down_blocks.1.attentions.0.transformer_blocks.1.attn1
down_blocks.1.attentions.0.transformer_blocks.1.attn2
down_blocks.1.attentions.1.transformer_blocks.0.attn1
down_blocks.1.attentions.1.transformer_blocks.0.attn2
down_blocks.1.attentions.1.transformer_blocks.1.attn1
down_blocks.1.attentions.1.transformer_blocks.1.attn2
down_blocks.2.attentions.0.transformer_blocks.0.attn1
down_blocks.2.attentions.0.transformer_blocks.0.attn2
down_blocks.2.attentions.0.transformer_blocks.1.attn1
down_blocks.2.attentions.0.transformer_blocks.1.attn2
down_blocks.2.attentions.0.transformer_blocks.2.attn1
down_blocks.2.attentions.0.transformer_blocks.2.attn2
down_blocks.2.attentions.0.transformer_blocks.3.attn1
down_blocks.2.attentions.0.transformer_blocks.3.attn2
down_blocks.2.attentions.0.transformer_blocks.4.attn1
down_blocks.2.attentions.0.transformer_blocks.4.attn2
down_blocks.2.attentions.0.transformer_blocks.5.attn1
down_blocks.2.attentions.0.transformer_blocks.5.attn2
down_blocks.2.attentions.0.transformer_blocks.6.attn1
down_blocks.2.attentions.0.transformer_blocks.6.attn2
down_blocks.2.attentions.0.transformer_blocks.7.attn1
down_blocks.2.attentions.0.transformer_blocks.7.attn2
down_blocks.2.attentions.0.transformer_blocks.8.attn1
down_blocks.2.attentions.0.transformer_blocks.8.attn2
down_blocks.2.attentions.0.transformer_blocks.9.attn1
down_blocks.2.attentions.0.transformer_blocks.9.attn2
down_blocks.2.attentions.1.transformer_blocks.0.attn1
down_blocks.2.attentions.1.transformer_blocks.0.attn2
down_blocks.2.attentions.1.transformer_blocks.1.attn1
down_blocks.2.attentions.1.transformer_blocks.1.attn2
down_blocks.2.attentions.1.transformer_blocks.2.attn1
down_blocks.2.attentions.1.transformer_blocks.2.attn2
down_blocks.2.attentions.1.transformer_blocks.3.attn1
down_blocks.2.attentions.1.transformer_blocks.3.attn2
down_blocks.2.attentions.1.transformer_blocks.4.attn1
down_blocks.2.attentions.1.transformer_blocks.4.attn2
down_blocks.2.attentions.1.transformer_blocks.5.attn1
down_blocks.2.attentions.1.transformer_blocks.5.attn2
down_blocks.2.attentions.1.transformer_blocks.6.attn1
down_blocks.2.attentions.1.transformer_blocks.6.attn2
down_blocks.2.attentions.1.transformer_blocks.7.attn1
down_blocks.2.attentions.1.transformer_blocks.7.attn2
down_blocks.2.attentions.1.transformer_blocks.8.attn1
down_blocks.2.attentions.1.transformer_blocks.8.attn2
down_blocks.2.attentions.1.transformer_blocks.9.attn1
down_blocks.2.attentions.1.transformer_blocks.9.attn2
up_blocks.0.attentions.0.transformer_blocks.0.attn1
up_blocks.0.attentions.0.transformer_blocks.0.attn2
up_blocks.0.attentions.0.transformer_blocks.1.attn1
up_blocks.0.attentions.0.transformer_blocks.1.attn2
up_blocks.0.attentions.0.transformer_blocks.2.attn1
up_blocks.0.attentions.0.transformer_blocks.2.attn2
up_blocks.0.attentions.0.transformer_blocks.3.attn1
up_blocks.0.attentions.0.transformer_blocks.3.attn2
up_blocks.0.attentions.0.transformer_blocks.4.attn1
up_blocks.0.attentions.0.transformer_blocks.4.attn2
up_blocks.0.attentions.0.transformer_blocks.5.attn1
up_blocks.0.attentions.0.transformer_blocks.5.attn2
up_blocks.0.attentions.0.transformer_blocks.6.attn1
up_blocks.0.attentions.0.transformer_blocks.6.attn2
up_blocks.0.attentions.0.transformer_blocks.7.attn1
up_blocks.0.attentions.0.transformer_blocks.7.attn2
up_blocks.0.attentions.0.transformer_blocks.8.attn1
up_blocks.0.attentions.0.transformer_blocks.8.attn2
up_blocks.0.attentions.0.transformer_blocks.9.attn1
up_blocks.0.attentions.0.transformer_blocks.9.attn2
up_blocks.0.attentions.1.transformer_blocks.0.attn1
up_blocks.0.attentions.1.transformer_blocks.0.attn2
up_blocks.0.attentions.1.transformer_blocks.1.attn1
up_blocks.0.attentions.1.transformer_blocks.1.attn2
up_blocks.0.attentions.1.transformer_blocks.2.attn1
up_blocks.0.attentions.1.transformer_blocks.2.attn2
up_blocks.0.attentions.1.transformer_blocks.3.attn1
up_blocks.0.attentions.1.transformer_blocks.3.attn2
up_blocks.0.attentions.1.transformer_blocks.4.attn1
up_blocks.0.attentions.1.transformer_blocks.4.attn2
up_blocks.0.attentions.1.transformer_blocks.5.attn1
up_blocks.0.attentions.1.transformer_blocks.5.attn2
up_blocks.0.attentions.1.transformer_blocks.6.attn1
up_blocks.0.attentions.1.transformer_blocks.6.attn2
up_blocks.0.attentions.1.transformer_blocks.7.attn1
up_blocks.0.attentions.1.transformer_blocks.7.attn2
up_blocks.0.attentions.1.transformer_blocks.8.attn1
up_blocks.0.attentions.1.transformer_blocks.8.attn2
up_blocks.0.attentions.1.transformer_blocks.9.attn1
up_blocks.0.attentions.1.transformer_blocks.9.attn2
up_blocks.0.attentions.2.transformer_blocks.0.attn1
up_blocks.0.attentions.2.transformer_blocks.0.attn2
up_blocks.0.attentions.2.transformer_blocks.1.attn1
up_blocks.0.attentions.2.transformer_blocks.1.attn2
up_blocks.0.attentions.2.transformer_blocks.2.attn1
up_blocks.0.attentions.2.transformer_blocks.2.attn2
up_blocks.0.attentions.2.transformer_blocks.3.attn1
up_blocks.0.attentions.2.transformer_blocks.3.attn2
up_blocks.0.attentions.2.transformer_blocks.4.attn1
up_blocks.0.attentions.2.transformer_blocks.4.attn2
up_blocks.0.attentions.2.transformer_blocks.5.attn1
up_blocks.0.attentions.2.transformer_blocks.5.attn2
up_blocks.0.attentions.2.transformer_blocks.6.attn1
up_blocks.0.attentions.2.transformer_blocks.6.attn2
up_blocks.0.attentions.2.transformer_blocks.7.attn1
up_blocks.0.attentions.2.transformer_blocks.7.attn2
up_blocks.0.attentions.2.transformer_blocks.8.attn1
up_blocks.0.attentions.2.transformer_blocks.8.attn2
up_blocks.0.attentions.2.transformer_blocks.9.attn1
up_blocks.0.attentions.2.transformer_blocks.9.attn2
up_blocks.1.attentions.0.transformer_blocks.0.attn1
up_blocks.1.attentions.0.transformer_blocks.0.attn2
up_blocks.1.attentions.0.transformer_blocks.1.attn1
up_blocks.1.attentions.0.transformer_blocks.1.attn2
up_blocks.1.attentions.1.transformer_blocks.0.attn1
up_blocks.1.attentions.1.transformer_blocks.0.attn2
up_blocks.1.attentions.1.transformer_blocks.1.attn1
up_blocks.1.attentions.1.transformer_blocks.1.attn2
up_blocks.1.attentions.2.transformer_blocks.0.attn1
up_blocks.1.attentions.2.transformer_blocks.0.attn2
up_blocks.1.attentions.2.transformer_blocks.1.attn1
up_blocks.1.attentions.2.transformer_blocks.1.attn2
mid_block.attentions.0.transformer_blocks.0.attn1
mid_block.attentions.0.transformer_blocks.0.attn2
mid_block.attentions.0.transformer_blocks.1.attn1
mid_block.attentions.0.transformer_blocks.1.attn2
mid_block.attentions.0.transformer_blocks.2.attn1
mid_block.attentions.0.transformer_blocks.2.attn2
mid_block.attentions.0.transformer_blocks.3.attn1
mid_block.attentions.0.transformer_blocks.3.attn2
mid_block.attentions.0.transformer_blocks.4.attn1
mid_block.attentions.0.transformer_blocks.4.attn2
mid_block.attentions.0.transformer_blocks.5.attn1
mid_block.attentions.0.transformer_blocks.5.attn2
mid_block.attentions.0.transformer_blocks.6.attn1
mid_block.attentions.0.transformer_blocks.6.attn2
mid_block.attentions.0.transformer_blocks.7.attn1
mid_block.attentions.0.transformer_blocks.7.attn2
mid_block.attentions.0.transformer_blocks.8.attn1
mid_block.attentions.0.transformer_blocks.8.attn2
mid_block.attentions.0.transformer_blocks.9.attn1
mid_block.attentions.0.transformer_blocks.9.attn2

feedforward layers
down_blocks.1.attentions.0.transformer_blocks.0.ff
down_blocks.1.attentions.0.transformer_blocks.1.ff
down_blocks.1.attentions.1.transformer_blocks.0.ff
down_blocks.1.attentions.1.transformer_blocks.1.ff
down_blocks.2.attentions.0.transformer_blocks.0.ff
down_blocks.2.attentions.0.transformer_blocks.1.ff
down_blocks.2.attentions.0.transformer_blocks.2.ff
down_blocks.2.attentions.0.transformer_blocks.3.ff
down_blocks.2.attentions.0.transformer_blocks.4.ff
down_blocks.2.attentions.0.transformer_blocks.5.ff
down_blocks.2.attentions.0.transformer_blocks.6.ff
down_blocks.2.attentions.0.transformer_blocks.7.ff
down_blocks.2.attentions.0.transformer_blocks.8.ff
down_blocks.2.attentions.0.transformer_blocks.9.ff
down_blocks.2.attentions.1.transformer_blocks.0.ff
down_blocks.2.attentions.1.transformer_blocks.1.ff
down_blocks.2.attentions.1.transformer_blocks.2.ff
down_blocks.2.attentions.1.transformer_blocks.3.ff
down_blocks.2.attentions.1.transformer_blocks.4.ff
down_blocks.2.attentions.1.transformer_blocks.5.ff
down_blocks.2.attentions.1.transformer_blocks.6.ff
down_blocks.2.attentions.1.transformer_blocks.7.ff
down_blocks.2.attentions.1.transformer_blocks.8.ff
down_blocks.2.attentions.1.transformer_blocks.9.ff
up_blocks.0.attentions.0.transformer_blocks.0.ff
up_blocks.0.attentions.0.transformer_blocks.1.ff
up_blocks.0.attentions.0.transformer_blocks.2.ff
up_blocks.0.attentions.0.transformer_blocks.3.ff
up_blocks.0.attentions.0.transformer_blocks.4.ff
up_blocks.0.attentions.0.transformer_blocks.5.ff
up_blocks.0.attentions.0.transformer_blocks.6.ff
up_blocks.0.attentions.0.transformer_blocks.7.ff
up_blocks.0.attentions.0.transformer_blocks.8.ff
up_blocks.0.attentions.0.transformer_blocks.9.ff
up_blocks.0.attentions.1.transformer_blocks.0.ff
up_blocks.0.attentions.1.transformer_blocks.1.ff
up_blocks.0.attentions.1.transformer_blocks.2.ff
up_blocks.0.attentions.1.transformer_blocks.3.ff
up_blocks.0.attentions.1.transformer_blocks.4.ff
up_blocks.0.attentions.1.transformer_blocks.5.ff
up_blocks.0.attentions.1.transformer_blocks.6.ff
up_blocks.0.attentions.1.transformer_blocks.7.ff
up_blocks.0.attentions.1.transformer_blocks.8.ff
up_blocks.0.attentions.1.transformer_blocks.9.ff
up_blocks.0.attentions.2.transformer_blocks.0.ff
up_blocks.0.attentions.2.transformer_blocks.1.ff
up_blocks.0.attentions.2.transformer_blocks.2.ff
up_blocks.0.attentions.2.transformer_blocks.3.ff
up_blocks.0.attentions.2.transformer_blocks.4.ff
up_blocks.0.attentions.2.transformer_blocks.5.ff
up_blocks.0.attentions.2.transformer_blocks.6.ff
up_blocks.0.attentions.2.transformer_blocks.7.ff
up_blocks.0.attentions.2.transformer_blocks.8.ff
up_blocks.0.attentions.2.transformer_blocks.9.ff
up_blocks.1.attentions.0.transformer_blocks.0.ff
up_blocks.1.attentions.0.transformer_blocks.1.ff
up_blocks.1.attentions.1.transformer_blocks.0.ff
up_blocks.1.attentions.1.transformer_blocks.1.ff
up_blocks.1.attentions.2.transformer_blocks.0.ff
up_blocks.1.attentions.2.transformer_blocks.1.ff
mid_block.attentions.0.transformer_blocks.0.ff
mid_block.attentions.0.transformer_blocks.1.ff
mid_block.attentions.0.transformer_blocks.2.ff
mid_block.attentions.0.transformer_blocks.3.ff
mid_block.attentions.0.transformer_blocks.4.ff
mid_block.attentions.0.transformer_blocks.5.ff
mid_block.attentions.0.transformer_blocks.6.ff
mid_block.attentions.0.transformer_blocks.7.ff
mid_block.attentions.0.transformer_blocks.8.ff
mid_block.attentions.0.transformer_blocks.9.ff

@a-r-r-o-w a-r-r-o-w requested review from asomoza and yiyixuxu August 3, 2025 09:03
@a-r-r-o-w a-r-r-o-w marked this pull request as ready for review August 3, 2025 09:03
@HuggingFaceDocBuilderDev
Copy link

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.

yiyixuxu
yiyixuxu approved these changes Aug 3, 2025
View reviewed changes
Copy link
Collaborator

@yiyixuxu yiyixuxu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks!
looks good to me
@asomoza let me know if this is what you need

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@yiyixuxu yiyixuxu yiyixuxu approved these changes

@asomoza asomoza Awaiting requested review from asomoza

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@a-r-r-o-w @HuggingFaceDocBuilderDev @yiyixuxu