ComfyUI
ComfyUI copied to clipboard
Add TPU/XLA devices support for ComfyUI
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.