Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add TPU/XLA devices support for ComfyUI #5657

Open
wants to merge 1 commit into
base: master
Choose a base branch
from

Conversation

radna0
Copy link

@radna0 radna0 commented Nov 17, 2024

This PR adds support for TPU/XLA devices in ComfyUI. users can enable by adding the command line arg --xla

  • Update README about support for TPU/XLA devices.
  • Detects TPU/XLA device, and memory information.
  • Add torch_xla.sync()/xm.mark_step() for the callback after every sampling step when using TPU devices. This is needed to execute the XLA current graph. *Note: Torch XLA sync is non blocking by default.

Tested on TPU v3 with ComfyUI-MochiWrapper models.

a reference image with sampling settings/seed/etc.
image

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant