Adding PyTorch XLA support for sdxl inference #5273
Conversation
… with instructions for xla
ssusie
changed the title
patrickvonplaten
added
community-examples
contributions-welcome
and removed
community-examples
contributions-welcome
labels
|
Hey @ssusie, Nice addition, generally I'm ok with having this addition. However, could you maybe add |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. |
|
Thanks for the feedback Patrick. Added the dependency in |
|
@ssusie could you run |
|
Thanks everyone for the review and comments. I ran the make style and added changes to the pr. |
| speedup, we need to call the pipe again on the input with the same length | ||
| as the original prompt to reuse the optimized graph and get the performance |
There was a problem hiding this comment.
Does the prompt length really matters here, as embeddings always have the same shape?
There was a problem hiding this comment.
Thanks for the great question. With the current setup we need the same length for the input for XLA, otherwise it will recompile the whole graph. We can potentially cut the graph with xm.mark_step after embeddings are calculated, but it is not address in this PR.
src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py
Show resolved
Hide resolved
|
Thanks for the reviews and comments. Is there anything else to change or address or can this be merged? |
|
Let's merge it - great job @ssusie! We don't have tests here yet, but I think this is fine to begin with :-) |
* Added mark_step for sdxl to run with pytorch xla. Also updated README with instructions for xla * adding soft dependency on torch_xla * fix some styling --------- Co-authored-by: Sayak Paul <[email protected]> Co-authored-by: Dhruv Nair <[email protected]>
* Added mark_step for sdxl to run with pytorch xla. Also updated README with instructions for xla * adding soft dependency on torch_xla * fix some styling --------- Co-authored-by: Sayak Paul <[email protected]> Co-authored-by: Dhruv Nair <[email protected]>
What does this PR do?
Enables users run sdxl inference in PyTorch XLA. README_sdxl.md guide is also updated.
@patrickvonplaten @sayakpaul
Thanks.