Skip to content

Commit ce1063a

Browse files
authored
[docs] add a snippet for compilation in the auraflow docs. (#11327)
* add a snippet for compilation in the auraflow docs. * include speedups.
1 parent 7212f35 commit ce1063a

File tree

1 file changed

+15
-0
lines changed

1 file changed

+15
-0
lines changed

docs/source/en/api/pipelines/aura_flow.md

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,21 @@ image = pipeline(prompt).images[0]
8989
image.save("auraflow.png")
9090
```
9191

92+
## Support for `torch.compile()`
93+
94+
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:
95+
96+
```diff
97+
+ torch.fx.experimental._config.use_duck_shape = False
98+
+ pipeline.transformer = torch.compile(
99+
pipeline.transformer, fullgraph=True, dynamic=True
100+
)
101+
```
102+
103+
This enables from 100% (on low resolutions) to a 30% (on 1536x1536 resolution) speed improvements.
104+
105+
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/)).
106+
92107
## AuraFlowPipeline
93108

94109
[[autodoc]] AuraFlowPipeline

0 commit comments

Comments
 (0)