Avoiding graph break by changing the way we infer dtype in vae.decoder

[ppadjinTT]

…sors

What does this PR do?

This PR addresses the problem disscused in #12501, where the usage of upscale_dtype = next(iter(self.up_blocks.parameters())).dtype to infer the dtype in the forward pass of the vae.decoder causes the graph break when compiling the model with torch.compile.

The issue is that the usage of next(iter(...)) forces the lazy tensors in the initial compiled model pass to materialize, resulting in graph break, which decreases performance.

This PR proposes a simple fix by infering the dtype as:

upscale_dtype = self.conv_out.weight.dtype

Fixes #12501

Who can review?

@sayakpaul

[sayakpaul]

@DN6 WDYT?

[HuggingFaceDocBuilderDev]

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.

[ppadjinTT]

I made sure all autoencoder tests are passing locally, I would be very thankful if you can take a look @DN6

[DN6]

Think current failing tests in the CI are due to the fact that not every decoder block has a norm1 with a weight. Hence the use of the generator here to avoid such cases.

@ppadjinTT I noticed you initially used self.conv_out.weight here? What was the issue you ran into with that?

[ppadjinTT]

okay, I will change that too, tnx! I intially changed the self.conv_out.weight because there are some tests that check what happens when conv_out and upscale_blocks have different dtypes

[DN6]

Could you point me to those tests? Seems like setting to conv_out is more robust.

[ppadjinTT]

Yup, these are the tests pytest -svvv tests/models/autoencoders/test_models_autoencoder_kl.py

This is one of the tests from this test set that fails tests/models/autoencoders/test_models_autoencoder_kl.py::AutoencoderKLTests::test_layerwise_casting_inference

[ppadjinTT]

I added better logic for inferring dtype, to capture the case where it doesn't work

[DN6]

Hmm I think we can remove upscale_type entirely here. I think all tests should still pass without it.

[ppadjinTT]

okay let's try that, i'm pushing the change

[ppadjinTT]

Do you think it's okay now? @DN6

[ppadjinTT]

Any chance you can take a look? Thanks for your effort @DN6

[DN6]

Thanks @ppadjinTT 👍🏽