xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-11-14 01:46:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add XLA/TPU support with sync callback for diffusion models

Browse Source
This commit is contained in:
radna0 2024-11-17 13:29:07 +00:00
parent 4ac401af2b
commit 812ceb49e7
4 changed files with 493 additions and 117 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
README.md
View File

@ -200,6 +200,25 @@ You can install ComfyUI in Apple Mac silicon (M1 or M2) with any recent macOS ve
```pip install torch-directml``` Then you can launch ComfyUI with: ```python main.py --directml```
#### TPU/XLA Devices
Users with TPU/XLA devices can install the PyTorch XLA stable build with the following command:
```pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html```
This is the command to install the nightly 2.6.0 which might have some performance improvements:
```
pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev-cp310-cp310-linux_x86_64.whl' -f https://storage.googleapis.com/libtpu-releases/index.html
```
<!-- tell them to install tpu-info if memory info is needed -->
To get memory info for TPU devices, install the [tpu-info](https://github.com/AI-Hypercomputer/cloud-accelerator-diagnostics/tree/main/tpu_info) package with the following command:
```pip install tpu-info```
# Running
```python main.py```

3
comfy/cli_args.py
View File

@ -138,6 +138,9 @@ parser.add_argument("--multi-user", action="store_true", help="Enables per-user
parser.add_argument("--verbose", default='INFO', const='DEBUG', nargs="?", choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], help='Set the logging level')
parser.add_argument("--xla", action="store_true", help="To use the XLA devices for everything.")
# parser.add_argument("--xla_fsdp", action="store_true", help="To use the XLA devices with FSDP for everything.")
# The default built-in provider hosted under web/
DEFAULT_VERSION_STRING = "comfyanonymous/ComfyUI@latest"

505
comfy/model_management.py
View File

File diff suppressed because it is too large Load Diff

83
latent_preview.py
View File

@ -11,12 +11,22 @@ import logging
MAX_PREVIEW_RESOLUTION = args.preview_size
def preview_to_image(latent_image):
latents_ubyte = (((latent_image + 1.0) / 2.0).clamp(0, 1) # change scale from -1..1 to 0..1
.mul(0xFF) # to 0..255
).to(device="cpu", dtype=torch.uint8, non_blocking=comfy.model_management.device_supports_non_blocking(latent_image.device))
return Image.fromarray(latents_ubyte.numpy())
def preview_to_image(latent_image):
latents_ubyte = (
((latent_image + 1.0) / 2.0)
.clamp(0, 1) # change scale from -1..1 to 0..1
.mul(0xFF) # to 0..255
).to(
device="cpu",
dtype=torch.uint8,
non_blocking=comfy.model_management.device_supports_non_blocking(
latent_image.device
),
)
return Image.fromarray(latents_ubyte.numpy())
class LatentPreviewer:
def decode_latent_to_preview(self, x0):
@ -26,6 +36,7 @@ class LatentPreviewer:
preview_image = self.decode_latent_to_preview(x0)
return ("JPEG", preview_image, MAX_PREVIEW_RESOLUTION)
class TAESDPreviewerImpl(LatentPreviewer):
def __init__(self, taesd):
self.taesd = taesd
@ -37,22 +48,34 @@ class TAESDPreviewerImpl(LatentPreviewer):
class Latent2RGBPreviewer(LatentPreviewer):
def __init__(self, latent_rgb_factors, latent_rgb_factors_bias=None):
self.latent_rgb_factors = torch.tensor(latent_rgb_factors, device="cpu").transpose(0, 1)
self.latent_rgb_factors = torch.tensor(
latent_rgb_factors, device="cpu"
).transpose(0, 1)
self.latent_rgb_factors_bias = None
if latent_rgb_factors_bias is not None:
self.latent_rgb_factors_bias = torch.tensor(latent_rgb_factors_bias, device="cpu")
self.latent_rgb_factors_bias = torch.tensor(
latent_rgb_factors_bias, device="cpu"
)
def decode_latent_to_preview(self, x0):
self.latent_rgb_factors = self.latent_rgb_factors.to(dtype=x0.dtype, device=x0.device)
self.latent_rgb_factors = self.latent_rgb_factors.to(
dtype=x0.dtype, device=x0.device
)
if self.latent_rgb_factors_bias is not None:
self.latent_rgb_factors_bias = self.latent_rgb_factors_bias.to(dtype=x0.dtype, device=x0.device)
self.latent_rgb_factors_bias = self.latent_rgb_factors_bias.to(
dtype=x0.dtype, device=x0.device
)
if x0.ndim == 5:
x0 = x0[0, :, 0]
else:
x0 = x0[0]
latent_image = torch.nn.functional.linear(x0.movedim(0, -1), self.latent_rgb_factors, bias=self.latent_rgb_factors_bias)
latent_image = torch.nn.functional.linear(
x0.movedim(0, -1),
self.latent_rgb_factors,
bias=self.latent_rgb_factors_bias,
)
# latent_image = x0[0].permute(1, 2, 0) @ self.latent_rgb_factors
return preview_to_image(latent_image)
@ -66,27 +89,44 @@ def get_previewer(device, latent_format):
taesd_decoder_path = None
if latent_format.taesd_decoder_name is not None:
taesd_decoder_path = next(
(fn for fn in folder_paths.get_filename_list("vae_approx")
if fn.startswith(latent_format.taesd_decoder_name)),
""
(
fn
for fn in folder_paths.get_filename_list("vae_approx")
if fn.startswith(latent_format.taesd_decoder_name)
),
"",
)
taesd_decoder_path = folder_paths.get_full_path(
"vae_approx", taesd_decoder_path
)
taesd_decoder_path = folder_paths.get_full_path("vae_approx", taesd_decoder_path)
if method == LatentPreviewMethod.Auto:
method = LatentPreviewMethod.Latent2RGB
if method == LatentPreviewMethod.TAESD:
if taesd_decoder_path:
taesd = TAESD(None, taesd_decoder_path, latent_channels=latent_format.latent_channels).to(device)
taesd = TAESD(
None,
taesd_decoder_path,
latent_channels=latent_format.latent_channels,
).to(device)
previewer = TAESDPreviewerImpl(taesd)
else:
logging.warning("Warning: TAESD previews enabled, but could not find models/vae_approx/{}".format(latent_format.taesd_decoder_name))
logging.warning(
"Warning: TAESD previews enabled, but could not find models/vae_approx/{}".format(
latent_format.taesd_decoder_name
)
)
if previewer is None:
if latent_format.latent_rgb_factors is not None:
previewer = Latent2RGBPreviewer(latent_format.latent_rgb_factors, latent_format.latent_rgb_factors_bias)
previewer = Latent2RGBPreviewer(
latent_format.latent_rgb_factors,
latent_format.latent_rgb_factors_bias,
)
return previewer
def prepare_callback(model, steps, x0_output_dict=None):
preview_format = "JPEG"
if preview_format not in ["JPEG", "PNG"]:
@ -95,6 +135,7 @@ def prepare_callback(model, steps, x0_output_dict=None):
previewer = get_previewer(model.load_device, model.model.latent_format)
pbar = comfy.utils.ProgressBar(steps)
def callback(step, x0, x, total_steps):
if x0_output_dict is not None:
x0_output_dict["x0"] = x0
@ -103,5 +144,11 @@ def prepare_callback(model, steps, x0_output_dict=None):
if previewer:
preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0)
pbar.update_absolute(step + 1, total_steps, preview_bytes)
return callback
if args.xla or args.xla_fsdp:
import torch_xla as xla
import torch_xla.core.xla_model as xm
logging.info(f"Step {step + 1}/{total_steps}")
xla.sync()
return callback