From ce1063acfa0cbc2168a7e9dddd4282ab8013b810 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 16 Apr 2025 11:12:09 +0530 Subject: [PATCH] [docs] add a snippet for compilation in the auraflow docs. (#11327) * add a snippet for compilation in the auraflow docs. * include speedups. --- docs/source/en/api/pipelines/aura_flow.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/source/en/api/pipelines/aura_flow.md b/docs/source/en/api/pipelines/aura_flow.md index 5d58690505..50a67bed29 100644 --- a/docs/source/en/api/pipelines/aura_flow.md +++ b/docs/source/en/api/pipelines/aura_flow.md @@ -89,6 +89,21 @@ image = pipeline(prompt).images[0] image.save("auraflow.png") ``` +## Support for `torch.compile()` + +AuraFlow can be compiled with `torch.compile()` to speed up inference latency even for different resolutions. First, install PyTorch nightly following the instructions from [here](https://pytorch.org/). The snippet below shows the changes needed to enable this: + +```diff ++ torch.fx.experimental._config.use_duck_shape = False ++ pipeline.transformer = torch.compile( + pipeline.transformer, fullgraph=True, dynamic=True +) +``` + +This enables from 100% (on low resolutions) to a 30% (on 1536x1536 resolution) speed improvements. + +Thanks to [AstraliteHeart](https://github.com/huggingface/diffusers/pull/11297/) who helped us rewrite the [`AuraFlowTransformer2DModel`] class so that the above works for different resolutions ([PR](https://github.com/huggingface/diffusers/pull/11297/)). + ## AuraFlowPipeline [[autodoc]] AuraFlowPipeline