ComfyUI icon indicating copy to clipboard operation
ComfyUI copied to clipboard

Add TPU/XLA devices support for ComfyUI

Open radna0 opened this issue 1 year ago • 0 comments

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

radna0 avatar Nov 17 '24 23:11 radna0