[tests] help us test torch.compile() for impactful models

[sayakpaul]

#11085 added a test for checking if there's any graph break or recompilation issue for torch.compiled model.

We should add this test to the most impactful models to ensure our code is torch.compile friendly and has the potential to benefit from it. So far, we test it for FluxTransformer2DModel. Below are some models I have in mind where this test should be added:

• HiDreamImageTransformer2DModel (currently doesn't have a test class like other models from src/diffusers/models)
• HunyuanVideoTransformer3DTests
• WanTransformer3DTests
• UNet2DConditionModelTests
• LTXVideoTransformer3DModel

Steps to contribute

1. Refer to [tests] add tests to check for graph breaks, recompilation, cuda syncs in pipelines during torch.compile() #11085 to under the changes needed in the test classes of the respective models.
2. To be specific, always attempt to just add TorchCompileTesterMixin to the respective model testing file. Like:

   diffusers/tests/models/transformers/test_models_transformer_flux.py

   Line 81 in fb29132

   class FluxTransformerTests(ModelTesterMixin, TorchCompileTesterMixin, unittest.TestCase):

3. Make the changes.
4. Run the tests locally with RUN_SLOW=1 RUN_COMPILE=1 pytest tests/<PATH_TO_TEST_FILE> -k "test_torch_compile_recompilation_and_graph_break". Make sure the machine has a GPU. Please confirm that this test is passing when opening the PR. In case it doesn't pass let us know.
5. Open the PR and mention this issue ([tests] help us test torch.compile() for impactful models #11430). Tag @sayakpaul and @DN6 for a review. Don't hesitate to ask for guidance/help if needed.

@DN6 any model classes I am missing?

[tongyu0924]

I would like to contribute by adding tests for HunyuanVideoTransformer3DTests

[olccihyeon]

@sayakpaul

I would like to contribute by adding tests for UNet2DConditionModelTests

[sayakpaul]

Let's go!

[cjfghk5697]

Hi @sayakpaul ,
Are there any other models that should also be covered for the torch.compile() tests?

[sayakpaul]

We're open to suggestions :)

[cjfghk5697]

@sayakpaul
Would it be okay if I work on adding the torch.compile test for the following models?

• DiTTransformer2DModel
• LatteTransformer3DModel

Let me know if they're suitable for this task or if you’d recommend focusing on other models.

[sayakpaul]

Thanks for being willing. Those two models are great but we want to be a bit mindful about these tests in the sense that we are aiming to initially only target the most impactful models. Sadly, those two haven't had much of an impact.

So, I would suggest:

diffusers/src/diffusers/models/transformers/transformer_ltx.py

Line 302 in fb29132

class LTXVideoTransformer3DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin, CacheMixin):

[cjfghk5697]

@sayakpaul
Thanks for the review PR has been submitted!

[olccihyeon]

@sayakpaul

I would appreciate it if you could take a look at the PR I submitted.