Added support to create asymmetrical U-Net structures #5400

Gothos · 2023-10-15T16:47:48Z

What does this PR do?

Adds UNetMidBlock2D as a supported mid_block_type in the UNet2DConditionModel class.
Adds basic support for specifying transformer_layers_per_block for every layer of the UNet2DConditionModel class.
Support is somewhat rigid at the moment, using this functionality requires specifying the reverse_transformer_layers_per_block attribute in a config (set to None by default).
Please comment on any style changes/bugs you find in the code, I'll fix it as soon as I can.

…-block init

sayakpaul

Thanks for working on this. The changes seem clean and simple to me.
@patrickvonplaten WDYT?

Updated the docstring for ```UNet2DConditionModel``` to include ```reverse_transformer_layers_per_block``` and updated checking for nested list type ```transformer_layers_per_block```

patrickvonplaten

Hmm I would need to see a clear use case for when asymmetrical U-Nets make sense. I'm a bit hesitant to adapt the Unet2DCondition class for use case that have not been verified to work yet.

Do we have a checkpoint that shows good results with an asymmetrical U-Net?

Gothos · 2023-10-16T13:51:40Z

The checkpoint we're using for the SDXL distillation is asymmetrical. We did put those results up on the slack channel.
The PR shouldn't affect anything existing at the current stage significantly, I'll be more than happy to do any needful changes quick.

patrickvonplaten

Cool, yes this design also works for me. Could we add a quick test here? :-)

Gothos · 2023-10-16T14:08:13Z

No problem! May I ask what exactly you have in mind? Instantiating 1-2 UNet from configs of the new type? Can I get a few examples for tests written for similar PRs?

sayakpaul · 2023-10-16T14:10:30Z

Hi @Gothos!

You can check out this test suite to get an idea:
https://github.com/huggingface/diffusers/blob/main/tests/models/test_models_unet_2d_condition.py

I'd instantiate an asymmetric UNet in that test suite and run dummy inference and perform assertions as you'd see for some of the tests being done in the test suite mentioned above.

Let me know if that makes sense?

Gothos · 2023-10-16T14:11:28Z

Yes, thanks! I'll get this done as quickly as I can.

Gothos · 2023-10-17T09:19:12Z

@sayakpaul I've written a basic test asserting input shape equals the output shape for an asymmetrical UNet. Please let me know if anything's wrong/ any other tests are required.

src/diffusers/models/unet_2d_blocks.py

src/diffusers/models/unet_2d_condition.py

Gothos · 2023-10-18T09:06:19Z

@sayakpaul I've done the needful.

Gothos · 2023-10-18T10:18:47Z

Running the commands above to enforce consistency seems to have attempted to make "UNetMidBlockFlat" a viable mid block type under the versatile diffusion pipeline in modeling_text_unet.py. However this class is not implemented in the file and this causes an error.

sayakpaul · 2023-10-18T10:27:51Z

We need to enforce repository consistency, sadly.

Is it possible to copy over the implementation in modeling_text_unet.py?

…ensure consistency

Gothos · 2023-10-18T10:39:53Z

We need to enforce repository consistency, sadly.

Is it possible to copy over the implementation in modeling_text_unet.py?

Sure! I've just done it and rerun the makes twice to make sure nothing popped up.

sayakpaul

Good work!

patrickvonplaten · 2023-10-18T20:31:27Z

src/diffusers/models/unet_2d_blocks.py

@@ -587,7 +587,7 @@ def __init__(
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

-    def forward(self, hidden_states, temb=None):
+    def forward(self, hidden_states, temb=None, *args, **kwargs):


Suggested change

def forward(self, hidden_states, temb=None, *args, **kwargs):

def forward(self, hidden_states, temb=None):

Is this really needed?

Can we remove this, I don't see exactly how this is required here

``` # 4. mid if self.mid_block is not None: sample = self.mid_block( sample, emb, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask, cross_attention_kwargs=cross_attention_kwargs, encoder_attention_mask=encoder_attention_mask, )```

These lines make the *args, **kwargs necessary. I'll go by the style elsewhere and use the has attention attribute instead

if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention: sample = self.mid_block( sample, emb, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask, cross_attention_kwargs=cross_attention_kwargs, encoder_attention_mask=encoder_attention_mask, ) else: sample = self.mid_block( sample, emb, )

Would following your suggestion above and then doing this work, @patrickvonplaten?

This sounds good to me!

All right, I'll commit the changes. Just wanted to check once.

tests/models/test_models_unet_2d_condition.py

sayakpaul · 2023-10-20T04:34:57Z

@patrickvonplaten could you give this a look too?

HuggingFaceDocBuilderDev · 2023-10-20T04:42:46Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

src/diffusers/models/unet_2d_condition.py

sayakpaul

Thanks a lot for working on this! PR is good to go from my end.

Wrap into single line Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* Added args, kwargs to ```U * Add UNetMidBlock2D as a supported mid block type * Fix extra init input for UNetMidBlock2D, change allowed types for Mid-block init * Update unet_2d_condition.py * Update unet_2d_condition.py * Update unet_2d_condition.py * Update unet_2d_condition.py * Update unet_2d_condition.py * Update unet_2d_condition.py * Update unet_2d_condition.py * Update unet_2d_condition.py * Update unet_2d_blocks.py * Update unet_2d_blocks.py * Update unet_2d_blocks.py * Update unet_2d_condition.py * Update unet_2d_blocks.py * Updated docstring, increased check strictness Updated the docstring for ```UNet2DConditionModel``` to include ```reverse_transformer_layers_per_block``` and updated checking for nested list type ```transformer_layers_per_block``` * Add basic shape-check test for asymmetrical unets * Update src/diffusers/models/unet_2d_blocks.py Removed blank line Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Update unet_2d_condition.py Remove blank space * Update unet_2d_condition.py Changed docstring for `mid_block_type` * Fixed docstring and wrong default value * Reformat with black * Reformat with necessary commands * Add UNetMidBlockFlat to versatile_diffusion/modeling_text_unet.py to ensure consistency * Removed args, kwargs, use on mid-block type * Make fix-copies * Update src/diffusers/models/unet_2d_condition.py Wrap into single line Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * make fix-copies --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

Gothos added 16 commits October 15, 2023 11:34

Added args, kwargs to ```U

b8dca9f

Add UNetMidBlock2D as a supported mid block type

577d1e2

Fix extra init input for UNetMidBlock2D, change allowed types for Mid…

2c8d804

…-block init

Update unet_2d_condition.py

e07c615

Update unet_2d_condition.py

8344b2d

Update unet_2d_condition.py

025b1ec

Update unet_2d_condition.py

6230214

Update unet_2d_condition.py

80b891b

Update unet_2d_condition.py

c4e4d40

Update unet_2d_condition.py

6176fd5

Update unet_2d_condition.py

73737f8

Update unet_2d_blocks.py

bc53a32

Update unet_2d_blocks.py

314730b

Update unet_2d_blocks.py

545a0b0

Update unet_2d_condition.py

bf77268

Update unet_2d_blocks.py

1998c17

sayakpaul requested a review from patrickvonplaten October 16, 2023 06:20

sayakpaul reviewed Oct 16, 2023

View reviewed changes

Updated docstring, increased check strictness

81d682e

Updated the docstring for ```UNet2DConditionModel``` to include ```reverse_transformer_layers_per_block``` and updated checking for nested list type ```transformer_layers_per_block```

patrickvonplaten reviewed Oct 16, 2023

View reviewed changes

patrickvonplaten approved these changes Oct 16, 2023

View reviewed changes

Add basic shape-check test for asymmetrical unets

93ce7d3

sayakpaul reviewed Oct 17, 2023

View reviewed changes

src/diffusers/models/unet_2d_blocks.py Outdated Show resolved Hide resolved

sayakpaul reviewed Oct 17, 2023

View reviewed changes

src/diffusers/models/unet_2d_condition.py Outdated Show resolved Hide resolved

sayakpaul reviewed Oct 17, 2023

View reviewed changes

src/diffusers/models/unet_2d_condition.py Outdated Show resolved Hide resolved

Reformat with necessary commands

196ab3e

sayakpaul requested a review from patrickvonplaten October 18, 2023 09:13

Merge branch 'huggingface:main' into main

a65ec27

Add UNetMidBlockFlat to versatile_diffusion/modeling_text_unet.py to …

bbebc23

…ensure consistency

sayakpaul approved these changes Oct 18, 2023

View reviewed changes

patrickvonplaten reviewed Oct 18, 2023

View reviewed changes

tests/models/test_models_unet_2d_condition.py Show resolved Hide resolved

sayakpaul and others added 4 commits October 19, 2023 10:38

Merge branch 'main' into main

b6c0ff6

Removed args, kwargs, use on mid-block type

0a73108

Merge remote-tracking branch 'refs/remotes/origin/main'

950d969

Make fix-copies

bfbf85c

sayakpaul requested a review from patrickvonplaten October 20, 2023 04:34

sayakpaul reviewed Oct 20, 2023

View reviewed changes

src/diffusers/models/unet_2d_condition.py Outdated Show resolved Hide resolved

sayakpaul approved these changes Oct 20, 2023

View reviewed changes

Gothos and others added 2 commits October 20, 2023 12:22

Update src/diffusers/models/unet_2d_condition.py

8389a85

Wrap into single line Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

make fix-copies

252ca85

patrickvonplaten merged commit 8dba180 into huggingface:main Oct 20, 2023
11 checks passed

Gothos mentioned this pull request Oct 25, 2023

Distilled SDXL from segmind doesn't work comfyanonymous/ComfyUI#1827

Closed

kaeru-shigure mentioned this pull request Oct 25, 2023

Support for SSD-1B kohya-ss/sd-scripts#897

Open

This was referenced Oct 26, 2023

Draft PR for SSD loading comfyanonymous/ComfyUI#1848

Closed

[Feature Request]: Support for SSD 1-B and other asymmetrical U-Nets AUTOMATIC1111/stable-diffusion-webui#13787

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Added support to create asymmetrical U-Net structures #5400

Added support to create asymmetrical U-Net structures #5400

Gothos commented Oct 15, 2023

sayakpaul left a comment

patrickvonplaten left a comment •

edited

Loading

Gothos commented Oct 16, 2023 •

edited

Loading

patrickvonplaten left a comment

Gothos commented Oct 16, 2023

sayakpaul commented Oct 16, 2023

Gothos commented Oct 16, 2023 •

edited

Loading

Gothos commented Oct 17, 2023

Gothos commented Oct 18, 2023

Gothos commented Oct 18, 2023

sayakpaul commented Oct 18, 2023

Gothos commented Oct 18, 2023

sayakpaul left a comment

patrickvonplaten Oct 18, 2023

patrickvonplaten Oct 18, 2023

Gothos Oct 19, 2023 •

edited

Loading

Gothos Oct 19, 2023 •

edited

Loading

sayakpaul Oct 19, 2023

Gothos Oct 19, 2023

sayakpaul commented Oct 20, 2023

HuggingFaceDocBuilderDev commented Oct 20, 2023

sayakpaul left a comment

	def forward(self, hidden_states, temb=None, args, *kwargs):
	def forward(self, hidden_states, temb=None):

Added support to create asymmetrical U-Net structures #5400

Added support to create asymmetrical U-Net structures #5400

Conversation

Gothos commented Oct 15, 2023

sayakpaul left a comment

Choose a reason for hiding this comment

patrickvonplaten left a comment • edited Loading

Choose a reason for hiding this comment

Gothos commented Oct 16, 2023 • edited Loading

patrickvonplaten left a comment

Choose a reason for hiding this comment

Gothos commented Oct 16, 2023

sayakpaul commented Oct 16, 2023

Gothos commented Oct 16, 2023 • edited Loading

Gothos commented Oct 17, 2023

Gothos commented Oct 18, 2023

Gothos commented Oct 18, 2023

sayakpaul commented Oct 18, 2023

Gothos commented Oct 18, 2023

sayakpaul left a comment

Choose a reason for hiding this comment

patrickvonplaten Oct 18, 2023

Choose a reason for hiding this comment

patrickvonplaten Oct 18, 2023

Choose a reason for hiding this comment

Gothos Oct 19, 2023 • edited Loading

Choose a reason for hiding this comment

Gothos Oct 19, 2023 • edited Loading

Choose a reason for hiding this comment

sayakpaul Oct 19, 2023

Choose a reason for hiding this comment

Gothos Oct 19, 2023

Choose a reason for hiding this comment

sayakpaul commented Oct 20, 2023

HuggingFaceDocBuilderDev commented Oct 20, 2023

sayakpaul left a comment

Choose a reason for hiding this comment

patrickvonplaten left a comment •

edited

Loading

Gothos commented Oct 16, 2023 •

edited

Loading

Gothos commented Oct 16, 2023 •

edited

Loading

Gothos Oct 19, 2023 •

edited

Loading

Gothos Oct 19, 2023 •

edited

Loading